diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu new file mode 100644 index 000000000000..6e28c8f4a358 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu @@ -0,0 +1,567 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3SelectBlocks.h" + +#include "tensorrt_llm/kernels/moeTopKFuncs.cuh" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +namespace cg = cooperative_groups; + +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) +{ + 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; + RedType localTopK[kTopK]; +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + localTopK[rank] = RedType{-INFINITY, RedType::kMaxIdx}; + } + + for (int32_t block = warp.thread_rank(); block < validBlocks; block += kWarpSize) + { + 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; + } + + RedType const candidate{score, block}; + if (candidate.compValIdx > localTopK[kTopK - 1].compValIdx) + { + int32_t insertion = kTopK - 1; +#pragma unroll + for (int32_t rank = kTopK - 2; rank >= 0; --rank) + { + if (candidate.compValIdx > localTopK[rank].compValIdx) + { + localTopK[rank + 1] = localTopK[rank]; + insertion = rank; + } + } + localTopK[insertion] = candidate; + } + } + + float localScores[kTopK]; + int32_t localIndices[kTopK]; +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + RedType::unpack(localScores[rank], localIndices[rank], localTopK[rank].compValIdx); + } + + selectFromCandidates( + warp, localScores, localIndices, output, query, kvHead, totalQueries, numKvHeads, numBlocks); +} + +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) + { + 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); + } +} + +} // namespace + +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, bool headMajorOutput, cudaStream_t stream) +{ + if (headMajorOutput) + { + 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); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h new file mode 100644 index 000000000000..45d068c303ab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +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, bool headMajorOutput, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index f64b815e4240..4604ccb28581 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -116,6 +116,7 @@ add_library( IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp sparseKvCacheCompactOp.cpp + minimaxM3SelectBlocksOp.cpp mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp new file mode 100644 index 000000000000..40218e195618 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3SelectBlocks.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Select the highest-scoring MiniMax-M3 MSA blocks for each query and KV head. +//! +//! @param scores Float32 CUDA tensor shaped [num_kv_heads, num_blocks, total_queries]. Strided tensors are supported. +//! @param nValidBlocks Contiguous int32 CUDA tensor shaped [total_queries], on the same device as scores. +//! @param topK Number of blocks to select. The operator currently requires topK == 16. +//! @param initBlocks Number of initial blocks forced into the selection. +//! @param localBlocks Number of blocks immediately preceding each query's valid-block boundary forced into selection. +//! @param headMajorOutput Whether to use head-major backing storage while preserving the logical return shape. +//! @return Int32 CUDA tensor shaped [total_queries, num_kv_heads, 16]. Valid block IDs are ascending, followed by -1 +//! padding when fewer than 16 valid blocks exist. +torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK, + int64_t initBlocks, int64_t localBlocks, bool headMajorOutput) +{ + constexpr int64_t kRequiredTopK = 16; + constexpr int64_t kMaxBlockIndex = 65'535; + + TORCH_CHECK(scores.is_cuda(), "minimax_m3_select_blocks expects CUDA scores"); + TORCH_CHECK(scores.scalar_type() == torch::kFloat32, "minimax_m3_select_blocks expects float32 scores, got ", + scores.scalar_type()); + TORCH_CHECK(scores.dim() == 3, "scores must be [num_kv_heads, num_blocks, total_queries]"); + TORCH_CHECK(scores.stride(0) >= 0 && scores.stride(1) >= 0 && scores.stride(2) >= 0, + "scores must have non-negative strides"); + + TORCH_CHECK(nValidBlocks.is_cuda(), "minimax_m3_select_blocks expects CUDA n_valid_blocks"); + TORCH_CHECK(nValidBlocks.device() == scores.device(), "scores and n_valid_blocks must be on the same device"); + TORCH_CHECK(nValidBlocks.scalar_type() == torch::kInt32, + "minimax_m3_select_blocks expects int32 n_valid_blocks, got ", nValidBlocks.scalar_type()); + TORCH_CHECK( + nValidBlocks.dim() == 1 && nValidBlocks.size(0) == scores.size(2), "n_valid_blocks must be [total_queries]"); + TORCH_CHECK(nValidBlocks.is_contiguous(), "n_valid_blocks must be contiguous"); + + TORCH_CHECK(topK == kRequiredTopK, "minimax_m3_select_blocks supports topk=16, got ", topK); + TORCH_CHECK(initBlocks >= 0, "init_blocks must be non-negative"); + TORCH_CHECK(localBlocks >= 0, "local_blocks must be non-negative"); + TORCH_CHECK(scores.size(1) <= kMaxBlockIndex, "minimax_m3_select_blocks supports at most ", kMaxBlockIndex, + " blocks, got ", scores.size(1)); + TORCH_CHECK(scores.size(0) <= std::numeric_limits::max() + && scores.size(1) <= std::numeric_limits::max() + && scores.size(2) <= std::numeric_limits::max(), + "minimax_m3_select_blocks dimensions exceed int32 range"); + TORCH_CHECK(scores.size(0) * scores.size(2) <= std::numeric_limits::max(), + "minimax_m3_select_blocks output rows exceed int32 range"); + TORCH_CHECK(initBlocks <= std::numeric_limits::max() && localBlocks <= std::numeric_limits::max(), + "minimax_m3_select_blocks forcing ranges exceed int32 range"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const outputOptions = scores.options().dtype(torch::kInt32); + auto const 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), + headMajorOutput, stream); + return output; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +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, bool head_major_output=False) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_select_blocks", &tensorrt_llm::torch_ext::minimaxM3SelectBlocks); +} 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 858348f914ba..edfb34ace8e5 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 @@ -824,6 +824,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_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim) @@ -861,6 +864,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 2cace5573064..961509af8409 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. @@ -183,6 +184,7 @@ def select_blocks( 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 262f5e10ec0d..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 @@ -11,7 +11,7 @@ import torch -from .common import _INIT_SCORE, _LOCAL_SCORE, write_kv_slots +from .common import write_kv_slots # fmha_sm100 ships only head_dim 128 variants and the MiniMax-M3 checkpoint # selects topk 16. Callers enforce these early so a misconfiguration fails @@ -187,42 +187,26 @@ 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. """ - num_kv_heads, n_blocks, total_q = max_score_kv.shape - device = max_score_kv.device - scores = max_score_kv.permute(2, 0, 1).to(torch.float32).clone() - block_ids = torch.arange(n_blocks, device=device, dtype=torch.long) - nvb = n_valid_blocks.to(device=device, dtype=torch.long) - - if init_blocks > 0: - init_mask = block_ids.view(1, 1, -1) < init_blocks - scores = torch.where(init_mask, torch.full_like(scores, _INIT_SCORE), scores) - if local_blocks > 0: - local_start = (nvb - local_blocks).clamp_min(0) - local_mask = (block_ids.view(1, -1) >= local_start.view(-1, 1)) & ( - block_ids.view(1, -1) < nvb.view(-1, 1) - ) - scores = torch.where(local_mask.unsqueeze(1), torch.full_like(scores, _LOCAL_SCORE), scores) - block_valid = block_ids.view(1, -1) < nvb.view(-1, 1) - scores = scores.masked_fill(~block_valid.unsqueeze(1), float("-inf")) - - k = min(topk, n_blocks) - vals, idx = scores.topk(k=k, dim=-1) - idx = torch.where(vals != float("-inf"), idx, torch.full_like(idx, -1)) - sort_key = torch.where(idx < 0, torch.full_like(idx, n_blocks), idx) - sort_key, _ = torch.sort(sort_key, dim=-1) - idx = torch.where(sort_key >= n_blocks, torch.full_like(sort_key, -1), sort_key) - if k < topk: - pad = torch.full((total_q, num_kv_heads, topk - k), -1, dtype=idx.dtype, device=device) - idx = torch.cat([idx, pad], dim=-1) - return idx.to(torch.int32) + 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, + head_major_output, + ) __all__ = [ diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 1cd00b3e9b6c..b6493fec0e26 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -276,6 +276,22 @@ def _(logits, # In-place operation, no return value (void function) pass + @torch.library.register_fake("trtllm::minimax_m3_select_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) + @torch.library.register_fake("trtllm::kda_decode") def _(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, conv_state_q, conv_state_k, conv_state_v, a_log, g, dt_bias, beta, 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 2e44c27e9bd3..f596ab7ab6db 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 @@ -165,6 +167,66 @@ def fake_select_blocks_from_maxscore(*args, **kwargs): assert result is expected +@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_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 new file mode 100644 index 000000000000..fdf456edeac9 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py @@ -0,0 +1,460 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Correctness tests for the fused MiniMax-M3 MSA block selector.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import _INIT_SCORE, _LOCAL_SCORE +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + select_blocks_from_maxscore, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference_select_blocks( + max_score_kv: torch.Tensor, + *, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, +) -> torch.Tensor: + num_kv_heads, n_blocks, total_q = max_score_kv.shape + device = max_score_kv.device + scores = max_score_kv.permute(2, 0, 1).to(torch.float32).clone() + block_ids = torch.arange(n_blocks, device=device, dtype=torch.long) + nvb = n_valid_blocks.to(device=device, dtype=torch.long) + + if init_blocks > 0: + init_mask = block_ids.view(1, 1, -1) < init_blocks + scores = torch.where(init_mask, torch.full_like(scores, _INIT_SCORE), scores) + if local_blocks > 0: + local_start = (nvb - local_blocks).clamp_min(0) + local_mask = (block_ids.view(1, -1) >= local_start.view(-1, 1)) & ( + block_ids.view(1, -1) < nvb.view(-1, 1) + ) + scores = torch.where(local_mask.unsqueeze(1), torch.full_like(scores, _LOCAL_SCORE), scores) + block_valid = block_ids.view(1, -1) < nvb.view(-1, 1) + scores = scores.masked_fill(~block_valid.unsqueeze(1), float("-inf")) + + k = min(topk, n_blocks) + vals, idx = scores.topk(k=k, dim=-1) + idx = torch.where(vals != float("-inf"), idx, torch.full_like(idx, -1)) + sort_key = torch.where(idx < 0, torch.full_like(idx, n_blocks), idx) + sort_key, _ = torch.sort(sort_key, dim=-1) + idx = torch.where(sort_key >= n_blocks, torch.full_like(sort_key, -1), sort_key) + if k < topk: + pad = torch.full( + (total_q, num_kv_heads, topk - k), + -1, + dtype=idx.dtype, + device=device, + ) + idx = torch.cat([idx, pad], dim=-1) + return idx.to(torch.int32) + + +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@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) + 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=0, + local_blocks=1, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + + 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 = torch.tensor( + [ + [-1] * 16, + list(range(8)) + [-1] * 8, + list(range(16)), + list(range(8)) + list(range(num_blocks - 13, num_blocks - 5)), + list(range(8)) + list(range(num_blocks - 12, num_blocks - 4)), + ], + device="cuda", + dtype=torch.int32, + ) + expected = expected[:, None, :].expand(-1, scores.shape[0], -1) + 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() + + +def test_fused_selector_head_major_output_through_msa_q2k_consumer(): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 (Blackwell) required") + + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_package_available, + ) + + if not msa_package_available(): + pytest.skip("fmha_sm100 (MSA) not importable") + + from fmha_sm100.sparse_fmha_adapter import _convert_kv_block_indexes_to_q2k + + total_q, num_kv_heads, num_blocks = 7, 4, 96 + scores = torch.randn(num_kv_heads, num_blocks, total_q, device="cuda") + n_valid_blocks = torch.arange(32, 32 + total_q, device="cuda", dtype=torch.int32) + selected = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + head_major_output=True, + ) + q2k = _convert_kv_block_indexes_to_q2k( + selected, + num_kv_heads=num_kv_heads, + num_qo_heads=num_kv_heads, + qhead_per_kv=1, + ) + + assert q2k.shape == (num_kv_heads, total_q, 16) + assert q2k.is_contiguous() + assert q2k.data_ptr() == selected.data_ptr() + assert q2k.untyped_storage().data_ptr() == selected.untyped_storage().data_ptr() + assert torch.equal(q2k.permute(1, 0, 2), selected) + + +@pytest.mark.parametrize( + ("init_blocks", "local_blocks"), + [(0, 0), (0, 1), (2, 3), (16, 1), (20, 0), (0, 20)], +) +def test_fused_selector_matches_reference_forced_and_padded(init_blocks, local_blocks): + scores = ( + torch.tensor( + [ + [ + float("-inf"), + 4.0, + 3.0, + 2.0, + 1.0, + 0.0, + -1.0, + -2.0, + -3.0, + -4.0, + -5.0, + -6.0, + -7.0, + -8.0, + -9.0, + -10.0, + -11.0, + -12.0, + -13.0, + -14.0, + ] + ], + device="cuda", + dtype=torch.float32, + ) + .unsqueeze(-1) + .expand(-1, -1, 5) + ) + n_valid_blocks = torch.tensor([0, 1, 7, 16, 20], dtype=torch.int32) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=init_blocks, + local_blocks=local_blocks, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=init_blocks, + local_blocks=local_blocks, + ) + + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("fill_value", [0.0, 1.0e30, 1.0e29]) +def test_fused_selector_matches_reference_equal_score_ties(fill_value): + scores = torch.full((2, 64, 3), fill_value, device="cuda", dtype=torch.float32) + n_valid_blocks = torch.tensor([15, 32, 64], dtype=torch.int32) + + expected = torch.tensor( + [ + list(range(15)) + [-1], + list(range(16)), + list(range(16)), + ], + device="cuda", + dtype=torch.int32, + ) + expected = expected[:, None, :].expand(-1, scores.shape[0], -1) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=20, + local_blocks=0, + ) + + assert torch.equal(actual, expected) + + +def test_fused_selector_matches_reference_nonfinite_and_validity_bounds(): + scores = ( + torch.tensor( + [ + float("nan"), + float("inf"), + float("-inf"), + -1.0, + 0.0, + 1.0, + float("nan"), + float("-inf"), + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 7.0, + 8.0, + 9.0, + 10.0, + 11.0, + 12.0, + 13.0, + ], + device="cuda", + dtype=torch.float32, + ) + .view(1, 20, 1) + .expand(-1, -1, 4) + ) + n_valid_blocks = torch.tensor([-3, 0, 17, 25], dtype=torch.int32) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + + assert torch.equal(actual, expected) + + +def test_fused_selector_supports_strided_scores_and_cuda_validity(): + generator = torch.Generator(device="cuda").manual_seed(7) + backing = torch.randn(2, 73, 22, generator=generator, device="cuda") + scores = backing[:, 1:72:2, ::2] + assert not scores.is_contiguous() + n_valid_blocks = torch.tensor( + [0, 1, 3, 8, 15, 16, 17, 20, 30, 35, 36], device="cuda", dtype=torch.int32 + ) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=2, + local_blocks=3, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=2, + local_blocks=3, + ) + + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize( + ("case", "match"), + [ + ("scores_dtype", "expects float32 scores"), + ("validity_dtype", "expects int32 n_valid_blocks"), + ("validity_contiguous", "n_valid_blocks must be contiguous"), + ("topk", "supports topk=16"), + ("negative_init", "init_blocks must be non-negative"), + ("negative_local", "local_blocks must be non-negative"), + ("too_many_blocks", "supports at most 65535 blocks"), + ], +) +def test_fused_selector_rejects_invalid_operator_contracts(case, match): + scores = torch.randn(1, 32, 2, device="cuda") + n_valid_blocks = torch.tensor([16, 32], device="cuda", dtype=torch.int32) + topk, init_blocks, local_blocks = 16, 0, 1 + + if case == "scores_dtype": + scores = scores.to(torch.bfloat16) + elif case == "validity_dtype": + n_valid_blocks = n_valid_blocks.to(torch.int64) + elif case == "validity_contiguous": + n_valid_blocks = torch.tensor([16, 0, 32, 0], device="cuda", dtype=torch.int32)[::2] + elif case == "topk": + topk = 8 + elif case == "negative_init": + init_blocks = -1 + elif case == "negative_local": + local_blocks = -1 + elif case == "too_many_blocks": + scores = torch.empty(1, 65_536, 2, device="cuda") + + with pytest.raises(RuntimeError, match=match): + torch.ops.trtllm.minimax_m3_select_blocks( + scores, + n_valid_blocks, + topk, + init_blocks, + local_blocks, + False, + ) + + +def test_fused_selector_cuda_graph_replay_updates_inputs(): + scores = torch.randn(1, 64, 4, device="cuda") + n_valid_blocks = torch.tensor([16, 32, 48, 64], device="cuda", dtype=torch.int32) + + for _ in range(3): + output = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + + scores.copy_(torch.arange(64, device="cuda", dtype=torch.float32).view(1, 64, 1)) + graph.replay() + torch.cuda.synchronize() + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=0, + local_blocks=1, + ) + assert torch.equal(output, expected)