From e733de8ba04d55b766b3f3562a2e0600f1b8feae Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:56:03 -0700 Subject: [PATCH] [None][perf] Optimize MiniMax-M3 CTX block selection Add warp-specialized selector kernels for up to 128 allocated blocks and emit a head-major backing during pure prefill so the selector-to-K2Q handoff is zero-copy. Keep mixed and decode batches Q-major and cover routing, strides, ties, forcing, and CUDA graph replay. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/minimaxM3SelectBlocks.cu | 486 ++++++++++++++++-- .../kernels/minimaxM3SelectBlocks.h | 2 +- .../thop/minimaxM3SelectBlocksOp.cpp | 11 +- .../sparse/minimax_m3/msa_backend.py | 14 +- .../sparse/minimax_m3/msa_indexer.py | 12 +- .../sparse/minimax_m3/msa_utils.py | 12 +- .../_torch/custom_ops/cpp_custom_ops.py | 12 +- .../sparse/test_minimax_m3_msa_backend.py | 63 +++ .../sparse/test_minimax_m3_msa_selector.py | 79 ++- 9 files changed, 631 insertions(+), 60 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu index 0490844cd82c..6e28c8f4a358 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu @@ -37,9 +37,403 @@ constexpr int kTopK = 16; constexpr int kWarpSize = 32; constexpr int kThreadsPerBlock = 256; constexpr int kWarpsPerBlock = kThreadsPerBlock / kWarpSize; +constexpr int kSmallMaxBlocks = 128; constexpr float kInitScore = 1.0e30F; constexpr float kLocalScore = 1.0e29F; +constexpr uint32_t kFullWarpMask = 0xFFFFFFFFU; +__forceinline__ __device__ bool candidateGreater( + uint32_t lhsScoreKey, int32_t lhsBlockId, uint32_t rhsScoreKey, int32_t rhsBlockId) +{ + return lhsScoreKey > rhsScoreKey || (lhsScoreKey == rhsScoreKey && lhsBlockId < rhsBlockId); +} + +__forceinline__ __device__ void warpBitonicSortDesc64( + uint32_t& scoreKey0, int32_t& blockId0, uint32_t& scoreKey1, int32_t& blockId1, int32_t lane) +{ +#pragma unroll + for (int32_t size = 2; size <= 2 * kWarpSize; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + if (stride == kWarpSize) + { + bool const firstGreater = candidateGreater(scoreKey0, blockId0, scoreKey1, blockId1); + uint32_t const greaterScoreKey = firstGreater ? scoreKey0 : scoreKey1; + int32_t const greaterBlockId = firstGreater ? blockId0 : blockId1; + uint32_t const lesserScoreKey = firstGreater ? scoreKey1 : scoreKey0; + int32_t const lesserBlockId = firstGreater ? blockId1 : blockId0; + scoreKey0 = greaterScoreKey; + blockId0 = greaterBlockId; + scoreKey1 = lesserScoreKey; + blockId1 = lesserBlockId; + } + else + { + uint32_t const partnerScoreKey0 = __shfl_xor_sync(kFullWarpMask, scoreKey0, stride); + int32_t const partnerBlockId0 = __shfl_xor_sync(kFullWarpMask, blockId0, stride); + uint32_t const partnerScoreKey1 = __shfl_xor_sync(kFullWarpMask, scoreKey1, stride); + int32_t const partnerBlockId1 = __shfl_xor_sync(kFullWarpMask, blockId1, stride); + bool const takeGreater0 = ((lane & size) == 0) == ((lane & stride) == 0); + int32_t const secondIndex = lane + kWarpSize; + bool const takeGreater1 = ((secondIndex & size) == 0) == ((secondIndex & stride) == 0); + bool const partnerGreater0 = candidateGreater(partnerScoreKey0, partnerBlockId0, scoreKey0, blockId0); + bool const currentGreater0 = candidateGreater(scoreKey0, blockId0, partnerScoreKey0, partnerBlockId0); + bool const partnerGreater1 = candidateGreater(partnerScoreKey1, partnerBlockId1, scoreKey1, blockId1); + bool const currentGreater1 = candidateGreater(scoreKey1, blockId1, partnerScoreKey1, partnerBlockId1); + if ((takeGreater0 && partnerGreater0) || (!takeGreater0 && currentGreater0)) + { + scoreKey0 = partnerScoreKey0; + blockId0 = partnerBlockId0; + } + if ((takeGreater1 && partnerGreater1) || (!takeGreater1 && currentGreater1)) + { + scoreKey1 = partnerScoreKey1; + blockId1 = partnerBlockId1; + } + } + } + } +} + +__forceinline__ __device__ void warpBitonicSortDesc128(uint32_t (&scoreKeys)[4], int32_t (&blockIds)[4], int32_t lane) +{ + // Virtual item slot * 32 + lane stays in registers. Short strides exchange + // lanes with shuffles; strides 32 and 64 exchange slots within each lane. +#pragma unroll + for (int32_t size = 2; size <= kSmallMaxBlocks; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + uint32_t previousScoreKeys[4]; + int32_t previousBlockIds[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + previousScoreKeys[slot] = scoreKeys[slot]; + previousBlockIds[slot] = blockIds[slot]; + } + +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + uint32_t partnerScoreKey; + int32_t partnerBlockId; + if (stride < kWarpSize) + { + partnerScoreKey = __shfl_xor_sync(kFullWarpMask, previousScoreKeys[slot], stride); + partnerBlockId = __shfl_xor_sync(kFullWarpMask, previousBlockIds[slot], stride); + } + else + { + int32_t const partnerSlot = slot ^ (stride / kWarpSize); + partnerScoreKey = previousScoreKeys[partnerSlot]; + partnerBlockId = previousBlockIds[partnerSlot]; + } + + int32_t const item = lane + slot * kWarpSize; + bool const takeGreater = ((item & size) == 0) == ((item & stride) == 0); + bool const partnerGreater = candidateGreater( + partnerScoreKey, partnerBlockId, previousScoreKeys[slot], previousBlockIds[slot]); + bool const currentGreater = candidateGreater( + previousScoreKeys[slot], previousBlockIds[slot], partnerScoreKey, partnerBlockId); + if ((takeGreater && partnerGreater) || (!takeGreater && currentGreater)) + { + scoreKeys[slot] = partnerScoreKey; + blockIds[slot] = partnerBlockId; + } + } + } + } +} + +template +__forceinline__ __device__ int64_t outputOffset( + int32_t query, int32_t kvHead, int32_t totalQueries, int32_t numKvHeads, int32_t rank) +{ + int64_t const outputRow = HeadMajorOutput ? static_cast(kvHead) * totalQueries + query + : static_cast(query) * numKvHeads + kvHead; + return outputRow * kTopK + rank; +} + +template +__forceinline__ __device__ void selectFromCandidates(cg::thread_block_tile const& warp, + float (&localScores)[NumCandidates], int32_t (&localIndices)[NumCandidates], int32_t* output, int32_t query, + int32_t kvHead, int32_t totalQueries, int32_t numKvHeads, int32_t numBlocks) +{ + float selectedScores[kTopK]; + int32_t selectedIndices[kTopK]; + reduce_topk::reduceTopK(warp, selectedScores, selectedIndices, localScores, localIndices, -INFINITY); + + if (warp.thread_rank() == 0) + { +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + if (selectedScores[rank] == -INFINITY) + { + selectedIndices[rank] = -1; + } + } + + // MSA consumes block IDs in ascending order. Sort the sixteen selected + // IDs in registers, treating -1 padding as greater than every valid ID. +#pragma unroll + for (int32_t rank = 1; rank < kTopK; ++rank) + { + int32_t const candidate = selectedIndices[rank]; + int32_t const candidateKey = candidate < 0 ? numBlocks : candidate; + int32_t insertion = rank; + while (insertion > 0) + { + int32_t const previous = selectedIndices[insertion - 1]; + int32_t const previousKey = previous < 0 ? numBlocks : previous; + if (previousKey <= candidateKey) + { + break; + } + selectedIndices[insertion] = previous; + --insertion; + } + selectedIndices[insertion] = candidate; + } + +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + output[outputOffset(query, kvHead, totalQueries, numKvHeads, rank)] + = selectedIndices[rank]; + } + } +} + +template +__global__ void minimaxM3SelectBlocksSmallKernel(float const* __restrict__ scores, int64_t headStride, + int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, + int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + static_assert(NumCandidates * kWarpSize <= kSmallMaxBlocks); + auto const warp = cg::tiled_partition(cg::this_thread_block()); + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + float localScores[NumCandidates]; + int32_t localIndices[NumCandidates]; +#pragma unroll + for (int32_t slot = 0; slot < NumCandidates; ++slot) + { + int32_t const block = warp.thread_rank() + slot * kWarpSize; + RedType candidate{-INFINITY, RedType::kMaxIdx}; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidate = RedType{score, block}; + } + RedType::unpack(localScores[slot], localIndices[slot], candidate.compValIdx); + } + + selectFromCandidates( + warp, localScores, localIndices, output, query, kvHead, totalQueries, numKvHeads, numBlocks); +} + +template +__global__ void minimaxM3SelectBlocks64Kernel(float const* __restrict__ scores, int64_t headStride, int64_t blockStride, + int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const lane = threadIdx.x % kWarpSize; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + RedType candidates[2]{{-INFINITY, RedType::kMaxIdx}, {-INFINITY, RedType::kMaxIdx}}; +#pragma unroll + for (int32_t slot = 0; slot < 2; ++slot) + { + int32_t const block = lane + slot * kWarpSize; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidates[slot] = RedType{score, block}; + } + } + + uint32_t scoreKey0 = static_cast(candidates[0].compValIdx >> RedType::kMoveBits); + int32_t blockId0 + = RedType::kMaxIdx - static_cast(static_cast(candidates[0].compValIdx) & 0xFFFFU); + uint32_t scoreKey1 = static_cast(candidates[1].compValIdx >> RedType::kMoveBits); + int32_t blockId1 + = RedType::kMaxIdx - static_cast(static_cast(candidates[1].compValIdx) & 0xFFFFU); + warpBitonicSortDesc64(scoreKey0, blockId0, scoreKey1, blockId1, lane); + + if (lane < kTopK) + { + RedType const negativeInfinity{-INFINITY, 0}; + uint32_t const negativeInfinityScoreKey + = static_cast(negativeInfinity.compValIdx >> RedType::kMoveBits); + if (scoreKey0 == negativeInfinityScoreKey) + { + blockId0 = numBlocks; + } + + // MSA consumes block IDs in ascending order. numBlocks is the sentinel + // so padding naturally follows every valid ID. +#pragma unroll + for (int32_t size = 2; size <= kTopK; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + int32_t const partnerBlockId = __shfl_xor_sync(0xFFFFU, blockId0, stride); + bool const takeMin = ((lane & size) == 0) == ((lane & stride) == 0); + blockId0 = takeMin ? min(blockId0, partnerBlockId) : max(blockId0, partnerBlockId); + } + } + + output[outputOffset(query, kvHead, totalQueries, numKvHeads, lane)] + = blockId0 == numBlocks ? -1 : blockId0; + } +} + +template +__global__ void minimaxM3SelectBlocks128Kernel(float const* __restrict__ scores, int64_t headStride, + int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, + int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const lane = threadIdx.x % kWarpSize; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + RedType candidates[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + candidates[slot] = RedType{-INFINITY, RedType::kMaxIdx}; + int32_t const block = lane + slot * kWarpSize; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidates[slot] = RedType{score, block}; + } + } + + uint32_t scoreKeys[4]; + int32_t blockIds[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + scoreKeys[slot] = static_cast(candidates[slot].compValIdx >> RedType::kMoveBits); + blockIds[slot] + = RedType::kMaxIdx - static_cast(static_cast(candidates[slot].compValIdx) & 0xFFFFU); + } + warpBitonicSortDesc128(scoreKeys, blockIds, lane); + + if (lane < kTopK) + { + RedType const negativeInfinity{-INFINITY, 0}; + uint32_t const negativeInfinityScoreKey + = static_cast(negativeInfinity.compValIdx >> RedType::kMoveBits); + if (scoreKeys[0] == negativeInfinityScoreKey) + { + blockIds[0] = numBlocks; + } + + // MSA consumes block IDs in ascending order. numBlocks is the sentinel + // so padding naturally follows every valid ID. +#pragma unroll + for (int32_t size = 2; size <= kTopK; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + int32_t const partnerBlockId = __shfl_xor_sync(0xFFFFU, blockIds[0], stride); + bool const takeMin = ((lane & size) == 0) == ((lane & stride) == 0); + blockIds[0] = takeMin ? min(blockIds[0], partnerBlockId) : max(blockIds[0], partnerBlockId); + } + } + + output[outputOffset(query, kvHead, totalQueries, numKvHeads, lane)] + = blockIds[0] == numBlocks ? -1 : blockIds[0]; + } +} + +template __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) @@ -109,48 +503,44 @@ __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, in RedType::unpack(localScores[rank], localIndices[rank], localTopK[rank].compValIdx); } - float selectedScores[kTopK]; - int32_t selectedIndices[kTopK]; - reduce_topk::reduceTopK(warp, selectedScores, selectedIndices, localScores, localIndices, -INFINITY); + selectFromCandidates( + warp, localScores, localIndices, output, query, kvHead, totalQueries, numKvHeads, numBlocks); +} - if (warp.thread_rank() == 0) +template +void launchMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, + int32_t initBlocks, int32_t localBlocks, cudaStream_t stream) +{ + int32_t const numOutputRows = totalQueries * numKvHeads; + if (numOutputRows == 0) { -#pragma unroll - for (int32_t rank = 0; rank < kTopK; ++rank) - { - if (selectedScores[rank] == -INFINITY) - { - selectedIndices[rank] = -1; - } - } - - // MSA consumes block IDs in ascending order. Sort the sixteen selected - // IDs in registers, treating -1 padding as greater than every valid ID. -#pragma unroll - for (int32_t rank = 1; rank < kTopK; ++rank) - { - int32_t const candidate = selectedIndices[rank]; - int32_t const candidateKey = candidate < 0 ? numBlocks : candidate; - int32_t insertion = rank; - while (insertion > 0) - { - int32_t const previous = selectedIndices[insertion - 1]; - int32_t const previousKey = previous < 0 ? numBlocks : previous; - if (previousKey <= candidateKey) - { - break; - } - selectedIndices[insertion] = previous; - --insertion; - } - selectedIndices[insertion] = candidate; - } - -#pragma unroll - for (int32_t rank = 0; rank < kTopK; ++rank) - { - output[static_cast(outputRow) * kTopK + rank] = selectedIndices[rank]; - } + return; + } + int32_t const gridSize = (numOutputRows + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (numBlocks <= kWarpSize) + { + minimaxM3SelectBlocksSmallKernel<1, HeadMajorOutput><<>>(scores, + headStride, blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else if (numBlocks <= 2 * kWarpSize) + { + minimaxM3SelectBlocks64Kernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else if (numBlocks <= kSmallMaxBlocks) + { + minimaxM3SelectBlocks128Kernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else + { + minimaxM3SelectBlocksKernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); } } @@ -158,16 +548,18 @@ __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, in void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, - int32_t initBlocks, int32_t localBlocks, cudaStream_t stream) + int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream) { - int32_t const numOutputRows = totalQueries * numKvHeads; - if (numOutputRows == 0) + if (headMajorOutput) { - return; + launchMinimaxM3SelectBlocks(scores, headStride, blockStride, queryStride, nValidBlocks, output, + numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks, stream); + } + else + { + launchMinimaxM3SelectBlocks(scores, headStride, blockStride, queryStride, nValidBlocks, output, + numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks, stream); } - int32_t const gridSize = (numOutputRows + kWarpsPerBlock - 1) / kWarpsPerBlock; - minimaxM3SelectBlocksKernel<<>>(scores, headStride, blockStride, queryStride, - nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h index 23aeaa0e1730..45d068c303ab 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h @@ -29,7 +29,7 @@ namespace kernels void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, - int32_t initBlocks, int32_t localBlocks, cudaStream_t stream); + int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp index 6b4c3e3a9a85..be2d4aebd917 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp @@ -27,7 +27,7 @@ namespace torch_ext { torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK, - int64_t initBlocks, int64_t localBlocks) + int64_t initBlocks, int64_t localBlocks, bool headMajorOutput) { constexpr int64_t kRequiredTopK = 16; constexpr int64_t kMaxBlockIndex = 65'535; @@ -61,13 +61,16 @@ torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor c TORCH_CHECK(initBlocks <= std::numeric_limits::max() && localBlocks <= std::numeric_limits::max(), "minimax_m3_select_blocks forcing ranges exceed int32 range"); - auto output = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32)); + auto const outputOptions = scores.options().dtype(torch::kInt32); + auto output = headMajorOutput + ? torch::empty({scores.size(0), scores.size(2), topK}, outputOptions).permute({1, 0, 2}) + : torch::empty({scores.size(2), scores.size(0), topK}, outputOptions); auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); tensorrt_llm::kernels::invokeMinimaxM3SelectBlocks(scores.data_ptr(), scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), output.data_ptr(), static_cast(scores.size(0)), static_cast(scores.size(1)), static_cast(scores.size(2)), static_cast(initBlocks), static_cast(localBlocks), - stream); + headMajorOutput, stream); return output; } @@ -79,7 +82,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "minimax_m3_select_blocks(Tensor scores, Tensor n_valid_blocks, int topk, int init_blocks, int " - "local_blocks) -> Tensor"); + "local_blocks, bool head_major_output=False) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index e80c76857c6a..cc325b59bc33 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 @@ -1025,6 +1025,9 @@ def run_indexer( config = self.m3_config idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 num_tokens = int(idx_q.shape[0]) + head_major_output = ( + int(metadata.num_contexts or 0) > 0 and int(metadata.num_generations or 0) == 0 + ) # idx_q and idx_k may be strided column-views of a fused buffer, so # reshape to keep them zero-copy. The proxy fmha_sm100 and the index-K # scatter below both honor the source strides. @@ -1067,12 +1070,18 @@ def run_indexer( max_score = None # The empty case was resolved on the host, so short-circuit here. if metadata.msa_eager_all_blocks_empty: - return torch.full( - (num_tokens, config.num_kv_heads, MSA_REQUIRED_TOPK), + output_shape = ( + (config.num_kv_heads, num_tokens, MSA_REQUIRED_TOPK) + if head_major_output + else (num_tokens, config.num_kv_heads, MSA_REQUIRED_TOPK) + ) + output = torch.full( + output_shape, -1, dtype=torch.int32, device=idx_q.device, ) + return output.permute(1, 0, 2) if head_major_output else output n_valid_blocks = metadata.msa_eager_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] @@ -1087,6 +1096,7 @@ def run_indexer( proxy_plan=proxy_plan, max_score=max_score, n_valid_blocks=n_valid_blocks, + head_major_output=head_major_output, ) def sparse_attn_predict( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py index 61893c6ee96a..482ddda18d38 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py @@ -131,6 +131,7 @@ def select_blocks( proxy_plan: Optional[tuple] = None, max_score: Optional[torch.Tensor] = None, n_valid_blocks: Optional[torch.Tensor] = None, + head_major_output: bool = False, ) -> torch.Tensor: """Return [total_q, num_kv_heads, topk] selected block indices. @@ -180,18 +181,25 @@ def select_blocks( # Empty-selection guard. n_valid_blocks is a host tensor on # this path, so the .item() read does not sync the device. if n_valid_blocks.numel() == 0 or int(n_valid_blocks.max().item()) <= 0: - return torch.full( - (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK), + output_shape = ( + (config.num_kv_heads, idx_q.shape[0], MSA_REQUIRED_TOPK) + if head_major_output + else (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK) + ) + output = torch.full( + output_shape, -1, dtype=torch.int32, device=idx_q.device, ) + return output.permute(1, 0, 2) if head_major_output else output return select_blocks_from_maxscore( max_score_kv, topk=MSA_REQUIRED_TOPK, n_valid_blocks=n_valid_blocks, init_blocks=config.init_blocks, local_blocks=config.local_blocks, + head_major_output=head_major_output, ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index 33547d84d54b..9cc666413da4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -187,17 +187,25 @@ def select_blocks_from_maxscore( n_valid_blocks: torch.Tensor, init_blocks: int, local_blocks: int, + head_major_output: bool = False, ) -> torch.Tensor: """Select per-query top-k blocks from per-KV-head block scores. Applies init and local forced blocks and per-query valid-block masking on the amax-reduced scores [num_kv_heads, n_blocks, total_q]. Returns [total_q, num_kv_heads, topk] int32 ascending block ids with -1 tail - padding. + padding. When ``head_major_output`` is set, the logical result uses a + head-major backing so ``result.permute(1, 0, 2)`` is contiguous without a + copy. """ nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() return torch.ops.trtllm.minimax_m3_select_blocks( - max_score_kv, nvb, topk, init_blocks, local_blocks + max_score_kv, + nvb, + topk, + init_blocks, + local_blocks, + head_major_output, ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index ce3e410c2d8d..b31a859843bb 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -277,8 +277,18 @@ def _(logits, pass @torch.library.register_fake("trtllm::minimax_m3_select_blocks") - def _(scores, n_valid_blocks, topk, init_blocks, local_blocks): + def _( + scores, + n_valid_blocks, + topk, + init_blocks, + local_blocks, + head_major_output=False, + ): del n_valid_blocks, init_blocks, local_blocks + if head_major_output: + return scores.new_empty((scores.shape[0], scores.shape[2], topk), + dtype=torch.int32).permute(1, 0, 2) return scores.new_empty((scores.shape[2], scores.shape[0], topk), dtype=torch.int32) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 3ce0bdd9033a..c961dbf57b89 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 @@ -7,6 +7,8 @@ by the SM100 integration accuracy test. """ +from types import SimpleNamespace + import pytest import torch @@ -247,6 +249,67 @@ def msa_idx_k_cache(self, layer_idx): assert "write" not in captured +@pytest.mark.parametrize( + ("num_contexts", "num_generations", "expected_head_major"), + [(2, 0, True), (1, 1, False), (0, 2, False)], +) +def test_run_indexer_routes_head_major_output_by_batch_mode( + num_contexts, num_generations, expected_head_major +): + num_tokens, num_index_heads, sparse_index_dim = 3, 4, 128 + captured = {} + + class FakeIndexer: + def select_blocks(self, *args, **kwargs): + del args + captured["head_major_output"] = kwargs["head_major_output"] + return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) + + class FakeMetadata: + msa_decode_proxy_plan = None + msa_eager_proxy_plan = ("eager",) + msa_eager_all_blocks_empty = False + msa_eager_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) + msa_kv_indices = torch.arange(num_tokens, dtype=torch.int32) + msa_qo_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) + msa_kv_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) + msa_qo_offset_cpu = torch.tensor([0], dtype=torch.int32) + + def __init__(self): + self.num_contexts = num_contexts + self.num_generations = num_generations + self.idx_k_cache = None + + def msa_write_idx_k(self, layer_idx, idx_k): + del layer_idx + self.idx_k_cache = idx_k + + def msa_idx_k_cache(self, layer_idx): + del layer_idx + return self.idx_k_cache + + attention = SimpleNamespace( + layer_idx=0, + m3_config=SimpleNamespace( + sparse_index_dim=sparse_index_dim, + num_index_heads=num_index_heads, + num_kv_heads=1, + ), + indexer=FakeIndexer(), + ) + metadata = FakeMetadata() + + result = MiniMaxM3MsaSparseAttention.run_indexer( + attention, + torch.zeros(num_tokens, num_index_heads * sparse_index_dim), + torch.zeros(num_tokens, sparse_index_dim), + metadata, + ) + + assert result.shape == (num_tokens, 1, 16) + assert captured["head_major_output"] is expected_head_major + + def test_msa_proxy_max_score_strided_index_k_matches_packed(): if not torch.cuda.is_available(): pytest.skip("CUDA required") diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py index d8617108fe8a..44e6a9b20aff 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py @@ -57,7 +57,9 @@ def _reference_select_blocks( @pytest.mark.parametrize("num_kv_heads", [1, 4]) -@pytest.mark.parametrize("num_blocks", [1, 8, 16, 17, 127, 1024, 1537]) +@pytest.mark.parametrize( + "num_blocks", [1, 8, 16, 17, 32, 33, 64, 65, 96, 127, 128, 129, 1024, 1537] +) def test_fused_selector_matches_reference_random(num_kv_heads, num_blocks): total_q = 19 generator = torch.Generator(device="cuda").manual_seed(num_blocks) @@ -95,7 +97,82 @@ def test_fused_selector_matches_reference_random(num_kv_heads, num_blocks): assert actual.dtype == torch.int32 assert actual.shape == (total_q, num_kv_heads, 16) + assert actual.stride() == (num_kv_heads * 16, 16, 1) + assert actual.is_contiguous() + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("num_blocks", [65, 96, 128]) +def test_fused_selector_128_path_matches_reference_ties_and_forcing(num_blocks): + scores = torch.zeros((2, num_blocks, 5), device="cuda", dtype=torch.float32) + n_valid_blocks = torch.tensor( + [0, 8, 16, num_blocks - 1, num_blocks], device="cuda", dtype=torch.int32 + ) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("num_blocks", [16, 48, 96, 129]) +def test_fused_selector_head_major_output_is_zero_copy_q2k(num_blocks): + total_q, num_kv_heads = 7, 4 + generator = torch.Generator(device="cuda").manual_seed(num_blocks) + scores = torch.randn( + num_kv_heads, + num_blocks, + total_q, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + n_valid_blocks = torch.randint( + 0, + num_blocks + 1, + (total_q,), + generator=generator, + device="cuda", + dtype=torch.int32, + ) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + head_major_output=True, + ) + q2k = actual.permute(1, 0, 2).contiguous().to(torch.int32) + assert torch.equal(actual, expected) + assert actual.shape == (total_q, num_kv_heads, 16) + assert actual.stride() == (16, total_q * 16, 1) + assert not actual.is_contiguous() + assert q2k.shape == (num_kv_heads, total_q, 16) + assert q2k.is_contiguous() + assert q2k.data_ptr() == actual.data_ptr() + assert q2k.untyped_storage().data_ptr() == actual.untyped_storage().data_ptr() @pytest.mark.parametrize(