From 8d8f469fadb695cfdd6309cf9afd5cf53dcbc831 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:42:23 -0700 Subject: [PATCH 1/7] [None][perf] Fuse MiniMax-M3 QKV and index projection Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 361 +++++++++++++ .../kernels/fusedQKNormRopeKernel.h | 23 +- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 228 ++++++++- docs/source/models/supported-models.md | 4 +- .../sparse/minimax_m3/cache_manager.py | 9 + .../sparse/minimax_m3/common.py | 15 +- .../sparse/minimax_m3/msa_backend.py | 205 +++++++- .../_torch/models/modeling_minimaxm3.py | 484 ++++++++++++++++-- tensorrt_llm/_torch/pyexecutor/_util.py | 50 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 12 + tensorrt_llm/_torch/speculative/eagle3.py | 15 +- tensorrt_llm/llmapi/llm_args.py | 21 +- .../defs/accuracy/references/gsm8k.yaml | 3 + .../defs/accuracy/references/mmlu.yaml | 3 + .../defs/accuracy/test_llm_api_pytorch.py | 93 +++- .../test_lists/qa/llm_function_core.txt | 2 + .../test_lists/test-db/l0_b200.yml | 3 + .../sparse/test_minimax_m3_msa_backend.py | 82 ++- .../unittest/_torch/models/test_minimax_m3.py | 98 ++++ ...test_minimax_m3_fp8_horizontal_producer.py | 279 ++++++++++ .../test_minimax_m3_fp8_main_kv_insert.py | 189 +++++++ 21 files changed, 2079 insertions(+), 100 deletions(-) create mode 100644 tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py create mode 100644 tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 2eaf4ba15fa6..8c7aa0a3e3e8 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -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 to a (possibly different-dtype) output buffer. // head_dim: the dimension of each head @@ -351,6 +366,295 @@ __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, int64_t numPages, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, float eps, + __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, float base, int const* positionIds) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * warpsPerBlock + warpId; + int const totalHeads = numHeadsQ + numHeadsK + numHeadsV; + int const tokenIdx = globalWarp / totalHeads; + int const localHead = globalWarp % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + + int const totalQKHeads = numHeadsQ + numHeadsK; + bool const isQ = localHead < numHeadsQ; + bool const isV = localHead >= totalQKHeads; + int const headIdx = isQ ? localHead : (isV ? localHead - totalQKHeads : localHead - numHeadsQ); + int64_t const inputOffset = (static_cast(tokenIdx) * totalHeads + localHead) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + + float elements[kMinimaxM3ElemsPerThread]; + float sumSquares = 0.0F; + constexpr int kVecSize = kMinimaxM3ElemsPerThread * sizeof(__nv_bfloat16) / 4; + using VecT = typename tensorrt_llm::common::packed_as::type; + VecT const packedInput = *reinterpret_cast(qkvInput + inputOffset); +#pragma unroll + for (int i = 0; i < kVecSize; ++i) + { + float2 const values = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packedInput) + i)); + if (!isV) + { + sumSquares += values.x * values.x; + sumSquares += values.y * values.y; + } + elements[2 * i] = values.x; + elements[2 * i + 1] = values.y; + } + + __nv_fp8_e4m3* output = qOutput; + int64_t outputOffset; + if (isQ) + { + outputOffset = (static_cast(tokenIdx) * numHeadsQ + headIdx) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + } + else + { + int slot = laneId == 0 ? outCacheLoc[tokenIdx] : 0; + slot = __shfl_sync(0xffffffff, slot, 0); + // CUDA-graph padding uses -1 for non-live cache destinations. Q is + // still produced for the padded row, but K/V must not address it. + if (slot < 0) + { + return; + } + int const page = slot >> 7; + if (page >= numPages) + { + return; + } + int const withinPage = slot & (kMinimaxM3PageSize - 1); + int const plane = isV ? 1 : 0; + output = kvCache; + outputOffset = static_cast(page) * pageStride + static_cast(plane) * planeStride + + static_cast(headIdx) * headStride + static_cast(withinPage) * tokenStride + + laneId * kMinimaxM3ElemsPerThread; + } + + // V is copy-cast only. + if (isV) + { + storeFp8HeadElements64(output, outputOffset, elements); + return; + } + + sumSquares = tensorrt_llm::common::warpReduceSum(sumSquares); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kMinimaxM3HeadDim) + eps); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + float const weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rmsReciprocal * (1.0F + weight); + } + + // MiniMax-M3 uses NeoX partial RoPE over the first 64 of 128 channels. + // Only lanes 0..7 calculate the 32 distinct angles; lanes 8..15 reuse + // them for the paired half, while lanes 16..31 bypass RoPE. + float pairedElements[kMinimaxM3ElemsPerThread]; + float cosineValues[kMinimaxM3ElemsPerThread] = {}; + float sineValues[kMinimaxM3ElemsPerThread] = {}; + __syncwarp(); + constexpr int kPairOffset = (kMinimaxM3RotaryDim / 2) / kMinimaxM3ElemsPerThread; + int positionId = laneId == 0 ? positionIds[tokenIdx] : 0; + positionId = __shfl_sync(0xffffffff, positionId, 0); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + pairedElements[i] = __shfl_xor_sync(0xffffffff, elements[i], kPairOffset); + if (laneId < kPairOffset) + { + pairedElements[i] = -pairedElements[i]; + } + + if (laneId < kPairOffset) + { + int const halfDim = dim; + float const frequency = powf(base, -2.0F * halfDim / static_cast(kMinimaxM3RotaryDim)); + __sincosf(static_cast(positionId) * frequency, &sineValues[i], &cosineValues[i]); + } + if (laneId < 2 * kPairOffset) + { + int const sourceLane = laneId % kPairOffset; + cosineValues[i] = __shfl_sync(0x0000ffff, cosineValues[i], sourceLane); + sineValues[i] = __shfl_sync(0x0000ffff, sineValues[i], sourceLane); + } + } + __syncwarp(); + +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + if (dim < kMinimaxM3RotaryDim) + { + elements[i] = elements[i] * cosineValues[i] + pairedElements[i] * sineValues[i]; + } + } + + storeFp8HeadElements64(output, outputOffset, elements); +} + +// Horizontal sparse producer for a packed [Q|K|V|index-Q|index-K] row. +// One warp owns one (token, head slot). All four norm/RoPE branches share the +// model's precomputed FP32 RoPE table, eliminating per-head powf/sincos work. +__global__ void minimaxM3Fp8QKVIndexerNormRopeKVInsertKernel(__nv_bfloat16 const* packedInput, __nv_fp8_e4m3* qOutput, + __nv_fp8_e4m3* indexQOutput, __nv_fp8_e4m3* kvCache, __nv_fp8_e4m3* indexKCache, int const* outCacheLoc, + int64_t kvPageStride, int64_t kvPlaneStride, int64_t kvHeadStride, int64_t kvTokenStride, int64_t indexPageStride, + int64_t indexTokenStride, int64_t numPages, int numTokens, int numHeadsQ, int numHeadsKV, int numHeadsIndex, + float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, __nv_bfloat16 const* indexQWeight, + __nv_bfloat16 const* indexKWeight, float const* rotaryCosSin, int const* positionIds) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * warpsPerBlock + warpId; + int const totalHeads = numHeadsQ + 2 * numHeadsKV + numHeadsIndex + 1; + int const tokenIdx = globalWarp / totalHeads; + int const localHead = globalWarp % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + + int const kBegin = numHeadsQ; + int const vBegin = kBegin + numHeadsKV; + int const indexQBegin = vBegin + numHeadsKV; + int const indexKHead = indexQBegin + numHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + int64_t const inputOffset = (static_cast(tokenIdx) * totalHeads + localHead) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + constexpr int kVecSize = kMinimaxM3ElemsPerThread * sizeof(__nv_bfloat16) / 4; + using VecT = typename tensorrt_llm::common::packed_as::type; + VecT const packed = *reinterpret_cast(packedInput + inputOffset); + + float elements[kMinimaxM3ElemsPerThread]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < kVecSize; ++pair) + { + float2 const values = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + elements[2 * pair] = values.x; + elements[2 * pair + 1] = values.y; + if (!isV) + { + sumSquares += values.x * values.x + values.y * values.y; + } + } + + if (!isV) + { + auto const* normWeight = isQ ? qWeight : (isK ? kWeight : (isIndexQ ? indexQWeight : indexKWeight)); + sumSquares = tensorrt_llm::common::warpReduceSum(sumSquares); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kMinimaxM3HeadDim) + eps); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + elements[i] *= rmsReciprocal * (1.0F + __bfloat162float(normWeight[dim])); + } + + __syncwarp(); + constexpr int kPairOffset = (kMinimaxM3RotaryDim / 2) / kMinimaxM3ElemsPerThread; + int positionId = laneId == 0 ? positionIds[tokenIdx] : 0; + positionId = __shfl_sync(0xffffffff, positionId, 0); + int64_t const ropeRow = static_cast(positionId) * kMinimaxM3RotaryDim; +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + float paired = __shfl_xor_sync(0xffffffff, elements[i], kPairOffset); + if (dim < kMinimaxM3RotaryDim) + { + bool const firstHalf = dim < kMinimaxM3RotaryDim / 2; + if (firstHalf) + { + paired = -paired; + } + int const coefficient = firstHalf ? dim : dim - kMinimaxM3RotaryDim / 2; + float const cosine = rotaryCosSin[ropeRow + coefficient]; + float const sine = rotaryCosSin[ropeRow + kMinimaxM3RotaryDim / 2 + coefficient]; + elements[i] = elements[i] * cosine + paired * sine; + } + } + __syncwarp(); + } + + if (isQ) + { + int const head = localHead; + int64_t const outputOffset = (static_cast(tokenIdx) * numHeadsQ + head) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(qOutput, outputOffset, elements); + return; + } + if (isIndexQ) + { + int const head = localHead - indexQBegin; + int64_t const outputOffset = (static_cast(tokenIdx) * numHeadsIndex + head) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + // Match vLLM's CUDA path: normalized/RoPE FP32 registers convert + // directly to saturating E4M3, without an intermediate BF16 round. + storeFp8HeadElements64(indexQOutput, outputOffset, elements); + return; + } + + int slot = laneId == 0 ? outCacheLoc[tokenIdx] : 0; + slot = __shfl_sync(0xffffffff, slot, 0); + if (slot < 0) + { + return; + } + int const page = slot >> 7; + if (page >= numPages) + { + return; + } + int const withinPage = slot & (kMinimaxM3PageSize - 1); + if (isIndexK) + { + int64_t const outputOffset = static_cast(page) * indexPageStride + + static_cast(withinPage) * indexTokenStride + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(indexKCache, outputOffset, elements); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int64_t const outputOffset = static_cast(page) * kvPageStride + static_cast(plane) * kvPlaneStride + + static_cast(head) * kvHeadStride + static_cast(withinPage) * kvTokenStride + + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(kvCache, outputOffset, elements); +} + +} // namespace + // Borrowed from // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568 #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ @@ -459,6 +763,63 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, void* qkv_out, int const num static_cast<__nv_bfloat16 const*>(k_weight), base, interleave, position_ids, factor, low, high, attention_factor, stream, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); } + +void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache, + int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, + int rotary_dim, float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(head_dim == kMinimaxM3HeadDim, "MiniMax-M3 FP8 main Q/K/V producer requires head_dim=128"); + TLLM_CHECK_WITH_INFO( + rotary_dim == kMinimaxM3RotaryDim, "MiniMax-M3 FP8 main Q/K/V producer requires rotary_dim=64"); + TLLM_CHECK_WITH_INFO(num_heads_q > 0, "MiniMax-M3 FP8 main Q/K/V producer requires query heads"); + TLLM_CHECK_WITH_INFO( + num_heads_k > 0 && num_heads_v > 0, "MiniMax-M3 FP8 main Q/K/V producer requires K and V heads"); + TLLM_CHECK_WITH_INFO(page_size == kMinimaxM3PageSize, "MiniMax-M3 FP8 main Q/K/V producer requires page_size=128"); + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int const totalWarps = num_tokens * (num_heads_q + num_heads_k + num_heads_v); + int const gridSize = common::divUp(totalWarps, kWarpsPerBlock); + minimaxM3Fp8QKNormRopeKVInsertKernel<<>>( + static_cast<__nv_bfloat16 const*>(qkv_input), static_cast<__nv_fp8_e4m3*>(q_output), + static_cast<__nv_fp8_e4m3*>(kv_cache), out_cache_loc, page_stride, plane_stride, head_stride, token_stride, + num_pages, num_tokens, num_heads_q, num_heads_k, num_heads_v, eps, static_cast<__nv_bfloat16 const*>(q_weight), + static_cast<__nv_bfloat16 const*>(k_weight), base, position_ids); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output, + void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride, + int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index, + int head_dim, int rotary_dim, float eps, void const* q_weight, void const* k_weight, void const* index_q_weight, + void const* index_k_weight, float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(head_dim == kMinimaxM3HeadDim, "MiniMax-M3 horizontal producer requires head_dim=128"); + TLLM_CHECK_WITH_INFO(rotary_dim == kMinimaxM3RotaryDim, "MiniMax-M3 horizontal producer requires rotary_dim=64"); + TLLM_CHECK_WITH_INFO(page_size == kMinimaxM3PageSize, "MiniMax-M3 horizontal producer requires page_size=128"); + TLLM_CHECK_WITH_INFO(num_heads_q > 0 && num_heads_kv > 0 && num_heads_index > 0, + "MiniMax-M3 horizontal producer requires Q, KV, and index heads"); + TLLM_CHECK_WITH_INFO( + num_heads_index == num_heads_kv, "MiniMax-M3 horizontal producer requires index heads to equal KV heads"); + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int const slotsPerToken = num_heads_q + 2 * num_heads_kv + num_heads_index + 1; + int const totalWarps = num_tokens * slotsPerToken; + int const gridSize = common::divUp(totalWarps, kWarpsPerBlock); + minimaxM3Fp8QKVIndexerNormRopeKVInsertKernel<<>>( + static_cast<__nv_bfloat16 const*>(packed_input), static_cast<__nv_fp8_e4m3*>(q_output), + static_cast<__nv_fp8_e4m3*>(index_q_output), static_cast<__nv_fp8_e4m3*>(kv_cache), + static_cast<__nv_fp8_e4m3*>(index_k_cache), out_cache_loc, kv_page_stride, kv_plane_stride, kv_head_stride, + kv_token_stride, index_page_stride, index_token_stride, num_pages, num_tokens, num_heads_q, num_heads_kv, + num_heads_index, eps, static_cast<__nv_bfloat16 const*>(q_weight), static_cast<__nv_bfloat16 const*>(k_weight), + static_cast<__nv_bfloat16 const*>(index_q_weight), static_cast<__nv_bfloat16 const*>(index_k_weight), + rotary_cos_sin, position_ids); + TLLM_CUDA_CHECK(cudaGetLastError()); +} } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index fd2401f592f1..4b41467d457c 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,8 @@ #pragma once #include "tensorrt_llm/common/config.h" + +#include #include TRTLLM_NAMESPACE_BEGIN @@ -61,6 +63,25 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, // BF16 input [num_tokens, t bool const interleave, int const* position_ids, float factor, float low, float high, float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2); +// MiniMax-M3-specific main-branch producer. It returns contiguous +// FP8 Q and inserts normalized/RoPE'd FP8 K plus copy-cast FP8 V directly into +// a paged HND pool [num_pages, 2, num_heads, page_size, head_dim]. +void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache, + int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, + int rotary_dim, float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids, + cudaStream_t stream); + +// MiniMax-M3 sparse producer for the packed [Q|K|V|index-Q|index-K] +// projection. It uses a precomputed FP32 RoPE table, emits compact FP8 Q and +// index-Q, and inserts main K/V plus index-K into their paged FP8 HND caches. +void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output, + void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride, + int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index, + int head_dim, int rotary_dim, float eps, void const* q_weight, void const* k_weight, void const* index_q_weight, + void const* index_k_weight, float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 12c124b85aaf..fe439d287132 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,9 +18,12 @@ #include "tensorrt_llm/thop/thUtils.h" #include +#include #include #include +#include + TRTLLM_NAMESPACE_BEGIN namespace torch_ext @@ -65,6 +68,35 @@ int64_t validateFusedQKNormRopeInputs(torch::Tensor const& qkv, torch::Tensor co return num_tokens; } +void checkMinimaxM3HndKVPool(torch::Tensor const& kvCache, int64_t numHeads, int64_t headDim) +{ + TORCH_CHECK(kvCache.is_cuda(), "kv_cache must be a CUDA tensor"); + TORCH_CHECK(kvCache.scalar_type() == at::ScalarType::Float8_e4m3fn, "kv_cache must use torch.float8_e4m3fn"); + TORCH_CHECK(kvCache.dim() == 5, "kv_cache must be HND [num_pages, 2, num_heads, page_size, head_dim]"); + TORCH_CHECK(kvCache.size(0) > 0 && kvCache.size(3) > 0, "kv_cache must have positive num_pages and page_size"); + TORCH_CHECK(kvCache.size(1) == 2, "kv_cache plane dimension must contain K and V"); + TORCH_CHECK(kvCache.size(2) == numHeads, "kv_cache num_heads mismatch"); + TORCH_CHECK(kvCache.size(4) == headDim, "kv_cache head_dim mismatch"); + TORCH_CHECK(kvCache.stride(4) == 1 && kvCache.stride(3) == headDim, + "kv_cache must have contiguous head_dim rows in HND layout"); + TORCH_CHECK(kvCache.stride(2) == kvCache.size(3) * kvCache.stride(3), + "kv_cache must have contiguous [page_size, head_dim] blocks in HND layout"); + TORCH_CHECK(kvCache.stride(1) >= kvCache.size(2) * kvCache.stride(2), "kv_cache K and V planes must not overlap"); + TORCH_CHECK(kvCache.stride(0) >= kvCache.size(1) * kvCache.stride(1), + "kv_cache page stride must not overlap adjacent HND pages"); + TORCH_CHECK(kvCache.stride(0) % 4 == 0 && kvCache.stride(1) % 4 == 0, + "kv_cache page and plane strides must preserve 32-bit FP8 store alignment"); +} + +void checkMinimaxM3Int32LaunchGeometry(int64_t numTokens, int64_t slotsPerToken) +{ + TORCH_CHECK(numTokens <= std::numeric_limits::max(), "MiniMax-M3 producer num_tokens exceeds int32"); + TORCH_CHECK(slotsPerToken > 0 && slotsPerToken <= std::numeric_limits::max(), + "MiniMax-M3 producer head geometry exceeds int32"); + TORCH_CHECK(numTokens == 0 || slotsPerToken <= std::numeric_limits::max() / numTokens, + "MiniMax-M3 producer launch geometry exceeds int32"); +} + } // namespace // Function for fused QK Norm and RoPE @@ -150,6 +182,187 @@ torch::Tensor fused_qk_norm_rope_to_fp8_meta(torch::Tensor const& qkv, int64_t n return torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn)); } +torch::Tensor minimaxM3Fp8QKNormRopeKVInsert(torch::Tensor const& qkv, torch::Tensor& kvCache, + torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t numHeadsK, int64_t numHeadsV, int64_t headDim, + int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, double base, bool isNeox, + torch::Tensor const& positionIds) +{ + constexpr int64_t kHeadDim = 128; + constexpr int64_t kRotaryDim = 64; + constexpr int64_t kPageSize = 128; + TORCH_CHECK(numHeadsQ > 0, "MiniMax-M3 FP8 main Q/K/V producer requires num_heads_q > 0"); + TORCH_CHECK(numHeadsK > 0 && numHeadsV > 0, "MiniMax-M3 FP8 main Q/K/V producer requires K and V heads"); + TORCH_CHECK(numHeadsK == numHeadsV, "MiniMax-M3 FP8 main Q/K/V producer requires equal K and V head counts"); + TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 FP8 main Q/K/V producer requires head_dim=128"); + TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 FP8 main Q/K/V producer requires rotary_dim=64"); + TORCH_CHECK(isNeox, "MiniMax-M3 FP8 main Q/K/V producer requires NeoX RoPE"); + TORCH_CHECK(std::isfinite(eps) && eps > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires finite eps > 0"); + TORCH_CHECK(std::isfinite(base) && base > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires finite RoPE base > 0"); + auto const epsFloat = static_cast(eps); + auto const baseFloat = static_cast(base); + TORCH_CHECK(std::isfinite(epsFloat) && epsFloat > 0.0F, + "MiniMax-M3 FP8 main Q/K/V producer eps must remain finite and positive in float32"); + TORCH_CHECK(std::isfinite(baseFloat) && baseFloat > 0.0F, + "MiniMax-M3 FP8 main Q/K/V producer RoPE base must remain finite and positive in float32"); + + TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); + TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional"); + TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional"); + TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1, "Q/K norm weights must be one-dimensional"); + + CHECK_INPUT(qkv, torch::kBFloat16); + CHECK_INPUT(outCacheLoc, torch::kInt32); + CHECK_INPUT(positionIds, torch::kInt32); + CHECK_INPUT(qWeight, torch::kBFloat16); + CHECK_INPUT(kWeight, torch::kBFloat16); + checkMinimaxM3HndKVPool(kvCache, numHeadsK, headDim); + TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 FP8 main Q/K/V producer requires page_size=128"); + + int64_t const numTokens = qkv.size(0); + int64_t const totalHeads = numHeadsQ + numHeadsK + numHeadsV; + checkMinimaxM3Int32LaunchGeometry(numTokens, totalHeads); + TORCH_CHECK(qkv.size(1) == totalHeads * headDim, + "QKV tensor width must equal (num_heads_q + num_heads_k + num_heads_v) * head_dim"); + TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens"); + TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens"); + TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim, "Q/K norm weight width must equal head_dim"); + TORCH_CHECK(reinterpret_cast(qkv.data_ptr()) % 8 == 0, + "QKV input must start at an 8-byte-aligned address for vectorized BF16 loads"); + TORCH_CHECK(reinterpret_cast(kvCache.data_ptr()) % 4 == 0, + "K/V cache must start at a 4-byte-aligned address for packed E4M3 stores"); + TORCH_CHECK(qkv.get_device() == kvCache.get_device() && qkv.get_device() == outCacheLoc.get_device() + && qkv.get_device() == positionIds.get_device() && qkv.get_device() == qWeight.get_device() + && qkv.get_device() == kWeight.get_device(), + "All MiniMax-M3 FP8 main Q/K/V producer tensors must be on the same CUDA device"); + + auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn)); + if (numTokens == 0) + { + return qOut; + } + + auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + tensorrt_llm::kernels::launchMinimaxM3Fp8QKNormRopeKVInsert(qkv.data_ptr(), qOut.data_ptr(), kvCache.data_ptr(), + outCacheLoc.data_ptr(), kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3), + kvCache.size(0), static_cast(kvCache.size(3)), static_cast(numTokens), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(headDim), + static_cast(rotaryDim), epsFloat, qWeight.data_ptr(), kWeight.data_ptr(), baseFloat, + positionIds.data_ptr(), stream); + return qOut; +} + +torch::Tensor minimaxM3Fp8QKNormRopeKVInsertMeta(torch::Tensor const& qkv, torch::Tensor& /*kvCache*/, + torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, int64_t /*numHeadsK*/, int64_t /*numHeadsV*/, + int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, torch::Tensor const& /*qWeight*/, + torch::Tensor const& /*kWeight*/, double /*base*/, bool /*isNeox*/, torch::Tensor const& /*positionIds*/) +{ + return torch::empty({qkv.size(0), numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn)); +} + +std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsert(torch::Tensor const& packed, + torch::Tensor& kvCache, torch::Tensor& indexKCache, torch::Tensor const& outCacheLoc, int64_t numHeadsQ, + int64_t numHeadsKV, int64_t numHeadsIndex, int64_t headDim, int64_t rotaryDim, double eps, + torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds) +{ + constexpr int64_t kHeadDim = 128; + constexpr int64_t kRotaryDim = 64; + constexpr int64_t kPageSize = 128; + TORCH_CHECK(numHeadsQ > 0 && numHeadsKV > 0 && numHeadsIndex > 0, + "MiniMax-M3 horizontal producer requires Q, KV, and index heads"); + TORCH_CHECK(numHeadsKV == numHeadsIndex, "MiniMax-M3 horizontal producer requires index heads to equal KV heads"); + TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 horizontal producer requires head_dim=128"); + TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 horizontal producer requires rotary_dim=64"); + TORCH_CHECK(std::isfinite(eps) && eps > 0.0, "MiniMax-M3 horizontal producer requires finite eps > 0"); + auto const epsFloat = static_cast(eps); + TORCH_CHECK(std::isfinite(epsFloat) && epsFloat > 0.0F, + "MiniMax-M3 horizontal producer eps must remain finite and positive in float32"); + + TORCH_CHECK(packed.dim() == 2, "Packed QKV+index tensor must be two-dimensional"); + TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional"); + TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional"); + CHECK_INPUT(packed, torch::kBFloat16); + CHECK_INPUT(outCacheLoc, torch::kInt32); + CHECK_INPUT(positionIds, torch::kInt32); + CHECK_INPUT(qWeight, torch::kBFloat16); + CHECK_INPUT(kWeight, torch::kBFloat16); + CHECK_INPUT(indexQWeight, torch::kBFloat16); + CHECK_INPUT(indexKWeight, torch::kBFloat16); + CHECK_INPUT(rotaryCosSin, torch::kFloat32); + checkMinimaxM3HndKVPool(kvCache, numHeadsKV, headDim); + TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 horizontal producer requires page_size=128"); + TORCH_CHECK(indexKCache.is_cuda() && indexKCache.scalar_type() == at::ScalarType::Float8_e4m3fn, + "Index-K cache must be CUDA torch.float8_e4m3fn"); + TORCH_CHECK(indexKCache.dim() == 4 && indexKCache.size(1) == 1 && indexKCache.size(2) == kPageSize + && indexKCache.size(3) == kHeadDim, + "Index-K cache must be HND [num_pages, 1, 128, 128]"); + TORCH_CHECK(indexKCache.stride(3) == 1 && indexKCache.stride(2) == kHeadDim, + "Index-K cache must have contiguous token rows"); + TORCH_CHECK(indexKCache.stride(1) >= indexKCache.size(2) * indexKCache.stride(2) + && indexKCache.stride(0) >= indexKCache.size(1) * indexKCache.stride(1), + "Index-K cache pages must not overlap"); + TORCH_CHECK(indexKCache.stride(0) % 4 == 0 && indexKCache.stride(1) % 4 == 0, + "Index-K cache page/head strides must preserve 32-bit FP8 store alignment"); + TORCH_CHECK( + indexKCache.size(0) == kvCache.size(0), "Main K/V and index-K caches must contain the same number of pages"); + TORCH_CHECK(rotaryCosSin.dim() == 3 && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == kRotaryDim / 2, + "rotary_cos_sin must be [max_positions, 2, rotary_dim/2]"); + + int64_t const numTokens = packed.size(0); + int64_t const totalHeads = numHeadsQ + 2 * numHeadsKV + numHeadsIndex + 1; + checkMinimaxM3Int32LaunchGeometry(numTokens, totalHeads); + TORCH_CHECK( + packed.size(1) == totalHeads * headDim, "Packed tensor width must equal (Q + 2*KV + index-Q + 1) * head_dim"); + TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens"); + TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens"); + TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim && indexQWeight.numel() == headDim + && indexKWeight.numel() == headDim, + "All norm weights must contain head_dim elements"); + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0, + "Packed input must start at an 8-byte-aligned address for vectorized BF16 loads"); + TORCH_CHECK(reinterpret_cast(kvCache.data_ptr()) % 4 == 0 + && reinterpret_cast(indexKCache.data_ptr()) % 4 == 0, + "Paged caches must start at 4-byte-aligned addresses for packed E4M3 stores"); + TORCH_CHECK(packed.get_device() == kvCache.get_device() && packed.get_device() == indexKCache.get_device() + && packed.get_device() == outCacheLoc.get_device() && packed.get_device() == positionIds.get_device() + && packed.get_device() == qWeight.get_device() && packed.get_device() == kWeight.get_device() + && packed.get_device() == indexQWeight.get_device() && packed.get_device() == indexKWeight.get_device() + && packed.get_device() == rotaryCosSin.get_device(), + "All MiniMax-M3 horizontal producer tensors must be on the same CUDA device"); + + auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn)); + auto indexQOut + = torch::empty({numTokens, numHeadsIndex, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn)); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + + auto stream = at::cuda::getCurrentCUDAStream(packed.get_device()); + tensorrt_llm::kernels::launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(packed.data_ptr(), qOut.data_ptr(), + indexQOut.data_ptr(), kvCache.data_ptr(), indexKCache.data_ptr(), outCacheLoc.data_ptr(), + kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3), indexKCache.stride(0), + indexKCache.stride(2), kvCache.size(0), static_cast(kvCache.size(3)), static_cast(numTokens), + static_cast(numHeadsQ), static_cast(numHeadsKV), static_cast(numHeadsIndex), + static_cast(headDim), static_cast(rotaryDim), epsFloat, qWeight.data_ptr(), kWeight.data_ptr(), + indexQWeight.data_ptr(), indexKWeight.data_ptr(), rotaryCosSin.data_ptr(), positionIds.data_ptr(), + stream); + return {qOut, indexQOut}; +} + +std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta(torch::Tensor const& packed, + torch::Tensor& /*kvCache*/, torch::Tensor& /*indexKCache*/, torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, + int64_t /*numHeadsKV*/, int64_t numHeadsIndex, int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, + torch::Tensor const& /*qWeight*/, torch::Tensor const& /*kWeight*/, torch::Tensor const& /*indexQWeight*/, + torch::Tensor const& /*indexKWeight*/, torch::Tensor const& /*rotaryCosSin*/, torch::Tensor const& /*positionIds*/) +{ + auto options = packed.options().dtype(at::ScalarType::Float8_e4m3fn); + return { + torch::empty({packed.size(0), numHeadsQ, headDim}, options), + torch::empty({packed.size(0), numHeadsIndex, headDim}, options), + }; +} + // Register the PyTorch operators TORCH_LIBRARY_FRAGMENT(trtllm, m) { @@ -164,6 +377,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float " "factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int " "mrope_section1, int mrope_section2) -> Tensor"); + m.def( + "minimax_m3_fp8_qk_norm_rope_kv_insert(Tensor qkv, Tensor(a!) kv_cache, Tensor out_cache_loc, int " + "num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim, float eps, Tensor q_weight, " + "Tensor k_weight, float base, bool is_neox, Tensor position_ids) -> Tensor"); + m.def( + "minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert(Tensor packed, Tensor(a!) kv_cache, Tensor(b!) " + "index_k_cache, Tensor out_cache_loc, int num_heads_q, int num_heads_kv, int num_heads_index, int head_dim, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, Tensor index_q_weight, Tensor index_k_weight, " + "Tensor rotary_cos_sin, Tensor position_ids) -> (Tensor, Tensor)"); } // Register the CUDA implementation @@ -171,12 +393,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/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 7e8fb78e6e3a..9b6dc9315cac 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -84,7 +84,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | Yes | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Gemma4UnifiedForConditionalGeneration` | Untested | Untested | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | -| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | +| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | Yes | No | No | Yes | Untested | No | N/A | Untested | Untested | [^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121 and in BF16/FP8 KV cache dtype. @@ -96,7 +96,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^9]: Audio modality only supported on E2B/E4B variants. [^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path. [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism. -[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. +[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. One-model linear EAGLE-3 is supported; combining it with CUDA graphs requires the MSA implementation on SM100. [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models). [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`). diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index 89a7595c52c9..00cc0c29995e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -150,6 +150,15 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): * ``sparse_index_dim`` — width of the index-K/V vectors. """ + # INDEX_KEY is coalesced into the target's V2 pool. Dense Eagle3 draft + # layers cannot consume the synthetic AttentionOp view of that layout, so + # they require their own ordinary KV manager even under attention DP. + supports_shared_draft_layers = False + + # MSA requires 128-token target pages. The dense Eagle3 generation kernel + # uses its validated 32-token page geometry in the separate draft manager. + draft_manager_tokens_per_block = 32 + def __init__( self, *args, 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 ea5ad1f7ee92..18ffbd8a194c 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 66871d0c8f7a..9fb76118a980 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 @@ -281,6 +281,11 @@ def __post_init__(self) -> None: super().__post_init__() params = self.sparse_metadata_params self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None + # Live geometry used to repair optimistic overlap-scheduler lengths. + self._msa_live_batch = 0 + self._msa_live_total_q = 0 + self._msa_page_size = 0 + self._msa_corrected_kv_lens_cpu: Optional[torch.Tensor] = None self._create_msa_buffers() @property @@ -294,11 +299,22 @@ def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: @property def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: - """Per-request KV length, cached plus new tokens (host int32).""" + """Per-request true attended KV length (host int32). + + The base ``kv_lens`` includes speculative draft-loop reserve slots. + Those slots are consumed by the generic executor but must not enter + MSA plans, cache-slot ladders, or page counts. + """ + if self._msa_corrected_kv_lens_cpu is not None: + return self._msa_corrected_kv_lens_cpu kv_lens = getattr(self, "kv_lens", None) if self.seq_lens is None or kv_lens is None: return None out = kv_lens[: self.num_seqs] + params = self.kv_cache_params + num_extra_kv_tokens = params.num_extra_kv_tokens if params is not None else 0 + if num_extra_kv_tokens: + out = out - num_extra_kv_tokens return out if out.dtype == torch.int32 else out.to(torch.int32) @property @@ -392,6 +408,39 @@ def _create_msa_buffers(self) -> None: dtype=torch.int32, capture_graph=capture_graph, ) + # Device staging used by on_update_kv_lens. The overlap scheduler can + # replace optimistic speculative lengths after prepare(), so cache + # slots and per-token causal bounds must be derivable without a host + # synchronization and without allocating inside CUDA graph replay. + tokens_per_block = int(kv_cache_manager.tokens_per_block) + self.msa_req_to_token = self.get_empty( + buffers, + (max_num_sequences, max_blocks_per_seq * tokens_per_block), + cache_name="msa_req_to_token", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_q_batch_row = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_q_batch_row", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_q_intra = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_q_intra", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_qo_lens_dev = self.get_empty( + buffers, + (max_num_sequences,), + cache_name="msa_qo_lens_dev", + dtype=torch.int32, + capture_graph=capture_graph, + ) # The proxy scratch needs the fmha_sm100 plan geometry. This metadata # exists only for the MSA backend, whose selection already required the # kernels, so a failed import here is a hard error rather than a reason @@ -399,45 +448,56 @@ 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, - max_batch=max_num_sequences, + num_index_heads=num_index_heads, + max_tokens=self._msa_max_decode_tokens(), max_k_tiles=max_k_tiles, capture_graph=capture_graph, ) self._msa_buffers_ready = True + def _msa_max_decode_tokens(self) -> int: + """Worst-case query-token count for a speculative verify step.""" + max_num_sequences = int(getattr(self, "max_num_sequences", 0) or 0) + max_num_tokens = int(getattr(self, "max_num_tokens", 0) or 0) + if max_num_tokens <= 0: + return max_num_sequences + # fmha_sm100 caps total_q * sharded query heads at 65536. M3 has four + # sharded index heads at TP4, so 16384 is the largest useful bound. + return max(max_num_sequences, min(max_num_tokens, 16384)) + def _alloc_msa_proxy_scratch( self, *, num_index_heads: int, - max_batch: int, + max_tokens: int, max_k_tiles: int, capture_graph: bool, ) -> None: """Allocate the flat proxy max-score store and the valid-block scratch. - The store is sized for the worst-case max_k_tiles so one allocation - serves every decode step. msa_proxy_max_score_view slices the per-step - shape out of it. + The store is sized for the worst-case max_k_tiles and query-token + count. Speculative verification contributes more than one query token + per request. """ buffers = self.cuda_graph_buffers self.msa_max_score = self.get_empty( buffers, - (num_index_heads * max_k_tiles * max_batch,), + (num_index_heads * max_k_tiles * max_tokens,), cache_name="msa_max_score", dtype=torch.float32, capture_graph=capture_graph, ) self.msa_n_valid_blocks = self.get_empty( buffers, - (max_batch,), + (max_tokens,), cache_name="msa_n_valid_blocks", dtype=torch.int32, capture_graph=capture_graph, @@ -452,14 +512,15 @@ def _ensure_msa_decode_scratch_buffers( required_max_k_tiles: int, ) -> None: """Ensure proxy scratch buffers exist and cover the current plan.""" - required_numel = num_index_heads * required_max_k_tiles * max_batch + max_tokens = max(int(max_batch), self._msa_max_decode_tokens()) + required_numel = num_index_heads * required_max_k_tiles * max_tokens if self.msa_max_score is not None: if self.msa_max_score.numel() < required_numel: raise ValueError( f"msa_max_score backing store ({self.msa_max_score.numel()} " f"elements) is smaller than the decode plan needs " f"({required_numel} = {num_index_heads} heads * " - f"{required_max_k_tiles} k-tiles * {max_batch} batch)." + f"{required_max_k_tiles} k-tiles * {max_tokens} tokens)." ) return @@ -481,7 +542,7 @@ def _ensure_msa_decode_scratch_buffers( ) self._alloc_msa_proxy_scratch( num_index_heads=num_index_heads, - max_batch=max_batch, + max_tokens=max_tokens, max_k_tiles=max_k_tiles, capture_graph=capture_graph, ) @@ -502,9 +563,78 @@ def _ensure_eager_n_valid_buffer(self, total_q: int, device: torch.device) -> to def prepare(self) -> None: super().prepare() + self._msa_corrected_kv_lens_cpu = None self._build_msa_fields() self._build_step_plans() + def on_update_kv_lens(self) -> None: + """Repair MSA state after overlap corrects speculative KV lengths. + + Pure-decode CUDA-graph steps use only device operations. Mixed eager + steps can safely copy the corrected lengths to the host and rebuild + their plans. The correction can only shrink the optimistic lengths, + so the decode worklist allocation and page-table capacity remain valid. + """ + super().on_update_kv_lens() + if not self._msa_fields_ready: + return + + batch_size = self._msa_live_batch + total_q = self._msa_live_total_q + if batch_size <= 0 or total_q <= 0: + return + + if self.msa_decode_proxy_plan is None: + if torch.cuda.is_current_stream_capturing(): + return + self._msa_corrected_kv_lens_cpu = self.kv_lens_cuda[:batch_size].to("cpu", torch.int32) + self._build_msa_fields() + self._build_step_plans() + return + + kv_lens = self.kv_lens_cuda[:batch_size] + q_batch_row = self.msa_q_batch_row[:total_q].to(torch.long) + qo_lens = self.msa_qo_lens_dev[:batch_size] + token_kv_lens = kv_lens[q_batch_row] + q_positions = token_kv_lens - qo_lens[q_batch_row] + self.msa_q_intra[:total_q] + + table_width = int(self.msa_req_to_token.shape[1]) + table_indices = q_positions.to(torch.long).clamp(min=0, max=table_width - 1) + cache_slots = self.msa_req_to_token.reshape(-1).index_select( + 0, q_batch_row * table_width + table_indices + ) + self.msa_out_cache_loc[:total_q].copy_(cache_slots) + + page_size = self._msa_page_size + n_valid_blocks = torch.div( + (q_positions + 1).clamp_min(1) + (page_size - 1), + page_size, + rounding_mode="floor", + ) + self.msa_n_valid_blocks[:total_q].copy_(n_valid_blocks.to(torch.int32)) + + request_offsets = (kv_lens - qo_lens).clamp_min(0) + for owner, expand_per_token in ( + (self._msa_proxy_plan, False), + (self._msa_gqa_plan, True), + (self._msa_dense_plan, False), + ): + if owner is None or owner.plan is None: + continue + decode_plan = owner.plan[3] + segment_lens = decode_plan.get("kv_segment_lens") + qo_offset = decode_plan.get("qo_offset") + if expand_per_token: + if segment_lens is not None: + segment_lens[:total_q].copy_(token_kv_lens.to(segment_lens.dtype)) + if qo_offset is not None: + qo_offset[:total_q].copy_(q_positions.clamp_min(0).to(qo_offset.dtype)) + else: + if segment_lens is not None: + segment_lens[:batch_size].copy_(kv_lens.to(segment_lens.dtype)) + if qo_offset is not None: + qo_offset[:batch_size].copy_(request_offsets.to(qo_offset.dtype)) + def _build_step_plans(self) -> None: """Build the three layer-invariant fmha_sm100 plans once per step. @@ -536,7 +666,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 @@ -546,7 +676,6 @@ def _build_step_plans(self) -> None: qo_offset_cpu = self.msa_qo_offset_cpu if qo_lens_cpu is None or kv_lens_cpu is None or qo_offset_cpu is None: return - batch = int(qo_lens_cpu.shape[0]) device = _cache_device(self) page_size = int(self.kv_cache_manager.tokens_per_block) capture_graph = self.is_cuda_graph @@ -627,27 +756,30 @@ def _build_step_plans(self) -> None: ) # Allocate the graph-safe plan owners once per metadata; later steps - # only refresh their contents below. + # only refresh their contents below. Speculative verification expands + # a request into one planner row per query token, so size the owners by + # the worst-case decode token count rather than request batch alone. if self._msa_proxy_plan is None: + max_plan_rows = max(max_batch, self._msa_max_decode_tokens()) num_ctas = torch.cuda.get_device_properties(device).multi_processor_count self._msa_proxy_plan = _MsaGraphSafePlan( self, "msa_proxy_plan", - max_batch=max_batch, + max_batch=max_plan_rows, num_ctas=num_ctas, capture_graph=capture_graph, ) self._msa_gqa_plan = _MsaGraphSafePlan( self, "msa_gqa_plan", - max_batch=max_batch, + max_batch=max_plan_rows, num_ctas=num_ctas, capture_graph=capture_graph, ) self._msa_dense_plan = _MsaGraphSafePlan( self, "msa_dense_plan", - max_batch=max_batch, + max_batch=max_plan_rows, num_ctas=num_ctas, capture_graph=capture_graph, ) @@ -661,7 +793,8 @@ def _build_step_plans(self) -> None: n_valid = per_token_valid_blocks( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) - self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True) + total_q = int(n_valid.shape[0]) + self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True) def _build_msa_fields(self) -> None: """Populate the MSA cache-write buffers for this step. @@ -687,14 +820,6 @@ def _build_msa_fields(self) -> None: cache_device = _cache_device(self) page_size = int(kv_cache_manager.tokens_per_block) - is_prefill = int(self.num_contexts or 0) > 0 - if not is_prefill and int(qo_lens_cpu.max().item()) > 1: - raise NotImplementedError( - "MiniMax-M3 MSA attention does not support speculative decoding " - "(multiple query tokens per decode step). Disable speculative " - "decoding or use the non-MSA MiniMax-M3 backend." - ) - # Built in prepare() (outside capture), so these transients are # fine: forwards read only the persistent buffers filled below. # qo_offset is the prefix length, so one build covers prefill @@ -723,6 +848,27 @@ def _build_msa_fields(self) -> None: self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) self.msa_kv_indices[:total_pages].copy_(kv_indices, non_blocking=True) + + # Keep the device-side geometry needed to repair optimistic + # speculative lengths after the overlap scheduler reports the actual + # accepted-token counts. + step_width = int(req_to_token.shape[1]) + self.msa_req_to_token[:batch_size, :step_width].copy_(req_to_token, non_blocking=True) + qo_lens_long = qo_lens_cpu.to(torch.long) + batch_rows = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32), qo_lens_long + ) + starts = torch.cumsum(qo_lens_long, 0) - qo_lens_long + intra = ( + torch.arange(total_new_tokens, dtype=torch.int64) + - torch.repeat_interleave(starts, qo_lens_long) + ).to(torch.int32) + self.msa_q_batch_row[:total_new_tokens].copy_(batch_rows) + self.msa_q_intra[:total_new_tokens].copy_(intra) + self.msa_qo_lens_dev[:batch_size].copy_(qo_lens_cpu) + self._msa_live_batch = batch_size + self._msa_live_total_q = total_new_tokens + self._msa_page_size = page_size self._msa_fields_ready = True def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @@ -876,6 +1022,9 @@ def run_indexer( ) idx_k_view = idx_k.view(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) # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized # index Q/K are cast directly and the proxy accumulates their QK scores # in FP32. Block ordering is invariant to the omitted positive scale. diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index ec226cce48a2..0c0b80860416 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -64,6 +64,7 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph +from ..speculative import SpecMetadata from ..utils import ( ActivationType, AuxStreamType, @@ -73,19 +74,178 @@ ) from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper -from .modeling_utils import ( - DecoderModel, - DecoderModelForCausalLM, - ModelConfig, - filter_weights, - register_auto_model, -) +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, ModelConfig, filter_weights, register_auto_model # Dense layers use SDPA with non-contiguous Q/K/V and a bool attn_mask. # Limit backends to memory-efficient and math; cuDNN SDPA fails for this layout, # and flash SDPA does not accept attn_mask. _DENSE_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH] + +class MiniMaxM3QKVIndexerLinear(Linear): + """Five-way MiniMax-M3 projection with vLLM-compatible TP sharding. + + Each rank emits ``[Q | K | V | index-Q | index-K]``. Q follows normal + attention head sharding, K/V/index-Q follow KV-head sharding (including + replication when TP exceeds the KV-head count), and the single index-K + head is replicated. The underlying :class:`Linear` remains the standard + quantized implementation; only checkpoint packing is model-specific. + """ + + _SHARD_NAMES = ("q", "k", "v", "index_q", "index_k") + + def __init__( + self, + *, + hidden_size: int, + head_dim: int, + total_num_heads: int, + total_num_kv_heads: int, + total_num_index_heads: int, + index_head_dim: int, + dtype: torch.dtype, + mapping: Mapping, + quant_config: Optional[QuantConfig], + skip_create_weights_in_init: bool, + force_dynamic_quantization: bool, + disable_deep_gemm: bool, + use_custom_cublas_mm: bool, + use_cute_dsl_bf16_gemm: bool, + use_cute_dsl_blockscaling_mm: bool, + ) -> None: + if total_num_index_heads != total_num_kv_heads: + raise ValueError( + "MiniMax-M3 fused QKV+index projection requires index heads " + f"({total_num_index_heads}) to equal KV heads ({total_num_kv_heads})." + ) + if index_head_dim != head_dim: + raise ValueError( + "MiniMax-M3 fused QKV+index projection requires index_head_dim " + f"({index_head_dim}) to equal head_dim ({head_dim})." + ) + + tp_size = int(mapping.tp_size) + if total_num_heads % tp_size != 0: + raise ValueError(f"Q heads ({total_num_heads}) must be divisible by TP ({tp_size}).") + if total_num_kv_heads >= tp_size: + if total_num_kv_heads % tp_size != 0: + raise ValueError( + f"KV heads ({total_num_kv_heads}) must be divisible by TP ({tp_size})." + ) + local_num_kv_heads = total_num_kv_heads // tp_size + else: + if tp_size % total_num_kv_heads != 0: + raise ValueError( + f"TP ({tp_size}) must be divisible by KV heads " + f"({total_num_kv_heads}) for replication." + ) + local_num_kv_heads = 1 + + self.total_num_heads = int(total_num_heads) + self.total_num_kv_heads = int(total_num_kv_heads) + self.total_num_index_heads = int(total_num_index_heads) + self.head_dim = int(head_dim) + self.index_head_dim = int(index_head_dim) + self.local_num_heads = total_num_heads // tp_size + self.local_num_kv_heads = local_num_kv_heads + self.local_num_index_heads = local_num_kv_heads + self.local_output_sizes = ( + self.local_num_heads * head_dim, + local_num_kv_heads * head_dim, + local_num_kv_heads * head_dim, + local_num_kv_heads * index_head_dim, + index_head_dim, + ) + local_out_features = sum(self.local_output_sizes) + + super().__init__( + hidden_size, + tp_size * local_out_features, + bias=False, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=quant_config, + weights_loading_config=WeightsLoadingConfig(weight_mode=WeightMode.FUSED_QKV_LINEAR), + reduce_output=False, + skip_create_weights_in_init=skip_create_weights_in_init, + force_dynamic_quantization=force_dynamic_quantization, + disable_deep_gemm=disable_deep_gemm, + use_custom_cublas_mm=use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + ) + + def _shard_geometry(self, shard_name: str) -> Tuple[int, int]: + """Return effective (world size, rank) for one checkpoint shard.""" + if shard_name == "q": + return self.tp_size, self.tp_rank + if shard_name == "index_k": + return 1, 0 + + total_heads = ( + self.total_num_index_heads if shard_name == "index_q" else self.total_num_kv_heads + ) + if self.tp_size <= total_heads: + return self.tp_size, self.tp_rank + replicas = self.tp_size // total_heads + return total_heads, self.tp_rank // replicas + + def load_five_way_weights(self, shards: Dict[str, Dict]) -> None: + """Load five checkpoint projections into this rank's packed MXFP8 matrix.""" + local_shards: Dict[str, Dict[str, torch.Tensor]] = {} + for shard_name in self._SHARD_NAMES: + shard = shards[shard_name] + if "weight" not in shard: + raise KeyError(f"Missing {shard_name} projection weight.") + shard_world, shard_rank = self._shard_geometry(shard_name) + local = {} + for key in ("weight", "weight_scale_inv", "weight_scale", "bias"): + if key in shard: + local[key] = load_weight_shard( + shard[key], + shard_world, + shard_rank, + TensorParallelMode.COLUMN, + device=torch.device("cuda"), + ) + local_shards[shard_name] = local + + combined: Dict[str, torch.Tensor] = { + "weight": torch.cat( + [local_shards[name]["weight"] for name in self._SHARD_NAMES], dim=0 + ).contiguous() + } + for key in ("weight_scale_inv", "weight_scale", "bias"): + present = [key in local_shards[name] for name in self._SHARD_NAMES] + if any(present): + if not all(present): + raise KeyError(f"Incomplete {key} across fused QKV+index shards.") + combined[key] = torch.cat( + [local_shards[name][key] for name in self._SHARD_NAMES], dim=0 + ).contiguous() + + # The checkpoint tensors above are already rank-local and packed. + # Temporarily select vanilla loading so Linear copies them without a + # second TP split or the three-shard QKV loader. + saved_tp_size = self.tp_size + saved_tp_rank = self.tp_rank + saved_tp_mode = self.tp_mode + saved_loading_config = self.weights_loading_config + try: + self.tp_size = 1 + self.tp_rank = 0 + self.tp_mode = None + self.weights_loading_config = WeightsLoadingConfig(weight_mode=WeightMode.VANILLA) + self.load_weights([combined]) + finally: + self.tp_size = saved_tp_size + self.tp_rank = saved_tp_rank + self.tp_mode = saved_tp_mode + self.weights_loading_config = saved_loading_config + + # --------------------------------------------------------------------------- # Config normalization helpers # --------------------------------------------------------------------------- @@ -660,8 +820,8 @@ def _extract_minimax_m3_attention_extra_attrs(layer_idx: str): @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) def minimax_m3_attn_custom_op_inplace( q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], layer_idx: str, @@ -672,8 +832,8 @@ def minimax_m3_attn_custom_op_inplace( num_tokens = attn_metadata.num_tokens attn_layer._dispatch_attention_backend( q[:num_tokens], - k[:num_tokens], - v[:num_tokens], + k[:num_tokens] if k is not None else None, + v[:num_tokens] if v is not None else None, idx_q[:num_tokens] if idx_q is not None else None, idx_k[:num_tokens] if idx_k is not None else None, attn_metadata, @@ -772,9 +932,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)) @@ -782,25 +949,45 @@ def __init__( self.sparse_local_block = int(sparse_cfg.get("sparse_local_block", 1)) self.sparse_score_type = str(sparse_cfg.get("sparse_score_type", "max")) - # Index Q and K are both replicated and project the same - # hidden_states, so fuse them into one GEMM with output - # [idx_q | idx_k]. idx_q holds all index heads; idx_k is a single K - # per token, broadcast across heads when scoring. + if self.enable_fused_qkv_index_projection: + old_qkv_proj = self.qkv_proj + self.qkv_proj = MiniMaxM3QKVIndexerLinear( + hidden_size=config.hidden_size, + head_dim=self.head_dim, + total_num_heads=config.num_attention_heads, + total_num_kv_heads=config.num_key_value_heads, + total_num_index_heads=total_num_index_heads, + index_head_dim=self.sparse_index_dim, + dtype=config.torch_dtype, + mapping=old_qkv_proj.mapping, + quant_config=old_qkv_proj.quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + force_dynamic_quantization=old_qkv_proj.force_dynamic_quantization, + disable_deep_gemm=old_qkv_proj.disable_deep_gemm, + use_custom_cublas_mm=old_qkv_proj.use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=old_qkv_proj.use_cute_dsl_bf16_gemm, + use_cute_dsl_blockscaling_mm=old_qkv_proj.use_cute_dsl_blockscaling_mm, + ) + self.sparse_num_index_heads = self.qkv_proj.local_num_index_heads + else: + # Compatibility path: index Q/K remain replicated and use a + # separate GEMM with output [index-Q | index-K]. + 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. @@ -1042,6 +1229,131 @@ def _fused_fp8_index_qk_norm_rope( position_ids.reshape(-1).contiguous().to(torch.int32), ).flatten(1) + def _fused_fp8_qkv_indexer_norm_rope_kv_insert( + self, + packed: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Run the vLLM-style horizontal producer for every sparse batch. + + The CUDA kernel is token-major and batch-type agnostic: per-token + positions and cache slots cover pure prefill, mixed aggregate batches, + and CUDA-graph decode. + """ + if ( + not self.enable_fused_qkv_index_projection + or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) + or not self._emit_fp8_main_qkv() + or self.attn.indexer_kv_dtype != "fp8" + ): + return None + if is_torch_compiling() 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 + norm_eps = self.q_norm.variance_epsilon + if any( + module.variance_epsilon != norm_eps + for module in (self.k_norm, self.index_q_norm, self.index_k_norm) + ): + return None + norm_weights = ( + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + ) + if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): + return None + + kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) + if kv_cache_manager is None: + return None + buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + num_tokens = int(packed.shape[0]) + supported_main_cache = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= self.num_key_value_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index_cache = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + rotary_cos_sin = self.rotary_emb.rotary_cos_sin + supported_rope_cache = ( + rotary_cos_sin.is_cuda + and rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main_cache + or not supported_index_cache + or not supported_rope_cache + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < num_tokens + ): + return None + + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers, + index_k_cache, + out_cache_loc[:num_tokens], + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + self.head_dim, + 64, + norm_eps, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + rotary_cos_sin, + position_ids.reshape(-1).contiguous().to(torch.int32), + ) + return q.flatten(1), index_q.flatten(1) + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: """Whether the fused kernel is expected to run instead of the fallback. @@ -1393,8 +1705,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, @@ -1421,8 +1733,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, @@ -1440,6 +1752,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) @@ -1449,8 +1762,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, @@ -1462,6 +1775,7 @@ def _msa_attention_core( FMHA forward; this layer selects the top-k blocks (sparse only) and builds the forward_args the FMHA reads. """ + assert (k is None) == (v is None) if self.is_sparse_attention_layer: assert idx_q is not None # Publish the selected blocks so the FMHA runs the sparse path. @@ -1525,11 +1839,31 @@ def _sparse_forward( "attn_metadata; received None." ) - # Project, norm, and apply RoPE for the main and index branches. Both - # read only hidden_states and write disjoint outputs, so they overlap on - # the aux stream and join before the attention core. + # The opt-in projection emits [Q|K|V|index-Q|index-K] with one GEMM. + # Its horizontal producer writes both paged caches directly and returns + # only compact Q/index-Q. Unsupported compile/geometry/layout cases + # split the packed output and retain the existing producer paths. + packed_qkv = None + packed_idx_qk = None + if self.enable_fused_qkv_index_projection: + packed = self.qkv_proj(hidden_states) + horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed, position_ids, attn_metadata + ) + if horizontal is not None: + q, idx_q = horizontal + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + main_size = self.q_size + 2 * self.kv_size + packed_qkv, packed_idx_qk = packed.split( + [main_size, self.index_q_size + self.index_k_size], dim=-1 + ) + + # Project, norm, and apply RoPE for the compatibility/fallback main and + # index branches. Both read only hidden_states and write disjoint + # outputs, so they overlap on the aux stream and join before attention. def _main_norm_rope(): - qkv = self.qkv_proj(hidden_states) + qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) fused_qkv = self._fused_qk_norm_rope( qkv, position_ids, @@ -1556,7 +1890,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. @@ -1809,6 +2145,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: # Layer-0 prologue only. For every subsequent layer the input_layernorm @@ -1834,10 +2171,22 @@ def forward( **kwargs, ) + capture_this_layer = spec_metadata is not None and spec_metadata.is_layer_capture( + self.layer_idx + ) + if capture_this_layer: + # Eagle3 needs the fully reduced residual stream before the next + # layer's RMSNorm. POST fusion would expose only the already + # normalized next-layer input plus the residual and make + # maybe_capture_hidden_states() add them a second time. + self.post_feed_forward_fusion = False + if self.block_sparse_moe is not None: - hidden_states, residual = self.forward_MoE(hidden_states, attn_metadata, residual) + hidden_states, residual = self.forward_MoE( + hidden_states, attn_metadata, residual, spec_metadata + ) else: - hidden_states, residual = self.forward_mlp(hidden_states, residual) + hidden_states, residual = self.forward_mlp(hidden_states, residual, spec_metadata) return hidden_states, residual @@ -1911,6 +2260,7 @@ def forward_MoE( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: torch.Tensor, + spec_metadata: Optional[SpecMetadata] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) @@ -1920,6 +2270,8 @@ def forward_MoE( final_all_reduce_params=self._feed_forward_all_reduce_params(), ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) return hidden_states, residual @@ -1927,6 +2279,7 @@ def forward_mlp( self, hidden_states: torch.Tensor, residual: torch.Tensor, + spec_metadata: Optional[SpecMetadata] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) @@ -1935,6 +2288,8 @@ def forward_mlp( final_all_reduce_params=self._feed_forward_all_reduce_params(), ) + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) return hidden_states, residual @@ -1994,6 +2349,7 @@ def forward( input_ids: Optional[torch.IntTensor] = None, position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if (input_ids is None) ^ (inputs_embeds is not None): @@ -2010,6 +2366,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, residual=residual, + spec_metadata=spec_metadata, ) # When setup_aliases has chained the final norm into the last decoder @@ -2051,6 +2408,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 @@ -2089,7 +2472,7 @@ def _fold_gemma_boundary_norm_weights(weights): @register_auto_model("MiniMaxM3SparseForCausalLM") -class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]): +class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" @classmethod @@ -2117,12 +2500,7 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): raw_pretrained = model_config.pretrained_config if is_minimax_m3_vl_config(raw_pretrained): model_config = get_text_model_config(model_config) - super().__init__( - MiniMaxM3Model(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size, - ) + super().__init__(MiniMaxM3Model(model_config), model_config) def load_weights( self, @@ -2131,8 +2509,9 @@ def load_weights( params_map: Optional[Dict[str, str]] = None, allow_partial_loading: bool = False, ) -> None: - # The generic loader has no rule for this fusion. The VL subclass routes - # its text weights through here, so both paths are covered. + # Pack the opt-in five-way projection with MiniMax-specific TP + # sharding before the generic mapper handles the compatibility path. + packed_projection_modules = _load_qkv_index_proj_weights(self, weights) _load_index_qk_proj_weights(self, weights) # Fold Gemma (1 + weight) into the layer-boundary RMSNorm weights so the # runtime norms can be plain (non-Gemma) and drive the fused @@ -2142,6 +2521,9 @@ def load_weights( weights = _fold_gemma_boundary_norm_weights(weights) if weight_mapper is None: weight_mapper = MiniMaxM3HfWeightMapper() + # The generic mapper understands three-way QKV fusion, not the + # already-loaded five-way Q/K/V/index-Q/index-K modules. + weight_mapper.add_skip_modules(packed_projection_modules) weight_mapper.init_model_and_config(self, self.model_config) merged_params_map = {**MINIMAX_M3_PARAMS_MAP, **(params_map or {})} super().load_weights( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 95c92473efa5..0fed71b42f22 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -765,10 +765,8 @@ def _get_kv_size_per_token(self, # External drafter: layers start from 0, normal PP distribution # Resolve draft manager class from draft config — may differ # from target (e.g. hybrid target + plain transformer draft). - draft_kv_cache_manager_cls = get_kv_cache_manager_cls( - effective_draft_config, - draft_kv_cache_config, - is_disagg=self._is_disagg) + draft_kv_cache_manager_cls = self._get_draft_kv_cache_manager_cls( + effective_draft_config, draft_kv_cache_config) total += self._per_manager_cache_cost( draft_kv_cache_manager_cls, effective_draft_config, draft_kv_cache_config) @@ -1434,10 +1432,11 @@ def _should_create_separate_draft_kv_cache(self) -> bool: in the target model and don't produce a separate ModelConfig. We fall back to the target model's config via _get_effective_draft_config(). """ - if self._mapping.enable_attention_dp: + if self._mapping.enable_attention_dp and getattr( + self._kv_cache_manager_cls, "supports_shared_draft_layers", + True): logger.info( - "Attention DP is enabled, separate draft KV cache is not supported." - ) + "Attention DP: draft layers share the target KV cache manager.") return False sparse_cfg = self._sparse_attention_config @@ -1478,6 +1477,21 @@ def _get_num_draft_layers(self) -> int: return self._draft_config.pretrained_config.num_hidden_layers return get_num_spec_layers(self._speculative_config) + def _get_draft_kv_cache_manager_cls(self, + effective_draft_config: ModelConfig, + draft_kv_config: KvCacheConfig): + """Resolve the draft manager, preserving a target V2 lifecycle.""" + draft_cls = get_kv_cache_manager_cls( + effective_draft_config, + draft_kv_config, + is_disagg=self._is_disagg, + cache_transceiver_config=self._cache_transceiver_config, + ) + if self._is_kv_cache_manager_v2 and not issubclass( + draft_cls, KVCacheManagerV2): + draft_cls = KVCacheManagerV2 + return draft_cls + def _get_draft_max_attention_window( self, max_seq_len: int, @@ -1550,8 +1564,8 @@ def _create_one_model_draft_kv_cache_manager( f"Derived draft KV cache max_attention_window for separate " f"draft manager: {draft_kv_config.max_attention_window}") # Get the appropriate KV cache manager class for the draft model - draft_kv_cache_manager_cls = get_kv_cache_manager_cls( - effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) + draft_kv_cache_manager_cls = self._get_draft_kv_cache_manager_cls( + effective_draft_config, draft_kv_config) draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2( draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config) @@ -1561,12 +1575,28 @@ def _create_one_model_draft_kv_cache_manager( # the sparse_attention_config. Get it from effective_draft_config which # falls back to the target model's config for MTP mode. sparse_attn_config = effective_draft_config.sparse_attention_config + # Sparse targets can require a target-only page geometry. MiniMax-M3 + # uses 128-token MSA pages while its dense Eagle3 draft manager uses + # the validated 32-token generation-kernel geometry. + draft_tokens_per_block = getattr( + self._kv_cache_manager_cls, + "draft_manager_tokens_per_block", + self._tokens_per_block, + ) + if draft_tokens_per_block != self._tokens_per_block: + logger.info( + "Draft KV cache manager uses tokens_per_block=%d " + "(target uses %d).", + draft_tokens_per_block, + self._tokens_per_block, + ) + draft_kv_config.tokens_per_block = draft_tokens_per_block return _create_kv_cache_manager( model_engine=None, kv_cache_manager_cls=draft_kv_cache_manager_cls, mapping=self._mapping, kv_cache_config=draft_kv_config, - tokens_per_block=self._tokens_per_block, + tokens_per_block=draft_tokens_per_block, max_seq_len=max_seq_len, max_batch_size=self._max_batch_size, spec_config=self._speculative_config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..4baea566b0cb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6987,6 +6987,13 @@ def _pad_attention_dp_dummy_request(self): key="attention_dp_dummy_insufficient_kv_capacity") return + # A separate one-model draft KV cache manager must see the same dummy + # request, otherwise its prepare_resources() lookup observes an + # unknown request id. MiniMax-M3 MSA always uses this separate dense + # Eagle3 manager, including under attention DP. + draft_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if (not self._enable_adp_dummy_fixes or self.kv_cache_transceiver is None): llm_request = self.kv_cache_manager.add_dummy_requests( @@ -6995,6 +7002,7 @@ def _pad_attention_dp_dummy_request(self): is_gen=self._adp_dummy_is_gen, prepare_resource=True, max_num_draft_tokens=self.max_total_draft_tokens, + draft_kv_cache_manager=draft_kv_cache_manager, )[0] llm_request.is_attention_dp_dummy = True spec_resource_manager = self.resource_manager.get_resource_manager( @@ -7018,6 +7026,7 @@ def _pad_attention_dp_dummy_request(self): is_gen=self._adp_dummy_is_gen, prepare_resource=True, max_num_draft_tokens=self.max_total_draft_tokens, + draft_kv_cache_manager=draft_kv_cache_manager, ) except OutOfPagesError: dummy_requests = None @@ -7091,6 +7100,8 @@ def _pad_empty_attention_dp_batch( return dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID] + draft_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) try: # Degrade to an empty batch rather than propagate: the ranks have # yet to agree on can_queue, so a rank-local raise would strand the @@ -7101,6 +7112,7 @@ def _pad_empty_attention_dp_batch( is_gen=True, prepare_resource=True, max_num_draft_tokens=self.max_total_draft_tokens, + draft_kv_cache_manager=draft_kv_cache_manager, ) except (OutOfPagesError, NoFreeSlotsError): dummy_requests = None diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d497862e0c8e..334d055c8830 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -653,8 +653,16 @@ def __init__(self, def max_draft_len(self) -> int: return self.spec_config.max_draft_len - def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + def _prepare_attn_metadata_for_spec_dec(self, attn_metadata, spec_metadata): + # Graph warmup runs more than once while the draft loop mutates + # kv_lens_cuda in place. Save/restore it during warmup; capture itself + # must record the mutation and therefore intentionally omits the save. + is_capturing = torch.cuda.is_current_stream_capturing() + if spec_metadata.is_cuda_graph and not is_capturing: + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda", + "kv_lens_cuda") + else: + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") batch_size = attn_metadata.num_seqs # Save spec-dec params that the drafting loop will overwrite. @@ -776,7 +784,8 @@ def _forward_impl(self, )) else: # Save the old attn_metadata and spec_metadata - self._prepare_attn_metadata_for_spec_dec(attn_metadata) + self._prepare_attn_metadata_for_spec_dec(attn_metadata, + spec_metadata) # Prepare inputs for the 1st draft model forward position_ids = position_ids.squeeze(0) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 645d715e721a..10045cfe22eb 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -716,7 +716,10 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): algorithm: Literal["minimax_m3"] = "minimax_m3" sparse_num_index_heads: PositiveInt = Field( default=4, - description="Number of index-attention heads (per TP rank's view).", + description= + "Checkpoint index-attention head count. The compatibility projection " + "replicates these heads on each TP rank; the fused QKV/index projection " + "shards them with the KV heads.", ) sparse_index_dim: int = Field( default=128, @@ -756,6 +759,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. MSA batches " + "also use a horizontal norm/RoPE/cache-insertion producer for prefill, " + "mixed, and CUDA-graph decode execution. The MiniMax-M3-specific path " + "requires the MSA implementation.", + status="prototype", + ) num_attention_heads: Optional[int] = Field( default=None, description= @@ -790,6 +803,10 @@ def _validate_msa_configuration(self): if self.indexer_kv_dtype == "fp8" and not self.sparse_disable_index_value: raise ValueError("MiniMax-M3 indexer_kv_dtype='fp8' requires " "sparse_disable_index_value=True.") + if self.fuse_qkv_index_projection and self.implementation != "msa": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True currently requires " + "the 'msa' implementation.") return self def supports_backend(self, backend: str) -> bool: @@ -813,6 +830,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): @@ -842,6 +860,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/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index cd2de8a50867..e57e454f94d3 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -513,6 +513,9 @@ nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION kv_cache_quant_algo: FP8 accuracy: 86 + - quant_algo: MIXED_PRECISION + spec_dec_algo: Eagle3 + accuracy: 88 nvidia/NVIDIA-Nemotron-Nano-9B-v2: - accuracy: 85.027 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index 314e04eb3052..b2e540868fb6 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -266,6 +266,9 @@ nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION kv_cache_quant_algo: FP8 accuracy: 81 + - quant_algo: MIXED_PRECISION + spec_dec_algo: Eagle3 + accuracy: 83 moonshotai/Kimi-K2-Instruct: - quant_algo: FP8_BLOCK_SCALES accuracy: 87.65 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index d89c33b7ee00..a103c3d3966f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -8430,7 +8430,8 @@ def test_nvfp4(self, use_msa): dtype="fp8" if use_msa else "auto") sparse_attention_config = MiniMaxM3SparseAttentionConfig( implementation="msa" if use_msa else "triton", - indexer_kv_dtype="fp8" if use_msa else "bf16") + indexer_kv_dtype="fp8" if use_msa else "bf16", + fuse_qkv_index_projection=use_msa) moe_config = MoeConfig(backend="CUTLASS") with LLM(model_path, tensor_parallel_size=tp_size, @@ -8446,6 +8447,96 @@ def test_nvfp4(self, use_msa): task = GSM8K(model_name) task.evaluate(llm) + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("cuda_graph", [True]) + @parametrize_with_ids("use_msa", [True]) + @parametrize_with_ids("overlap_scheduler", [False, True]) + @parametrize_with_ids("attention_dp", [False, True]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, + overlap_scheduler, use_msa, cuda_graph): + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import \ + msa_package_available + + if not msa_package_available(): + pytest.skip("MSA kernels (fmha_sm100) not available") + + model_name = "nvidia/MiniMax-M3-NVFP4" + model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" + spec_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3", + ) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + enable_block_reuse=False) + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa", + fuse_qkv_index_projection=True, + ) + + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + sparse_attention_config=sparse_attention_config, + moe_config=MoeConfig(backend="CUTLASS"), + max_seq_len=4096, + max_batch_size=256 if attention_dp else 512, + speculative_config=spec_config, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, + max_batch_size=64 if attention_dp else 128, + ) if cuda_graph else None, + disable_overlap_scheduler=not overlap_scheduler, + enable_attention_dp=attention_dp, + enable_iter_perf_stats=True, + trust_remote_code=True) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + + def drain_spec_stats(llm): + drafted = accepted = steps = 0 + for stats in llm.get_stats(timeout=2): + stats = json.loads(stats) if isinstance(stats, + str) else stats + spec_stats = stats.get("specDecodingStats") or {} + drafted += spec_stats.get("numDraftTokens", 0) + accepted += spec_stats.get("numAcceptedTokens", 0) + steps += spec_stats.get("numRequestsWithDraftTokens", 0) + return drafted, accepted, steps + + MMLU(model_name).evaluate(llm) + GSM8K(model_name).evaluate(llm) + + questions = [ + row["question"] + for row in load_dataset("gsm8k", "main", split="test") + ][:200] + chat_prompts = [ + llm.tokenizer.apply_chat_template([{ + "role": "user", + "content": question, + }], + tokenize=False, + add_generation_prompt=True) + for question in questions + ] + drain_spec_stats(llm) + llm.generate(chat_prompts, + SamplingParams(max_tokens=512, temperature=0)) + drafted, accepted, steps = drain_spec_stats(llm) + assert steps > 0, "no speculative iterations recorded" + chat_rate = accepted / drafted + chat_length = 1 + accepted / steps + print("MiniMax-M3 Eagle3 chat-GSM8K acceptance: " + f"rate={chat_rate:.3f}, mean acceptance length=" + f"{chat_length:.3f} ({steps} spec iterations)") + assert chat_rate > 0.78, ( + f"Eagle3 chat-GSM8K acceptance rate too low: {chat_rate:.3f}") + assert chat_length > 3.3, ( + "Eagle3 chat-GSM8K acceptance length too low: " + f"{chat_length:.3f}") + @skip_pre_blackwell class TestGLM5FP8(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d0e17e71119d..e53a25015695 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -649,6 +649,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEO accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6f0b5c320cc9..fb17a9ac2b71 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -78,6 +78,9 @@ l0_b200: - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False] - test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] - unittest/_torch/attention + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py - unittest/_torch/compilation - unittest/_torch/debugger - unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py 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 a6a543b99724..a1cc2dfd50ab 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 @@ -18,6 +18,7 @@ MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, ) +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import msa_paged_kv from tensorrt_llm._torch.attention_backend.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 @@ -90,6 +91,45 @@ def test_msa_fp8_indexer_config_is_explicit_and_lowered() -> None: MiniMaxM3SparseAttentionConfig(sparse_index_dim=sparse_index_dim) +def test_fused_qkv_index_projection_is_explicit_and_shards_index_heads() -> None: + cfg = MiniMaxM3SparseAttentionConfig( + implementation="msa", + fuse_qkv_index_projection=True, + num_attention_heads=64, + num_key_value_heads=4, + ) + sparse_params = cfg.to_sparse_params() + metadata_params = cfg.to_sparse_metadata_params() + mapping = SimpleNamespace(tp_size=2, enable_attention_dp=False) + + assert sparse_params.fuse_qkv_index_projection is True + assert metadata_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 + + kernel_cfg = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, + num_q_heads=32, + num_kv_heads=2, + head_dim=128, + ) + assert kernel_cfg.num_index_heads == 2 + + compatibility_cfg = MiniMaxM3SparseAttentionConfig( + implementation="msa", + num_attention_heads=64, + num_key_value_heads=4, + ) + compatibility_metadata = compatibility_cfg.to_sparse_metadata_params() + assert compatibility_metadata.sharded_index_head_count(mapping) == 4 + + with pytest.raises(ValueError, match=r"requires the 'msa' implementation"): + MiniMaxM3SparseAttentionConfig( + implementation="triton", + fuse_qkv_index_projection=True, + ) + + @pytest.mark.parametrize( ( "configured_sparse_index_dim", @@ -489,7 +529,7 @@ def test_run_indexer_routes_head_major_output_by_batch_mode( class FakeIndexer: def select_blocks(self, *args: object, **kwargs: object) -> torch.Tensor: - del args + captured["index_k_cache"] = args[1] captured["head_major_output"] = kwargs["head_major_output"] return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) @@ -541,6 +581,7 @@ def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: assert result.shape == (num_tokens, 1, 16) assert captured["head_major_output"] is expected_head_major + assert captured["index_k_cache"] is metadata.idx_k_cache @pytest.mark.parametrize( @@ -608,3 +649,42 @@ def test_msa_proxy_max_score_strided_index_k_matches_packed( assert not index_k_strided.is_contiguous() assert index_k_strided.stride(0) == coalescing_scale * page_size * head_dim assert torch.equal(strided_scores, packed_scores) + + +def test_msa_scratch_sizing_covers_spec_verify_tokens() -> None: + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.kv_cache_manager = None + metadata.max_num_sequences = 2 + metadata.max_num_tokens = 8 + metadata.msa_max_score = torch.zeros(4 * 16 * 2) + + with pytest.raises(ValueError, match=r"msa_max_score backing store"): + metadata._ensure_msa_decode_scratch_buffers( + num_index_heads=4, + max_batch=2, + capture_graph=False, + required_max_k_tiles=16, + ) + + +def test_per_token_valid_blocks_multi_token_decode() -> None: + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + per_token_valid_blocks, + ) + + qo_lens = torch.tensor([4], dtype=torch.int32) + kv_lens = torch.tensor([10], dtype=torch.int32) + qo_offset = torch.tensor([6], dtype=torch.int32) + n_valid = per_token_valid_blocks(qo_lens, kv_lens, qo_offset, causal=True, block_size=2) + assert n_valid.tolist() == [4, 4, 5, 5] + + qo_lens = torch.tensor([1, 3], dtype=torch.int32) + kv_lens = torch.tensor([9, 6], dtype=torch.int32) + n_valid = per_token_valid_blocks(qo_lens, kv_lens, kv_lens - qo_lens, causal=True, block_size=4) + assert n_valid.tolist() == [3, 1, 2, 2] + + +def test_msa_target_requires_separate_32_token_draft_cache() -> None: + assert MiniMaxM3KVCacheManagerV2.supports_shared_draft_layers is False + assert MiniMaxM3KVCacheManagerV2.draft_manager_tokens_per_block == 32 diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 981dc469eca5..b91837bf4199 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -38,8 +38,12 @@ ) from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, + MiniMaxM3DecoderLayer, + MiniMaxM3ForCausalLM, MiniMaxM3Model, + MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, + _load_qkv_index_proj_weights, _minimax_m3_swiglu_oai, _strip_language_model_prefix, _validate_sparse_attention_runtime_config, @@ -50,6 +54,7 @@ get_text_config, is_minimax_m3_vl_config, ) +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl_v2 from tensorrt_llm._torch.modules.fused_moe.routing import ( MiniMaxM2MoeRoutingMethod, @@ -94,6 +99,49 @@ def test_validate_sparse_attention_runtime_config_accepts_minimax_m3() -> None: _validate_sparse_attention_runtime_config(model_config) +def test_minimax_m3_uses_one_engine_speculative_base() -> None: + assert issubclass(MiniMaxM3ForCausalLM, SpecDecOneEngineForCausalLM) + + +def test_eagle_capture_precedes_next_layer_norm() -> None: + class CaptureMetadata: + def __init__(self) -> None: + self.captured = None + + def is_layer_capture(self, layer_idx: int) -> bool: + return layer_idx == 25 + + def maybe_capture_hidden_states(self, layer_idx, hidden_states, residual) -> None: + self.captured = (layer_idx, hidden_states.clone(), residual.clone()) + + layer = SimpleNamespace( + layer_idx=25, + _apply_pre_feed_forward_norm=lambda hidden, residual: (hidden + 1, residual + 2), + block_sparse_moe=lambda hidden, unused_metadata, **unused_kwargs: hidden + 3, + _feed_forward_all_reduce_params=lambda: None, + _apply_next_layer_layernorm=lambda hidden, residual: (hidden + 10, residual + 20), + ) + spec_metadata = CaptureMetadata() + hidden_states = torch.tensor([1.0]) + residual = torch.tensor([2.0]) + + output, output_residual = MiniMaxM3DecoderLayer.forward_MoE( + layer, + hidden_states, + SimpleNamespace(), + residual, + spec_metadata, + ) + + assert spec_metadata.captured is not None + layer_idx, captured_hidden, captured_residual = spec_metadata.captured + assert layer_idx == 25 + torch.testing.assert_close(captured_hidden, torch.tensor([5.0])) + torch.testing.assert_close(captured_residual, torch.tensor([4.0])) + torch.testing.assert_close(output, torch.tensor([15.0])) + torch.testing.assert_close(output_residual, torch.tensor([24.0])) + + def test_model_init_validates_sparse_attention_runtime_config() -> None: model_config = ModelConfig( pretrained_config=_make_text_config(), @@ -573,6 +621,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..cee422f5399a --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py @@ -0,0 +1,279 @@ +# 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 + ) + # Keep the parity reference slots valid: the legacy separate main-K/V + # producer does not support negative slots. Negative-slot handling is + # exercised below using horizontal eager execution versus graph replay. + rope_cache = _rope_cache(max(256, num_tokens)) + + main_width = (num_heads_q + 2 * num_kv_heads) * 128 + main_input = packed[:, :main_width].contiguous() + index_input = packed[:, main_width:].contiguous() + reference_main_cache = _main_cache(num_pages, num_kv_heads) + reference_index_cache = _index_cache(num_pages) + q_reference = torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + main_input, + reference_main_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + ) + index_q_reference = torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + index_input, + reference_index_cache, + slots, + num_index_heads, + 128, + 64, + 1e-5, + index_q_weight, + index_k_weight, + 5_000_000.0, + position_ids, + ) + + main_cache = _main_cache(num_pages, num_kv_heads) + index_cache = _index_cache(num_pages) + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + main_cache, + index_cache, + slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + position_ids, + ) + + valid = slots >= 0 + pages = slots[valid].long() // 128 + within = slots[valid].long() % 128 + 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, + ) + + # Aggregate decode captures this producer in a CUDA graph. The operator + # allocates compact Q/index-Q outputs while writing graph-stable paged + # caches through a graph-stable slot mapping, so exercise both capture and + # replay for decode-sized (1) and larger mixed/prefill token counts. + graph_packed = packed.clone() + graph_positions = position_ids.clone() + graph_slots = slots.clone() + graph_main_cache = _main_cache(num_pages, num_kv_heads) + graph_index_cache = _index_cache(num_pages) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_q, graph_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + graph_packed, + graph_main_cache, + graph_index_cache, + graph_slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + graph_positions, + ) + + # Replay with different projection values, nonuniform positions, and new + # cache destinations. This proves replay reads the refreshed graph buffers + # rather than retaining capture-time values or slots. + replay_packed = torch.randn_like(packed) + replay_positions = ( + torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 7 + 3 + ) % rope_cache.shape[0] + replay_slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 53 + 11) % ( + (num_pages - 1) * 128 + ) + if num_tokens > 1: + replay_slots[-1] = -1 + graph_packed.copy_(replay_packed) + graph_positions.copy_(replay_positions) + graph_slots.copy_(replay_slots) + graph_main_cache.zero_() + graph_index_cache.zero_() + + replay_main_cache = _main_cache(num_pages, num_kv_heads) + replay_index_cache = _index_cache(num_pages) + replay_q, replay_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + replay_packed, + replay_main_cache, + replay_index_cache, + replay_slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + replay_positions, + ) + graph.replay() + torch.cuda.synchronize() + + replay_valid = replay_slots >= 0 + replay_pages = replay_slots[replay_valid].long() // 128 + replay_within = replay_slots[replay_valid].long() % 128 + assert torch.equal(graph_q.view(torch.uint8), replay_q.view(torch.uint8)) + torch.testing.assert_close( + graph_index_q.float(), + replay_index_q.float(), + rtol=0.0, + atol=0.0, + ) + assert torch.equal( + graph_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8), + replay_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8), + ) + torch.testing.assert_close( + graph_index_cache[replay_pages, :, replay_within, :].float(), + replay_index_cache[replay_pages, :, replay_within, :].float(), + rtol=0.0, + atol=0.0, + ) + + +def test_minimax_m3_horizontal_producer_ignores_out_of_range_cache_slot(): + num_heads_q = 8 + num_kv_heads = 2 + total_heads = num_heads_q + 3 * num_kv_heads + 1 + packed = torch.randn(1, total_heads * 128, dtype=torch.bfloat16, device="cuda") + weights = [torch.randn(128, dtype=torch.bfloat16, device="cuda") for _ in range(4)] + positions = torch.zeros(1, dtype=torch.int32, device="cuda") + slots = torch.tensor([2 * 128], dtype=torch.int32, device="cuda") + main_cache = _main_cache(2, num_kv_heads) + index_cache = _index_cache(2) + main_before = main_cache.clone() + index_before = index_cache.clone() + + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + main_cache, + index_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + *weights, + _rope_cache(1), + positions, + ) + + assert q.shape == (1, num_heads_q, 128) + assert index_q.shape == (1, num_kv_heads, 128) + assert torch.equal(main_cache.view(torch.uint8), main_before.view(torch.uint8)) + assert torch.equal(index_cache.view(torch.uint8), index_before.view(torch.uint8)) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py new file mode 100644 index 000000000000..5d17670704d8 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(qkv, num_heads_q, num_kv_heads, q_weight, k_weight, position_ids): + output = torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + return output.view(qkv.shape[0], num_heads_q + 2 * num_kv_heads, 128).split( + [num_heads_q, num_kv_heads, num_kv_heads], dim=1 + ) + + +def _strided_kv_cache(num_pages, num_kv_heads, page_size=128, stride_scale=3): + backing = torch.zeros( + num_pages * stride_scale, + 2, + num_kv_heads, + page_size, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + return backing[::stride_scale] + + +def _inputs(num_tokens, num_heads_q, num_kv_heads): + qkv = torch.randn( + num_tokens, + (num_heads_q + 2 * num_kv_heads) * 128, + dtype=torch.bfloat16, + device="cuda", + ) + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + 8192 + return qkv, q_weight, k_weight, position_ids + + +def _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, num_heads_q, num_kv_heads): + return torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + qkv, + kv_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + ) + + +@pytest.mark.parametrize(("num_heads_q", "num_kv_heads"), [(8, 1), (8, 8), (64, 4)]) +@pytest.mark.parametrize("num_tokens", [1, 16, 129]) +def test_minimax_m3_fp8_main_kv_insert_matches_materialize_then_scatter( + num_tokens, num_heads_q, num_kv_heads +): + torch.manual_seed(1234) + page_size = 128 + num_pages = max(4, (num_tokens + page_size - 1) // page_size + 2) + qkv, q_weight, k_weight, position_ids = _inputs(num_tokens, num_heads_q, num_kv_heads) + slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 37) % ( + (num_pages - 1) * page_size + ) + kv_cache = _strided_kv_cache(num_pages, num_kv_heads, page_size) + guard_page = kv_cache[-1].clone() + + q_out = _run( + qkv, + kv_cache, + slots, + q_weight, + k_weight, + position_ids, + num_heads_q, + num_kv_heads, + ) + q_ref, k_ref, v_ref = _reference( + qkv, num_heads_q, num_kv_heads, q_weight, k_weight, position_ids + ) + pages = slots.long() // page_size + within = slots.long() % page_size + + # The specialized kernel uses powf while fused_qk_norm_rope_to_fp8 uses + # the exp2f/log2f equivalent, so values at an FP8 boundary can round to + # adjacent E4M3 values. + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + torch.testing.assert_close( + kv_cache[:, 0][pages, :, within, :].float(), + k_ref.float(), + rtol=0.13, + atol=0.05, + ) + assert torch.equal( + kv_cache[:, 1][pages, :, within, :].view(torch.uint8), + v_ref.contiguous().view(torch.uint8), + ) + assert torch.equal(kv_cache[-1].view(torch.uint8), guard_page.view(torch.uint8)) + + +def test_minimax_m3_fp8_main_kv_insert_uses_64bit_cache_offsets(): + """Exercise a real paged-cache address beyond INT32_MAX elements. + + The smallest contiguous HND pool whose page-65536 K row starts at + 2**31 FP8 elements is about 2 GiB. The former implicit int conversion in + the store helper wrapped this address negative; this test writes and reads + that real allocation so arithmetic-only tests cannot mask the bug. + """ + required_bytes = (65537 * 2 * 128 * 128) + (1 << 30) + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < required_bytes: + pytest.skip("64-bit cache-offset test requires about 3 GiB free GPU memory") + + torch.manual_seed(4321) + page = 65536 + kv_cache = torch.empty( + page + 1, + 2, + 1, + 128, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + qkv, q_weight, k_weight, position_ids = _inputs(1, 8, 1) + slots = torch.tensor([page * 128], dtype=torch.int32, device="cuda") + + q_out = _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, 8, 1) + q_ref, k_ref, v_ref = _reference(qkv, 8, 1, q_weight, k_weight, position_ids) + torch.cuda.synchronize() + + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + torch.testing.assert_close( + kv_cache[page, 0, :, 0, :].float(), + k_ref[0].float(), + rtol=0.13, + atol=0.05, + ) + assert torch.equal( + kv_cache[page, 1, :, 0, :].view(torch.uint8), + v_ref[0].contiguous().view(torch.uint8), + ) + + +@pytest.mark.parametrize("invalid_slot", [-1, 2 * 128]) +def test_minimax_m3_fp8_main_kv_insert_ignores_invalid_slot(invalid_slot): + torch.manual_seed(5678) + qkv, q_weight, k_weight, position_ids = _inputs(1, 8, 1) + kv_cache = _strided_kv_cache(2, 1) + kv_cache.fill_(1.0) + before = kv_cache.clone() + slots = torch.tensor([invalid_slot], dtype=torch.int32, device="cuda") + + q_out = _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, 8, 1) + q_ref, _, _ = _reference(qkv, 8, 1, q_weight, k_weight, position_ids) + + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + assert torch.equal(kv_cache.view(torch.uint8), before.view(torch.uint8)) From 6e42493d5e4440d5b610e067cb8f41801ff9eb79 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 02:33:43 -0700 Subject: [PATCH 2/7] [None][fix] Support fused MiniMax-M3 projection with piecewise graphs Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../sparse/minimax_m3/msa_backend.py | 45 +++++ .../_torch/models/modeling_minimaxm3.py | 99 ++++++++++- .../usage/llm_args_golden_manifest.json | 7 + .../sparse/test_minimax_m3_msa_backend.py | 49 ++++++ .../unittest/_torch/models/test_minimax_m3.py | 155 ++++++++++++++++++ 5 files changed, 350 insertions(+), 5 deletions(-) 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 9fb76118a980..05dc9995f14c 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 @@ -1087,6 +1087,51 @@ 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 horizontal producer inserted main K/V. + + ``TrtllmAttention.forward`` interprets ``k=None`` as a packed fused-QKV + buffer, so it cannot represent compact Q with prewritten paged K/V. + Dispatch the same MSA paged-GQA helper directly; it already skips its + cache write when live K/V tensors are absent. + """ + output = forward_args.output + if output is None: + raise RuntimeError( + f"{type(self).__name__}.forward_prepopulated_kv requires an output buffer." + ) + + sparse_backend_args = forward_args.sparse_backend_args + kv_block_indexes = ( + sparse_backend_args.topk_indices if sparse_backend_args is not None else None + ) + 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/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 0c0b80860416..2f4e7e416be7 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -817,19 +817,78 @@ def _extract_minimax_m3_attention_extra_attrs(layer_idx: str): return metadata, attn_layer +@torch.library.custom_op("trtllm::minimax_m3_qkv_index_proj", mutates_args=()) +def minimax_m3_qkv_index_proj( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> torch.Tensor: + """Run the five-way projection as one CUDA-graph-capturable custom op.""" + del position_ids # Symbolic token-shape carrier; projection itself is position agnostic. + _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + return attn_layer.qkv_proj(hidden_states) + + +@minimax_m3_qkv_index_proj.register_fake +def _minimax_m3_qkv_index_proj_fake( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> torch.Tensor: + """Preserve the symbolic token dimension across the opaque projection.""" + _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + qkv_proj = attn_layer.qkv_proj + if not isinstance(qkv_proj, MiniMaxM3QKVIndexerLinear): + raise RuntimeError( + "MiniMax-M3 fused QKV/index projection custom op requires " + f"MiniMaxM3QKVIndexerLinear, got {type(qkv_proj).__name__}." + ) + # The real projection is token-major over ``hidden_states``. Keep that + # exact output contract here so generation CUDA-graph batch sizes can + # share Dynamo's dynamic-shape specialization. ``position_ids`` remains + # an explicit custom-op input to carry the unpadded token symbol into + # piecewise context segments; the attention output below uses that symbol. + return hidden_states.new_empty((hidden_states.shape[0], sum(qkv_proj.local_output_sizes))) + + @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) def minimax_m3_attn_custom_op_inplace( - q: torch.Tensor, + q: Optional[torch.Tensor], k: Optional[torch.Tensor], v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], + packed_qkv_index: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], layer_idx: str, output: torch.Tensor, ) -> None: - """Run MiniMax-M3 cache and attention work behind a compile boundary.""" + """Run MiniMax-M3 cache and attention work behind a compile boundary. + + The horizontal producer needs live paged-cache tensors and cache-slot + metadata, which are intentionally resolved inside this opaque attention + boundary rather than traced through Dynamo. Projection remains in the + captured segment; only the cache-writing producer and MSA attention stay + on the eager side of the existing piecewise boundary. + """ attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) num_tokens = attn_metadata.num_tokens + if packed_qkv_index is not None: + horizontal = attn_layer._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed_qkv_index[:num_tokens], + position_ids[..., :num_tokens] if position_ids is not None else None, + attn_metadata, + ) + if horizontal is None: + raise RuntimeError( + "MiniMax-M3 fused horizontal producer is required inside the " + f"piecewise attention boundary for layer {layer_idx}, but its " + "runtime geometry or cache layout was unsupported." + ) + q, idx_q = horizontal + k = v = idx_k = None + if q is None: + raise RuntimeError(f"MiniMax-M3 attention layer {layer_idx} received no query tensor.") attn_layer._dispatch_attention_backend( q[:num_tokens], k[:num_tokens] if k is not None else None, @@ -1248,8 +1307,6 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( or self.attn.indexer_kv_dtype != "fp8" ): return None - if is_torch_compiling() or int(packed.shape[0]) != int(attn_metadata.num_tokens): - return None if ( packed.dtype != torch.bfloat16 or position_ids is None @@ -1723,6 +1780,8 @@ def _forward_attention_core( v, idx_q, idx_k, + None, + None, self.layer_idx_str, output, ) @@ -1788,7 +1847,10 @@ def _msa_attention_core( assert idx_q is None and idx_k is None # 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 k is None: + 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( @@ -1846,6 +1908,33 @@ def _sparse_forward( packed_qkv = None packed_idx_qk = None if self.enable_fused_qkv_index_projection: + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + # Keep the projection in the captured segment while hiding + # its shape-specializing MXFP8 internals behind a symbolic + # fake implementation. Cache insertion and MSA remain in the + # existing eager attention boundary below. + packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( + hidden_states, position_ids, self.layer_idx_str + ) + num_tokens = ( + position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] + ) + output = packed.new_empty( + (num_tokens, self.num_heads * self.head_dim), + dtype=self.attn_activation_dtype, + ) + maybe_bcg_minimax_m3_attn_custom_op_inplace( + None, + None, + None, + None, + None, + packed, + position_ids, + self.layer_idx_str, + output, + ) + return self.o_proj(output, all_reduce_params=all_reduce_params) packed = self.qkv_proj(hidden_states) horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( packed, position_ids, attn_metadata diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 0430e464a01d..bf014ad713b5 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1566,6 +1566,13 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.fuse_qkv_index_projection" + }, { "allowed_values": [ "triton", 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 a1cc2dfd50ab..90a4fee7a9a0 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 @@ -14,12 +14,14 @@ import pytest import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import ( MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, ) from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import msa_paged_kv +from tensorrt_llm._torch.attention_backend.sparse.params import SparseBackendForwardArgs from tensorrt_llm._torch.attention_backend.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.bindings import DataType @@ -584,6 +586,53 @@ def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: assert captured["index_k_cache"] is metadata.idx_k_cache +@pytest.mark.parametrize("is_sparse", [False, True], ids=["dense", "sparse"]) +@pytest.mark.parametrize("use_decode_plan", [False, True], ids=["eager", "decode"]) +def test_forward_prepopulated_kv_dispatches_compact_q_without_cache_rewrite( + monkeypatch, + is_sparse: bool, + use_decode_plan: bool, +) -> None: + from tensorrt_llm._torch.attention_backend.fmha import msa_sparse_gqa + + captured = {} + + def fake_run_msa_paged_gqa(*args, **kwargs) -> None: + captured["args"] = args + captured["kwargs"] = kwargs + + monkeypatch.setattr(msa_sparse_gqa, "run_msa_paged_gqa", fake_run_msa_paged_gqa) + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + q = torch.randn(3, 8) + output = torch.empty_like(q) + topk_indices = torch.zeros(3, 1, 16, dtype=torch.int32) if is_sparse else None + metadata = SimpleNamespace( + msa_decode_gqa_plan=("decode_sparse",) if use_decode_plan else None, + msa_eager_gqa_plan=("eager_sparse",), + msa_decode_dense_plan=("decode_dense",) if use_decode_plan else None, + msa_eager_dense_plan=("eager_dense",), + ) + sparse_args = SparseBackendForwardArgs(topk_indices=topk_indices) if is_sparse else None + + attention.forward_prepopulated_kv( + q, + metadata, + AttentionForwardArgs(output=output, sparse_backend_args=sparse_args), + ) + + args = captured["args"] + assert args[0] is attention + assert args[1] is q + assert args[2] is None + assert args[3] is None + assert args[4] is metadata + assert args[5] is output + assert captured["kwargs"]["kv_block_indexes"] is topk_indices + plan_kind = "sparse" if is_sparse else "dense" + plan_phase = "decode" if use_decode_plan else "eager" + assert captured["kwargs"]["plan"] == (f"{plan_phase}_{plan_kind}",) + + @pytest.mark.parametrize( "indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn], diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index b91837bf4199..dcdf15a3a3f0 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -31,6 +31,7 @@ from transformers import AutoConfig from utils.llm_data import llm_models_root +import tensorrt_llm._torch.models.modeling_minimaxm3 as modeling_minimaxm3 from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( @@ -142,6 +143,160 @@ def maybe_capture_hidden_states(self, layer_idx, hidden_states, residual) -> Non torch.testing.assert_close(output_residual, torch.tensor([24.0])) +def test_piecewise_attention_boundary_runs_horizontal_producer(monkeypatch) -> None: + class FakeAttentionLayer: + def __init__(self) -> None: + self.producer_shapes = None + + def _fused_fp8_qkv_indexer_norm_rope_kv_insert(self, packed, position_ids, attn_metadata): + self.producer_shapes = ( + tuple(packed.shape), + tuple(position_ids.shape), + attn_metadata.num_tokens, + ) + return packed[:, :3].clone(), packed[:, :1].clone() + + def _dispatch_attention_backend(self, q, k, v, idx_q, idx_k, attn_metadata, output) -> None: + assert k is None and v is None and idx_k is None + assert idx_q.shape == (attn_metadata.num_tokens, 1) + output.copy_(q) + + metadata = SimpleNamespace(num_tokens=2) + layer = FakeAttentionLayer() + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (metadata, layer), + ) + packed = torch.arange(20, dtype=torch.float32).reshape(4, 5) + position_ids = torch.arange(4, dtype=torch.int32).reshape(1, 4) + output = torch.full((4, 3), -1.0) + + modeling_minimaxm3.minimax_m3_attn_custom_op_inplace( + None, + None, + None, + None, + None, + packed, + position_ids, + "3", + output, + ) + + assert layer.producer_shapes == ((2, 5), (1, 2), 2) + torch.testing.assert_close(output[:2], packed[:2, :3]) + torch.testing.assert_close(output[2:], torch.full((2, 3), -1.0)) + + +def test_piecewise_projection_fake_preserves_padded_hidden_rows(monkeypatch) -> None: + projection = object.__new__(MiniMaxM3QKVIndexerLinear) + nn.Module.__init__(projection) + projection.local_output_sizes = (3, 4) + layer = SimpleNamespace(qkv_proj=projection) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (SimpleNamespace(), layer), + ) + hidden_states = torch.randn(256, 5) + position_ids = torch.arange(6).reshape(1, 6) + + packed = modeling_minimaxm3._minimax_m3_qkv_index_proj_fake(hidden_states, position_ids, "3") + + # The real GEMM projects every padded hidden row. Position IDs remain an + # input solely to carry the unpadded token symbol to the piecewise segment. + assert packed.shape == (hidden_states.shape[0], 7) + + +def test_piecewise_fused_projection_preserves_input_token_dimension(monkeypatch) -> None: + """Do not inherit a bucket-specialized token dimension from the GEMM output.""" + packed = torch.randn(2, 7) + captured = {} + + def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, output): + assert q is None and k is None and v is None + assert idx_q is None and idx_k is None + captured["packed"] = packed_arg + captured["position_ids"] = position_ids + captured["output_shape"] = tuple(output.shape) + output.zero_() + + layer = SimpleNamespace( + enable_fused_qkv_index_projection=True, + qkv_proj=lambda hidden_states: packed, + register_to_config=True, + num_heads=1, + head_dim=3, + attn_activation_dtype=torch.float32, + layer_idx_str="3", + o_proj=lambda output, all_reduce_params: output, + ) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (SimpleNamespace(), layer), + ) + monkeypatch.setattr(modeling_minimaxm3, "is_torch_compiling", lambda: True) + monkeypatch.setattr( + modeling_minimaxm3, + "maybe_bcg_minimax_m3_attn_custom_op_inplace", + fake_boundary, + ) + hidden_states = torch.randn(4, 5) + position_ids = torch.arange(6).reshape(1, 6) + + result = MiniMaxM3Attention._sparse_forward( + layer, + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + ) + + assert captured["packed"] is packed + assert captured["position_ids"] is position_ids + assert captured["output_shape"] == (position_ids.shape[-1], 3) + assert result.shape == (position_ids.shape[-1], 3) + + +def test_msa_attention_core_routes_compact_q_to_prewritten_kv_entrypoint() -> None: + selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) + + class FakeMsaBackend: + def __init__(self) -> None: + self.prepopulated_call = None + + def run_indexer(self, idx_q, idx_k, metadata): + assert idx_k is None + assert metadata is attn_metadata + return selected_blocks + + def forward_prepopulated_kv(self, q, metadata, forward_args) -> None: + self.prepopulated_call = (q, metadata, forward_args) + + def forward(self, *unused_args, **unused_kwargs) -> None: + raise AssertionError("compact Q must not enter TrtllmAttention.forward") + + layer = MiniMaxM3Attention.__new__(MiniMaxM3Attention) + backend = FakeMsaBackend() + layer.attn = backend + layer.is_sparse_attention_layer = True + q = torch.randn(2, 8) + idx_q = torch.randn(2, 4) + attn_metadata = SimpleNamespace() + output = torch.empty_like(q) + + result = layer._msa_attention_core(q, None, None, idx_q, None, attn_metadata, output) + + assert result is output + assert backend.prepopulated_call is not None + called_q, called_metadata, forward_args = backend.prepopulated_call + assert called_q is q + assert called_metadata is attn_metadata + assert forward_args.output is output + assert forward_args.sparse_backend_args.topk_indices is selected_blocks + + def test_model_init_validates_sparse_attention_runtime_config() -> None: model_config = ModelConfig( pretrained_config=_make_text_config(), From cb5edf86b81c6b4980b2dceb0a1624d3350f867e Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 14:49:23 -0700 Subject: [PATCH 3/7] [None][fix] Address MiniMax-M3 fused-producer review feedback Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 6 +- .../sparse/minimax_m3/msa_backend.py | 56 ++++++++++---- .../_torch/models/modeling_minimaxm3.py | 29 +++++++- tensorrt_llm/_torch/pyexecutor/_util.py | 7 +- tensorrt_llm/_torch/speculative/eagle3.py | 6 +- tensorrt_llm/llmapi/llm_args.py | 7 +- .../sparse/test_minimax_m3_msa_backend.py | 73 ++++++++++++++++++- .../unittest/_torch/models/test_minimax_m3.py | 44 +++++++++++ 8 files changed, 200 insertions(+), 28 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 8c7aa0a3e3e8..607485edfb1f 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -93,6 +93,7 @@ __device__ __forceinline__ void storeFp8HeadElements64( __nv_fp8_e4m3* out, int64_t offsetThread, float const (&elements)[numElemsPerThread]) { static_assert(numElemsPerThread == 4, "MiniMax-M3 FP8 store expects four elements per thread"); + static_assert(sizeof(__nv_fp8x2_storage_t) == 2, "MiniMax-M3 FP8 pair storage must be 16 bits"); // Form the final pointer with 64-bit arithmetic before one aligned 32-bit // store. Production coalesced paged-cache offsets can exceed INT32_MAX // FP8 elements even though each individual head row is small. @@ -372,6 +373,7 @@ namespace constexpr int kMinimaxM3HeadDim = 128; constexpr int kMinimaxM3RotaryDim = 64; constexpr int kMinimaxM3PageSize = 128; +static_assert((kMinimaxM3PageSize & (kMinimaxM3PageSize - 1)) == 0, "page size must be a power of two"); constexpr int kMinimaxM3ElemsPerThread = kMinimaxM3HeadDim / 32; // MiniMax-M3-only direct-cache specialization for eager pure prefill. The @@ -437,7 +439,7 @@ __global__ void minimaxM3Fp8QKNormRopeKVInsertKernel(__nv_bfloat16 const* qkvInp { return; } - int const page = slot >> 7; + int const page = slot / kMinimaxM3PageSize; if (page >= numPages) { return; @@ -631,7 +633,7 @@ __global__ void minimaxM3Fp8QKVIndexerNormRopeKVInsertKernel(__nv_bfloat16 const { return; } - int const page = slot >> 7; + int const page = slot / kMinimaxM3PageSize; if (page >= numPages) { return; 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 05dc9995f14c..93b2b5b83b17 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 @@ -250,6 +250,14 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # Device-only overlap-correction geometry. Page starts/counts replace a + # token-granular [max_sequences, max_seq_len] map, avoiding hundreds of + # MiB of persistent graph memory at long context lengths. + msa_kv_page_starts: Optional[torch.Tensor] = None + msa_kv_page_counts: Optional[torch.Tensor] = None + msa_q_batch_row: Optional[torch.Tensor] = None + msa_q_intra: Optional[torch.Tensor] = None + msa_qo_lens_dev: Optional[torch.Tensor] = None # _msa_buffers_ready gates the once-only device buffers; # _msa_fields_ready marks that the current step's buffers are populated. @@ -412,11 +420,17 @@ def _create_msa_buffers(self) -> None: # replace optimistic speculative lengths after prepare(), so cache # slots and per-token causal bounds must be derivable without a host # synchronization and without allocating inside CUDA graph replay. - tokens_per_block = int(kv_cache_manager.tokens_per_block) - self.msa_req_to_token = self.get_empty( + self.msa_kv_page_starts = self.get_empty( buffers, - (max_num_sequences, max_blocks_per_seq * tokens_per_block), - cache_name="msa_req_to_token", + (max_num_sequences,), + cache_name="msa_kv_page_starts", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_kv_page_counts = self.get_empty( + buffers, + (max_num_sequences,), + cache_name="msa_kv_page_counts", dtype=torch.int32, capture_graph=capture_graph, ) @@ -598,14 +612,25 @@ def on_update_kv_lens(self) -> None: token_kv_lens = kv_lens[q_batch_row] q_positions = token_kv_lens - qo_lens[q_batch_row] + self.msa_q_intra[:total_q] - table_width = int(self.msa_req_to_token.shape[1]) - table_indices = q_positions.to(torch.long).clamp(min=0, max=table_width - 1) - cache_slots = self.msa_req_to_token.reshape(-1).index_select( - 0, q_batch_row * table_width + table_indices + page_size = self._msa_page_size + page_starts = self.msa_kv_page_starts[:batch_size].to(torch.long) + page_counts = self.msa_kv_page_counts[:batch_size].to(torch.long) + token_page_counts = page_counts[q_batch_row] + # Overlap correction only shrinks optimistic lengths, so valid rows are + # already in range. Keep the fallback entirely device-side to preserve + # CUDA-graph replay and prevent invalid metadata from becoming an OOB + # page-table access that poisons the process. + safe_positions = q_positions.to(torch.long).clamp_min(0) + safe_positions = torch.minimum( + safe_positions, + (token_page_counts * page_size - 1).clamp_min(0), ) + logical_pages = torch.div(safe_positions, page_size, rounding_mode="floor") + page_table_rows = page_starts[q_batch_row] + logical_pages + physical_pages = self.msa_kv_indices.index_select(0, page_table_rows) + cache_slots = physical_pages * page_size + torch.remainder(safe_positions, page_size) self.msa_out_cache_loc[:total_q].copy_(cache_slots) - page_size = self._msa_page_size n_valid_blocks = torch.div( (q_positions + 1).clamp_min(1) + (page_size - 1), page_size, @@ -852,8 +877,14 @@ def _build_msa_fields(self) -> None: # Keep the device-side geometry needed to repair optimistic # speculative lengths after the overlap scheduler reports the actual # accepted-token counts. - step_width = int(req_to_token.shape[1]) - self.msa_req_to_token[:batch_size, :step_width].copy_(req_to_token, non_blocking=True) + page_counts = torch.div( + kv_lens_cpu.to(torch.int64) + page_size - 1, + page_size, + rounding_mode="floor", + ).to(torch.int32) + page_starts = torch.cumsum(page_counts, 0) - page_counts + self.msa_kv_page_starts[:batch_size].copy_(page_starts, non_blocking=True) + self.msa_kv_page_counts[:batch_size].copy_(page_counts, non_blocking=True) qo_lens_long = qo_lens_cpu.to(torch.long) batch_rows = torch.repeat_interleave( torch.arange(batch_size, dtype=torch.int32), qo_lens_long @@ -1022,9 +1053,6 @@ def run_indexer( ) idx_k_view = idx_k.view(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) # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized # index Q/K are cast directly and the proxy accumulates their QK scores # in FP32. Block ordering is invariant to the omitted positive scale. diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 2f4e7e416be7..8290722ade15 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -338,6 +338,24 @@ def _validate_sparse_attention_runtime_config( "Set the following in the LLM API configuration:\n" "sparse_attention_config:\n algorithm: minimax_m3" ) + if getattr(sparse_config, "fuse_qkv_index_projection", False): + if getattr(sparse_config, "implementation", None) != "msa": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires the 'msa' implementation." + ) + if getattr(sparse_config, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires indexer_kv_dtype='fp8'." + ) + quant_config = model_config.quant_config + if ( + quant_config is None + or quant_config.quant_mode is None + or not quant_config.quant_mode.has_fp8_kv_cache() + ): + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires an FP8 main KV cache." + ) def get_sparse_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[int]]: @@ -1307,6 +1325,8 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( or self.attn.indexer_kv_dtype != "fp8" ): return None + rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None + rotary_dim = int(rope.dim) if rope is not None else 0 if ( packed.dtype != torch.bfloat16 or position_ids is None @@ -1317,8 +1337,8 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( 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 rope is None + or rotary_dim != 64 ): return None norm_eps = self.q_norm.variance_epsilon @@ -1400,7 +1420,7 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( self.num_key_value_heads, self.sparse_num_index_heads, self.head_dim, - 64, + rotary_dim, norm_eps, self.q_norm.weight, self.k_norm.weight, @@ -2508,7 +2528,7 @@ def _load_qkv_index_proj_weights(model: nn.Module, weights) -> List[str]: 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) + for shard_name, checkpoint_name in zip(shard_names, checkpoint_names, strict=True) } module.load_five_way_weights(shards) loaded_modules.append(name) @@ -2630,6 +2650,7 @@ def setup_aliases(self) -> None: layer's input_layernorm; the last layer chains the final model norm so its output AllReduce folds the final normalization too. """ + super().setup_aliases() layers = self.model.layers num_layers = len(layers) for idx, layer in enumerate(layers): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0fed71b42f22..fb47100459b6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1585,11 +1585,8 @@ def _create_one_model_draft_kv_cache_manager( ) if draft_tokens_per_block != self._tokens_per_block: logger.info( - "Draft KV cache manager uses tokens_per_block=%d " - "(target uses %d).", - draft_tokens_per_block, - self._tokens_per_block, - ) + f"Draft KV cache manager uses tokens_per_block={draft_tokens_per_block} " + f"(target uses {self._tokens_per_block}).") draft_kv_config.tokens_per_block = draft_tokens_per_block return _create_kv_cache_manager( model_engine=None, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 334d055c8830..398eb05f2611 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -653,7 +653,11 @@ def __init__(self, def max_draft_len(self) -> int: return self.spec_config.max_draft_len - def _prepare_attn_metadata_for_spec_dec(self, attn_metadata, spec_metadata): + def _prepare_attn_metadata_for_spec_dec( + self, + attn_metadata: AttentionMetadata, + spec_metadata: Eagle3OneModelSpecMetadata, + ) -> None: # Graph warmup runs more than once while the draft loop mutates # kv_lens_cuda in place. Save/restore it during warmup; capture itself # must record the mutation and therefore intentionally omits the save. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 10045cfe22eb..bc7b883ea614 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -766,7 +766,8 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): "is sharded with the KV heads and index-K is replicated. MSA batches " "also use a horizontal norm/RoPE/cache-insertion producer for prefill, " "mixed, and CUDA-graph decode execution. The MiniMax-M3-specific path " - "requires the MSA implementation.", + "requires the MSA implementation, indexer_kv_dtype='fp8', and an FP8 " + "main KV cache.", status="prototype", ) num_attention_heads: Optional[int] = Field( @@ -807,6 +808,10 @@ def _validate_msa_configuration(self): raise ValueError( "MiniMax-M3 fuse_qkv_index_projection=True currently requires " "the 'msa' implementation.") + if self.fuse_qkv_index_projection and self.indexer_kv_dtype != "fp8": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True currently requires " + "indexer_kv_dtype='fp8'.") return self def supports_backend(self, backend: str) -> bool: 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 90a4fee7a9a0..26d0f32b4b6b 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 @@ -23,6 +23,7 @@ from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import msa_paged_kv from tensorrt_llm._torch.attention_backend.sparse.params import SparseBackendForwardArgs from tensorrt_llm._torch.attention_backend.sparse.registry import _resolve_minimax_m3_backend_cls +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -96,6 +97,7 @@ def test_msa_fp8_indexer_config_is_explicit_and_lowered() -> None: def test_fused_qkv_index_projection_is_explicit_and_shards_index_heads() -> None: cfg = MiniMaxM3SparseAttentionConfig( implementation="msa", + indexer_kv_dtype="fp8", fuse_qkv_index_projection=True, num_attention_heads=64, num_key_value_heads=4, @@ -109,6 +111,10 @@ def test_fused_qkv_index_projection_is_explicit_and_shards_index_heads() -> None assert metadata_params.sharded_head_counts(mapping) == (32, 2) assert metadata_params.sharded_index_head_count(mapping) == 2 + attention_dp_mapping = SimpleNamespace(tp_size=4, enable_attention_dp=True) + assert metadata_params.sharded_head_counts(attention_dp_mapping) == (64, 4) + assert metadata_params.sharded_index_head_count(attention_dp_mapping) == 4 + kernel_cfg = MiniMaxM3SparseConfig.from_sparse_params( sparse_params, num_q_heads=32, @@ -130,6 +136,71 @@ def test_fused_qkv_index_projection_is_explicit_and_shards_index_heads() -> None implementation="triton", fuse_qkv_index_projection=True, ) + with pytest.raises(ValueError, match=r"requires indexer_kv_dtype='fp8'"): + MiniMaxM3SparseAttentionConfig( + implementation="msa", + fuse_qkv_index_projection=True, + ) + + +def test_overlap_kv_correction_derives_slots_from_compact_page_table(monkeypatch) -> None: + """Decode correction stays device-only without a token-granular page map.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + monkeypatch.setattr(TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) + + metadata._msa_fields_ready = True + metadata._msa_live_batch = 2 + metadata._msa_live_total_q = 3 + metadata._msa_page_size = 128 + metadata.kv_lens_cuda = torch.tensor([130, 257], dtype=torch.int32) + metadata.msa_q_batch_row = torch.tensor([0, 1, 1], dtype=torch.int32) + metadata.msa_q_intra = torch.tensor([0, 0, 1], dtype=torch.int32) + metadata.msa_qo_lens_dev = torch.tensor([1, 2], dtype=torch.int32) + metadata.msa_kv_page_starts = torch.tensor([0, 2], dtype=torch.int32) + metadata.msa_kv_page_counts = torch.tensor([2, 3], dtype=torch.int32) + metadata.msa_kv_indices = torch.tensor([10, 20, 30, 40, 50], dtype=torch.int32) + metadata.msa_out_cache_loc = torch.empty(3, dtype=torch.int32) + metadata.msa_n_valid_blocks = torch.empty(3, dtype=torch.int32) + + def owner(rows: int) -> SimpleNamespace: + return SimpleNamespace( + plan=( + False, + 0, + rows, + { + "kv_segment_lens": torch.zeros(rows, dtype=torch.int32), + "qo_offset": torch.zeros(rows, dtype=torch.int32), + }, + None, + ) + ) + + metadata._msa_proxy_plan = owner(2) + metadata._msa_gqa_plan = owner(3) + metadata._msa_dense_plan = owner(2) + + metadata.on_update_kv_lens() + + torch.testing.assert_close( + metadata.msa_out_cache_loc, + torch.tensor([20 * 128 + 1, 40 * 128 + 127, 50 * 128], dtype=torch.int32), + ) + torch.testing.assert_close( + metadata.msa_n_valid_blocks, torch.tensor([2, 2, 3], dtype=torch.int32) + ) + torch.testing.assert_close( + metadata._msa_proxy_plan.plan[3]["kv_segment_lens"], metadata.kv_lens_cuda + ) + torch.testing.assert_close( + metadata._msa_gqa_plan.plan[3]["kv_segment_lens"], + torch.tensor([130, 257, 257], dtype=torch.int32), + ) + torch.testing.assert_close( + metadata._msa_dense_plan.plan[3]["qo_offset"], + torch.tensor([129, 255], dtype=torch.int32), + ) @pytest.mark.parametrize( @@ -556,7 +627,7 @@ def __init__(self) -> None: def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: del layer_idx - self.idx_k_cache = idx_k + self.idx_k_cache.copy_(idx_k) def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: del layer_idx diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index dcdf15a3a3f0..106acb0911b7 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -64,6 +64,8 @@ from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.llmapi import MiniMaxM3SparseAttentionConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo # --------------------------------------------------------------------------- # Fixtures @@ -100,10 +102,52 @@ def test_validate_sparse_attention_runtime_config_accepts_minimax_m3() -> None: _validate_sparse_attention_runtime_config(model_config) +def test_validate_fused_projection_requires_fp8_main_kv_cache() -> None: + sparse_config = MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + ) + model_config = ModelConfig( + pretrained_config=_make_text_config(), + sparse_attention_config=sparse_config, + ) + with pytest.raises(ValueError, match="requires an FP8 main KV cache"): + _validate_sparse_attention_runtime_config(model_config) + + model_config.quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + _validate_sparse_attention_runtime_config(model_config) + + def test_minimax_m3_uses_one_engine_speculative_base() -> None: assert issubclass(MiniMaxM3ForCausalLM, SpecDecOneEngineForCausalLM) +def test_setup_aliases_preserves_one_engine_draft_weight_loading() -> None: + loaded = [] + + class DraftModel: + shares_target_kv_cache = True + + def load_weights_from_target_model(self, target) -> None: + loaded.append(target) + + target = MiniMaxM3ForCausalLM.__new__(MiniMaxM3ForCausalLM) + layers = [ + SimpleNamespace(input_layernorm=object()), + SimpleNamespace(input_layernorm=object()), + ] + final_norm = object() + object.__setattr__(target, "draft_model", DraftModel()) + object.__setattr__(target, "model", SimpleNamespace(layers=layers, norm=final_norm)) + + target.setup_aliases() + + assert loaded == [target] + assert layers[0].next_layer_layernorm is layers[1].input_layernorm + assert layers[1].next_layer_layernorm is final_norm + + def test_eagle_capture_precedes_next_layer_norm() -> None: class CaptureMetadata: def __init__(self) -> None: From 78acf9d5a563008fb490f5c11b77a9656856ae69 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:41:55 -0700 Subject: [PATCH 4/7] [None][fix] Free draft KV for rejected ADP dummies Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 ++++- .../_torch/executor/test_py_executor.py | 60 ++++++++++++++++++- 2 files changed, 68 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8a8dc45ebcb8..00327dc726b1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3416,6 +3416,14 @@ def _revert_gen_alloc(self, scheduled_batch): continue self.kv_cache_manager.revert_allocate_generation(req) + def _free_adp_dummy_kv_resources(self, dummy_request: LlmRequest) -> None: + """Release target and independent draft KV allocated for an ADP dummy.""" + self.kv_cache_manager.free_resources(dummy_request) + draft_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(dummy_request) + def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: """Commit or roll back this iteration's tentative ADP dummy. @@ -3447,7 +3455,7 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: ResourceManagerType.SPEC_RESOURCE_MANAGER) if spec_resource_manager is not None: spec_resource_manager.free_resources(dummy_request) - self.kv_cache_manager.free_resources(dummy_request) + self._free_adp_dummy_kv_resources(dummy_request) self.active_requests.remove(dummy_request) def _revert_ctx_alloc(self, dropped_context_requests): @@ -7098,7 +7106,7 @@ def _pad_attention_dp_dummy_request(self): try: spec_resource_manager.add_dummy_requests(dummy_request_ids) except NoFreeSlotsError: - self.kv_cache_manager.free_resources(dummy_request) + self._free_adp_dummy_kv_resources(dummy_request) return assert dummy_request is not None @@ -7186,7 +7194,7 @@ def _pad_empty_attention_dp_batch( try: spec_resource_manager.add_dummy_requests(dummy_request_ids) except NoFreeSlotsError: - self.kv_cache_manager.free_resources(dummy_request) + self._free_adp_dummy_kv_resources(dummy_request) return dummy_request.is_attention_dp_dummy = True diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index d29ffdf8dcc9..8d2b3884407f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1776,6 +1776,9 @@ def _add_dummy(**kwargs): self.resource_manager = Mock() self.resource_manager.get_resource_manager.return_value = None + self._free_adp_dummy_kv_resources = types.MethodType( + PyExecutor._free_adp_dummy_kv_resources, self + ) def _run_pad(stub): @@ -1793,6 +1796,19 @@ def _run_update_role(stub, candidates): PyExecutor._update_adp_dummy_role(stub, candidates) +def _set_adp_resource_managers( + stub: _StubADPExecutor, + *, + spec_resource_manager: Mock | None = None, + draft_kv_cache_manager: Mock | None = None, +) -> None: + managers = { + ResourceManagerType.SPEC_RESOURCE_MANAGER: spec_resource_manager, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: draft_kv_cache_manager, + } + stub.resource_manager.get_resource_manager.side_effect = managers.get + + def test_adp_dummy_role_set_to_ctx_on_context_only_request(): from tensorrt_llm.bindings.internal.batch_manager import LlmRequestType @@ -2142,20 +2158,31 @@ def test_pad_dummy_spec_allocation_failure_rolls_back_kv_candidate(): stub.active_requests = [terminal_request] stub.expected_num_active_requests = 2 spec_resource_manager = Mock() + draft_kv_cache_manager = Mock() spec_resource_manager.add_dummy_requests.side_effect = NoFreeSlotsError("No free slots") - stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) _run_pad(stub) assert stub.active_requests == [terminal_request] assert stub._pending_adp_dummy_request is None stub.kv_cache_manager.free_resources.assert_called_once() + draft_kv_cache_manager.free_resources.assert_called_once() def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): stub = _StubADPExecutor() spec_resource_manager = Mock() - stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + draft_kv_cache_manager = Mock() + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) stub.active_requests = [terminal_request] stub.expected_num_active_requests = 2 @@ -2173,6 +2200,7 @@ def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): assert stub.active_requests == [terminal_request] spec_resource_manager.free_resources.assert_called_once_with(first_dummy) stub.kv_cache_manager.free_resources.assert_called_once_with(first_dummy) + draft_kv_cache_manager.free_resources.assert_called_once_with(first_dummy) stub.dist.tp_allgather.return_value = [1, 1] _run_pad(stub) @@ -2189,6 +2217,7 @@ def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): assert spec_resource_manager.add_dummy_requests.call_count == 2 spec_resource_manager.free_resources.assert_called_once_with(first_dummy) stub.kv_cache_manager.free_resources.assert_called_once_with(first_dummy) + draft_kv_cache_manager.free_resources.assert_called_once_with(first_dummy) def test_adp_dummy_rollback_only_frees_pending_candidate(): @@ -2455,6 +2484,25 @@ def test_pad_empty_batch_degrades_on_allocation_error(error): assert len(stub.active_requests) == 1 +def test_pad_empty_batch_spec_failure_rolls_back_target_and_draft_kv(): + stub, scheduled_batch = _unfittable_rank() + spec_resource_manager = Mock() + draft_kv_cache_manager = Mock() + spec_resource_manager.add_dummy_requests.side_effect = NoFreeSlotsError("No free slots") + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + + _run_pad_empty(stub, scheduled_batch) + + assert scheduled_batch.batch_size == 0 + assert len(stub.active_requests) == 1 + stub.kv_cache_manager.free_resources.assert_called_once() + draft_kv_cache_manager.free_resources.assert_called_once() + + @pytest.mark.parametrize("enable_adp_dummy_fixes", [False, True]) def test_pad_empty_batch_dummy_rolled_back_when_fleet_still_cannot_queue( enable_adp_dummy_fixes, @@ -2465,7 +2513,12 @@ def test_pad_empty_batch_dummy_rolled_back_when_fleet_still_cannot_queue( stub, scheduled_batch = _unfittable_rank(enable_adp_dummy_fixes=enable_adp_dummy_fixes) active_request = stub.active_requests[0] spec_resource_manager = Mock() - stub.resource_manager.get_resource_manager.return_value = spec_resource_manager + draft_kv_cache_manager = Mock() + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) _run_pad_empty(stub, scheduled_batch) dummy = stub._pending_adp_dummy_request @@ -2477,6 +2530,7 @@ def test_pad_empty_batch_dummy_rolled_back_when_fleet_still_cannot_queue( assert stub._pending_adp_dummy_request is None spec_resource_manager.free_resources.assert_called_once_with(dummy) stub.kv_cache_manager.free_resources.assert_called_once_with(dummy) + draft_kv_cache_manager.free_resources.assert_called_once_with(dummy) def test_pad_empty_batch_dummy_kept_when_fleet_can_queue(): From 6393f42bceb79b06276c7820658da8a762a162d9 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:35:31 -0700 Subject: [PATCH 5/7] [None][fix] Complete ADP dummy rollback paths Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 55 ++++--- .../_torch/executor/test_py_executor.py | 155 +++++++++++++++++- 2 files changed, 182 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fbb1bf1b31d3..16078c11fdf1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2659,6 +2659,7 @@ def _executor_loop_pp(self): can_queue = False if not can_queue: self._revert_gen_alloc(scheduled_batch) + self._finalize_adp_dummy_allocation(can_queue) if not can_queue: logger.debug( f"microbatch {microbatch_id} cannot be queued, skipping" @@ -3618,7 +3619,7 @@ def _check_disagg_transfer_progress_when_idle(self) -> None: # A single-rank CTX worker cannot diverge on a collective. Reap # completed sends while it is idle so their pinned KV blocks can # be reused by the next context requests. - if (is_idle and self._dist_size(self.dist, "world_size") == 1 and + if (self._dist_size(self.dist, "world_size") == 1 and self.async_transfer_manager.has_any_inflight_requests()): self._check_disagg_ctx_cache_transfer_status(0) return @@ -7045,29 +7046,43 @@ def _pad_attention_dp_dummy_request(self): draft_kv_cache_manager = self.resource_manager.get_resource_manager( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + has_live_adp_dummy = any( + request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID + for request in self.active_requests) + if has_live_adp_dummy: + return + assert self._pending_adp_dummy_request is None + if (not self._enable_adp_dummy_fixes or self.kv_cache_transceiver is None): - llm_request = self.kv_cache_manager.add_dummy_requests( - request_ids=dummy_request_ids, - token_nums=token_nums, - is_gen=self._adp_dummy_is_gen, - prepare_resource=True, - max_num_draft_tokens=self.max_total_draft_tokens, - draft_kv_cache_manager=draft_kv_cache_manager, - )[0] - llm_request.is_attention_dp_dummy = True + try: + dummy_requests = self.kv_cache_manager.add_dummy_requests( + request_ids=dummy_request_ids, + token_nums=token_nums, + is_gen=self._adp_dummy_is_gen, + prepare_resource=True, + max_num_draft_tokens=self.max_total_draft_tokens, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + except (OutOfPagesError, NoFreeSlotsError): + dummy_requests = None + if not dummy_requests: + logger.warning("Cannot allocate ADP pad dummy; rank schedules " + "an empty batch and the fleet will retry.") + return + + dummy_request = dummy_requests[0] spec_resource_manager = self.resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) if spec_resource_manager is not None: - spec_resource_manager.add_dummy_requests(dummy_request_ids) - self.active_requests.append(llm_request) - return - - assert self._pending_adp_dummy_request is None - has_live_adp_dummy = any( - request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID - for request in self.active_requests) - if has_live_adp_dummy: + try: + spec_resource_manager.add_dummy_requests(dummy_request_ids) + except NoFreeSlotsError: + self._free_adp_dummy_kv_resources(dummy_request) + return + dummy_request.is_attention_dp_dummy = True + self.active_requests.append(dummy_request) + self._pending_adp_dummy_request = dummy_request return try: @@ -7079,7 +7094,7 @@ def _pad_attention_dp_dummy_request(self): max_num_draft_tokens=self.max_total_draft_tokens, draft_kv_cache_manager=draft_kv_cache_manager, ) - except OutOfPagesError: + except (OutOfPagesError, NoFreeSlotsError): dummy_requests = None if not dummy_requests: logger.warning("Cannot allocate ADP pad dummy; rank schedules " diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 4363274fff29..16efde85fc1f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1040,14 +1040,7 @@ def test_sync_single_rank_ctx_reaps_idle_transfer( executor._check_disagg_gen_cache_transfer_status = Mock() executor._check_disagg_ctx_cache_transfer_status = Mock() - PyExecutor._check_disagg_transfer_progress_when_idle( - executor, - num_fitting_reqs=0, - fitting_disagg_gen_init_requests=[], - wait_for_disagg_gen_transfer_progress=True, - all_gen_first=False, - is_idle=True, - ) + PyExecutor._check_disagg_transfer_progress_when_idle(executor) executor.dist.allreduce.assert_not_called() executor.dist.tp_allreduce.assert_not_called() @@ -1328,6 +1321,68 @@ def stop_after_schedule(requests, inflight_req_ids): ] +def test_pp_loop_finalizes_pending_adp_dummy_after_queue_decision(monkeypatch): + class StopAfterFinalize(RuntimeError): + pass + + executor = object.__new__(PyExecutor) + executor.dist = Mock(pp_rank=0, rank=0) + executor.device_id = 0 + profiler = MagicMock() + profiler.__enter__.return_value = Mock() + executor._profiler = Mock(return_value=profiler) + executor.hang_detector = MagicMock() + executor.enable_iter_perf_stats = False + executor._uses_kv_manager_v2 = Mock(return_value=False) + executor._pp_rebalance_drain_iters = None + executor._handle_disagg_cache_errors_synced = Mock() + executor._fetch_and_activate_new_requests = Mock(return_value=[]) + executor.should_stop_processing = False + executor._handle_control_request = Mock() + executor.kv_cache_transceiver = None + executor._pad_attention_dp_dummy_request = Mock() + scheduled_batch = Mock( + batch_size=1, + encoder_requests=[], + num_encoder_requests=0, + num_context_requests=0, + num_generation_requests=1, + ) + executor._pp_schedule_and_propagate = Mock(return_value=(scheduled_batch, [], 0, False)) + executor._mm_encoder_item_scheduling_enabled = False + executor._is_kv_manager_v2 = False + executor.active_requests = [] + executor.num_scheduled_requests = 0 + executor.iter_counter = 0 + executor._can_queue = Mock(return_value=(False, False)) + executor._revert_gen_alloc = Mock() + + def stop_after_finalize(can_queue): + assert can_queue is False + raise StopAfterFinalize + + executor._finalize_adp_dummy_allocation = Mock(side_effect=stop_after_finalize) + + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.set_device", + Mock(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.cudart.cudaSetDevice", + Mock(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.CUASSERT", + Mock(), + ) + + with pytest.raises(StopAfterFinalize): + PyExecutor._executor_loop_pp(executor) + + executor._revert_gen_alloc.assert_called_once_with(scheduled_batch) + executor._finalize_adp_dummy_allocation.assert_called_once_with(False) + + def test_schedule_prepares_snapshot_points_before_scheduling(): class StopSchedule(RuntimeError): pass @@ -2111,6 +2166,35 @@ def test_pad_dummy_allocation_failure_skips_padding(): assert not any(r.is_attention_dp_dummy for r in stub.active_requests) +@pytest.mark.parametrize("error", [OutOfPagesError("no pages"), NoFreeSlotsError("no slots")]) +def test_pad_dummy_allocation_error_skips_padding(error): + """Rank-local KV allocation errors must not strand collective peers.""" + stub = _StubADPExecutor() + stub.active_requests = [_make_adp_request(_STATE_GENERATION_TO_COMPLETE)] + stub.expected_num_active_requests = 2 + stub.kv_cache_manager.add_dummy_requests.side_effect = error + + _run_pad(stub) + + assert len(stub.active_requests) == 1 + assert stub._pending_adp_dummy_request is None + + +def test_pad_dummy_reuses_live_singleton_without_reallocating(): + """A live dummy wins over stale expected-count input on a repeated call.""" + stub = _StubADPExecutor() + live_dummy = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + live_dummy.is_attention_dp_dummy = True + stub.active_requests = [live_dummy] + stub.expected_num_active_requests = 2 + stub._pending_adp_dummy_request = live_dummy + + _run_pad(stub) + + assert stub.add_dummy_calls == [] + assert stub._pending_adp_dummy_request is live_dummy + + def test_adp_pad_dummy_checks_full_context_capacity(): stub = _StubADPExecutor( max_num_tokens=4096, @@ -2170,6 +2254,7 @@ def test_pad_dummy_spec_allocation_failure_rolls_back_kv_candidate(): _run_pad(stub) + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager assert stub.active_requests == [terminal_request] assert stub._pending_adp_dummy_request is None stub.kv_cache_manager.free_resources.assert_called_once() @@ -2192,6 +2277,7 @@ def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): _run_pad(stub) first_dummy = stub._pending_adp_dummy_request assert first_dummy is not None + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager stub.dist.tp_allgather.side_effect = None stub.dist.tp_allgather.return_value = [1, 0] @@ -2222,6 +2308,57 @@ def test_adp_dummy_peer_empty_rolls_back_and_retry_succeeds(): draft_kv_cache_manager.free_resources.assert_called_once_with(first_dummy) +def test_legacy_pad_dummy_is_transactional_across_pp_queue_veto(): + stub = _StubADPExecutor(enable_adp_dummy_fixes=False) + terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + stub.active_requests = [terminal_request] + stub.expected_num_active_requests = 2 + spec_resource_manager = Mock() + draft_kv_cache_manager = Mock() + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + + _run_pad(stub) + + dummy = stub._pending_adp_dummy_request + assert dummy is not None + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager + + PyExecutor._finalize_adp_dummy_allocation(stub, can_queue=False) + + assert stub.active_requests == [terminal_request] + assert stub._pending_adp_dummy_request is None + spec_resource_manager.free_resources.assert_called_once_with(dummy) + stub.kv_cache_manager.free_resources.assert_called_once_with(dummy) + draft_kv_cache_manager.free_resources.assert_called_once_with(dummy) + + +def test_legacy_pad_dummy_spec_failure_rolls_back_target_and_draft_kv(): + stub = _StubADPExecutor(enable_adp_dummy_fixes=False) + terminal_request = _make_adp_request(_STATE_GENERATION_TO_COMPLETE) + stub.active_requests = [terminal_request] + stub.expected_num_active_requests = 2 + spec_resource_manager = Mock() + draft_kv_cache_manager = Mock() + spec_resource_manager.add_dummy_requests.side_effect = NoFreeSlotsError("No free slots") + _set_adp_resource_managers( + stub, + spec_resource_manager=spec_resource_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + + _run_pad(stub) + + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager + assert stub.active_requests == [terminal_request] + assert stub._pending_adp_dummy_request is None + stub.kv_cache_manager.free_resources.assert_called_once() + draft_kv_cache_manager.free_resources.assert_called_once() + + def test_adp_dummy_rollback_only_frees_pending_candidate(): stub = _StubADPExecutor() prior_dummy = _make_adp_request( @@ -2499,6 +2636,7 @@ def test_pad_empty_batch_spec_failure_rolls_back_target_and_draft_kv(): _run_pad_empty(stub, scheduled_batch) + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager assert scheduled_batch.batch_size == 0 assert len(stub.active_requests) == 1 stub.kv_cache_manager.free_resources.assert_called_once() @@ -2525,6 +2663,7 @@ def test_pad_empty_batch_dummy_rolled_back_when_fleet_still_cannot_queue( _run_pad_empty(stub, scheduled_batch) dummy = stub._pending_adp_dummy_request assert dummy is not None + assert stub.add_dummy_calls[0]["draft_kv_cache_manager"] is draft_kv_cache_manager PyExecutor._finalize_adp_dummy_allocation(stub, False) From c36ac82eb99498a566e0038fa5820a0006f995de Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:35:37 -0700 Subject: [PATCH 6/7] [None][test] Harden MiniMax-M3 fused path coverage Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../defs/accuracy/references/gsm8k.yaml | 1 + .../defs/accuracy/references/mmlu.yaml | 1 + .../defs/accuracy/test_llm_api_pytorch.py | 4 +- .../integration/test_lists/test-db/l0_cpu.yml | 10 +++ .../sparse/test_minimax_m3_msa_backend.py | 63 +++++++++++++++++++ .../_torch/speculative/test_eagle3.py | 44 +++++++++++++ 6 files changed, 122 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index e57e454f94d3..1c3ffff016c4 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -514,6 +514,7 @@ nvidia/MiniMax-M3-NVFP4: kv_cache_quant_algo: FP8 accuracy: 86 - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 spec_dec_algo: Eagle3 accuracy: 88 nvidia/NVIDIA-Nemotron-Nano-9B-v2: diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index b2e540868fb6..a7e144940348 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -267,6 +267,7 @@ nvidia/MiniMax-M3-NVFP4: kv_cache_quant_algo: FP8 accuracy: 81 - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 spec_dec_algo: Eagle3 accuracy: 83 moonshotai/Kimi-K2-Instruct: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2157eeb11eab..446682c2d387 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -8481,9 +8481,11 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3", ) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, - enable_block_reuse=False) + enable_block_reuse=False, + dtype="fp8") sparse_attention_config = MiniMaxM3SparseAttentionConfig( implementation="msa", + indexer_kv_dtype="fp8", fuse_qkv_index_projection=True, ) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 2b46f28c6873..2c20f7525b6e 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -30,7 +30,17 @@ l0_cpu: - unittest/_torch/peft - unittest/_torch/memory - unittest/_torch/modeling + - unittest/_torch/models/test_minimax_m3.py::test_eagle_capture_precedes_next_layer_norm + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_uses_one_engine_speculative_base + - unittest/_torch/models/test_minimax_m3.py::test_msa_attention_core_routes_compact_q_to_prewritten_kv_entrypoint + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_runs_horizontal_producer + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_fused_projection_preserves_input_token_dimension + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_projection_fake_preserves_padded_hidden_rows + - unittest/_torch/models/test_minimax_m3.py::test_setup_aliases_preserves_one_engine_draft_weight_loading + - unittest/_torch/models/test_minimax_m3.py::test_validate_fused_projection_requires_fp8_main_kv_cache - unittest/_torch/models/checkpoints - unittest/_torch/modules - unittest/_torch/multimodal 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 a94e9c2f77ab..ed11f05762fd 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 @@ -203,6 +203,58 @@ def owner(rows: int) -> SimpleNamespace: ) +def test_overlap_kv_correction_rebuilds_eager_plans(monkeypatch) -> None: + """Mixed/eager overlap correction rebuilds plans from corrected lengths.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + monkeypatch.setattr(TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + metadata._msa_fields_ready = True + metadata._msa_live_batch = 2 + metadata._msa_live_total_q = 3 + metadata._msa_proxy_plan = None + metadata.kv_lens_cuda = torch.tensor([129, 257], dtype=torch.int32) + rebuild_fields = Mock() + rebuild_plans = Mock() + monkeypatch.setattr(metadata, "_build_msa_fields", rebuild_fields) + monkeypatch.setattr(metadata, "_build_step_plans", rebuild_plans) + + metadata.on_update_kv_lens() + + torch.testing.assert_close( + metadata._msa_corrected_kv_lens_cpu, + torch.tensor([129, 257], dtype=torch.int32), + ) + rebuild_fields.assert_called_once_with() + rebuild_plans.assert_called_once_with() + + +def test_overlap_kv_correction_does_not_rebuild_during_capture(monkeypatch) -> None: + """A missing decode plan cannot introduce host work during graph capture.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + monkeypatch.setattr(TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + metadata._msa_fields_ready = True + metadata._msa_live_batch = 1 + metadata._msa_live_total_q = 1 + metadata._msa_proxy_plan = None + metadata._msa_corrected_kv_lens_cpu = None + metadata.kv_lens_cuda = torch.tensor([129], dtype=torch.int32) + rebuild_fields = Mock() + rebuild_plans = Mock() + monkeypatch.setattr(metadata, "_build_msa_fields", rebuild_fields) + monkeypatch.setattr(metadata, "_build_step_plans", rebuild_plans) + + metadata.on_update_kv_lens() + + assert metadata._msa_corrected_kv_lens_cpu is None + rebuild_fields.assert_not_called() + rebuild_plans.assert_not_called() + + @pytest.mark.parametrize( ( "configured_sparse_index_dim", @@ -704,6 +756,17 @@ def fake_run_msa_paged_gqa(*args, **kwargs) -> None: assert captured["kwargs"]["plan"] == (f"{plan_phase}_{plan_kind}",) +def test_forward_prepopulated_kv_requires_output_buffer() -> None: + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + + with pytest.raises(RuntimeError, match="requires an output buffer"): + attention.forward_prepopulated_kv( + torch.empty(1, 8), + SimpleNamespace(), + AttentionForwardArgs(output=None), + ) + + @pytest.mark.parametrize( "indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn], diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index e76cf655b01c..45b807478e5a 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -50,6 +50,50 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +@pytest.mark.parametrize( + ("is_cuda_graph", "is_capturing", "expected_fields"), + [ + (False, False, ("_seq_lens", "_seq_lens_cuda")), + (True, False, ("_seq_lens", "_seq_lens_cuda", "kv_lens_cuda")), + (True, True, ("_seq_lens", "_seq_lens_cuda")), + ], + ids=["eager", "graph-warmup", "graph-capture"], +) +def test_eagle3_spec_dec_preserves_kv_lens_only_during_graph_warmup( + monkeypatch: pytest.MonkeyPatch, + is_cuda_graph: bool, + is_capturing: bool, + expected_fields: tuple[str, ...], +) -> None: + """Warmup restores KV lengths, while capture records their mutation.""" + from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + + metadata = SimpleNamespace( + prepare_for_spec_dec=MagicMock(), + num_seqs=2, + spec_decoding_packed_mask=None, + spec_decoding_position_offsets=None, + spec_decoding_generation_lengths=None, + ) + worker = object.__new__(Eagle3OneModelWorker) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: is_capturing, + ) + + worker._prepare_attn_metadata_for_spec_dec( + metadata, + SimpleNamespace(is_cuda_graph=is_cuda_graph), + ) + + metadata.prepare_for_spec_dec.assert_called_once_with(*expected_fields) + assert worker._saved_packed_mask is None + assert worker._saved_position_offsets is None + assert worker._saved_position_offsets_cpp is None + assert worker._saved_generation_lengths is None + + def test_mtp_eagle_refreshes_dsa_metadata_before_draft_forward() -> None: """Refresh DSA mappings after switching to the draft cache.""" events = [] From 86c2a3132deef86da1dc81e1bab662e8276e3aac Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:00:03 -0700 Subject: [PATCH 7/7] [None][fix] Harden MiniMax-M3 fused path review fixes Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 2 ++ tensorrt_llm/_torch/pyexecutor/_util.py | 12 +++---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +++-- .../defs/accuracy/test_llm_api_pytorch.py | 2 +- .../_torch/executor/test_py_executor.py | 26 +++++++++++++-- ...test_minimax_m3_fp8_horizontal_producer.py | 32 +++++++++++++++++++ 6 files changed, 70 insertions(+), 12 deletions(-) diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index fe439d287132..933ed62432c2 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -315,6 +315,8 @@ std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsert( packed.size(1) == totalHeads * headDim, "Packed tensor width must equal (Q + 2*KV + index-Q + 1) * head_dim"); TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens"); TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens"); + TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1 && indexQWeight.dim() == 1 && indexKWeight.dim() == 1, + "All norm weights must be one-dimensional"); TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim && indexQWeight.numel() == headDim && indexKWeight.numel() == headDim, "All norm weights must contain head_dim elements"); diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dea56153d0a4..00f471d9aa63 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -69,9 +69,9 @@ use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KVCacheCompressionManager, KVCacheManager, - PeftCacheManager, ResourceManager, - ResourceManagerType) +from .resource_manager import (BaseResourceManager, KVCacheCompressionManager, + KVCacheManager, PeftCacheManager, + ResourceManager, ResourceManagerType) from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, TRTLLMSampler) from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, @@ -1480,9 +1480,9 @@ def _get_num_draft_layers(self) -> int: return self._draft_config.pretrained_config.num_hidden_layers return get_num_spec_layers(self._speculative_config) - def _get_draft_kv_cache_manager_cls(self, - effective_draft_config: ModelConfig, - draft_kv_config: KvCacheConfig): + def _get_draft_kv_cache_manager_cls( + self, effective_draft_config: ModelConfig, + draft_kv_config: KvCacheConfig) -> type[BaseResourceManager]: """Resolve the draft manager, preserving a target V2 lifecycle.""" draft_cls = get_kv_cache_manager_cls( effective_draft_config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 16078c11fdf1..20f027c416fe 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3411,11 +3411,13 @@ def _revert_gen_alloc(self, scheduled_batch): def _free_adp_dummy_kv_resources(self, dummy_request: LlmRequest) -> None: """Release target and independent draft KV allocated for an ADP dummy.""" - self.kv_cache_manager.free_resources(dummy_request) draft_kv_cache_manager = self.resource_manager.get_resource_manager( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(dummy_request) + try: + self.kv_cache_manager.free_resources(dummy_request) + finally: + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(dummy_request) def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: """Commit or roll back this iteration's tentative ADP dummy. diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 446682c2d387..5dafc736d07c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -8524,7 +8524,7 @@ def drain_spec_stats(llm): questions = [ row["question"] - for row in load_dataset("gsm8k", "main", split="test") + for row in load_dataset(GSM8K.DATASET_DIR, "main", split="test") ][:200] chat_prompts = [ llm.tokenizer.apply_chat_template([{ diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 16efde85fc1f..62d6e364d03b 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1321,7 +1321,7 @@ def stop_after_schedule(requests, inflight_req_ids): ] -def test_pp_loop_finalizes_pending_adp_dummy_after_queue_decision(monkeypatch): +def test_pp_loop_finalizes_pending_adp_dummy_after_queue_decision(monkeypatch) -> None: class StopAfterFinalize(RuntimeError): pass @@ -1337,7 +1337,7 @@ class StopAfterFinalize(RuntimeError): executor._pp_rebalance_drain_iters = None executor._handle_disagg_cache_errors_synced = Mock() executor._fetch_and_activate_new_requests = Mock(return_value=[]) - executor.should_stop_processing = False + executor.is_shutdown = False executor._handle_control_request = Mock() executor.kv_cache_transceiver = None executor._pad_attention_dp_dummy_request = Mock() @@ -1866,6 +1866,28 @@ def _set_adp_resource_managers( stub.resource_manager.get_resource_manager.side_effect = managers.get +def test_adp_dummy_kv_cleanup_attempts_draft_after_target_failure() -> None: + stub = _StubADPExecutor() + dummy_request = _make_adp_request( + _STATE_GENERATION_IN_PROGRESS, + request_id=ATTENTION_DP_DUMMY_REQUEST_ID, + is_dummy_request=True, + ) + draft_kv_cache_manager = Mock() + _set_adp_resource_managers( + stub, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + target_error = RuntimeError("target KV cleanup failed") + stub.kv_cache_manager.free_resources.side_effect = target_error + + with pytest.raises(RuntimeError, match="target KV cleanup failed"): + stub._free_adp_dummy_kv_resources(dummy_request) + + stub.kv_cache_manager.free_resources.assert_called_once_with(dummy_request) + draft_kv_cache_manager.free_resources.assert_called_once_with(dummy_request) + + def test_adp_dummy_role_set_to_ctx_on_context_only_request(): from tensorrt_llm.bindings.internal.batch_manager import LlmRequestType 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 index cee422f5399a..d74c78dd5c3b 100644 --- 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 @@ -277,3 +277,35 @@ def test_minimax_m3_horizontal_producer_ignores_out_of_range_cache_slot(): assert index_q.shape == (1, num_kv_heads, 128) assert torch.equal(main_cache.view(torch.uint8), main_before.view(torch.uint8)) assert torch.equal(index_cache.view(torch.uint8), index_before.view(torch.uint8)) + + +@pytest.mark.parametrize("invalid_weight_index", range(4)) +def test_minimax_m3_horizontal_producer_rejects_non_vector_norm_weight( + invalid_weight_index: int, +) -> None: + num_heads_q = 8 + num_kv_heads = 2 + num_index_heads = num_kv_heads + total_heads = num_heads_q + 2 * num_kv_heads + num_index_heads + 1 + packed = torch.randn(1, total_heads * 128, dtype=torch.bfloat16, device="cuda") + weights = [torch.randn(128, dtype=torch.bfloat16, device="cuda") for _ in range(4)] + weights[invalid_weight_index] = weights[invalid_weight_index].reshape(1, 128) + positions = torch.zeros(1, dtype=torch.int32, device="cuda") + slots = torch.zeros(1, dtype=torch.int32, device="cuda") + + with pytest.raises(RuntimeError, match="norm weights must be one-dimensional"): + torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + _main_cache(1, num_kv_heads), + _index_cache(1), + slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + *weights, + _rope_cache(1), + positions, + )