diff --git a/.gitmodules b/.gitmodules index e69de29bb2d1..627760b34da2 100644 --- a/.gitmodules +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "3rdparty/MSA"] + path = 3rdparty/MSA + url = https://gitlab.com/nvidia/tensorrt-llm/oss-components/msa.git diff --git a/3rdparty/MSA b/3rdparty/MSA new file mode 160000 index 000000000000..e2ebe7656649 --- /dev/null +++ b/3rdparty/MSA @@ -0,0 +1 @@ +Subproject commit e2ebe7656649f619af0ad1d457b534283034655e diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch new file mode 100644 index 000000000000..743587c5608b --- /dev/null +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -0,0 +1,184 @@ +diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py +index d72b17a..e27a74f 100644 +--- a/python/fmha_sm100/cute/interface.py ++++ b/python/fmha_sm100/cute/interface.py +@@ -136,6 +136,32 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): + return k, v + + ++def _prepare_paged_hnd_input(tensor: torch.Tensor, blk_kv: int) -> torch.Tensor: ++ """Keep TMA-compatible paged HND views strided across physical pages. ++ ++ The sparse prefill kernel imports the runtime tensor strides through ++ DLPack. It requires each ``[page_size, head_dim]`` head plane to be packed, ++ but the outer physical-page stride may include other coalesced cache roles. ++ Materialize every other layout to preserve the public API's old contract. ++ """ ++ if tensor.ndim != 4 or int(tensor.shape[2]) != int(blk_kv): ++ return tensor.contiguous() ++ ++ head_dim = int(tensor.shape[3]) ++ page_size = int(tensor.shape[2]) ++ packed_within_page = ( ++ tensor.stride(3) == 1 ++ and tensor.stride(2) == head_dim ++ and tensor.stride(1) == page_size * head_dim ++ ) ++ alignment_bytes = 16 ++ aligned_for_tma = ( ++ tensor.data_ptr() % alignment_bytes == 0 ++ and tensor.stride(0) * tensor.element_size() % alignment_bytes == 0 ++ ) ++ return tensor if packed_within_page and aligned_for_tma else tensor.contiguous() ++ ++ + def _validate_cu_seqlens( + cu_seqlens: torch.Tensor, + *, +@@ -736,10 +762,21 @@ def sparse_atten_func( + max_seqlen_q = int(max_seqlen_q) + max_seqlen_k = int(max_seqlen_k) + ++ k_input = ( ++ k.contiguous() ++ if page_table is None ++ else _prepare_paged_hnd_input(k, blk_kv) ++ ) ++ v_input = ( ++ v.contiguous() ++ if page_table is None ++ else _prepare_paged_hnd_input(v, blk_kv) ++ ) ++ + return _sparse_atten_csr_varlen_forward( + q.contiguous(), +- k.contiguous(), +- v.contiguous(), ++ k_input, ++ v_input, + k2q_row_ptr.contiguous(), + k2q_q_indices.contiguous(), + int(topK), +diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py +index 21c777e..c22beef 100644 +--- a/python/fmha_sm100/cute/test_sparse_atten.py ++++ b/python/fmha_sm100/cute/test_sparse_atten.py +@@ -61,6 +61,43 @@ DECODE_DIM = 128 + DECODE_KV_TOKEN_SWEEP = tuple(2**exp for exp in range(3, 21)) + + ++def test_prepare_paged_hnd_input_keeps_aligned_outer_page_stride() -> None: ++ pages, roles, heads, page_size, head_dim = 5, 3, 2, 128, 128 ++ pool = torch.empty( ++ pages, ++ roles, ++ heads, ++ page_size, ++ head_dim, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = pool[:, 1] ++ ++ assert not view.is_contiguous() ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) ++ assert prepared.data_ptr() == view.data_ptr() ++ assert prepared.stride() == view.stride() ++ ++ ++def test_prepare_paged_hnd_input_materializes_unpacked_tokens() -> None: ++ pages, heads, page_size, head_dim = 5, 2, 128, 128 ++ storage = torch.empty( ++ pages, ++ heads, ++ page_size * 2, ++ head_dim, ++ dtype=torch.float8_e4m3fn, ++ device="cuda", ++ ) ++ view = storage[:, :, ::2, :] ++ ++ prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) ++ assert prepared.is_contiguous() ++ assert prepared.data_ptr() != view.data_ptr() ++ torch.testing.assert_close(prepared, view, rtol=0, atol=0) ++ ++ + @contextmanager + def _nvtx_range(message: str): + torch.cuda.nvtx.range_push(message) +@@ -1786,6 +1823,74 @@ def test_sparse_page_atten( + + _assert_forward_close(out, out_ref, out_pt.float(), lse, lse_ref) + ++ ++def test_sparse_page_atten_strided_outer_page_matches_packed() -> None: ++ inputs = _build_paged_inputs( ++ batch=1, ++ seqlen_q=2048, ++ seqlen_kv=2048, ++ head_kv=2, ++ qhead_per_kv=16, ++ dim=128, ++ topk=16, ++ blk_kv=128, ++ causal=True, ++ page_size=128, ++ seqused_trim=0, ++ dtype=torch.float8_e4m3fn, ++ ) ++ k_packed = inputs["k_paged"].detach().clone() ++ v_packed = inputs["v_paged"].detach().clone() ++ pool = torch.empty( ++ k_packed.shape[0], ++ 4, ++ *k_packed.shape[1:], ++ dtype=k_packed.dtype, ++ device=k_packed.device, ++ ) ++ k_strided = pool[:, 1] ++ v_strided = pool[:, 3] ++ k_strided.copy_(k_packed) ++ v_strided.copy_(v_packed) ++ ++ assert not k_strided.is_contiguous() ++ assert not v_strided.is_contiguous() ++ assert ( ++ sparse_interface._prepare_paged_hnd_input(k_strided, inputs["blk_kv"]).data_ptr() ++ == k_strided.data_ptr() ++ ) ++ assert ( ++ sparse_interface._prepare_paged_hnd_input(v_strided, inputs["blk_kv"]).data_ptr() ++ == v_strided.data_ptr() ++ ) ++ ++ def run(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ++ return sparse_atten_func( ++ inputs["q"], ++ k, ++ v, ++ inputs["k2q_row_ptr"], ++ inputs["k2q_q_indices"], ++ 16, ++ blk_kv=inputs["blk_kv"], ++ causal=True, ++ softmax_scale=inputs["softmax_scale"], ++ return_softmax_lse=True, ++ cu_seqlens_q=inputs["cu_seqlens_q"], ++ cu_seqlens_k=inputs["cu_seqlens_k"], ++ max_seqlen_q=inputs["max_seqlen_q"], ++ max_seqlen_k=inputs["max_seqlen_k"], ++ page_table=inputs["page_table"], ++ seqused_k=inputs["seqused_k"], ++ schedule=inputs["schedule"], ++ ) ++ ++ packed_out, packed_lse = run(k_packed, v_packed) ++ strided_out, strided_lse = run(k_strided, v_strided) ++ torch.testing.assert_close(strided_out, packed_out, rtol=0, atol=0) ++ torch.testing.assert_close(strided_lse, packed_lse, rtol=0, atol=0) ++ ++ + @pytest.mark.parametrize("paged", [False, True]) + @pytest.mark.parametrize("causal", [True]) + @pytest.mark.parametrize("batch", [3]) diff --git a/LICENSE b/LICENSE index 8ba867f30567..ed5ca3b261ba 100644 --- a/LICENSE +++ b/LICENSE @@ -19,6 +19,13 @@ Original Source: https://github.com/Dao-AILab/causal-conv1d Copyright (c) 2024, Tri Dao. Licensed under the BSD 3-Clause License +-------------------------------------------------------------------------------- +CUTLASS +-------------------------------------------------------------------------------- +Original Source: https://github.com/NVIDIA/cutlass +Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Licensed under the BSD 3-Clause License + -------------------------------------------------------------------------------- flash-attention -------------------------------------------------------------------------------- @@ -34,6 +41,14 @@ Original Source: https://github.com/fla-org/flash-linear-attention Copyright (c) 2023-2025 Songlin Yang Licensed under the MIT License +-------------------------------------------------------------------------------- +FlashInfer +-------------------------------------------------------------------------------- +Original Source: https://github.com/flashinfer-ai/flashinfer +Copyright 2025-2026 NVIDIA +Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) +Licensed under the Apache License 2.0 + -------------------------------------------------------------------------------- InstructEval -------------------------------------------------------------------------------- @@ -59,6 +74,14 @@ Original Source: https://github.com/state-spaces/mamba Copyright 2023 Tri Dao, Albert Gu Licensed under the Apache License 2.0 +-------------------------------------------------------------------------------- +MSA (MiniMax Sparse Attention) +-------------------------------------------------------------------------------- +Original Source: https://github.com/MiniMax-AI/MSA +Copyright (c) 2026 MiniMax +Licensed under the MIT License + + -------------------------------------------------------------------------------- Quack -------------------------------------------------------------------------------- diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index e1685c7c4ba5..5f175a33723c 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -216,12 +216,20 @@ struct VmmDescSplitter /// For non-VRAM or addresses not in the map, descs pass through unchanged. [[nodiscard]] static MemoryDescs splitDescsWithRegionMap(MemoryDescs const& descs, VramRegionMap const& regionMap); - /// @brief Split paired src/dst descs using local and remote region maps. - /// src is split by localRegionMap, dst is split by remoteRegionMap. - /// The final piece size is min(srcPiece, dstPiece, remaining). - [[nodiscard]] static std::pair splitTransferDescsWithRegionMaps( - MemoryDescs const& srcDescs, MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap, - VramRegionMap const& remoteRegionMap); + /// @brief Split paired src/dst descs at chunk boundaries, then coalesce contiguous pieces. + /// src is split by localRegionMap, dst is split by remoteRegionMap; each piece size is + /// min(srcPiece, dstPiece, remaining). Pairs are sorted by src address, and adjacent pieces + /// whose src AND dst are both contiguous (same deviceId) are merged — but a merged desc never + /// crosses a chunk boundary on either side, and never spans two distinct regions, so every + /// output desc stays within a single registered memory region. Merging requires region + /// metadata: a piece whose address misses the region map on either side is never merged, + /// because two unknown regions are indistinguishable and a merge could cross a chunk or + /// registration boundary. With no region metadata the result is split-only. Non-kVRAM descs + /// pass through unchanged (no region info is available to bound the merge). + /// @param enableCoalesce When false, only split at chunk boundaries without merging pieces. + [[nodiscard]] static std::pair splitAndCoalesceTransferDescs(MemoryDescs const& srcDescs, + MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap, VramRegionMap const& remoteRegionMap, + bool enableCoalesce = true); /// @brief Split VRAM descs at VMM chunk boundaries detected via cuMemGetAddressRange. /// For cudaMalloc memory (single allocation), descs pass through unchanged. diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 14c637699584..9dcf6c81634e 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -517,10 +517,10 @@ uint16_t getEnvNixlPort() return nixlPort; } -bool getEnvNixlEnableCoalesce() +bool getEnvNixlDisableCoalesce() { - static bool const enableCoalesce = getBoolEnv("TRTLLM_NIXL_ENABLE_COALESCE"); - return enableCoalesce; + static bool const disableCoalesce = getBoolEnv("TRTLLM_NIXL_DISABLE_COALESCE"); + return disableCoalesce; } bool getEnvDisaggBenchmarkGenOnly() diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 37b5934b4515..13ad0399d574 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -151,7 +151,8 @@ bool getEnvKVCachePoolUseFabricMemory(); uint16_t getEnvNixlPort(); -bool getEnvNixlEnableCoalesce(); +// Whether to disable coalescing of contiguous NIXL transfer descriptors (coalescing is on by default). +bool getEnvNixlDisableCoalesce(); bool getEnvDisaggBenchmarkGenOnly(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 711750f84f54..5ab589d7ca75 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -347,164 +347,6 @@ NixlTransferStatus::~NixlTransferStatus() noexcept } } -[[nodiscard]] MemoryDescs NixlHelper::coalesceMemoryDescs(MemoryDescs const& descs) -{ - auto const& descVec = descs.getDescs(); - - // If empty or single element, return as-is - if (descVec.size() <= 1) - { - return descs; - } - - size_t const numDescs = descVec.size(); - - // Create index array and sort by address - std::vector sortedIndices(numDescs); - std::iota(sortedIndices.begin(), sortedIndices.end(), 0); - - std::sort(sortedIndices.begin(), sortedIndices.end(), - [&descVec](size_t lhs, size_t rhs) - { - // Sort by deviceId first, then by address - if (descVec[lhs].getDeviceId() != descVec[rhs].getDeviceId()) - { - return descVec[lhs].getDeviceId() < descVec[rhs].getDeviceId(); - } - return descVec[lhs].getAddr() < descVec[rhs].getAddr(); - }); - - std::vector coalesced; - coalesced.reserve(numDescs); - - // Start with the first entry - size_t firstIdx = sortedIndices[0]; - uintptr_t currentAddr = descVec[firstIdx].getAddr(); - size_t currentLen = descVec[firstIdx].getLen(); - uint32_t currentDeviceId = descVec[firstIdx].getDeviceId(); - - for (size_t idx = 1; idx < numDescs; ++idx) - { - size_t sortedIdx = sortedIndices[idx]; - auto const& desc = descVec[sortedIdx]; - - // Check if current can be coalesced with previous - bool isContiguous = (currentAddr + currentLen == desc.getAddr()) && (currentDeviceId == desc.getDeviceId()); - - if (isContiguous) - { - // Coalesce: extend the current region - currentLen += desc.getLen(); - } - else - { - // Cannot coalesce: save the current region and start a new one - coalesced.emplace_back(currentAddr, currentLen, currentDeviceId); - - currentAddr = desc.getAddr(); - currentLen = desc.getLen(); - currentDeviceId = desc.getDeviceId(); - } - } - - // Add the last region - coalesced.emplace_back(currentAddr, currentLen, currentDeviceId); - - TLLM_LOG_DEBUG("NixlHelper::coalesceMemoryDescs: coalesced %zu -> %zu entries", descVec.size(), coalesced.size()); - - return MemoryDescs{descs.getType(), std::move(coalesced)}; -} - -[[nodiscard]] std::pair NixlHelper::coalesceTransferDescs( - TransferDescs const& srcDescs, TransferDescs const& dstDescs) -{ - auto const& srcVec = srcDescs.getDescs(); - auto const& dstVec = dstDescs.getDescs(); - - // If sizes don't match or empty, return as-is - if (srcVec.size() != dstVec.size() || srcVec.empty()) - { - return {srcDescs, dstDescs}; - } - - size_t const numDescs = srcVec.size(); - - // Create index array and sort by src address - // This allows us to find contiguous regions even if the original order is scattered - std::vector sortedIndices(numDescs); - std::iota(sortedIndices.begin(), sortedIndices.end(), 0); - - std::sort(sortedIndices.begin(), sortedIndices.end(), - [&srcVec](size_t lhs, size_t rhs) - { - // Sort by deviceId first, then by address - if (srcVec[lhs].getDeviceId() != srcVec[rhs].getDeviceId()) - { - return srcVec[lhs].getDeviceId() < srcVec[rhs].getDeviceId(); - } - return srcVec[lhs].getAddr() < srcVec[rhs].getAddr(); - }); - - std::vector coalescedSrc; - std::vector coalescedDst; - coalescedSrc.reserve(numDescs); - coalescedDst.reserve(numDescs); - - // Start with the first entry (using sorted order) - size_t firstIdx = sortedIndices[0]; - uintptr_t currentSrcAddr = srcVec[firstIdx].getAddr(); - size_t currentSrcLen = srcVec[firstIdx].getLen(); - uint32_t currentSrcDeviceId = srcVec[firstIdx].getDeviceId(); - - uintptr_t currentDstAddr = dstVec[firstIdx].getAddr(); - size_t currentDstLen = dstVec[firstIdx].getLen(); - uint32_t currentDstDeviceId = dstVec[firstIdx].getDeviceId(); - - for (size_t idx = 1; idx < numDescs; ++idx) - { - size_t sortedIdx = sortedIndices[idx]; - auto const& src = srcVec[sortedIdx]; - auto const& dst = dstVec[sortedIdx]; - - // Check if current src and dst can be coalesced with previous - bool srcContiguous - = (currentSrcAddr + currentSrcLen == src.getAddr()) && (currentSrcDeviceId == src.getDeviceId()); - bool dstContiguous - = (currentDstAddr + currentDstLen == dst.getAddr()) && (currentDstDeviceId == dst.getDeviceId()); - - if (srcContiguous && dstContiguous) - { - // Coalesce: extend the current region - currentSrcLen += src.getLen(); - currentDstLen += dst.getLen(); - } - else - { - // Cannot coalesce: save the current region and start a new one - coalescedSrc.emplace_back(currentSrcAddr, currentSrcLen, currentSrcDeviceId); - coalescedDst.emplace_back(currentDstAddr, currentDstLen, currentDstDeviceId); - - currentSrcAddr = src.getAddr(); - currentSrcLen = src.getLen(); - currentSrcDeviceId = src.getDeviceId(); - - currentDstAddr = dst.getAddr(); - currentDstLen = dst.getLen(); - currentDstDeviceId = dst.getDeviceId(); - } - } - - // Don't forget to add the last region - coalescedSrc.emplace_back(currentSrcAddr, currentSrcLen, currentSrcDeviceId); - coalescedDst.emplace_back(currentDstAddr, currentDstLen, currentDstDeviceId); - - TLLM_LOG_DEBUG( - "NixlHelper::coalesceTransferDescs: coalesced %zu -> %zu transfer entries", srcVec.size(), coalescedSrc.size()); - - return {MemoryDescs{srcDescs.getType(), std::move(coalescedSrc)}, - MemoryDescs{dstDescs.getType(), std::move(coalescedDst)}}; -} - TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { auto startTime = std::chrono::steady_clock::now(); @@ -693,12 +535,8 @@ void NixlTransferAgent::registerMemory(RegisterDescs const& descs) auto detectedRegionMap = VmmDescSplitter::detectVramRegionMap(descs); mLocalVramRegionInfo.merge(detectedRegionMap); - // Coalesce contiguous memory regions to reduce registration overhead (disabled by default) - // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization - auto coalescedDescs = common::getEnvNixlEnableCoalesce() ? NixlHelper::coalesceMemoryDescs(splitDescs) : splitDescs; - nixl_status_t status; - status = mRawAgent->registerMem(NixlHelper::convertRegDlist(coalescedDescs), &mExtraParams); + status = mRawAgent->registerMem(NixlHelper::convertRegDlist(splitDescs), &mExtraParams); TLLM_CHECK(status == NIXL_SUCCESS); std::string localMD; @@ -713,12 +551,8 @@ void NixlTransferAgent::deregisterMemory(RegisterDescs const& descs) // Split using per-region registry info to match what was registered auto splitDescs = VmmDescSplitter::splitDescsWithRegionMap(descs, mLocalVramRegionInfo); - // Coalesce contiguous memory regions to match what was registered (disabled by default) - // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization - auto coalescedDescs = common::getEnvNixlEnableCoalesce() ? NixlHelper::coalesceMemoryDescs(splitDescs) : splitDescs; - nixl_status_t status; - status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(coalescedDescs), &mExtraParams); + status = mRawAgent->deregisterMem(NixlHelper::convertRegDlist(splitDescs), &mExtraParams); TLLM_CHECK(status == NIXL_SUCCESS); // Remove entries from registry @@ -743,7 +577,7 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, AgentDesc const name == remoteName, "loadRemoteAgent gets error agent name: %s != %s", name.c_str(), remoteName.c_str()); // Store remote VMM region info for chunk boundary calculations in - // VmmDescSplitter::splitTransferDescsWithRegionMaps. Per-agent map because different remote agents may have + // VmmDescSplitter::splitAndCoalesceTransferDescs. Per-agent map because different remote agents may have // overlapping virtual addresses. auto const& regions = agentDesc.getVramRegions(); if (!regions.empty()) @@ -764,14 +598,13 @@ AgentDesc NixlTransferAgent::getLocalAgentDesc() nixl_status_t status = mRawAgent->getLocalMD(nixlBlob); TLLM_CHECK(status == NIXL_SUCCESS); - // Pack local VMM region info so remote agents can compute chunk boundaries. + // Pack ALL local region info (VMM multi-chunk and single-allocation alike) so remote agents can + // compute chunk boundaries and never coalesce transfer descs across separately registered regions. std::vector regions; + regions.reserve(mLocalVramRegionInfo.size()); for (auto const& [base, info] : mLocalVramRegionInfo) { - if (info.chunkSize > 0) - { - regions.push_back({base, info.totalLen, info.chunkSize}); - } + regions.push_back({base, info.totalLen, info.chunkSize}); } return AgentDesc{nixlBlob, std::move(regions)}; @@ -809,32 +642,25 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) { reqParams.hasNotif = false; } - // Split transfer descriptors at VMM chunk boundaries to match registered memory. - // Both src and dst are split at chunk boundaries to ensure each descriptor - // falls within a single registered memory region on both local and remote sides. - // Find remote agent's VMM region map (empty map if not found). + // Split transfer descriptors at VMM chunk boundaries to match registered memory, then coalesce + // contiguous pieces. A coalesced descriptor never crosses a chunk boundary or a registered + // region boundary on either side, so every descriptor still falls within a single registered + // memory region on both local and remote sides. Set TRTLLM_NIXL_DISABLE_COALESCE=1 to fall back + // to split-only descriptors. Find remote agent's region map (empty map if not found — e.g. the + // peer's AgentDesc carried no region info; addresses missing from a map are never coalesced, + // so an empty remote map degrades to split-only rather than risking merges across unknown + // remote chunk/registration boundaries). static VramRegionMap const kEmptyMap; auto remoteIt = mRemoteVramRegionInfo.find(request.getRemoteName()); auto const& remoteRegionMap = (remoteIt != mRemoteVramRegionInfo.end()) ? remoteIt->second : kEmptyMap; - auto [splitSrc, splitDst] = VmmDescSplitter::splitTransferDescsWithRegionMaps( - request.getSrcDescs(), request.getDstDescs(), mLocalVramRegionInfo, remoteRegionMap); + auto [xferSrc, xferDst] = VmmDescSplitter::splitAndCoalesceTransferDescs(request.getSrcDescs(), + request.getDstDescs(), mLocalVramRegionInfo, remoteRegionMap, !common::getEnvNixlDisableCoalesce()); - // Coalesce contiguous memory regions to reduce transfer count (disabled by default) - // This matches the coalescing done during registerMemory() - // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization - if (common::getEnvNixlEnableCoalesce()) - { - NVTX3_SCOPED_RANGE(coalesceTransferDescs_CreateXferReq); - auto [coalescedSrc, coalescedDst] = NixlHelper::coalesceTransferDescs(splitSrc, splitDst); - status - = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(coalescedSrc), - NixlHelper::convertXferDist(coalescedDst), request.getRemoteName(), handle, &reqParams); - } - else { - status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(splitSrc), - NixlHelper::convertXferDist(splitDst), request.getRemoteName(), handle, &reqParams); + NVTX3_SCOPED_RANGE(createXferReq); + status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(xferSrc), + NixlHelper::convertXferDist(xferDst), request.getRemoteName(), handle, &reqParams); } TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index cb371f02439e..31e62fd6d822 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -39,21 +39,6 @@ struct NixlHelper [[nodiscard]] static nixl_xfer_dlist_t convertXferDist(FileDescs const& descs); static void posixGpuToFileFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs); static void posixFileToGpuFallback(MemoryDescs const& memoryDesc, FileDescs const& fileDescs); - - /// @brief Coalesce contiguous memory regions to reduce memory registration overhead. - /// Adjacent memory regions with the same deviceId will be merged into a single region. - /// @param descs Memory descriptors to coalesce - /// @return Coalesced MemoryDescs - [[nodiscard]] static MemoryDescs coalesceMemoryDescs(MemoryDescs const& descs); - - /// @brief Coalesce contiguous memory regions in src and dst to reduce transfer count. - /// If src[i] and src[i+1] are contiguous, and dst[i] and dst[i+1] are also contiguous - /// (with same deviceId), they will be merged into a single transfer. - /// @param srcDescs Source memory descriptors - /// @param dstDescs Destination memory descriptors - /// @return Pair of coalesced (src, dst) MemoryDescs - [[nodiscard]] static std::pair coalesceTransferDescs( - TransferDescs const& srcDescs, TransferDescs const& dstDescs); }; class NixlTransferStatus final : public TransferStatus diff --git a/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp index fd60651dc6a0..147c99b37329 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -19,9 +19,12 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/executor/serializeUtils.h" +#include #include #include +#include #include +#include namespace tensorrt_llm::executor::kv_cache { @@ -152,8 +155,9 @@ MemoryDescs VmmDescSplitter::splitDescsWithRegionMap(MemoryDescs const& descs, V return MemoryDescs{descs.getType(), std::move(result)}; } -std::pair VmmDescSplitter::splitTransferDescsWithRegionMaps(MemoryDescs const& srcDescs, - MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap, VramRegionMap const& remoteRegionMap) +std::pair VmmDescSplitter::splitAndCoalesceTransferDescs(MemoryDescs const& srcDescs, + MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap, VramRegionMap const& remoteRegionMap, + bool enableCoalesce) { if (srcDescs.getType() != MemoryType::kVRAM) return {srcDescs, dstDescs}; @@ -161,54 +165,143 @@ std::pair VmmDescSplitter::splitTransferDescsWithRegio auto const& srcVec = srcDescs.getDescs(); auto const& dstVec = dstDescs.getDescs(); TLLM_CHECK(srcVec.size() == dstVec.size()); + if (srcVec.empty()) + return {srcDescs, dstDescs}; - std::vector splitSrc, splitDst; - splitSrc.reserve(srcVec.size()); - splitDst.reserve(dstVec.size()); - - for (size_t i = 0; i < srcVec.size(); ++i) + // Sort pair indices by (src deviceId, src addr) so pairs that are contiguous in memory become + // adjacent, maximizing coalescing regardless of input order. Pairs are independent transfers, + // so reordering is safe. The sort only serves coalescing; with it disabled, keep input order. + std::vector order(srcVec.size()); + std::iota(order.begin(), order.end(), 0); + if (enableCoalesce) { - auto [srcChunkSize, srcBase] = lookupChunkInfo(srcVec[i].getAddr(), localRegionMap); - auto [dstChunkSize, dstBase] = lookupChunkInfo(dstVec[i].getAddr(), remoteRegionMap); + std::sort(order.begin(), order.end(), + [&srcVec](size_t lhs, size_t rhs) + { + if (srcVec[lhs].getDeviceId() != srcVec[rhs].getDeviceId()) + { + return srcVec[lhs].getDeviceId() < srcVec[rhs].getDeviceId(); + } + return srcVec[lhs].getAddr() < srcVec[rhs].getAddr(); + }); + } + + std::vector outSrc, outDst; + outSrc.reserve(srcVec.size()); + outDst.reserve(dstVec.size()); + + // Region info of the last emitted piece. Invariant: every emitted desc lies within a single + // chunk on both sides, so a merge is legal iff the new piece is contiguous, in the same region, + // and does not start on a chunk boundary (starting on a boundary means the merge would cross it). + size_t prevSrcChunkSize = 0, prevDstChunkSize = 0; + uintptr_t prevSrcBase = 0, prevDstBase = 0; - // If neither side is multi-chunk VMM, no splitting is needed. - if (srcChunkSize == 0 && dstChunkSize == 0) + auto emitPiece + = [&](uintptr_t srcAddr, uintptr_t dstAddr, size_t len, uint32_t srcDev, uint32_t dstDev, size_t srcChunkSize, + uintptr_t srcBase, size_t dstChunkSize, uintptr_t dstBase, bool regionsKnown) + { + if (enableCoalesce && regionsKnown && !outSrc.empty()) { - splitSrc.push_back(srcVec[i]); - splitDst.push_back(dstVec[i]); - continue; + auto const& lastSrc = outSrc.back(); + auto const& lastDst = outDst.back(); + bool contiguous = lastSrc.getAddr() + lastSrc.getLen() == srcAddr && lastSrc.getDeviceId() == srcDev + && lastDst.getAddr() + lastDst.getLen() == dstAddr && lastDst.getDeviceId() == dstDev; + bool sameSrcRegion = srcChunkSize == prevSrcChunkSize && srcBase == prevSrcBase; + bool sameDstRegion = dstChunkSize == prevDstChunkSize && dstBase == prevDstBase; + bool srcWithinChunk = srcChunkSize == 0 || (srcAddr - srcBase) % srcChunkSize != 0; + bool dstWithinChunk = dstChunkSize == 0 || (dstAddr - dstBase) % dstChunkSize != 0; + if (contiguous && sameSrcRegion && sameDstRegion && srcWithinChunk && dstWithinChunk) + { + outSrc.back() = MemoryDesc{lastSrc.getAddr(), lastSrc.getLen() + len, srcDev}; + outDst.back() = MemoryDesc{lastDst.getAddr(), lastDst.getLen() + len, dstDev}; + return; + } + } + outSrc.emplace_back(srcAddr, len, srcDev); + outDst.emplace_back(dstAddr, len, dstDev); + prevSrcChunkSize = srcChunkSize; + prevSrcBase = srcBase; + prevDstChunkSize = dstChunkSize; + prevDstBase = dstBase; + }; + + // One-entry region cache per side: after sorting, consecutive pairs almost always fall in the + // same region (typically one KV pool), so the O(log R) map lookup is skipped on cache hits. + struct RegionCache + { + uintptr_t base = 0; + size_t totalLen = 0; + size_t chunkSize = 0; + bool valid = false; + }; + + // Returns {chunkSize, regionBase, found}. A miss means the address is not covered by any + // region metadata (e.g. the peer did not send its region info). Two misses both look like + // {0, 0} yet may belong to two distinct regions, so pieces with a missed lookup on either + // side are never merged — a merge could silently cross a chunk or registration boundary. + auto cachedLookup = [](uintptr_t addr, VramRegionMap const& regionMap, RegionCache& cache) + { + if (cache.valid && addr >= cache.base && addr - cache.base < cache.totalLen) + { + return std::tuple{cache.chunkSize, cache.base, true}; + } + auto it = regionMap.upper_bound(addr); + if (it != regionMap.begin()) + { + --it; + if (addr >= it->first && addr - it->first < it->second.totalLen) + { + cache = {it->first, it->second.totalLen, it->second.chunkSize, true}; + return std::tuple{cache.chunkSize, cache.base, true}; + } } + return std::tuple{0, 0, false}; + }; + + RegionCache srcCache, dstCache; + size_t numPieces = 0; + for (size_t idx : order) + { + auto const& src = srcVec[idx]; + auto const& dst = dstVec[idx]; + auto [srcChunkSize, srcBase, srcFound] = cachedLookup(src.getAddr(), localRegionMap, srcCache); + auto [dstChunkSize, dstBase, dstFound] = cachedLookup(dst.getAddr(), remoteRegionMap, dstCache); - uintptr_t srcAddr = srcVec[i].getAddr(); - uintptr_t dstAddr = dstVec[i].getAddr(); - size_t remaining = srcVec[i].getLen(); + uintptr_t srcAddr = src.getAddr(); + uintptr_t dstAddr = dst.getAddr(); + size_t remaining = src.getLen(); while (remaining > 0) { size_t srcPieceSize = remaining; if (srcChunkSize > 0) { - size_t srcOffsetInChunk = static_cast((srcAddr - srcBase) % srcChunkSize); - srcPieceSize = srcChunkSize - srcOffsetInChunk; + srcPieceSize = srcChunkSize - static_cast((srcAddr - srcBase) % srcChunkSize); } size_t dstPieceSize = remaining; if (dstChunkSize > 0) { - size_t dstOffsetInChunk = static_cast((dstAddr - dstBase) % dstChunkSize); - dstPieceSize = dstChunkSize - dstOffsetInChunk; + dstPieceSize = dstChunkSize - static_cast((dstAddr - dstBase) % dstChunkSize); } size_t pieceSize = std::min({remaining, srcPieceSize, dstPieceSize}); - splitSrc.emplace_back(srcAddr, pieceSize, srcVec[i].getDeviceId()); - splitDst.emplace_back(dstAddr, pieceSize, dstVec[i].getDeviceId()); + emitPiece(srcAddr, dstAddr, pieceSize, src.getDeviceId(), dst.getDeviceId(), srcChunkSize, srcBase, + dstChunkSize, dstBase, srcFound && dstFound); srcAddr += pieceSize; dstAddr += pieceSize; remaining -= pieceSize; + ++numPieces; } } - return {MemoryDescs{srcDescs.getType(), std::move(splitSrc)}, MemoryDescs{dstDescs.getType(), std::move(splitDst)}}; + if (outSrc.size() != srcVec.size()) + { + TLLM_LOG_DEBUG("VmmDescSplitter::splitAndCoalesceTransferDescs: %zu pairs -> %zu pieces -> %zu transfers", + srcVec.size(), numPieces, outSrc.size()); + } + + return {MemoryDescs{srcDescs.getType(), std::move(outSrc)}, MemoryDescs{dstDescs.getType(), std::move(outDst)}}; } MemoryDescs VmmDescSplitter::splitVmmDescs(MemoryDescs const& descs, size_t& detectedChunkSize) diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu new file mode 100644 index 000000000000..0490844cd82c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu @@ -0,0 +1,175 @@ +/* + * 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 float kInitScore = 1.0e30F; +constexpr float kLocalScore = 1.0e29F; + +__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); + } + + 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[static_cast(outputRow) * kTopK + rank] = selectedIndices[rank]; + } + } +} + +} // 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, cudaStream_t stream) +{ + int32_t const numOutputRows = totalQueries * numKvHeads; + if (numOutputRows == 0) + { + return; + } + int32_t const gridSize = (numOutputRows + kWarpsPerBlock - 1) / kWarpsPerBlock; + minimaxM3SelectBlocksKernel<<>>(scores, headStride, blockStride, queryStride, + nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks); +} + +} // 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..23aeaa0e1730 --- /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, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index 39021ce642b1..635407433c90 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -174,9 +174,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 } else if (routingMethodType == RoutingMethodType::MiniMax2) { - // MiniMaxM2: sigmoid(logit) + bias → topK → renormalize un-biased sigmoid scores. - // Similar to DeepSeek no-groups but with routeScale = 1.0 and epsilon = 1e-20 - // to match the Python reference: weight / (sum + 1e-20). + // MiniMaxM2/M3: sigmoid(logit) + bias → topK → renormalize un-biased sigmoid scores. + // MiniMaxM2 uses the default routeScale = 1.0, while MiniMaxM3 supplies its + // model-specific scale. Both use epsilon = 1e-20 to match the Python reference. moe::dev::routing::routingCustom::Data routingData; // @@ -189,7 +189,7 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mPostprocessType = moe::dev::routing::RoutingPostprocessType::ScaledSumNormalize; routingData.mPtrRoutingBias = routingBias; routingData.mDtypeBias = dtypeRoutingBias; - routingData.mRouteScale = 1.0f; + routingData.mRouteScale = routedScalingFactor; routingData.mSumEpsilon = 1e-20f; // Pass-through raw pointer; kernels will cast to the proper InputT based on routing method diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index c628a9a152c0..855147742214 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -110,6 +110,7 @@ add_library( IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.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..6b4c3e3a9a85 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp @@ -0,0 +1,88 @@ +/* + * 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 + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK, + int64_t initBlocks, int64_t localBlocks) +{ + 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"); + + auto output = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32)); + 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); + 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) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_select_blocks", &tensorrt_llm::torch_ext::minimaxM3SelectBlocks); +} diff --git a/cpp/tests/unit_tests/executor/CMakeLists.txt b/cpp/tests/unit_tests/executor/CMakeLists.txt index a51baa6ed00f..a67264ed9c29 100644 --- a/cpp/tests/unit_tests/executor/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-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 @@ -34,6 +34,7 @@ add_gtest(requestWithIdTest requestWithIdTest.cpp) add_gtest(loraConfigTest loraConfigTest.cpp) add_gtest(intervalSetTest intervalSetTest.cpp) add_gtest(dynamicBatchTunerTest dynamicBatchTunerTest.cpp) +add_gtest(coalesceTest coalesceTest.cpp) add_gtest(genUniqueAgentNameTest genUniqueAgentNameTest.cpp) target_link_libraries(genUniqueAgentNameTest PRIVATE ${Python3_LIBRARIES}) add_gtest(ucxCommTest ucxCommTest.cpp) @@ -56,10 +57,6 @@ if(NIXL_ROOT OR (MOONCAKE_ROOT AND NOT IS_ROCKY8)) ${Python3_LIBRARIES}) target_compile_definitions(transferAgentTest PRIVATE TEST_NIXL_BACKEND=1) target_compile_definitions(agentCommTest PRIVATE TEST_NIXL_BACKEND=1) - - add_gtest(coalesceTest coalesceTest.cpp) - target_link_libraries(coalesceTest PRIVATE tensorrt_llm_nixl_wrapper - NIXL::nixl) endif() if(MOONCAKE_ROOT) diff --git a/cpp/tests/unit_tests/executor/coalesceTest.cpp b/cpp/tests/unit_tests/executor/coalesceTest.cpp index 8675d93fc3b4..f115d0fd06db 100644 --- a/cpp/tests/unit_tests/executor/coalesceTest.cpp +++ b/cpp/tests/unit_tests/executor/coalesceTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,154 +15,60 @@ * limitations under the License. */ -#include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" +#include "tensorrt_llm/executor/transferAgent.h" #include using namespace tensorrt_llm::executor::kv_cache; -// ==================== coalesceMemoryDescs tests ==================== - -TEST(CoalesceMemoryDescsTest, EmptyInput) -{ - MemoryDescs descs{MemoryType::kVRAM, {}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - EXPECT_EQ(result.getDescs().size(), 0); -} - -TEST(CoalesceMemoryDescsTest, SingleEntry) -{ - MemoryDescs descs{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 1); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 256); - EXPECT_EQ(result.getDescs()[0].getDeviceId(), 0); -} - -TEST(CoalesceMemoryDescsTest, TwoContiguous) -{ - // [0x1000, 256) then [0x1100, 256) — adjacent, should merge into one - MemoryDescs descs{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 1); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 512); - EXPECT_EQ(result.getDescs()[0].getDeviceId(), 0); -} - -TEST(CoalesceMemoryDescsTest, TwoNonContiguous) +namespace { - // Gap between blocks — should stay as two - MemoryDescs descs{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x2000, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 2); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 256); - EXPECT_EQ(result.getDescs()[1].getAddr(), 0x2000); - EXPECT_EQ(result.getDescs()[1].getLen(), 256); -} +VramRegionMap const kEmptyMap; -TEST(CoalesceMemoryDescsTest, ThreeContiguous) +// Merging requires region metadata: an address that misses its region map is never coalesced. +// Tests exercising merge behavior therefore provide maps covering their addresses; a flat +// region (chunkSize=0) imposes no chunk boundaries within it. +VramRegionMap flatRegion(uintptr_t base, size_t len) { - MemoryDescs descs{ - MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x1200, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 1); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 768); + VramRegionMap map; + map[base] = {len, 0}; + return map; } -TEST(CoalesceMemoryDescsTest, UnsortedInput) +std::pair run(TransferDescs const& src, TransferDescs const& dst, + VramRegionMap const& localMap = kEmptyMap, VramRegionMap const& remoteMap = kEmptyMap) { - // Same three contiguous blocks but in reverse order — sorting should fix it - MemoryDescs descs{ - MemoryType::kVRAM, {MemoryDesc{0x1200, 256, 0}, MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 1); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 768); + return VmmDescSplitter::splitAndCoalesceTransferDescs(src, dst, localMap, remoteMap); } +} // namespace -TEST(CoalesceMemoryDescsTest, DifferentDevices) -{ - // Contiguous addresses but different devices — should NOT merge - MemoryDescs descs{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 1}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 2); -} +// ==================== coalescing within flat regions (chunkSize=0) ==================== -TEST(CoalesceMemoryDescsTest, MixedContiguousAndGaps) -{ - // First two are contiguous, then a gap before the third - MemoryDescs descs{ - MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x3000, 128, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 2); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 512); - EXPECT_EQ(result.getDescs()[1].getAddr(), 0x3000); - EXPECT_EQ(result.getDescs()[1].getLen(), 128); -} - -TEST(CoalesceMemoryDescsTest, MultipleDevicesEachContiguous) -{ - // Two contiguous on device 0, two contiguous on device 1 - MemoryDescs descs{MemoryType::kVRAM, - {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x2000, 128, 1}, - MemoryDesc{0x2080, 128, 1}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 2); - EXPECT_EQ(result.getDescs()[0].getAddr(), 0x1000); - EXPECT_EQ(result.getDescs()[0].getLen(), 512); - EXPECT_EQ(result.getDescs()[0].getDeviceId(), 0); - EXPECT_EQ(result.getDescs()[1].getAddr(), 0x2000); - EXPECT_EQ(result.getDescs()[1].getLen(), 256); - EXPECT_EQ(result.getDescs()[1].getDeviceId(), 1); -} - -TEST(CoalesceMemoryDescsTest, PreservesMemoryType) -{ - MemoryDescs descs{MemoryType::kDRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - EXPECT_EQ(result.getType(), MemoryType::kDRAM); -} - -TEST(CoalesceMemoryDescsTest, AllSeparate) -{ - // Nothing can be merged — all have gaps - MemoryDescs descs{MemoryType::kVRAM, - {MemoryDesc{0x1000, 100, 0}, MemoryDesc{0x2000, 100, 0}, MemoryDesc{0x3000, 100, 0}, - MemoryDesc{0x4000, 100, 0}}}; - auto result = NixlHelper::coalesceMemoryDescs(descs); - ASSERT_EQ(result.getDescs().size(), 4); -} - -// ==================== coalesceTransferDescs tests ==================== - -TEST(CoalesceTransferDescsTest, EmptyInput) +TEST(SplitAndCoalesceTest, EmptyInput) { TransferDescs src{MemoryType::kVRAM, {}}; TransferDescs dst{MemoryType::kVRAM, {}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst); EXPECT_EQ(resSrc.getDescs().size(), 0); EXPECT_EQ(resDst.getDescs().size(), 0); } -TEST(CoalesceTransferDescsTest, SinglePair) +TEST(SplitAndCoalesceTest, SinglePair) { TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst); ASSERT_EQ(resSrc.getDescs().size(), 1); ASSERT_EQ(resDst.getDescs().size(), 1); + EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x1000); + EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x5000); } -TEST(CoalesceTransferDescsTest, BothSidesContiguous) +TEST(SplitAndCoalesceTest, BothSidesContiguous) { - // src contiguous AND dst contiguous — should merge into one transfer + // src contiguous AND dst contiguous within known regions — should merge into one transfer TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 1); ASSERT_EQ(resDst.getDescs().size(), 1); EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x1000); @@ -171,98 +77,79 @@ TEST(CoalesceTransferDescsTest, BothSidesContiguous) EXPECT_EQ(resDst.getDescs()[0].getLen(), 512); } -TEST(CoalesceTransferDescsTest, SrcContiguousDstNot) +TEST(SplitAndCoalesceTest, SrcContiguousDstNot) { - // src is contiguous but dst has a gap — can't merge TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x6000, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x2000)); ASSERT_EQ(resSrc.getDescs().size(), 2); ASSERT_EQ(resDst.getDescs().size(), 2); } -TEST(CoalesceTransferDescsTest, DstContiguousSrcNot) +TEST(SplitAndCoalesceTest, DstContiguousSrcNot) { - // dst is contiguous but src has a gap — can't merge TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x2000, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x2000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 2); ASSERT_EQ(resDst.getDescs().size(), 2); } -TEST(CoalesceTransferDescsTest, DifferentDevicesOnSrc) +TEST(SplitAndCoalesceTest, DifferentDevicesOnSrc) { - // src addresses look contiguous but are on different devices — can't merge TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 1}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 0}, MemoryDesc{0x5100, 256, 0}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 2); ASSERT_EQ(resDst.getDescs().size(), 2); } -TEST(CoalesceTransferDescsTest, DifferentDevicesOnDst) +TEST(SplitAndCoalesceTest, DifferentDevicesOnDst) { - // dst addresses look contiguous but are on different devices — can't merge TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 0}, MemoryDesc{0x5100, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 2); ASSERT_EQ(resDst.getDescs().size(), 2); } -TEST(CoalesceTransferDescsTest, MismatchedSizes) -{ - // src has 2 entries, dst has 1 — sizes don't match, return as-is - TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; - TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); - ASSERT_EQ(resSrc.getDescs().size(), 2); - ASSERT_EQ(resDst.getDescs().size(), 1); -} - -TEST(CoalesceTransferDescsTest, ThreePairsAllContiguous) +TEST(SplitAndCoalesceTest, ThreePairsAllContiguous) { TransferDescs src{ MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x1200, 256, 0}}}; TransferDescs dst{ MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}, MemoryDesc{0x5200, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 1); ASSERT_EQ(resDst.getDescs().size(), 1); EXPECT_EQ(resSrc.getDescs()[0].getLen(), 768); EXPECT_EQ(resDst.getDescs()[0].getLen(), 768); } -TEST(CoalesceTransferDescsTest, PartialMerge) +TEST(SplitAndCoalesceTest, PartialMerge) { - // First two pairs: both sides contiguous — merge - // Third pair: src contiguous but dst has gap — stays separate + // First two pairs merge; third pair's dst has a gap — stays separate TransferDescs src{ MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x1200, 256, 0}}}; TransferDescs dst{ MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}, MemoryDesc{0x9000, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x8000)); ASSERT_EQ(resSrc.getDescs().size(), 2); ASSERT_EQ(resDst.getDescs().size(), 2); - // Merged pair EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x1000); EXPECT_EQ(resSrc.getDescs()[0].getLen(), 512); EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x5000); EXPECT_EQ(resDst.getDescs()[0].getLen(), 512); - // Separate pair EXPECT_EQ(resSrc.getDescs()[1].getAddr(), 0x1200); - EXPECT_EQ(resSrc.getDescs()[1].getLen(), 256); EXPECT_EQ(resDst.getDescs()[1].getAddr(), 0x9000); - EXPECT_EQ(resDst.getDescs()[1].getLen(), 256); } -TEST(CoalesceTransferDescsTest, UnsortedInput) +TEST(SplitAndCoalesceTest, UnsortedInput) { - // Same as BothSidesContiguous but in reverse order — sorting should fix it + // Same as BothSidesContiguous but in reverse order — sorting by src addr should fix it TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1100, 256, 0}, MemoryDesc{0x1000, 256, 0}}}; TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5100, 256, 1}, MemoryDesc{0x5000, 256, 1}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000), flatRegion(0x5000, 0x1000)); ASSERT_EQ(resSrc.getDescs().size(), 1); ASSERT_EQ(resDst.getDescs().size(), 1); EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x1000); @@ -271,11 +158,259 @@ TEST(CoalesceTransferDescsTest, UnsortedInput) EXPECT_EQ(resDst.getDescs()[0].getLen(), 512); } -TEST(CoalesceTransferDescsTest, PreservesMemoryType) +TEST(SplitAndCoalesceTest, NonVramPassthrough) { - TransferDescs src{MemoryType::kDRAM, {MemoryDesc{0x1000, 256, 0}}}; - TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 0}}}; - auto [resSrc, resDst] = NixlHelper::coalesceTransferDescs(src, dst); + // Non-kVRAM descs pass through unchanged: no region info exists to bound a merge + TransferDescs src{MemoryType::kDRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; + TransferDescs dst{MemoryType::kDRAM, {MemoryDesc{0x5000, 256, 0}, MemoryDesc{0x5100, 256, 0}}}; + auto [resSrc, resDst] = run(src, dst); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); EXPECT_EQ(resSrc.getType(), MemoryType::kDRAM); - EXPECT_EQ(resDst.getType(), MemoryType::kVRAM); +} + +// ==================== unknown regions never merge ==================== + +TEST(SplitAndCoalesceTest, NoMergeWithoutRegionMetadata) +{ + // Contiguous on both sides, but neither map covers the addresses: two unknown regions are + // indistinguishable (both look up as a miss), so merging is disabled and pairs stay separate. + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}}}; + auto [resSrc, resDst] = run(src, dst); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); +} + +TEST(SplitAndCoalesceTest, NoMergeWhenRemoteRegionUnknown) +{ + // Local map covers src, but the remote side sent no region info (e.g. an older peer): + // dst lookups miss, so nothing merges even though both sides are contiguous. + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}}}; + auto [resSrc, resDst] = run(src, dst, flatRegion(0x1000, 0x1000)); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); +} + +TEST(SplitAndCoalesceTest, NoMergeWhenLocalRegionUnknown) +{ + // Remote map covers dst, but src addresses miss the local map: no merging. + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 256, 0}, MemoryDesc{0x1100, 256, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x5000, 256, 1}, MemoryDesc{0x5100, 256, 1}}}; + auto [resSrc, resDst] = run(src, dst, kEmptyMap, flatRegion(0x5000, 0x1000)); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); +} + +// ==================== chunk-boundary-constrained coalescing ==================== + +TEST(SplitAndCoalesceTest, MergeStopsAtSrcChunkBoundary) +{ + // Local VMM region: base=0x100000, 4MB total, 2MB chunks → boundary at 0x300000. + VramRegionMap localMap; + localMap[0x100000] = {0x400000, 0x200000}; + + // Four contiguous 1MB pairs covering 4MB on both sides; dst is one flat remote region. + std::vector srcVec, dstVec; + for (size_t i = 0; i < 4; ++i) + { + srcVec.emplace_back(0x100000 + i * 0x100000, 0x100000, 0); + dstVec.emplace_back(0x900000 + i * 0x100000, 0x100000, 1); + } + TransferDescs src{MemoryType::kVRAM, srcVec}; + TransferDescs dst{MemoryType::kVRAM, dstVec}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x400000)); + // Merged per src chunk: two 2MB transfers, split exactly at the chunk boundary. + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); + EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x200000); + EXPECT_EQ(resSrc.getDescs()[1].getAddr(), 0x300000); + EXPECT_EQ(resSrc.getDescs()[1].getLen(), 0x200000); + EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x900000); + EXPECT_EQ(resDst.getDescs()[0].getLen(), 0x200000); + EXPECT_EQ(resDst.getDescs()[1].getAddr(), 0xB00000); + EXPECT_EQ(resDst.getDescs()[1].getLen(), 0x200000); +} + +TEST(SplitAndCoalesceTest, MergeStopsAtDstChunkBoundary) +{ + // Remote VMM region: base=0x800000, 1MB chunks → boundary at 0x900000. + VramRegionMap remoteMap; + remoteMap[0x800000] = {0x400000, 0x100000}; + + // Two contiguous pairs whose dst junction sits exactly on the remote chunk boundary. + auto localMap = flatRegion(0x1000, 0x100000); + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x1000, 0x80000, 0}, MemoryDesc{0x81000, 0x80000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x880000, 0x80000, 1}, MemoryDesc{0x900000, 0x80000, 1}}}; + + // With a flat remote region they would merge into one transfer... + { + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x800000, 0x400000)); + ASSERT_EQ(resSrc.getDescs().size(), 1); + ASSERT_EQ(resDst.getDescs().size(), 1); + } + // ...but the dst chunk boundary must block the merge. + auto [resSrc, resDst] = run(src, dst, localMap, remoteMap); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); + EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x880000); + EXPECT_EQ(resDst.getDescs()[1].getAddr(), 0x900000); +} + +TEST(SplitAndCoalesceTest, SplitPiecesAreNotRemerged) +{ + // A single pair spanning two src chunks stays split even though the pieces are contiguous. + VramRegionMap localMap; + localMap[0x100000] = {0x400000, 0x100000}; + + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x100000, 0x200000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x900000, 0x200000, 1}}}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x400000)); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[1].getLen(), 0x100000); +} + +TEST(SplitAndCoalesceTest, UnalignedRegionBaseBoundary) +{ + // Chunk boundaries are relative to the region base, not absolute alignment. + // base=0x180000, 1MB chunks → boundaries at 0x280000, 0x380000, ... + VramRegionMap localMap; + localMap[0x180000] = {0x300000, 0x100000}; + + // Three contiguous 512KB pairs: first two share the first chunk, third starts a new chunk. + std::vector srcVec, dstVec; + for (size_t i = 0; i < 3; ++i) + { + srcVec.emplace_back(0x180000 + i * 0x80000, 0x80000, 0); + dstVec.emplace_back(0x900000 + i * 0x80000, 0x80000, 1); + } + TransferDescs src{MemoryType::kVRAM, srcVec}; + TransferDescs dst{MemoryType::kVRAM, dstVec}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x300000)); + ASSERT_EQ(resSrc.getDescs().size(), 2); + EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x180000); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[1].getAddr(), 0x280000); + EXPECT_EQ(resSrc.getDescs()[1].getLen(), 0x80000); +} + +TEST(SplitAndCoalesceTest, NoMergeAcrossRegions) +{ + // Two VA-adjacent but distinct local regions (cudaMalloc-style, chunkSize=0): + // contiguous descs must not merge across the region boundary. + VramRegionMap localMap; + localMap[0x100000] = {0x100000, 0}; + localMap[0x200000] = {0x100000, 0}; + + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x100000, 0x100000, 0}, MemoryDesc{0x200000, 0x100000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x900000, 0x100000, 1}, MemoryDesc{0xA00000, 0x100000, 1}}}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x200000)); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); +} + +TEST(SplitAndCoalesceTest, MergeWithinSingleRegion) +{ + // Control for NoMergeAcrossRegions: same layout as one region (chunkSize=0) merges freely. + VramRegionMap localMap; + localMap[0x100000] = {0x200000, 0}; + + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x100000, 0x100000, 0}, MemoryDesc{0x200000, 0x100000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x900000, 0x100000, 1}, MemoryDesc{0xA00000, 0x100000, 1}}}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x200000)); + ASSERT_EQ(resSrc.getDescs().size(), 1); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x200000); + EXPECT_EQ(resDst.getDescs()[0].getLen(), 0x200000); +} + +TEST(SplitAndCoalesceTest, NoMergeAcrossRemoteRegions) +{ + // The remote side registered two discrete but VA-adjacent buffers (chunkSize=0 each): + // contiguous dst descs must not merge across the remote registration boundary. + VramRegionMap remoteMap; + remoteMap[0x900000] = {0x100000, 0}; + remoteMap[0xA00000] = {0x100000, 0}; + + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x100000, 0x100000, 0}, MemoryDesc{0x200000, 0x100000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x900000, 0x100000, 1}, MemoryDesc{0xA00000, 0x100000, 1}}}; + + auto [resSrc, resDst] = run(src, dst, flatRegion(0x100000, 0x200000), remoteMap); + ASSERT_EQ(resSrc.getDescs().size(), 2); + ASSERT_EQ(resDst.getDescs().size(), 2); +} + +TEST(SplitAndCoalesceTest, CoalesceDisabledSplitsOnly) +{ + // With coalescing disabled, contiguous pairs stay separate and chunk splitting still applies. + VramRegionMap localMap; + localMap[0x100000] = {0x400000, 0x100000}; + + // Two contiguous 512KB pairs within one chunk plus one pair spanning a chunk boundary. + TransferDescs src{MemoryType::kVRAM, + {MemoryDesc{0x100000, 0x80000, 0}, MemoryDesc{0x180000, 0x80000, 0}, MemoryDesc{0x200000, 0x200000, 0}}}; + TransferDescs dst{MemoryType::kVRAM, + {MemoryDesc{0x900000, 0x80000, 1}, MemoryDesc{0x980000, 0x80000, 1}, MemoryDesc{0xA00000, 0x200000, 1}}}; + + auto [resSrc, resDst] + = VmmDescSplitter::splitAndCoalesceTransferDescs(src, dst, localMap, kEmptyMap, /*enableCoalesce=*/false); + // No merging: pair 1, pair 2, and pair 3 split into two chunk pieces → 4 descs. + ASSERT_EQ(resSrc.getDescs().size(), 4); + ASSERT_EQ(resDst.getDescs().size(), 4); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x80000); + EXPECT_EQ(resSrc.getDescs()[1].getLen(), 0x80000); + EXPECT_EQ(resSrc.getDescs()[2].getLen(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[3].getLen(), 0x100000); +} + +TEST(SplitAndCoalesceTest, CoalesceDisabledPreservesInputOrder) +{ + // With coalescing disabled there is no sorting either: descs come out in input order, + // matching the historical split-only behavior. + TransferDescs src{MemoryType::kVRAM, {MemoryDesc{0x2000, 256, 0}, MemoryDesc{0x1000, 256, 0}}}; + TransferDescs dst{MemoryType::kVRAM, {MemoryDesc{0x6000, 256, 1}, MemoryDesc{0x5000, 256, 1}}}; + + auto [resSrc, resDst] + = VmmDescSplitter::splitAndCoalesceTransferDescs(src, dst, kEmptyMap, kEmptyMap, /*enableCoalesce=*/false); + ASSERT_EQ(resSrc.getDescs().size(), 2); + EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x2000); + EXPECT_EQ(resSrc.getDescs()[1].getAddr(), 0x1000); + EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x6000); + EXPECT_EQ(resDst.getDescs()[1].getAddr(), 0x5000); +} + +TEST(SplitAndCoalesceTest, SplitAndMergeScatteredBlocks) +{ + // Scattered input order + chunked src region: blocks are sorted, merged per chunk. + // base=0x100000, 1MB chunks; four 512KB blocks given out of order. + VramRegionMap localMap; + localMap[0x100000] = {0x400000, 0x100000}; + + std::vector perm{2, 0, 3, 1}; + std::vector srcVec, dstVec; + for (size_t i : perm) + { + srcVec.emplace_back(0x100000 + i * 0x80000, 0x80000, 0); + dstVec.emplace_back(0x900000 + i * 0x80000, 0x80000, 1); + } + TransferDescs src{MemoryType::kVRAM, srcVec}; + TransferDescs dst{MemoryType::kVRAM, dstVec}; + + auto [resSrc, resDst] = run(src, dst, localMap, flatRegion(0x900000, 0x400000)); + // 2MB of contiguous data over two 1MB chunks → one transfer per chunk. + ASSERT_EQ(resSrc.getDescs().size(), 2); + EXPECT_EQ(resSrc.getDescs()[0].getAddr(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[0].getLen(), 0x100000); + EXPECT_EQ(resSrc.getDescs()[1].getAddr(), 0x200000); + EXPECT_EQ(resSrc.getDescs()[1].getLen(), 0x100000); + EXPECT_EQ(resDst.getDescs()[0].getAddr(), 0x900000); + EXPECT_EQ(resDst.getDescs()[1].getAddr(), 0xA00000); } diff --git a/cpp/tests/unit_tests/executor/transferAgentTest.cpp b/cpp/tests/unit_tests/executor/transferAgentTest.cpp index b915ec3bb9c4..dcf6dd5a6e1e 100644 --- a/cpp/tests/unit_tests/executor/transferAgentTest.cpp +++ b/cpp/tests/unit_tests/executor/transferAgentTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -702,8 +702,7 @@ TEST(VmmDescSplitterTest, SplitTransferDescsDifferentChunkSizes) MemoryDescs srcInput{MemoryType::kVRAM, srcDescs}; MemoryDescs dstInput{MemoryType::kVRAM, dstDescs}; - auto [splitSrc, splitDst] - = VmmDescSplitter::splitTransferDescsWithRegionMaps(srcInput, dstInput, localMap, remoteMap); + auto [splitSrc, splitDst] = VmmDescSplitter::splitAndCoalesceTransferDescs(srcInput, dstInput, localMap, remoteMap); // dst has smaller chunks (512KB), so we get 4 pieces: 512K, 512K, 512K, 512K ASSERT_EQ(splitSrc.getDescs().size(), 4); @@ -735,8 +734,7 @@ TEST(VmmDescSplitterTest, SplitTransferDescsUnalignedBothSides) MemoryDescs srcInput{MemoryType::kVRAM, srcDescs}; MemoryDescs dstInput{MemoryType::kVRAM, dstDescs}; - auto [splitSrc, splitDst] - = VmmDescSplitter::splitTransferDescsWithRegionMaps(srcInput, dstInput, localMap, remoteMap); + auto [splitSrc, splitDst] = VmmDescSplitter::splitAndCoalesceTransferDescs(srcInput, dstInput, localMap, remoteMap); // src has 4MB chunks starting at srcBase → 1 piece from src side // dst has 2MB chunks starting at dstBase → 2 pieces from dst side @@ -760,7 +758,7 @@ TEST(VmmDescSplitterTest, SplitTransferDescsNoDstRegion) MemoryDescs dstInput{MemoryType::kVRAM, dstDescs}; auto [splitSrc, splitDst] - = VmmDescSplitter::splitTransferDescsWithRegionMaps(srcInput, dstInput, localMap, emptyRemoteMap); + = VmmDescSplitter::splitAndCoalesceTransferDescs(srcInput, dstInput, localMap, emptyRemoteMap); // Only src boundaries: 2 pieces of 1MB each ASSERT_EQ(splitSrc.getDescs().size(), 2); diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 46f90a9df31c..7e05e60017fa 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -78,7 +78,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | No | 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 SM100/SM103. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103 and in BF16/FP8 KV cache dtype. @@ -90,7 +90,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 kernels (`sparse_use_msa=True`, SM100). [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models). # Multimodal Feature Support Matrix (PyTorch Backend) diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index caadaca639b3..8d3db9cb63e3 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -384,7 +384,7 @@ def runLLMBuild(pipeline, buildFlags, tarName, is_linux_x86_64) sh "ccache -sv" sh "rm -rf **/*.xml *.tar.gz" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) if (env.alternativeTRT) { sh "cd ${LLM_ROOT} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } @@ -471,7 +471,7 @@ def buildWheelInContainer(pipeline, libraries=[], triple=X86_64_TRIPLE, clean=fa sh "cat ${CCACHE_DIR}/ccache.conf" // Step 1: cloning tekit source code - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) if (env.alternativeTRT) { trtllm_utils.replaceWithAlternativeTRT(env.alternativeTRT, cpver) sh "cd ${LLM_ROOT} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy index c41b999b8920..415eadca35e3 100644 --- a/jenkins/BuildDockerImage.groovy +++ b/jenkins/BuildDockerImage.groovy @@ -300,7 +300,7 @@ def buildImage(config, imageKeyToTag) stage (config.stageName) { // Step 1: Clone TRT-LLM source codes // If using a forked repo, svc_tensorrt needs to have the access to the forked repo. - trtllm_utils.checkoutSource(LLM_REPO, LLM_COMMIT_OR_BRANCH, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, LLM_COMMIT_OR_BRANCH, LLM_ROOT, true, true) } // Step 2: Build the images diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 57bc4123c2d5..ea7571878238 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -325,10 +325,10 @@ def setupPipelineEnvironment(pipeline, testFilter, globalVars) // NB: getContainerURIs reads files in ${LLM_ROOT}/jenkins/ if (env.gitlabMergeRequestLastCommit) { env.gitlabCommit = env.gitlabMergeRequestLastCommit - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) } else { branch = env.gitlabBranch ? env.gitlabBranch : "main" - trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, true, true) checkoutCommit = sh (script: "cd ${LLM_ROOT} && git rev-parse HEAD",returnStdout: true).trim() env.gitlabCommit = checkoutCommit } @@ -451,7 +451,7 @@ def launchReleaseCheck(pipeline, globalVars) sh "pip3 config set global.break-system-packages true" sh "git config --global --add safe.directory \"*\"" // Step 1: Clone TRT-LLM source codes - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) sh "cd ${LLM_ROOT} && git config --unset-all core.hooksPath" // Step 2: Run guardwords scan @@ -1197,7 +1197,7 @@ def collectTestResults(pipeline, testFilter, globalVars) echo "Result File Number: ${resultFileNumber}, Downloaded: ${resultFileDownloadedNumber}" sh "find . -name results-\\*.tar.gz -type f -exec tar -zxvf {} \\; || true" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) if (testFilter[(IS_POST_MERGE)]) { try { sh "python3 llm/scripts/generate_duration.py --duration-file=new_test_duration.json" diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 8f246f744287..e090074bc7c6 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2852,7 +2852,7 @@ def runLLMDocBuild(pipeline, config) sh "pwd && ls -alh" sh "env | sort" // allow to checkout from forked repo, svc_tensorrt needs to have access to the repo, otherwise clone will fail - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) sh "mkdir TensorRT-LLM" sh "cp -r ${LLM_ROOT}/ TensorRT-LLM/src/" trtllm_utils.llmExecStepWithRetry(pipeline, script: "git config --global --add safe.directory \"*\"") @@ -4276,7 +4276,7 @@ def runLLMBuild( sh "env | sort" sh "ccache -sv" - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, "tensorrt_llm", true, true) if (env.alternativeTRT) { sh "cd tensorrt_llm/ && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } @@ -5367,7 +5367,7 @@ def launchTestJobs(pipeline, testFilter) trtllm_utils.llmExecStepWithRetry(pipeline, script: 'rm -rf $(python3 -c "import site; print(site.getsitepackages()[0])")/nvidia_cutlass_dsl*') } trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y python3-pip git rsync curl wget") - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 config set global.break-system-packages true") trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install requests") trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 uninstall -y tensorrt") diff --git a/jenkins/TensorRT_LLM_PLC.groovy b/jenkins/TensorRT_LLM_PLC.groovy index dc3c1735e2d1..651e431ea96f 100644 --- a/jenkins/TensorRT_LLM_PLC.groovy +++ b/jenkins/TensorRT_LLM_PLC.groovy @@ -140,7 +140,7 @@ def checkoutSource () def LLM_REPO = getLLMRepo() sh "git config --global --add safe.directory ${env.WORKSPACE}" def ref = params.ref - trtllm_utils.checkoutSource(LLM_REPO, ref, env.WORKSPACE, false, true) + trtllm_utils.checkoutSource(LLM_REPO, ref, env.WORKSPACE, true, true) } def getPulseToken(serviceId, scopes) { diff --git a/jenkins/runPerfSanityTriage.groovy b/jenkins/runPerfSanityTriage.groovy index 86591da57cc7..bc31ad53e318 100644 --- a/jenkins/runPerfSanityTriage.groovy +++ b/jenkins/runPerfSanityTriage.groovy @@ -89,7 +89,7 @@ pipeline { container("trt-llm") { script { sh "pwd && ls -alh" - trtllm_utils.checkoutSource(LLM_REPO, params.BRANCH, LLM_ROOT, false, false) + trtllm_utils.checkoutSource(LLM_REPO, params.BRANCH, LLM_ROOT, true, false) def commandsBase64 = params.COMMANDS.bytes.encodeBase64().toString() sh """ cd ${LLM_ROOT}/jenkins/scripts/perf && python3 perf_sanity_triage.py \ diff --git a/requirements.txt b/requirements.txt index dab577614f57..148460a26d8a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -74,6 +74,8 @@ blobfile openai-harmony==0.0.4 nvidia-cutlass-dsl[cu13]==4.5.0; python_version >= "3.10" nvidia-matmul-heuristics==0.1.0.27; python_version >= "3.10" # analytic GEMM heuristics for CuTe DSL autotuner tactic pruning +quack-kernels>=0.2.10; python_version >= "3.10" # required for MinimaxM3 MSA +jinja2 # required for MinimaxM3 MSA plotly numexpr partial_json_parser diff --git a/scripts/attribution/scan/metadata/msa.yml b/scripts/attribution/scan/metadata/msa.yml new file mode 100644 index 000000000000..ab212fc1b664 --- /dev/null +++ b/scripts/attribution/scan/metadata/msa.yml @@ -0,0 +1,5 @@ +name: msa +description: MiniMax Sparse Attention (fmha_sm100) kernels for SM100 sparse attention +source: submodule +directory_matches: +- 3rdparty/MSA diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 263f061d14b6..8b5dbdb32357 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -80,6 +80,37 @@ def get_build_dir(build_dir, build_type): return build_dir +def stage_msa_package(project_dir: Path, build_dir: Path) -> Path: + """Copy pinned MSA sources and apply TensorRT-LLM's downstream patch.""" + msa_source_dir = project_dir / "3rdparty" / "MSA" + msa_package_dir = msa_source_dir / "python" / "fmha_sm100" + msa_patch = project_dir / "3rdparty" / "patches" / "msa_strided_paged_kv.patch" + if not msa_package_dir.is_dir(): + raise FileNotFoundError( + f"MSA sources are missing at {msa_package_dir}; initialize 3rdparty/MSA" + ) + + staging_dir = build_dir / "msa_patched" + if staging_dir.exists(): + rmtree(staging_dir) + copytree(msa_source_dir, staging_dir, ignore=shutil.ignore_patterns(".git")) + git_env = os.environ.copy() + git_env["GIT_CEILING_DIRECTORIES"] = str(staging_dir.parent.resolve()) + run( + ["git", "apply", "--check", str(msa_patch)], + cwd=staging_dir, + env=git_env, + check=True, + ) + run( + ["git", "apply", str(msa_patch)], + cwd=staging_dir, + env=git_env, + check=True, + ) + return staging_dir / "python" / "fmha_sm100" + + def clear_folder(folder_path): for item in os.listdir(folder_path): item_path = os.path.join(folder_path, item) @@ -1173,10 +1204,14 @@ def get_binding_lib(subdirectory, name): f"Copied auto-generated attributions to {project_dir / 'ATTRIBUTIONS.md'}" ) + msa_package_dir = stage_msa_package(project_dir, build_dir) + wheel_env = os.environ.copy() + wheel_env["TRTLLM_MSA_PACKAGE_DIR"] = str(msa_package_dir) + build_run( - f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"' - ) - env = os.environ.copy() + f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {extra_wheel_build_args} --no-isolation --wheel --outdir "{dist_dir}"', + env=wheel_env) + env = wheel_env.copy() if mypyc: env["TRTLLM_ENABLE_MYPYC"] = "1" else: diff --git a/setup.py b/setup.py index 7a3b7d2af8f7..b6e3f3e95bc1 100644 --- a/setup.py +++ b/setup.py @@ -406,6 +406,12 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], # internal absolute imports (e.g., "from triton_kernels.foo import bar") work. packages += find_packages(include=["triton_kernels", "triton_kernels.*"]) +msa_package_dir = { + "fmha_sm100": + os.environ.get("TRTLLM_MSA_PACKAGE_DIR", "3rdparty/MSA/python/fmha_sm100") +} +packages += ["fmha_sm100"] + # https://setuptools.pypa.io/en/latest/references/keywords.html setup( name='tensorrt_llm', @@ -420,6 +426,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], url="https://github.com/NVIDIA/TensorRT-LLM", download_url="https://github.com/NVIDIA/TensorRT-LLM/tags", packages=packages, + package_dir=msa_package_dir, exclude_package_data=exclude_package_data, # TODO Add windows support for python bindings. classifiers=[ @@ -432,8 +439,17 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], license="Apache License 2.0", keywords="nvidia tensorrt deeplearning inference", package_data={ - 'tensorrt_llm': package_data, + 'tensorrt_llm': + package_data, 'triton_kernels': ['LICENSE', 'VERSION', 'README.md'], + 'fmha_sm100': [ + '*.py', + 'csrc/**/*', + 'cute/**/*', + 'cutlass/include/**/*', + 'cutlass/tools/util/include/**/*', + 'cutlass/LICENSE.txt', + ], }, license_files=get_license(), entry_points={ diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py index 1c3981abcf91..b2cd1e75ec7e 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -16,6 +16,7 @@ from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha +from .msa_sparse_gqa import MsaSparseGqaFmha from .phased import FmhaParams, PhasedFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes @@ -27,6 +28,7 @@ "Fmha", "FmhaCls", "FmhaParams", + "MsaSparseGqaFmha", "PhasedFmha", "get_enabled_fmha_lib_classes", ] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py new file mode 100644 index 000000000000..5578f36e3ca8 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse GQA FMHA backed by MSA's fmha_sm100 kernel. + +MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and +participates in the standard TrtllmAttention.forward dispatch loop. The +owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the +per-query KV blocks and publishes them on forward_args.sparse_prediction; +this class attends over them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._utils import is_sm_100f + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +def run_msa_sparse_gqa( + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + kv_block_indexes: Optional[torch.Tensor] = None, + *, + kv_indices: torch.Tensor, + sm_scale: float, + qo_lens_cpu: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, + qo_offset_cpu: Optional[torch.Tensor] = None, + causal: bool = True, + head_dim: int = 128, + plan: Optional[tuple] = None, + out: Optional[torch.Tensor] = None, + use_fp8: bool = False, +) -> None: + """Run fmha_sm100 paged GQA (plan/run split). + + `kv_block_indexes`: if set, sparse top-k mode (fixed `kv_block_num=topk`); + if None, dense mode attending all pages in `kv_indices`. + `plan`: prebuilt execution plan; if None, built inline from the CPU length + tensors (eager prefill/tests vs. CUDA-graph decode). + `out`: destination buffer the kernel writes in place. + `use_fp8`: FP8 KV cache. The caller must pass FP8 `q` to match the FP8 paged + K/V, since the kernel variant shares one dtype across q/k/v. Also selects the + FP8 AOT kernels for an inline sparse-prefill plan; no-op for the decode planner. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import require_msa_module + + fmha_sm100 = require_msa_module() + + if q.dim() != 3: + raise ValueError( + f"MsaSparseGqaFmha expects q [total_q, num_qo_heads, head_dim]; got {tuple(q.shape)}." + ) + if q.shape[-1] != head_dim: + raise NotImplementedError( + f"MsaSparseGqaFmha supports head_dim={head_dim}; got {q.shape[-1]}." + ) + if k_paged.dim() != 4 or v_paged.dim() != 4: + raise ValueError( + "MsaSparseGqaFmha expects paged KV [num_pages, num_kv_heads, page_size, head_dim]; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + if k_paged.shape != v_paged.shape: + raise ValueError( + f"MsaSparseGqaFmha requires k and v to share shape; " + f"got k={tuple(k_paged.shape)}, v={tuple(v_paged.shape)}." + ) + + if plan is None: + # kv_block_num is planned only for the sparse (block-indexed) path; + # dense paged GQA leaves it unset and attends the full page table. + kv_block_num = int(kv_block_indexes.shape[-1]) if kv_block_indexes is not None else -1 + plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + int(q.shape[1]), # num query heads. + num_kv_heads=int(k_paged.shape[1]), + qo_offset=qo_offset_cpu, + page_size=int(k_paged.shape[2]), + kv_block_num=kv_block_num, + causal=causal, + num_kv_splits=1, + use_fp8_kvcache=use_fp8, + ) + fmha_sm100.fmha_sm100( + q, + k_paged, + v_paged, + plan, + kv_indices=kv_indices, + kv_block_indexes=kv_block_indexes, + out=out, + sm_scale=sm_scale, + output_maxscore=False, + ) + + +def run_msa_paged_gqa( + attn: "TrtllmAttention", + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + output: torch.Tensor, + *, + kv_block_indexes: Optional[torch.Tensor], + plan: Optional[tuple], +) -> None: + """Write the new-token main K/V, then run paged GQA into output in place. + + Shared by the sparse layers (kv_block_indexes is the per-query top-k table, + with the sparse plan) and the dense layers (kv_block_indexes None, with the + dense plan, attending the full page table). fmha_sm100 reads the paged cache + directly, so the new-token K/V must be resident before the run. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_paged_kv, + write_msa_main_kv, + ) + + layer_idx = attn.layer_idx + head_dim = attn.head_dim + kv_cache_manager = metadata.kv_cache_manager + num_tokens = int(q.shape[0]) + if k is not None and v is not None: + write_msa_main_kv( + kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v + ) + + q_view = q.view(num_tokens, attn.num_heads, head_dim) + out_view = output.view(num_tokens, attn.num_heads, head_dim) + k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx) + sm_scale = (head_dim**-0.5) / float(attn.q_scaling) + + # The fmha_sm100 variant is chosen from q.dtype and shares one dtype across + # q/k/v, so q must be FP8 to match an FP8 paged K/V. MiniMax-M3 has no + # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. + use_fp8 = k_paged.dtype == torch.float8_e4m3fn + if use_fp8: + q_view = q_view.to(torch.float8_e4m3fn) + + run_msa_sparse_gqa( + q_view, + k_paged, + v_paged, + kv_block_indexes, + kv_indices=metadata.msa_kv_indices, + sm_scale=sm_scale, + qo_lens_cpu=metadata.msa_qo_lens_cpu, + kv_lens_cpu=metadata.msa_kv_lens_cpu, + qo_offset_cpu=metadata.msa_qo_offset_cpu, + causal=True, + head_dim=head_dim, + plan=plan, + out=out_view, + use_fp8=use_fp8, + ) + + +class MsaSparseGqaFmha(Fmha): + """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. + + Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's + selected KV block indices on forward_args.sparse_prediction.sparse_attn_indices + and attend those blocks; dense layers leave the indices None and attend the + full page table. + + Inherits Fmha rather than PhasedFmha: fmha_sm100 takes a single plan and + the selected block indices span the whole batch, so it handles a mixed + context and generation batch in one call and there is no + context/generation split from PhasedFmha to reuse. Requires head_dim 128 + and 4-D HND paged K/V. + """ + + @classmethod + def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: + # fmha_sm100 runs only on the SM100 family and ships in the MSA git + # submodule, so it is unavailable off SM100 or without the package. + # Imported lazily because the minimax_m3 package init imports the trtllm + # attention classes, which a module-scope import here would cycle with. + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + msa_package_available, + ) + + if not is_sm_100f() or not msa_package_available(): + return False + # Only the MiniMax-M3 MSA layer uses this library. Matching the lowered + # sparse algorithm lets the base create_fmha_libs add it to that layer + # alone, so no create_fmha_libs override is needed. Dense layers (e.g. + # an Eagle3 draft model) have no sparse_params. + return attn.sparse_params is not None and attn.sparse_params.algorithm == "minimax_m3" + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + ) -> None: + output = forward_args.output + if output is None: + raise RuntimeError(f"{type(self).__name__} requires an output buffer.") + + # Sparse layers attend the per-query top-k blocks with the sparse plan; + # dense layers leave the indices None and attend the full page table + # with the dense plan. + kv_block_indexes = forward_args.sparse_prediction.sparse_attn_indices + if kv_block_indexes is not None: + plan = metadata.msa_decode_gqa_plan + if plan is None: + plan = getattr(metadata, "msa_eager_gqa_plan", None) + else: + plan = metadata.msa_decode_dense_plan + if plan is None: + plan = getattr(metadata, "msa_eager_dense_plan", None) + run_msa_paged_gqa( + self.attn, + q, + k, + v, + metadata, + output, + kv_block_indexes=kv_block_indexes, + plan=plan, + ) + + +__all__ = ["MsaSparseGqaFmha"] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py index 97467657e9b5..2b3ef3fc7ff3 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/registry.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -22,10 +22,24 @@ FmhaCls: TypeAlias = type[Fmha] -FMHA_LIBS: dict[str, FmhaCls] = { - "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, - "fallback": FallbackFmha, -} + +def init_fmha_libs() -> dict[str, "FmhaCls"]: + """Build the ordered FMHA library registry. + + Backend classes are imported inside this factory rather than at module + scope, so backends can import trtllm attention classes at module scope + without an import cycle. + """ + from .msa_sparse_gqa import MsaSparseGqaFmha + + return { + "msa_sparse_gqa": MsaSparseGqaFmha, + "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, + "fallback": FallbackFmha, + } + + +FMHA_LIBS: dict[str, FmhaCls] = init_fmha_libs() DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(FMHA_LIBS) @@ -78,4 +92,5 @@ def get_enabled_fmha_lib_classes() -> list[FmhaCls]: "FMHA_LIBS", "FmhaCls", "get_enabled_fmha_lib_classes", + "init_fmha_libs", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py index 0b2994441bf0..f293f9547506 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py @@ -1,13 +1,3 @@ -# yapf: disable -from .minimax_m3 import (MiniMaxM3SparseAttention, - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, MiniMaxM3SparseIndexCache, - allocate_minimax_m3_static_buffers, - build_runtime_metadata_from_kv_manager, - get_minimax_m3_attention_backend_cls, - get_minimax_m3_kv_cache_manager_cls, - minimax_m3_sparse_decode, minimax_m3_sparse_prefill) -# yapf: enable from .utils import (get_flashinfer_sparse_attn_attention_backend, get_sparse_attn_kv_cache_manager, get_trtllm_sparse_attn_attention_backend, @@ -18,14 +8,4 @@ "get_vanilla_sparse_attn_attention_backend", "get_trtllm_sparse_attn_attention_backend", "get_flashinfer_sparse_attn_attention_backend", - "MiniMaxM3SparseAttention", - "MiniMaxM3SparseAttentionMetadata", - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseIndexCache", - "allocate_minimax_m3_static_buffers", - "build_runtime_metadata_from_kv_manager", - "get_minimax_m3_attention_backend_cls", - "get_minimax_m3_kv_cache_manager_cls", - "minimax_m3_sparse_decode", - "minimax_m3_sparse_prefill", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py index c248ffd91119..522bf2d5b9ee 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.py @@ -4,24 +4,37 @@ Layered as: - * :mod:`.kernels` -- OpenAI Triton kernels (per-block max - score, masked softmax for sparse GQA). - * :mod:`.metadata` -- ``MiniMaxM3SparseConfig`` / - ``MiniMaxM3SparseAttentionMetadata`` - dataclasses, CUDA-graph-stable buffer - allocator + builder, and the - :class:`AttentionMetadata` subclass - factory. - * :mod:`.cache_manager` -- standalone side index cache used by tests - and the :class:`KVCacheManagerV2` - subclass factory. - * :mod:`.backend` -- the algorithm itself (vectorized - paged-cache helpers, prefill / decode - entry points, the thin - :class:`MiniMaxM3SparseAttention` - orchestrator) and the - :class:`AttentionBackend` subclass - factory. + * :mod:`.triton_metadata` -- ``MiniMaxM3TritonSparseAttentionMetadata`` + dataclass, CUDA-graph-stable buffer + allocator + builder, and the + :class:`AttentionMetadata` subclass + factory for the Triton reference path. + * :mod:`.cache_manager` -- standalone side index cache used by tests + and the :class:`KVCacheManagerV2` + subclass factory. Shared by both backends. + * :mod:`.common` -- backend-neutral config bundles, the paged + KV-slot writer, block-priority sentinels, and + the paged-cache slot mapping builder shared by + both backends. + * :mod:`.msa_utils` -- MSA-only (fmha_sm100) helpers: import guard, + kernel precondition constants, HND paged-cache + adapters, main-KV writer, page-table builder, + valid-block counting, and top-k selection. + * :mod:`.triton_kernels` -- OpenAI Triton kernels (per-block max + score, masked softmax for sparse GQA). + * :mod:`.triton_backend` -- the Triton reference algorithm (vectorized + paged-cache helpers, prefill / decode + entry points, the thin + :class:`MiniMaxM3TritonSparseAttention` + orchestrator) and its + :class:`AttentionBackend` subclass + factory. + * :mod:`.msa_backend` -- the MSA (fmha_sm100) backend, its flat + metadata, and the backend factory. + * :mod:`.msa_indexer` -- the MSA proxy scoring + top-k block + selection submodule. + * :mod:`.msa_availability`-- SM100 and fmha_sm100 gating for the MSA + path. This package's public surface re-exports the names callers historically imported from ``...sparse.minimax_m3`` so external @@ -29,48 +42,20 @@ working unchanged. """ -# Re-export the algorithm-internal helpers focused unit tests reach -# into so the package preserves the surface the monolithic module -# exposed. These are not part of ``__all__`` (still package-private) -# but stay importable as ``from ...minimax_m3 import _write_main_kv_slots``. -from .backend import ( # noqa: F401 - MiniMaxM3SparseAttention, - _compute_index_attn_chunk_q, - _compute_sparse_gqa_chunk_q, +# The dense Triton oracle in the model imports these paged-cache helpers, so +# they stay importable from the package. They are package-private and are not +# part of __all__. Every other backend/metadata/config symbol is imported +# directly from its defining submodule by the code that needs it. +from .cache_manager import MiniMaxM3KVCacheManagerV2 +from .msa_backend import MiniMaxM3MsaSparseAttention +from .triton_backend import ( # noqa: F401 + MiniMaxM3SparseRuntimeBackend, _gather_paged_batched, - _index_attention_and_select, - _write_main_kv_slots, _write_main_kv_slots_to_pool, - get_minimax_m3_attention_backend_cls, - minimax_m3_sparse_decode, - minimax_m3_sparse_prefill, -) -from .cache_manager import ( - MiniMaxM3KVCacheManagerV2, - MiniMaxM3SparseIndexCache, - get_minimax_m3_kv_cache_manager_cls, -) -from .metadata import ( - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, - allocate_minimax_m3_static_buffers, - build_runtime_metadata_from_kv_manager, - get_minimax_m3_attention_metadata_cls, - replace_metadata, ) __all__ = [ "MiniMaxM3KVCacheManagerV2", - "MiniMaxM3SparseAttention", - "MiniMaxM3SparseAttentionMetadata", - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseIndexCache", - "allocate_minimax_m3_static_buffers", - "build_runtime_metadata_from_kv_manager", - "get_minimax_m3_attention_backend_cls", - "get_minimax_m3_attention_metadata_cls", - "get_minimax_m3_kv_cache_manager_cls", - "minimax_m3_sparse_decode", - "minimax_m3_sparse_prefill", - "replace_metadata", + "MiniMaxM3MsaSparseAttention", + "MiniMaxM3SparseRuntimeBackend", ] 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 b9badca6ea35..fd0f72159f1f 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 @@ -1,5 +1,17 @@ # 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. """KV cache management for MiniMax-M3 sparse attention. Provides: @@ -13,21 +25,17 @@ from __future__ import annotations -from typing import List, Optional +from typing import List, Optional, Sequence import torch -from tensorrt_llm._utils import ( - TensorWrapper, - binding_to_torch_dtype, - convert_to_torch_tensor, - prefer_pinned, -) +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, LayerId +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX -from tensorrt_llm.runtime.kv_cache_manager_v2._utils import typed_range +from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole from ....pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role @@ -142,6 +150,20 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): * ``sparse_index_dim`` — width of the index-K/V vectors. """ + # The AttentionOp-facing tensors this manager builds are synthetic + # placeholders over INDEX_KEY-coalesced pools, so one-model speculative + # draft layers must live in a separate manager even under attention DP + # (read by ``_should_create_separate_draft_kv_cache``). + supports_shared_draft_layers = False + + # Workaround: the Eagle3 drafter's trtllm-gen generation kernels hit an + # illegal memory access at tokens_per_block=128 (base regression), while + # the MSA target requires 128. The dense draft layers have no page-size + # constraint, so the separate draft manager runs at 32 (read by + # ``_create_one_model_draft_kv_cache_manager``). Remove once the kernel + # bug is fixed. + draft_manager_tokens_per_block = 32 + def __init__( self, *args, @@ -176,6 +198,15 @@ def __init__( super().__init__(*args, **kwargs) + index_v_layer_ids = set(self.sparse_layer_ids) - self.disable_index_value_layer_ids + if self.is_disagg and index_v_layer_ids: + raise ValueError( + "MiniMax M3 disaggregated serving requires disable_index_value=True " + "for every sparse layer because the optional test-only index-V cache " + "is not managed or transferred by KVCacheManagerV2; enabled layers=" + f"{sorted(index_v_layer_ids)}" + ) + # Optional plain-tensor index-V cache for non-disabled sparse # layers (test-only; production has disable_index_value=True # on every sparse layer). @@ -212,6 +243,13 @@ def _extra_buffers_per_layer(self, *, tokens_per_block): if layer_id in self.layer_offsets } + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """Declare MiniMax M3's token-major K/V and replicated index-K.""" + return { + Role.ALL: MapperKind.NHD, + Role.INDEX_KEY: MapperKind.REPLICATED, + } + def _compute_num_total_slots(self) -> int: """Total token slots across all blocks in the main K pool. @@ -234,18 +272,18 @@ def _torch_dtype_for_index_cache(self) -> torch.dtype: return torch.float32 return torch.bfloat16 - def get_index_k_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: + def get_index_k_buffer(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: """Return the V2-managed paged index-K view for ``layer_idx``. - Shape: ``[num_pages, tokens_per_block, 1, sparse_index_dim]``. - Reads/writes decompose ``slot = (page, within)`` and use - multi-dim fancy indexing; writes propagate to pool storage. + NHD shape is ``[num_pages, tokens_per_block, 1, sparse_index_dim]``; + HND shape is ``[num_pages, 1, tokens_per_block, sparse_index_dim]``. """ return super().get_index_k_buffer( layer_idx, num_heads=1, head_dim=self.sparse_index_dim, dtype=self._torch_dtype_for_index_cache(), + kv_layout=kv_layout, ) def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: @@ -260,12 +298,12 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. The base :meth:`KVCacheManagerV2.get_buffers` produces a ``[num_pages, kv_factor, ...]`` view with contiguous strides - that assume each slot contains exactly K+V. With INDEX_KEY - registered, sparse layers may have ``scale > 2`` per-slot - buffers (e.g. M3 TP=8 coalesces K, V, INDEX_K into one pool - where ``scale == 3 * num_sparse + 2 * num_dense``), and the - base view's dim-0 stride no longer reaches the next slot's K - for this layer. + that assume the slot holds exactly one layer's K+V. In M3's + pool the slot packs K+V for *all* layers of the group + (``scale >= 2 * num_layers_in_group``), so the base view's + dim-0 stride does not reach the next slot's K for this layer. + (When INDEX_KEY's per-block size coincides with K/V's, it is + coalesced into the same pool and contributes to ``scale`` too.) The override builds a ``[num_slots, scale, ...]`` view rooted at K's base, then slices ``[:, :2]`` to extract K+V. The slice @@ -341,75 +379,20 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) return full_view[:, :2] - def _build_pool_mapping_tensors(self): - """Compute pool-mapping offsets from layer position in the pool group. - - The base method does ``exact_div(addr_offset, key_bytes * - kv_factor * tokens_per_block)``, which assumes each layer - contributes exactly K+V. When INDEX_KEY coincidentally shares - the same per-block size as K/V (M3 production at TP=8: all - three are 256 B/token), V2 coalesces all three into one pool - and the per-layer stride becomes ``3 * single_buffer_size`` — - the base ``exact_div`` then asserts. - - Compute ``offset`` directly from - ``self.impl.layer_grouping[group_id]`` so the formula stays - correct regardless of how many extra buffers coalesce with - K/V. The M3 forward path uses :meth:`get_buffers` / + def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int: + """Pool-mapping offset from layer position in the pool group. + + The base formula ``exact_div(addr_offset, key_bytes * kv_factor * + tokens_per_block)`` assumes each layer contributes exactly K+V to + its pool slot. When index-K coalesces into the K/V pool the layer + stride is non-uniform (sparse layers add an INDEX_KEY sub-page), + so derive ``offset`` from ``self.impl.layer_grouping`` instead. + The M3 forward path uses :meth:`get_buffers` / :meth:`get_index_k_buffer` rather than this mapping, so the offset just needs to be consistent (layer position in group). """ - kv_cache_pool_pointers = torch.tensor( - [ - [ - self.impl.get_mem_pool_base_address( - self.impl.layer_grouping[pool_id][0], Role.KEY - ), - 0, - ] - for pool_id in range(self.num_pools) - ], - dtype=torch.int64, - device="cpu", - pin_memory=prefer_pinned(), - ) - - if self.dtype == DataType.NVFP4: - kv_cache_pool_pointers = torch.stack( - [ - kv_cache_pool_pointers, - torch.tensor( - [ - [ - self.impl.get_mem_pool_base_address( - self.impl.layer_grouping[pool_id][0], Role.KEY_BLOCK_SCALE - ), - 0, - ] - for pool_id in range(self.num_pools) - ], - dtype=torch.int64, - device="cpu", - pin_memory=prefer_pinned(), - ), - ], - dim=-1, - ) - - kv_cache_pool_mapping_list = [] - for layer_id in typed_range(LayerId(self.num_local_layers)): - layer_group_id = self.impl.get_layer_group_id(layer_id) - layers_in_group = list(self.impl.layer_grouping[int(layer_group_id)]) - offset = layers_in_group.index(int(layer_id)) - kv_cache_pool_mapping_list.append([int(layer_group_id), offset]) - - kv_cache_pool_mapping = torch.tensor( - kv_cache_pool_mapping_list, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - return kv_cache_pool_pointers, kv_cache_pool_mapping + layers_in_group = list(self.impl.layer_grouping[int(layer_group_id)]) + return layers_in_group.index(int(layer_id)) def _get_batch_cache_indices_by_pool_id( self, @@ -417,34 +400,50 @@ def _get_batch_cache_indices_by_pool_id( *, pool_id: int = 0, is_kv_aggregate: bool = True, + num_blocks_per_seq: Optional[Sequence[int]] = None, + index_scale: Optional[int] = None, ): - """Return per-request slot ids in ``[0, num_slots)`` directly. + """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1). The base method converts slot ids to V1-style block ids via ``base_idx * index_scales[pool_id] // kv_factor``, which is - only correct when each layer contributes exactly K+V. With - INDEX_KEY-coalesced sparse pools (M3 production), the scale - breaks the V1 conversion and produces out-of-bounds block ids - during V2 warmup. + only correct when each layer contributes exactly K+V. M3's slot + packs K+V for all layers of the group, so the scale breaks the + V1 conversion and produces out-of-bounds block ids during V2 + warmup. + + ``index_scale`` (per-layer scale supplied by the base + :meth:`get_batch_cache_indices`) is accepted for + signature compatibility but intentionally ignored: this + override bypasses the scale conversion entirely. Bypass the conversion: the M3 forward path indexes paged views (built by :meth:`get_buffers` / :meth:`get_index_k_buffer`) directly by slot id. - ``BAD_PAGE_INDEX`` slots stay as 0 to match the legacy - padding contract. + ``BAD_PAGE_INDEX`` slots remain ``-1`` here because disaggregation's + :class:`KVRegionExtractorV1` filters ``region_ids >= 0``. + :meth:`get_block_ids_per_seq` maps them to zero for the attention + metadata's padded tensor. + + Args: + request_ids: Request IDs whose page-index rows are returned. + pool_id: V2 pool whose page indices are requested. + is_kv_aggregate: Kept for compatibility with the base virtual method. + num_blocks_per_seq: Optional per-request truncation limits. When + omitted, preserve the full padded width required by MiniMax + CUDA-graph metadata initialization. + index_scale: Kept for compatibility with the base virtual method; + M3 bypasses the V1 block-id conversion entirely, so any + caller-supplied scale is ignored alongside ``index_scales``. """ res = [] - for req_id in request_ids: - idx_tensor = torch.as_tensor(self.kv_cache_map[req_id].get_base_page_indices(pool_id)) - res.append( - ( - torch.where( - idx_tensor != BAD_PAGE_INDEX, - idx_tensor, - torch.full_like(idx_tensor, BAD_PAGE_INDEX), - ) - ).tolist() - ) + for req_idx, req_id in enumerate(request_ids): + kv_cache = self.kv_cache_map[req_id] + base_page_indices = kv_cache.get_base_page_indices(pool_id) + if num_blocks_per_seq is not None: + num_blocks = min(kv_cache.num_blocks, num_blocks_per_seq[req_idx]) + base_page_indices = base_page_indices[:num_blocks] + res.append(list(base_page_indices)) return res def get_block_ids_per_seq(self, request_ids): diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py new file mode 100644 index 000000000000..9f2863846df2 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared building blocks for the MiniMax-M3 sparse attention backends. + +Both the Triton reference and the MSA (fmha_sm100) path share these +backend-neutral pieces: the lowered parameter and per-rank kernel config +bundles, block-priority sentinels, KV-slot writers, and the paged-cache +slot mapping builder. MSA-only helpers live in :mod:`.msa_utils`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, Literal, Optional, Tuple + +import torch + +from ..params import SparseMetadataParams, SparseParams + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + +# Sentinel scores that force init and local blocks into the top-k regardless +# of their computed score. Init outranks local. +_INIT_SCORE = 1e30 +_LOCAL_SCORE = 1e29 + + +@dataclass(frozen=True) +class MiniMaxM3SparseParams(SparseParams): + """Lowered runtime parameters for the MiniMax-M3 sparse backend.""" + + algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") + num_index_heads: int = 4 + sparse_index_dim: int = 128 + block_size: int = 128 + topk: int = 16 + init_blocks: int = 0 + local_blocks: int = 1 + score_type: str = "max" + disable_index_value: bool = True + implementation: Literal["triton", "msa"] = "triton" + + @property + def indices_block_size(self) -> int: + """Block granularity of the selected sparse indices. + + Read by the shared TrtllmAttention forward when publishing the + sparse prediction. It equals the per-block scoring size. + """ + return self.block_size + + +@dataclass(frozen=True) +class MiniMaxM3SparseMetadataParams(SparseMetadataParams): + """Metadata-facing MiniMax-M3 sparse geometry.""" + + global_num_q_heads: int = 0 + global_num_kv_heads: int = 0 + num_index_heads: int = 4 + topk: int = 16 + + def sharded_head_counts(self, mapping: Optional["Mapping"] = None) -> Tuple[int, int]: + """Return per-rank (num_q_heads, num_kv_heads) for mapping. + + Matches the model's attention sharding: no split under attention data + parallelism, otherwise split by tp_size. + """ + if mapping is not None and not getattr(mapping, "enable_attention_dp", False): + tp_size = int(getattr(mapping, "tp_size", 1) or 1) + else: + tp_size = 1 + + def _shard(num_heads: int) -> int: + return (int(num_heads) + tp_size - 1) // tp_size + + return _shard(self.global_num_q_heads), _shard(self.global_num_kv_heads) + + +@dataclass(frozen=True) +class MiniMaxM3SparseConfig: + """Per-rank kernel parameter bundle for MiniMax-M3 sparse attention. + + This is **not** a user-facing config (use + :class:`tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig` + for that). It is the layer-invariant, post-TP-shard parameter bundle + that backend kernels and reference helpers consume. The user knobs + come from :class:`MiniMaxM3SparseParams`; ``num_q_heads`` / + ``num_kv_heads`` / ``head_dim`` come from the per-rank model + geometry and must be supplied by the caller (typically via + :meth:`from_sparse_params`). + """ + + num_q_heads: int + num_kv_heads: int + head_dim: int + num_index_heads: int + sparse_index_dim: int + block_size: int + topk: int + init_blocks: int = 0 + local_blocks: int = 1 + score_type: str = "max" + + def __post_init__(self) -> None: + if self.num_q_heads % self.num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({self.num_q_heads}) must be divisible by " + f"num_kv_heads ({self.num_kv_heads})" + ) + if self.num_index_heads % self.num_kv_heads != 0: + raise ValueError( + f"num_index_heads ({self.num_index_heads}) must be divisible " + f"by num_kv_heads ({self.num_kv_heads})" + ) + if self.block_size <= 0: + raise ValueError(f"block_size must be > 0, got {self.block_size}") + if self.topk <= 0: + raise ValueError(f"topk must be > 0, got {self.topk}") + if self.init_blocks < 0: + raise ValueError(f"init_blocks must be >= 0, got {self.init_blocks}") + if self.local_blocks < 0: + raise ValueError(f"local_blocks must be >= 0, got {self.local_blocks}") + if self.score_type != "max": + # SGLang exposes only "max" today and that is what the MiniMax-M3 + # checkpoint config specifies. Reject anything else explicitly so + # a config drift surfaces immediately. + raise ValueError( + f"score_type={self.score_type!r} is not supported " + "(only 'max' matches the SGLang reference)" + ) + + @classmethod + def from_sparse_params( + cls, + sparse_params: "MiniMaxM3SparseParams", + *, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + ) -> "MiniMaxM3SparseConfig": + """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` + and the per-rank model geometry. + """ + return cls( + num_q_heads=int(num_q_heads), + num_kv_heads=int(num_kv_heads), + head_dim=int(head_dim), + num_index_heads=int(sparse_params.num_index_heads), + sparse_index_dim=int(sparse_params.sparse_index_dim), + block_size=int(sparse_params.block_size), + topk=int(sparse_params.topk), + init_blocks=int(sparse_params.init_blocks), + local_blocks=int(sparse_params.local_blocks), + score_type=str(sparse_params.score_type), + ) + + +def write_kv_slots( + cache: torch.Tensor, + out_cache_loc: torch.Tensor, + values: torch.Tensor, + *, + layout: Literal["NHD", "HND"] = "NHD", +) -> None: + """Write per-token values into a K, V, or index-K cache at given slots. + + Handles a 3-D flat-slot cache and a 4-D paged view. `layout` sets the paged + axis order: "NHD" is [num_pages, tokens_per_block, num_heads, channel], + "HND" is [num_pages, num_heads, tokens_per_block, channel]. The paged view + is non-contiguous, so the slot id is split into (page, within) and written + by multi-dim assignment. `values` is always [num_tokens, num_heads, channel]. + """ + with torch.no_grad(): + if cache.ndim >= 4: + token_axis = 2 if layout == "HND" else 1 + tokens_per_block = int(cache.shape[token_axis]) + out_long = out_cache_loc.to(torch.long) + page = out_long // tokens_per_block + within = out_long % tokens_per_block + if layout == "HND": + # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and + # move front, giving a [num_tokens, num_heads, channel] target. + cache[page, :, within, :] = values.to(cache.dtype) + else: + cache[page, within] = values.to(cache.dtype) + else: + cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + + +def build_paged_kv_slot_mapping( + *, + kv_cache_manager, + request_ids, + qo_lens_cpu: torch.Tensor, + qo_offset_cpu: torch.Tensor, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the backend-neutral paged-cache slot mapping. + + Returns (req_to_token, slot_ids, out_cache_loc), derived only from the paged + KV cache manager and the per-request query geometry, with no dependency on + any backend-specific metadata. + + req_to_token is the [batch, max_kv_len] int32 map from (request, position) + to a global slot id, expanded from get_block_ids_per_seq with + tokens_per_block as block_id * tokens_per_block + offset_within_block. + slot_ids is the [batch] identity row index into req_to_token. out_cache_loc + lists the per-new-token slot ids in flattened query order: request b + contributes positions qo_offset[b] through qo_offset[b] + qo_lens[b] - 1. + That one formula covers prefill (qo_offset is the prefix length) and decode + (qo_offset is kv_len - 1 with qo_len 1). + + The req_to_token reads that build out_cache_loc sync the host, so call this + only from prepare(), never from the forward path. + """ + tokens_per_block = int(kv_cache_manager.tokens_per_block) + # block_ids_per_seq is a [batch, max_blocks_per_seq] tensor; row b holds the + # block ids assigned to request_ids[b] in order. + block_ids = kv_cache_manager.get_block_ids_per_seq(list(request_ids)) + batch = int(qo_lens_cpu.shape[0]) + max_blocks = int(block_ids.shape[1]) + max_kv_len = max_blocks * tokens_per_block + + # Expand block ids -> per-token slot ids. + block_ids_dev = block_ids.to(device).to(torch.int64) + within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) + # Outer product per batch entry: [batch, max_blocks, tokens_per_block] + slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block + req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) + slot_ids = torch.arange(batch, device=device, dtype=torch.int32) + + # out_cache_loc: per-new-token slot ids, in flattened query-token order. + req_to_token_cpu = req_to_token.to("cpu") + qo_lens_list = qo_lens_cpu.to(torch.long).tolist() + qo_offset_list = qo_offset_cpu.to(torch.long).tolist() + out_cache_loc_list: List[int] = [] + for b in range(batch): + start = int(qo_offset_list[b]) + for offset in range(int(qo_lens_list[b])): + out_cache_loc_list.append(int(req_to_token_cpu[b, start + offset].item())) + out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) + return req_to_token, slot_ids, out_cache_loc + + +__all__ = [ + "MiniMaxM3SparseConfig", + "MiniMaxM3SparseMetadataParams", + "MiniMaxM3SparseParams", + "build_paged_kv_slot_mapping", + "write_kv_slots", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py new file mode 100644 index 000000000000..c4751cdec2e4 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Availability checks for the MiniMax-M3 MSA sparse attention kernels. + +The MSA kernels are provided by the fmha_sm100 package from the MSA git +submodule at 3rdparty/MSA and run only on the SM100 architecture family +(SM100 and SM103). These helpers gate backend selection so a request for the +MSA path fails early with a clear message on unsupported systems. +""" + +from __future__ import annotations + +from tensorrt_llm._utils import get_sm_version, is_sm_100f + +from .msa_utils import msa_package_available + +# fmha_sm100 runs on the SM100 architecture family (SM100 and SM103). Other +# architectures, including SM120, are not supported. +MSA_PACKAGE = "fmha_sm100" + + +def ensure_msa_available() -> None: + """Raise RuntimeError if the MSA sparse attention path cannot run here.""" + if not msa_package_available(): + raise RuntimeError( + f"MiniMax-M3 MSA sparse attention requires the {MSA_PACKAGE} kernels " + "from the MSA git submodule at 3rdparty/MSA. Initialize it with " + "'git submodule update --init --recursive'." + ) + if not is_sm_100f(): + sm_version = get_sm_version() + raise RuntimeError( + "MiniMax-M3 MSA sparse attention requires an SM100 or SM103 device, " + f"but the current device reports SM version {sm_version}." + ) + + +__all__ = [ + "MSA_PACKAGE", + "ensure_msa_available", +] 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 new file mode 100644 index 000000000000..7ebc4768bc0f --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -0,0 +1,1017 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA-backed MiniMax-M3 sparse attention on the TrtllmAttention stack. + + * MiniMaxM3MsaSparseAttention subclasses TrtllmAttention and reuses its + inherited forward, overriding only the sparse hooks and owning an + MsaIndexer. + * The main sparse GQA runs through the registered MsaSparseGqaFmha. + * The indexer calls fmha_sm100 directly to produce the per-query selected + block indices, which the model layer threads through + forward_args.topk_indices. + * MiniMaxM3MsaSparseAttentionMetadata subclasses TrtllmAttentionMetadata and + stores its per-forward MSA tensors in CUDA-graph-stable buffers. + The buffers are allocated once in __post_init__ via + get_empty(capture_graph=...), and prepare() copies the per-step values + into them. The standard CUDAGraphRunner clones one metadata per graph + batch size (create_cuda_graph_metadata), so no per-batch-size cache is + needed here. + +The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at +module scope. This is cycle-free because the fmha registry defers its +MsaSparseGqaFmha import (see fmha/registry.py), so trtllm's import chain does +not reach this module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata + +from .common import ( + MiniMaxM3SparseConfig, + MiniMaxM3SparseMetadataParams, + build_paged_kv_slot_mapping, + write_kv_slots, +) +from .msa_indexer import MsaIndexer +from .msa_utils import ( + MSA_REQUIRED_HEAD_DIM, + MSA_REQUIRED_TOPK, + build_kv_page_indices, + per_token_valid_blocks, + require_msa_module, +) + + +def _cache_device(meta) -> torch.device: + """Device hosting the paged KV buffers, else the current CUDA device.""" + kv_cache_manager = meta.kv_cache_manager + if kv_cache_manager is not None: + try: + return kv_cache_manager.get_buffers(0).device + except Exception: + pass + return torch.device(f"cuda:{torch.cuda.current_device()}") + + +def _worst_case_proxy_max_k_tiles( + fmha_sm100, + *, + num_index_heads: int, + kv_cache_manager, + max_batch: int, +) -> int: + """Return max_k_tiles for a proxy plan at the manager's max KV length.""" + page_size = int(kv_cache_manager.tokens_per_block) + max_kv_len = int(kv_cache_manager.max_blocks_per_seq) * page_size + qo_lens = torch.ones(max_batch, dtype=torch.int32) + kv_lens = torch.full((max_batch,), max_kv_len, dtype=torch.int32) + qo_offset = kv_lens - qo_lens + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens, + kv_lens, + num_index_heads, + num_kv_heads=1, + qo_offset=qo_offset, + page_size=page_size, + output_maxscore=True, + num_kv_splits=1, + causal=True, + ) + return int(proxy_plan[3]["max_k_tiles"]) + + +# Per-step fmha_sm100 plan tensors that must live in CUDA-graph-stable buffers. +# At num_kv_splits=1 the plan carries no split-KV workspaces, and +# cute_workspace_buffer is the vendor's cached scratch (kept by reference, not +# copied). +_MSA_PLAN_STABLE_KEYS = ( + "packed_work_range", + "packed_work_info", + "qo_segment_offsets", + "kv_segment_offsets", + "kv_page_indptr", + "qo_segment_lens", + "kv_segment_lens", + "qo_offset", +) +_MSA_PLAN_INT64_KEYS = ("packed_work_range", "packed_work_info") +# fmha_sm100 sizes packed_work_info at 131072 * max(num_kv_splits, 1); forcing +# num_kv_splits=1 pins this worklist width. +_MSA_PACKED_WORK_INFO_LEN = 131072 +_MSA_SPLIT_KV_KEYS = ( + "kv_tile_begin_indices", + "kv_tile_end_indices", + "kv_split_indices", + "num_kv_splits_per_row", + "workspace_o", + "workspace_lse", +) + + +class _MsaGraphSafePlan: + """CUDA-graph-stable mirror of one fmha_sm100 decode plan. + + Owns fixed device buffers for the per-step plan worklists. refresh() copies + a freshly built plan into them and returns a plan tuple pointing at the + stable buffers, so the captured fmha_sm100 run reads addresses that do not + change across replays. Mirrors FlashInfer's fixed indptr/indices buffers. + + Only valid at num_kv_splits=1: the plan then has no split-KV workspaces + (refresh() asserts this), and cute_workspace_buffer and the scalar fields + pass through unchanged. + """ + + def __init__(self, metadata, name: str, *, max_batch: int, num_ctas: int, capture_graph: bool): + buffers = metadata.cuda_graph_buffers + self._buf = {} + # Set by refresh(), read through the plan property. + self._plan: Optional[tuple] = None + # cute_workspace_buffer must keep a fixed address across steps for the + # captured graph to replay correctly. Pin it on first use and fail if + # it moves. + self._ws_ptr: Optional[int] = None + for key in _MSA_PLAN_STABLE_KEYS: + if key == "packed_work_range": + shape = (num_ctas,) + elif key == "packed_work_info": + shape = (_MSA_PACKED_WORK_INFO_LEN,) + elif key in ("qo_segment_offsets", "kv_segment_offsets", "kv_page_indptr"): + shape = (max_batch + 1,) + else: + shape = (max_batch,) + dtype = torch.int64 if key in _MSA_PLAN_INT64_KEYS else torch.int32 + self._buf[key] = metadata.get_empty( + buffers, + shape, + cache_name=f"{name}_{key}", + dtype=dtype, + capture_graph=capture_graph, + ) + + @property + def plan(self) -> Optional[tuple]: + """The current graph-safe plan tuple, or None if no decode plan is live.""" + return self._plan + + def reset(self) -> None: + """Drop the live plan tuple (e.g. for a prefill/mixed or captured step).""" + self._plan = None + + def refresh(self, plan_tuple) -> tuple: + has_mixed, split, batch, decode, prefill = plan_tuple + if has_mixed: + raise RuntimeError( + "MSA decode expects a single (non-mixed) fmha_sm100 plan; a decode " + "batch must be pure decode." + ) + for key in _MSA_SPLIT_KV_KEYS: + if decode.get(key) is not None: + raise RuntimeError( + f"MSA decode plan used split-KV workspace {key!r}; num_kv_splits=1 " + "is required for graph-safe decode." + ) + ws = decode.get("cute_workspace_buffer") + if ws is not None: + if self._ws_ptr is None: + self._ws_ptr = ws.data_ptr() + elif ws.data_ptr() != self._ws_ptr: + raise RuntimeError( + "cute_workspace_buffer moved across steps; the fmha_sm100 plan " + "is not CUDA-graph safe." + ) + rebuilt = dict(decode) + for key in _MSA_PLAN_STABLE_KEYS: + src = decode.get(key) + if src is None: + continue + n = int(src.shape[0]) + dst = self._buf[key] + if n > dst.shape[0]: + raise ValueError( + f"MSA plan buffer {key} ({dst.shape[0]}) is smaller than the plan tensor ({n})." + ) + dst[:n].copy_(src, non_blocking=True) + rebuilt[key] = dst[:n] + self._plan = (has_mixed, split, batch, rebuilt, prefill) + return self._plan + + +@dataclass(init=False) +class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): + """TrtllmAttentionMetadata for MiniMax-M3 MSA sparse layers. + + Tensors read inside the captured forward are CUDA-graph-stable: the + cache slots (msa_out_cache_loc), page table (msa_kv_indices), and proxy + scratch (msa_max_score, msa_n_valid_blocks) are allocated once from the + manager's worst-case geometry. msa_out_cache_loc, msa_kv_indices, and + msa_n_valid_blocks are refreshed via copy_, while the fmha_sm100 proxy pass + writes msa_max_score directly (see msa_proxy_max_score_view). Decode-plan + worklists live on _MsaGraphSafePlan owners, surfaced via msa_decode_*_plan. + + Length inputs to fmha_sm100_plan (msa_qo_lens_cpu, msa_kv_lens_cpu, + msa_qo_offset_cpu) are host properties of the base seq_lens/kv_lens, + read only while building plans in prepare() (outside capture), so they + need no graph-stable storage. Plans are built in _build_step_plans: + pure-decode batches use the graph-safe owners (msa_decode_*_plan) while + prefill/mixed batches keep plain eager tuples (msa_eager_*_plan). + """ + + # Graph-stable buffers; consumers slice to the live count at the call + # site. Filled once the current step's cache write is prepared. + msa_out_cache_loc: Optional[torch.Tensor] = None + msa_kv_indices: Optional[torch.Tensor] = None + msa_max_score: Optional[torch.Tensor] = None + msa_n_valid_blocks: 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. + _msa_buffers_ready: bool = False + _msa_fields_ready: bool = False + # Sparse geometry the decode plans need. + _msa_params: Optional[MiniMaxM3SparseMetadataParams] = None + # Plan owners, created lazily when the decode plans are first built and + # reused across steps. Each owns its graph-safe plan buffers and the + # current refreshed plan tuple. + _msa_proxy_plan: Optional["_MsaGraphSafePlan"] = None + _msa_gqa_plan: Optional["_MsaGraphSafePlan"] = None + _msa_dense_plan: Optional["_MsaGraphSafePlan"] = None + # Eager (prefill/mixed) plans, plain tuples with no graph-stable buffers + # since prefill runs eagerly and is never CUDA-graph captured. Built once + # per step in prepare() and reused by every layer. + _msa_eager_proxy_plan: Optional[tuple] = None + _msa_eager_gqa_plan: Optional[tuple] = None + _msa_eager_dense_plan: Optional[tuple] = None + + def __post_init__(self) -> None: + super().__post_init__() + params = self.sparse_metadata_params + self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None + # See on_update_kv_lens. + 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 + def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: + """Per-request query length (host int32), from the base seq_lens.""" + seq_lens = self.seq_lens + if seq_lens is None: + return None + out = seq_lens[: self.num_seqs] + return out if out.dtype == torch.int32 else out.to(torch.int32) + + @property + def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: + """Per-request KV length, cached plus new tokens (host int32). + + The base ``kv_lens`` includes ``num_extra_kv_tokens`` (speculative + draft-loop slots consumed by the C++ kernels); the MSA plans, ladder + slots and page counts need the true attended length, so it is + excluded here. + """ + 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 + extra = params.num_extra_kv_tokens if params is not None else 0 + if extra: + out = out - extra + return out if out.dtype == torch.int32 else out.to(torch.int32) + + @property + def msa_qo_offset_cpu(self) -> Optional[torch.Tensor]: + """Per-request causal offset (kv_len - qo_len), the cached prefix length.""" + qo = self.msa_qo_lens_cpu + kv = self.msa_kv_lens_cpu + if qo is None or kv is None: + return None + return kv - qo + + @property + def msa_decode_proxy_plan(self) -> Optional[tuple]: + """Proxy (max-score) plan tuple, or None outside decode.""" + plan = self._msa_proxy_plan + return plan.plan if plan is not None else None + + @property + def msa_decode_gqa_plan(self) -> Optional[tuple]: + """Sparse GQA plan tuple, or None outside decode.""" + plan = self._msa_gqa_plan + return plan.plan if plan is not None else None + + @property + def msa_decode_dense_plan(self) -> Optional[tuple]: + """Dense GQA plan tuple, shared by dense layers 0 to 2.""" + plan = self._msa_dense_plan + return plan.plan if plan is not None else None + + @property + def msa_eager_proxy_plan(self) -> Optional[tuple]: + """Prebuilt indexer proxy plan for the eager (prefill/mixed) path.""" + return self._msa_eager_proxy_plan + + @property + def msa_eager_gqa_plan(self) -> Optional[tuple]: + """Prebuilt sparse GQA plan for the eager (prefill/mixed) path.""" + return self._msa_eager_gqa_plan + + @property + def msa_eager_dense_plan(self) -> Optional[tuple]: + """Prebuilt dense GQA plan for the eager (prefill/mixed) path.""" + return self._msa_eager_dense_plan + + def _msa_main_kv_is_fp8(self) -> bool: + """Whether the main paged K/V cache is stored as FP8 E4M3. + + The eager GQA and dense plans must pass use_fp8_kvcache so the inline + sparse-prefill path selects the FP8 AOT kernels; it is a no-op for the + decode planner. Mirrors the k_paged.dtype check in run_msa_paged_gqa. + """ + kv_cache_manager = self.kv_cache_manager + if kv_cache_manager is None: + return False + try: + buffers = kv_cache_manager.get_buffers(0, kv_layout="HND") + except TypeError: + buffers = kv_cache_manager.get_buffers(0) + except Exception: + return False + + try: + return buffers[:, 0].dtype == torch.float8_e4m3fn + except Exception: + return False + + def _create_msa_buffers(self) -> None: + """Allocate the CUDA-graph-stable MSA device buffers. + + Buffers come from the shared graph buffer pool so they are reserved + under capture. Sizing follows the worst-case graph geometry: + max_num_tokens for cache slots, max_num_sequences * max_blocks_per_seq + for the page table, and worst-case max_k_tiles for proxy scratch. + """ + kv_cache_manager = self.kv_cache_manager + self._msa_buffers_ready = False + if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): + return + capture_graph = self.is_cuda_graph + buffers = self.cuda_graph_buffers + max_num_sequences = int(self.max_num_sequences) + max_blocks_per_seq = int(kv_cache_manager.max_blocks_per_seq) + max_total_pages = max_num_sequences * max_blocks_per_seq + max_num_tokens = int(self.max_num_tokens) + + self.msa_out_cache_loc = self.get_empty( + buffers, + (max_num_tokens,), + cache_name="msa_out_cache_loc", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.msa_kv_indices = self.get_empty( + buffers, + (max_total_pages,), + cache_name="msa_kv_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Staging for on_update_kv_lens: re-derives slots/bounds on device + # from the corrected kv_lens_cuda, sync-free. + 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 + # to skip allocation. + params = self._msa_params + if params is not None: + fmha_sm100 = require_msa_module() + max_k_tiles = _worst_case_proxy_max_k_tiles( + fmha_sm100, + num_index_heads=params.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_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 decode-step query tokens (spec verify emits 1 + draft_len + per row), bounded by max_num_tokens. getattr fallbacks cover metadata + built via ``__new__`` in structural tests. + """ + max_seqs = int(getattr(self, "max_num_sequences", 0) or 0) + max_toks = int(getattr(self, "max_num_tokens", 0) or 0) + if max_toks <= 0: + return max_seqs + # 16384 keeps total_q * num_qo_heads under the fmha_sm100 planner cap + # (65536, fmha_sm100/api.py) at 4 sharded index heads. + return max(max_seqs, min(max_toks, 16384)) + + def _alloc_msa_proxy_scratch( + self, + *, + num_index_heads: 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 and the worst-case + per-step query-token count (which exceeds the batch size under + speculative multi-token verify), so one allocation serves every + decode step. msa_proxy_max_score_view slices the per-step shape out + of it. + """ + buffers = self.cuda_graph_buffers + self.msa_max_score = self.get_empty( + buffers, + (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_tokens,), + cache_name="msa_n_valid_blocks", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + def _ensure_msa_decode_scratch_buffers( + self, + *, + num_index_heads: int, + max_batch: int, + capture_graph: bool, + required_max_k_tiles: int, + ) -> None: + """Ensure proxy scratch buffers exist and cover the current plan.""" + 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_tokens} tokens)." + ) + return + + kv_cache_manager = self.kv_cache_manager + if kv_cache_manager is None: + return + + fmha_sm100 = require_msa_module() + max_k_tiles = _worst_case_proxy_max_k_tiles( + fmha_sm100, + num_index_heads=num_index_heads, + kv_cache_manager=kv_cache_manager, + max_batch=max_batch, + ) + if max_k_tiles < required_max_k_tiles: + raise ValueError( + f"Worst-case max_k_tiles ({max_k_tiles}) is less than the " + f"decode plan ({required_max_k_tiles})." + ) + self._alloc_msa_proxy_scratch( + num_index_heads=num_index_heads, + max_tokens=max_tokens, + max_k_tiles=max_k_tiles, + capture_graph=capture_graph, + ) + + 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: + """Re-derive length-dependent MSA state from the corrected kv_lens_cuda. + + The overlap scheduler corrects kv_lens_cuda on device after prepare() + staged optimistic (full-acceptance) lens. Decode steps are patched + with pure device ops (capture-safe); the correction only shrinks + lengths, so the host-baked plan worklists and the page-table layout + stay valid. Eager mixed batches install corrected host lens and + rebuild (small D2H sync). Idempotent. + """ + super().on_update_kv_lens() + if not self._msa_fields_ready: + return + batch = self._msa_live_batch + total_q = self._msa_live_total_q + if batch <= 0 or total_q <= 0: + return + if self.msa_decode_proxy_plan is None: + # Prefill/mixed step: the eager plans and the page-table layout + # were built from the optimistic host mirrors, so install the + # corrected lens and rebuild both. Mixed steps are never captured. + if torch.cuda.is_current_stream_capturing(): + return + self._msa_corrected_kv_lens_cpu = self.kv_lens_cuda[:batch].to("cpu", torch.int32) + self._build_msa_fields() + self._build_step_plans() + return + + kv_true = self.kv_lens_cuda[:batch] + qbr = self.msa_q_batch_row[:total_q].to(torch.long) + qo_dev = self.msa_qo_lens_dev[:batch] + kv_true_tok = kv_true[qbr] + pos = kv_true_tok - qo_dev[qbr] + self.msa_q_intra[:total_q] + + # KV/idx-K write slots: slot[j] = req_to_token[request[j], pos[j]]. + width = int(self.msa_req_to_token.shape[1]) + idx = pos.to(torch.long).clamp(min=0, max=width - 1) + slots = self.msa_req_to_token.reshape(-1).index_select(0, qbr * width + idx) + self.msa_out_cache_loc[:total_q].copy_(slots) + + # Per-token valid-block counts for top-k selection. clamp_min(1) + # keeps degenerate (graph-padding) rows from -inf-masking every + # block, which would NaN the fully-masked GQA row. + page = self._msa_page_size + n_valid = torch.div((pos + 1).clamp_min(1) + (page - 1), page, rounding_mode="floor") + self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32)) + + # Plan mirrors: proxy and dense keep one row per request; the sparse + # GQA plan is row-expanded per token with the ladder in the per-row + # offset. qo_offset must stay non-negative: negative values hit the + # kernel's packed-length sentinel fallback. + off_req = (kv_true - qo_dev).clamp_min(0) + for owner, expanded in ( + (self._msa_proxy_plan, False), + (self._msa_gqa_plan, True), + (self._msa_dense_plan, False), + ): + decode = owner.plan[3] + seg_lens = decode.get("kv_segment_lens") + qo_off = decode.get("qo_offset") + if expanded: + if seg_lens is not None: + seg_lens[:total_q].copy_(kv_true_tok.to(seg_lens.dtype)) + if qo_off is not None: + qo_off[:total_q].copy_(pos.clamp_min(0).to(qo_off.dtype)) + else: + if seg_lens is not None: + seg_lens[:batch].copy_(kv_true.to(seg_lens.dtype)) + if qo_off is not None: + qo_off[:batch].copy_(off_req.to(qo_off.dtype)) + + def _build_step_plans(self) -> None: + """Build the three layer-invariant fmha_sm100 plans once per step. + + Runs in prepare(), outside CUDA graph capture. The proxy, GQA, and + dense plans depend only on the per-step sparse geometry (qo/kv lengths, + head counts, topk, page size), never on the layer, so they are built + once here and reused by every layer: + + * Pure-decode batches mirror the plans into the CUDA-graph-stable + _MsaGraphSafePlan buffers (surfaced by msa_decode_*_plan), because + decode is captured and the plan worklists must keep a fixed address + across replays. + * Prefill, chunked-prefill, and mixed batches run eagerly (never + captured), so the plans are stored as plain tuples (msa_eager_*_plan) + that every sparse and dense layer reuses. + """ + # Drop any plan tuples from the previous step; the msa_decode_*_plan and + # msa_eager_*_plan properties then report None until rebuilt below. + for plan in (self._msa_proxy_plan, self._msa_gqa_plan, self._msa_dense_plan): + if plan is not None: + plan.reset() + self._msa_eager_proxy_plan = None + self._msa_eager_gqa_plan = None + self._msa_eager_dense_plan = None + if not self._msa_fields_ready: + return + # Geometry is captured in __post_init__; skip when it is unavailable. + params = self._msa_params + if params is None: + return + num_index_heads = params.num_index_heads + num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) + topk = params.topk + + fmha_sm100 = require_msa_module() + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + 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 + device = _cache_device(self) + page_size = int(self.kv_cache_manager.tokens_per_block) + capture_graph = self.is_cuda_graph + max_batch = int(self.max_num_sequences) + # A decode batch is pure generation (no context requests). Only that + # path is CUDA-graph captured and uses the graph-stable plan buffers. + is_decode = int(self.num_contexts or 0) == 0 + # The main-attention GQA and dense plans need use_fp8_kvcache so the + # eager (inline sparse-prefill) kernel selection matches an FP8 paged + # cache; it is a no-op for the decode planner. The proxy runs over the + # bf16 index-K cache, so it never needs the flag. + use_fp8 = self._msa_main_kv_is_fp8() + + # Proxy plan: MQA (num_kv_heads=1) max-score pass over the index + # branch; output_maxscore feeds the indexer's top-k block selection. + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_index_heads, + num_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=page_size, + output_maxscore=True, + num_kv_splits=1, + causal=True, + ) + # Sparse-layer plan: kv_block_num=topk limits attention to top-k blocks. + gqa_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_q_heads, + num_kv_heads=num_kv_heads, + qo_offset=qo_offset_cpu, + page_size=page_size, + kv_block_num=topk, + num_kv_splits=1, + causal=True, + use_fp8_kvcache=use_fp8, + ) + # Dense-layer plan: no kv_block_num, so it attends the full page table. + dense_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + num_q_heads, + num_kv_heads=num_kv_heads, + qo_offset=qo_offset_cpu, + page_size=page_size, + num_kv_splits=1, + causal=True, + use_fp8_kvcache=use_fp8, + ) + + if not is_decode: + # Prefill and mixed batches run eagerly, so keep the plain plan + # tuples and leave the graph-safe owners reset. The indexer computes + # its own per-query valid-block count on the eager path. + self._msa_eager_proxy_plan = proxy_plan + self._msa_eager_gqa_plan = gqa_plan + self._msa_eager_dense_plan = dense_plan + return + + required_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + self._ensure_msa_decode_scratch_buffers( + num_index_heads=num_index_heads, + max_batch=max_batch, + capture_graph=capture_graph, + required_max_k_tiles=required_max_k_tiles, + ) + + # Allocate the graph-safe plan owners once per metadata; later steps + # only refresh their contents below. The plan worklists are sized per + # expanded row (the planner splits qo_len > 1 requests into per-token + # rows under speculative multi-token verify), so use the worst-case + # decode-step token count rather than the batch size. + 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_plan_rows, + num_ctas=num_ctas, + capture_graph=capture_graph, + ) + self._msa_gqa_plan = _MsaGraphSafePlan( + self, + "msa_gqa_plan", + 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_plan_rows, + num_ctas=num_ctas, + capture_graph=capture_graph, + ) + + # refresh() stores each plan tuple on its owner, surfaced by the + # msa_decode_*_plan properties. + self._msa_proxy_plan.refresh(proxy_plan) + self._msa_gqa_plan.refresh(gqa_plan) + self._msa_dense_plan.refresh(dense_plan) + + n_valid = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + # One entry per query token (qo_len > 1 under spec verify). + 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. + + The page table and per-new-token cache slots are derived via the + build_paged_kv_slot_mapping helper, then copied into the persistent + buffers. The transient builder tensors are discarded. + """ + self._msa_fields_ready = False + if not self._msa_buffers_ready: + return + request_ids = self.request_ids + qo_lens_cpu = self.msa_qo_lens_cpu + kv_lens_cpu = self.msa_kv_lens_cpu + qo_offset_cpu = self.msa_qo_offset_cpu + if request_ids is None or qo_lens_cpu is None: + return + batch_size = int(qo_lens_cpu.shape[0]) + if batch_size == 0: + return + + kv_cache_manager = self.kv_cache_manager + cache_device = _cache_device(self) + page_size = int(kv_cache_manager.tokens_per_block) + + # 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 + # (num_cached) and decode (kv_len - 1 with qo_len 1). + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=cache_device, + ) + kv_indices = build_kv_page_indices(req_to_token, slot_ids, kv_lens_cpu, page_size) + + total_new_tokens = int(out_cache_loc.shape[0]) + total_pages = int(kv_indices.shape[0]) + if total_new_tokens > self.msa_out_cache_loc.shape[0]: + raise ValueError( + f"MSA out_cache_loc buffer ({self.msa_out_cache_loc.shape[0]}) is " + f"smaller than the step's new-token count ({total_new_tokens})." + ) + if total_pages > self.msa_kv_indices.shape[0]: + raise ValueError( + f"MSA kv_indices buffer ({self.msa_kv_indices.shape[0]}) is " + f"smaller than the step's page count ({total_pages})." + ) + + 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) + + # Staging for on_update_kv_lens. + 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_long = qo_lens_cpu.to(torch.long) + batch_row_cpu = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32), qo_long + ) + starts = torch.cumsum(qo_long, 0) - qo_long + intra_cpu = ( + torch.arange(total_new_tokens, dtype=torch.int64) + - torch.repeat_interleave(starts, qo_long) + ).to(torch.int32) + self.msa_q_batch_row[:total_new_tokens].copy_(batch_row_cpu) + self.msa_q_intra[:total_new_tokens].copy_(intra_cpu) + 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: + """Return the paged index-K cache in the HND layout MSA consumes.""" + return self.kv_cache_manager.get_index_k_buffer(layer_idx, kv_layout="HND") + + def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: + """Write the new-token index-K into the side cache at out_cache_loc.""" + cache = self.msa_idx_k_cache(layer_idx) + sparse_index_dim = int(cache.shape[-1]) + num_tokens = int(idx_k.shape[0]) + write_kv_slots( + cache, + self.msa_out_cache_loc[:num_tokens], + idx_k.reshape(num_tokens, 1, sparse_index_dim), + layout="HND", + ) + + def msa_proxy_max_score_view( + self, num_index_heads: int, plan_max_k_tiles: int, num_tokens: int + ) -> torch.Tensor: + """Return a contiguous [num_index_heads, plan_max_k_tiles, num_tokens] view. + + fmha_sm100 ignores the passed tensor's strides and writes a contiguous + [num_index_heads, plan_max_k_tiles, total_q] block sized by the current + decode plan, so it must receive a tensor contiguous in exactly that + shape. The view is taken from the flat store's prefix starting at offset + 0, so its data_ptr is stable for CUDA graph replay. Capture builds the + decode plan at the worst-case max_k_tiles, so replays only shrink it. + """ + store = self.msa_max_score + numel = num_index_heads * plan_max_k_tiles * num_tokens + if numel > store.numel(): + raise ValueError( + f"msa_max_score backing store ({store.numel()} elements) is " + f"smaller than the proxy view needs ({numel} = {num_index_heads} " + f"heads * {plan_max_k_tiles} k-tiles * {num_tokens} tokens)." + ) + return store[:numel].view(num_index_heads, plan_max_k_tiles, num_tokens) + + +class MiniMaxM3MsaSparseAttention(TrtllmAttention): + """MSA-backed MiniMax-M3 sparse attention.""" + + Metadata = MiniMaxM3MsaSparseAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config=None, + *, + sparse_params, + **kwargs, + ): + TrtllmAttention.__init__( + self, + layer_idx, + num_heads, + head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + sparse_params=sparse_params, + **kwargs, + ) + self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, + num_q_heads=num_heads, + num_kv_heads=num_kv_heads or num_heads, + head_dim=head_dim, + ) + self.disable_index_value = bool(sparse_params.disable_index_value) + self._validate_msa_preconditions() + self.indexer = MsaIndexer(self.m3_config) + + def _validate_msa_preconditions(self) -> None: + config = self.m3_config + if not self.disable_index_value: + raise NotImplementedError( + "MSA backend requires disable_index_value=True; the proxy pass " + "consumes only the max score and has no index-V path." + ) + if config.head_dim != MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires head_dim={MSA_REQUIRED_HEAD_DIM}, got {config.head_dim}." + ) + if config.sparse_index_dim != MSA_REQUIRED_HEAD_DIM: + raise NotImplementedError( + f"MSA backend requires sparse_index_dim={MSA_REQUIRED_HEAD_DIM}, " + f"got {config.sparse_index_dim}." + ) + if config.topk != MSA_REQUIRED_TOPK: + raise NotImplementedError( + f"MSA backend requires topk={MSA_REQUIRED_TOPK}, got {config.topk}." + ) + + @classmethod + def support_fused_rope(cls) -> bool: + # The MiniMax-M3 model layer applies partial RoPE to the main and + # index branches explicitly. + return False + + def run_indexer( + self, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + metadata, + *, + idx_sm_scale: Optional[float] = None, + ) -> torch.Tensor: + """Write the index-K cache and return the selected block indices. + + The model layer runs this before forward and threads the result through + forward_args.topk_indices. Returns [total_q, num_kv_heads, topk]. + Decode uses the prebuilt graph-safe proxy plan; prefill and mixed + batches use the prebuilt eager proxy plan. + """ + 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]) + 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) + + metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) + + # One selection path: decode passes the prebuilt graph-safe proxy plan + # plus the proxy scratch shaped to the live query count. Prefill and + # mixed batches pass the prebuilt eager proxy plan and let the kernel + # allocate the score buffer and the indexer compute the per-query + # valid-block count. Only when neither is present (for example a + # standalone test that skips prepare) does select_blocks build the proxy + # plan inline. + proxy_plan = metadata.msa_decode_proxy_plan + if proxy_plan is not None: + # proxy_plan is (has_mixed, split, batch, decode_dict, prefill); + # decode_dict carries max_k_tiles for the contiguous score view. + plan_max_k_tiles = int(proxy_plan[3]["max_k_tiles"]) + max_score = metadata.msa_proxy_max_score_view( + config.num_index_heads, plan_max_k_tiles, num_tokens + ) + n_valid_blocks = metadata.msa_n_valid_blocks[:num_tokens] + else: + proxy_plan = metadata.msa_eager_proxy_plan + max_score = None + n_valid_blocks = None + return self.indexer.select_blocks( + idx_q_view, + idx_k_cache, + idx_sm_scale=idx_sm_scale, + kv_indices=metadata.msa_kv_indices, + qo_lens_cpu=metadata.msa_qo_lens_cpu, + kv_lens_cpu=metadata.msa_kv_lens_cpu, + qo_offset_cpu=metadata.msa_qo_offset_cpu, + proxy_plan=proxy_plan, + max_score=max_score, + n_valid_blocks=n_valid_blocks, + ) + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata, + forward_args: "AttentionForwardArgs", + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + # The model layer runs run_indexer and passes the selected block + # indices through forward_args.topk_indices. Publish them as the + # sparse attention indices MsaSparseGqaFmha reads. + return forward_args.topk_indices, None + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata, + forward_args: "AttentionForwardArgs", + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + return None, None + + +__all__ = [ + "MiniMaxM3MsaSparseAttention", + "MiniMaxM3MsaSparseAttentionMetadata", +] 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 new file mode 100644 index 000000000000..b4a836d03ffa --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MSA sparse-attention indexer. + +Mirrors the DSA indexer pattern: a submodule owned by the sparse backend +that runs the predictor pass and returns the per-query selected KV block +indices the main attention consumes. It calls fmha_sm100 directly in +output_maxscore mode, reduces the per-index-head max score to KV-head +granularity, and selects the top-k blocks per query. + +Results are [total_q, num_kv_heads, topk] int32, ascending with -1 padding. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from .msa_utils import ( + MSA_REQUIRED_TOPK, + per_token_valid_blocks, + require_msa_module, + select_blocks_from_maxscore, +) + +if TYPE_CHECKING: + from .common import MiniMaxM3SparseConfig + + +def _proxy_max_score( + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + kv_indices: torch.Tensor, + sm_scale: float, + causal: bool, +) -> torch.Tensor: + """Run the fmha_sm100 MQA proxy pass and return the per-block max score. + + Follows MSA's two-call pattern: fmha_sm100_plan builds the plan with + output_maxscore and num_kv_heads 1, then fmha_sm100 runs with output_o + disabled so only the per-block max score is produced. Returns + [num_index_heads, max_k_tiles, total_q] float32. + """ + fmha_sm100 = require_msa_module() + + if idx_q.dim() != 3: + raise ValueError( + "MsaIndexer expects idx_q [total_q, num_index_heads, head_dim]; " + f"got {tuple(idx_q.shape)}." + ) + if idx_k_paged.dim() != 4 or idx_k_paged.shape[1] != 1: + raise ValueError( + "MsaIndexer expects MQA paged index-K [num_pages, 1, page_size, head_dim]; " + f"got {tuple(idx_k_paged.shape)}." + ) + + page_size = int(idx_k_paged.shape[2]) + proxy_plan = fmha_sm100.fmha_sm100_plan( + qo_lens_cpu, + kv_lens_cpu, + idx_q.shape[1], + num_kv_heads=1, + qo_offset=qo_offset_cpu, + page_size=page_size, + output_maxscore=True, + causal=causal, + num_kv_splits=1, + ) + _, max_score = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + sm_scale=sm_scale, + ) + return max_score + + +def _group_max_reduce( + max_score: torch.Tensor, + config: "MiniMaxM3SparseConfig", +) -> torch.Tensor: + """Reduce per-index-head max score to per-KV-head granularity by amax. + + Index heads are assumed to be grouped contiguously per KV head, so head h + maps to KV group h // group. + """ + group, rem = divmod(config.num_index_heads, config.num_kv_heads) + if rem != 0: + raise ValueError( + "num_index_heads must be divisible by num_kv_heads for group max " + f"reduce; got num_index_heads={config.num_index_heads}, " + f"num_kv_heads={config.num_kv_heads}." + ) + if group > 1: + return max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + return max_score + + +class MsaIndexer: + """Predictor submodule: proxy MQA scoring and top-k block selection. + + Owned by the MSA attention layer. Stateless in eager mode: it reads the + per-forward page table and lengths from the attention metadata and calls + the kernel directly. + """ + + def __init__(self, config: "MiniMaxM3SparseConfig"): + self.config = config + + def select_blocks( + self, + idx_q: torch.Tensor, + idx_k_paged: torch.Tensor, + *, + idx_sm_scale: float, + kv_indices: torch.Tensor, + qo_lens_cpu: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, + qo_offset_cpu: Optional[torch.Tensor] = None, + proxy_plan: Optional[tuple] = None, + max_score: Optional[torch.Tensor] = None, + n_valid_blocks: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Return [total_q, num_kv_heads, topk] selected block indices. + + Plan/run split, mirroring the sparse GQA. When `proxy_plan` is None + (prefill and focused tests) the proxy plan is built inline and the + per-query valid-block count is derived here; when provided (CUDA-graph + decode) the proxy runs from the prebuilt plan into the preallocated + `max_score` buffer with a precomputed `n_valid_blocks`, so there is + no host sync inside the captured region. The same top-k selection serves + both, and generation is the one-query-token-per-request special case. + """ + config = self.config + + if proxy_plan is None: + max_score = _proxy_max_score( + idx_q, + idx_k_paged, + qo_lens_cpu=qo_lens_cpu, + kv_lens_cpu=kv_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + kv_indices=kv_indices, + sm_scale=idx_sm_scale, + causal=True, + ) + else: + fmha_sm100 = require_msa_module() + _, max_score = fmha_sm100.fmha_sm100( + idx_q, + idx_k_paged, + idx_k_paged, + proxy_plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + max_score=max_score, + sm_scale=idx_sm_scale, + ) + max_score_kv = _group_max_reduce(max_score, config) + + if n_valid_blocks is None: + n_valid_blocks = per_token_valid_blocks( + qo_lens_cpu, + kv_lens_cpu, + qo_offset_cpu, + causal=True, + block_size=int(idx_k_paged.shape[2]), + ) + # The empty-selection guard uses a host sync, so it only runs on the + # eager path; a decode batch always has valid blocks. + 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), + -1, + dtype=torch.int32, + device=idx_q.device, + ) + 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, + ) + + +__all__ = ["MsaIndexer"] 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 new file mode 100644 index 000000000000..33547d84d54b --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA (fmha_sm100) specific helpers for the MiniMax-M3 sparse backend.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from typing import Optional, Tuple + +import torch + +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 +# with a clear message rather than a cryptic shape error inside the kernel. +MSA_REQUIRED_TOPK = 16 +MSA_REQUIRED_HEAD_DIM = 128 + +# Path of the fmha_sm100 package inside the MSA git submodule relative to the +# repository root (see 3rdparty/MSA/LICENSE and 3rdparty/MSA/NOTICE). +_MSA_PYTHON_RELPATH = Path("3rdparty") / "MSA" / "python" + + +def _find_msa_python_dir() -> Optional[Path]: + """Locate the fmha_sm100 package dir by walking up from this file. + + Returns None in installed layouts where the 3rdparty submodule is not + shipped. Walking up avoids hardcoding this module's depth below the + repository root. + """ + for parent in Path(__file__).resolve().parents: + candidate = parent / _MSA_PYTHON_RELPATH + if candidate.is_dir(): + return candidate + return None + + +def _ensure_msa_on_path() -> None: + """Prepend the MSA python package directory to sys.path if present.""" + msa_python = _find_msa_python_dir() + if msa_python is not None and str(msa_python) not in sys.path: + sys.path.insert(0, str(msa_python)) + + +def msa_package_available() -> bool: + """True if fmha_sm100 can be imported (submodule checkout or installed).""" + _ensure_msa_on_path() + return importlib.util.find_spec("fmha_sm100") is not None + + +def require_msa_module(): + """Import fmha_sm100 from the MSA submodule or raise a clear error. + + The import is deferred to first kernel use so the MSA backend can be + advertised in the config schema on systems where the kernels cannot load. + The 3rdparty/MSA/python directory is added to sys.path first, so a source + checkout with the submodule initialized resolves without a separate install. + A missing package is a hard error, never a silent fallback to another backend. + """ + _ensure_msa_on_path() + try: + import fmha_sm100 + except ImportError as exc: + raise RuntimeError( + "MiniMax-M3 MSA attention requires the fmha_sm100 kernels from the " + "MSA git submodule at 3rdparty/MSA. Initialize it with " + "'git submodule update --init --recursive', or install fmha_sm100." + ) from exc + return fmha_sm100 + + +def msa_paged_kv(kv_cache_manager, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Return per-layer paged K and V in fmha_sm100 HND layout, zero-copy. + + The cache is stored head-major (see `write_msa_main_kv`), so the "HND" + buffer view is already the [num_slots, num_kv_heads, page_size, head_dim] + layout fmha_sm100 expects. The kernel reads the page and head strides at + runtime and needs only each page's [page_size, head_dim] block to be + contiguous, which this view satisfies, so no copy is required. + """ + buffers = kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + return buffers[:, 0], buffers[:, 1] + + +def write_msa_main_kv( + kv_cache_manager, + layer_idx: int, + out_cache_loc: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + """Write new-token K and V into the paged main cache at out_cache_loc. + + fmha_sm100 reads the paged cache directly, so the new-token K and V must be + resident before the sparse GQA runs. The write uses the head-major HND view + so `msa_paged_kv` can return a zero-copy view. + """ + buffers = kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + k_view, v_view = buffers[:, 0], buffers[:, 1] + num_kv_heads = int(k_view.shape[1]) + head_dim = int(k_view.shape[3]) + num_tokens = int(k.shape[0]) + write_kv_slots( + k_view, out_cache_loc, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + write_kv_slots( + v_view, out_cache_loc, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + + +def build_kv_page_indices( + req_to_token: torch.Tensor, + slot_ids: torch.Tensor, + kv_lens_cpu: torch.Tensor, + page_size: int, +) -> torch.Tensor: + """Build the flattened per-request page table fmha_sm100 consumes. + + Returns int32 global page ids concatenated per request. A request's + pages come from the first slot of each page in its req_to_token row. + Page ids are global and non-contiguous in production, so they are not + clamped to a per-request bound. + """ + device = req_to_token.device + req_rows = req_to_token.index_select(0, slot_ids.to(torch.long)).to(torch.long) + batch = int(req_rows.shape[0]) + kv_lens_list = kv_lens_cpu.to(torch.long).tolist() + + page_lists = [] + for b in range(batch): + kv_len = int(kv_lens_list[b]) + if kv_len <= 0: + continue + num_pages = (kv_len + page_size - 1) // page_size + page_starts = torch.arange(num_pages, device=device, dtype=torch.long) * page_size + page_ids = req_rows[b].gather(0, page_starts) // page_size + page_lists.append(page_ids.to(torch.int32)) + + if page_lists: + return torch.cat(page_lists, dim=0) + return torch.empty(0, dtype=torch.int32, device=device) + + +def per_token_valid_blocks( + qo_lens_cpu: torch.Tensor, + kv_lens_cpu: torch.Tensor, + qo_offset_cpu: Optional[torch.Tensor], + *, + causal: bool, + block_size: int, +) -> torch.Tensor: + """Return the per-query number of valid KV blocks, on CPU. + + Expands per-request lengths and offsets to a per-token vector so block + selection can honour each query token's own causal extent. + """ + qo = qo_lens_cpu.to(torch.long) + kv = kv_lens_cpu.to(torch.long) + batch = int(qo.shape[0]) + total = int(qo.sum().item()) + if total == 0: + return torch.zeros(0, dtype=torch.long) + batch_row = torch.repeat_interleave(torch.arange(batch, dtype=torch.long), qo) + starts = torch.zeros(batch, dtype=torch.long) + if batch > 1: + starts[1:] = torch.cumsum(qo, 0)[:-1] + intra = torch.arange(total, dtype=torch.long) - starts[batch_row] + kv_per = kv[batch_row] + if causal: + if qo_offset_cpu is not None: + off = qo_offset_cpu.to(torch.long)[batch_row] + else: + off = (kv - qo)[batch_row] + eff = torch.minimum(off + intra + 1, kv_per) + else: + eff = kv_per + return (eff + block_size - 1) // block_size + + +def select_blocks_from_maxscore( + max_score_kv: torch.Tensor, + *, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, +) -> 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. + """ + 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 + ) + + +__all__ = [ + "MSA_REQUIRED_HEAD_DIM", + "MSA_REQUIRED_TOPK", + "build_kv_page_indices", + "msa_package_available", + "msa_paged_kv", + "per_token_valid_blocks", + "require_msa_module", + "select_blocks_from_maxscore", + "write_msa_main_kv", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py similarity index 63% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py index 090c6d7956fd..844fd979469c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py @@ -46,7 +46,7 @@ ----------------- All scalar max lengths (``max_seqlen_q``, ``max_seqlen_k``) are -pre-computed CPU-side in :meth:`MiniMaxM3SparseAttentionMetadata.prepare` +pre-computed CPU-side in :meth:`MiniMaxM3TritonSparseAttentionMetadata.prepare` and stored as plain Python ints. The hot path uses only batched tensor ops with static shapes derived from those CPU-side scalars; no ``.item()`` or other GPU-CPU sync runs inside the forward @@ -56,29 +56,29 @@ from __future__ import annotations -import functools import os from dataclasses import dataclass from typing import TYPE_CHECKING, Optional, Tuple import torch -from .kernels import triton_block_max_score, triton_sparse_softmax -from .metadata import ( - MiniMaxM3SparseAttentionMetadata, - MiniMaxM3SparseConfig, +from ...interface import ( + AttentionBackend, + AttentionForwardArgs, + AttentionMetadata, + merge_attention_forward_args, +) +from .common import _INIT_SCORE, _LOCAL_SCORE, MiniMaxM3SparseConfig, write_kv_slots +from .triton_kernels import triton_block_max_score, triton_sparse_softmax +from .triton_metadata import ( + MiniMaxM3AttentionMetadata, + MiniMaxM3TritonSparseAttentionMetadata, ensure_metadata_on_device, - get_minimax_m3_attention_metadata_cls, ) if TYPE_CHECKING: from .cache_manager import MiniMaxM3SparseIndexCache - from .metadata import MiniMaxM3SparseParams - -# Sentinel block score for blocks that init / local priority forces into -# the top-k regardless of their numerical score. -_INIT_SCORE = 1e30 -_LOCAL_SCORE = 1e29 + from .common import MiniMaxM3SparseParams # --------------------------------------------------------------------------- @@ -143,53 +143,6 @@ def _gather_paged_batched( return flat.view(batch, max_k, *cache.shape[1:]) -def _assert_paged_write_in_bounds( - name: str, - cache: torch.Tensor, - page: torch.Tensor, - within: torch.Tensor, -) -> None: - """Optional CPU-side bounds check for paged-cache writes. - - The runtime computes per-token slot ids from - ``KVCacheManagerV2``'s block ids; if the runtime ever produces a - block id that does not fit the per-layer view's dim-0 the write - falls into another layer's coalesced memory and corrupts the - cache, or fires the CUDA ``IndexKernel.cu`` device-side assert - during fancy indexing. Both are far enough away from the root - cause to be hard to triage. - - When ``TRTLLM_MINIMAX_M3_DEBUG_BOUNDS`` is set this check runs a - CPU-side max/min comparison against the cache's dim-0 and dim-2 - bounds, surfacing the misindex with the exact tensor names and - values instead of a device-side assert spam. It is opt-in - because the comparison forces a CPU sync. - """ - if not os.environ.get("TRTLLM_MINIMAX_M3_DEBUG_BOUNDS"): - return - num_pages = int(cache.shape[0]) - tokens_per_block = int(cache.shape[1]) if cache.ndim == 4 else int(cache.shape[2]) - page_max = int(page.max().item()) if page.numel() else -1 - page_min = int(page.min().item()) if page.numel() else 0 - within_max = int(within.max().item()) if within.numel() else -1 - within_min = int(within.min().item()) if within.numel() else 0 - assert 0 <= page_min and page_max < num_pages, ( - f"{name}: page index out of bounds — page.min={page_min} " - f"page.max={page_max} but cache.shape[0]={num_pages} " - f"(shape={tuple(cache.shape)}). This usually means the " - f"runtime's get_block_ids_per_seq produced a block id wider " - f"than the per-layer paged view's dim-0; check that the M3 " - f"override path returns slot ids in [0, num_slots)." - ) - assert 0 <= within_min and within_max < tokens_per_block, ( - f"{name}: within-page offset out of bounds — within.min=" - f"{within_min} within.max={within_max} but tokens_per_block=" - f"{tokens_per_block} (shape={tuple(cache.shape)}). This " - f"usually means out_cache_loc was computed with a different " - f"tokens_per_block than the cache was allocated with." - ) - - def _write_main_kv_slots_to_pool( pool: torch.Tensor, kv_index: int, @@ -201,40 +154,12 @@ def _write_main_kv_slots_to_pool( ``pool`` is the 5-D main K/V pool returned by :meth:`KVCacheManagerV2.get_buffers` with the NHD layout ``[num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]``. - ``values`` has shape ``[num_new_tokens, num_kv_heads, head_dim]`` - and ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of - flat slot ids the caller wants to update. - - The write decomposes each flat slot id into - ``(page = s // tokens_per_block, within = s % tokens_per_block)`` - and uses multi-dim fancy-index assignment so the writes propagate - to the underlying pool storage. The previously used pattern - ``pool[:, kv_index].reshape(-1, num_kv_heads, head_dim) - .index_copy_(0, ...)`` instead wrote into a silent copy (see - :func:`_gather_paged_batched`), so the next forward call read - zeros for the prefilled positions. - - The optional CPU-side bounds assertion (enabled when the - ``TRTLLM_MINIMAX_M3_DEBUG_BOUNDS`` env var is set) catches - block_ids overflowing the pool's dim-0 before the device-side - ``IndexKernel.cu`` assert fires deep inside the kernel. The - assertion is a CPU sync, so the env var keeps it opt-in for - production runs that need a clean fast path. + ``values`` has shape ``[num_new_tokens, num_kv_heads, head_dim]`` and + ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of flat slot + ids to update. ``pool[:, kv_index]`` is a storage-sharing view, so the + shared :func:`common.write_kv_slots` propagates the write to the pool. """ - tokens_per_block = int(pool.shape[2]) - out_long = out_cache_loc.to(torch.long) - page = out_long // tokens_per_block - within = out_long % tokens_per_block - _assert_paged_write_in_bounds("pool", pool, page, within) - # KV-cache writes never need to participate in autograd. Wrap the - # fancy-index assignment in ``torch.no_grad()`` so callers that - # enter this path with an active grad context (e.g. unit tests - # exercising :class:`MiniMaxM3Attention` without ``inference_mode``) - # do not trip the "leaf Variable that requires grad is being used - # in an in-place operation" autograd guard on the view chain. - with torch.no_grad(): - # Multi-dim fancy assignment writes into the underlying pool buffer. - pool[page, kv_index, within] = values.to(pool.dtype) + write_kv_slots(pool[:, kv_index], out_cache_loc, values) def _write_main_kv_slots( @@ -242,39 +167,13 @@ def _write_main_kv_slots( out_cache_loc: torch.Tensor, values: torch.Tensor, ) -> None: - """Layout-aware writer for K (or V) caches used by the M3 backend. - - Supports two layouts, mirroring :func:`_gather_paged_batched`: - - * **3-D flat-slot** ``[num_slots, num_kv_heads, head_dim]``: used - by focused unit tests that allocate the cache as a contiguous - flat-slot tensor. ``index_copy_(0, ...)`` writes propagate - because the tensor IS the storage. - * **4-D multi-dim paged** ``[num_pages, tokens_per_block, - num_kv_heads, head_dim]``: used when the cache is a view of - ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``. The view is - non-contiguous (its dim-0 stride is 2× the contiguous stride - because dim 1 separates K from V in the pool), so - ``index_copy_(0, ...)`` would silently fork a copy and the - write would be lost. Decompose the flat slot id into - ``(page, within)`` and use multi-dim fancy assignment so the - write propagates through the view to the underlying pool. + """Write per-new-token K or V into a cache view via the shared writer. + + Delegates to :func:`common.write_kv_slots`, which handles both the 3-D + flat-slot layout used by focused unit tests and the 4-D paged view of + ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``. """ - # KV-cache writes never need to participate in autograd. Wrap both - # branches in ``torch.no_grad()`` so callers that enter this path - # with an active grad context (e.g. unit tests exercising - # :class:`MiniMaxM3Attention` without ``inference_mode``) do not - # trip the autograd in-place guard on the cache view chain. - with torch.no_grad(): - if cache.ndim >= 4: - tokens_per_block = int(cache.shape[1]) - out_long = out_cache_loc.to(torch.long) - page = out_long // tokens_per_block - within = out_long % tokens_per_block - _assert_paged_write_in_bounds("cache", cache, page, within) - cache[page, within] = values.to(cache.dtype) - else: - cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + write_kv_slots(cache, out_cache_loc, values) def _scatter_topk_to_block_mask( @@ -699,7 +598,7 @@ def minimax_m3_sparse_decode( v_cache: torch.Tensor, idx_k_cache: torch.Tensor, idx_v_cache: Optional[torch.Tensor], - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, config: MiniMaxM3SparseConfig, *, disable_index_value: bool, @@ -794,7 +693,7 @@ def minimax_m3_sparse_prefill( idx_q: torch.Tensor, idx_k_cache: torch.Tensor, idx_v_cache: Optional[torch.Tensor], - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, config: MiniMaxM3SparseConfig, *, disable_index_value: bool, @@ -890,18 +789,8 @@ def minimax_m3_sparse_prefill( # --------------------------------------------------------------------------- -# Lazy import alias to avoid a circular import at module load — the side -# cache class lives in ``cache_manager`` which imports from this module -# is fine (no cycle), but keeping the indirection explicit makes the -# dependency direction in this file clearer. -def _import_index_cache_cls(): - from .cache_manager import MiniMaxM3SparseIndexCache - - return MiniMaxM3SparseIndexCache - - @dataclass -class MiniMaxM3SparseAttention: +class MiniMaxM3TritonSparseAttention: """Thin orchestrator for :func:`minimax_m3_sparse_prefill` and :func:`minimax_m3_sparse_decode`. @@ -909,7 +798,7 @@ class MiniMaxM3SparseAttention: :class:`MiniMaxM3SparseIndexCache`. The caller is responsible for routing the projected Q, K, V, ``idx_q``, ``idx_k`` (and optional ``idx_v``) tensors plus the populated - :class:`MiniMaxM3SparseAttentionMetadata`. + :class:`MiniMaxM3TritonSparseAttentionMetadata`. """ config: MiniMaxM3SparseConfig @@ -946,7 +835,7 @@ def forward( idx_q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - metadata: MiniMaxM3SparseAttentionMetadata, + metadata: MiniMaxM3TritonSparseAttentionMetadata, *, disable_index_value: bool, sm_scale: Optional[float] = None, @@ -1001,322 +890,304 @@ def forward( # --------------------------------------------------------------------------- -@functools.lru_cache(maxsize=1) -def get_minimax_m3_attention_backend_cls(): - """Return :class:`MiniMaxM3SparseRuntimeBackend` (lazy import). +class MiniMaxM3SparseRuntimeBackend(AttentionBackend[AttentionMetadata]): + """:class:`AttentionBackend` for MiniMax-M3 sparse layers. - Deferring the :class:`AttentionBackend` import keeps the algorithm - module usable from test paths that do not need the runtime backend. - """ - from ...interface import ( - AttentionBackend, - AttentionForwardArgs, - AttentionMetadata, - merge_attention_forward_args, - ) + Constructed under the standard ``create_attention(...)`` dispatch + when ``SparseAttentionConfig(algorithm='minimax_m3', ...)`` is + configured. Drives the MiniMax-M3 sparse algorithm directly via + :func:`minimax_m3_sparse_prefill` and + :func:`minimax_m3_sparse_decode`. - metadata_cls = get_minimax_m3_attention_metadata_cls() - - class MiniMaxM3SparseRuntimeBackend(AttentionBackend[AttentionMetadata]): - """:class:`AttentionBackend` for MiniMax-M3 sparse layers. - - Constructed under the standard ``create_attention(...)`` dispatch - when ``SparseAttentionConfig(algorithm='minimax_m3', ...)`` is - configured. Drives the MiniMax-M3 sparse algorithm directly via - :func:`minimax_m3_sparse_prefill` and - :func:`minimax_m3_sparse_decode`. - - The standard :class:`AttentionForwardArgs` surface does not - carry ``idx_q`` / ``idx_k`` slots, so the model layer threads - the index branch through ``**kwargs`` of :meth:`forward`. When - ``forward`` is called without ``idx_q`` it raises - :class:`NotImplementedError` with a pointer at the model layer - — the backend's ``forward`` is **executable**, but it is not a - substitute for the MiniMax-specific projection / norm / RoPE - steps the model layer is responsible for. - - The backend exposes - :meth:`forward_sparse` for callers that want a name-explicit - entry point (the model layer calls it directly); :meth:`forward` - is the standard contract entry point and routes to - :meth:`forward_sparse` when ``idx_q`` is supplied. - """ + The standard :class:`AttentionForwardArgs` surface does not + carry ``idx_q`` / ``idx_k`` slots, so the model layer threads + the index branch through ``**kwargs`` of :meth:`forward`. When + ``forward`` is called without ``idx_q`` it raises + :class:`NotImplementedError` with a pointer at the model layer + — the backend's ``forward`` is **executable**, but it is not a + substitute for the MiniMax-specific projection / norm / RoPE + steps the model layer is responsible for. + + The backend exposes + :meth:`forward_sparse` for callers that want a name-explicit + entry point (the model layer calls it directly); :meth:`forward` + is the standard contract entry point and routes to + :meth:`forward_sparse` when ``idx_q`` is supplied. + """ - Metadata = metadata_cls + Metadata = MiniMaxM3AttentionMetadata - def __init__( - self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config=None, - sparse_params: Optional["MiniMaxM3SparseParams"] = None, + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config=None, + sparse_params: Optional["MiniMaxM3SparseParams"] = None, + **kwargs, + ): + if sparse_params is None: + raise ValueError("sparse_params is required for MiniMaxM3SparseRuntimeBackend") + super().__init__( + layer_idx, + num_heads, + head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + sparse_params=sparse_params, **kwargs, - ): - if sparse_params is None: - raise ValueError("sparse_params is required for MiniMaxM3SparseRuntimeBackend") - super().__init__( - layer_idx, - num_heads, - head_dim, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - sparse_params=sparse_params, - **kwargs, + ) + self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, + num_q_heads=num_heads, + num_kv_heads=num_kv_heads or num_heads, + head_dim=head_dim, + ) + self.disable_index_value = bool(sparse_params.disable_index_value) + + @staticmethod + def support_fused_rope() -> bool: + # The MiniMax-M3 model layer applies RoPE explicitly because + # both the main and the index branches need partial RoPE, + # and the standard fused-RoPE attention op does not have a + # hook for the index branch. + return False + + def forward_sparse( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + idx_v: Optional[torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_k_cache: torch.Tensor, + idx_v_cache: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + m3_metadata: "MiniMaxM3TritonSparseAttentionMetadata", + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Execute the MiniMax-M3 sparse path end-to-end. + + Inputs: + ``q``, ``k``, ``v`` : new-token projections, already + per-head norm + RoPE applied. + ``idx_q``, ``idx_k`` : index-branch projections, already + per-head norm + RoPE applied. + ``idx_v`` : index-V projection (only when + ``disable_index_value=False``). + ``k_cache``, ``v_cache``: flat-slot view of the paged + main K/V cache, + ``[num_slots, num_kv_heads, head_dim]``. + ``idx_k_cache`` : side index-K cache, + ``[num_slots, 1, sparse_index_dim]``. + ``idx_v_cache`` : side index-V cache (or ``None``). + ``out_cache_loc`` : ``[num_new_tokens]`` int slot + indices to write the new + token's K/V/idx_K to. + ``m3_metadata`` : populated + :class:`MiniMaxM3TritonSparseAttentionMetadata`. + ``output`` : optional preallocated final output, + ``[num_tokens, num_q_heads * head_dim]``. + + Returns ``[num_tokens, num_q_heads * head_dim]``. + """ + num_kv_heads = self.m3_config.num_kv_heads + head_dim = self.m3_config.head_dim + sparse_index_dim = self.m3_config.sparse_index_dim + num_idx_heads = self.m3_config.num_index_heads + + num_tokens = int(q.shape[0]) + q_view = q.view(num_tokens, self.num_heads, head_dim) + k_view = k.view(num_tokens, num_kv_heads, head_dim) + v_view = v.view(num_tokens, num_kv_heads, head_dim) + idx_q_view = idx_q.view(num_tokens, num_idx_heads, sparse_index_dim) + idx_k_view = idx_k.view(num_tokens, 1, sparse_index_dim) + + # Production paths build the M3 metadata on the cache + # device in ``MiniMaxM3AttentionMetadata.prepare`` (called + # outside any CUDA-graph capture window), and test paths + # construct it directly on the desired device. So all + # metadata tensors should already live on ``k_cache.device`` + # by this point. We keep a same-device pass for resilience + # against legacy test callers that produce metadata on a + # different device, but it must not introduce CPU->GPU + # copies inside the capture window. ``ensure_metadata_on_device`` + # is a no-op when each tensor is already on the target + # device; under capture that no-op path is the contract. + cache_device = k_cache.device + if any( + t is not None and t.device != cache_device + for t in ( + m3_metadata.req_to_token, + m3_metadata.slot_ids, + m3_metadata.seq_lens, + m3_metadata.prefix_lens, + m3_metadata.cu_seqlens_q, + m3_metadata.q_batch_row, + m3_metadata.q_positions, ) - self.m3_config = MiniMaxM3SparseConfig.from_sparse_params( - sparse_params, - num_q_heads=num_heads, - num_kv_heads=num_kv_heads or num_heads, - head_dim=head_dim, + ): + m3_metadata = ensure_metadata_on_device(m3_metadata, cache_device) + + # Write new K/V/idx_K to the configured slots. + # ``out_cache_loc`` comes from the pre-built attachment so + # it already lives on the cache device. The write goes + # through :func:`_write_main_kv_slots`, which is layout- + # aware: + # + # * 4-D multi-dim paged caches (the production V2 path: + # main K/V is ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``, and + # ``idx_k_cache`` is the V2 4-D paged view ``[num_pages, + # tokens_per_block, 1, sparse_index_dim]``): decomposes + # each slot id into + # ``(page, within)`` and uses multi-dim fancy + # assignment so the write propagates to the underlying + # pool. A plain ``index_copy_(0, ...)`` would either + # silently fork a copy (non-contiguous main K/V view) + # or raise a shape mismatch (4-D index-K view), but + # the layout-aware helper sidesteps both failure + # modes. + # * 3-D flat-slot caches (focused unit tests that + # allocate plain ``torch.zeros((num_slots, num_heads, + # channel))`` tensors): falls back to + # ``index_copy_(0, ...)`` because the tensor IS the + # storage. + _write_main_kv_slots(k_cache, out_cache_loc, k_view) + _write_main_kv_slots(v_cache, out_cache_loc, v_view) + _write_main_kv_slots(idx_k_cache, out_cache_loc, idx_k_view) + if idx_v is not None and idx_v_cache is not None: + idx_v_view = idx_v.view(num_tokens, 1, sparse_index_dim) + _write_main_kv_slots(idx_v_cache, out_cache_loc, idx_v_view) + + if m3_metadata.is_prefill: + _, o = minimax_m3_sparse_prefill( + q_view, + k_cache, + v_cache, + idx_q_view, + idx_k_cache, + None if self.disable_index_value else idx_v_cache, + m3_metadata, + self.m3_config, + disable_index_value=self.disable_index_value, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + output=output, ) - self.disable_index_value = bool(sparse_params.disable_index_value) - - @staticmethod - def support_fused_rope() -> bool: - # The MiniMax-M3 model layer applies RoPE explicitly because - # both the main and the index branches need partial RoPE, - # and the standard fused-RoPE attention op does not have a - # hook for the index branch. - return False - - def forward_sparse( - self, - *, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - idx_q: torch.Tensor, - idx_k: torch.Tensor, - idx_v: Optional[torch.Tensor], - k_cache: torch.Tensor, - v_cache: torch.Tensor, - idx_k_cache: torch.Tensor, - idx_v_cache: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - m3_metadata: "MiniMaxM3SparseAttentionMetadata", - sm_scale: Optional[float] = None, - idx_sm_scale: Optional[float] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Execute the MiniMax-M3 sparse path end-to-end. - - Inputs: - ``q``, ``k``, ``v`` : new-token projections, already - per-head norm + RoPE applied. - ``idx_q``, ``idx_k`` : index-branch projections, already - per-head norm + RoPE applied. - ``idx_v`` : index-V projection (only when - ``disable_index_value=False``). - ``k_cache``, ``v_cache``: flat-slot view of the paged - main K/V cache, - ``[num_slots, num_kv_heads, head_dim]``. - ``idx_k_cache`` : side index-K cache, - ``[num_slots, 1, sparse_index_dim]``. - ``idx_v_cache`` : side index-V cache (or ``None``). - ``out_cache_loc`` : ``[num_new_tokens]`` int slot - indices to write the new - token's K/V/idx_K to. - ``m3_metadata`` : populated - :class:`MiniMaxM3SparseAttentionMetadata`. - ``output`` : optional preallocated final output, - ``[num_tokens, num_q_heads * head_dim]``. - - Returns ``[num_tokens, num_q_heads * head_dim]``. - """ - num_kv_heads = self.m3_config.num_kv_heads - head_dim = self.m3_config.head_dim - sparse_index_dim = self.m3_config.sparse_index_dim - num_idx_heads = self.m3_config.num_index_heads - - num_tokens = int(q.shape[0]) - q_view = q.view(num_tokens, self.num_heads, head_dim) - k_view = k.view(num_tokens, num_kv_heads, head_dim) - v_view = v.view(num_tokens, num_kv_heads, head_dim) - idx_q_view = idx_q.view(num_tokens, num_idx_heads, sparse_index_dim) - idx_k_view = idx_k.view(num_tokens, 1, sparse_index_dim) - - # Production paths build the M3 metadata on the cache - # device in ``MiniMaxM3AttentionMetadata.prepare`` (called - # outside any CUDA-graph capture window), and test paths - # construct it directly on the desired device. So all - # metadata tensors should already live on ``k_cache.device`` - # by this point. We keep a same-device pass for resilience - # against legacy test callers that produce metadata on a - # different device, but it must not introduce CPU->GPU - # copies inside the capture window. ``ensure_metadata_on_device`` - # is a no-op when each tensor is already on the target - # device; under capture that no-op path is the contract. - cache_device = k_cache.device - if any( - t is not None and t.device != cache_device - for t in ( - m3_metadata.req_to_token, - m3_metadata.slot_ids, - m3_metadata.seq_lens, - m3_metadata.prefix_lens, - m3_metadata.cu_seqlens_q, - m3_metadata.q_batch_row, - m3_metadata.q_positions, - ) - ): - m3_metadata = ensure_metadata_on_device(m3_metadata, cache_device) - - # Write new K/V/idx_K to the configured slots. - # ``out_cache_loc`` comes from the pre-built attachment so - # it already lives on the cache device. The write goes - # through :func:`_write_main_kv_slots`, which is layout- - # aware: - # - # * 4-D multi-dim paged caches (the production V2 path: - # main K/V is ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``, and - # ``idx_k_cache`` is the V2 4-D paged view ``[num_pages, - # tokens_per_block, 1, sparse_index_dim]``): decomposes - # each slot id into - # ``(page, within)`` and uses multi-dim fancy - # assignment so the write propagates to the underlying - # pool. A plain ``index_copy_(0, ...)`` would either - # silently fork a copy (non-contiguous main K/V view) - # or raise a shape mismatch (4-D index-K view), but - # the layout-aware helper sidesteps both failure - # modes. - # * 3-D flat-slot caches (focused unit tests that - # allocate plain ``torch.zeros((num_slots, num_heads, - # channel))`` tensors): falls back to - # ``index_copy_(0, ...)`` because the tensor IS the - # storage. - _write_main_kv_slots(k_cache, out_cache_loc, k_view) - _write_main_kv_slots(v_cache, out_cache_loc, v_view) - _write_main_kv_slots(idx_k_cache, out_cache_loc, idx_k_view) - if idx_v is not None and idx_v_cache is not None: - idx_v_view = idx_v.view(num_tokens, 1, sparse_index_dim) - _write_main_kv_slots(idx_v_cache, out_cache_loc, idx_v_view) - - if m3_metadata.is_prefill: - _, o = minimax_m3_sparse_prefill( - q_view, - k_cache, - v_cache, - idx_q_view, - idx_k_cache, - None if self.disable_index_value else idx_v_cache, - m3_metadata, - self.m3_config, - disable_index_value=self.disable_index_value, - sm_scale=sm_scale, - idx_sm_scale=idx_sm_scale, - output=output, - ) - else: - _, o = minimax_m3_sparse_decode( - q_view, - idx_q_view, - k_cache, - v_cache, - idx_k_cache, - None if self.disable_index_value else idx_v_cache, - m3_metadata, - self.m3_config, - disable_index_value=self.disable_index_value, - sm_scale=sm_scale, - idx_sm_scale=idx_sm_scale, - output=output, - ) - return o - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata=None, - forward_args: Optional[AttentionForwardArgs] = None, - *, - output: Optional[torch.Tensor] = None, - idx_q: Optional[torch.Tensor] = None, - idx_k: Optional[torch.Tensor] = None, - idx_v: Optional[torch.Tensor] = None, - k_cache: Optional[torch.Tensor] = None, - v_cache: Optional[torch.Tensor] = None, - idx_k_cache: Optional[torch.Tensor] = None, - idx_v_cache: Optional[torch.Tensor] = None, - out_cache_loc: Optional[torch.Tensor] = None, - m3_metadata: Optional["MiniMaxM3SparseAttentionMetadata"] = None, - sm_scale: Optional[float] = None, - idx_sm_scale: Optional[float] = None, - **kwargs, - ) -> torch.Tensor: - """Standard ``AttentionBackend.forward`` entry point. - - The MiniMax-M3 sparse path needs the index branch projection - and the M3-shaped metadata; both arrive through keyword - arguments because the standard - :class:`AttentionForwardArgs` surface does not carry them. - - When ``idx_q`` is omitted, this method raises - :class:`NotImplementedError` to make the misuse loud — a - generic AttentionBackend dispatch site cannot drive this - backend without supplying the index branch. - """ - forward_args = merge_attention_forward_args(forward_args, kwargs) - if ( - output is not None - and forward_args.output is not None - and output is not forward_args.output - ): - raise ValueError("output was supplied both directly and through forward_args") - if output is None: - output = forward_args.output - if idx_q is None or idx_k is None or m3_metadata is None: - raise NotImplementedError( - f"MiniMaxM3SparseRuntimeBackend.forward (layer " - f"{self.layer_idx}) requires the M3 index branch and " - "metadata to be passed as keyword arguments " - "(`idx_q`, `idx_k`, `m3_metadata`, " - "`out_cache_loc`, `k_cache`, `v_cache`, " - "`idx_k_cache`). The standard AttentionForwardArgs " - "surface does not carry them; the model layer " - "(`MiniMaxM3Attention.forward`) supplies them when " - "calling this backend." - ) - if ( - k is None - or v is None - or k_cache is None - or v_cache is None - or idx_k_cache is None - or out_cache_loc is None - ): - raise ValueError( - "MiniMaxM3SparseRuntimeBackend.forward requires k, v, " - "k_cache, v_cache, idx_k_cache, and out_cache_loc to " - "be supplied alongside idx_q / idx_k / m3_metadata." - ) - return self.forward_sparse( - q=q, - k=k, - v=v, - idx_q=idx_q, - idx_k=idx_k, - idx_v=idx_v, - k_cache=k_cache, - v_cache=v_cache, - idx_k_cache=idx_k_cache, - idx_v_cache=idx_v_cache, - out_cache_loc=out_cache_loc, - m3_metadata=m3_metadata, + else: + _, o = minimax_m3_sparse_decode( + q_view, + idx_q_view, + k_cache, + v_cache, + idx_k_cache, + None if self.disable_index_value else idx_v_cache, + m3_metadata, + self.m3_config, + disable_index_value=self.disable_index_value, sm_scale=sm_scale, idx_sm_scale=idx_sm_scale, output=output, ) + return o - return MiniMaxM3SparseRuntimeBackend + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata=None, + forward_args: Optional[AttentionForwardArgs] = None, + *, + output: Optional[torch.Tensor] = None, + idx_q: Optional[torch.Tensor] = None, + idx_k: Optional[torch.Tensor] = None, + idx_v: Optional[torch.Tensor] = None, + k_cache: Optional[torch.Tensor] = None, + v_cache: Optional[torch.Tensor] = None, + idx_k_cache: Optional[torch.Tensor] = None, + idx_v_cache: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + m3_metadata: Optional["MiniMaxM3TritonSparseAttentionMetadata"] = None, + sm_scale: Optional[float] = None, + idx_sm_scale: Optional[float] = None, + **kwargs, + ) -> torch.Tensor: + """Standard ``AttentionBackend.forward`` entry point. + + The MiniMax-M3 sparse path needs the index branch projection + and the M3-shaped metadata; both arrive through keyword + arguments because the standard + :class:`AttentionForwardArgs` surface does not carry them. + + When ``idx_q`` is omitted, this method raises + :class:`NotImplementedError` to make the misuse loud — a + generic AttentionBackend dispatch site cannot drive this + backend without supplying the index branch. + """ + forward_args = merge_attention_forward_args(forward_args, kwargs) + if ( + output is not None + and forward_args.output is not None + and output is not forward_args.output + ): + raise ValueError("output was supplied both directly and through forward_args") + if output is None: + output = forward_args.output + if idx_q is None or idx_k is None or m3_metadata is None: + raise NotImplementedError( + f"MiniMaxM3SparseRuntimeBackend.forward (layer " + f"{self.layer_idx}) requires the M3 index branch and " + "metadata to be passed as keyword arguments " + "(`idx_q`, `idx_k`, `m3_metadata`, " + "`out_cache_loc`, `k_cache`, `v_cache`, " + "`idx_k_cache`). The standard AttentionForwardArgs " + "surface does not carry them; the model layer " + "(`MiniMaxM3Attention.forward`) supplies them when " + "calling this backend." + ) + if ( + k is None + or v is None + or k_cache is None + or v_cache is None + or idx_k_cache is None + or out_cache_loc is None + ): + raise ValueError( + "MiniMaxM3SparseRuntimeBackend.forward requires k, v, " + "k_cache, v_cache, idx_k_cache, and out_cache_loc to " + "be supplied alongside idx_q / idx_k / m3_metadata." + ) + return self.forward_sparse( + q=q, + k=k, + v=v, + idx_q=idx_q, + idx_k=idx_k, + idx_v=idx_v, + k_cache=k_cache, + v_cache=v_cache, + idx_k_cache=idx_k_cache, + idx_v_cache=idx_v_cache, + out_cache_loc=out_cache_loc, + m3_metadata=m3_metadata, + sm_scale=sm_scale, + idx_sm_scale=idx_sm_scale, + output=output, + ) __all__ = [ - "MiniMaxM3SparseAttention", - "get_minimax_m3_attention_backend_cls", + "MiniMaxM3SparseRuntimeBackend", + "MiniMaxM3TritonSparseAttention", "minimax_m3_sparse_decode", "minimax_m3_sparse_prefill", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/kernels.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py similarity index 100% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/kernels.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py similarity index 53% rename from tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py rename to tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py index 925ca6c0189b..a57cf5eac2e1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py @@ -1,131 +1,31 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""MiniMax-M3 sparse attention configuration + per-forward metadata. +"""MiniMax-M3 Triton reference per-forward metadata. Contains: - * :class:`MiniMaxM3SparseConfig` -- post-TP-shard kernel - parameter bundle. - * :class:`MiniMaxM3SparseAttentionMetadata` -- per-forward metadata - dataclass with a - CUDA-graph-safe - :meth:`prepare`. + * :class:`MiniMaxM3TritonSparseAttentionMetadata` -- per-forward metadata + dataclass with a CUDA-graph-safe :meth:`prepare`. * Helpers to migrate metadata across devices, build it from a real - :class:`KVCacheManagerV2`, and pre-allocate CUDA-graph-stable - buffers. - * :func:`get_minimax_m3_attention_metadata_cls` -- lazy factory for - the :class:`AttentionMetadata` subclass the pyexecutor wires into - the M3 sparse layer's forward path. + :class:`KVCacheManagerV2`, and pre-allocate CUDA-graph-stable buffers. + * :class:`MiniMaxM3AttentionMetadata` -- the :class:`AttentionMetadata` + subclass the pyexecutor wires into the M3 sparse layer's forward path. """ from __future__ import annotations import dataclasses -import functools from dataclasses import dataclass, field -from typing import List, Literal, Optional, Tuple +from typing import List, Optional, Tuple import torch -from ..params import SparseParams - - -@dataclass(frozen=True) -class MiniMaxM3SparseParams(SparseParams): - """Lowered runtime parameters for the MiniMax-M3 sparse backend.""" - - algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") - num_index_heads: int = 4 - sparse_index_dim: int = 128 - block_size: int = 128 - topk: int = 16 - init_blocks: int = 0 - local_blocks: int = 1 - score_type: str = "max" - disable_index_value: bool = True - - -@dataclass(frozen=True) -class MiniMaxM3SparseConfig: - """Per-rank kernel parameter bundle for MiniMax-M3 sparse attention. - - This is **not** a user-facing config (use - :class:`tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig` - for that). It is the layer-invariant, post-TP-shard parameter bundle - that backend kernels and reference helpers consume. The user knobs - come from :class:`MiniMaxM3SparseParams`; ``num_q_heads`` / - ``num_kv_heads`` / ``head_dim`` come from the per-rank model - geometry and must be supplied by the caller (typically via - :meth:`from_sparse_params`). - """ - - num_q_heads: int - num_kv_heads: int - head_dim: int - num_index_heads: int - sparse_index_dim: int - block_size: int - topk: int - init_blocks: int = 0 - local_blocks: int = 1 - score_type: str = "max" - - def __post_init__(self) -> None: - if self.num_q_heads % self.num_kv_heads != 0: - raise ValueError( - f"num_q_heads ({self.num_q_heads}) must be divisible by " - f"num_kv_heads ({self.num_kv_heads})" - ) - if self.num_index_heads % self.num_kv_heads != 0: - raise ValueError( - f"num_index_heads ({self.num_index_heads}) must be divisible " - f"by num_kv_heads ({self.num_kv_heads})" - ) - if self.block_size <= 0: - raise ValueError(f"block_size must be > 0, got {self.block_size}") - if self.topk <= 0: - raise ValueError(f"topk must be > 0, got {self.topk}") - if self.init_blocks < 0: - raise ValueError(f"init_blocks must be >= 0, got {self.init_blocks}") - if self.local_blocks < 0: - raise ValueError(f"local_blocks must be >= 0, got {self.local_blocks}") - if self.score_type != "max": - # SGLang exposes only "max" today and that is what the MiniMax-M3 - # checkpoint config specifies. Reject anything else explicitly so - # a config drift surfaces immediately. - raise ValueError( - f"score_type={self.score_type!r} is not supported " - "(only 'max' matches the SGLang reference)" - ) - - @classmethod - def from_sparse_params( - cls, - sparse_params: "MiniMaxM3SparseParams", - *, - num_q_heads: int, - num_kv_heads: int, - head_dim: int, - ) -> "MiniMaxM3SparseConfig": - """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` - and the per-rank model geometry. - """ - return cls( - num_q_heads=int(num_q_heads), - num_kv_heads=int(num_kv_heads), - head_dim=int(head_dim), - num_index_heads=int(sparse_params.num_index_heads), - sparse_index_dim=int(sparse_params.sparse_index_dim), - block_size=int(sparse_params.block_size), - topk=int(sparse_params.topk), - init_blocks=int(sparse_params.init_blocks), - local_blocks=int(sparse_params.local_blocks), - score_type=str(sparse_params.score_type), - ) +from ...trtllm import TrtllmAttentionMetadata +from .common import build_paged_kv_slot_mapping @dataclass -class MiniMaxM3SparseAttentionMetadata: +class MiniMaxM3TritonSparseAttentionMetadata: """Per-forward metadata for MiniMax-M3 sparse attention. Mirrors the shape of SGLang's @@ -193,6 +93,10 @@ class MiniMaxM3SparseAttentionMetadata: prefix_lens: Optional[torch.Tensor] = None cu_seqlens_q: Optional[torch.Tensor] = None extend_seq_lens_cpu: Optional[List[int]] = None + # Query tokens per request on decode-shaped metadata: 1 normally, + # 1 + draft_len under one-model Eagle3 verify. Consumed by the dense + # SDPA decode ladder mask. + decode_qo_len: int = 1 q_batch_row: Optional[torch.Tensor] = None q_positions: Optional[torch.Tensor] = None max_seqlen_q: int = field(default=1) @@ -265,16 +169,20 @@ def prepare(self) -> None: if batch_size == 0: self.max_seqlen_k = 1 else: - self.max_seqlen_k = int(self.seq_lens_cpu[:batch_size].max().item()) + max_k = int(self.seq_lens_cpu[:batch_size].max().item()) + # Optimistic overlap+spec lengths can overhang the page table at + # a page boundary; SDPA consumes max_seqlen_k as the exact + # mask/gather width, so clamp to the table. + self.max_seqlen_k = min(max_k, int(self.req_to_token.shape[1])) def ensure_metadata_on_device( - metadata: "MiniMaxM3SparseAttentionMetadata", + metadata: "MiniMaxM3TritonSparseAttentionMetadata", device: torch.device, -) -> "MiniMaxM3SparseAttentionMetadata": +) -> "MiniMaxM3TritonSparseAttentionMetadata": """Return ``metadata`` with every GPU-consumed tensor on ``device``. - Constructs a new :class:`MiniMaxM3SparseAttentionMetadata` whose + Constructs a new :class:`MiniMaxM3TritonSparseAttentionMetadata` whose tensor fields are migrated to ``device`` when they live elsewhere. The CPU-side mirror ``seq_lens_cpu`` is preserved because the algorithm reads it only when explicitly noted (e.g. for @@ -306,18 +214,6 @@ def _move(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: ) -def replace_metadata( - metadata: MiniMaxM3SparseAttentionMetadata, - **changes, -) -> MiniMaxM3SparseAttentionMetadata: - """Helper around :func:`dataclasses.replace` for ``metadata``. - - Provided so callers can build a decode metadata from a prefill - metadata without manually re-typing every field. - """ - return dataclasses.replace(metadata, **changes) - - def allocate_minimax_m3_static_buffers( *, max_num_sequences: int, @@ -405,6 +301,120 @@ def allocate_minimax_m3_static_buffers( } +def _build_runtime_metadata_fresh( + *, + kv_cache_manager, + request_ids, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + is_prefill: bool, + prefix_lens: Optional[torch.Tensor], + extend_seq_lens_cpu: Optional[List[int]], + device: torch.device, +) -> Tuple[MiniMaxM3TritonSparseAttentionMetadata, torch.Tensor]: + """Fresh-allocation build of the Triton reference metadata. + + Delegates the backend-neutral req_to_token, slot_ids and out_cache_loc + derivation to common.build_paged_kv_slot_mapping and adds the Triton-only + fields: cu_seqlens_q, prefix_lens on device, and the scalar max_seqlen + values and per-query-token tensors that prepare() computes. Used when no + graph-stable static_buffers are supplied; the static_buffers path in + build_runtime_metadata_from_kv_manager keeps its own in-place buffer writes. + """ + batch = int(seq_lens.shape[0]) + seq_lens_dev = seq_lens.to(device) if seq_lens.device != device else seq_lens + + if is_prefill: + if extend_seq_lens_cpu is None: + raise ValueError("prefill metadata requires extend_seq_lens_cpu") + if prefix_lens is None: + raise ValueError("prefill metadata requires prefix_lens") + prefix_lens_dev = prefix_lens.to(device) if prefix_lens.device != device else prefix_lens + qo_lens_cpu = torch.tensor([int(x) for x in extend_seq_lens_cpu], dtype=torch.int32) + qo_offset_cpu = prefix_lens.detach().to(device="cpu", dtype=torch.int32) + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=device, + ) + cu_q: List[int] = [0] + for ext in extend_seq_lens_cpu: + cu_q.append(cu_q[-1] + int(ext)) + cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) + meta = MiniMaxM3TritonSparseAttentionMetadata( + is_prefill=True, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_dev, + seq_lens_cpu=seq_lens_cpu, + prefix_lens=prefix_lens_dev, + cu_seqlens_q=cu_seqlens_q, + extend_seq_lens_cpu=list(extend_seq_lens_cpu), + q_batch_row=None, + q_positions=None, + ) + else: + # Decode: the new token sits at position seq_lens[b] - 1. + qo_lens_cpu = torch.ones(batch, dtype=torch.int32) + qo_offset_cpu = seq_lens_cpu.detach().to(device="cpu", dtype=torch.int32) - 1 + req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=device, + ) + meta = MiniMaxM3TritonSparseAttentionMetadata( + is_prefill=False, + req_to_token=req_to_token, + slot_ids=slot_ids, + seq_lens=seq_lens_dev, + seq_lens_cpu=seq_lens_cpu, + ) + meta.prepare() + return meta, out_cache_loc + + +def derive_q_positions_and_cache_slots( + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + q_batch_row: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-Q-token K-side positions and KV slot ids, on device, sync-free. + + Shared by the metadata builder and the ``on_update_kv_lens`` + re-derivation so the two cannot drift. + """ + total_q = int(q_batch_row.shape[0]) + qbr = q_batch_row.to(torch.long) + tok = torch.arange(total_q, dtype=torch.int32, device=q_batch_row.device) + q_positions = prefix_lens[qbr] + (tok - cu_seqlens_q[qbr]) + # Optimistic prefix_lens can overhang the last allocated page; the + # overhanging slots are placeholders (on_update_kv_lens re-derives them + # before any forward reads them) but the gather must stay in bounds. + # Clamp only the table index — non-inplace, since ``.to`` aliases int64 + # inputs. + idx = q_positions.to(torch.long).clamp(min=0, max=req_to_token.shape[1] - 1) + flat = qbr * req_to_token.shape[1] + idx + return q_positions, req_to_token.reshape(-1).index_select(0, flat) + + +def derive_decode_cache_slots(req_to_token: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor: + """Decode-row KV slot ids (new token at position ``seq_lens[b] - 1``). + + Same in-bounds-placeholder clamp contract as + :func:`derive_q_positions_and_cache_slots`; the ``min=0`` floor + additionally covers zero-length dummy rows indexing ``-1``. + """ + rows = torch.arange(seq_lens.shape[0], device=seq_lens.device, dtype=torch.long) + idx = (seq_lens.to(torch.long) - 1).clamp_(min=0, max=req_to_token.shape[1] - 1) + flat = rows * req_to_token.shape[1] + idx + return req_to_token.reshape(-1).index_select(0, flat) + + def build_runtime_metadata_from_kv_manager( *, kv_cache_manager, @@ -416,8 +426,8 @@ def build_runtime_metadata_from_kv_manager( extend_seq_lens_cpu: Optional[List[int]] = None, device: Optional[torch.device] = None, static_buffers: Optional[dict] = None, -) -> Tuple[MiniMaxM3SparseAttentionMetadata, torch.Tensor]: - """Build a :class:`MiniMaxM3SparseAttentionMetadata` from a real +) -> Tuple[MiniMaxM3TritonSparseAttentionMetadata, torch.Tensor]: + """Build a :class:`MiniMaxM3TritonSparseAttentionMetadata` from a real :class:`MiniMaxM3KVCacheManagerV2`. Returns the populated metadata plus an ``out_cache_loc`` tensor @@ -452,6 +462,21 @@ def build_runtime_metadata_from_kv_manager( the end-to-end runtime path without going through the full LLM forward. """ + if device is None: + device = seq_lens.device + + if static_buffers is None: + return _build_runtime_metadata_fresh( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + is_prefill=is_prefill, + prefix_lens=prefix_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + device=device, + ) + tokens_per_block = int(kv_cache_manager.tokens_per_block) # block_ids_per_seq is a [batch_size, max_blocks_per_seq] tensor; row b # holds the block ids assigned to request_ids[b] in order. @@ -459,8 +484,6 @@ def build_runtime_metadata_from_kv_manager( batch = int(seq_lens.shape[0]) max_blocks = int(block_ids.shape[1]) max_kv_len = max_blocks * tokens_per_block - if device is None: - device = seq_lens.device if static_buffers is not None: if static_buffers.get("device") != device: @@ -562,21 +585,13 @@ def build_runtime_metadata_from_kv_manager( req_to_token = req_to_token_fresh slot_ids = torch.arange(batch, device=device, dtype=torch.int32) - # Compute out_cache_loc: per-new-token slot ids, in flattened order - # matching the q-token order the model layer projects. The Python - # loops below run on CPU lists derived from the CPU-resident - # ``seq_lens_cpu`` / ``prefix_lens`` / ``extend_seq_lens_cpu``, so - # no GPU sync is needed at this point. The resulting - # ``out_cache_loc`` tensor is constructed directly on ``device``. - # The ``int(...item())`` reads against ``req_to_token`` are a CPU - # sync but only ever run from ``prepare()`` (outside any CUDA-graph - # capture window) — they are not in the forward path. + # out_cache_loc must be flattened in the q-token order the model layer + # projects, or K/V lands in the wrong requests' slots. if is_prefill: if extend_seq_lens_cpu is None: raise ValueError("prefill metadata requires extend_seq_lens_cpu") if prefix_lens is None: raise ValueError("prefill metadata requires prefix_lens") - prefix_lens_cpu = prefix_lens.to("cpu").tolist() if static_buffers is not None: prefix_buf = static_buffers["prefix_lens"] prefix_src = prefix_lens.to(device=device, dtype=torch.int32, non_blocking=True) @@ -586,17 +601,18 @@ def build_runtime_metadata_from_kv_manager( prefix_lens_dev = ( prefix_lens.to(device) if prefix_lens.device != device else prefix_lens ) - out_cache_loc_list: List[int] = [] cu_q: List[int] = [0] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pref = int(prefix_lens_cpu[b]) - ext = int(extend_seq_lens_cpu[b]) - for offset in range(ext): - slot = int(req_to_token_cpu[b, pref + offset].item()) - out_cache_loc_list.append(slot) - cu_q.append(cu_q[-1] + ext) + for ext in extend_seq_lens_cpu: + cu_q.append(cu_q[-1] + int(ext)) total_q = cu_q[-1] + cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) + q_batch_row_src = torch.repeat_interleave( + torch.arange(batch, device=device, dtype=torch.int32), + torch.tensor(extend_seq_lens_cpu, dtype=torch.int64, device=device), + ) + q_positions_src, out_cache_loc_src = derive_q_positions_and_cache_slots( + req_to_token, prefix_lens_dev, cu_seqlens_q_src, q_batch_row_src + ) if static_buffers is not None: if total_q > static_buffers["max_num_tokens"]: raise ValueError( @@ -604,34 +620,23 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current total_q={total_q}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:total_q].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:total_q] cu_seqlens_q_buf = static_buffers["cu_seqlens_q"] - cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device) cu_seqlens_q_buf[: batch + 1].copy_(cu_seqlens_q_src, non_blocking=True) cu_seqlens_q = cu_seqlens_q_buf[: batch + 1] - # Populate persistent q_batch_row / q_positions in-place so - # the inner metadata's prepare() can leave them alone. q_batch_row_buf = static_buffers["q_batch_row"] q_positions_buf = static_buffers["q_positions"] - for b in range(batch): - start, end = cu_q[b], cu_q[b + 1] - if end > start: - q_batch_row_buf[start:end] = b - pref = int(prefix_lens_cpu[b]) - offsets = ( - torch.arange(start, end, device=device, dtype=torch.int32) - start + pref - ) - q_positions_buf[start:end].copy_(offsets, non_blocking=True) + q_batch_row_buf[:total_q].copy_(q_batch_row_src, non_blocking=True) + q_positions_buf[:total_q].copy_(q_positions_src, non_blocking=True) q_batch_row = q_batch_row_buf[:total_q] q_positions = q_positions_buf[:total_q] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) - q_batch_row = None - q_positions = None - meta = MiniMaxM3SparseAttentionMetadata( + out_cache_loc = out_cache_loc_src + cu_seqlens_q = cu_seqlens_q_src + q_batch_row = q_batch_row_src + q_positions = q_positions_src + meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=True, req_to_token=req_to_token, slot_ids=slot_ids, @@ -645,12 +650,7 @@ def build_runtime_metadata_from_kv_manager( ) else: # Decode: the new token sits at position seq_lens[b] - 1. - seq_lens_cpu_list = seq_lens_cpu.to("cpu").tolist() - out_cache_loc_list = [] - req_to_token_cpu = req_to_token_fresh.to("cpu") - for b in range(batch): - pos = int(seq_lens_cpu_list[b]) - 1 - out_cache_loc_list.append(int(req_to_token_cpu[b, pos].item())) + out_cache_loc_src = derive_decode_cache_slots(req_to_token, seq_lens_dev) if static_buffers is not None: if batch > static_buffers["max_num_tokens"]: raise ValueError( @@ -658,12 +658,11 @@ def build_runtime_metadata_from_kv_manager( f"is smaller than current batch={batch}" ) out_cache_loc_buf = static_buffers["out_cache_loc"] - out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) out_cache_loc_buf[:batch].copy_(out_cache_loc_src, non_blocking=True) out_cache_loc = out_cache_loc_buf[:batch] else: - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - meta = MiniMaxM3SparseAttentionMetadata( + out_cache_loc = out_cache_loc_src + meta = MiniMaxM3TritonSparseAttentionMetadata( is_prefill=False, req_to_token=req_to_token, slot_ids=slot_ids, @@ -674,242 +673,270 @@ def build_runtime_metadata_from_kv_manager( return meta, out_cache_loc -@functools.lru_cache(maxsize=1) -def get_minimax_m3_attention_metadata_cls(): - """Return :class:`MiniMaxM3AttentionMetadata` (lazy import). +class MiniMaxM3AttentionMetadata(TrtllmAttentionMetadata): + """:class:`TrtllmAttentionMetadata` that pre-builds MiniMax-M3 metadata. + + Subclasses :class:`TrtllmAttentionMetadata` (precedent: + ``DSAtrtllmAttentionMetadata``): one-model Eagle3 draft layers run + :class:`TrtllmAttention` against this shared per-step metadata, and the + engine gates spec-dec plumbing on the TRTLLM ``isinstance``. + + Overrides :meth:`prepare` so the M3-sparse + :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token + ``out_cache_loc`` are built **once per scheduler step**, on the + cache device, before the model forward runs. The result is + stored as ``self.minimax_m3 = {"metadata": m3_meta, + "out_cache_loc": out_cache_loc}`` so the model layer's + ``_dense_forward`` and ``_sparse_forward`` can read it without + any device migration. - The class extends :class:`AttentionMetadata` so the pyexecutor's - metadata-creation/prepare hooks (model_engine.py) drive M3 metadata - construction outside the CUDA-graph capture window. Building the - M3-sparse ``req_to_token`` / ``slot_ids`` / ``out_cache_loc`` - tensors during ``prepare()`` lands them on the GPU **before** the - forward call; the forward path then reads from the pre-built - attachment and performs no CPU->GPU copies, which is required for - CUDA-graph capture safety (``cudaErrorStreamCaptureUnsupported`` - fires for CPU->GPU ``memcpyAsync`` calls inside a captured stream). + Test paths that build their own metadata can short-circuit by + attaching ``attn_metadata.minimax_m3`` directly before calling + the forward; those paths do not go through :meth:`prepare`. """ - from ...interface import AttentionMetadata - - class MiniMaxM3AttentionMetadata(AttentionMetadata): - """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata. - - Overrides :meth:`prepare` so the M3-sparse - :class:`MiniMaxM3SparseAttentionMetadata` and the per-new-token - ``out_cache_loc`` are built **once per scheduler step**, on the - cache device, before the model forward runs. The result is - stored as ``self.minimax_m3 = {"metadata": m3_meta, - "out_cache_loc": out_cache_loc}`` so the model layer's - ``_dense_forward`` and ``_sparse_forward`` can read it without - any device migration. - - Test paths that build their own metadata can short-circuit by - attaching ``attn_metadata.minimax_m3`` directly before calling - the forward; those paths do not go through :meth:`prepare`. + + minimax_m3: Optional[dict] = None + # Lazily allocated dict of persistent device buffers used to keep + # ``MiniMaxM3TritonSparseAttentionMetadata`` tensor addresses stable + # across CUDA-graph capture/replay. None until the first + # ``prepare()`` call decides to use them (``is_cuda_graph`` / + # graph-stable mode). + _m3_static_buffers: Optional[dict] = None + + def _maybe_get_m3_static_buffers( + self, cache_device: torch.device, kv_cache_manager + ) -> Optional[dict]: + """Return persistent M3 buffers when graph stability is + required. + + Allocates the persistent buffer dict the first time it is + needed and caches it on ``self._m3_static_buffers``. We + allocate the buffers under two conditions: + + * ``self.is_cuda_graph`` is True -- the captured graph + requires stable ``data_ptr()`` across replays; OR + * the previous prepare() already allocated buffers -- + we keep using them so the algorithm sees the same + addresses even between non-graph and graph-mode calls + (which can happen when the model engine alternates + between eager warmup and graph replay). + + Returns ``None`` when no static buffers should be used (e.g. + eager-only test paths that rely on per-call allocations). """ + need_static = ( + bool(getattr(self, "is_cuda_graph", False)) or self._m3_static_buffers is not None + ) + if not need_static: + return None + if self._m3_static_buffers is not None: + bufs = self._m3_static_buffers + if bufs.get("device") == cache_device: + return bufs + + # First-time use: return an empty placeholder dict. + # ``build_runtime_metadata_from_kv_manager`` performs the + # actual allocation lazily on the first call where the + # current scheduler step's geometry (max_kv_len from the + # manager's block-id table, total_q from extend_seq_lens, + # actual batch size after CUDA-graph padding) is known. + # That removes the need to predict the warmup geometry up + # front. The first allocation pins the buffer addresses + # for the rest of this metadata instance's lifetime, so all + # subsequent prepare() calls reuse the same ``data_ptr()``s + # and CUDA graph capture/replay stays valid. + placeholder: dict = { + "device": cache_device, + # Caller-provided hints used by the lazy allocator below + # when it sizes the persistent buffers on the first real + # prepare() call. + "max_num_sequences_hint": int( + getattr(self, "max_num_sequences", None) or self.max_num_requests + ), + "max_num_tokens_hint": int( + getattr(self, "max_num_tokens", None) + or (int(getattr(self, "max_num_sequences", None) or self.max_num_requests)) + ), + } + self._m3_static_buffers = placeholder + return placeholder - minimax_m3: Optional[dict] = None - # Lazily allocated dict of persistent device buffers used to keep - # ``MiniMaxM3SparseAttentionMetadata`` tensor addresses stable - # across CUDA-graph capture/replay. None until the first - # ``prepare()`` call decides to use them (``is_cuda_graph`` / - # graph-stable mode). - _m3_static_buffers: Optional[dict] = None - - def _maybe_get_m3_static_buffers( - self, cache_device: torch.device, kv_cache_manager - ) -> Optional[dict]: - """Return persistent M3 buffers when graph stability is - required. - - Allocates the persistent buffer dict the first time it is - needed and caches it on ``self._m3_static_buffers``. We - allocate the buffers under two conditions: - - * ``self.is_cuda_graph`` is True -- the captured graph - requires stable ``data_ptr()`` across replays; OR - * the previous prepare() already allocated buffers -- - we keep using them so the algorithm sees the same - addresses even between non-graph and graph-mode calls - (which can happen when the model engine alternates - between eager warmup and graph replay). - - Returns ``None`` when no static buffers should be used (e.g. - eager-only test paths that rely on per-call allocations). - """ - need_static = ( - bool(getattr(self, "is_cuda_graph", False)) or self._m3_static_buffers is not None + def prepare(self) -> None: + super().prepare() + + # Always rebuild the M3 metadata block on each prepare() + # call so it reflects the current scheduler step's seq_lens + # / request_ids / num_cached_tokens. Production + # ``model_engine`` invokes ``prepare()`` outside any CUDA + # graph capture window, so the (potentially expensive) build + # is safe to perform here. + # + # When CUDA graph is enabled the inner ``build_runtime_metadata_from_kv_manager`` + # call writes into the persistent ``_m3_static_buffers`` so + # the captured graph keeps reading from stable ``data_ptr()``s + # across replays. Without this the captured ``index_select`` + # over ``req_to_token``/``slot_ids`` reads from freed warmup + # memory and either produces wrong tokens or fires + # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. + self.minimax_m3 = None + + # Production path: build the M3 metadata from the standard + # AttentionMetadata fields. Requires kv_cache_manager + the + # M3 sparse-cache contract. + kv_cache_manager = getattr(self, "kv_cache_manager", None) + if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): + # Not an M3 KV cache manager: nothing to build. The + # forward path will raise a clear error if the M3 + # backend ends up dispatched without the M3 cache. + return + request_ids = getattr(self, "request_ids", None) + seq_lens = self.seq_lens + if request_ids is None or seq_lens is None: + return + num_contexts = int(getattr(self, "num_contexts", 0) or 0) + batch_size = int(seq_lens.shape[0]) + if batch_size == 0: + return + + # The cache device hosts every paged buffer; this is the + # device the forward path consumes. + try: + layer_buf = kv_cache_manager.get_buffers(0) + cache_device = layer_buf.device + except Exception: + cache_device = torch.device(f"cuda:{torch.cuda.current_device()}") + + seq_lens_cpu = ( + getattr(self, "seq_lens_cpu", None) + if hasattr(self, "seq_lens_cpu") + else seq_lens.detach().to("cpu") + ) + if seq_lens_cpu is None: + seq_lens_cpu = seq_lens.detach().to("cpu") + + kv_cache_params = getattr(self, "kv_cache_params", None) + num_cached_per_seq = ( + kv_cache_params.num_cached_tokens_per_seq + if kv_cache_params is not None + else [0] * batch_size + ) + + # ``attn_metadata.seq_lens`` from the PyExecutor is the + # per-step new-token count. The M3 sparse-attention algorithm + # consumes a *cumulative* kv length: ``minimax_m3_sparse_*`` + # masks reads against ``metadata.seq_lens`` as the per-request + # K-side extent. Compute that cumulative kv length per request + # and feed it into the algorithm metadata builder. + kv_lens_cpu_list = [ + int(num_cached_per_seq[b]) + int(seq_lens_cpu[b].item()) for b in range(batch_size) + ] + kv_lens_cpu = torch.tensor(kv_lens_cpu_list, dtype=torch.int32) + kv_lens_dev = kv_lens_cpu.to(device=cache_device, non_blocking=True) + + static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager) + + # Any batch containing a context (prefill or chunked extend) + # request takes the extend path. For prefill rows + # ``num_cached_per_seq`` is ``prefix_lens`` and the full new + # chunk is ``extend_seq_len``; for decode rows + # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is + # 1, so the same builder produces the correct one-slot + # entry. Pure-decode batches (``num_contexts == 0``) still + # take the decode optimization for CUDA-graph warmup + # geometry. + # + # Mixed prefill+decode batches always take the extend path: + # the prefill kernel handles decode rows as 1-slot extends. + # The decode branch below is a pure-decode-only perf + # specialization. (iter-131 regression: previously a wrong + # predicate routed mixed batches into the decode branch and + # crashed in index_copy_.) + # Multi-token gen rows (spec verify) also route through the extend + # path as prefix+window extends; decode stays one-token-per-row. + is_extend = num_contexts > 0 or int(seq_lens_cpu[:batch_size].max().item()) > 1 + if is_extend: + prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] + extend_seq_lens_cpu = [ + kv_lens_cpu_list[b] - prefix_lens_list[b] for b in range(batch_size) + ] + prefix_lens = torch.tensor( + prefix_lens_list, + dtype=torch.int32, + device=cache_device, ) - if not need_static: - return None - if self._m3_static_buffers is not None: - bufs = self._m3_static_buffers - if bufs.get("device") == cache_device: - return bufs - - # First-time use: return an empty placeholder dict. - # ``build_runtime_metadata_from_kv_manager`` performs the - # actual allocation lazily on the first call where the - # current scheduler step's geometry (max_kv_len from the - # manager's block-id table, total_q from extend_seq_lens, - # actual batch size after CUDA-graph padding) is known. - # That removes the need to predict the warmup geometry up - # front. The first allocation pins the buffer addresses - # for the rest of this metadata instance's lifetime, so all - # subsequent prepare() calls reuse the same ``data_ptr()``s - # and CUDA graph capture/replay stays valid. - placeholder: dict = { - "device": cache_device, - # Caller-provided hints used by the lazy allocator below - # when it sizes the persistent buffers on the first real - # prepare() call. - "max_num_sequences_hint": int( - getattr(self, "max_num_sequences", None) or self.max_num_requests - ), - "max_num_tokens_hint": int( - getattr(self, "max_num_tokens", None) - or (int(getattr(self, "max_num_sequences", None) or self.max_num_requests)) - ), - } - self._m3_static_buffers = placeholder - return placeholder - - def prepare(self) -> None: - super().prepare() - - # Always rebuild the M3 metadata block on each prepare() - # call so it reflects the current scheduler step's seq_lens - # / request_ids / num_cached_tokens. Production - # ``model_engine`` invokes ``prepare()`` outside any CUDA - # graph capture window, so the (potentially expensive) build - # is safe to perform here. - # - # When CUDA graph is enabled the inner ``build_runtime_metadata_from_kv_manager`` - # call writes into the persistent ``_m3_static_buffers`` so - # the captured graph keeps reading from stable ``data_ptr()``s - # across replays. Without this the captured ``index_select`` - # over ``req_to_token``/``slot_ids`` reads from freed warmup - # memory and either produces wrong tokens or fires - # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. - self.minimax_m3 = None - - # Production path: build the M3 metadata from the standard - # AttentionMetadata fields. Requires kv_cache_manager + the - # M3 sparse-cache contract. - kv_cache_manager = getattr(self, "kv_cache_manager", None) - if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): - # Not an M3 KV cache manager: nothing to build. The - # forward path will raise a clear error if the M3 - # backend ends up dispatched without the M3 cache. - return - request_ids = getattr(self, "request_ids", None) - seq_lens = self.seq_lens - if request_ids is None or seq_lens is None: - return - num_contexts = int(getattr(self, "num_contexts", 0) or 0) - batch_size = int(seq_lens.shape[0]) - if batch_size == 0: - return - - # The cache device hosts every paged buffer; this is the - # device the forward path consumes. - try: - layer_buf = kv_cache_manager.get_buffers(0) - cache_device = layer_buf.device - except Exception: - cache_device = torch.device(f"cuda:{torch.cuda.current_device()}") - - seq_lens_cpu = ( - getattr(self, "seq_lens_cpu", None) - if hasattr(self, "seq_lens_cpu") - else seq_lens.detach().to("cpu") + m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=kv_lens_dev, + seq_lens_cpu=kv_lens_cpu, + is_prefill=True, + prefix_lens=prefix_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + device=cache_device, + static_buffers=static_buffers, ) - if seq_lens_cpu is None: - seq_lens_cpu = seq_lens.detach().to("cpu") - - kv_cache_params = getattr(self, "kv_cache_params", None) - num_cached_per_seq = ( - kv_cache_params.num_cached_tokens_per_seq - if kv_cache_params is not None - else [0] * batch_size + else: + m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + seq_lens=kv_lens_dev, + seq_lens_cpu=kv_lens_cpu, + is_prefill=False, + device=cache_device, + static_buffers=static_buffers, ) - # ``attn_metadata.seq_lens`` from the PyExecutor is the - # per-step new-token count. The M3 sparse-attention algorithm - # consumes a *cumulative* kv length: ``minimax_m3_sparse_*`` - # masks reads against ``metadata.seq_lens`` as the per-request - # K-side extent. Compute that cumulative kv length per request - # and feed it into the algorithm metadata builder. - kv_lens_cpu_list = [ - int(num_cached_per_seq[b]) + int(seq_lens_cpu[b].item()) for b in range(batch_size) - ] - kv_lens_cpu = torch.tensor(kv_lens_cpu_list, dtype=torch.int32) - kv_lens_dev = kv_lens_cpu.to(device=cache_device, non_blocking=True) - - static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager) - - # Any batch containing a context (prefill or chunked extend) - # request takes the extend path. For prefill rows - # ``num_cached_per_seq`` is ``prefix_lens`` and the full new - # chunk is ``extend_seq_len``; for decode rows - # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is - # 1, so the same builder produces the correct one-slot - # entry. Pure-decode batches (``num_contexts == 0``) still - # take the decode optimization for CUDA-graph warmup - # geometry. - # - # Mixed prefill+decode batches always take the extend path: - # the prefill kernel handles decode rows as 1-slot extends. - # The decode branch below is a pure-decode-only perf - # specialization. (iter-131 regression: previously a wrong - # predicate routed mixed batches into the decode branch and - # crashed in index_copy_.) - is_extend = num_contexts > 0 - if is_extend: - prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)] - extend_seq_lens_cpu = [ - kv_lens_cpu_list[b] - prefix_lens_list[b] for b in range(batch_size) - ] - prefix_lens = torch.tensor( - prefix_lens_list, - dtype=torch.int32, - device=cache_device, - ) - m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - seq_lens=kv_lens_dev, - seq_lens_cpu=kv_lens_cpu, - is_prefill=True, - prefix_lens=prefix_lens, - extend_seq_lens_cpu=extend_seq_lens_cpu, - device=cache_device, - static_buffers=static_buffers, - ) - else: - m3_meta, out_cache_loc = build_runtime_metadata_from_kv_manager( - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - seq_lens=kv_lens_dev, - seq_lens_cpu=kv_lens_cpu, - is_prefill=False, - device=cache_device, - static_buffers=static_buffers, - ) - - self.minimax_m3 = { - "metadata": m3_meta, - "out_cache_loc": out_cache_loc, - } - - return MiniMaxM3AttentionMetadata + self.minimax_m3 = { + "metadata": m3_meta, + "out_cache_loc": out_cache_loc, + } + + def on_update_kv_lens(self) -> None: + """Re-derive the M3 attachment from the corrected ``kv_lens_cuda``. + + Under the overlap scheduler + speculative decoding, prepare() + runs with optimistic cached counts (full draft acceptance) and + the engine corrects ``kv_lens_cuda`` on device before invoking + this hook (same pattern as ``DSAtrtllmAttentionMetadata``). + On-device, sync-free, and idempotent; ``seq_lens_cpu`` / + ``max_seqlen_k`` keep the optimistic values — they only bound + arange widths that the kernels mask by ``seq_lens``. + """ + super().on_update_kv_lens() + attachment = self.minimax_m3 + if not attachment: + return + meta = attachment["metadata"] + out_cache_loc = attachment["out_cache_loc"] + batch = int(meta.slot_ids.shape[0]) + kv_lens = self.kv_lens_cuda[:batch] + meta.seq_lens[:batch].copy_(kv_lens) + if meta.is_prefill: + # Only the K-side prefix moves with rejections; the Q-side + # structure (cu_seqlens_q, q_batch_row) is fixed per step. + total_q = int(meta.q_positions.shape[0]) + cu = meta.cu_seqlens_q + meta.prefix_lens[:batch].copy_(kv_lens - (cu[1 : batch + 1] - cu[:batch])) + q_positions, cache_slots = derive_q_positions_and_cache_slots( + meta.req_to_token, + meta.prefix_lens[:batch], + cu, + meta.q_batch_row[:total_q], + ) + meta.q_positions[:total_q].copy_(q_positions) + out_cache_loc[:total_q].copy_(cache_slots) + else: + # Identity today; keeps 0-draft corrections right if they + # become reachable. + out_cache_loc[:batch].copy_(derive_decode_cache_slots(meta.req_to_token, kv_lens)) __all__ = [ - "MiniMaxM3SparseConfig", - "MiniMaxM3SparseAttentionMetadata", + "MiniMaxM3AttentionMetadata", + "MiniMaxM3TritonSparseAttentionMetadata", "allocate_minimax_m3_static_buffers", "build_runtime_metadata_from_kv_manager", + "derive_decode_cache_slots", + "derive_q_positions_and_cache_slots", "ensure_metadata_on_device", - "get_minimax_m3_attention_metadata_cls", - "replace_metadata", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index f438d998939f..1762268bf242 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -42,14 +42,30 @@ def get_sparse_attn_kv_cache_manager( ) +def _resolve_minimax_m3_backend_cls( + sparse_params: "SparseParams") -> Type["AttentionBackend"]: + """Select the MiniMax-M3 sparse backend from the lowered params. + + The Triton reference is the default. When implementation is 'msa' the MSA + (fmha_sm100) backend is used instead, gated on SM100 availability so an + unsupported system fails early rather than at kernel launch. + """ + from .minimax_m3 import MiniMaxM3SparseRuntimeBackend + if getattr(sparse_params, "implementation", "triton") == "msa": + from .minimax_m3 import MiniMaxM3MsaSparseAttention + from .minimax_m3.msa_availability import ensure_msa_available + ensure_msa_available() + return MiniMaxM3MsaSparseAttention + return MiniMaxM3SparseRuntimeBackend + + def get_vanilla_sparse_attn_attention_backend( sparse_params: "SparseParams") -> Type["AttentionBackend"]: - from .minimax_m3 import get_minimax_m3_attention_backend_cls from .rocket import RocketVanillaAttention if sparse_params.algorithm == "rocket": return RocketVanillaAttention elif sparse_params.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_params) else: raise ValueError( f"Unsupported sparse attention algorithm in vanilla attention backend: {sparse_params.algorithm}" @@ -62,7 +78,6 @@ def get_trtllm_sparse_attn_attention_backend( from .deepseek_v4 import DeepseekV4TrtllmAttention from .dsa import DSATrtllmAttention - from .minimax_m3 import get_minimax_m3_attention_backend_cls from .rocket import RocketTrtllmAttention if sparse_params.algorithm == "rocket": return RocketTrtllmAttention @@ -78,7 +93,7 @@ def get_trtllm_sparse_attn_attention_backend( # `create_attention(...)` dispatch in `Attention.__init__` # returns an instantiable AttentionBackend under the trtllm # attention backend slot. - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_params) else: raise ValueError( f"Unsupported sparse attention algorithm in trtllm attention backend: {sparse_params.algorithm}" @@ -87,9 +102,8 @@ def get_trtllm_sparse_attn_attention_backend( def get_flashinfer_sparse_attn_attention_backend( sparse_params: "SparseParams") -> Type["AttentionBackend"]: - from .minimax_m3 import get_minimax_m3_attention_backend_cls if sparse_params.algorithm == "minimax_m3": - return get_minimax_m3_attention_backend_cls() + return _resolve_minimax_m3_backend_cls(sparse_params) raise ValueError( f"Unsupported sparse attention algorithm in flashinfer attention backend: {sparse_params.algorithm}" ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index e86993ef8f67..851a55c32f3f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -276,6 +276,12 @@ 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): + del n_valid_blocks, init_blocks, local_blocks + return scores.new_empty((scores.shape[2], scores.shape[0], topk), + dtype=torch.int32) + @torch.library.register_fake("trtllm::userbuffers_allreduce_finalize") def _(input, force_applying_finalize): return torch.empty_like(input) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 714c77be63ab..f3b0dc0476d3 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2041,7 +2041,9 @@ def _(a, b, a_scale, b_scale, tune_max_num_tokens=4096): def silu_and_mul(x: torch.Tensor, scale: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = None, - swiglu_limit: Optional[float] = None) -> torch.Tensor: + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None) -> torch.Tensor: b, n = x.shape assert n % 2 == 0 @@ -2061,6 +2063,8 @@ def grid(meta: Mapping[str, int]) -> tuple[int, int]: x_stride=x.stride(0), d=d, swiglu_limit=swiglu_limit or 0.0, + swiglu_alpha=swiglu_alpha if swiglu_alpha is not None else 1.0, + swiglu_beta=swiglu_beta if swiglu_beta is not None else 0.0, BLOCK_SIZE=1024, HAS_O_SCALE=scale is not None, HAS_SWIGLU_LIMIT=swiglu_limit is not None and swiglu_limit > 0.0, @@ -2075,6 +2079,8 @@ def _( scale: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = None, swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, ) -> torch.Tensor: b, n = x.shape diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 12d7cadfee4f..405810661fba 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -1,3 +1,20 @@ +# 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. + +from collections.abc import Sequence + import numpy as np from tensorrt_llm import logger @@ -12,107 +29,124 @@ from tensorrt_llm._utils import nvtx_range -class IdentityMapper(RegionMapperBase): +class IntactMapper(RegionMapperBase): """ - ---- mapper_identity ---- + ---- mapper_intact ---- + + Copy the selected layers' class regions between slots. Consumes slot-base + pointers from the extractor and expands them with slot-relative per-layer + byte offsets taken from the logical view's buffer entries, so it handles + non-uniform layer strides (a slot may interleave other role classes + between this class's layers). - Pass-through mapping. Do not change pointers or sizes. + src slot: [ base ]--+off(L2)--> [L2 region] --+off(L3)--> [L3 region] ... + dst slot: [ base ]--+off'(L2)-> [L2 region] --+off'(L3)-> [L3 region] ... - src_ptrs: [ S0 ] [ S1 ] [ S2 ] ... - | | | - v v v - dst_ptrs: [ D0 ] [ D1 ] [ D2 ] ... + Layers whose regions are contiguous on BOTH sides are merged into one + fragment, so a fully contiguous class (e.g. K/V in a dedicated pool) + degrades to a single whole-region copy per block. """ - @nvtx_range("IdentityMapper.map") - def map(self, src_regions: SpecRegion, dst_regions: SpecRegion) -> SpecRegionPair: + def __init__( + self, + src_layer_offsets: Sequence[int] | np.ndarray, + dst_layer_offsets: Sequence[int] | np.ndarray, + self_bytes_per_layer: int, + peer_bytes_per_layer: int, + *, + mapper_name: str = "Intact", + ) -> None: + if self_bytes_per_layer != peer_bytes_per_layer: + raise ValueError( + f"{mapper_name} cache region size mismatch: " + f"local={self_bytes_per_layer}, peer={peer_bytes_per_layer}" + ) + src = np.asarray(src_layer_offsets, dtype=np.int64) + dst = np.asarray(dst_layer_offsets, dtype=np.int64) + if src.size == 0 or src.size != dst.size: + raise ValueError( + f"{mapper_name} layer offsets must be non-empty and equal-length: " + f"src={src.size}, dst={dst.size}" + ) + self._runs = self._merge_contiguous(src, dst, self_bytes_per_layer) + + @staticmethod + def _merge_contiguous( + src: np.ndarray, dst: np.ndarray, bytes_per_layer: int + ) -> list[tuple[int, int, int]]: + runs: list[tuple[int, int, int]] = [] + run_start = 0 + for i in range(1, src.size + 1): + if ( + i == src.size + or src[i] != src[i - 1] + bytes_per_layer + or dst[i] != dst[i - 1] + bytes_per_layer + ): + run_layers = i - run_start + runs.append( + (int(src[run_start]), int(dst[run_start]), run_layers * bytes_per_layer) + ) + run_start = i + return runs + + @nvtx_range("IntactMapper.map") + def map(self, src_regions: SpecRegion, dst_regions: SpecRegion): src_group = src_regions.memory dst_group = dst_regions.memory - assert src_group.ptrs.size == dst_group.ptrs.size, ( - f"Number of regions of src({src_group.ptrs.size}) and dst({dst_group.ptrs.size}) must match" - ) - return SpecRegionPair( - src=SpecRegion(memory=src_group, spec=src_regions.spec), - dst=SpecRegion(memory=dst_group, spec=dst_regions.spec), - ) + if src_group.ptrs.size != dst_group.ptrs.size: + raise ValueError( + f"Number of regions of src({src_group.ptrs.size}) and " + f"dst({dst_group.ptrs.size}) must match" + ) + pairs = [ + SpecRegionPair( + src=SpecRegion( + memory=MemRegionGroup( + ptrs=src_group.ptrs + src_off, bytes_per_region=run_bytes + ), + spec=src_regions.spec, + ), + dst=SpecRegion( + memory=MemRegionGroup( + ptrs=dst_group.ptrs + dst_off, bytes_per_region=run_bytes + ), + spec=dst_regions.spec, + ), + ) + for src_off, dst_off, run_bytes in self._runs + ] + return pairs[0] if len(pairs) == 1 else pairs -class HeadMatchMapper(RegionMapperBase): - """ - ---- mapper_head_match ---- - - Move/copy entire contiguous block(s) (multi-layer fragment) as a single chunk. - Align by whole fragment size (frag_size) and apply a constant source/destination block offset. - - src_ptrs: [ S0 ] [ S1 ] ... - | | - + src_off + src_off - | | - [ S0 + src_off ] [ S1 + src_off ] -> (each points to a frag of size frag_size) - copy whole frag - | | - v v - [ D0 + dst_off ] [ D1 + dst_off ] -> (destination frags) - - Contiguous-layer assumption: - This mapper assumes that ``transfer_layers`` consecutive layers - starting at ``src_layer_off`` (and ``dst_layer_off``) are laid out - contiguously within each slot. This holds because - ``buffer_attributes()`` in the storage config assigns buffer - offsets sequentially from 0 for each layer_group (life cycle), - and each PoolDescriptor only contains layers belonging to a single - layer_group. Even when multiple layer_groups share the same - physical storage pool_group, each layer_group independently - occupies the full slot (offsets start from 0), so the contiguous - layout is preserved. +class ReplicatedMapper(IntactMapper): + """Copy TP-replicated per-layer regions without KV-head remapping. + + Every TP rank holds identical bytes per layer (MiniMax M3 index-key, DSA + indexer K), so no head slicing applies. Layer selection under partial PP + overlap happens through the per-layer offsets, and fan-in routing (one + owning sender per destination) is decided upstream by + ``PeerRegistrar.should_send_pool``. """ def __init__( self, - transfer_layers: int, - src_layer_off: int, - dst_layer_off: int, - self_ri: RankInfo, - peer_ri: RankInfo, - slot_size_per_layer: int, - ): - if not isinstance(slot_size_per_layer, int): - raise TypeError( - f"slot_size_per_layer must be int, got {type(slot_size_per_layer).__name__} " - f"(value={slot_size_per_layer}). Use // instead of / for integer division." - ) - self._kv_factor = self_ri.attention.kv_factor - self._frag_size = self._block_size(transfer_layers, slot_size_per_layer=slot_size_per_layer) - self._src_block_off = self._block_size( - src_layer_off, slot_size_per_layer=slot_size_per_layer - ) - self._dst_block_off = self._block_size( - dst_layer_off, slot_size_per_layer=slot_size_per_layer + src_layer_offsets: Sequence[int] | np.ndarray, + dst_layer_offsets: Sequence[int] | np.ndarray, + self_bytes_per_layer: int, + peer_bytes_per_layer: int, + ) -> None: + super().__init__( + src_layer_offsets, + dst_layer_offsets, + self_bytes_per_layer, + peer_bytes_per_layer, + mapper_name="Replicated", ) - @nvtx_range("HeadMatchMapper.map") - def map(self, src_regions: SpecRegion, dst_regions: SpecRegion) -> SpecRegionPair: - src_group = src_regions.memory - dst_group = dst_regions.memory - assert src_group.ptrs.size == dst_group.ptrs.size, ( - f"Number of regions of src({src_group.ptrs.size}) and dst({dst_group.ptrs.size}) must match" - ) - new_src_ptrs = src_group.ptrs + self._src_block_off - new_dst_ptrs = dst_group.ptrs + self._dst_block_off - new_src = MemRegionGroup(ptrs=new_src_ptrs, bytes_per_region=self._frag_size) - new_dst = MemRegionGroup(ptrs=new_dst_ptrs, bytes_per_region=self._frag_size) - return SpecRegionPair( - src=SpecRegion(memory=new_src, spec=src_regions.spec), - dst=SpecRegion(memory=new_dst, spec=dst_regions.spec), - ) - - def _block_size(self, layer_num: int, slot_size_per_layer: int) -> int: - return layer_num * slot_size_per_layer - -class HeadMismatchMapper(RegionMapperBase): +class HNDHeadMismatchMapper(RegionMapperBase): """ - ---- mapper_head_mismatch ---- + ---- mapper_hnd_head_mismatch ---- Fine-grained mapping when head counts or TP/DP partitioning differ. Split layers into per-head (or contiguous-heads) fragments and map them individually. @@ -134,26 +168,60 @@ class HeadMismatchMapper(RegionMapperBase): def __init__( self, - transfer_layers: int, - src_layer_off: int, - peer_layer_off: int, + *, + src_layer_offsets: "Sequence[int] | np.ndarray", + dst_layer_offsets: "Sequence[int] | np.ndarray", self_ri: RankInfo, peer_ri: RankInfo, + self_bytes_per_layer: int, + peer_bytes_per_layer: int, + self_buffers_per_layer: int, + peer_buffers_per_layer: int, ): self._ri = self_ri self._peer_ri = peer_ri - kv_factor = self_ri.attention.kv_factor self_tp_per_dp = self_ri.tp_size_per_dp_group peer_tp_per_dp = peer_ri.tp_size_per_dp_group self_tp_rank = self_ri.tp_rank peer_tp_rank = peer_ri.tp_rank - bytes_per_head = ( - self._ri.attention.tokens_per_block - * self._ri.attention.dims_per_head - * self._ri.attention.element_bytes + if self_buffers_per_layer != peer_buffers_per_layer: + raise ValueError( + "HND buffer count per layer mismatch: " + f"local={self_buffers_per_layer}, peer={peer_buffers_per_layer}" + ) + src_offsets = np.asarray(src_layer_offsets, dtype=np.int64) + dst_offsets = np.asarray(dst_layer_offsets, dtype=np.int64) + if src_offsets.size == 0 or src_offsets.size != dst_offsets.size: + raise ValueError( + "HND layer offsets must be non-empty and equal-length: " + f"src={src_offsets.size}, dst={dst_offsets.size}" + ) + + # Byte geometry is derived from the per-layer region size registered + # by storage (always whole bytes) rather than element_bytes x dims + # arithmetic, so sub-byte dtypes (e.g. NVFP4) stay exact-integer. + src_buffer_bytes = self._bytes_per_buffer( + bytes_per_layer=self_bytes_per_layer, + buffers_per_layer=self_buffers_per_layer, + side="local", + ) + dst_buffer_bytes = self._bytes_per_buffer( + bytes_per_layer=peer_bytes_per_layer, + buffers_per_layer=peer_buffers_per_layer, + side="peer", + ) + bytes_per_head = self._bytes_per_head( + src_buffer_bytes, self._ri.attention.kv_heads_per_rank, side="local" ) + peer_bytes_per_head = self._bytes_per_head( + dst_buffer_bytes, peer_ri.attention.kv_heads_per_rank, side="peer" + ) + if bytes_per_head != peer_bytes_per_head: + raise ValueError( + f"HND bytes per head mismatch: local={bytes_per_head}, peer={peer_bytes_per_head}" + ) self._bytes_cont_heads = ( min(self._ri.attention.kv_heads_per_rank, peer_ri.attention.kv_heads_per_rank) * bytes_per_head @@ -168,55 +236,41 @@ def __init__( peer_kv_heads=peer_ri.attention.kv_heads_per_rank, bytes_per_head=bytes_per_head, ) - self._peer_layer_off = peer_layer_off - - # --- Pre-compute flat 1D offset arrays --- - # - # Each KV cache block (slot) is laid out as: - # - # block_base ──► [layer_0 kv_0] [layer_0 kv_1] [layer_1 kv_0] [layer_1 kv_1] ... - # ◄─ layer_kv ─► ◄─ layer_kv ─► - # ◄────── layer_num (= layer_kv * kv_factor) ──────► - # - # To address fragment (layer=j, kv=k) within a block at base_ptr: - # - # frag_ptr = base_ptr - # + layer_num * (layer_off + j) # skip to the right layer - # + layer_kv * k # skip to key or value - # + head_off # head offset for TP mismatch - # - # The original code computed this as a 3D broadcast in map(): - # bases[:, None, None] + layer_offsets[None, :, None] - # + kv_offsets[None, None, :] + head_off - # producing shape (n_blocks, transfer_layers, kv_factor) then .ravel(). - # - # Optimization: since the (layer, kv) offsets are independent of the - # per-call block base pointers, we pre-compute them here as a flat 1D - # array of length (transfer_layers * kv_factor). At map() time we only - # need np.add.outer(bases, flat_offsets).ravel(), which produces the - # same result in the same C-order traversal (blocks outer, offsets inner) - # but with fewer intermediate allocations. - layer_indices = np.arange(transfer_layers, dtype=np.int64) - kv_indices = np.arange(kv_factor, dtype=np.int64) - - src_layer_kv_num = self._get_layer_kv_num(self._ri) - src_layer_num = src_layer_kv_num * kv_factor - # Shape (transfer_layers, kv_factor) → ravel to 1D - self._src_flat_offsets = ( - src_layer_num * (src_layer_off + layer_indices)[:, None] - + src_layer_kv_num * kv_indices[None, :] - + self._src_head_off - ).ravel() - dst_layer_kv_num = self._get_layer_kv_num(self._peer_ri) - dst_layer_num = dst_layer_kv_num * kv_factor - self._dst_flat_offsets = ( - dst_layer_num * (peer_layer_off + layer_indices)[:, None] - + dst_layer_kv_num * kv_indices[None, :] - + self._dst_head_off + # Pre-compute flat 1D offset arrays: one fragment per (layer, buffer) + # where buffers are the layer's K/V (or scale) buffers laid out + # back-to-back within the layer's region. Layer starts come from the + # view's buffer entries, so interleaved role classes (non-uniform + # layer strides) are handled the same way as everywhere else. At + # map() time a single np.add.outer(bases, flat_offsets) expands the + # per-block base pointers. + self._src_flat_offsets = self._build_flat_offsets( + layer_offsets=src_offsets, + buffers_per_layer=self_buffers_per_layer, + buffer_bytes=src_buffer_bytes, + head_offset=self._src_head_off, + ) + self._dst_flat_offsets = self._build_flat_offsets( + layer_offsets=dst_offsets, + buffers_per_layer=peer_buffers_per_layer, + buffer_bytes=dst_buffer_bytes, + head_offset=self._dst_head_off, + ) + + @staticmethod + def _build_flat_offsets( + *, + layer_offsets: np.ndarray, + buffers_per_layer: int, + buffer_bytes: int, + head_offset: int, + ) -> np.ndarray: + buffer_indices = np.arange(buffers_per_layer, dtype=np.int64) + return ( + layer_offsets[:, None] + buffer_bytes * buffer_indices[None, :] + head_offset ).ravel() - @nvtx_range("HeadMismatchMapper.map") + @nvtx_range("HNDHeadMismatchMapper.map") def map(self, src_regions: SpecRegion, dst_regions: SpecRegion) -> SpecRegionPair: src_group = src_regions.memory dst_group = dst_regions.memory @@ -261,57 +315,167 @@ def _compute_head_offsets( return 0, dst_head_idx * bytes_per_head @staticmethod - def _get_layer_kv_num(ri: RankInfo) -> int: - return ( - ri.attention.kv_heads_per_rank - * ri.attention.tokens_per_block - * ri.attention.dims_per_head - * ri.attention.element_bytes - ) + def _bytes_per_buffer(*, bytes_per_layer: int, buffers_per_layer: int, side: str) -> int: + """Bytes of one buffer (K or V) within a layer's region.""" + if buffers_per_layer <= 0 or bytes_per_layer % buffers_per_layer != 0: + raise ValueError( + f"HND layer geometry is not evenly divisible ({side}): " + f"bytes_per_layer={bytes_per_layer}, buffers_per_layer={buffers_per_layer}" + ) + return bytes_per_layer // buffers_per_layer + @staticmethod + def _bytes_per_head(layer_kv_bytes: int, heads: int, *, side: str) -> int: + """Bytes of one head's rows within a K/V buffer. + + Byte-granular head slicing is only valid when a head lands on a byte + boundary; sub-byte dtypes (e.g. NVFP4) satisfy this whenever the + per-head element count covers whole bytes, which this divisibility + check enforces without any fractional arithmetic. + """ + if heads <= 0 or layer_kv_bytes % heads != 0: + raise ValueError( + f"HND head slicing is not byte-aligned ({side}): " + f"layer_kv_bytes={layer_kv_bytes}, kv_heads={heads}" + ) + return layer_kv_bytes // heads -class IndexerKCacheHeadMatchMapper(RegionMapperBase): - """ - Mapper for indexer K cache when head counts match. - Moves contiguous block(s) as a single chunk, aligned by block_size_per_layer, - with constant source/destination block offsets. +class NHDHeadMismatchMapper(HNDHeadMismatchMapper): + """Map heterogeneous KV heads stored token-major as ``[N, H, D]``. + + ``HNDHeadMismatchMapper`` selects one contiguous head range per K/V buffer, + which is correct for HND storage. In NHD storage, the selected head range + is contiguous only within one token, so this mapper emits one fragment per + ``(layer, K/V, token)``. Only offset precomputation differs from the + parent; the inherited :meth:`map` consumes ``_src_flat_offsets``, + ``_dst_flat_offsets``, and ``_bytes_cont_heads``. """ def __init__( self, - transfer_layers: int, - src_layer_off: int, - dst_layer_off: int, + *, + src_layer_offsets: Sequence[int] | np.ndarray, + dst_layer_offsets: Sequence[int] | np.ndarray, self_ri: RankInfo, peer_ri: RankInfo, - block_size_per_layer: int, - ): - if not isinstance(block_size_per_layer, int): - raise TypeError( - f"block_size_per_layer must be int, got {type(block_size_per_layer).__name__} " - f"(value={block_size_per_layer}). Use // instead of / for integer division." + self_bytes_per_layer: int, + peer_bytes_per_layer: int, + self_buffers_per_layer: int, + peer_buffers_per_layer: int, + ) -> None: + # Deliberately do not call HNDHeadMismatchMapper.__init__: its offsets + # assume HND-contiguous heads. Initialize the three attributes consumed + # by the inherited map() with NHD token-granular geometry instead. + self_tpb = self_ri.attention.tokens_per_block + peer_tpb = peer_ri.attention.tokens_per_block + if self_tpb != peer_tpb: + raise ValueError( + "NHDHeadMismatchMapper requires equal tokens_per_block; " + f"local={self_tpb}, peer={peer_tpb}" ) - self._frag_size = block_size_per_layer * transfer_layers - self._src_block_off = block_size_per_layer * src_layer_off - self._dst_block_off = block_size_per_layer * dst_layer_off - @nvtx_range("IndexerKCacheHeadMatchMapper.map") - def map(self, src_regions: SpecRegion, dst_regions: SpecRegion) -> SpecRegionPair: - src_group = src_regions.memory - dst_group = dst_regions.memory - assert src_group.ptrs.size == dst_group.ptrs.size, ( - f"Number of regions of src({src_group.ptrs.size}) and dst({dst_group.ptrs.size}) must match" + self_heads = self_ri.attention.kv_heads_per_rank + peer_heads = peer_ri.attention.kv_heads_per_rank + if self_buffers_per_layer != peer_buffers_per_layer: + raise ValueError( + "NHD buffer count per layer mismatch: " + f"local={self_buffers_per_layer}, peer={peer_buffers_per_layer}" + ) + src_offsets = np.asarray(src_layer_offsets, dtype=np.int64) + dst_offsets = np.asarray(dst_layer_offsets, dtype=np.int64) + if src_offsets.size == 0 or src_offsets.size != dst_offsets.size: + raise ValueError( + "NHD layer offsets must be non-empty and equal-length: " + f"src={src_offsets.size}, dst={dst_offsets.size}" + ) + + self_bytes_per_token_head = self._bytes_per_token_head( + bytes_per_layer=self_bytes_per_layer, + buffers_per_layer=self_buffers_per_layer, + tokens_per_block=self_tpb, + heads=self_heads, ) - new_src_ptrs = src_group.ptrs + self._src_block_off - new_dst_ptrs = dst_group.ptrs + self._dst_block_off - new_src = MemRegionGroup(ptrs=new_src_ptrs, bytes_per_region=self._frag_size) - new_dst = MemRegionGroup(ptrs=new_dst_ptrs, bytes_per_region=self._frag_size) - return SpecRegionPair( - src=SpecRegion(memory=new_src, spec=src_regions.spec), - dst=SpecRegion(memory=new_dst, spec=dst_regions.spec), + peer_bytes_per_token_head = self._bytes_per_token_head( + bytes_per_layer=peer_bytes_per_layer, + buffers_per_layer=peer_buffers_per_layer, + tokens_per_block=peer_tpb, + heads=peer_heads, + ) + if self_bytes_per_token_head != peer_bytes_per_token_head: + raise ValueError( + "NHD bytes per token/head mismatch: " + f"local={self_bytes_per_token_head}, peer={peer_bytes_per_token_head}" + ) + self._bytes_cont_heads = min(self_heads, peer_heads) * self_bytes_per_token_head + + src_head_off, dst_head_off = HNDHeadMismatchMapper._compute_head_offsets( + self_ri.tp_size_per_dp_group, + peer_ri.tp_size_per_dp_group, + self_ri.tp_rank, + peer_ri.tp_rank, + self_kv_heads=self_heads, + peer_kv_heads=peer_heads, + bytes_per_head=self_bytes_per_token_head, ) + self._src_flat_offsets = self._build_flat_offsets( + layer_offsets=src_offsets, + buffers_per_layer=self_buffers_per_layer, + tokens_per_block=self_tpb, + heads=self_heads, + bytes_per_token_head=self_bytes_per_token_head, + head_offset=src_head_off, + ) + self._dst_flat_offsets = self._build_flat_offsets( + layer_offsets=dst_offsets, + buffers_per_layer=peer_buffers_per_layer, + tokens_per_block=peer_tpb, + heads=peer_heads, + bytes_per_token_head=peer_bytes_per_token_head, + head_offset=dst_head_off, + ) + + @staticmethod + def _bytes_per_token_head( + *, + bytes_per_layer: int, + buffers_per_layer: int, + tokens_per_block: int, + heads: int, + ) -> int: + denominator = buffers_per_layer * tokens_per_block * heads + if denominator <= 0 or bytes_per_layer % denominator != 0: + raise ValueError( + "NHD region geometry is not evenly divisible: " + f"bytes_per_layer={bytes_per_layer}, " + f"buffers_per_layer={buffers_per_layer}, " + f"tokens_per_block={tokens_per_block}, kv_heads={heads}, " + f"denominator={denominator}" + ) + return bytes_per_layer // denominator + + @staticmethod + def _build_flat_offsets( + *, + layer_offsets: np.ndarray, + buffers_per_layer: int, + tokens_per_block: int, + heads: int, + bytes_per_token_head: int, + head_offset: int, + ) -> np.ndarray: + buffer_indices = np.arange(buffers_per_layer, dtype=np.int64) + token_indices = np.arange(tokens_per_block, dtype=np.int64) + token_bytes = heads * bytes_per_token_head + buffer_bytes = tokens_per_block * token_bytes + return ( + layer_offsets[:, None, None] + + buffer_bytes * buffer_indices[None, :, None] + + token_bytes * token_indices[None, None, :] + + head_offset + ).ravel() + class AttentionPolicy: def __init__(self, self_rank_info: RankInfo): @@ -335,9 +499,29 @@ def _mismatch(self, field: str, local, peer) -> bool: local != peer, f"{field} mismatch", field=field, local=local, peer=peer ) - def _tpb_check(self, local: int, peer: int) -> bool: + @staticmethod + def _uses_exact_tpb_mapper(ri: RankInfo) -> bool: + """NHD / replicated pools address bytes inside a block, so their + geometry only lines up when both sides use the same tokens_per_block.""" + if ri.page_table is None: + return False + return any( + pool_view.mapper_kind in (MapperKind.NHD, MapperKind.REPLICATED) + for layer_group in ri.page_table.layer_groups + for pool_view in getattr(layer_group, "pool_views", ()) + ) + + def _tpb_check(self, local: int, peer: int, peer_ri: RankInfo) -> bool: if local == peer: return False + if self._uses_exact_tpb_mapper(self._ri) or self._uses_exact_tpb_mapper(peer_ri): + logger.warning( + "AttentionPolicy: incompatible: tokens_per_block mismatch for " + "NHD/replicated pools; local=%d peer=%d", + local, + peer, + ) + return True larger, smaller = max(local, peer), min(local, peer) if larger % smaller != 0: logger.warning( @@ -367,7 +551,15 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer=peer_ri.cp_size, ) or self._mismatch("element_bytes", a.element_bytes, b.element_bytes) - or self._tpb_check(a.tokens_per_block, b.tokens_per_block) + or self._fail_if( + not self.head_match(peer_ri)[0] + and not float(a.tokens_per_block * a.dims_per_head * a.element_bytes).is_integer(), + "sub-byte head slicing is not byte-aligned", + tokens_per_block=a.tokens_per_block, + dims_per_head=a.dims_per_head, + element_bytes=a.element_bytes, + ) + or self._tpb_check(a.tokens_per_block, b.tokens_per_block, peer_ri) or self._mismatch("dims_per_head", a.dims_per_head, b.dims_per_head) or self._fail_if( a.is_mla and (a.kv_heads_per_rank != 1 or b.kv_heads_per_rank != 1), @@ -405,53 +597,66 @@ def build_kv_mapper( *, peer_ri: RankInfo, mapper_kind: MapperKind, - transfer_layers: int, - self_layer_offset: int, - peer_layer_offset: int, - self_pool_num_layers: int, - peer_pool_num_layers: int, - self_pool_slot_bytes: int, - peer_pool_slot_bytes: int, + self_layer_offsets: "Sequence[int] | np.ndarray", + peer_layer_offsets: "Sequence[int] | np.ndarray", + self_bytes_per_layer: int, + peer_bytes_per_layer: int, + self_buffers_per_layer: int = 1, + peer_buffers_per_layer: int = 1, ) -> RegionMapperBase: - head_match, _ = self.head_match(peer_ri) - - if head_match and transfer_layers == self_pool_num_layers == peer_pool_num_layers: - return IdentityMapper() + """Pick the mapper for one view pair. + + Every view is entries-driven: layer selection always uses explicit + slot-relative per-layer byte offsets, never positional arithmetic + (other role classes may interleave between this class's layers). + The kind only decides the two irreducible semantic differences: + + - REPLICATED skips head matching entirely (bytes are identical on + every TP rank; fan-in ownership is decided upstream). + - Under head mismatch, HND (INDEXED) slices one contiguous head + range per K/V buffer, while NHD must slice inside every token. + + Head-matched views of any kind collapse into IntactMapper, + whose run merging degrades to a single whole-region copy per block + when the selected layers are contiguous on both sides. + """ + if mapper_kind == MapperKind.REPLICATED: + return ReplicatedMapper( + self_layer_offsets, + peer_layer_offsets, + self_bytes_per_layer, + peer_bytes_per_layer, + ) + head_match, _ = self.head_match(peer_ri) if head_match: - if mapper_kind == MapperKind.FLAT: - block_size_per_layer = self_pool_slot_bytes // self_pool_num_layers - return IndexerKCacheHeadMatchMapper( - transfer_layers=transfer_layers, - src_layer_off=self_layer_offset, - dst_layer_off=peer_layer_offset, - self_ri=self._ri, - peer_ri=peer_ri, - block_size_per_layer=block_size_per_layer, - ) - - slot_size_per_layer = self_pool_slot_bytes // self_pool_num_layers - peer_size_per_layer = peer_pool_slot_bytes // peer_pool_num_layers - assert slot_size_per_layer == peer_size_per_layer, ( - f"slot_size_per_layer mismatch between self ({slot_size_per_layer}) " - f"and peer ({peer_size_per_layer}) for HeadMatchMapper" + return IntactMapper( + self_layer_offsets, + peer_layer_offsets, + self_bytes_per_layer, + peer_bytes_per_layer, + mapper_name=mapper_kind.name, ) - return HeadMatchMapper( - transfer_layers=transfer_layers, - src_layer_off=self_layer_offset, - dst_layer_off=peer_layer_offset, + + if mapper_kind == MapperKind.NHD: + return NHDHeadMismatchMapper( + src_layer_offsets=self_layer_offsets, + dst_layer_offsets=peer_layer_offsets, self_ri=self._ri, peer_ri=peer_ri, - slot_size_per_layer=slot_size_per_layer, + self_bytes_per_layer=self_bytes_per_layer, + peer_bytes_per_layer=peer_bytes_per_layer, + self_buffers_per_layer=self_buffers_per_layer, + peer_buffers_per_layer=peer_buffers_per_layer, ) - if mapper_kind == MapperKind.FLAT: - raise ValueError("IndexerKCacheHeadMatchMapper is not supported for head mismatch case") - - return HeadMismatchMapper( - transfer_layers=transfer_layers, - src_layer_off=self_layer_offset, - peer_layer_off=peer_layer_offset, + return HNDHeadMismatchMapper( + src_layer_offsets=self_layer_offsets, + dst_layer_offsets=peer_layer_offsets, self_ri=self._ri, peer_ri=peer_ri, + self_bytes_per_layer=self_bytes_per_layer, + peer_bytes_per_layer=peer_bytes_per_layer, + self_buffers_per_layer=self_buffers_per_layer, + peer_buffers_per_layer=peer_buffers_per_layer, ) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.py index 4c3db7cc8638..931eef38dc16 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.py @@ -1,3 +1,18 @@ +# 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. + from dataclasses import asdict, dataclass @@ -7,7 +22,7 @@ class AttentionInfo: kv_heads_per_rank: int tokens_per_block: int dims_per_head: int - element_bytes: int + element_bytes: int | float enable_attention_dp: bool is_mla: bool diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 1902c9d14f4f..8b03e3a7b1e9 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -1,19 +1,38 @@ +# 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. + +from collections import Counter from dataclasses import dataclass, field from typing import Dict, List, Tuple +import numpy as np + from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import ( + AttentionLayerGroup, + MapperKind, + PoolView, +) from tensorrt_llm._torch.disaggregation.resource.utils import ( - get_global_layer_ids, - get_layer_group_num_layers, + get_layer_byte_ranges, get_layer_to_layer_group, - get_physical_pool, get_pool_view_global_layer_ids, - get_pool_view_num_layers, ) # Type alias for (lg_idx, pool_idx) pair @@ -57,6 +76,33 @@ def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): extractor = KVRegionExtractorV1(peer_ri.page_table) self._peer_ext_cache[key] = extractor + head_match, _ = self._attention_policy.head_match(peer_ri) + if not head_match: + self_page_table = self._self_ext_cache.page_table + nhd_fragments_per_token = sum( + len( + self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].buffer_entries + ) + for layer_group_id, pool_idx in self.get_pool_mapping(peer_ri) + if self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].mapper_kind + == MapperKind.NHD + ) + if nhd_fragments_per_token: + local_heads = self._ri.attention.kv_heads_per_rank + peer_heads = peer_ri.attention.kv_heads_per_rank + logger.warning_once( + "NHD head-mismatched disaggregated KV transfer has no " + "contiguous staging path and will emit approximately " + f"{nhd_fragments_per_token} NIXL descriptors per transferred " + "token per peer, excluding block-level replicated pools " + f"(local_kv_heads={local_heads}, peer_kv_heads={peer_heads}). " + "Long-context TEP/DEP transfers may have high latency.", + key=( + "native-nhd-head-mismatch-" + f"{local_heads}-{peer_heads}-{nhd_fragments_per_token}" + ), + ) + def peer_extractor(self, peer_name: str, peer_rank: int) -> KVRegionExtractorV1: return self._peer_ext_cache[self._unique_key(peer_name, peer_rank)] @@ -118,6 +164,13 @@ def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: Layer-overlap is required: a peer pool with the same pool_role but zero layer overlap with self is *not* a match — the two pools cover disjoint layers and have nothing to transfer. + + A self layer group never matches multiple peer layer groups, so the + result is one peer pool per self pool. Layer groups partition each + rank's layers by attention/life-cycle class, which both sides derive + from the same model config; PP only changes which layers overlap (the + fan-out across peer PP ranks is handled by calling this method once + per peer rank). Step 1 raises if this invariant is ever violated. """ key = self._unique_key(peer_ri.instance_name, peer_ri.instance_rank) if key in self._lg_pool_mapping_cache: @@ -140,27 +193,40 @@ def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: if not isinstance(self_lg, AttentionLayerGroup): continue for self_pi, self_pv in enumerate(self_lg.pool_views): - # The only place mapper_kind affects pool matching: - # INDEXED → pool may cover a subset of the LG; read - # buffer_entries to find the exact layer set. - # FLAT → pool covers the entire LG by convention; - # use the LG's layer ids directly. - self_is_flat = self_pv.mapper_kind == MapperKind.FLAT - pv_global_ids = ( - get_global_layer_ids(self_lg) - if self_is_flat - else get_pool_view_global_layer_ids(self_pv, self_lg) - ) + # Every view carries buffer_entries, so a view's exact layer + # set always comes from its entries (a view may cover a + # subset of the LG when V2 splits an LG into multiple pools + # by buffer-size class, or when a role class exists only on + # some layers, e.g. sparse-layer index-K). + pv_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) if not pv_global_ids: continue - # Step 1: find peer layer_group via any overlapping global_layer_id. - peer_lg_idx = next( - (peer_layer_to_group[g] for g in pv_global_ids if g in peer_layer_to_group), - None, - ) - if peer_lg_idx is None: + # Step 1: find the peer layer_group via overlapping global_layer_ids. + # A self layer group (hence each of its pool views) never matches + # multiple peer layer groups: layer groups partition a rank's + # layers by attention/life-cycle class, both sides derive that + # class from the same model config, and global ids are + # PP-invariant. So PP only changes WHICH layers overlap — layers + # the peer doesn't hold are a legal skip (PP slices, one-sided + # MTP layers) — never how many peer LGs they land in; the PP + # fan-out is handled by per-peer-rank calls of this method. A + # multi-LG hit therefore means the two peers group layers + # differently (unsupported topology), and we fail loudly instead + # of silently transferring only the first LG's overlap. + peer_lg_indices = { + peer_layer_to_group[g] for g in pv_global_ids if g in peer_layer_to_group + } + if not peer_lg_indices: continue + if len(peer_lg_indices) > 1: + raise ValueError( + "PeerRegistrar.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}) spans multiple peer " + f"layer groups {sorted(peer_lg_indices)}; mismatched layer " + "grouping between peers is not supported" + ) + peer_lg_idx = next(iter(peer_lg_indices)) peer_lg = peer_pt.layer_groups[peer_lg_idx] # Step 2: pick the first peer pool with the same pool_role @@ -181,11 +247,7 @@ def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: for peer_pi, peer_pv in enumerate(peer_lg.pool_views): if peer_pv.pool_role != self_pv.pool_role: continue - peer_global_ids = ( - get_global_layer_ids(peer_lg) - if peer_pv.mapper_kind == MapperKind.FLAT - else get_pool_view_global_layer_ids(peer_pv, peer_lg) - ) + peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) if not set(peer_global_ids) & self_layer_set: continue if peer_pv.mapper_kind != self_pv.mapper_kind: @@ -240,53 +302,81 @@ def get_kv_map( f"(local={self_pv.mapper_kind.name}, peer={peer_pv.mapper_kind.name})" ) - # FLAT pools carry no per-buffer layer info, so layer ids and - # layer count come from the layer_group itself. - # - # Sort by global_layer_id so that ``.index(first_overlap_layer)`` - # below returns the layer's slot position. This relies on the - # convention that managers (V1 / V2 / DSv4) assign global_layer_id - # monotonically with the layer's byte offset in the slot. - if self_pv.mapper_kind == MapperKind.FLAT: - self_global_ids = sorted(get_global_layer_ids(self_lg)) - peer_global_ids = sorted(get_global_layer_ids(peer_lg)) - self_num_layers = get_layer_group_num_layers(self_lg) - peer_num_layers = get_layer_group_num_layers(peer_lg) - else: - self_global_ids = sorted(get_pool_view_global_layer_ids(self_pv, self_lg)) - peer_global_ids = sorted(get_pool_view_global_layer_ids(peer_pv, peer_lg)) - self_num_layers = get_pool_view_num_layers(self_pv) - peer_num_layers = get_pool_view_num_layers(peer_pv) - + # Every view is entries-driven: resolve the overlap layers to + # slot-relative byte offsets on each side from the views' buffer + # entries. Layer selection is explicit, so mappers never assume a + # uniform layer stride (other role classes may interleave), and no + # convention about global-id/byte-offset ordering is needed. + self_global_ids = get_pool_view_global_layer_ids(self_pv, self_lg) + peer_global_ids = get_pool_view_global_layer_ids(peer_pv, peer_lg) overlapping_layers = sorted(set(self_global_ids) & set(peer_global_ids)) - transfer_layers = len(overlapping_layers) - - if transfer_layers > 0: - first_overlap_layer = overlapping_layers[0] - self_layer_offset = self_global_ids.index(first_overlap_layer) - peer_layer_offset = peer_global_ids.index(first_overlap_layer) - else: - self_layer_offset = 0 - peer_layer_offset = 0 - self_phys = get_physical_pool(self_pt, self_lg_idx, self_pv.pool_idx) - peer_phys = get_physical_pool(peer_pt, peer_lg_idx, peer_pv.pool_idx) + self_starts, self_bytes_per_layer = get_layer_byte_ranges(self_pv) + peer_starts, peer_bytes_per_layer = get_layer_byte_ranges(peer_pv) + self_g2l = {ll.global_layer_id: ll.local_layer_id for ll in self_lg.local_layers} + peer_g2l = {ll.global_layer_id: ll.local_layer_id for ll in peer_lg.local_layers} + self_layer_offsets = np.array( + [self_starts[self_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 + ) + peer_layer_offsets = np.array( + [peer_starts[peer_g2l[gid]] for gid in overlapping_layers], dtype=np.int64 + ) + # Per-layer buffer count (K and V are separate buffers within a + # layer's region); head-mismatch mappers slice heads inside each. + self_buffers_per_layer = self._get_buffers_per_layer( + self_pv, + layer_group_id=self_lg_idx, + pool_idx=self_pi, + ) + peer_buffers_per_layer = self._get_buffers_per_layer( + peer_pv, + layer_group_id=peer_lg_idx, + pool_idx=peer_pi, + ) mapper = self._attention_policy.build_kv_mapper( peer_ri=peer_ri, mapper_kind=self_pv.mapper_kind, - transfer_layers=transfer_layers, - self_layer_offset=self_layer_offset, - peer_layer_offset=peer_layer_offset, - self_pool_num_layers=self_num_layers, - peer_pool_num_layers=peer_num_layers, - self_pool_slot_bytes=self_phys.slot_bytes, - peer_pool_slot_bytes=peer_phys.slot_bytes, + self_layer_offsets=self_layer_offsets, + peer_layer_offsets=peer_layer_offsets, + self_bytes_per_layer=self_bytes_per_layer, + peer_bytes_per_layer=peer_bytes_per_layer, + self_buffers_per_layer=self_buffers_per_layer, + peer_buffers_per_layer=peer_buffers_per_layer, ) self._kv_map_cache[cache_key] = mapper return mapper + @staticmethod + def _get_buffers_per_layer( + pool_view: PoolView, + *, + layer_group_id: int, + pool_idx: int, + ) -> int: + """Per-layer buffer count of a view (e.g. K+V -> 2, key-only -> 1). + + Views are bucketed per (layer group, pool, mapper kind) at page-table + build time, so every layer in a view carries the same role set and + hence the same entry count — a skewed distribution should never occur. + Still verify it per layer rather than via total-count divisibility: + e.g. 1 + 3 entries over two layers passes ``total % layers == 0`` yet + would make head-slicing mappers split every layer at wrong offsets. + """ + entries = pool_view.buffer_entries + if len(entries) == 0: + return 1 + counts = Counter(int(e["local_layer_id"]) for e in entries) + distinct = set(counts.values()) + if len(distinct) != 1: + raise ValueError( + "PoolView buffer entries are not evenly distributed across layers: " + f"layer_group={layer_group_id}, pool={pool_idx}, " + f"per-layer entry counts={sorted(counts.items())}" + ) + return distinct.pop() + @staticmethod def _find_overlap(self_val, peer_val, self_rank, peer_rank=None): if self_val <= peer_val: @@ -368,6 +458,44 @@ def should_send_kv(self, peer_overlap: PeerOverlap, peer_rank_info: RankInfo) -> self_tp_rank_in_dp_group % dup_head_factor ) + def _owns_tp_fan_in(self, peer_rank_info: RankInfo) -> bool: + """Elect one owner when replicated bytes fan in across TP ranks. + + A peer with fewer TP shards receives identical replicated data from + several local ranks. Rotate the elected owner by the destination's + DP rank (mirroring ``should_send_kv``'s head-duplication pairing) so + that with a multi-DP-group generation side the extra replicated + traffic spreads across local ranks instead of always landing on the + first rank of each fan-in group. + """ + ratio = max( + 1, + self._ri.tp_size_per_dp_group // peer_rank_info.tp_size_per_dp_group, + ) + self_tp_rank = self._ri.tp_rank % self._ri.tp_size_per_dp_group + return self_tp_rank % ratio == peer_rank_info.dp_rank % ratio + + def should_send_pool( + self, + peer_overlap: PeerOverlap, + peer_rank_info: RankInfo, + layer_group_id: int, + pool_idx: int, + ) -> bool: + """Return whether this rank owns the transfer of one view pair. + + ``pool_idx`` indexes the layer group's ``pool_views`` list (one view + per role class; several views may share a physical pool). Each view + is kind-homogeneous, so ownership is a single per-view decision: + replicated views use one sender per fan-in group, sharded views + retain head-duplication routing. + """ + layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] + pool_view = layer_group.pool_views[pool_idx] + if pool_view.mapper_kind == MapperKind.REPLICATED: + return self._owns_tp_fan_in(peer_rank_info) + return self.should_send_kv(peer_overlap, peer_rank_info) + def should_send_aux(self, peer_rank_info: RankInfo) -> bool: # to ensure the transfer aux is not duplicated diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 12a614ca2dad..8620efebb684 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -1,3 +1,18 @@ +# 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. + from dataclasses import asdict, dataclass from typing import List, Optional @@ -59,6 +74,14 @@ def from_kv_cache_manager( m = kv_cache_manager.mapping kvm = kv_cache_manager enable_attention_dp = m.enable_attention_dp + # Eight is the smallest element count guaranteed to occupy whole bytes + # for every supported sub-byte cache dtype (including NVFP4). + bytes_for_eight_elements = get_size_in_bytes(8, kvm.dtype) + element_bytes = ( + bytes_for_eight_elements // 8 + if bytes_for_eight_elements % 8 == 0 + else bytes_for_eight_elements / 8 + ) return cls( instance_name=instance_name, instance_rank=m.rank, @@ -80,7 +103,7 @@ def from_kv_cache_manager( kv_heads_per_rank=kvm.num_kv_heads_per_layer[0], tokens_per_block=kvm.tokens_per_block, dims_per_head=kvm.head_dim, - element_bytes=get_size_in_bytes(1, kvm.dtype), + element_bytes=element_bytes, enable_attention_dp=enable_attention_dp, is_mla=kvm.kv_factor == 1, ), diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 228a321ec23f..8e0e004dda6c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1,3 +1,18 @@ +# 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. + from __future__ import annotations import os @@ -48,6 +63,7 @@ from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -766,85 +782,86 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write peer_extractor = self._registrar.peer_extractor( peer_ri.instance_name, peer_ri.instance_rank ) - if self._registrar.should_send_kv(targets, peer_ri): - pool_mapping = self._registrar.get_pool_mapping(peer_ri) - dst_block_ids_per_groups = req_info.block_ids_per_layer_groups - src_block_ids_per_groups = task._slice.block_ids_per_layer_groups - - # Aggregate fragments from all matching pools using numpy concatenation - for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): - src_block_ids = src_block_ids_per_groups[self_lg] - dst_block_ids = dst_block_ids_per_groups[peer_lg] - - # Both sides trim block lists to ceil(prompt_len / tpb) in - # _create_kv_slice, so dst must never exceed src. A smaller dst - # (generation prefix-cache reuse) is handled via dst_start below. - block_diff = dst_block_ids.size - src_block_ids.size - if block_diff > 0: - raise ValueError( - f"src/dst block count mismatch: {src_block_ids.size} vs " - f"{dst_block_ids.size} (dst must not exceed src)" - ) - tpb = extractor.page_table.tokens_per_block - token_range = task._slice.token_range - lg_info = extractor.page_table.layer_groups[self_lg] - window_size = getattr(lg_info, "sliding_window_size", None) - - # Block lists are the suffix of [..., slice_end); cached prefix - # is implicit in their size. token_start = (total_blocks - n) * tpb. - slice_end = token_range.end if token_range is not None else 0 - total_blocks = (slice_end + tpb - 1) // tpb - src_beam0_blocks = Sender._beam0_block_count( - src_block_ids, total_blocks, task._beam_width - ) - dst_beam0_blocks = Sender._beam0_block_count( - dst_block_ids, total_blocks, task._beam_width - ) - assert src_beam0_blocks <= total_blocks, ( - f"src beam-0 block list ({src_beam0_blocks}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" - ) - assert dst_beam0_blocks <= total_blocks, ( - f"dst beam-0 block list ({dst_beam0_blocks}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + pool_mapping = self._registrar.get_pool_mapping(peer_ri) + dst_block_ids_per_groups = req_info.block_ids_per_layer_groups + src_block_ids_per_groups = task._slice.block_ids_per_layer_groups + + # Aggregate fragments from all matching pools using numpy concatenation. + # Send ownership is per pool: replicated pools elect one fan-in + # owner, sharded pools keep head-duplication routing. + for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): + if not self._registrar.should_send_pool(targets, peer_ri, self_lg, self_pi): + continue + src_block_ids = src_block_ids_per_groups[self_lg] + dst_block_ids = dst_block_ids_per_groups[peer_lg] + + # Both sides trim block lists to ceil(prompt_len / tpb) in + # _create_kv_slice, so dst must never exceed src. A smaller dst + # (generation prefix-cache reuse) is handled via dst_start below. + block_diff = dst_block_ids.size - src_block_ids.size + if block_diff > 0: + raise ValueError( + f"src/dst block count mismatch: {src_block_ids.size} vs " + f"{dst_block_ids.size} (dst must not exceed src)" ) - src_start = (total_blocks - src_beam0_blocks) * tpb - dst_start = (total_blocks - dst_beam0_blocks) * tpb - if req_info.dst_start_token is not None: - dst_start = max(dst_start, req_info.dst_start_token) - if window_size is not None: - # SWA stale_end uses the request prompt_len (not slice_end — - # they differ for non-final slices). prompt_len must be plumbed - # via the session; falling back to slice_end is wrong on - # non-final slices. - assert task._prompt_len is not None, ( - "SWA layer requires session.prompt_len; " - "set TxSession(prompt_len=request.prompt_len)." - ) - stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb) - src_start = max(stale_end * tpb, src_start) - dst_start = max(stale_end * tpb, dst_start) - src_block_ids, dst_block_ids = Sender._align_kv_blocks( - src_block_ids, - dst_block_ids, - src_token_start=src_start, - dst_token_start=dst_start, - tokens_per_block=tpb, + tpb = extractor.page_table.tokens_per_block + token_range = task._slice.token_range + lg_info = extractor.page_table.layer_groups[self_lg] + window_size = getattr(lg_info, "sliding_window_size", None) + + # Block lists are the suffix of [..., slice_end); cached prefix + # is implicit in their size. token_start = (total_blocks - n) * tpb. + slice_end = token_range.end if token_range is not None else 0 + total_blocks = (slice_end + tpb - 1) // tpb + src_beam0_blocks = Sender._beam0_block_count( + src_block_ids, total_blocks, task._beam_width + ) + dst_beam0_blocks = Sender._beam0_block_count( + dst_block_ids, total_blocks, task._beam_width + ) + assert src_beam0_blocks <= total_blocks, ( + f"src beam-0 block list ({src_beam0_blocks}) exceeds total slice " + f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + ) + assert dst_beam0_blocks <= total_blocks, ( + f"dst beam-0 block list ({dst_beam0_blocks}) exceeds total slice " + f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + ) + src_start = (total_blocks - src_beam0_blocks) * tpb + dst_start = (total_blocks - dst_beam0_blocks) * tpb + if req_info.dst_start_token is not None: + dst_start = max(dst_start, req_info.dst_start_token) + if window_size is not None: + # SWA stale_end uses the request prompt_len (not slice_end — + # they differ for non-final slices). prompt_len must be plumbed + # via the session; falling back to slice_end is wrong on + # non-final slices. + assert task._prompt_len is not None, ( + "SWA layer requires session.prompt_len; " + "set TxSession(prompt_len=request.prompt_len)." ) + stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb) + src_start = max(stale_end * tpb, src_start) + dst_start = max(stale_end * tpb, dst_start) + src_block_ids, dst_block_ids = Sender._align_kv_blocks( + src_block_ids, + dst_block_ids, + src_token_start=src_start, + dst_token_start=dst_start, + tokens_per_block=tpb, + ) - src_region = extractor.extract( - src_block_ids, layer_group_id=self_lg, pool_idx=self_pi - ) - dst_region = peer_extractor.extract( - dst_block_ids, layer_group_id=peer_lg, pool_idx=peer_pi - ) - mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) - region_pair = mapper.map(src_region, dst_region) - region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] - for rp in region_pairs: - src_frag_parts.append(rp.src.memory.ptrs) - dst_frag_parts.append(rp.dst.memory.ptrs) - size_specs.append((rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region)) + src_region = extractor.extract(src_block_ids, layer_group_id=self_lg, pool_idx=self_pi) + dst_region = peer_extractor.extract( + dst_block_ids, layer_group_id=peer_lg, pool_idx=peer_pi + ) + mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) + region_pair = mapper.map(src_region, dst_region) + region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] + for rp in region_pairs: + src_frag_parts.append(rp.src.memory.ptrs) + dst_frag_parts.append(rp.dst.memory.ptrs) + size_specs.append((rp.src.memory.ptrs.size, rp.src.memory.bytes_per_region)) if src_frag_parts: src_frags = np.concatenate(src_frag_parts) @@ -1530,6 +1547,14 @@ def _fanin_bounce_safe(overlap, peer_ri) -> bool: lpp = getattr(peer_ri, "layer_num_per_pp", None) if not lpp or len(lpp) < overlap.overlap_pp_size or len(set(lpp)) != 1: return False + # Replicated pools (e.g. MiniMax M3 index-key) are sent by one elected + # fan-in owner only, so with multiple writers their contributions + # differ in size and the equal split is invalid. + if len(overlap.ranks) > 1 and peer_ri.page_table is not None: + for layer_group in peer_ri.page_table.layer_groups: + for pool_view in getattr(layer_group, "pool_views", ()): + if pool_view.mapper_kind == MapperKind.REPLICATED: + return False return True def dispatch_task(self, task: KVRecvTask): diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 92f0a9c2f4ad..4aa76c5b9605 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -1,3 +1,19 @@ +# 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. + +from collections import defaultdict from typing import Dict, List import numpy as np @@ -20,12 +36,28 @@ PhysicalPoolGroup, PoolView, ) -from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool +from tensorrt_llm._torch.disaggregation.resource.utils import ( + compute_layer_byte_ranges, + get_physical_pool, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes, nvtx_range from tensorrt_llm.bindings import DataType +# Mapper kinds a V2 manager may declare via get_disagg_role_mapper_kinds(). +# A physical pool may mix kinds (V2 storage coalesces buffers purely by +# size within a life cycle); the page-table builder emits one PoolView per +# (pool, kind) so each view stays kind-homogeneous. +_V2_ROLE_MAPPER_KINDS = frozenset( + { + MapperKind.INDEXED, + MapperKind.REPLICATED, + MapperKind.NHD, + } +) + class KVRegionExtractorV1(RegionExtractorBase): """ @@ -60,8 +92,9 @@ def extract( described by region_ids. For KV cache: each ptr = base_address + slot_id * slot_bytes, pointing - to the start of a full slot. The slot contains buffer entries for all - layers in this layer_group laid out contiguously from offset 0. + to the start of a full slot. Sub-slot selection (layers, role classes, + heads) is the mappers' responsibility; logical views carry that + geometry in their buffer entries. Args: layer_group_id: The layer group index (= life cycle index). @@ -193,11 +226,15 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), pool_role=frozenset(kv_role_names), mapper_kind=MapperKind.INDEXED, + bytes_per_layer=stride, ) physical_pools = [kv_physical] pool_views = [kv_view] - # Indexer K cache support + # Indexer K cache support. The DSA indexer K cache is identical on + # every TP rank (single index head), so its view is REPLICATED with + # one synthesized buffer entry per local layer: the slot packs the + # layers equal-sized in local-layer order. if getattr(kv_cache_manager, "enable_indexer_k_cache", False): indexer_pool = kv_cache_manager.impl.get_indexer_k_cache_pool() # indexer_pool shape: (numBlocks, numLayers, kvFactor, blockSize), dtype=UINT8 @@ -211,11 +248,19 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: slot_bytes=indexer_slot_bytes, num_slots=num_blocks, ) + indexer_bytes_per_layer = indexer_slot_bytes // len(local_layer_ids) indexer_view = PoolView( pool_idx=1, - buffer_entries=np.array([], dtype=BUFFER_ENTRY_DTYPE), + buffer_entries=np.array( + [ + (lid, i * indexer_bytes_per_layer, indexer_bytes_per_layer) + for i, lid in enumerate(local_layer_ids) + ], + dtype=BUFFER_ENTRY_DTYPE, + ), pool_role=frozenset({"indexer_k"}), - mapper_kind=MapperKind.FLAT, + mapper_kind=MapperKind.REPLICATED, + bytes_per_layer=indexer_bytes_per_layer, ) physical_pools.append(indexer_physical) pool_views.append(indexer_view) @@ -288,18 +333,48 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. - Uses KVCacheManagerV2's public pool_group_descs layout API. A physical - pool group may be shared by several layer groups; layer_groups remains - indexed by layer_group_id while pool_group_idx points at the shared - physical pool group entry. - - Each PoolView is stamped with the manager's native role-name strings + Uses KVCacheManagerV2's public ``pool_group_descs`` layout API and + stamps each PoolView with the manager's native role-name strings (``pool_role``) plus the closed-set ``mapper_kind`` discriminator used by ``build_kv_mapper``. + + A physical pool group may be shared by several layer groups (life + cycles whose coalesced-buffer sizes are identical); each layer group + is exactly one ``SlotDescVariant`` of one pool group, so iterating + variants visits every layer group once. ``layer_groups`` stays indexed + by layer_group_id while ``pool_group_idx`` points at the shared + physical pool group entry, so per-window transfer logic keeps working. """ config = manager.impl.init_config pool_group_descs = manager.impl.pool_group_descs + # Every V2 manager declares how native roles map to the closed set of + # disaggregation mapper kinds; Role.ALL is the required fallback. + role_mapper_kinds = manager.get_disagg_role_mapper_kinds() + if Role.ALL not in role_mapper_kinds: + raise ValueError("Disaggregation role mapping must define Role.ALL") + for role, mapper_kind in role_mapper_kinds.items(): + if not isinstance(mapper_kind, MapperKind): + raise ValueError( + f"Invalid disaggregation mapper kind {mapper_kind!r} for role {role!s}" + ) + if mapper_kind not in _V2_ROLE_MAPPER_KINDS: + supported = ", ".join(kind.name for kind in sorted(_V2_ROLE_MAPPER_KINDS)) + raise ValueError( + f"Unsupported V2 disaggregation mapper kind {mapper_kind.name} " + f"for role {role!s}; supported kinds: {supported}" + ) + # INDEXED is the whole-manager legacy default, not a per-role + # choice: it may only appear as the Role.ALL fallback. Side-cache + # roles (e.g. INDEX_KEY) may declare their own non-INDEXED kind + # alongside it. + if mapper_kind is MapperKind.INDEXED and role != Role.ALL: + raise ValueError( + f"MapperKind.INDEXED is only valid as the Role.ALL mapping; " + f"got it for role {role!s}" + ) + default_mapper_kind = role_mapper_kinds[Role.ALL] + def _window_size_for_layer(internal_layer_id: int): if internal_layer_id < len(config.layers): return getattr(config.layers[internal_layer_id], "window_size", None) @@ -336,6 +411,11 @@ def _window_size_for_layer(internal_layer_id: int): ) ) + # Each variant is one layer group (life cycle) drawing slots from + # this pool group. Multiple layer groups share a pool group when + # their coalesced-buffer sizes are identical; within a slot, each + # layer group's buffer offsets start from 0 independently — the + # memory is reused, not concatenated. for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) @@ -346,31 +426,78 @@ def _window_size_for_layer(internal_layer_id: int): for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) ] - pool_views = [] + # Bucket buffer entries by (pool, mapper kind). One PoolView is + # emitted per bucket and spans every layer of that role class, + # so the view count per layer group is bounded by the number of + # role classes — never by the layer count. A physical pool may + # hold several classes (V2 storage coalesces buffers purely by + # size within a layer group, so e.g. MiniMax M3's index-K shares + # the K/V pool when their per-block sizes coincide); each class + # still gets its own view, which keeps peer matching independent + # of that physical coalescing decision. ``pool_role`` stays the + # manager-supplied equivalence label used for peer matching + # without enumerating role names. Buffer offsets within a slot + # follow ``buffer_ids`` order: the i-th buffer of a coalesced + # buffer lives at ``i * single_buffer_size``. + bucket_entries: Dict[tuple, list] = defaultdict(list) + bucket_roles: Dict[tuple, set] = defaultdict(set) for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): - entries = [] - # Native role-name strings for this pool — used as - # ``PoolView.pool_role``, the manager-supplied equivalence - # label that disagg uses to match pools across peers without - # enumerating roles. - native_roles: set = set() - offset = 0 single_buffer_size = int(coalesced_buffer.single_buffer_size) + offset = 0 for buffer_id in coalesced_buffer.buffer_ids: - entries.append((int(buffer_id.layer_id), offset, single_buffer_size)) - native_roles.add(str(buffer_id.role)) + kind = role_mapper_kinds.get(buffer_id.role, default_mapper_kind) + bucket_key = (pool_idx, kind) + bucket_entries[bucket_key].append( + (int(buffer_id.layer_id), offset, single_buffer_size) + ) + bucket_roles[bucket_key].add(str(buffer_id.role)) offset += single_buffer_size - if entries: - pool_views.append( - PoolView( - pool_idx=pool_idx, - buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), - pool_role=frozenset(native_roles), - mapper_kind=MapperKind.INDEXED, - ) + # Emit this layer group's views: one per (pool, mapper-kind + # class of roles). Roles sharing a kind share a view + # (KEY+VALUE); roles with different kinds in the same physical + # pool get separate views (M3 coalesced index-K). + # All ordering below is canonicalization — the page table is + # serialized and matched against peers, so view order (pool, + # then lowest slot offset), entry order (slot offset), and role + # text must not depend on dict/set iteration order. + pool_views = [] + lg_bucket_keys = sorted( + bucket_entries, + key=lambda key: (key[0], min(entry[1] for entry in bucket_entries[key])), + ) + for bucket_key in lg_bucket_keys: + pool_idx, mapper_kind = bucket_key + roles = frozenset(bucket_roles[bucket_key]) + entries = np.array( + sorted(bucket_entries[bucket_key], key=lambda entry: entry[1]), + dtype=BUFFER_ENTRY_DTYPE, + ) + # Fail fast on invalid geometry and record the uniform + # per-layer region size on the wire. Every kind is + # entries-driven, so the contiguous-layer-region / + # uniform-size invariants apply to all views uniformly. + _, bytes_per_layer = compute_layer_byte_ranges( + entries, + context=( + f"View(layer_group={layer_group_id}, pool={pool_idx}, " + f"kind={mapper_kind.name}, role={sorted(roles)})" + ), + ) + pool_views.append( + PoolView( + pool_idx=pool_idx, + buffer_entries=entries, + pool_role=roles, + mapper_kind=mapper_kind, + bytes_per_layer=bytes_per_layer, ) + ) + # Determine layer group metadata. + # For managers with virtual layers, internal layer_ids + # may exceed the length of num_kv_heads_per_layer. Use index 0 as + # all layers within a pool group share the same kv_heads count. first_local_layer = all_internal_layer_ids[0] if first_local_layer < len(manager.num_kv_heads_per_layer): num_kv_heads = manager.num_kv_heads_per_layer[first_local_layer] diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 81514c06d6a0..259561cfd354 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -1,3 +1,18 @@ +# 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. + from __future__ import annotations from dataclasses import dataclass, field @@ -16,23 +31,38 @@ class MapperKind(IntEnum): - """Slot metadata shape — selects how disagg derives the pool's layer set. - - INDEXED: PoolView.buffer_entries lists ``(local_layer_id, offset, size)`` - per buffer. Disagg reads ``local_layer_id`` to know *which* layers - from the LG live in this pool (a pool may cover a subset when V2 - splits an LG into multiple pools by buffer-size class). The - ``offset`` / ``size`` columns are carried for future use but are not - currently consumed at byte-transfer time. - FLAT: PoolView.buffer_entries is empty. Disagg assumes the pool - covers *all* layers of the LG, packed equal-sized in - ``local_layers`` order. Used today by the DSA (DeepSeek Sparse - Attention, v3.2) indexer K cache pool, whose slot layout is a dense - ``(numLayers, kvFactor, blockSize)`` array. - - Byte arithmetic is the same for both kinds: per-layer stride is - ``slot_bytes // num_layers``. The kind only affects how disagg discovers - the pool's layer set during pool matching. + """Transfer semantics of one physical pool's bytes. + + Every PoolView carries ``buffer_entries`` listing + ``(local_layer_id, offset, size)`` per buffer; the view's exact layer + set always comes from those entries (a view may cover a subset of the + LG when V2 splits an LG into multiple pools by buffer-size class, or + when a role class exists only on some layers). The kind selects how + bytes move between heterogeneous topologies: + + INDEXED: Head-major (HND) K/V — the layout written by the TRTLLM + attention kernels and the default for V1 and standard V2 managers. + Heterogeneous-head transfer selects one contiguous head-major range + per K/V buffer. + REPLICATED: The pool holds bytes that are identical on every TP rank + (MiniMax M3 index-key, DSA indexer K). Copied without KV-head + remapping using per-layer strides; fan-in routing elects one owning + sender per destination so each peer receives exactly one copy. + NHD: Ordinary K/V whose per-buffer storage is token-major + ``[token, head, dim]``. Heterogeneous-head transfer must select the + corresponding head slice inside every token rather than a single + contiguous head-major range. + + A physical pool may hold roles of different kinds: V2 storage coalesces + buffers purely by ``(life_cycle, buffer size)``, so e.g. MiniMax M3's + replicated index-K shares the K/V pool at TP degrees where their + per-block sizes coincide. The page-table builder therefore emits one + PoolView per ``(physical pool, mapper kind)`` — a view covers exactly + the bytes of one role class, and its per-layer byte ranges come from + ``buffer_entries`` (offsets are per layer because another class may + interleave between layers; only the per-layer size is uniform, recorded + in ``bytes_per_layer``). View count per layer group is bounded by the + number of role classes, never by layer count. Mamba state pools do not use this enum: Mamba's transfer is dispatched through :class:`MambaPolicy` which hard-codes the ``is_conv`` switch and @@ -40,7 +70,8 @@ class MapperKind(IntEnum): """ INDEXED = 0 - FLAT = 1 + REPLICATED = 1 + NHD = 2 @dataclass @@ -107,7 +138,7 @@ class PoolView: pool_idx: Index of the physical pool within its pool group. buffer_entries: Structured array using ``BUFFER_ENTRY_DTYPE``. Each entry records a buffer's ``local_layer_id`` and its byte ``offset`` - and ``size`` within the pool slot. FLAT pools have no entries. + and ``size`` within the pool slot. pool_role: Set of native role-name strings (whatever the cache manager uses, e.g. ``"key"`` / ``"value"`` / ``"deepseek_v4_swa"``) that live in this pool. Used as the *equivalence label* for peer-to-peer @@ -115,12 +146,19 @@ class PoolView: are equal. Disagg never enumerates the role-name vocabulary — adding a new role on the manager side requires no disagg change. mapper_kind: Closed-set discriminator for picking the Mapper family. + bytes_per_layer: Uniform byte size of one layer's region within the + slot. The per-layer *offsets* live in ``buffer_entries``; only + the size is uniform, because a slot may interleave other role + classes between layers, making the layer stride non-uniform. + Set for every kind; ``None`` only in tables serialized by older + builders, where consumers re-derive it from the entries. """ pool_idx: int buffer_entries: np.ndarray # dtype=BUFFER_ENTRY_DTYPE pool_role: FrozenSet[str] = field(default_factory=frozenset) mapper_kind: MapperKind = MapperKind.INDEXED + bytes_per_layer: Optional[int] = None def to_dict(self) -> dict: return { @@ -128,6 +166,9 @@ def to_dict(self) -> dict: "buffer_entries": self.buffer_entries.tolist(), "pool_role": sorted(self.pool_role), "mapper_kind": int(self.mapper_kind), + "bytes_per_layer": ( + int(self.bytes_per_layer) if self.bytes_per_layer is not None else None + ), } @staticmethod @@ -143,6 +184,9 @@ def from_dict(data: dict) -> "PoolView": ), pool_role=frozenset(data["pool_role"]), mapper_kind=MapperKind(int(data["mapper_kind"])), + bytes_per_layer=( + int(data["bytes_per_layer"]) if data.get("bytes_per_layer") is not None else None + ), ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 21c4d98bd2aa..6f58bc2782e4 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -1,3 +1,18 @@ +# 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. + from __future__ import annotations from typing import Dict, List, Set @@ -26,6 +41,64 @@ def get_slot_address(pool: PhysicalPool, slot_id: int) -> int: # ------------------------------------------------------------------------- +def compute_layer_byte_ranges( + buffer_entries, + *, + declared_bytes_per_layer: "int | None" = None, + context: str = "PoolView", +) -> tuple[Dict[int, int], int]: + """Per-layer slot-relative byte offsets from raw buffer entries. + + Returns ``({local_layer_id: start_offset}, bytes_per_layer)``. A layer's + region is the concatenation of its buffer entries, which must be + contiguous within the slot; the region size must be uniform across + layers (the slot may interleave other role classes between layers, so + only the *size* is uniform — offsets are per layer). ``context`` labels + error messages; ``declared_bytes_per_layer`` cross-checks a size that + was recorded elsewhere. + """ + starts: Dict[int, int] = {} + totals: Dict[int, int] = {} + entries_by_layer: Dict[int, list] = {} + for entry in buffer_entries: + entries_by_layer.setdefault(int(entry["local_layer_id"]), []).append( + (int(entry["offset"]), int(entry["size"])) + ) + if not entries_by_layer: + raise ValueError(f"{context} has no buffer entries; per-layer byte ranges are undefined") + for layer_id, spans in entries_by_layer.items(): + spans.sort() + for (off, size), (next_off, _) in zip(spans, spans[1:]): + if off + size != next_off: + raise ValueError( + f"{context} layer {layer_id} buffers are " + f"not contiguous: [{off}, {off + size}) is followed by offset {next_off}" + ) + starts[layer_id] = spans[0][0] + totals[layer_id] = sum(size for _, size in spans) + distinct_totals = set(totals.values()) + if len(distinct_totals) != 1: + raise ValueError( + f"{context} per-layer region sizes are not uniform: {sorted(totals.items())}" + ) + bytes_per_layer = distinct_totals.pop() + if declared_bytes_per_layer is not None and declared_bytes_per_layer != bytes_per_layer: + raise ValueError( + f"{context} declares bytes_per_layer={declared_bytes_per_layer} but buffer " + f"entries sum to {bytes_per_layer} per layer" + ) + return starts, bytes_per_layer + + +def get_layer_byte_ranges(pool_view: PoolView) -> tuple[Dict[int, int], int]: + """Per-layer byte ranges of a view; see :func:`compute_layer_byte_ranges`.""" + return compute_layer_byte_ranges( + pool_view.buffer_entries, + declared_bytes_per_layer=pool_view.bytes_per_layer, + context=f"PoolView(pool_idx={pool_view.pool_idx}, role={sorted(pool_view.pool_role)})", + ) + + def get_unique_layers(pool_view: PoolView) -> Set[int]: """Unique local layer IDs in *pool_view*.""" return {int(e["local_layer_id"]) for e in pool_view.buffer_entries} @@ -144,13 +217,24 @@ def get_unique_pool_memory_descs( def get_layer_to_layer_group(page_table: KVCachePageTable) -> Dict[int, int]: """ - Build ``{global_layer_id: lg_idx}`` mapping + Build ``{global_layer_id: lg_idx}`` mapping. + + Layer groups must partition a rank's attention layers: every + global_layer_id belongs to exactly one group. Peer matching relies on + this, so a duplicate raises instead of silently keeping the last group. """ out: Dict[int, int] = {} for lg_idx, lg in enumerate(page_table.layer_groups): if isinstance(lg, AttentionLayerGroup): for ll in lg.local_layers: - out[int(ll.global_layer_id)] = int(lg_idx) + gid = int(ll.global_layer_id) + if gid in out: + raise ValueError( + f"global_layer_id {gid} appears in layer groups " + f"{out[gid]} and {lg_idx}; layer groups must partition " + "a rank's attention layers" + ) + out[gid] = int(lg_idx) return out diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py new file mode 100644 index 000000000000..bb6d58957fd3 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py @@ -0,0 +1,57 @@ +# 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. + +from torch import Tensor, nn + +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + +MINIMAX_M3_PARAMS_MAP = { + r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": r"\1.gate.e_score_correction_bias", +} + + +@register_mapper("HF", "MiniMaxM3SparseForCausalLM") +@register_mapper("HF", "MiniMaxM3SparseForConditionalGeneration") +class MiniMaxM3HfWeightMapper(HfWeightMapper): + """Handle M3 gate naming and MXFP8 GQA duplication for loader v2.""" + + def __init__(self) -> None: + super().__init__() + self.params_map = MINIMAX_M3_PARAMS_MAP + + def _duplicate_kv_weights( + self, module: nn.Module, new_name: str, weights: dict[str, Tensor] + ) -> dict[str, Tensor]: + if new_name not in ["k_proj", "v_proj"]: + return weights + + duplicated_keys = ["weight", "bias"] + quant_config = getattr(module, "quant_config", None) + if quant_config is not None: + quant_mode = quant_config.quant_mode + if quant_mode.has_nvfp4(): + duplicated_keys.append("weight_scale") + if quant_mode.has_mxfp8(): + duplicated_keys.extend(["weight_scale", "weight_scale_inv"]) + + return { + key: self._duplicate_kv( + weight=value[:], num_kv_heads=self._num_kv_heads, tensor_parallel_size=self._tp_size + ) + if key in duplicated_keys + else value + for key, value in weights.items() + } diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e38b9f8be115..cd69ec4bba0b 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -21,7 +21,7 @@ import copy import dataclasses -from functools import partial +import os from typing import Any, Dict, List, Optional, Tuple from typing import Mapping as TMapping @@ -30,21 +30,40 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from transformers import PretrainedConfig +from tensorrt_llm._utils import nvtx_range_debug from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata -from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams -from ..distributed import AllReduce, AllReduceParams, MiniMaxAllReduceRMS +from ..attention_backend.interface import ( + AttentionForwardArgs, + PositionalEmbeddingParams, + RopeParams, +) +from ..attention_backend.sparse.minimax_m3 import ( + MiniMaxM3MsaSparseAttention, + MiniMaxM3SparseRuntimeBackend, + _gather_paged_batched, + _write_main_kv_slots_to_pool, +) +from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams, MiniMaxAllReduceRMS from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import MiniMaxM3MoeRoutingMethod, create_moe from ..modules.gated_mlp import GatedMLP -from ..modules.linear import Linear, TensorParallelMode, copy_weight, load_weight_shard +from ..modules.linear import ( + Linear, + TensorParallelMode, + WeightMode, + WeightsLoadingConfig, + copy_weight, + load_weight_shard, +) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..speculative import SpecMetadata from ..utils import ( ActivationType, AuxStreamType, @@ -52,7 +71,10 @@ get_model_extra_attrs, is_torch_compiling, ) -from .modeling_utils import DecoderModel, DecoderModelForCausalLM, ModelConfig, register_auto_model +from .checkpoints.base_weight_mapper import BaseWeightMapper +from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper +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, @@ -234,11 +256,18 @@ def _build_swiglu_oai_dense_mlp( # MoE composition) carries its own reduction unless ADP collapses # TP to 1. reduce_output = False if is_shared_expert else (not enable_adp) + # SwiGLU-OAI is plain SwiGLU with an alpha gain and (up + 1) offset, so it + # routes through the fused silu_and_mul kernel (one launch, optional fp8 + # epilogue) instead of the eager elementwise fallback. Mirrors the + # routed-expert SwigluBias path. See _minimax_m3_swiglu_oai for the math. return GatedMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, bias=False, - activation=partial(_minimax_m3_swiglu_oai, alpha=swiglu_alpha, limit=swiglu_limit), + activation=torch.nn.functional.silu, + swiglu_alpha=swiglu_alpha, + swiglu_beta=1.0, + swiglu_limit=swiglu_limit, dtype=config.torch_dtype, config=model_config, overridden_tp_size=1 if enable_adp else None, @@ -617,7 +646,7 @@ def minimax_m3_attn_custom_op_inplace( """Run MiniMax-M3 cache and attention work behind a compile boundary.""" attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) num_tokens = attn_metadata.num_tokens - attn_layer._attention_core( + attn_layer._dispatch_attention_backend( q[:num_tokens], k[:num_tokens], v[:num_tokens], @@ -633,12 +662,10 @@ class MiniMaxM3Attention(Attention): Both branches share the same dense GQA scaffolding (``qkv_proj`` + ``o_proj`` + per-head Gemma Q/K norm + partial RoPE). Sparse layers - additionally carry the MiniMax index branch (``index_q_proj``, - ``index_k_proj`` and their per-head norms). The index value/output - branch is omitted because the M3 checkpoint sets - ``sparse_disable_index_value=True`` on every sparse layer; if a - future config variant flips that flag the gate will catch the - unmapped keys. + additionally carry the MiniMax index branch: a fused index_qk_proj (output + [idx_q | idx_k]) plus per-head index norms. The index value/output branch + is omitted because the M3 checkpoint sets sparse_disable_index_value=True on + every sparse layer; the gate catches the unmapped keys if that ever flips. """ def __init__( @@ -648,6 +675,7 @@ def __init__( layer_idx: Optional[int] = None, is_sparse_attention_layer: bool = False, disable_index_value: bool = False, + aux_stream: Optional[torch.cuda.Stream] = None, ): config = model_config.pretrained_config self.pretrained_config = config @@ -679,6 +707,11 @@ def __init__( getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) ) + # Dtype of the q/k/v activations fed into norm+RoPE. This is always the + # model compute dtype (bf16) for every M3 config. KV-cache quantization + # (fp8/fp4) only changes cache storage, not the activation dtype. + self.attn_activation_dtype = config.torch_dtype + # Per-head Gemma RMSNorm — one set of weights shared across heads. self.q_norm = RMSNorm( hidden_size=self.head_dim_value, @@ -693,6 +726,15 @@ def __init__( use_gemma=self.use_gemma_norm, ) + # Aux stream and events to run the per-head q-norm and k-norm + # concurrently on the separate-norm fallback path. The stream is shared + # across layers via aux_stream_dict, matching DeepSeekV3; it falls back + # to a private stream when constructed standalone (for example in tests). + # The bf16 fast path fuses norm and RoPE into one kernel and does not + # use these. + self.aux_stream = aux_stream if aux_stream is not None else torch.cuda.Stream() + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + self.is_sparse_attention_layer = bool(is_sparse_attention_layer) self.disable_index_value = bool(disable_index_value) if self.is_sparse_attention_layer: @@ -705,33 +747,23 @@ 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_proj is **replicated** across TP ranks. The sparse - # forward reshapes idx_q to - # ``[num_tokens, sparse_num_index_heads, sparse_index_dim]``, - # which requires the rank-local idx_q to carry all heads. - index_q_total = self.sparse_num_index_heads * self.sparse_index_dim - self.index_q_proj = Linear( + # Index Q and K are both replicated (no head sharding) and project + # the same hidden_states, so fuse them into one index_qk_proj GEMM + # with output [idx_q | idx_k]. idx_q holds all num_index_heads heads; + # idx_k is a single K per token, broadcast across heads when scoring. + 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, - index_q_total, - bias=False, - dtype=config.torch_dtype, - mapping=model_config.mapping, - tensor_parallel_mode=None, - quant_config=None, - skip_create_weights_in_init=model_config.skip_create_weights_in_init, - ) - # index_k_proj is also replicated across TP ranks and - # outputs ``sparse_index_dim`` channels — a single K per - # token (not per-head), broadcast across index heads when - # scoring blocks. - self.index_k_proj = Linear( - config.hidden_size, - self.sparse_index_dim, + 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``; @@ -762,8 +794,23 @@ def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, """ q_shape = q.shape k_shape = k.shape - q = self.q_norm(q.reshape(-1, self.head_dim_value)).reshape(q_shape) - k = self.k_norm(k.reshape(-1, self.head_dim_value)).reshape(k_shape) + + def _q_norm(): + return self.q_norm(q.reshape(-1, self.head_dim_value)).reshape(q_shape) + + def _k_norm(): + return self.k_norm(k.reshape(-1, self.head_dim_value)).reshape(k_shape) + + # Run the independent q-norm and k-norm concurrently on the aux stream. + # Falls back to sequential execution under torch.compile. + q, k = maybe_execute_in_parallel( + _q_norm, + _k_norm, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) return q, k def apply_index_qk_norm( @@ -790,32 +837,104 @@ def apply_index_qk_norm( ) idx_q_shape = idx_q.shape idx_k_shape = idx_k.shape - idx_q = self.index_q_norm(idx_q.reshape(-1, self.sparse_index_dim)).reshape(idx_q_shape) - idx_k = self.index_k_norm(idx_k.reshape(-1, self.sparse_index_dim)).reshape(idx_k_shape) + + def _idx_q_norm(): + return self.index_q_norm(idx_q.reshape(-1, self.sparse_index_dim)).reshape(idx_q_shape) + + def _idx_k_norm(): + return self.index_k_norm(idx_k.reshape(-1, self.sparse_index_dim)).reshape(idx_k_shape) + + # Run the independent index q-norm and k-norm concurrently on the aux + # stream. Falls back to sequential execution under torch.compile. + idx_q, idx_k = maybe_execute_in_parallel( + _idx_q_norm, + _idx_k_norm, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) return idx_q, idx_k - def apply_rope( + def _fused_qk_norm_rope( self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - position_ids: torch.Tensor, - ): - """Run per-head QK norm before partial RoPE. + qkv: torch.Tensor, + position_ids: Optional[torch.Tensor], + *, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + q_norm: RMSNorm, + k_norm: RMSNorm, + ) -> Optional[torch.Tensor]: + """Fuse per-head Gemma RMSNorm and partial RoPE into one kernel. + + Runs torch.ops.trtllm.fused_qk_norm_rope in place over the fused qkv + (Q heads then K heads, optional V heads left untouched) and returns the + mutated tensor. The kernel norms the full head_dim and rotates only the + first rotary_dim channels read from RopeParams.dim, which matches M3's + whole-head norm with front partial RoPE, and applies the same Gemma + (1 + weight) scaling as apply_qk_norm. + + Returns None, leaving qkv untouched, when the fused path does not apply + so callers fall back to separate norm and RoPE. This happens when + activations are not bf16 (the kernel is bf16-only), when RoPE has no + position_ids, or when no partial-RoPE rotary_emb exists. + """ + if position_ids is None or qkv.dtype != torch.bfloat16: + return None + if ( + self.rotary_emb is None + or self.pos_embd_params is None + or self.pos_embd_params.rope is None + ): + return None + + # Partial-RoPE dim comes from RopeParams (M3 rotates 64 of 128). + rotary_dim = int(self.pos_embd_params.rope.dim) + # The kernel assumes a contiguous [num_tokens, total_heads * head_dim]. + qkv = qkv.contiguous() + torch.ops.trtllm.fused_qk_norm_rope( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + q_norm.variance_epsilon, + q_norm.weight, + k_norm.weight, + self.pos_embd_params.rope.theta, + self.pos_embd_params.is_neox, + position_ids.reshape(-1).contiguous().to(torch.int32), + 1.0, # factor: no YARN (M3 has no rope_scaling) + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + self.use_gemma_norm, # use_gemma + False, # use_mrope + 0, # mrope_section1 + 0, # mrope_section2 + ) + return qkv - The base ``Attention.apply_rope`` consumes split q/k/v. We split, - apply per-head QK norm, then defer to the base partial-RoPE - implementation (driven by ``RopeParams.dim < head_dim``). + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: + """Whether the fused path must run rather than fall back. + + The fused kernel runs when attention activations are bf16 and RoPE has + position_ids. Every M3 config keeps bf16 attention activations, so a + fallback is a regression. Guards the forward asserts. """ - q, k, v = self.split_qkv(q, k, v) - q, k = self.apply_qk_norm(q, k) - return super().apply_rope(q, k, v, position_ids) + return self.attn_activation_dtype == torch.bfloat16 and position_ids is not None def forward( self, position_ids: Optional[torch.IntTensor] = None, hidden_states: Optional[torch.Tensor] = None, attn_metadata: Optional[AttentionMetadata] = None, + all_reduce_params: Optional[AllReduceParams] = None, **kwargs, ): """Dispatch sparse layers to the MiniMax-M3 sparse algorithm. @@ -835,14 +954,27 @@ def forward( side index-K buffer. """ if not self.is_sparse_attention_layer: - return self._dense_forward(position_ids, hidden_states, attn_metadata, **kwargs) - return self._sparse_forward(position_ids, hidden_states, attn_metadata, **kwargs) + return self._dense_forward( + position_ids, + hidden_states, + attn_metadata, + all_reduce_params=all_reduce_params, + **kwargs, + ) + return self._sparse_forward( + position_ids, + hidden_states, + attn_metadata, + all_reduce_params=all_reduce_params, + **kwargs, + ) def _dense_forward( self, position_ids: torch.IntTensor, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, **kwargs, ) -> torch.Tensor: """Dense MiniMax-M3 attention for layers 0-2. @@ -858,9 +990,10 @@ def _dense_forward( Steps: 1. Project Q/K/V via fused ``qkv_proj``. - 2. Apply per-head Gemma RMSNorm to Q/K (same as - :meth:`_sparse_forward` step 2 minus the index branch). - 3. Apply partial RoPE. + 2-3. Apply per-head Gemma RMSNorm and partial RoPE to Q/K. The bf16 + fast path fuses both into one fused_qk_norm_rope kernel over the + fused qkv; the non-bf16-activation fallback runs apply_qk_norm then + rotary_emb separately. 4. Pull the paged main K/V cache from the M3 cache manager. 5. Read the pre-built :class:`MiniMaxM3SparseAttentionMetadata` from ``attn_metadata.minimax_m3``. Production code paths @@ -885,24 +1018,49 @@ def _dense_forward( "attn_metadata; received None." ) - # 1. Projections (no index branch). + # Projections (no index branch). qkv = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - - # 2. Per-head Gemma RMSNorm on Q/K (no index norm). - q, k = self.apply_qk_norm(q, k) - # 3. Partial RoPE on Q/K (no index branch). - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) + # Per-head Gemma RMSNorm and partial RoPE on Q/K. The bf16 fast + # path fuses both into one kernel over the fused qkv; otherwise fall + # back to separate norm and RoPE. + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + ) + if fused_qkv is not None: + q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # Match the fallback contiguity; V stays a column-slice view. + q, k = q.contiguous(), k.contiguous() + else: + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) # Keep token-wise projections and the output projection visible to # torch.compile. Only the metadata/cache-dependent attention core is # hidden behind the inplace custom op. o = self._forward_attention_core(q, k, v, None, None, attn_metadata) - return self.o_proj(o) + # all_reduce_params lets the decoder defer the o_proj output AllReduce so + # it can be fused with post_attention_layernorm (RESIDUAL_RMS_NORM). + # Passing None preserves the standalone o_proj reduction used by the + # single-GPU, attention-DP, and fusion-disabled paths. + return self.o_proj(o, all_reduce_params=all_reduce_params) - def _dense_attention_core( + def _sdpa_dense_attention_core( self, q: torch.Tensor, k: torch.Tensor, @@ -910,12 +1068,14 @@ def _dense_attention_core( attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: - """Run dense cache updates and attention into ``output``.""" - from ..attention_backend.sparse.minimax_m3 import ( - _gather_paged_batched, - _write_main_kv_slots_to_pool, - ) + """Run dense cache updates and attention into ``output`` (SDPA path). + This is the non-MSA dense reference path: it runs standard GQA via + ``torch.nn.functional.scaled_dot_product_attention`` (memory-efficient + / math backends). It is selected only when ``self.attn`` is **not** a + :class:`MiniMaxM3MsaSparseAttention`; the MSA backend handles dense + layers itself inside :meth:`_msa_attention_core`. + """ kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) if kv_cache_manager is None: raise RuntimeError( @@ -967,7 +1127,17 @@ def _dense_attention_core( # 7. Gather padded K/V for every batch row and run dense GQA. batch = int(m3_meta.slot_ids.shape[0]) - max_k = int(m3_meta.max_seqlen_k) + # Graph capture bakes the gather/mask width, and max_seqlen_k is a + # per-step host bound later replays can outgrow. Bake + # min(page-table width, engine max_seq_len): the raw table width + # alone is inflated far past max_seq_len by the KV-estimation pass + # and would OOM the [batch, max_k, heads] gather. The seq_lens mask + # below invalidates the slack. + if attn_metadata.is_cuda_graph: + capacity = int(m3_meta.req_to_token.shape[1]) + max_k = min(capacity, int(attn_metadata.max_seq_len or capacity)) + else: + max_k = int(m3_meta.max_seqlen_k) if max_k <= 0: max_k = 1 # ``_gather_paged_batched`` decomposes the flat slot id into @@ -994,15 +1164,11 @@ def _dense_attention_core( f"by num_key_value_heads ({self.num_key_value_heads})" ) group = self.num_heads // max(self.num_key_value_heads, 1) - if group > 1: - k_padded = k_padded.repeat_interleave(group, dim=2) - v_padded = v_padded.repeat_interleave(group, dim=2) # Build the per-query attention mask that masks out padded KV # positions beyond each sequence's true ``seq_lens`` and (for # prefill) preserves causality. ``q_positions`` from the prefill - # metadata names each Q token's K-side position; for decode - # there is one Q token per request at position ``seq_lens - 1``. + # metadata names each Q token's K-side position. # The metadata tensors are produced by # :meth:`MiniMaxM3AttentionMetadata.prepare` on the cache # device, so ``.to(dtype=torch.long)`` is a same-device dtype @@ -1011,6 +1177,9 @@ def _dense_attention_core( kv_positions = torch.arange(max_k, device=q.device).unsqueeze(0) # [1, max_k] if m3_meta.is_prefill: + if group > 1: + k_padded = k_padded.repeat_interleave(group, dim=2) + v_padded = v_padded.repeat_interleave(group, dim=2) # Prefill: build [total_q, max_k] mask using q_positions / q_batch_row. # Prefill never runs inside the CUDA-graph capture window # (capture is decode-only), so the per-batch Python loop and @@ -1051,35 +1220,42 @@ def _dense_attention_core( ) # [1, H, q, d] output_view[start:end].copy_(out_b.squeeze(0).transpose(0, 1)) else: - # Decode: one Q token per request at position seq_lens - 1. - # Every input tensor here is already on q.device (set up by - # prepare()), so SDPA captures cleanly. - valid = kv_positions < seq_lens_dev.unsqueeze(-1) # [batch, max_k] - q_b = q_view.unsqueeze(1).transpose(1, 2) # [batch, H, 1, d] - k_b = k_padded.transpose(1, 2) # [batch, H, k, d] - v_b = v_padded.transpose(1, 2) # [batch, H, k, d] - mask_b = valid.unsqueeze(1).unsqueeze(1) # [batch, 1, 1, k] + # Decode: qo_len query tokens per request; token t of request b + # attends seq_lens[b] - qo_len + t + 1 positions (the causal + # ladder; qo_len=1 is the classic one-token mask). Every input + # tensor here is already on q.device (set up by prepare()), so + # SDPA captures cleanly. + qo_len = int(m3_meta.decode_qo_len) + ladder = torch.arange(1 - qo_len, 1, device=q.device, dtype=torch.long) + # eff[b, t] = attendable position count for token t of row b. + eff = seq_lens_dev.unsqueeze(-1) + ladder # [batch, qo] + valid = kv_positions.unsqueeze(1) < eff.unsqueeze(-1) # [batch, qo, max_k] + q_b = q_view.view(batch, qo_len, self.num_heads, self.head_dim).transpose( + 1, 2 + ) # [batch, H, qo, d] + mask_b = valid.unsqueeze(1) # [batch, 1, qo, k] + # Expand K/V one KV head at a time: the all-heads transient is + # O(batch * max_k * num_heads) inside the CUDA-graph pool and + # exceeds the pool budget under attention DP (unsharded heads); + # with TP-sharded KV heads this is one iteration. + out_b = q.new_empty(batch, self.num_heads, qo_len, self.head_dim) with sdpa_kernel(_DENSE_SDPA_BACKENDS): - out_b = torch.nn.functional.scaled_dot_product_attention( - q_b.to(q.dtype), - k_b.to(q.dtype), - v_b.to(q.dtype), - attn_mask=mask_b, - dropout_p=0.0, - is_causal=False, - ) # [batch, H, 1, d] - # Drop the singleton Q-length axis and write the resulting - # ``[batch, num_heads, head_dim]`` tensor into the final buffer. - # The prior ``.transpose(1, 2).reshape(batch, H, d)`` pattern - # was wrong: with ``H != head_dim`` (M3 TP=8 has H=8, d=128) - # the non-contiguous transpose forces ``reshape`` to copy the - # data in C-order under its current ``[batch, d, H]`` shape, - # then reinterpret as ``[batch, H, d]`` — which scrambles - # ``(head, head_dim)`` ordering and feeds permuted activations - # into ``o_proj``. Prefill is unaffected because its - # ``transpose(0, 1)`` runs between q-len and num_heads axes - # which the per-batch loop already laid out correctly. - output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2)) + for h in range(max(self.num_key_value_heads, 1)): + qh = slice(h * group, (h + 1) * group) + k_h = k_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) + v_h = v_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) + out_b[:, qh] = torch.nn.functional.scaled_dot_product_attention( + q_b[:, qh].to(q.dtype), + k_h.transpose(1, 2).to(q.dtype), + v_h.transpose(1, 2).to(q.dtype), + attn_mask=mask_b, + dropout_p=0.0, + is_causal=False, + ) # [batch, group, qo, d] + # Copy through a token-major [batch, qo, H, dh] view rather + # than transpose(1, 2).reshape, which (with H != head_dim) + # copies in C-order and scrambles (head, head_dim) into o_proj. + output.view(batch, qo_len, self.num_heads, self.head_dim).copy_(out_b.transpose(1, 2)) return output @@ -1104,10 +1280,10 @@ def _forward_attention_core( output, ) else: - self._attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + self._dispatch_attention_backend(q, k, v, idx_q, idx_k, attn_metadata, output) return output - def _attention_core( + def _dispatch_attention_backend( self, q: torch.Tensor, k: torch.Tensor, @@ -1117,17 +1293,59 @@ def _attention_core( attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: + """Route the attention core to the configured backend. + + ``self.attn`` is either a :class:`MiniMaxM3MsaSparseAttention` (MSA + backend, which handles both dense and sparse layers) or a + :class:`MiniMaxM3SparseRuntimeBackend` (Triton backend). This method + only selects the backend-specific core: + + * MSA → :meth:`_msa_attention_core` + * Triton sparse → :meth:`_triton_sparse_attention_core` + * SDPA dense → :meth:`_sdpa_dense_attention_core` + """ + if isinstance(self.attn, MiniMaxM3MsaSparseAttention): + return self._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) if self.is_sparse_attention_layer: assert idx_q is not None and idx_k is not None - return self._sparse_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + return self._triton_sparse_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) assert idx_q is None and idx_k is None - return self._dense_attention_core(q, k, v, attn_metadata, output) + return self._sdpa_dense_attention_core(q, k, v, attn_metadata, output) + + def _msa_attention_core( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + idx_q: Optional[torch.Tensor], + idx_k: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + output: torch.Tensor, + ) -> torch.Tensor: + """Run the MSA backend (:class:`MiniMaxM3MsaSparseAttention`). + + The MSA backend runs the sparse GQA or dense paged GQA through its + inherited FMHA forward; this layer selects the top-k blocks (sparse + only) and builds the ``forward_args`` the FMHA reads. + """ + if self.is_sparse_attention_layer: + assert idx_q is not None and idx_k is not None + # Publish the selected blocks so the FMHA runs the sparse path. + kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata) + forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) + else: + 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) + return output def _sparse_forward( self, position_ids: torch.IntTensor, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, **kwargs, ) -> torch.Tensor: """Run a MiniMax-M3 sparse attention forward end-to-end. @@ -1135,13 +1353,15 @@ def _sparse_forward( Steps: 1. Project ``hidden_states`` to Q/K/V (fused ``qkv_proj``) plus index Q (per-head) and index K (single replicated). - 2. Apply per-head Gemma RMSNorm to both branches. - 3. Apply partial RoPE (``rotary_dim`` channels of ``head_dim``) - to both branches. + 2-3. Apply per-head Gemma RMSNorm and partial RoPE to the main and + index branches. The bf16 fast path fuses each branch into one + fused_qk_norm_rope kernel (the index branch passes num_heads_v=0 + and norms only the concatenated idx_q/idx_k); the non-bf16-activation + fallback runs the norm helpers then rotary_emb. 4. Pull paged main K/V cache (reshaped to flat-slot view) and paged side index-K cache from the :class:`MiniMaxM3KVCacheManagerV2`. - 5. Build a :class:`MiniMaxM3SparseAttentionMetadata` from the + 5. Build a :class:`MiniMaxM3TritonSparseAttentionMetadata` from the standard :class:`AttentionMetadata` (using ``request_ids`` + ``seq_lens`` + ``num_cached_tokens_per_seq``). 6. Write the new token's K/V/idx_K to the slots named by the @@ -1154,7 +1374,7 @@ def _sparse_forward( Production callers (the LLM API path) drive :meth:`MiniMaxM3AttentionMetadata.prepare` outside any CUDA-graph capture window; that method attaches a pre-built - :class:`MiniMaxM3SparseAttentionMetadata` and an + :class:`MiniMaxM3TritonSparseAttentionMetadata` and an ``out_cache_loc`` tensor as ``attn_metadata.minimax_m3``. Test callers attach the same dict manually. This forward path always reads the pre-built attachment and never builds metadata @@ -1168,27 +1388,86 @@ def _sparse_forward( f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires " "attn_metadata; received None." ) - # 1. Projections. - qkv = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - idx_q = self.index_q_proj(hidden_states) - idx_k = self.index_k_proj(hidden_states) - # 2. Per-head Gemma RMSNorm on both branches. - q, k = self.apply_qk_norm(q, k) - idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) - - # 3. Partial RoPE on both branches. The base ``Attention`` - # constructor created ``self.rotary_emb`` for the configured - # partial ``rotary_dim`` because ``rope_fusion=False``. - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + # Project, per-head Gemma RMSNorm, and partial RoPE for the main and + # index branches. Each branch fuses norm and RoPE into one bf16 kernel, + # falling back to separate norm and RoPE otherwise. The branches read + # only hidden_states and write disjoint outputs, so each runs its + # projection, norm, and RoPE concurrently on the aux stream when + # multi-stream is enabled, joining before the attention core. + def _main_norm_rope(): + qkv = self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + ) + if fused_qkv is not None: + q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + return q.contiguous(), k.contiguous(), v + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + return q, k, v + + def _index_norm_rope(): + idx_qk = self.index_qk_proj(hidden_states) + fused_idx = self._fused_qk_norm_rope( + idx_qk, + position_ids, + num_heads_q=self.sparse_num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=self.sparse_index_dim, + q_norm=self.index_q_norm, + k_norm=self.index_k_norm, + ) + if fused_idx is not None: + idx_q, idx_k = fused_idx.split( + [self.sparse_num_index_heads * self.sparse_index_dim, self.sparse_index_dim], + dim=-1, + ) + return idx_q.contiguous(), idx_k.contiguous() + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, index_dim=" + f"{self.sparse_index_dim}) but fell back to the separate path; idx " + f"dtype is {idx_qk.dtype} (expected {self.attn_activation_dtype})." + ) + idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) + if self.rotary_emb is not None and position_ids is not None: + idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + return idx_q, idx_k + + (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( + _main_norm_rope, + _index_norm_rope, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) - return self.o_proj(o) + # all_reduce_params lets the decoder defer the o_proj output AllReduce so + # it can be fused with post_attention_layernorm (RESIDUAL_RMS_NORM). + # Passing None preserves the standalone o_proj reduction used by the + # single-GPU, attention-DP, and fusion-disabled paths. + return self.o_proj(o, all_reduce_params=all_reduce_params) - def _sparse_attention_core( + def _triton_sparse_attention_core( self, q: torch.Tensor, k: torch.Tensor, @@ -1198,7 +1477,13 @@ def _sparse_attention_core( attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: - """Run sparse cache updates and attention into ``output``.""" + """Run sparse cache updates and attention into ``output`` (Triton path). + + This is the non-MSA sparse path: it dispatches to + :class:`MiniMaxM3SparseRuntimeBackend` (the Triton sparse runtime). It + is selected only when ``self.attn`` is that backend; the MSA backend + handles sparse layers itself inside :meth:`_msa_attention_core`. + """ kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) if kv_cache_manager is None: raise RuntimeError( @@ -1268,10 +1553,7 @@ def _sparse_attention_core( # registers :class:`MiniMaxM3SparseRuntimeBackend` as # ``self.attn``; any other backend on a sparse layer is a # configuration error. - from ..attention_backend.sparse.minimax_m3 import get_minimax_m3_attention_backend_cls - - m3_backend_cls = get_minimax_m3_attention_backend_cls() - if not isinstance(self.attn, m3_backend_cls): + if not isinstance(self.attn, MiniMaxM3SparseRuntimeBackend): raise RuntimeError( f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires " f"self.attn to be a MiniMaxM3SparseRuntimeBackend, got " @@ -1312,6 +1594,7 @@ def __init__( self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.mapping = model_config.mapping + self.enable_attention_dp = self.mapping.enable_attention_dp _, sparse_layer_ids = get_sparse_layer_ids(config) disable_index_value_ids = set(get_sparse_disable_index_value_layer_ids(config)) @@ -1323,6 +1606,7 @@ def __init__( layer_idx=layer_idx, is_sparse_attention_layer=is_sparse, disable_index_value=disable_index_value, + aux_stream=aux_stream_dict[AuxStreamType.Attention], ) _, moe_layer_ids = get_moe_layer_ids(config) @@ -1342,45 +1626,206 @@ def __init__( ) self.block_sparse_moe = None + # Layer-boundary RMSNorms are plain (non-Gemma) norms so they can drive + # the fused AllReduce+residual+RMSNorm epilogue + # (AllReduceFusionOp.RESIDUAL_RMS_NORM), whose kernel applies a plain + # weight * x scaling with no Gemma (1 + weight) offset. When the + # checkpoint stores Gemma norms (use_gemma_norm=True), the loader folds + # (1 + weight) into the stored weight at load time (see + # _fold_gemma_boundary_norm_weights), so the runtime norm is numerically + # identical to the original Gemma norm on every path. The per-head + # q/k/index norms keep use_gemma because they are consumed by the + # separate fused_qk_norm_rope kernel, which handles Gemma directly. self.input_layernorm = RMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype, - use_gemma=bool(getattr(config, "use_gemma_norm", False)), + use_gemma=False, ) self.post_attention_layernorm = RMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype, - use_gemma=bool(getattr(config, "use_gemma_norm", False)), + use_gemma=False, ) + # DeepSeek-V3-style layer-boundary AllReduce fusion. Each layer folds + # the attention o_proj output AllReduce into post_attention_layernorm + # (PRE fusion) and the MoE/MLP output AllReduce into the next layer's + # input_layernorm (POST fusion, wired via next_layer_layernorm in + # setup_aliases). Fusion is only meaningful when there is a real + # cross-rank reduction to fold, i.e. TP>1 and not attention-DP (each DP + # rank owns independent tokens, so no attention/MoE AllReduce happens + # there). An env override matches the DeepSeek-V3 escape hatch. + self.enable_fusion = os.environ.get("TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED", "0") == "0" + self.enable_fusion &= (not self.enable_attention_dp) and self.mapping.tp_size > 1 + self.pre_feed_forward_fusion = self.enable_fusion + self.post_feed_forward_fusion = self.enable_fusion + + self.allreduce = None + if not self.enable_attention_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, + ) + + # Wired by MiniMaxM3ForCausalLM.setup_aliases after weight load to the + # next layer's input_layernorm (or the final model norm for the last + # layer). None disables POST fusion and boundary-norm folding. + self.next_layer_layernorm: Optional[RMSNorm] = None + def forward( self, position_ids: torch.IntTensor, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: + # NVTX markers below are emitted only when TLLM_NVTX_DEBUG=1 (or + # TLLM_LLMAPI_ENABLE_NVTX=1) is set; otherwise they are no-ops. + attn_kind = "sparse_attn" if self.self_attn.is_sparse_attention_layer else "dense_attn" + + # Layer-0 prologue only. For every subsequent layer the input_layernorm + # (an add+RMSNorm at the layer boundary) was already applied by the + # previous layer as its next_layer_layernorm, so residual is not None + # here and this block is skipped (matches DeepSeek-V3). if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - - hidden_states = self.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - **kwargs, + with nvtx_range_debug(f"layer{self.layer_idx}.input_layernorm"): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # When PRE fusion is active the attention defers its o_proj AllReduce so + # it can be fused into post_attention_layernorm below; otherwise the + # o_proj reduces as usual (all_reduce_params=None preserves the + # single-GPU, attention-DP, and fusion-disabled behavior exactly). + attn_all_reduce_params = ( + AllReduceParams(enable_allreduce=False) if self.pre_feed_forward_fusion else None ) + with nvtx_range_debug(f"layer{self.layer_idx}.{attn_kind}"): + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=attn_all_reduce_params, + **kwargs, + ) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) if self.block_sparse_moe is not None: - hidden_states = self.block_sparse_moe(hidden_states, attn_metadata) + hidden_states, residual = self.forward_MoE(hidden_states, attn_metadata, residual) else: - hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.forward_mlp(hidden_states, residual) + + # hidden_states is fully TP-reduced at layer exit (no cross-layer + # allreduce+norm fusion). + 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) + + return hidden_states, residual + + def _apply_pre_feed_forward_norm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """AllReduce(+residual+RMSNorm) between attention and the feed-forward. + + On the PRE-fusion path the deferred attention o_proj AllReduce is fused + with post_attention_layernorm into one kernel; otherwise it is the plain + add+RMSNorm and the attention already reduced its own output. + """ + if self.pre_feed_forward_fusion: + return self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + return self.post_attention_layernorm(hidden_states, residual) + + def _apply_next_layer_layernorm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply the next layer's input_layernorm at the layer boundary. + + On the POST-fusion path the deferred feed-forward output AllReduce is + fused with next_layer_layernorm (the next layer's input_layernorm, or + the final model norm for the last layer). Off the fusion path it is the + plain add+RMSNorm. When next_layer_layernorm has not been wired (e.g. a + standalone unit test that never ran setup_aliases) the + (hidden_states, residual) pair is returned unchanged so the model can + apply the final norm itself. + """ + if self.next_layer_layernorm is None: + return hidden_states, residual + if self.post_feed_forward_fusion: + return self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + return self.next_layer_layernorm(hidden_states, residual) + + def _feed_forward_all_reduce_params(self) -> Optional[AllReduceParams]: + """AllReduce params handed to the MoE or dense-MLP output projection. + + Disables the module's internal output AllReduce when POST fusion will + fold it into next_layer_layernorm; otherwise None preserves the module's + own reduction (single-GPU, attention-DP, fusion-disabled). + """ + if self.post_feed_forward_fusion: + return AllReduceParams(enable_allreduce=False) + return None + + def forward_MoE( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + with nvtx_range_debug(f"layer{self.layer_idx}.post_attention_layernorm"): + hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) + + with nvtx_range_debug(f"layer{self.layer_idx}.moe"): + hidden_states = self.block_sparse_moe( + hidden_states, + attn_metadata, + final_all_reduce_params=self._feed_forward_all_reduce_params(), + ) + + with nvtx_range_debug(f"layer{self.layer_idx}.next_layer_layernorm"): + hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) + return hidden_states, residual + + def forward_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + with nvtx_range_debug(f"layer{self.layer_idx}.post_attention_layernorm"): + hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) + + with nvtx_range_debug(f"layer{self.layer_idx}.mlp"): + hidden_states = self.mlp( + hidden_states, + final_all_reduce_params=self._feed_forward_all_reduce_params(), + ) + + with nvtx_range_debug(f"layer{self.layer_idx}.next_layer_layernorm"): + hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual) return hidden_states, residual @@ -1397,10 +1842,12 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): model_config.pretrained_config.torch_dtype = torch.bfloat16 config = model_config.pretrained_config self.vocab_size = config.vocab_size - # Two aux streams: one for MoE shared/routed parallel execution, - # one for MoE chunking overlap inside the fused MoE kernel. - # Matches the DeepSeekV3 convention. + # Aux streams shared across layers, matching the DeepSeekV3 convention: + # one for the per-head q/k norm overlap in attention, one for MoE + # shared/routed parallel execution, and one for MoE chunking overlap + # inside the fused MoE kernel. self.aux_stream_dict = { + AuxStreamType.Attention: torch.cuda.Stream(), AuxStreamType.MoeShared: torch.cuda.Stream(), AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), } @@ -1418,11 +1865,17 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): for layer_idx in range(config.num_hidden_layers) ] ) + # Final norm is a plain (non-Gemma) RMSNorm for the same reason as the + # layer-boundary norms (see MiniMaxM3DecoderLayer.__init__): it doubles + # as the last layer's next_layer_layernorm, so the last MoE/MLP output + # AllReduce folds into it via RESIDUAL_RMS_NORM. The Gemma (1 + weight) + # offset is folded into the stored weight at load time (see + # _fold_gemma_boundary_norm_weights). self.norm = RMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype, - use_gemma=bool(getattr(config, "use_gemma_norm", False)), + use_gemma=False, ) def forward( @@ -1431,6 +1884,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): @@ -1441,54 +1895,147 @@ def forward( hidden_states = inputs_embeds residual = None - for decoder_layer in self.layers: - hidden_states, residual = decoder_layer( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=residual, - ) + for layer_idx, decoder_layer in enumerate(self.layers): + # Per-layer NVTX range (layer0, layer1, ...). Emitted only when + # TLLM_NVTX_DEBUG=1 (or TLLM_LLMAPI_ENABLE_NVTX=1) is set. + with nvtx_range_debug(f"MiniMaxM3.layer{layer_idx}"): + hidden_states, residual = decoder_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + ) - hidden_states, _ = self.norm(hidden_states, residual) + # When setup_aliases has chained the final norm into the last decoder + # layer (next_layer_layernorm = self.norm), the last layer's boundary + # step already applied it (fused or plain), so hidden_states is normed + # and this is skipped. The fallback covers paths that never ran + # setup_aliases (e.g. standalone unit tests): there the last layer + # returns the unnormed (hidden_states, residual) pair and the final + # add+RMSNorm is applied here. + if self.layers[-1].next_layer_layernorm is None: + hidden_states, _ = self.norm(hidden_states, residual) return hidden_states -# HF MiniMax-M3 stores the routed score-correction bias one level above -# the router weight (``block_sparse_moe.e_score_correction_bias``, -# sibling of ``block_sparse_moe.gate.weight``). The TRT-LLM module tree -# binds it to :class:`MiniMaxM3Gate`, so the generic loader expects to -# see it at ``block_sparse_moe.gate.e_score_correction_bias``. The -# regex below moves the key into the gate's prefix before the loader -# dispatches; this lets ``mark_consumed("...gate")`` cleanly remove -# both the weight and the bias together without disturbing the sibling -# ``block_sparse_moe.experts.*`` backend subtree. -_M3_GATE_BIAS_RENAME_MAP = { - r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": (r"\1.gate.e_score_correction_bias"), -} +def _load_index_qk_proj_weights(model: nn.Module, weights) -> None: + """Fuse checkpoint index_q_proj and index_k_proj into index_qk_proj. + + The shared weight loader only auto-fuses qkv_proj and gate_up_proj by + name, so walk each fused index module and load the sibling checkpoint + tensors through the FUSED_GATE_UP_LINEAR row-cat path. Sources are + marked consumed so the generic pass does not treat them as unused. + """ + for name, module in model.named_modules(): + if name.split(".")[-1] != "index_qk_proj": + continue + parent = name.rsplit(".", 1)[0] + q_weights = filter_weights(f"{parent}.index_q_proj", weights) + k_weights = filter_weights(f"{parent}.index_k_proj", weights) + # Missing sources make Linear.load_weights assert rather than leave + # the fused module silently uninitialized. + module.load_weights(weights=[q_weights, k_weights]) + if hasattr(weights, "mark_consumed"): + weights.mark_consumed(f"{parent}.index_q_proj") + weights.mark_consumed(f"{parent}.index_k_proj") + else: + for key in list(weights.keys()): + if key.startswith(f"{parent}.index_q_proj.") or key.startswith( + f"{parent}.index_k_proj." + ): + del weights[key] + + +# 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 +# the norms that drive the DeepSeek-V3-style fused AllReduce+residual+RMSNorm +# epilogue, whose kernel has no Gemma offset. The per-head q/k/index norms are +# intentionally excluded: they feed the separate fused_qk_norm_rope kernel, +# which handles Gemma directly and stays use_gemma=True. +_M3_BOUNDARY_NORM_SUFFIXES = ( + ".input_layernorm.weight", + ".post_attention_layernorm.weight", +) +_M3_FINAL_NORM_KEY = "model.norm.weight" + + +def _fold_gemma_boundary_norm_weights(weights): + """Fold Gemma (1 + weight) into the layer-boundary RMSNorm weights. + + MiniMax-M3 stores every RMSNorm as a Gemma norm (use_gemma_norm=True), which + computes (1 + weight) * x. The layer-boundary norms are constructed as plain + norms (use_gemma=False, weight * x) so they can drive the fused + AllReduce+RMSNorm kernels, so their stored weights must be pre-incremented by + 1.0. This is a numerically exact, load-time-only rewrite; the resulting norm + is identical to the original Gemma norm on every path. + + Only the decoder input_layernorm / post_attention_layernorm and the final + model.norm are touched. A no-op for keys that are absent (partial load) so it + is safe to call unconditionally, but it must only run when the checkpoint + actually uses Gemma norms (guarded by the caller). + """ + for key in list(weights.keys()): + if key.endswith(_M3_BOUNDARY_NORM_SUFFIXES) or key == _M3_FINAL_NORM_KEY: + w = weights[key] + w = w[:] if hasattr(w, "__getitem__") else w + weights[key] = w + 1.0 + return weights @register_auto_model("MiniMaxM3SparseForCausalLM") -class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]): +class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" 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, + weights: Dict, + weight_mapper: Optional[BaseWeightMapper] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, + ) -> None: + # Fuse index_q/index_k into each index_qk_proj module (also covers + # the VL path, which routes text weights through here). + _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 + # AllReduce+residual+RMSNorm epilogue. Only when the checkpoint actually + # stores Gemma norms; otherwise the boundary norms are already plain. + if bool(getattr(self.config, "use_gemma_norm", False)): + weights = _fold_gemma_boundary_norm_weights(weights) + if weight_mapper is None: + weight_mapper = MiniMaxM3HfWeightMapper() + weight_mapper.init_model_and_config(self, self.model_config) + merged_params_map = {**MINIMAX_M3_PARAMS_MAP, **(params_map or {})} + super().load_weights( + weights=weights, + weight_mapper=weight_mapper, + params_map=merged_params_map, + allow_partial_loading=allow_partial_loading, ) - def load_weights(self, weights, *args, **kwargs): - # Merge the M3-specific gate-bias rename into any caller- - # supplied ``params_map`` so the VL wrapper and any downstream - # tooling that already passes one keep working. - params_map = kwargs.pop("params_map", None) or {} - merged = {**_M3_GATE_BIAS_RENAME_MAP, **params_map} - return super().load_weights(weights, *args, params_map=merged, **kwargs) + def setup_aliases(self) -> None: + """Chain each decoder layer's next_layer_layernorm for POST fusion. + + Wired after weight load (the generic loader skips next_layer_layernorm + aliases). Each layer's MoE/MLP output AllReduce is fused into the next + layer's input_layernorm; the last layer chains the final model norm so + its output AllReduce folds the final normalization too. + """ + layers = self.model.layers + num_layers = len(layers) + for idx, layer in enumerate(layers): + if idx == num_layers - 1: + layer.next_layer_layernorm = self.model.norm + else: + layer.next_layer_layernorm = layers[idx + 1].input_layernorm def _strip_language_model_prefix( @@ -1624,7 +2171,13 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): self.last_loaded_vision_keys = [] self.last_missing_vision_keys = [] - def load_weights(self, weights, *args, **kwargs): + def load_weights( + self, + weights: Dict, + weight_mapper: Optional[BaseWeightMapper] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, + ) -> None: text_cfg = self.config if is_minimax_m3_vl_config(text_cfg): text_cfg = get_text_config(text_cfg) @@ -1647,7 +2200,12 @@ def load_weights(self, weights, *args, **kwargs): self.last_loaded_vision_keys = loaded self.last_missing_vision_keys = missing - return super().load_weights(text_weights, *args, **kwargs) + super().load_weights( + weights=text_weights, + weight_mapper=weight_mapper, + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) def forward( self, diff --git a/tensorrt_llm/_torch/models/modeling_step3p7.py b/tensorrt_llm/_torch/models/modeling_step3p7.py index 0a17e7fd0227..e9e8378c14fa 100644 --- a/tensorrt_llm/_torch/models/modeling_step3p7.py +++ b/tensorrt_llm/_torch/models/modeling_step3p7.py @@ -377,8 +377,8 @@ class Step3p7MoeRoutingMethod(MiniMaxM2MoeRoutingMethod): Inherits from ``MiniMaxM2MoeRoutingMethod`` so the TRTLLMGen ``_extract_routing_params`` helper recognises us via ``isinstance`` and - feeds the bias pointer to the kernel. The MiniMax2 C++ routing path - hard-codes ``routeScale = 1.0f`` (see ``runner.cu``), so + feeds the bias pointer to the kernel. The generic MiniMax2 metadata does + not supply a route scale and therefore defaults to ``1.0f``, so ``routed_scaling_factor`` is applied to the MoE output in ``Step3p7MoE.forward`` instead of inside the kernel. """ @@ -903,7 +903,7 @@ def forward( all_rank_num_tokens=attn_metadata.all_rank_num_tokens, use_dp_padding=False, ) - # TRTLLMGen MiniMax2 kernel hard-codes routeScale=1.0, so apply + # Step3p7 uses the generic MiniMax2 metadata with routeScale=1.0, so apply # ``routed_scaling_factor`` to the MoE output here (mathematically # equivalent to scaling each topk weight). if self.routed_scaling_factor != 1.0: diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index e6b192a3f5ed..6e0f9aa90e32 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -47,7 +47,7 @@ # isort: on from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, DeepSeekV4MoeRoutingMethod, DefaultMoeRoutingMethod, - MiniMaxM2MoeRoutingMethod) + MiniMaxM2MoeRoutingMethod, MiniMaxM3MoeRoutingMethod) @dataclass @@ -649,6 +649,14 @@ def _extract_routing_params(self) -> RoutingParams: routed_scaling_factor=self.routing_method.routing_impl. routed_scaling_factor, ) + elif isinstance(self.routing_method, MiniMaxM3MoeRoutingMethod): + return RoutingParams( + top_k=self.routing_method.top_k, + routing_bias=self.routing_method.e_score_correction_bias, + n_group=None, + topk_group=None, + routed_scaling_factor=self.routing_method.routed_scaling_factor, + ) elif isinstance(self.routing_method, MiniMaxM2MoeRoutingMethod): return RoutingParams( top_k=self.routing_method.top_k, diff --git a/tensorrt_llm/_torch/modules/fused_moe/routing.py b/tensorrt_llm/_torch/modules/fused_moe/routing.py index 3273269c4ed1..28dd1c2bcc0a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/routing.py +++ b/tensorrt_llm/_torch/modules/fused_moe/routing.py @@ -676,12 +676,33 @@ def __init__( output_dtype=output_dtype, ) self.routed_scaling_factor = float(routed_scaling_factor) + self.routing_impl = Deepseekv3RoutingImpl( + top_k=top_k, + n_group=1, + topk_group=1, + routed_scaling_factor=self.routed_scaling_factor, + ) def apply( self, router_logits: torch.Tensor, input_ids: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: + # M3 produces contiguous FP32 logits and bias. Keep the eager path for + # alternate dtype, device, shape, or layout contracts because the native + # op consumes raw contiguous CUDA buffers and emits FP32 weights. + if (router_logits.is_cuda and router_logits.dtype == torch.float32 + and router_logits.ndim == 2 and router_logits.is_contiguous() + and self.output_dtype == torch.float32): + routing_bias = self.e_score_correction_bias + if (routing_bias.is_cuda + and routing_bias.device == router_logits.device + and routing_bias.dtype == torch.float32 + and routing_bias.ndim == 1 + and routing_bias.numel() == router_logits.shape[1] + and routing_bias.is_contiguous()): + return self.routing_impl.apply(router_logits, routing_bias) + # ``moe_scheduler.MoEScheduler`` calls ``routing_method.apply(router_logits, # input_ids)`` positionally; mirror the parent's signature so the # MiniMax-M3 routing accepts (but ignores) the optional ``input_ids`` diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index fb7f43f689b1..ee18d726a9e2 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -35,6 +35,8 @@ def __init__( use_custom_cublas_mm: bool = False, is_shared_expert: bool = False, swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, ): super().__init__() @@ -45,6 +47,13 @@ def __init__( self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm self.swiglu_limit = float( swiglu_limit) if swiglu_limit is not None else None + # SwiGLU-OAI parameters (alpha gain inside the sigmoid, (up + beta) + # offset). Left None for plain SwiGLU, where the fused kernel defaults + # to alpha=1.0 / beta=0.0. + self.swiglu_alpha = float( + swiglu_alpha) if swiglu_alpha is not None else None + self.swiglu_beta = float( + swiglu_beta) if swiglu_beta is not None else None config = config or ModelConfig() use_cute_dsl_bf16_gemm = getattr(config, "use_cute_dsl_bf16_gemm", @@ -146,14 +155,22 @@ def _apply_activation(self, x, *, has_lora: bool = False): logger.warning( f"GatedMLP._apply_activation: LoRA path active; forcing non-FP8 activation dtype bf16/fp16, layer_idx={self.layer_idx}" ) - return swiglu(x, swiglu_limit=self.swiglu_limit) + return swiglu(x, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta) else: return swiglu(x, quant_scale=self.down_proj.input_scale, quant_type=torch.float8_e4m3fn, - swiglu_limit=self.swiglu_limit) + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta) else: - return swiglu(x, swiglu_limit=self.swiglu_limit) + return swiglu(x, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta) elif callable(self.activation): return self.activation(x) elif self.activation is None: @@ -163,17 +180,26 @@ def _apply_activation(self, x, *, has_lora: bool = False): f"Activation {self.activation} not yet implemented for fused GatedMLP" ) + def _is_plain_swiglu(self): + """True when the SwiGLU has no alpha gain or (up + beta) offset. + + The fused CuteDSL GEMM+SwiGLU epilogue implements plain silu(gate) * up + only; a swigluoai-style alpha/beta must use the Triton swiglu kernel. + """ + return ((self.swiglu_alpha is None or self.swiglu_alpha == 1.0) + and (self.swiglu_beta is None or self.swiglu_beta == 0.0)) + def _can_fuse_gate_up_swiglu(self): """Check if fused GEMM + SwiGLU path is available. Returns True when all conditions are met: - CuteDSL blockscaling mode is enabled (implies Blackwell + CuteDSL) - - Activation is SwiGLU (F.silu) + - Activation is SwiGLU (F.silu) with no alpha gain / (up + beta) offset - gate_up_proj uses NVFP4 quantization - gate_up_proj has no bias (bias not supported in fused kernel) """ return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu - and self.gate_up_proj.has_nvfp4 + and self._is_plain_swiglu() and self.gate_up_proj.has_nvfp4 and not self.gate_up_proj.has_bias) def _can_fuse_gate_up_swiglu_fp4out(self): diff --git a/tensorrt_llm/_torch/modules/swiglu.py b/tensorrt_llm/_torch/modules/swiglu.py index 328bb4e9fdcb..b76e558ab46a 100644 --- a/tensorrt_llm/_torch/modules/swiglu.py +++ b/tensorrt_llm/_torch/modules/swiglu.py @@ -27,7 +27,8 @@ def scale_and_clamp(x, scale, dtype): @triton.jit def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, - swiglu_limit: tl.constexpr, BLOCK_SIZE: tl.constexpr, + swiglu_limit: tl.constexpr, swiglu_alpha: tl.constexpr, + swiglu_beta: tl.constexpr, BLOCK_SIZE: tl.constexpr, HAS_O_SCALE: tl.constexpr, HAS_SWIGLU_LIMIT: tl.constexpr) -> None: i = tl.program_id(axis=0).to(tl.int64) @@ -43,10 +44,16 @@ def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, b = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) if HAS_SWIGLU_LIMIT: + # Gate clamp is upper-side only; up clamp is symmetric. This is the + # shape both plain SwiGLU-with-limit and the MiniMax-M3 / SGLang + # swigluoai activation require. a = tl.minimum(a, swiglu_limit) b = tl.clamp(b, -swiglu_limit, swiglu_limit) - result = tl.sigmoid(a) * a * b + # swiglu_alpha=1.0, swiglu_beta=0.0 recovers plain silu_and_mul; + # swiglu_alpha=1.702, swiglu_beta=1.0 is the swigluoai shape (alpha gain + # inside the sigmoid, (up + 1) offset). + result = a * tl.sigmoid(swiglu_alpha * a) * (b + swiglu_beta) if HAS_O_SCALE: o_scale = tl.load(o_scale_ptr) @@ -58,7 +65,9 @@ def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, def swiglu(x, quant_scale: Optional[torch.Tensor] = None, quant_type=None, - swiglu_limit: Optional[float] = None): + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None): if quant_scale is not None: assert quant_type is not None return torch.ops.trtllm.silu_and_mul( @@ -66,6 +75,11 @@ def swiglu(x, scale=quant_scale, dtype=quant_type, swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, ) - return torch.ops.trtllm.silu_and_mul(x, swiglu_limit=swiglu_limit) + return torch.ops.trtllm.silu_and_mul(x, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9049b1b90a73..35fbfb71741f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -479,10 +479,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, - kv_cache_config, - is_disagg=self._is_disagg) + draft_kv_cache_manager_cls = self._get_draft_kv_cache_manager_cls( + effective_draft_config, kv_cache_config) total += self._per_manager_cache_cost( draft_kv_cache_manager_cls, effective_draft_config, kv_cache_config) @@ -1047,10 +1045,17 @@ 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: - logger.info( - "Attention DP is enabled, separate draft KV cache is not supported." - ) + if self._mapping.enable_attention_dp and getattr( + self._kv_cache_manager_cls, 'supports_shared_draft_layers', + True): + # Under attention DP, draft layers share the target manager (the + # layout existing deployments were validated with). A manager can + # opt out: MiniMax-M3's coalesces an index-K pool into its KV + # pages and exposes only synthetic AttentionOp tensors, which the + # dense Eagle3 drafter cannot attend against, so it requires the + # separate draft manager even under attention DP. + logger.info("Attention DP: draft layers share the target KV " + "cache manager.") return False return should_use_separate_draft_kv_cache(self._speculative_config) @@ -1081,6 +1086,17 @@ 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, + draft_kv_config): + """Resolve the draft manager class from the draft config, promoted + to V2 when the target manager is V2.""" + draft_kv_cache_manager_cls = get_kv_cache_manager_cls( + effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) + if self._is_kv_cache_manager_v2 and not issubclass( + draft_kv_cache_manager_cls, KVCacheManagerV2): + draft_kv_cache_manager_cls = KVCacheManagerV2 + return draft_kv_cache_manager_cls + def _create_one_model_draft_kv_cache_manager( self, estimating_kv_cache: bool = False, @@ -1145,8 +1161,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._fallback_if_unsupported_kv_cache_manager_v2( draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config) @@ -1156,12 +1172,23 @@ 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 + # A target manager class may request a different page size for the + # separate draft manager (e.g. MiniMax-M3, see + # draft_manager_tokens_per_block there for the rationale). + draft_tpb = getattr(self._kv_cache_manager_cls, + 'draft_manager_tokens_per_block', + self._tokens_per_block) + if draft_tpb != self._tokens_per_block: + logger.info( + f"Draft KV cache manager uses tokens_per_block={draft_tpb} " + f"(target uses {self._tokens_per_block}).") + draft_kv_config.tokens_per_block = draft_tpb 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_tpb, max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, spec_config=self._speculative_config, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 5f1e0cbc0a26..0a7421d036cc 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -96,6 +96,11 @@ def is_mla(config): return False +def is_minimax_m3(sparse_attention_config): + """True when the sparse attention config selects the MiniMax-M3 algorithm.""" + return sparse_attention_config.algorithm == "minimax_m3" + + def is_qwen3_next(config): return hasattr( config, 'architectures' diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 46d3c55a0771..ba19d7ece768 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -23,6 +23,7 @@ import torch from strenum import StrEnum +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp from tensorrt_llm._utils import ( TensorWrapper, @@ -1059,23 +1060,11 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: for layer_id in typed_range(LayerId(self.num_local_layers)): layer_group_id = self.impl.get_layer_group_id(layer_id) - if self.dtype != DataType.NVFP4: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ) - - key_base_addr - ) - else: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + offset = self._kv_pool_mapping_offset(layer_id, layer_group_id, key_base_addr) + + if self.dtype == DataType.NVFP4: block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ) - - key_base_addr - ) block_scale_addr_offset = ( self.impl.get_mem_pool_base_address( layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED @@ -1088,14 +1077,6 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: * self.kv_factor * self.tokens_per_block, ) - offset = exact_div( - addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.kv_factor - * self.tokens_per_block, - ) - - if self.dtype == DataType.NVFP4: assert block_scale_offset == offset, ( "Block scale offset and offset should be the same" ) @@ -1156,6 +1137,29 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if self.enable_swa_scratch_reuse: self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) + def _kv_pool_mapping_offset( + self, layer_id: LayerId, layer_group_id: int, key_base_addr: int + ) -> int: + """Per-layer offset recorded in ``kv_cache_pool_mapping``. + + The default derives the layer's position within its pool from the K + base address, assuming every layer contributes exactly K(+V) to the + pool slot so the layer stride is uniform. Managers whose pool slots + may interleave extra per-layer buffers between layers (non-uniform + layer strides — e.g. MiniMax M3 when the index-K buffer coalesces + into the K/V pool) must override this with a positional formula. + """ + addr_offset = ( + self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) + - key_base_addr + ) + return exact_div( + addr_offset, + self.get_layer_bytes_per_token(layer_id, Role.KEY) + * self.kv_factor + * self.tokens_per_block, + ) + def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: layer_sizes = [] attention_windows = [] @@ -1588,6 +1592,38 @@ def _extra_buffers_per_layer( """ return None + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """Map native cache roles to disaggregation mapper kinds. + + ``Role.ALL`` is the required fallback for roles without an explicit + entry. The default is the head-major (HND) ``INDEXED`` layout written + by the TRTLLM attention kernels — correct for V1 and standard V2 + managers. ``Role.INDEX_KEY`` defaults to ``REPLICATED``: every + index-key side cache shipped so far (DSA indexer-K on V1, MiniMax M3 + on V2) computes its projection replicated across TP ranks, so the + cache bytes are identical per rank; the entry is inert unless a + subclass actually registers ``INDEX_KEY`` buffers via + ``_extra_buffers_per_layer``. Model-specific managers may declare + logical layouts without requiring the shared extractor to inspect + private attributes or role names. MiniMax M3, for example, maps + ordinary K/V to ``NHD`` and keeps index-key ``REPLICATED``. + + This declaration does not influence storage pooling: V2 storage + coalesces buffers purely by ``(life_cycle, buffer size)``, so roles + with different transfer semantics may share a pool slot when their + per-block sizes coincide (e.g. MiniMax M3 at TP degrees where + K == V == INDEX_KEY bytes per block). The disagg page-table builder + splits each physical pool into one logical view per mapper kind, so + transfer correctness never depends on the coalescing outcome. + + Pool memory is layout-agnostic; this declaration describes what the + manager's paired attention backend actually writes. A static mapping + is valid only for a fixed manager/backend pair. A manager whose backend + selects the layout at runtime must derive the mapping from that + backend's configuration. + """ + return {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED} + @property def blocks_in_primary_pool(self) -> int: """ @@ -1648,18 +1684,19 @@ def get_index_k_buffer( num_heads: int = 1, head_dim: int, dtype: Union[torch.dtype, "DataType", str] = torch.bfloat16, + kv_layout: str = "NHD", ) -> Optional[torch.Tensor]: """Return a torch view over the V2-managed paged ``Role.INDEX_KEY`` buffer for ``layer_idx``, or ``None`` when the layer has no INDEX_KEY buffer registered (e.g. dense layers in a sparse model, or non-local layers on the current PP rank). - The view has shape ``[num_pages, tokens_per_block, num_heads, - head_dim]`` where ``num_pages == impl.get_page_index_upper_bound( - layer_idx, Role.INDEX_KEY)``. Sparse modeling code addresses - entries by ``(page, within_page, head, dim)`` after decomposing - the per-token slot id used by the main paged K/V cache into - ``(page, within_page)``. + For ``kv_layout="NHD"``, the view has shape ``[num_pages, + tokens_per_block, num_heads, head_dim]``. For ``kv_layout="HND"``, + it has shape ``[num_pages, num_heads, tokens_per_block, head_dim]``. + Sparse modeling code decomposes the per-token slot id used by the + main paged K/V cache into ``(page, within_page)`` and indexes the + token axis selected by ``kv_layout``. Because :class:`BufferConfig` only carries an opaque byte ``size`` per block, the dtype and head shape are caller-side contracts. @@ -1678,6 +1715,8 @@ def get_index_k_buffer( propagate to the pool, and successive calls return views over the same backing storage. """ + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") if layer_idx not in self.layer_offsets: return None layer_offset = self.layer_offsets[layer_idx] @@ -1732,21 +1771,21 @@ def get_index_k_buffer( f"{num_slots_total} is not divisible by scale = {scale}." ) num_slots = num_slots_total // scale + if kv_layout == "NHD": + page_shape = [self.tokens_per_block, num_heads, head_dim] + else: + page_shape = [num_heads, self.tokens_per_block, head_dim] if scale == 1: # Non-coalesced INDEX_KEY pool: the per-buffer stride is - # the entire page, so ``[page_upper, tokens_per_block, - # num_heads, head_dim]`` is the correct contiguous view. - shape = [page_upper, self.tokens_per_block, num_heads, head_dim] + # the entire page, so either layout is a contiguous view. + shape = [page_upper, *page_shape] return convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, shape)) - # Coalesced pool: build a ``[num_slots, scale, tokens_per_block, - # num_heads, head_dim]`` view at INDEX_KEY's base, then slice - # ``[:, 0]`` to extract this layer's INDEX_KEY data. The slice - # preserves dim-0 stride = ``scale * page_stride`` bytes, so - # ``view[s, w, h, d]`` lands on the correct byte for any - # ``s`` in [0, num_slots). - full_slot_shape = [num_slots, scale, self.tokens_per_block, num_heads, head_dim] + # Coalesced pool: include the buffers-per-slot dimension, then + # slice ``[:, 0]`` to extract this layer's INDEX_KEY data while + # preserving dim-0 stride = ``scale * page_stride`` bytes. + full_slot_shape = [num_slots, scale, *page_shape] full_view = convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, full_slot_shape)) return full_view[:, 0] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3ab3b8d4f6d9..59f2fdf218e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5728,6 +5728,11 @@ def _pad_attention_dp_dummy_request(self): and self.max_num_tokens is not None): token_nums = [self.max_num_tokens] + # A separate draft KV cache manager must also see the dummy, or + # its prepare_resources hits an unknown request id. + draft_kv_cache_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if (not self._enable_dsv4_adp_dummy_fixes or self.kv_cache_transceiver is None): llm_request = self.kv_cache_manager.add_dummy_requests( @@ -5736,6 +5741,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( @@ -5760,6 +5766,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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a4c2e48ddc53..57aa45e606f6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -40,7 +40,7 @@ from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, create_py_executor_instance, instantiate_sampler, is_mla, validate_feature_combination) -from .config_utils import is_hybrid_linear +from .config_utils import is_hybrid_linear, is_minimax_m3 from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder @@ -411,6 +411,13 @@ def create_py_executor( if llm_args.attn_backend == "VANILLA": tokens_per_block = max_num_tokens + # The MSA kernels require a page size of 128; the Triton reference uses TRT-LLM's default + # of 32. + m3_sparse_config = llm_args.sparse_attention_config + if is_minimax_m3(m3_sparse_config): + tokens_per_block = 128 if m3_sparse_config.implementation == "msa" else 32 + kv_cache_config.tokens_per_block = tokens_per_block + if llm_args.attn_backend in ["FLASHINFER", "FLASHINFER_STAR_ATTENTION"]: # Workaround for flashinfer and star attention if kv_cache_config.enable_block_reuse: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 5bac35ab06d0..74b0ba32a6e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -307,6 +307,9 @@ def __init__( self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type + # Consumed by the disaggregation page-table builder to expose the DSA + # indexer K cache pool as a REPLICATED pool view. + self.enable_indexer_k_cache = enable_indexer_k_cache self.spec_config = spec_config self.pp_layers, self.num_layers = get_pp_layers( num_layers, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d9b5be1ecc1e..d09bc91c3e8e 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -628,8 +628,17 @@ 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 twice before capture while the draft loop + # mutates kv_lens_cuda in place — save/restore it during warmup + # (same as DFlash/PARD). During capture the mutation must be + # recorded, so skip the save there. + 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. @@ -737,7 +746,7 @@ def forward(self, ) # 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 57c3de25dcfe..317272c9e2c8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -680,20 +680,18 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): Drives the two-step sparse attention used by MiniMax-M3 layers 3..N: - 1. An index attention branch projects a per-head Q vector and a - **single replicated** K vector, scores main K/V cache blocks, - and selects the top-``topk`` blocks per ``(num_kv_heads, q_token)`` - pair (with ``init_blocks`` forced at the head and ``local_blocks`` - forced at the tail). + 1. An index attention branch projects a per-head Q vector and a single + replicated K vector, scores main K/V cache blocks, and selects the + top-k blocks per (num_kv_heads, q_token) pair, with init_blocks forced + at the head and local_blocks forced at the tail. 2. A sparse GQA attention runs only over the selected blocks. - The selected backend at runtime uses - :class:`tensorrt_llm._torch.attention_backend.sparse.minimax_m3.MiniMaxM3SparseAttention` - on top of a :class:`MiniMaxM3KVCacheManagerV2` that allocates a - paged side index-K cache (``[num_slots, 1, sparse_index_dim]``) - parallel to the main K/V cache. The M3 checkpoint sets - ``disable_index_value=True`` on every sparse layer so no index V - cache is allocated for the bring-up. + At runtime one of the MiniMax-M3 sparse attention backends under + tensorrt_llm._torch.attention_backend.sparse.minimax_m3 is selected. The + chosen backend runs on top of a MiniMaxM3KVCacheManagerV2 that allocates a + paged side index-K cache of shape [num_slots, 1, sparse_index_dim] parallel + to the main K/V cache. The M3 checkpoint sets disable_index_value=True on + every sparse layer, so no index V cache is allocated. """ algorithm: Literal["minimax_m3"] = "minimax_m3" @@ -730,6 +728,34 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): default=True, description="If True, skip the index V branch (M3 checkpoint default).", ) + num_attention_heads: Optional[int] = Field( + default=None, + description= + "Global number of attention (query) heads. When unset, it falls back " + "to pretrained_config.num_attention_heads.", + ) + num_key_value_heads: Optional[int] = Field( + default=None, + description= + "Global number of key/value heads. When unset, it falls back to " + "pretrained_config.num_key_value_heads, then to num_attention_heads.", + ) + implementation: Literal["triton", "msa"] = Field( + default="triton", + description= + "Sparse attention implementation: 'triton' reference (default) or 'msa' " + "(fmha_sm100 kernels). The 'msa' implementation requires an SM100 GPU, " + "the fmha_sm100 package, and sparse_block_size == 128.", + status="prototype", + ) + + @model_validator(mode="after") + def _validate_msa_block_size(self): + if self.implementation == "msa" and self.sparse_block_size != 128: + raise ValueError( + "MiniMax-M3 'msa' implementation requires sparse_block_size == " + f"128, got {self.sparse_block_size}.") + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -738,7 +764,7 @@ def get_indices_block_size(self) -> int: return self.sparse_block_size def to_sparse_params(self, **kwargs): - from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.metadata import \ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import \ MiniMaxM3SparseParams return MiniMaxM3SparseParams( @@ -750,6 +776,36 @@ def to_sparse_params(self, **kwargs): local_blocks=self.sparse_local_blocks, score_type=self.sparse_score_type, disable_index_value=self.sparse_disable_index_value, + implementation=self.implementation, + ) + + def to_sparse_metadata_params(self, **kwargs): + """Lower into MiniMaxM3SparseMetadataParams for the attention metadata. + + Head counts resolve as this config, then pretrained_config, then a + default; num_key_value_heads falls back to num_attention_heads. Setting + them on the config lets tests skip building a pretrained_config. + """ + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import \ + MiniMaxM3SparseMetadataParams + + pretrained_config = kwargs.get("pretrained_config", None) + + def _value(name: str, default=None): + value = getattr(self, name) + if value is not None: + return value + if pretrained_config is not None: + return getattr(pretrained_config, name, default) + return default + + num_attention_heads = int(_value("num_attention_heads", 0)) + num_kv_heads = int(_value("num_key_value_heads", num_attention_heads)) + return MiniMaxM3SparseMetadataParams( + global_num_q_heads=num_attention_heads, + global_num_kv_heads=num_kv_heads, + num_index_heads=self.sparse_num_index_heads, + topk=self.sparse_topk_blocks, ) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ca6e762490cf..066389d7be27 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1389,6 +1389,16 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "allowed_values": [ + "triton", + "msa" + ], + "annotation": "Literal['triton', 'msa']", + "converter": "", + "kind": "categorical", + "path": "sparse_attention_config.implementation" + }, { "allowed_values": [], "annotation": "Optional[int]", @@ -1451,6 +1461,20 @@ "kind": "categorical", "path": "sparse_attention_config.kt_cache_dtype" }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.num_attention_heads" + }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.num_key_value_heads" + }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index facfca89379b..f61bef4ba921 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -482,6 +482,12 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 88 + - 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 cbce62591221..c332eb57e35b 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -284,6 +284,12 @@ MiniMaxAI/MiniMax-M3-MXFP8: nvidia/MiniMax-M3-NVFP4: - quant_algo: MIXED_PRECISION accuracy: 83 + - 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 2305c0852664..13a8ad792165 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7570,11 +7570,12 @@ def test_auto_dtype(self, tp_size, ep_size): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) - def test_mxfp8(self, tp_size, ep_size): + @parametrize_with_ids("use_msa", [False, True]) + def test_mxfp8(self, use_msa): # MXFP8 checkpoint: weights are MXFP8 (e4m3 + UE8M0 1x32 block # scales) with MXFP8 dynamic activations; the KV cache stays in - # BF16 and the sparse attention path is unchanged from BF16. + # BF16. + tp_size = ep_size = 4 model_name = "MiniMaxAI/MiniMax-M3-MXFP8" model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8" # Halving TP from the BF16 reference (TP=8) doubles per-rank @@ -7583,7 +7584,8 @@ def test_mxfp8(self, tp_size, ep_size): # under the PyTorch cap. kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, enable_block_reuse=False) - sparse_attention_config = MiniMaxM3SparseAttentionConfig() + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton") with LLM(model_path, tensor_parallel_size=tp_size, moe_expert_parallel_size=ep_size, @@ -7634,16 +7636,19 @@ def test_mxfp8_piecewise_cuda_graph(self, tp_size, ep_size): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) - def test_nvfp4(self, tp_size, ep_size): + @parametrize_with_ids("use_msa", [False, True]) + def test_nvfp4(self, use_msa): # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts - # (MIXED_PRECISION checkpoint); the KV cache stays in BF16 and the - # sparse attention path is unchanged from BF16. + # (MIXED_PRECISION checkpoint). The MSA path runs an FP8 KV cache; the + # Triton path keeps the KV cache in BF16. + tp_size = ep_size = 4 model_name = "nvidia/MiniMax-M3-NVFP4" model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, - enable_block_reuse=False) - sparse_attention_config = MiniMaxM3SparseAttentionConfig() + enable_block_reuse=False, + dtype="fp8" if use_msa else "auto") + sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton") moe_config = MoeConfig(backend="CUTLASS") with LLM(model_path, tensor_parallel_size=tp_size, @@ -7659,6 +7664,105 @@ def test_nvfp4(self, tp_size, ep_size): 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): + if use_msa: + 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" + max_draft_len = 3 + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3", + ) + # The runtime forces tokens_per_block per implementation (128 MSA / 32 + # reference). + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + enable_block_reuse=False) + with LLM( + model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa" if use_msa else "triton"), + moe_config=MoeConfig(backend="CUTLASS"), + max_seq_len=4096, + # The fmha_sm100 decode planner caps total_q x num_qo_heads at + # 65536; with 1 + draft_len = 4 verify tokens per row that + # bounds the batch at 1024 / TP-sharded heads (256 unsharded + # under attention DP). + max_batch_size=256 if attention_dp else 512, + speculative_config=spec_config, + # Graphs + spec requires the MSA path: its verify batches + # are decode-shaped and capture-safe (the reference path + # rejects graphs+spec at creation). + 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 s in llm.get_stats(timeout=2): + s = json.loads(s) if isinstance(s, str) else s + sd = s.get("specDecodingStats") or {} + drafted += sd.get("numDraftTokens", 0) + accepted += sd.get("numAcceptedTokens", 0) + steps += sd.get("numRequestsWithDraftTokens", 0) + return drafted, accepted, steps + + task = MMLU(model_name) + task.evaluate(llm) + task = GSM8K(model_name) + task.evaluate(llm) + + # Chat-format acceptance — the drafter's training distribution + # (Inferact/MiniMax-M3-EAGLE3 card: 0.839 / 3.518). Reuses the + # live engine and the cached dataset; ~20 s under CUDA graphs. + questions = [ + r["question"] + for r in load_dataset("gsm8k", "main", split="test") + ][:200] + chat_prompts = [ + llm.tokenizer.apply_chat_template([{ + "role": "user", + "content": q + }], + tokenize=False, + add_generation_prompt=True) + for q 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(f"MiniMax-M3 Eagle3 chat-GSM8K acceptance: rate=" + f"{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, \ + f"Eagle3 chat-GSM8K acceptance length too low: {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 8c36ff420a4a..b8de89d10aa7 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -684,9 +684,13 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=True] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[tp_size=8-ep_size=8] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] 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 (180) +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_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index a6f112843fcf..f5c50285a2f4 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -68,6 +68,7 @@ l0_a10: - unittest/disaggregated/test_cluster_storage.py - unittest/disaggregated/test_extractor.py - unittest/disaggregated/test_peer.py + - unittest/disaggregated/test_cache_reuse_adapter.py - unittest/disaggregated/test_bounce.py - unittest/disaggregated/region/test_block.py - unittest/disaggregated/test_mamba_transfer.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d5d316518f22..1c7cd7b4f40a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -55,7 +55,7 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_nixl[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60) - unittest/_torch/modeling/test_modeling_deepseekv4.py - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) @@ -275,6 +275,7 @@ l0_dgx_b200: backend: pytorch orchestrator: mpi tests: + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-CUTLASS] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 67fb75c29aaf..6add1fd6a59b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -86,6 +86,10 @@ l0_h100: - unittest/disaggregated/test_request_id.py - unittest/disaggregated/test_kv_transfer.py - unittest/disaggregated/test_kv_transfer_mp.py + - unittest/disaggregated/test_transceiver_bounded_polling.py + - unittest/disaggregated/test_pool_matching.py + - unittest/disaggregated/test_deepseek_v4_kv_transfer.py + - unittest/disaggregated/test_minimax_m3_kv_transfer.py # Split the large cache transceiver parametrization into readable chunks. # Main transfer matrix. - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver -k "v1 and no_window" @@ -100,6 +104,8 @@ l0_h100: # Boundary request lengths. - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_boundary_lengths -k "v1" - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_boundary_lengths -k "v2" + # DSA indexer-K side cache (V1, REPLICATED). + - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_dsa_indexer - unittest/disaggregated/test_cache_transceiver_harness_report.py - unittest/disaggregated/test_cache_transceiver_harness.py - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[PYTHON-mha-ctx_fp16_gen_fp16] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 32e9f557da81..49f7faaa9b0a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -214,8 +214,8 @@ full:B200/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_ full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6422343) full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6384747) -full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) -full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) +full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) +full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6422318) full:B200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6344107) @@ -237,9 +237,9 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gp full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency_qsplit] SKIP (https://nvbugs/6423866) full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343) full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375) -full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6424188) +full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[tp_size=8-ep_size=8] SKIP (https://nvbugs/6442594) -full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6445375) +full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6445375) full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6422318) full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_5_VL_7B::test_auto_dtype SKIP (https://nvbugs/6316983) @@ -273,8 +273,8 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6388153) full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] SKIP (https://nvbugs/6422502) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] SKIP (https://nvbugs/6422502) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6422502) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6422502) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6422318) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) 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 new file mode 100644 index 000000000000..1c777d593ae9 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Structural tests for the MiniMax-M3 MSA sparse attention backend. + +These validate backend selection and decode scratch-buffer sizing without +launching kernels. Numerical parity against the Triton reference is covered +by the SM100 integration accuracy test. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + +def test_resolver_selects_msa_backend_when_available(monkeypatch): + import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_availability as avail + + monkeypatch.setattr(avail, "ensure_msa_available", lambda: None) + params = MiniMaxM3SparseAttentionConfig(implementation="msa").to_sparse_params() + assert _resolve_minimax_m3_backend_cls(params) is MiniMaxM3MsaSparseAttention + + +def test_msa_requires_block_size_128(): + # The MSA implementation is fixed to a 128-token page size; a mismatched + # sparse_block_size must fail loudly at config construction rather than being + # silently overridden at runtime. + with pytest.raises(ValueError, match=r"sparse_block_size == 128"): + MiniMaxM3SparseAttentionConfig(implementation="msa", sparse_block_size=64) + + # The Triton reference is unaffected by the constraint. + cfg = MiniMaxM3SparseAttentionConfig(implementation="triton", sparse_block_size=64) + assert cfg.sparse_block_size == 64 + + +def test_msa_metadata_rejects_undersized_max_score_buffer(): + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + # Flat backing store sized for 4 heads * 8 k-tiles * 2 batch = 64 elements, + # too small for the plan's required 4 * 16 * 2 = 128. + metadata.msa_max_score = torch.zeros(4 * 8 * 2) + metadata.kv_cache_manager = None + + 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_msa_proxy_max_score_view_is_contiguous_over_stable_store(): + """The proxy view fed to fmha_sm100 must be contiguous in the exact + [num_index_heads, plan_max_k_tiles, num_tokens] shape the kernel writes, + backed by a stable store so its data_ptr survives CUDA graph replay. + """ + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + # Worst-case store: 4 heads * 16 k-tiles * 8 batch. + num_index_heads, worst_k, max_batch = 4, 16, 8 + metadata.msa_max_score = torch.zeros(num_index_heads * worst_k * max_batch) + store_ptr = metadata.msa_max_score.data_ptr() + + # A smaller live step still yields a contiguous view sized to that step, + # which is what the kernel's stride-agnostic write requires. + view = metadata.msa_proxy_max_score_view(num_index_heads, 5, 3) + assert view.shape == (num_index_heads, 5, 3) + assert view.is_contiguous() + assert view.data_ptr() == store_ptr + + # Oversized requests are rejected rather than silently corrupting memory. + with pytest.raises(ValueError, match=r"msa_max_score backing store"): + metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1) + + +def test_msa_index_k_uses_hnd_cache_view_and_writer(): + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + num_pages, coalescing_scale, page_size, head_dim = 2, 7, 8, 16 + pool = torch.zeros( + num_pages, + coalescing_scale, + 1, + page_size, + head_dim, + dtype=torch.bfloat16, + ) + hnd_cache = pool[:, 0] + + class FakeCacheManager: + def __init__(self): + self.calls = [] + + def get_index_k_buffer(self, layer_idx, kv_layout="NHD"): + self.calls.append((layer_idx, kv_layout)) + return hnd_cache + + manager = FakeCacheManager() + metadata.kv_cache_manager = manager + metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32) + values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim) + + returned = metadata.msa_idx_k_cache(3) + metadata.msa_write_idx_k(3, values) + + assert returned.data_ptr() == hnd_cache.data_ptr() + assert not returned.is_contiguous() + assert manager.calls == [(3, "HND"), (3, "HND")] + torch.testing.assert_close(hnd_cache[0, 0, 2], values[0, 0].to(torch.bfloat16)) + torch.testing.assert_close(hnd_cache[1, 0, 5], values[1, 0].to(torch.bfloat16)) + + +def test_msa_indexer_preserves_strided_hnd_index_k(monkeypatch): + import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer as indexer_module + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig + + config = MiniMaxM3SparseConfig( + num_q_heads=4, + num_kv_heads=1, + head_dim=128, + num_index_heads=4, + sparse_index_dim=128, + block_size=128, + topk=16, + ) + indexer = indexer_module.MsaIndexer(config) + pool = torch.randn(2, 7, 1, 128, 128, dtype=torch.bfloat16) + idx_k_paged = pool[:, 0] + captured = {} + + def fake_proxy_max_score(idx_q, passed_idx_k, **kwargs): + del kwargs + captured["idx_k"] = passed_idx_k + return torch.zeros(4, 2, idx_q.shape[0]) + + expected = torch.zeros(1, 1, 16, dtype=torch.int32) + + def fake_select_blocks_from_maxscore(*args, **kwargs): + del args, kwargs + return expected + + monkeypatch.setattr(indexer_module, "_proxy_max_score", fake_proxy_max_score) + monkeypatch.setattr( + indexer_module, + "select_blocks_from_maxscore", + fake_select_blocks_from_maxscore, + ) + + result = indexer.select_blocks( + torch.zeros(1, 4, 128, dtype=torch.bfloat16), + idx_k_paged, + idx_sm_scale=128**-0.5, + kv_indices=torch.arange(2, dtype=torch.int32), + qo_lens_cpu=torch.tensor([1], dtype=torch.int32), + kv_lens_cpu=torch.tensor([256], dtype=torch.int32), + qo_offset_cpu=torch.tensor([255], dtype=torch.int32), + ) + + assert captured["idx_k"] is idx_k_paged + assert captured["idx_k"].data_ptr() == idx_k_paged.data_ptr() + assert not captured["idx_k"].is_contiguous() + assert result is expected + + +def test_msa_proxy_max_score_strided_index_k_matches_packed(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 (Blackwell) required") + + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer import _proxy_max_score + 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") + + page_size = head_dim = 128 + num_index_heads = 4 + coalescing_scale = 57 + kv_lens_cpu = torch.tensor([1, 130, 257, 128, 511, 1024, 33, 900], dtype=torch.int32) + pages_per_sequence = (kv_lens_cpu + page_size - 1) // page_size + num_pages = int(pages_per_sequence.sum().item()) + + generator = torch.Generator(device="cuda").manual_seed(0) + index_k_pool = torch.randn( + num_pages, + coalescing_scale, + 1, + page_size, + head_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + index_k_strided = index_k_pool[:, 0] + index_k_packed = index_k_strided.contiguous() + index_q = torch.randn( + kv_lens_cpu.numel(), + num_index_heads, + head_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + kwargs = { + "qo_lens_cpu": torch.ones_like(kv_lens_cpu), + "kv_lens_cpu": kv_lens_cpu, + "qo_offset_cpu": kv_lens_cpu - 1, + "kv_indices": torch.arange(num_pages, device="cuda", dtype=torch.int32), + "sm_scale": head_dim**-0.5, + "causal": True, + } + + strided_scores = _proxy_max_score(index_q, index_k_strided, **kwargs) + packed_scores = _proxy_max_score(index_q, index_k_packed, **kwargs) + torch.cuda.synchronize() + + 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(): + """Under one-model Eagle3 spec verify a decode step carries + 1 + draft_len query tokens per request, so the proxy scratch must be + sized by the worst-case decode TOKEN count, not the batch size. + """ + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.kv_cache_manager = None + # 2 sequences, 4 tokens each (draft_len=3): 8 decode tokens per step. + metadata.max_num_sequences = 2 + metadata.max_num_tokens = 8 + # Store sized for batch-only sizing (4 heads * 16 k-tiles * 2), which is + # too small once tokens are accounted for (4 * 16 * 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(): + """Spec-verify decode rows expose one entry per query TOKEN, walking the + causal ladder within the verify window.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + per_token_valid_blocks, + ) + + # One request verifying 4 tokens against kv_len 10 (offset 6): token t + # attends 7 + t positions; with 2-token blocks that is ceil((7+t)/2). + qo = torch.tensor([4], dtype=torch.int32) + kv = torch.tensor([10], dtype=torch.int32) + off = torch.tensor([6], dtype=torch.int32) + n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=2) + assert n_valid.tolist() == [4, 4, 5, 5] + + # Mixed batch: an ordinary decode row (qo=1) alongside a verify row. + qo = torch.tensor([1, 3], dtype=torch.int32) + kv = torch.tensor([9, 6], dtype=torch.int32) + off = kv - qo + n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=4) + # Row 0: 9 positions -> 3 blocks. Row 1 tokens attend 4, 5, 6 -> 1, 2, 2. + assert n_valid.tolist() == [3, 1, 2, 2] 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..d8617108fe8a --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py @@ -0,0 +1,294 @@ +# 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, 127, 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 torch.equal(actual, expected) + + +@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 = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=20, + local_blocks=0, + ) + 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) + + +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) diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index be1f2936abf8..535049f53f77 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -1,5 +1,17 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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. from types import SimpleNamespace @@ -104,3 +116,19 @@ def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: assert kv_cache.committed_tokens == [4, 5, 6, 7, 8, 9] assert kv_cache.num_committed_tokens == 10 assert kv_cache.stopped_committing + + +def test_disagg_role_mapper_kinds_default_to_indexed(): + from tensorrt_llm._torch.disaggregation.resource.page import MapperKind + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role + + manager = object.__new__(KVCacheManagerV2) + + # K/V default to the TRTLLM head-major layout; the index-key side cache + # defaults to REPLICATED (every shipped index-K — DSA V1, MiniMax M3 — + # is TP-replicated). The INDEX_KEY entry is inert unless a subclass + # registers such buffers. + assert manager.get_disagg_role_mapper_kinds() == { + Role.ALL: MapperKind.INDEXED, + Role.INDEX_KEY: MapperKind.REPLICATED, + } diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py index d23a677dad47..073190eacf71 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py @@ -75,6 +75,15 @@ def _extra_buffers_per_layer(self, *, tokens_per_block): } +class _CoalescedIndexKeyV2(_IndexKeyOnSparseLayersV2): + """Mirror M3's pool mapping when INDEX_KEY coalesces with K/V.""" + + def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr): + del key_base_addr + layers_in_group = list(self.impl.layer_grouping[int(layer_group_id)]) + return layers_in_group.index(int(layer_id)) + + class _DuplicateRoleV2(KVCacheManagerV2): """Negative-control subclass: register Role.KEY as an "extra" so the standard buffer + extra duplicate. Must raise.""" @@ -238,11 +247,10 @@ def test_duplicate_role_against_standard_buffer_asserts(self): class TestIndexKeyBufferAccessor(unittest.TestCase): """CUDA/GPU regressions for :meth:`KVCacheManagerV2.get_index_k_buffer`. - The accessor returns a paged torch view shaped - ``[num_pages, tokens_per_block, num_heads, head_dim]`` over the - managed ``Role.INDEX_KEY`` pool, returns ``None`` for dense layers, - rejects wiring mismatches against the V2-reported page stride, and - is zero-copy (writes propagate; ``data_ptr`` stable across calls). + The accessor returns NHD or HND paged torch views over the managed + ``Role.INDEX_KEY`` pool, returns ``None`` for dense layers, rejects + wiring mismatches against the V2-reported page stride, and is zero-copy + (writes propagate; ``data_ptr`` stable across calls). """ NUM_HEADS = 1 @@ -353,6 +361,62 @@ def test_accessor_view_is_zero_copy_over_pool(self): mgr.shutdown() del mgr + def test_accessor_hnd_and_nhd_alias_coalesced_pool(self): + # Match MiniMax-M3's one-head Index-K byte size to one main K/V + # head so V2 coalesces all roles. Both layout views must retain the + # larger physical page stride while addressing the same bytes. + mgr = _CoalescedIndexKeyV2( + sparse_layer_ids=[0, 1, 2, 3], + **_make_kwargs( + num_layers=4, + num_kv_heads=1, + dtype=DataType.BF16, + ), + ) + try: + layer_idx = 3 + nhd = mgr.get_index_k_buffer( + layer_idx, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + ) + hnd = mgr.get_index_k_buffer( + layer_idx, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + kv_layout="HND", + ) + self.assertIsNotNone(nhd) + self.assertIsNotNone(hnd) + + converter = mgr.impl.get_page_index_converter( + mgr.layer_offsets[layer_idx], Role.INDEX_KEY + ) + self.assertGreater(int(converter.scale), 1) + self.assertEqual( + nhd.shape, + (nhd.shape[0], mgr.tokens_per_block, self.NUM_HEADS, self.HEAD_DIM), + ) + self.assertEqual( + hnd.shape, + (hnd.shape[0], self.NUM_HEADS, mgr.tokens_per_block, self.HEAD_DIM), + ) + self.assertEqual(nhd.data_ptr(), hnd.data_ptr()) + self.assertEqual(nhd.stride(0), hnd.stride(0)) + self.assertEqual(hnd.stride(2), self.HEAD_DIM) + self.assertEqual(hnd.stride(3), 1) + self.assertFalse(hnd.is_contiguous()) + + sentinel = torch.tensor(37.0, dtype=torch.bfloat16, device="cuda") + hnd[0, 0, 3, 7] = sentinel + torch.cuda.synchronize() + self.assertEqual(nhd[0, 3, 0, 7].item(), float(sentinel)) + finally: + mgr.shutdown() + del mgr + def test_accessor_pointer_matches_v2_pool_base(self): # Verify the wrapper points at the exact V2-managed pool base # for INDEX_KEY, so the view participates in the same lifecycle @@ -416,6 +480,21 @@ def test_accessor_rejects_wrong_dtype(self): mgr.shutdown() del mgr + def test_accessor_rejects_unknown_layout(self): + mgr = self._make_sparse_mgr(sparse_layer_ids=(3,)) + try: + with self.assertRaisesRegex(ValueError, "Unsupported kv_layout"): + mgr.get_index_k_buffer( + 3, + num_heads=self.NUM_HEADS, + head_dim=self.HEAD_DIM, + dtype=torch.bfloat16, + kv_layout="HDN", + ) + finally: + mgr.shutdown() + del mgr + def test_accessor_returns_none_on_base_v2_default_manager(self): # The base KVCacheManagerV2 (no extra buffers) must report None # for every local layer when queried for INDEX_KEY, because the diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py new file mode 100644 index 000000000000..896539beaa71 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -0,0 +1,198 @@ +# 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. + +import inspect +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( + MiniMaxM3HfWeightMapper, +) +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + MiniMaxM3VLForConditionalGeneration, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +_NUM_KV_HEADS = 4 +_ROWS_PER_HEAD = 2 + + +def _make_mapper(tp_size: int = 8, num_kv_heads: int = _NUM_KV_HEADS) -> MiniMaxM3HfWeightMapper: + config = SimpleNamespace(num_key_value_heads=num_kv_heads, num_attention_heads=8) + model_config = ModelConfig( + pretrained_config=config, + mapping=Mapping(world_size=tp_size, rank=0, tp_size=tp_size), + ) + model = SimpleNamespace(model_config=model_config, config=config) + mapper = MiniMaxM3HfWeightMapper() + mapper.init_model_and_config(model, model_config) + return mapper + + +def _duplicate_heads(tensor: torch.Tensor, repetitions: int) -> torch.Tensor: + return ( + tensor.reshape(_NUM_KV_HEADS, _ROWS_PER_HEAD, -1) + .repeat_interleave(repetitions, dim=0) + .reshape(_NUM_KV_HEADS * repetitions * _ROWS_PER_HEAD, -1) + ) + + +@pytest.mark.parametrize( + "architecture", + [ + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + ], +) +def test_mapper_registration_and_mx_fallback(architecture: str) -> None: + assert isinstance(AutoCheckpointMapper.get("HF", architecture), MiniMaxM3HfWeightMapper) + assert isinstance(AutoCheckpointMapper.get("MX", architecture), MiniMaxM3HfWeightMapper) + + +@pytest.mark.parametrize( + "model_class", + [MiniMaxM3ForCausalLM, MiniMaxM3VLForConditionalGeneration], +) +def test_load_weights_exposes_weight_mapper(model_class: type) -> None: + assert "weight_mapper" in inspect.signature(model_class.load_weights).parameters + + +@pytest.mark.parametrize("scale_name", ["weight_scale_inv", "weight_scale"]) +def test_tp8_mxfp8_duplicates_kv_weight_and_scale_via_callbacks(scale_name: str) -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.MXFP8)) + prefix = "model.layers.0.self_attn" + + weights = {} + sources = {} + for offset, projection in enumerate(("k_proj", "v_proj")): + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + offset * 100 + scale = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + offset * 32 + weights[f"{prefix}.{projection}.weight"] = weight + weights[f"{prefix}.{projection}.{scale_name}"] = scale + sources[projection] = {"weight": weight, scale_name: scale} + + mapped = mapper.apply_callbacks(module, "qkv_proj", prefix.split("."), weights) + + assert mapped[0] == {} + for projection, projected in zip(("k_proj", "v_proj"), mapped[1:]): + for name, source in sources[projection].items(): + torch.testing.assert_close(projected[name], _duplicate_heads(source, repetitions=2)) + + +def test_nvfp4_scale_behavior_is_preserved() -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4)) + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + weight_scale = torch.arange(16, dtype=torch.float32).reshape(8, 2) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "k_proj", + { + "weight": weight, + "weight_scale": weight_scale, + "weight_scale_inv": scale_inv, + }, + ) + + torch.testing.assert_close(mapped["weight"], _duplicate_heads(weight, repetitions=2)) + torch.testing.assert_close( + mapped["weight_scale"], _duplicate_heads(weight_scale, repetitions=2) + ) + assert mapped["weight_scale_inv"] is scale_inv + + +def test_quant_config_none_is_guarded() -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=None) + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "k_proj", + { + "weight": weight, + "weight_scale_inv": scale_inv, + }, + ) + + torch.testing.assert_close(mapped["weight"], _duplicate_heads(weight, repetitions=2)) + assert mapped["weight_scale_inv"] is scale_inv + + +def test_kv_scale_is_not_expanded_when_kv_heads_cover_tp() -> None: + mapper = _make_mapper(tp_size=2) + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.MXFP8)) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "v_proj", + {"weight_scale_inv": scale_inv}, + ) + + torch.testing.assert_close(mapped["weight_scale_inv"], scale_inv) + + +def test_gate_bias_params_map() -> None: + mapper = MiniMaxM3HfWeightMapper() + source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias" + target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias" + bias = torch.arange(4, dtype=torch.float32) + gate_weight = torch.ones(4, 4) + + renamed = mapper.rename_by_params_map( + mapper.params_map, + { + source_name: bias, + "model.layers.3.block_sparse_moe.gate.weight": gate_weight, + }, + ) + + assert source_name not in renamed + assert renamed[target_name] is bias + assert renamed["model.layers.3.block_sparse_moe.gate.weight"] is gate_weight + + +def test_load_weights_accepts_base_mapper_without_params_map() -> None: + config = SimpleNamespace(num_key_value_heads=_NUM_KV_HEADS, num_attention_heads=8) + model_config = ModelConfig(pretrained_config=config, mapping=Mapping()) + model = object.__new__(MiniMaxM3ForCausalLM) + torch.nn.Module.__init__(model) + model.model_config = model_config + mapper = HfWeightMapper() + source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias" + target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias" + bias = torch.arange(4, dtype=torch.float32) + + with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: + model.load_weights({source_name: bias}, weight_mapper=mapper) + + call_kwargs = base_load_weights.call_args.kwargs + assert call_kwargs["weight_mapper"] is mapper + renamed = mapper.rename_by_params_map(call_kwargs["params_map"], {source_name: bias}) + assert renamed[target_name] is bias diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index ca2ed046e57c..ea97c53985c1 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -32,9 +32,13 @@ from utils.llm_data import llm_models_root from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( + MiniMaxM3HfWeightMapper, +) from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, _build_swiglu_oai_dense_mlp, + _minimax_m3_swiglu_oai, _strip_language_model_prefix, _wrap_dict_as_config, get_moe_layer_ids, @@ -43,7 +47,7 @@ get_text_config, is_minimax_m3_vl_config, ) -from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl +from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl_v2 from tensorrt_llm._torch.modules.fused_moe.routing import ( MiniMaxM2MoeRoutingMethod, MiniMaxM3MoeRoutingMethod, @@ -276,17 +280,12 @@ def test_get_moe_layer_ids_length_mismatch_raises(): # ``use_gemma=True`` and ``hidden_size=head_dim``; the # :meth:`apply_qk_norm` reshape matches an independent hand-written # reference. -# * Sparse index branch: ``index_q_proj`` is column-parallel and -# projects to ``num_index_heads * sparse_index_dim``; -# ``index_k_proj`` is **replicated** (tp_mode is None) and projects -# to **only** ``sparse_index_dim`` (single K per token, broadcast -# across all index heads for block-selection scoring) — this is the -# SGLang reference contract, confirmed by the M3 checkpoint shape -# ``(sparse_index_dim, hidden_size)``. +# * Sparse index branch: fused replicated (tp_mode None) index_qk_proj with +# output [idx_q | idx_k] = num_index_heads * sparse_index_dim + sparse_index_dim +# (idx_k is one K per token). Source weights (512, 6144) + (128, 6144) are +# concatenated into (640, 6144) at load time. # * Dense layers do not expose any index branch attributes (negative # control). -# * Real M3 checkpoint shape for ``index_k_proj.weight`` is -# ``(sparse_index_dim, hidden_size)`` = ``(128, 6144)``. def _make_attention_test_config(): @@ -401,6 +400,7 @@ def test_minimax_m3_attention_dense_construction_matches_config(): for name in ( "index_q_proj", "index_k_proj", + "index_qk_proj", "index_q_norm", "index_k_norm", ): @@ -480,16 +480,12 @@ def test_minimax_m3_attention_apply_qk_norm_matches_reference(): def test_minimax_m3_attention_sparse_construction_matches_config(): """Sparse layer adds index branch with the SGLang-correct shapes. - Verifies the **bug fix** from the iter-4 work: - * ``index_q_proj`` is column-parallel and projects to - ``num_index_heads * sparse_index_dim``. - * ``index_k_proj`` is replicated (``tp_mode is None``) and projects - to **only** ``sparse_index_dim`` — a single replicated K per - token, *not* per-head. This matches SGLang's ``ReplicatedLinear`` - and the M3 checkpoint's ``index_k_proj.weight`` shape - ``(sparse_index_dim, hidden_size)``. - * ``index_q_norm`` / ``index_k_norm`` are per-head Gemma RMSNorm - of width ``sparse_index_dim``. + Verifies the fused index projection: + * index_qk_proj is replicated (tp_mode None), out = + num_index_heads * sparse_index_dim (idx_q) + sparse_index_dim (idx_k), + where idx_k is one K per token (SGLang ReplicatedLinear contract). + * index_q_norm / index_k_norm are per-head Gemma RMSNorm of width + sparse_index_dim. """ text_cfg, model_cfg = _make_attention_test_config() sparse_cfg = text_cfg.sparse_attention_config @@ -506,33 +502,19 @@ def test_minimax_m3_attention_sparse_construction_matches_config(): assert attn.is_sparse_attention_layer is True assert attn.disable_index_value is True - # index_q_proj: per-head Q for the index branch. As of iter-15 this - # is **replicated** (tp_mode=None) across TP ranks, not - # column-parallel: the sparse forward consumes ``idx_q`` reshaped to - # ``[num_tokens, num_index_heads, sparse_index_dim]`` and a - # column-parallel split would slice the head dimension (breaking the - # reshape at any ``tp_size > num_index_heads`` geometry, including - # the TP=8 configuration the real-checkpoint smoke test now uses). - # The replicated weight is small (~3 MiB BF16) so the per-rank - # memory cost is negligible. - assert attn.index_q_proj.in_features == hidden - assert attn.index_q_proj.out_features == num_index_heads * sparse_index_dim - assert attn.index_q_proj.tp_mode is None, ( - f"index_q_proj must be replicated (tp_mode=None) so the sparse " - f"forward's `idx_q.view(num_tokens, num_index_heads, sparse_index_dim)` " - f"reshape is well-defined at any TP geometry, got " - f"{attn.index_q_proj.tp_mode!r}" - ) - - # index_k_proj: REPLICATED, only sparse_index_dim outputs. - assert attn.index_k_proj.in_features == hidden - assert attn.index_k_proj.out_features == sparse_index_dim, ( - f"index_k_proj.out_features must be sparse_index_dim={sparse_index_dim}, " - f"got {attn.index_k_proj.out_features} (regression of the iter-4 fix)" - ) - assert attn.index_k_proj.tp_mode is None, ( - f"index_k_proj must be replicated (tp_mode=None), got {attn.index_k_proj.tp_mode!r}" + # index_qk_proj: fused [idx_q | idx_k], replicated (tp_mode None) so the + # idx_q -> [num_tokens, num_index_heads, sparse_index_dim] reshape holds at + # any TP geometry (column-parallel would slice the head dim). + assert attn.index_q_size == num_index_heads * sparse_index_dim + assert attn.index_k_size == sparse_index_dim + assert attn.index_qk_proj.in_features == hidden + assert attn.index_qk_proj.out_features == num_index_heads * sparse_index_dim + sparse_index_dim + assert attn.index_qk_proj.tp_mode is None, ( + f"index_qk_proj must be replicated, got {attn.index_qk_proj.tp_mode!r}" ) + # The separate projections must be gone after fusion. + assert not hasattr(attn, "index_q_proj") + assert not hasattr(attn, "index_k_proj") # Per-head Gemma RMSNorm of width sparse_index_dim. assert attn.index_q_norm.use_gemma is True @@ -620,22 +602,254 @@ def test_minimax_m3_attention_dense_apply_index_qk_norm_raises(): attn.apply_index_qk_norm(idx_q, idx_k) +def _make_fused_qk_norm_rope_test_config(): + """Return (text_config, ModelConfig) with a kernel-supported head_dim. + + fused_qk_norm_rope only compiles for head_dim in {64, 128, 256}, so the + scaled-down _make_attention_test_config (head_dim=32) cannot drive the + fused path. This variant uses the real M3 head_dim=128, sparse_index_dim=128 + and rotary_dim=64 geometry with a small head count so the tensors stay tiny. + """ + n_layers = 4 + sparse_cfg = { + "use_sparse_attention": True, + "sparse_index_dim": 128, + "sparse_num_index_heads": 2, + "sparse_topk_blocks": 4, + "sparse_block_size": 16, + "sparse_init_block": 0, + "sparse_local_block": 1, + "sparse_score_type": "max", + "sparse_disable_index_value": [0, 1, 1, 1], + "sparse_attention_freq": [0, 1, 1, 1], + } + text_cfg = _wrap_dict_as_config( + { + "hidden_size": 512, + "intermediate_size": 128, + "num_hidden_layers": n_layers, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 128, + "vocab_size": 256, + "max_position_embeddings": 256, + "rms_norm_eps": 1e-6, + "use_gemma_norm": True, + "rope_theta": 5000000.0, + "rotary_dim": 64, + "partial_rotary_factor": 0.5, + "qk_norm_type": "per_head", + "use_qk_norm": True, + "sparse_attention_config": sparse_cfg, + "torch_dtype": torch.bfloat16, + } + ) + model_cfg = ModelConfig( + pretrained_config=text_cfg, + mapping=Mapping(), + skip_create_weights_in_init=True, + ) + return text_cfg, model_cfg + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_main_matches_separate(): + """Fused main-branch helper matches separate norm plus partial RoPE. + + Exercises the layer helper end-to-end so the wiring (reading rope.dim for + the partial rotary dim, rope.theta, is_neox, use_gemma_norm, per-head norm + weights, and variance_epsilon) is validated against the fallback + apply_qk_norm plus rotary_emb path the fused kernel replaces. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + device = torch.device("cuda") + dtype = torch.bfloat16 + head_dim = attn.head_dim + + torch.manual_seed(0) + attn.q_norm.weight = torch.nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device) * 0.2) + attn.k_norm.weight = torch.nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device) * 0.2) + + seq = 6 + qkv = torch.randn(seq, attn.q_size + 2 * attn.kv_size, dtype=dtype, device=device) + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + 3 + + # Fused path. + fused = attn._fused_qk_norm_rope( + qkv.clone(), + position_ids, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + assert fused is not None, "bf16 qkv + position_ids must take the fused path" + q_f, k_f, v_f = fused.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + + # Separate fallback path. + q_s, k_s, v_s = qkv.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + q_s, k_s = attn.apply_qk_norm(q_s, k_s) + q_s, k_s = attn.rotary_emb(position_ids, [q_s, k_s]) + + torch.testing.assert_close(q_f.contiguous(), q_s.contiguous(), rtol=5e-2, atol=1e-1) + torch.testing.assert_close(k_f.contiguous(), k_s.contiguous(), rtol=5e-2, atol=1e-1) + # V is untouched by both paths. + torch.testing.assert_close(v_f.contiguous(), v_s.contiguous(), rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_index_matches_separate(): + """Fused index-branch helper (num_heads_v=0) matches the separate path. + + The index branch concatenates idx_q (sparse_num_index_heads heads) and + idx_k (1 replicated head), passes num_heads_v=0 so the kernel norms and + rotates only the Q/K segments, and splits back. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + device = torch.device("cuda") + dtype = torch.bfloat16 + sparse_index_dim = attn.sparse_index_dim + num_index_heads = attn.sparse_num_index_heads + + torch.manual_seed(1) + attn.index_q_norm.weight = torch.nn.Parameter( + torch.randn(sparse_index_dim, dtype=dtype, device=device) * 0.3 + ) + attn.index_k_norm.weight = torch.nn.Parameter( + torch.randn(sparse_index_dim, dtype=dtype, device=device) * 0.3 + ) + + seq = 5 + idx_q = torch.randn(seq, num_index_heads * sparse_index_dim, dtype=dtype, device=device) + idx_k = torch.randn(seq, sparse_index_dim, dtype=dtype, device=device) + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + 7 + + # Fused path over concatenated [idx_q, idx_k]. + fused = attn._fused_qk_norm_rope( + torch.cat([idx_q, idx_k], dim=-1), + position_ids, + num_heads_q=num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=sparse_index_dim, + q_norm=attn.index_q_norm, + k_norm=attn.index_k_norm, + ) + assert fused is not None + iq_f, ik_f = fused.split([num_index_heads * sparse_index_dim, sparse_index_dim], dim=-1) + + # Separate fallback path. + iq_s, ik_s = attn.apply_index_qk_norm(idx_q, idx_k) + iq_s, ik_s = attn.rotary_emb(position_ids, [iq_s, ik_s]) + + torch.testing.assert_close(iq_f.contiguous(), iq_s.contiguous(), rtol=5e-2, atol=1e-1) + torch.testing.assert_close(ik_f.contiguous(), ik_s.contiguous(), rtol=5e-2, atol=1e-1) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_fallbacks(): + """The fused helper returns None (fallback) for non-bf16 or no position_ids.""" + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + device = torch.device("cuda") + seq = 3 + total = attn.q_size + 2 * attn.kv_size + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + + # No position_ids means RoPE cannot run, so fall back. + qkv_bf16 = torch.randn(seq, total, dtype=torch.bfloat16, device=device) + assert ( + attn._fused_qk_norm_rope( + qkv_bf16, + None, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=attn.head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + is None + ) + + # Non-bf16 activations hit the bf16-only guard, so fall back. + qkv_fp16 = torch.randn(seq, total, dtype=torch.float16, device=device) + assert ( + attn._fused_qk_norm_rope( + qkv_fp16, + position_ids, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=attn.head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + is None + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") +def test_minimax_m3_expect_fused_qk_norm_rope_predicate(): + """The fused-path expectation guarding the forward asserts. + + For every M3 quant flavor with bf16 attention activations (plain BF16, + MXFP8, and the NVFP4 checkpoint), the fused kernel is expected, so a runtime + fallback trips the assertion. Non-bf16 activations or missing position_ids + relax the expectation. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + device = torch.device("cuda") + position_ids = torch.arange(4, dtype=torch.int32, device=device) + + # bf16 activations require fusion. + assert attn.attn_activation_dtype == torch.bfloat16 + assert attn._expect_fused_qk_norm_rope(position_ids) is True + + # No position_ids means RoPE cannot run, so a fallback is allowed. + assert attn._expect_fused_qk_norm_rope(None) is False + + # Non-bf16 activations allow a fallback with no assertion. + attn.attn_activation_dtype = torch.float16 + assert attn._expect_fused_qk_norm_rope(position_ids) is False + + @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") def test_minimax_m3_attention_real_config_index_branch_shapes(): """Real M3 config → sparse-layer index branch has the checkpoint's shapes. - Asserts the iter-4 fix in numbers: - * ``index_q_proj.out_features == 512`` (= 4 * 128 - = ``num_index_heads * sparse_index_dim``). - * ``index_k_proj.out_features == 128`` (= ``sparse_index_dim``) - and ``tp_mode is None`` (replicated). The real - ``index_k_proj.weight`` in the checkpoint has shape - ``(128, 6144)``; a regression to the old - ``num_index_heads * sparse_index_dim`` (512) would break weight - loading at runtime. - * ``index_q_norm.weight.shape == (128,)`` and - ``index_k_norm.weight.shape == (128,)``. + Asserts the fused index projection in numbers: + * index_qk_proj.out_features == 640 (4 * 128 + 128), replicated. Source + weights (512, 6144) + (128, 6144) merge into (640, 6144) at load time. + * index_q_norm / index_k_norm weights have shape (128,). """ pytest.importorskip("transformers") cfg = AutoConfig.from_pretrained(_checkpoint_path(), trust_remote_code=True) @@ -668,22 +882,14 @@ def test_minimax_m3_attention_real_config_index_branch_shapes(): assert num_index_heads == 4 assert sparse_index_dim == 128 - # index_q_proj: 4 * 128 = 512 out, replicated (tp_mode=None) as of - # iter-15. The downstream sparse forward reshapes ``idx_q`` to - # ``[num_tokens, num_index_heads, sparse_index_dim]``; a - # column-parallel split would slice the head dimension and break - # that reshape at any ``tp_size > num_index_heads`` geometry - # (including TP=8 used by the real-checkpoint smoke test). The - # replicated weight is ~3 MiB BF16 — the per-rank memory cost is - # negligible. - assert attn.index_q_proj.in_features == int(text_cfg.hidden_size) - assert attn.index_q_proj.out_features == num_index_heads * sparse_index_dim - assert attn.index_q_proj.tp_mode is None - - # index_k_proj: 128 out (NOT 512), replicated. - assert attn.index_k_proj.in_features == int(text_cfg.hidden_size) - assert attn.index_k_proj.out_features == sparse_index_dim - assert attn.index_k_proj.tp_mode is None + # index_qk_proj: fused [idx_q | idx_k] = 4 * 128 + 128 = 640 out, replicated. + assert attn.index_q_size == num_index_heads * sparse_index_dim + assert attn.index_k_size == sparse_index_dim + assert attn.index_qk_proj.in_features == int(text_cfg.hidden_size) + assert attn.index_qk_proj.out_features == num_index_heads * sparse_index_dim + sparse_index_dim + assert attn.index_qk_proj.tp_mode is None + assert not hasattr(attn, "index_q_proj") + assert not hasattr(attn, "index_k_proj") # Per-head Gemma index norms: width sparse_index_dim. assert tuple(attn.index_q_norm.weight.shape) == (sparse_index_dim,) @@ -761,10 +967,87 @@ def test_minimax_m3_routing_method_default_scale_is_identity(): torch.testing.assert_close(same_weights, base_weights, rtol=0, atol=0) +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="fused MiniMax-M3 routing requires CUDA") +@pytest.mark.parametrize("num_tokens", [1, 64, 8192]) +def test_minimax_m3_fused_routing_matches_reference(num_tokens, monkeypatch): + num_experts = 128 + top_k = 4 + routed_scaling_factor = 2.0 + logits = torch.full((num_tokens, num_experts), -4.0, device="cuda", dtype=torch.float32) + token_offsets = torch.arange(num_tokens, device="cuda", dtype=torch.int64).unsqueeze(1) + expert_offsets = torch.arange(top_k, device="cuda", dtype=torch.int64).unsqueeze(0) + selected_experts = (token_offsets + expert_offsets) % num_experts + selected_logits = torch.tensor([4.0, 3.0, 2.0, 1.0], device="cuda").expand(num_tokens, -1) + logits.scatter_(1, selected_experts, selected_logits) + bias = torch.linspace(-0.01, 0.01, num_experts, device="cuda", dtype=torch.float32) + + fused = MiniMaxM3MoeRoutingMethod( + top_k=top_k, + num_experts=num_experts, + callable_e_score_correction_bias=lambda: bias, + routed_scaling_factor=routed_scaling_factor, + ) + + scores = torch.sigmoid(logits) + _, reference_idx = torch.topk(scores + bias, k=top_k, dim=-1, sorted=False) + reference_weights = scores.gather(1, reference_idx) + reference_idx = reference_idx.to(torch.int32) + reference_weights = ( + reference_weights / (reference_weights.sum(dim=-1, keepdim=True) + 1e-20) + ) * routed_scaling_factor + + monkeypatch.setattr( + MiniMaxM2MoeRoutingMethod, + "apply", + lambda *_args, **_kwargs: pytest.fail("production FP32 routing used the PyTorch fallback"), + ) + fused_idx, fused_weights = fused.apply(logits) + + reference_order = reference_idx.argsort(dim=-1) + fused_order = fused_idx.argsort(dim=-1) + reference_idx = reference_idx.gather(1, reference_order) + reference_weights = reference_weights.gather(1, reference_order) + fused_idx = fused_idx.gather(1, fused_order) + fused_weights = fused_weights.gather(1, fused_order) + + torch.testing.assert_close(fused_idx, reference_idx, rtol=0, atol=0) + torch.testing.assert_close(fused_weights, reference_weights, rtol=1e-5, atol=1e-6) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="fused MiniMax-M3 routing requires CUDA") +def test_minimax_m3_fused_routing_cuda_graph_replay_tracks_inputs(): + num_tokens = 64 + num_experts = 128 + torch.manual_seed(1) + logits = torch.empty(num_tokens, num_experts, device="cuda", dtype=torch.float32) + bias = torch.randn(num_experts, device="cuda", dtype=torch.float32) * 0.1 + routing = MiniMaxM3MoeRoutingMethod( + top_k=4, + num_experts=num_experts, + callable_e_score_correction_bias=lambda: bias, + routed_scaling_factor=2.0, + ) + + logits.normal_() + routing.apply(logits) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_idx, graph_weights = routing.apply(logits) + + logits.normal_(mean=1.0, std=2.0) + reference_idx, reference_weights = routing.apply(logits.clone()) + graph.replay() + + torch.testing.assert_close(graph_idx, reference_idx, rtol=0, atol=0) + torch.testing.assert_close(graph_weights, reference_weights, rtol=0, atol=0) + + @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 needs CUDA") -def test_text_norm_weights_real_loader_smoke(): - """real ``_load_weights_impl`` populates norm parameters. +def test_text_norm_weights_real_loader_smoke(monkeypatch: pytest.MonkeyPatch): + """real ``_load_weights_impl_v2`` populates norm parameters. Constructs a memory-safe stub containing the top-level ``model.norm`` and the first decoder layer's ``input_layernorm`` and @@ -772,7 +1055,7 @@ def test_text_norm_weights_real_loader_smoke(): :class:`RMSNorm`, ~12 KB on CUDA), reads the corresponding tensors from the real checkpoint via ``safetensors``, strips the ``language_model.`` prefix exactly as the M3 VL wrapper does, and - invokes :func:`_load_weights_impl` end-to-end. The test fails if any + invokes :func:`_load_weights_impl_v2` end-to-end. The test fails if any target parameter remains at its zero-initialisation, proving the canonical loader walks the module tree and copies the correct source keys for these BF16 parameters. @@ -871,17 +1154,21 @@ def __init__(self) -> None: "model.layers.0.post_attention_layernorm.weight", } - # Invoke the canonical loader. `_load_weights_impl` walks the stub's + # Invoke the canonical loader. `_load_weights_impl_v2` walks the stub's # module tree and uses the generic per-parameter copy fallback because # RMSNorm does not define ``load_weights``. Disable the parallel # executor so a failure surfaces immediately rather than as a thread # traceback (the parallel path is exercised in production; for this # tiny 3-module slice the serial walk is what the test should observe). - os.environ["TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL"] = "True" - try: - _load_weights_impl(stub, text_weights, allow_partial_loading=True) - finally: - os.environ.pop("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", None) + monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") + weight_mapper = MiniMaxM3HfWeightMapper() + weight_mapper.init_model_and_config(stub, stub.model_config) + _load_weights_impl_v2( + stub, + text_weights, + weight_mapper, + allow_partial_loading=True, + ) # The three norms should now hold the source tensors' values. torch.testing.assert_close( @@ -971,3 +1258,47 @@ def test_minimax_m3_swiglu_oai_dense_mlp_under_adp_is_replicated(): "ADP down_proj must skip the cross-rank all-reduce; otherwise it " "mixes outputs across independent rank-local token sets" ) + + +# --------------------------------------------------------------------------- +# Fused SwiGLU-OAI numeric equivalence +# --------------------------------------------------------------------------- +# +# The dense MLP and MoE shared expert run the swigluoai activation, expressed +# as plain SwiGLU with (alpha, beta, limit) so it routes through the fused +# silu_and_mul kernel. This pins that the fused kernel with alpha=1.702, +# beta=1.0, limit=7.0 matches the reference _minimax_m3_swiglu_oai, and that +# alpha=1.0, beta=0.0 recovers plain SwiGLU. + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="fused silu_and_mul Triton kernel requires CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_minimax_m3_swiglu_oai_fused_matches_reference(dtype): + """The fused silu_and_mul kernel with (alpha, beta, limit) must match the + eager _minimax_m3_swiglu_oai reference, and reduce to plain SwiGLU when + alpha=1 and beta=0. + """ + from tensorrt_llm._torch.modules.swiglu import swiglu + + torch.manual_seed(0) + alpha, limit = 1.702, 7.0 + # Wide range so both clamp branches (gate upper, up symmetric) are hit. + gate_up = torch.randn(64, 2 * 128, device="cuda", dtype=dtype) * 6.0 + + ref = _minimax_m3_swiglu_oai(gate_up, alpha=alpha, limit=limit) + fused = swiglu(gate_up, swiglu_alpha=alpha, swiglu_beta=1.0, swiglu_limit=limit) + + assert fused.shape == ref.shape + assert fused.dtype == gate_up.dtype + # fp32 accumulation inside the kernel; loose tol for bf16 rounding. + atol = 2e-2 if dtype == torch.bfloat16 else 1e-4 + torch.testing.assert_close(fused.float(), ref.float(), atol=atol, rtol=1e-2) + + # alpha=1 / beta=0 reduces to plain SwiGLU: silu(gate_clamped) * up_clamped. + gate, up = gate_up.chunk(2, dim=-1) + gate_c = gate.clamp(max=limit) + up_c = up.clamp(min=-limit, max=limit) + plain_ref = torch.nn.functional.silu(gate_c) * up_c + plain = swiglu(gate_up, swiglu_limit=limit) + torch.testing.assert_close(plain.float(), plain_ref.float(), atol=atol, rtol=1e-2) diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index f9fcaec58a9c..40b4f5d109dc 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -78,6 +78,7 @@ DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, MiniMaxM2MoeRoutingMethod, + MiniMaxM3MoeRoutingMethod, RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, SigmoidRenormMoeRoutingMethod, @@ -411,15 +412,18 @@ def _create_routing_method(routing_method_cls, top_k, num_experts, dtype, model_ is_fused=False, # Use PyTorch implementation for testing ) - # MiniMaxM2 routing method requires special parameters - if routing_method_cls == MiniMaxM2MoeRoutingMethod: + # MiniMax routing methods require the correction bias and expert count. + if routing_method_cls in (MiniMaxM2MoeRoutingMethod, MiniMaxM3MoeRoutingMethod): # Create e_score_correction_bias as a zero tensor (no bias correction in test) e_score_correction_bias = torch.zeros(num_experts, dtype=dtype, device="cuda") - return routing_method_cls( + kwargs = dict( top_k=top_k, num_experts=num_experts, callable_e_score_correction_bias=lambda: e_score_correction_bias, ) + if routing_method_cls == MiniMaxM3MoeRoutingMethod: + kwargs["routed_scaling_factor"] = 2.0 + return routing_method_cls(**kwargs) # SigmoidRenorm routing method requires num_experts if routing_method_cls == SigmoidRenormMoeRoutingMethod: @@ -1664,6 +1668,40 @@ def test_trtllm_gen_fp32_routing_bias(routing_method_cls, moe_model_config, quan ) +def test_trtllm_gen_minimax_m3_integrated_routing_scale(): + """TRTLLM-Gen integrated routing must preserve M3's routed scale.""" + if not torch.cuda.is_available() or get_sm_version() not in (100, 103): + pytest.skip("TRTLLM-Gen NVFP4 integrated routing requires Blackwell") + + model_config = MoeModelConfig(128, 4, 512, 512) + skip_reason = should_skip_trtllm( + MoeBackendType.TRTLLM, + QuantAlgo.NVFP4, + model_config, + routing_method_cls=MiniMaxM3MoeRoutingMethod, + ) + if skip_reason: + pytest.skip(skip_reason) + + # A single-rank, non-attention-DP mapping makes ConfigurableMoE pass + # router_logits into TRTLLM-Gen, exercising the integrated MiniMax2/M3 + # routing kernel rather than the scheduler's separated routing path. + _test_moe_worker( + moe_backend=MoeBackendType.TRTLLM.value, + dtype=torch.bfloat16, + quant_algo=QuantAlgo.NVFP4, + model_config=model_config, + seq_len=8, + enable_autotune=True, + routing_method_cls=MiniMaxM3MoeRoutingMethod, + dtype_routing_logits=torch.float32, + swiglu_alpha=1.702, + swiglu_beta=1.0, + swiglu_limit=7.0, + bias_dtype=torch.float32, + ) + + # ============================================================================ # MoE Multi-GPU Tests # ============================================================================ diff --git a/tests/unittest/disaggregated/kv_transfer_harness.py b/tests/unittest/disaggregated/kv_transfer_harness.py new file mode 100644 index 000000000000..5b1089f4e099 --- /dev/null +++ b/tests/unittest/disaggregated/kv_transfer_harness.py @@ -0,0 +1,514 @@ +# 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. + +"""Threaded single-process NIXL harness for V2 disaggregated KV transfer tests. + +Creates one ``KvCacheTransceiverV2`` per rank inside a single process using +threads plus a Barrier-based ``Distributed`` mock, then drives a full +ctx-send / gen-receive transfer and hands verification back to the caller. +Model specifics (cache-manager construction, pool initialization, post-transfer +verification) are injected as hooks; see ``test_deepseek_v4_kv_transfer.py`` +and ``test_minimax_m3_kv_transfer.py`` for the two current users. +""" + +import os +import threading +import uuid +from typing import Dict, List, Optional, Protocol, Sequence, TypeVar + +import tensorrt_llm +import tensorrt_llm.bindings +import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 +from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig + +# Reduce NIXL threads for unit test: default 8 threads per agent causes heavy +# contention when creating multiple agents on a single GPU in the same process. +os.environ.setdefault("TRTLLM_NIXL_NUM_THREADS", "0") + + +# --------------------------------------------------------------------------- +# Harness-scale constants shared by all model-specific manager factories +# --------------------------------------------------------------------------- +TOKENS_PER_BLOCK = 128 +MAX_SEQ_LEN = 512 +MAX_BATCH_SIZE = 16 +VOCAB_SIZE = 129280 + +_CacheManagerT = TypeVar("_CacheManagerT", bound=KVCacheManagerV2) +_CacheManagerT_co = TypeVar("_CacheManagerT_co", bound=KVCacheManagerV2, covariant=True) +_CacheManagerT_contra = TypeVar("_CacheManagerT_contra", bound=KVCacheManagerV2, contravariant=True) + + +class ManagerFactory(Protocol[_CacheManagerT_co]): + """Build one cache manager per rank of a (tp x pp) instance. + + Model-specific knobs (dtype, compress ratios, sparse layers, ...) are + closed over by the caller rather than threaded through the harness. + """ + + def __call__( + self, + tp: int, + pp: int, + enable_dp: bool, + /, + ) -> Sequence[_CacheManagerT_co]: ... + + +class CacheInitializer(Protocol[_CacheManagerT_contra]): + def __call__( + self, + managers: Sequence[_CacheManagerT_contra], + tp: int, + /, + *, + seed_base: int = 0, + fill_random: bool = True, + ) -> None: ... + + +class CacheVerifier(Protocol[_CacheManagerT_contra]): + def __call__( + self, + *, + request_lengths: List[int], + ctx_managers: Sequence[_CacheManagerT_contra], + gen_managers: Sequence[_CacheManagerT_contra], + ctx_tp: int, + ctx_pp: int, + gen_tp: int, + gen_pp: int, + ctx_enable_dp: bool, + gen_enable_dp: bool, + ctx_request_ids: List[int], + gen_request_ids: List[int], + ) -> None: ... + + +# --------------------------------------------------------------------------- +# ThreadSafeDistributed: threading.Barrier-based Distributed mock +# --------------------------------------------------------------------------- +class ThreadSafeDistributed: + """Distributed mock using threading.Barrier for single-process multi-rank testing. + + Provides the same interface as TorchDistributedWrapper from test_py_cache_transceiver_mp.py + but uses Barrier + Lock + shared dict instead of torch.distributed. + """ + + def __init__( + self, + local_rank: int, + world_size: int, + tp_size: int, + pp_size: int, + tp_rank: int, + pp_rank: int, + shared: dict, + ): + self.rank = local_rank + self._world_size = world_size + self._tp_size = tp_size + self._pp_size = pp_size + self._tp_rank = tp_rank + self._pp_rank = pp_rank + self._s = shared + self._bcast_idx = 0 + self._ag_idx = 0 + self._pp_ag_idx = 0 + self._tp_ag_idx = 0 + + @property + def tp_size(self): + return self._tp_size + + @property + def pp_size(self): + return self._pp_size + + @property + def world_size(self): + return self._world_size + + def broadcast(self, obj, root=0): + idx = self._bcast_idx + self._bcast_idx += 1 + key = f"bcast_{idx}" + if self.rank == root: + self._s[key] = obj + self._s["barrier"].wait() + result = self._s[key] + self._s["barrier"].wait() + return result + + def allgather(self, obj): + idx = self._ag_idx + self._ag_idx += 1 + key = f"ag_{idx}" + with self._s["lock"]: + if key not in self._s: + self._s[key] = [None] * self._world_size + self._s[key][self.rank] = obj + self._s["barrier"].wait() + result = list(self._s[key]) + self._s["barrier"].wait() + return result + + def pp_allgather(self, obj): + idx = self._pp_ag_idx + self._pp_ag_idx += 1 + key = f"pp_ag_{idx}_tp{self._tp_rank}" + with self._s["lock"]: + if key not in self._s: + self._s[key] = [None] * self._pp_size + self._s[key][self._pp_rank] = obj + self._s["barrier"].wait() + result = list(self._s[key]) + self._s["barrier"].wait() + return result + + def tp_allgather(self, obj): + idx = self._tp_ag_idx + self._tp_ag_idx += 1 + key = f"tp_ag_{idx}_pp{self._pp_rank}" + with self._s["lock"]: + if key not in self._s: + self._s[key] = [None] * self._tp_size + self._s[key][self._tp_rank] = obj + self._s["barrier"].wait() + result = list(self._s[key]) + self._s["barrier"].wait() + return result + + +# --------------------------------------------------------------------------- +# Threading helpers +# --------------------------------------------------------------------------- +def run_concurrent(items, fn): + """Run fn(item) for each item concurrently in threads and propagate errors.""" + errors = [None] * len(items) + results = [None] * len(items) + + def _worker(idx, item): + try: + results[idx] = fn(item) + except Exception as e: + errors[idx] = e + + threads = [threading.Thread(target=_worker, args=(i, item)) for i, item in enumerate(items)] + for t in threads: + t.start() + for t in threads: + t.join() + for i, err in enumerate(errors): + if err is not None: + raise err + return results + + +def _create_transceiver_in_thread(rank, mapping, cache_manager, dist_mock, config, results, errors): + """Thread target: create one KvCacheTransceiverV2.""" + try: + tc = KvCacheTransceiverV2( + mapping=mapping, + dist=dist_mock, + kv_cache_manager=cache_manager, + cache_transceiver_config=config, + ) + results[rank] = tc + except Exception as e: + errors[rank] = e + + +def create_instance_transceivers( + tp: int, + pp: int, + enable_dp: bool, + cache_managers: Sequence[KVCacheManagerV2], + config: CacheTransceiverConfig, +) -> List[KvCacheTransceiverV2]: + """Create KvCacheTransceiverV2 for all ranks via threaded init.""" + world_size = tp * pp + shared = {"barrier": threading.Barrier(world_size), "lock": threading.Lock()} + results = [None] * world_size + errors = [None] * world_size + threads = [] + + for rank in range(world_size): + pp_rank = rank // tp + tp_rank = rank % tp + mapping = Mapping( + world_size=world_size, + rank=rank, + tp_size=tp, + pp_size=pp, + enable_attention_dp=enable_dp, + ) + dist_mock = ThreadSafeDistributed(rank, world_size, tp, pp, tp_rank, pp_rank, shared) + t = threading.Thread( + target=_create_transceiver_in_thread, + args=( + rank, + mapping, + cache_managers[rank], + dist_mock, + config, + results, + errors, + ), + ) + threads.append(t) + + for t in threads: + t.start() + for t in threads: + t.join() + + for rank, err in enumerate(errors): + if err is not None: + raise err + + return results + + +def get_ctx_info_endpoint(tc: KvCacheTransceiverV2) -> Optional[str]: + """Extract the context_info_endpoint from a transceiver's disaggregated params.""" + endpoints = tc.get_disaggregated_params().get("ctx_info_endpoint") or [] + return endpoints[0] if endpoints else None + + +def get_layers_per_pp(num_layers: int, pp_size: int) -> List[int]: + """Return a list of layer counts per PP rank (mirrors C++ getLayerNumPPRank). + + When num_layers is not evenly divisible by pp_size, the first + (num_layers % pp_size) ranks get one extra layer. + Matches Mapping.pp_layers / torch.tensor_split behaviour. + """ + base = num_layers // pp_size + extra = num_layers % pp_size + return [base + (1 if r < extra else 0) for r in range(pp_size)] + + +def run_kv_transfer_test( + ctx_tp: int, + ctx_pp: int, + gen_tp: int, + gen_pp: int, + ctx_enable_dp: bool, + gen_enable_dp: bool, + update_before_transfer: bool = True, + *, + manager_factory: ManagerFactory[_CacheManagerT], + init_fn: CacheInitializer[_CacheManagerT], + verify_fn: CacheVerifier[_CacheManagerT], +) -> None: + """Run one ctx->gen KV transfer with injectable model-specific cache hooks.""" + ctx_world = ctx_tp * ctx_pp + gen_world = gen_tp * gen_pp + + # Mix of block-aligned and non-aligned lengths for boundary testing. + # TOKENS_PER_BLOCK=128: 65=half+1, 256=2x exact, 129=1x+1, 383=3x-1 + request_lengths = [65, 256, 129, 383] + + # ===== 1. Create cache managers ===== + ctx_managers = list(manager_factory(ctx_tp, ctx_pp, ctx_enable_dp)) + gen_managers = list(manager_factory(gen_tp, gen_pp, gen_enable_dp)) + + # ===== 2. Initialize data ===== + # ctx: random data, seed=pp_rank (same across TP, different across PP) + init_fn(ctx_managers, ctx_tp, seed_base=1000, fill_random=True) + # gen: zeros + init_fn(gen_managers, gen_tp, fill_random=False) + + # ===== 3. Create KvCacheTransceiverV2 instances (threaded init) ===== + config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + max_tokens_in_buffer=512, + ) + ctx_tcs = create_instance_transceivers(ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config) + gen_tcs = create_instance_transceivers(gen_tp, gen_pp, gen_enable_dp, gen_managers, config) + + try: + ctx_info_endpoint = get_ctx_info_endpoint(ctx_tcs[0]) + + # ===== 4. Create requests and determine handle map ===== + # handle_map: rank -> [(req_idx, ctx_request, gen_request)] + ctx_handle_map: Dict[int, List] = {r: [] for r in range(ctx_world)} + gen_handle_map: Dict[int, List] = {r: [] for r in range(gen_world)} + ctx_request_ids: List[int] = [] + gen_request_ids: List[int] = [] + + sampling_params = SamplingParams() + + for req_idx, req_len in enumerate(request_lengths): + unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + ctx_rid = req_idx * 2 + gen_rid = req_idx * 2 + 1 + ctx_request_ids.append(ctx_rid) + gen_request_ids.append(gen_rid) + + ctx_dp_rank = req_idx % ctx_tp if ctx_enable_dp else 0 + + ctx_request = LlmRequest( + request_id=ctx_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config() + ), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + ) + ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + + gen_request = LlmRequest( + request_id=gen_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config() + ), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + gen_request.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=ctx_rid, + ctx_dp_rank=ctx_dp_rank, + ctx_info_endpoint=ctx_info_endpoint, + disagg_request_id=unique_rid, + ) + + for rank in range(ctx_world): + tp_rank = rank % ctx_tp + should_handle = (not ctx_enable_dp) or (req_idx % ctx_tp == tp_rank) + if should_handle: + ctx_handle_map[rank].append((req_idx, ctx_request)) + + for rank in range(gen_world): + tp_rank = rank % gen_tp + should_handle = (not gen_enable_dp) or (req_idx % gen_tp == tp_rank) + if should_handle: + gen_handle_map[rank].append((req_idx, gen_request)) + + # ===== 5. Allocate KV cache for all ranks ===== + # prepare_resources is a no-op for non-draft KVCacheManagerV2. + # All ranks must allocate BEFORE mutating shared request objects + # (add_new_token changes is_first_context_chunk). + # + # Gen ranks take the disagg-gen-init path: prepare_disagg_gen_init + # sizes the cache for the full prompt and pre-declares + # history_length=prompt_len, matching what the V2 scheduler's + # _try_schedule_disagg_gen_init does in production so the + # transceiver's TRANS_COMPLETE contract check is satisfied. + # Ctx ranks take the regular prefill path (prepare_context + + # resize_context). + gen_batches: Dict[int, ScheduledRequests] = {} + for rank in range(gen_world): + reqs = [req for _, req in gen_handle_map[rank]] + if reqs: + batch = ScheduledRequests() + batch.context_requests_last_chunk = reqs + for req in reqs: + gen_managers[rank].prepare_disagg_gen_init(req) + gen_batches[rank] = batch + + ctx_batches: Dict[int, ScheduledRequests] = {} + for rank in range(ctx_world): + reqs = [req for _, req in ctx_handle_map[rank]] + if reqs: + batch = ScheduledRequests() + batch.context_requests_last_chunk = reqs + for req in reqs: + ctx_managers[rank].prepare_context(req) + ctx_managers[rank].resize_context(req, req.context_chunk_size) + ctx_batches[rank] = batch + + # ===== 5.5. context_current_position + add_new_token ===== + # Set position on each unique request once (needed for transfer metadata). + seen: set = set() + for rank in range(ctx_world): + for _, req in ctx_handle_map[rank]: + if req.py_request_id not in seen: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + seen.add(req.py_request_id) + + seen = set() + for rank in range(gen_world): + for _, req in gen_handle_map[rank]: + if req.py_request_id not in seen: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + seen.add(req.py_request_id) + + # ===== 5.6. update_resources BEFORE transfer (mode: update_before) ===== + if update_before_transfer: + for rank, batch in ctx_batches.items(): + ctx_managers[rank].update_resources(batch) + for rank, batch in gen_batches.items(): + gen_managers[rank].update_resources(batch) + + # ===== 6. gen receive + ctx send ===== + for rank in range(gen_world): + for _, req in gen_handle_map[rank]: + gen_tcs[rank].request_and_receive_async(req) + for rank in range(ctx_world): + for _, req in ctx_handle_map[rank]: + ctx_tcs[rank].respond_and_send_async(req) + + # ===== 7. Wait for completion (threaded, dist calls inside) ===== + run_concurrent( + ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) + ) + run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) + + # ===== 7.5. update_resources AFTER transfer (mode: update_after) ===== + if not update_before_transfer: + for rank, batch in ctx_batches.items(): + ctx_managers[rank].update_resources(batch) + for rank, batch in gen_batches.items(): + gen_managers[rank].update_resources(batch) + + # ===== 8. Verify ===== + verify_fn( + request_lengths=request_lengths, + ctx_managers=ctx_managers, + gen_managers=gen_managers, + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + ctx_request_ids=ctx_request_ids, + gen_request_ids=gen_request_ids, + ) + + finally: + for tc in ctx_tcs + gen_tcs: + try: + tc.shutdown() + except Exception: + pass + for mgr in ctx_managers + gen_managers: + try: + mgr.shutdown() + except Exception: + pass diff --git a/tests/unittest/disaggregated/region/test_block.py b/tests/unittest/disaggregated/region/test_block.py index cc6cc3323608..0ccac385ad09 100644 --- a/tests/unittest/disaggregated/region/test_block.py +++ b/tests/unittest/disaggregated/region/test_block.py @@ -1,3 +1,18 @@ +# 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. + import numpy as np from tensorrt_llm._torch.disaggregation.base.region import ( @@ -6,9 +21,8 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ( - HeadMatchMapper, - HeadMismatchMapper, - IdentityMapper, + HNDHeadMismatchMapper, + IntactMapper, ) from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -81,12 +95,13 @@ def test_spec_region_and_spec_region_pair(): assert pair.dst.spec == "spec_dst" -def test_identity_mapper(): +def test_intact_mapper_identity_degenerate(): + """Full contiguous overlap degrades to a whole-region pass-through copy.""" src_group = MemRegionGroup(ptrs=np.array([100, 200], dtype=np.int64), bytes_per_region=32) dst_group = MemRegionGroup(ptrs=np.array([300, 400], dtype=np.int64), bytes_per_region=32) src_spec = SpecRegion(memory=src_group, spec="a") dst_spec = SpecRegion(memory=dst_group, spec="b") - mapper = IdentityMapper() + mapper = IntactMapper([0, 16], [0, 16], 16, 16) result = mapper.map(src_spec, dst_spec) assert isinstance(result, SpecRegionPair) np.testing.assert_array_equal(result.src.memory.ptrs, [100, 200]) @@ -95,14 +110,11 @@ def test_identity_mapper(): assert result.dst.memory.bytes_per_region == 32 -def test_head_match_mapper(): +def test_intact_mapper_partial_layers(): + """Selecting a contiguous layer subset yields one shifted fragment.""" self_ri = make_rankinfo(kv_heads_per_rank=2) - peer_ri = make_rankinfo(kv_heads_per_rank=2) - transfer_layers = 2 - src_layer_off = 1 - dst_layer_off = 1 - # slot_size_per_layer = kv_factor * kv_heads * tokens_per_block * dims_per_head * element_bytes - slot_size_per_layer = ( + # bytes_per_layer = kv_factor * kv_heads * tokens_per_block * dims_per_head * element_bytes + bytes_per_layer = ( self_ri.attention.kv_factor * self_ri.attention.kv_heads_per_rank * self_ri.attention.tokens_per_block @@ -113,39 +125,43 @@ def test_head_match_mapper(): dst_group = MemRegionGroup(ptrs=np.array([30, 40], dtype=np.int64), bytes_per_region=1) src_spec = SpecRegion(memory=src_group, spec="srcspec") dst_spec = SpecRegion(memory=dst_group, spec="dstspec") - mapper = HeadMatchMapper( - transfer_layers, - src_layer_off, - dst_layer_off, - self_ri, - peer_ri, - slot_size_per_layer=slot_size_per_layer, - ) + # Layers 1..2 of a 3-layer slot on both sides. + offsets = [bytes_per_layer, 2 * bytes_per_layer] + mapper = IntactMapper(offsets, offsets, bytes_per_layer, bytes_per_layer) result = mapper.map(src_spec, dst_spec) - expected_off = transfer_layers * slot_size_per_layer + expected_bytes = 2 * bytes_per_layer np.testing.assert_array_equal( - result.src.memory.ptrs, [10 + mapper._src_block_off, 20 + mapper._src_block_off] + result.src.memory.ptrs, [10 + bytes_per_layer, 20 + bytes_per_layer] ) np.testing.assert_array_equal( - result.dst.memory.ptrs, [30 + mapper._dst_block_off, 40 + mapper._dst_block_off] + result.dst.memory.ptrs, [30 + bytes_per_layer, 40 + bytes_per_layer] ) - assert result.src.memory.bytes_per_region == expected_off - assert result.dst.memory.bytes_per_region == expected_off + assert result.src.memory.bytes_per_region == expected_bytes + assert result.dst.memory.bytes_per_region == expected_bytes def test_head_mismatch_mapper(): self_ri = make_rankinfo(kv_heads_per_rank=2, tp_size=2, tp_rank=1) peer_ri = make_rankinfo(kv_heads_per_rank=4, tp_size=4, tp_rank=2) - transfer_layers = 1 - src_layer_off = 0 - peer_layer_off = 1 + # buffer bytes = heads * tokens_per_block * dims_per_head * element_bytes + self_bytes_per_layer = 2 * (2 * 4 * 2 * 1) # kv_factor x K-buffer bytes + peer_bytes_per_layer = 2 * (4 * 4 * 2 * 1) src_group = MemRegionGroup(ptrs=np.array([111], dtype=np.int64), bytes_per_region=32) dst_group = MemRegionGroup(ptrs=np.array([222], dtype=np.int64), bytes_per_region=32) src_spec = SpecRegion(memory=src_group, spec="srcspec") dst_spec = SpecRegion(memory=dst_group, spec="dstspec") - mapper = HeadMismatchMapper(transfer_layers, src_layer_off, peer_layer_off, self_ri, peer_ri) + mapper = HNDHeadMismatchMapper( + src_layer_offsets=[0], + dst_layer_offsets=[peer_bytes_per_layer], # layer 1 on the peer side + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=self_bytes_per_layer, + peer_bytes_per_layer=peer_bytes_per_layer, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) result = mapper.map(src_spec, dst_spec) - expected_frag_count = self_ri.attention.kv_factor * transfer_layers + expected_frag_count = 2 # one fragment per (layer, K/V buffer) assert isinstance(result, SpecRegionPair) assert len(result.src.memory.ptrs) == expected_frag_count assert len(result.dst.memory.ptrs) == expected_frag_count diff --git a/tests/unittest/disaggregated/region/test_page.py b/tests/unittest/disaggregated/region/test_page.py index 14042299fa3b..2e5d76782ba4 100644 --- a/tests/unittest/disaggregated/region/test_page.py +++ b/tests/unittest/disaggregated/region/test_page.py @@ -1,3 +1,18 @@ +# 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. + import numpy as np from tensorrt_llm._torch.disaggregation.resource.page import ( @@ -6,6 +21,7 @@ KVCachePageTable, LayerGroup, LocalLayer, + MapperKind, PhysicalPool, PhysicalPoolGroup, PoolView, @@ -49,6 +65,24 @@ def test_pool_view_roundtrip(): assert restored.buffer_entries[1]["offset"] == 128 +def test_pool_view_kind_roundtrip(): + """REPLICATED / NHD kinds survive serialization; default stays INDEXED.""" + entries = _make_buffer_entries() + for kind in (MapperKind.REPLICATED, MapperKind.NHD): + view = PoolView( + pool_idx=2, + buffer_entries=entries, + pool_role=frozenset({"index_key"}), + mapper_kind=kind, + ) + restored = PoolView.from_dict(view.to_dict()) + assert restored.mapper_kind == kind + assert restored.pool_role == frozenset({"index_key"}) + + legacy = PoolView(pool_idx=0, buffer_entries=entries).to_dict() + assert PoolView.from_dict(legacy).mapper_kind == MapperKind.INDEXED + + def test_local_layer_roundtrip(): ll = LocalLayer(local_layer_id=0, global_layer_id=5) d = ll.to_dict() @@ -100,3 +134,94 @@ def test_kv_cache_page_table_roundtrip(): assert len(restored.layer_groups) == 1 assert len(restored.pool_groups) == 1 assert restored.pool_groups[0].pools[0].base_address == 0x10000 + + +# --------------------------------------------------------------------------- +# bytes_per_layer serialization and get_layer_byte_ranges geometry +# --------------------------------------------------------------------------- + + +def test_pool_view_bytes_per_layer_roundtrip(): + entries = _make_buffer_entries() + view = PoolView( + pool_idx=1, + buffer_entries=entries, + pool_role=frozenset({"key", "value"}), + mapper_kind=MapperKind.NHD, + bytes_per_layer=256, + ) + restored = PoolView.from_dict(view.to_dict()) + assert restored.bytes_per_layer == 256 + + # None (INDEXED views) survives the roundtrip too. + legacyless = PoolView(pool_idx=0, buffer_entries=entries) + assert PoolView.from_dict(legacyless.to_dict()).bytes_per_layer is None + + +def _view(entries, bytes_per_layer=None): + return PoolView( + pool_idx=0, + buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), + pool_role=frozenset({"key", "value"}), + mapper_kind=MapperKind.NHD, + bytes_per_layer=bytes_per_layer, + ) + + +def test_layer_byte_ranges_uniform_stride(): + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + # Dedicated K/V pool: dense layers, uniform stride. + starts, bytes_per_layer = get_layer_byte_ranges( + _view([(0, 0, 128), (0, 128, 128), (1, 256, 128), (1, 384, 128)]) + ) + assert starts == {0: 0, 1: 256} + assert bytes_per_layer == 256 + + +def test_layer_byte_ranges_non_uniform_stride(): + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + # Coalesced slot: another role class interleaves 64B after layer 0, + # so the layer stride is non-uniform while sizes stay uniform. + starts, bytes_per_layer = get_layer_byte_ranges( + _view([(0, 0, 128), (0, 128, 128), (1, 320, 128), (1, 448, 128)]) + ) + assert starts == {0: 0, 1: 320} + assert bytes_per_layer == 256 + + +def test_layer_byte_ranges_noncontiguous_layer_raises(): + import pytest + + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + with pytest.raises(ValueError, match="not contiguous"): + get_layer_byte_ranges(_view([(0, 0, 128), (0, 192, 128)])) + + +def test_layer_byte_ranges_nonuniform_size_raises(): + import pytest + + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + with pytest.raises(ValueError, match="not uniform"): + get_layer_byte_ranges(_view([(0, 0, 128), (1, 128, 64)])) + + +def test_layer_byte_ranges_declared_size_mismatch_raises(): + import pytest + + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + with pytest.raises(ValueError, match="declares bytes_per_layer"): + get_layer_byte_ranges(_view([(0, 0, 128)], bytes_per_layer=256)) + + +def test_layer_byte_ranges_empty_view_raises(): + import pytest + + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + with pytest.raises(ValueError, match="no buffer entries"): + get_layer_byte_ranges(_view([])) diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 8faadb0147b4..b58739d889a3 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -196,13 +196,19 @@ def test_fanin_bounce_safe_gate(): expected_transfers) must fall back to the per-fragment path. """ tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + from tensorrt_llm._torch.disaggregation.resource.page import MapperKind + safe = tfr.Receiver._fanin_bounce_safe - def ov(dup, pp): - return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp) + def ov(dup, pp, ranks=(0,)): + return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp, ranks=list(ranks)) + + def ri(lpp, page_table=None): + return SimpleNamespace(layer_num_per_pp=lpp, page_table=page_table) - def ri(lpp): - return SimpleNamespace(layer_num_per_pp=lpp) + def pt(mapper_kind): + view = SimpleNamespace(mapper_kind=mapper_kind) + return SimpleNamespace(layer_groups=[SimpleNamespace(pool_views=[view])]) # single PP stage (overlap_pp_size <= 1): only duplicate_head_factor matters assert safe(ov(1, 1), ri([24])) is True @@ -216,6 +222,12 @@ def ri(lpp): assert safe(ov(1, 4), ri([20])) is False # duplicate heads blocks even an otherwise-even PP split assert safe(ov(2, 4), ri([20, 20, 20, 20])) is False + # replicated views (one elected sender per destination) make multi-writer + # contributions unequal -> fall back; single-writer overlap stays safe, + # and sharded-only view schemes are unaffected + assert safe(ov(1, 1, ranks=(0, 1)), ri([24], pt(MapperKind.REPLICATED))) is False + assert safe(ov(1, 1, ranks=(0,)), ri([24], pt(MapperKind.REPLICATED))) is True + assert safe(ov(1, 1, ranks=(0, 1)), ri([24], pt(MapperKind.NHD))) is True # --------------------------------------------------------------------------- # diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index c6828bac8684..c131183461be 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -117,20 +117,27 @@ class _FakeMgr: def __init__(self): self.beam_width = None + self.pool_indices_window = None def get_batch_cache_indices(self, request_ids, layer_idx=None, beam_width=1): self.beam_width = beam_width return [[10, 11, 12, 13]] + def get_memory_pool_block_indices(self, block_ids, window_size): + # Identity translation: nothing offloaded, block_id == pool slot. + self.pool_indices_window = window_size + return block_ids + req = _FakeReq(prompt_len=7) req.py_request_id = 1 req.py_beam_width = 4 req.sampling_config = _FakeSamplingConfig(beam_width=1) mgr = _FakeMgr() - block_ids = _CacheReuseAdapterV1(mgr).get_block_ids(req, 0, _lg()) + block_ids = _CacheReuseAdapterV1(mgr).get_block_ids(req, 0, _lg(window=512)) assert mgr.beam_width == 4 + assert mgr.pool_indices_window == 512 np.testing.assert_array_equal(block_ids, [10, 11, 12, 13]) def test_pack_beam_cache_indices_single_block_prompt_keeps_all_beams(self): diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 1e9e9a4e3899..651a347b2f08 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -1,3 +1,18 @@ +# 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. + """Single-process test for KVCacheManager (V1/V2) + KvCacheTransceiverV2. Uses threading + ThreadSafeDistributed to create KvCacheTransceiverV2 instances @@ -47,6 +62,7 @@ # Constants # --------------------------------------------------------------------------- NUM_LAYERS = 4 +INDEXER_HEAD_DIM = 128 NUM_KV_HEADS = 4 # 1 for MLA HEAD_DIM = 128 TOKENS_PER_BLOCK = 8 @@ -309,8 +325,10 @@ def _create_cache_manager( max_attention_window_vec: Optional[List[int]] = None, num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, + enable_indexer_k_cache: bool = False, ) -> "KVCacheManager | KVCacheManagerV2": """Create a KVCacheManager (V1) or KVCacheManagerV2 for the given mapping.""" + assert not (enable_indexer_k_cache and use_v2), "DSA indexer K cache is V1-only" num_kv_heads = 1 if is_mla else NUM_KV_HEADS cache_type = CacheTypeCpp.SELFKONLY if is_mla else CacheTypeCpp.SELF @@ -377,6 +395,9 @@ def _create_cache_manager( mapping=mapping, dtype=DataType.FLOAT, model_config=model_config, + enable_indexer_k_cache=enable_indexer_k_cache, + indexer_k_cache_quant_block_size=128, + indexer_k_cache_index_head_dim=INDEXER_HEAD_DIM if enable_indexer_k_cache else 0, ) @@ -389,6 +410,7 @@ def _create_managers_for_instance( max_attention_window_vec: Optional[List[int]] = None, num_layers: int = NUM_LAYERS, max_batch_size: int = MAX_BATCH_SIZE, + enable_indexer_k_cache: bool = False, ) -> List: """Create cache managers for all ranks in an instance.""" managers = [] @@ -404,7 +426,13 @@ def _create_managers_for_instance( ) managers.append( _create_cache_manager( - mapping, is_mla, use_v2, max_attention_window_vec, num_layers, max_batch_size + mapping, + is_mla, + use_v2, + max_attention_window_vec, + num_layers, + max_batch_size, + enable_indexer_k_cache, ) ) return managers @@ -439,6 +467,27 @@ def _init_pool_data_v1( else: pool_tensor.zero_() + if getattr(mgr, "enable_indexer_k_cache", False): + # DSA indexer K is TP-replicated: seed by PP stage only so every + # TP rank of a stage holds identical bytes. + indexer_pool = mgr.impl.get_indexer_k_cache_pool().view(torch.uint8) + if fill_random: + generator = torch.Generator(device=indexer_pool.device).manual_seed( + seed_base + 7000 + pp_rank + ) + indexer_pool.copy_( + torch.randint( + 0, + 256, + indexer_pool.shape, + dtype=torch.uint8, + device=indexer_pool.device, + generator=generator, + ) + ) + else: + indexer_pool.zero_() + def _init_pool_data_v2( managers: List[KVCacheManagerV2], @@ -816,6 +865,77 @@ def verify_all_requests( ) +def _get_indexer_block_data(mgr, request_id, layer_idx, num_layers, pp, tp, enable_dp, req_idx): + """Per-request indexer-K bytes for one global layer on the owning rank. + + Indexer K is TP-replicated, so any TP rank of the layer's PP stage works; + with attention DP only the request's DP group holds the request. + """ + pp_rank = _pp_rank_of_layer(layer_idx, num_layers, pp) + tp_rank = req_idx % tp if enable_dp else 0 + owner = mgr[pp_rank * tp + tp_rank] + block_indices = owner.get_batch_cache_indices([request_id], layer_idx)[0] + valid = [idx for idx in block_indices if idx >= 0] + if not valid: + return None + local_layer = layer_idx - _pp_layer_start(pp_rank, num_layers, pp) + # Pool shape: (numBlocks, numLayers, kvFactor, blockSize), dtype uint8. + pool = owner.impl.get_indexer_k_cache_pool().view(torch.uint8) + return pool[valid, local_layer] + + +def _verify_indexer_k_all_requests( + request_lengths: List[int], + ctx_managers: List, + gen_managers: List, + ctx_tp: int, + ctx_pp: int, + gen_tp: int, + gen_pp: int, + ctx_enable_dp: bool, + gen_enable_dp: bool, + ctx_request_ids: List[int], + gen_request_ids: List[int], + num_layers: int, +): + """Compare the transferred DSA indexer K bytes for every request/layer.""" + for req_idx, _req_len in enumerate(request_lengths): + for layer_idx in range(num_layers): + ctx_data = _get_indexer_block_data( + ctx_managers, + ctx_request_ids[req_idx], + layer_idx, + num_layers, + ctx_pp, + ctx_tp, + ctx_enable_dp, + req_idx, + ) + gen_data = _get_indexer_block_data( + gen_managers, + gen_request_ids[req_idx], + layer_idx, + num_layers, + gen_pp, + gen_tp, + gen_enable_dp, + req_idx, + ) + if ctx_data is None or gen_data is None: + continue + assert ctx_data.shape == gen_data.shape, ( + f"Indexer shape mismatch at req={req_idx} layer={layer_idx}: " + f"ctx={ctx_data.shape} gen={gen_data.shape}" + ) + torch.testing.assert_close( + gen_data, + ctx_data, + rtol=0, + atol=0, + msg=lambda m: (f"Indexer data mismatch at req={req_idx} layer={layer_idx}: {m}"), + ) + + # --------------------------------------------------------------------------- # Main test orchestrator # --------------------------------------------------------------------------- @@ -831,6 +951,7 @@ def run_transfer_test( max_attention_window_vec: Optional[List[int]] = None, num_layers: int = NUM_LAYERS, request_lengths: Optional[List[int]] = None, + enable_indexer_k_cache: bool = False, ): """Run a full KV transfer test using KvCacheTransceiverV2.""" if request_lengths is None: @@ -849,6 +970,7 @@ def run_transfer_test( max_attention_window_vec, num_layers, max_batch_size, + enable_indexer_k_cache, ) gen_managers = _create_managers_for_instance( gen_tp, @@ -859,6 +981,7 @@ def run_transfer_test( max_attention_window_vec, num_layers, max_batch_size, + enable_indexer_k_cache, ) # 2. Initialize data: random for ctx, zeros for gen @@ -987,6 +1110,21 @@ def run_transfer_test( max_attention_window_vec=max_attention_window_vec, num_layers=num_layers, ) + if enable_indexer_k_cache: + _verify_indexer_k_all_requests( + request_lengths=request_lengths, + ctx_managers=ctx_managers, + gen_managers=gen_managers, + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + ctx_request_ids=ctx_request_ids, + gen_request_ids=gen_request_ids, + num_layers=num_layers, + ) # 9. Cleanup if use_v2: @@ -1311,6 +1449,50 @@ def test_cache_transceiver_boundary_lengths( print("PASSED") +# DSA (DeepSeek V3.2) indexer K cache: V1-only, MLA, TP-replicated single +# index head. Covers the REPLICATED pool view end to end through the real +# python transceiver: fan-in owner election, fan-out, and PP layer subsets +# (layer-strided ReplicatedMapper offsets). +DSA_INDEXER_CONFIGS = [ + # (ctx_tp, ctx_pp, gen_tp, gen_pp, ctx_dp, gen_dp, test_id) + (1, 1, 1, 1, False, False, "tp1_to_tp1"), + (2, 1, 1, 1, False, False, "tp2_to_tp1_fanin"), + (1, 1, 2, 1, False, False, "tp1_to_tp2_fanout"), + (1, 2, 1, 1, False, False, "pp2_to_pp1"), + (1, 1, 1, 2, False, False, "pp1_to_pp2"), + (2, 2, 1, 1, False, False, "tp2pp2_to_tp1"), + (2, 1, 2, 1, False, True, "tp2_to_dep2"), +] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,ctx_pp,gen_tp,gen_pp,ctx_enable_dp,gen_enable_dp", + [c[:6] for c in DSA_INDEXER_CONFIGS], + ids=[c[6] for c in DSA_INDEXER_CONFIGS], +) +def test_cache_transceiver_v1_dsa_indexer( + ctx_tp, + ctx_pp, + gen_tp, + gen_pp, + ctx_enable_dp, + gen_enable_dp, +): + """V1 KVCacheManager + DSA indexer K cache through KvCacheTransceiverV2.""" + run_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + is_mla=True, + use_v2=False, + enable_indexer_k_cache=True, + ) + + if __name__ == "__main__": # Quick smoke test run_transfer_test(1, 1, 1, 1, False, False, False, False) diff --git a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py index 0fbe520ee067..b5514c5f116a 100644 --- a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py +++ b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py @@ -1,22 +1,40 @@ +# 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. + """Test KV Transfer for DeepseekV4CacheManager with KvCacheTransceiverV2. -Uses threading + ThreadSafeDistributed to create KvCacheTransceiverV2 instances -in a single process. Validates all DeepseekV4AttentionType cache transfers across -different TP/PP/DP configurations. +Drives the shared threaded single-process harness (``kv_transfer_harness``) +with DeepSeek-V4 cache managers. Validates all DeepseekV4AttentionType cache +transfers across different TP/PP/DP configurations. """ -import os -import threading -import uuid -from typing import Dict, List, Optional, Tuple +import functools +from typing import Dict, List, Optional, Sequence, Tuple import pytest import torch +from kv_transfer_harness import ( + MAX_BATCH_SIZE, + MAX_SEQ_LEN, + TOKENS_PER_BLOCK, + VOCAB_SIZE, + get_layers_per_pp, + run_kv_transfer_test, +) -import tensorrt_llm -import tensorrt_llm.bindings -import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 -from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm import Mapping from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO, @@ -25,22 +43,10 @@ ) from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool, get_pool_bytes -from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.llmapi.llm_args import ( - CacheTransceiverConfig, - DeepSeekV4SparseAttentionConfig, - KvCacheConfig, -) - -# Reduce NIXL threads for unit test: default 8 threads per agent causes heavy -# contention when creating multiple agents on a single GPU in the same process. -os.environ.setdefault("TRTLLM_NIXL_NUM_THREADS", "0") - +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig # --------------------------------------------------------------------------- # Constants matching DeepseekV4CacheManager defaults @@ -48,200 +54,14 @@ HEAD_DIM = 256 # Reduced from 512: test validates transfer, not attention correctness INDEX_HEAD_DIM = 128 WINDOW_SIZE = 128 -TOKENS_PER_BLOCK = 128 -MAX_SEQ_LEN = 512 -MAX_BATCH_SIZE = 16 -VOCAB_SIZE = 129280 NUM_KV_HEADS = 1 INDEXER_QUANT_BLOCK_SIZE = 128 - # DeepSeek-V4 specific ratios (mirrors module constants) SPARSE_RATIO = DEEPSEEK_V4_SPARSE_RATIO OVERLAP_COMPRESSOR_RATIO = DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO -# --------------------------------------------------------------------------- -# ThreadSafeDistributed: threading.Barrier-based Distributed mock -# --------------------------------------------------------------------------- -class ThreadSafeDistributed: - """Distributed mock using threading.Barrier for single-process multi-rank testing. - - Provides the same interface as TorchDistributedWrapper from test_py_cache_transceiver_mp.py - but uses Barrier + Lock + shared dict instead of torch.distributed. - """ - - def __init__( - self, - local_rank: int, - world_size: int, - tp_size: int, - pp_size: int, - tp_rank: int, - pp_rank: int, - shared: dict, - ): - self.rank = local_rank - self._world_size = world_size - self._tp_size = tp_size - self._pp_size = pp_size - self._tp_rank = tp_rank - self._pp_rank = pp_rank - self._s = shared - self._bcast_idx = 0 - self._ag_idx = 0 - self._pp_ag_idx = 0 - self._tp_ag_idx = 0 - - @property - def tp_size(self): - return self._tp_size - - @property - def pp_size(self): - return self._pp_size - - @property - def world_size(self): - return self._world_size - - def broadcast(self, obj, root=0): - idx = self._bcast_idx - self._bcast_idx += 1 - key = f"bcast_{idx}" - if self.rank == root: - self._s[key] = obj - self._s["barrier"].wait() - result = self._s[key] - self._s["barrier"].wait() - return result - - def allgather(self, obj): - idx = self._ag_idx - self._ag_idx += 1 - key = f"ag_{idx}" - with self._s["lock"]: - if key not in self._s: - self._s[key] = [None] * self._world_size - self._s[key][self.rank] = obj - self._s["barrier"].wait() - result = list(self._s[key]) - self._s["barrier"].wait() - return result - - def pp_allgather(self, obj): - idx = self._pp_ag_idx - self._pp_ag_idx += 1 - key = f"pp_ag_{idx}_tp{self._tp_rank}" - with self._s["lock"]: - if key not in self._s: - self._s[key] = [None] * self._pp_size - self._s[key][self._pp_rank] = obj - self._s["barrier"].wait() - result = list(self._s[key]) - self._s["barrier"].wait() - return result - - def tp_allgather(self, obj): - idx = self._tp_ag_idx - self._tp_ag_idx += 1 - key = f"tp_ag_{idx}_pp{self._pp_rank}" - with self._s["lock"]: - if key not in self._s: - self._s[key] = [None] * self._tp_size - self._s[key][self._tp_rank] = obj - self._s["barrier"].wait() - result = list(self._s[key]) - self._s["barrier"].wait() - return result - - -# --------------------------------------------------------------------------- -# Threading helpers -# --------------------------------------------------------------------------- -def run_concurrent(items, fn): - """Run fn(item) for each item concurrently in threads and propagate errors.""" - errors = [None] * len(items) - results = [None] * len(items) - - def _worker(idx, item): - try: - results[idx] = fn(item) - except Exception as e: - errors[idx] = e - - threads = [threading.Thread(target=_worker, args=(i, item)) for i, item in enumerate(items)] - for t in threads: - t.start() - for t in threads: - t.join() - for i, err in enumerate(errors): - if err is not None: - raise err - return results - - -def _create_transceiver_in_thread(rank, mapping, cache_manager, dist_mock, config, results, errors): - """Thread target: create one KvCacheTransceiverV2.""" - try: - tc = KvCacheTransceiverV2( - mapping=mapping, - dist=dist_mock, - kv_cache_manager=cache_manager, - cache_transceiver_config=config, - ) - results[rank] = tc - except Exception as e: - errors[rank] = e - - -def create_instance_transceivers( - tp: int, pp: int, enable_dp: bool, cache_managers: List, config: CacheTransceiverConfig -) -> List[KvCacheTransceiverV2]: - """Create KvCacheTransceiverV2 for all ranks via threaded init.""" - world_size = tp * pp - shared = {"barrier": threading.Barrier(world_size), "lock": threading.Lock()} - results = [None] * world_size - errors = [None] * world_size - threads = [] - - for rank in range(world_size): - pp_rank = rank // tp - tp_rank = rank % tp - mapping = Mapping( - world_size=world_size, - rank=rank, - tp_size=tp, - pp_size=pp, - enable_attention_dp=enable_dp, - ) - dist_mock = ThreadSafeDistributed(rank, world_size, tp, pp, tp_rank, pp_rank, shared) - t = threading.Thread( - target=_create_transceiver_in_thread, - args=( - rank, - mapping, - cache_managers[rank], - dist_mock, - config, - results, - errors, - ), - ) - threads.append(t) - - for t in threads: - t.start() - for t in threads: - t.join() - - for rank, err in enumerate(errors): - if err is not None: - raise err - - return results - - # --------------------------------------------------------------------------- # DeepseekV4CacheManager creation helpers # --------------------------------------------------------------------------- @@ -302,7 +122,12 @@ def _create_managers_for_instance( return managers -def _init_pool_data(managers: List, tp: int, seed_base: int = 0, fill_random: bool = True): +def _init_pool_data( + managers: Sequence[DeepseekV4CacheManager], + tp: int, + seed_base: int = 0, + fill_random: bool = True, +) -> None: """Initialize pool data for all managers. Uses half-precision view of pool memory for initialization. @@ -466,7 +291,7 @@ def _read_cache_data( def _find_ctx_rank_for_layer( layer_idx: int, - ctx_managers: List[DeepseekV4CacheManager], + ctx_managers: Sequence[DeepseekV4CacheManager], ctx_tp: int, ctx_enable_dp: bool, req_idx: int, @@ -491,8 +316,8 @@ def _find_ctx_rank_for_layer( def verify_all_requests( request_lengths: List[int], compress_ratios: List[int], - ctx_managers: List[DeepseekV4CacheManager], - gen_managers: List[DeepseekV4CacheManager], + ctx_managers: Sequence[DeepseekV4CacheManager], + gen_managers: Sequence[DeepseekV4CacheManager], ctx_tp: int, ctx_pp: int, gen_tp: int, @@ -581,12 +406,6 @@ def verify_all_requests( # --------------------------------------------------------------------------- # Main test function # --------------------------------------------------------------------------- -def _get_ctx_info_endpoint(tc: KvCacheTransceiverV2) -> Optional[str]: - """Extract the context_info_endpoint from a transceiver's disaggregated params.""" - endpoints = tc.get_disaggregated_params().get("ctx_info_endpoint") or [] - return endpoints[0] if endpoints else None - - def run_deepseek_v4_transfer_test( ctx_tp: int, ctx_pp: int, @@ -596,202 +415,22 @@ def run_deepseek_v4_transfer_test( gen_enable_dp: bool, compress_ratios: List[int], update_before_transfer: bool = True, -): - """Run a full DeepSeek-V4 KV transfer test.""" - ctx_world = ctx_tp * ctx_pp - gen_world = gen_tp * gen_pp - - # Mix of block-aligned and non-aligned lengths for boundary testing. - # TOKENS_PER_BLOCK=128: 65=half+1, 256=2x exact, 129=1x+1, 383=3x-1 - request_lengths = [65, 256, 129, 383] - - # ===== 1. Create DeepseekV4CacheManagers ===== - ctx_managers = _create_managers_for_instance(ctx_tp, ctx_pp, ctx_enable_dp, compress_ratios) - gen_managers = _create_managers_for_instance(gen_tp, gen_pp, gen_enable_dp, compress_ratios) - - # ===== 2. Initialize data ===== - # ctx: random data, seed=pp_rank (same across TP, different across PP) - _init_pool_data(ctx_managers, ctx_tp, seed_base=1000, fill_random=True) - # gen: zeros - _init_pool_data(gen_managers, gen_tp, fill_random=False) - - # ===== 3. Create KvCacheTransceiverV2 instances (threaded init) ===== - config = CacheTransceiverConfig( - backend="NIXL", - transceiver_runtime="PYTHON", - max_tokens_in_buffer=512, +) -> None: + """Run the shared transfer harness with DeepSeek-V4 cache hooks.""" + run_kv_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + update_before_transfer=update_before_transfer, + manager_factory=lambda tp, pp, enable_dp: _create_managers_for_instance( + tp, pp, enable_dp, compress_ratios + ), + init_fn=_init_pool_data, + verify_fn=functools.partial(verify_all_requests, compress_ratios=compress_ratios), ) - ctx_tcs = create_instance_transceivers(ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config) - gen_tcs = create_instance_transceivers(gen_tp, gen_pp, gen_enable_dp, gen_managers, config) - - try: - ctx_info_endpoint = _get_ctx_info_endpoint(ctx_tcs[0]) - - # ===== 4. Create requests and determine handle map ===== - # handle_map: rank -> [(req_idx, ctx_request, gen_request)] - ctx_handle_map: Dict[int, List] = {r: [] for r in range(ctx_world)} - gen_handle_map: Dict[int, List] = {r: [] for r in range(gen_world)} - ctx_request_ids: List[int] = [] - gen_request_ids: List[int] = [] - - sampling_params = SamplingParams() - - for req_idx, req_len in enumerate(request_lengths): - unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF - ctx_rid = req_idx * 2 - gen_rid = req_idx * 2 + 1 - ctx_request_ids.append(ctx_rid) - gen_request_ids.append(gen_rid) - - ctx_dp_rank = req_idx % ctx_tp if ctx_enable_dp else 0 - - ctx_request = LlmRequest( - request_id=ctx_rid, - max_new_tokens=1, - input_tokens=list(range(req_len)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config() - ), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, - ) - ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) - - gen_request = LlmRequest( - request_id=gen_rid, - max_new_tokens=1, - input_tokens=list(range(req_len)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config() - ), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, - ) - gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, - ctx_dp_rank=ctx_dp_rank, - ctx_info_endpoint=ctx_info_endpoint, - disagg_request_id=unique_rid, - ) - - for rank in range(ctx_world): - tp_rank = rank % ctx_tp - should_handle = (not ctx_enable_dp) or (req_idx % ctx_tp == tp_rank) - if should_handle: - ctx_handle_map[rank].append((req_idx, ctx_request)) - - for rank in range(gen_world): - tp_rank = rank % gen_tp - should_handle = (not gen_enable_dp) or (req_idx % gen_tp == tp_rank) - if should_handle: - gen_handle_map[rank].append((req_idx, gen_request)) - - # ===== 5. Allocate KV cache for all ranks ===== - # prepare_resources is a no-op for non-draft KVCacheManagerV2. - # All ranks must allocate BEFORE mutating shared request objects - # (add_new_token changes is_first_context_chunk). - # - # Gen ranks take the disagg-gen-init path: prepare_disagg_gen_init - # sizes the cache for the full prompt and pre-declares - # history_length=prompt_len, matching what the V2 scheduler's - # _try_schedule_disagg_gen_init does in production so the - # transceiver's TRANS_COMPLETE contract check is satisfied. - # Ctx ranks take the regular prefill path (prepare_context + - # resize_context). - gen_batches: Dict[int, ScheduledRequests] = {} - for rank in range(gen_world): - reqs = [req for _, req in gen_handle_map[rank]] - if reqs: - batch = ScheduledRequests() - batch.context_requests_last_chunk = reqs - for req in reqs: - gen_managers[rank].prepare_disagg_gen_init(req) - gen_batches[rank] = batch - - ctx_batches: Dict[int, ScheduledRequests] = {} - for rank in range(ctx_world): - reqs = [req for _, req in ctx_handle_map[rank]] - if reqs: - batch = ScheduledRequests() - batch.context_requests_last_chunk = reqs - for req in reqs: - ctx_managers[rank].prepare_context(req) - ctx_managers[rank].resize_context(req, req.context_chunk_size) - ctx_batches[rank] = batch - - # ===== 5.5. context_current_position + add_new_token ===== - # Set position on each unique request once (needed for transfer metadata). - seen: set = set() - for rank in range(ctx_world): - for _, req in ctx_handle_map[rank]: - if req.py_request_id not in seen: - req.context_current_position = req.prompt_len - req.add_new_token(req.prompt_len, 0) - seen.add(req.py_request_id) - - seen = set() - for rank in range(gen_world): - for _, req in gen_handle_map[rank]: - if req.py_request_id not in seen: - req.context_current_position = req.prompt_len - req.add_new_token(req.prompt_len, 0) - seen.add(req.py_request_id) - - # ===== 5.6. update_resources BEFORE transfer (mode: update_before) ===== - if update_before_transfer: - for rank, batch in ctx_batches.items(): - ctx_managers[rank].update_resources(batch) - for rank, batch in gen_batches.items(): - gen_managers[rank].update_resources(batch) - - # ===== 6. gen receive + ctx send ===== - for rank in range(gen_world): - for _, req in gen_handle_map[rank]: - gen_tcs[rank].request_and_receive_async(req) - for rank in range(ctx_world): - for _, req in ctx_handle_map[rank]: - ctx_tcs[rank].respond_and_send_async(req) - - # ===== 7. Wait for completion (threaded, dist calls inside) ===== - run_concurrent( - ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) - ) - run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) - - # ===== 7.5. update_resources AFTER transfer (mode: update_after) ===== - if not update_before_transfer: - for rank, batch in ctx_batches.items(): - ctx_managers[rank].update_resources(batch) - for rank, batch in gen_batches.items(): - gen_managers[rank].update_resources(batch) - - # ===== 8. Verify ===== - verify_all_requests( - request_lengths=request_lengths, - compress_ratios=compress_ratios, - ctx_managers=ctx_managers, - gen_managers=gen_managers, - ctx_tp=ctx_tp, - ctx_pp=ctx_pp, - gen_tp=gen_tp, - gen_pp=gen_pp, - ctx_enable_dp=ctx_enable_dp, - gen_enable_dp=gen_enable_dp, - ctx_request_ids=ctx_request_ids, - gen_request_ids=gen_request_ids, - ) - - finally: - for tc in ctx_tcs + gen_tcs: - try: - tc.shutdown() - except Exception: - pass - for mgr in ctx_managers + gen_managers: - try: - mgr.shutdown() - except Exception: - pass # --------------------------------------------------------------------------- @@ -862,21 +501,6 @@ def test_deepseek_v4_kv_transfer( print("PASSED") -# --------------------------------------------------------------------------- -# PP layer distribution helpers (mirrors C++ getLayerNumPPRank) -# --------------------------------------------------------------------------- -def _get_layers_per_pp(num_layers: int, pp_size: int) -> List[int]: - """Return a list of layer counts per PP rank. - - When num_layers is not evenly divisible by pp_size, the first - (num_layers % pp_size) ranks get one extra layer. - Matches Mapping.pp_layers / torch.tensor_split behaviour. - """ - base = num_layers // pp_size - extra = num_layers % pp_size - return [base + (1 if r < extra else 0) for r in range(pp_size)] - - # --------------------------------------------------------------------------- # Uneven PP layer test configurations # --------------------------------------------------------------------------- @@ -945,8 +569,8 @@ def test_deepseek_v4_kv_transfer_uneven_pp( f"ctx_tp={ctx_tp} ctx_pp={ctx_pp} gen_tp={gen_tp} gen_pp={gen_pp} " f"ctx_dp={ctx_enable_dp} gen_dp={gen_enable_dp} " f"num_layers={num_layers} compress_ratios={compress_ratios} " - f"layers_per_pp(ctx)={_get_layers_per_pp(num_layers, ctx_pp)} " - f"layers_per_pp(gen)={_get_layers_per_pp(num_layers, gen_pp)}" + f"layers_per_pp(ctx)={get_layers_per_pp(num_layers, ctx_pp)} " + f"layers_per_pp(gen)={get_layers_per_pp(num_layers, gen_pp)}" ) run_deepseek_v4_transfer_test( diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 6fc2f68fef68..9cd9c1698f1e 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -1,3 +1,18 @@ +# 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. + import numpy as np import pytest @@ -5,16 +20,19 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( KVRegionExtractorV1, build_page_table, + build_page_table_from_manager, ) from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import ( get_global_layer_ids, + get_layer_byte_ranges, get_layer_to_layer_group, get_num_layer_groups, get_num_layers, get_physical_pool, get_unique_layers, ) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, @@ -187,6 +205,142 @@ def test_build_page_table(): manager.shutdown() +def _make_v1_dsa_manager(pp_size: int = 1, pp_rank: int = 0) -> KVCacheManager: + """V1 KVCacheManager with the DSA indexer K cache enabled (MLA-style).""" + return KVCacheManager( + kv_cache_config=KvCacheConfig( + max_tokens=512, + enable_block_reuse=False, + event_buffer_max_size=0, + ), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=4, + num_kv_heads=1, + head_dim=64, + tokens_per_block=32, + max_seq_len=256, + max_batch_size=2, + mapping=Mapping(world_size=pp_size, rank=pp_rank, tp_size=1, pp_size=pp_size), + dtype=DataType.HALF, + enable_indexer_k_cache=True, + indexer_k_cache_quant_block_size=128, + indexer_k_cache_index_head_dim=128, + ) + + +def _byte_view(ptr: int, nbytes: int): + from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor + + return convert_to_torch_tensor(TensorWrapper(int(ptr), DataType.INT8, [int(nbytes)])) + + +@pytest.mark.cuda +def test_v1_dsa_indexer_page_table_is_replicated_with_per_layer_entries(): + manager = _make_v1_dsa_manager() + try: + page_table = build_page_table(manager) + lg = page_table.layer_groups[0] + assert len(lg.pool_views) == 2 + + kv_view, idx_view = lg.pool_views + assert kv_view.mapper_kind == MapperKind.INDEXED + assert idx_view.mapper_kind == MapperKind.REPLICATED + assert idx_view.pool_role == frozenset({"indexer_k"}) + + # One synthesized entry per LG layer, equal-sized, contiguous from 0. + assert get_unique_layers(idx_view) == get_unique_layers(kv_view) + idx_pool = get_physical_pool(page_table, 0, idx_view.pool_idx) + sizes = {int(entry["size"]) for entry in idx_view.buffer_entries} + assert len(sizes) == 1 + per_layer = sizes.pop() + assert per_layer * len(idx_view.buffer_entries) == idx_pool.slot_bytes + offsets = sorted(int(entry["offset"]) for entry in idx_view.buffer_entries) + assert offsets == [i * per_layer for i in range(len(offsets))] + finally: + manager.shutdown() + + +@pytest.mark.cuda +def test_v1_dsa_indexer_replicated_transfer_across_pp(): + """PP1 ctx sends the DSA indexer K cache into two PP2 gen ranks. + + Exercises the full python path on real V1 managers: page-table build, + role-set matching, ReplicatedMapper layer-strided offsets, and a + byte-level copy that must land each gen rank's layer subset at the + right offsets. + """ + import torch + + from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ReplicatedMapper + from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + + ctx = _make_v1_dsa_manager() + gens = [_make_v1_dsa_manager(pp_size=2, pp_rank=r) for r in range(2)] + try: + ctx_extractor = KVRegionExtractorV1(ctx) + ctx_ri = RankInfo.from_kv_cache_manager("ctx", ctx, device_id=0) + registrar = PeerRegistrar(ctx_ri, ctx_extractor) + + # Fill ctx indexer pool with position-dependent bytes; zero gen pools. + ctx_pool_tensor = ctx.impl.get_indexer_k_cache_pool() + flat = ctx_pool_tensor.view(torch.uint8).flatten() + flat.copy_(torch.arange(flat.numel(), dtype=torch.int64, device=flat.device) % 251) + for gen in gens: + gen.impl.get_indexer_k_cache_pool().view(torch.uint8).zero_() + + block_ids = np.array([0, 2, 3], dtype=np.int64) + ctx_pt = ctx_extractor.page_table + idx_pool = get_physical_pool(ctx_pt, 0, 1) + layers_per_gen = 2 + per_layer = idx_pool.slot_bytes // 4 + + for gen_pp_rank, gen in enumerate(gens): + gen_ri = RankInfo.from_kv_cache_manager("gen", gen, device_id=0) + registrar.register("gen", gen_ri.instance_rank, gen_ri) + mapping = registrar.get_pool_mapping(gen_ri) + # KV pool and indexer pool each match 1:1. + assert mapping == {(0, 0): (0, 0), (0, 1): (0, 1)} + + mapper = registrar.get_kv_map(gen_ri, (0, 1), (0, 1)) + assert isinstance(mapper, ReplicatedMapper) + + gen_extractor = registrar.peer_extractor("gen", gen_ri.instance_rank) + pair = mapper.map( + ctx_extractor.extract(block_ids, layer_group_id=0, pool_idx=1), + gen_extractor.extract(block_ids, layer_group_id=0, pool_idx=1), + ) + # This gen rank holds 2 of the 4 layers; the fragment is the + # contiguous 2-layer range at this PP stage's offset within + # the ctx slot. + assert pair.src.memory.bytes_per_region == layers_per_gen * per_layer + expected_src_off = gen_pp_rank * layers_per_gen * per_layer + base_ptrs = idx_pool.base_address + block_ids * idx_pool.slot_bytes + np.testing.assert_array_equal(pair.src.memory.ptrs, base_ptrs + expected_src_off) + + # Emulate the transfer with raw byte copies, then verify content. + for src_ptr, dst_ptr in zip(pair.src.memory.ptrs, pair.dst.memory.ptrs): + _byte_view(dst_ptr, pair.dst.memory.bytes_per_region).copy_( + _byte_view(src_ptr, pair.src.memory.bytes_per_region) + ) + + gen_pool_tensor = gen.impl.get_indexer_k_cache_pool() + gen_bytes = gen_pool_tensor.view(torch.uint8).reshape(gen_pool_tensor.shape[0], -1) + ctx_bytes = ctx_pool_tensor.view(torch.uint8).reshape(ctx_pool_tensor.shape[0], -1) + src_lo = gen_pp_rank * layers_per_gen * per_layer + src_hi = src_lo + layers_per_gen * per_layer + torch.testing.assert_close( + gen_bytes[block_ids], + ctx_bytes[block_ids, src_lo:src_hi], + rtol=0, + atol=0, + ) + finally: + ctx.shutdown() + for gen in gens: + gen.shutdown() + + def test_layer_group_meta_serialization(): import numpy as np @@ -232,6 +386,198 @@ def test_layer_group_meta_serialization(): assert len(restored_lg.pool_views[0].buffer_entries) == 2 +def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_list=(640,)): + """Build a minimal duck-typed KVCacheManagerV2 for _build_page_table_v2. + + ``attrs`` keeps the storage-layer shape ``{(layer_id, role): + (life_cycle_id, pool_index, offset, size)}``; the fake synthesizes the + public ``pool_group_descs`` view from it. Buffer order within a + coalesced buffer follows ascending offset, mirroring how the real + storage config assigns offsets in ``buffer_ids`` order. + """ + from types import SimpleNamespace + + base = 0x1000 + + per_pool: dict[int, list] = {} + for (layer_id, role), attr in attrs.items(): + per_pool.setdefault(int(attr.pool_index), []).append( + (int(attr.offset), layer_id, role, int(attr.size)) + ) + + coalesced_buffers = [] + for pool_idx in range(num_pools): + buffers = sorted(per_pool.get(pool_idx, [])) + sizes = {size for _, _, _, size in buffers} + assert len(sizes) == 1, "buffers in a coalesced buffer are uniform-size by construction" + coalesced_buffers.append( + SimpleNamespace( + single_buffer_size=sizes.pop(), + buffer_ids=[ + SimpleNamespace(layer_id=layer_id, role=role) + for _, layer_id, role, _ in buffers + ], + ) + ) + + pg_desc = SimpleNamespace( + pool_group_index=0, + num_slots=4, + pools=[ + SimpleNamespace(pool_index=pi, base_address=base, slot_bytes=sb) + for pi, sb in enumerate(slot_bytes_list) + ], + slot_desc=SimpleNamespace( + variants=[SimpleNamespace(layer_group_id=0, coalesced_buffers=coalesced_buffers)] + ), + ) + + impl = SimpleNamespace( + layer_grouping=((0, 1),), + init_config=SimpleNamespace( + tokens_per_block=16, + layers=[SimpleNamespace(window_size=None), SimpleNamespace(window_size=None)], + ), + pool_group_descs=[pg_desc], + ) + return SimpleNamespace( + impl=impl, + pp_layers=[0, 1], + num_kv_heads_per_layer=[1, 1], + get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, + ) + + +def test_v2_builder_stamps_homogeneous_pool_kinds(): + """Split KV / INDEX_KEY pools get NHD and REPLICATED views respectively.""" + from types import SimpleNamespace + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + (1, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=256, size=128), + (1, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=384, size=128), + (1, Role.INDEX_KEY): SimpleNamespace(life_cycle_id=0, pool_index=1, offset=0, size=128), + } + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.NHD, Role.INDEX_KEY: MapperKind.REPLICATED}, + num_pools=2, + slot_bytes_list=(512, 128), + ) + + views = build_page_table_from_manager(manager).layer_groups[0].pool_views + assert len(views) == 2 + assert views[0].pool_role == frozenset({"key", "value"}) + assert views[0].mapper_kind == MapperKind.NHD + assert get_unique_layers(views[0]) == {0, 1} + assert views[1].pool_role == frozenset({"index_key"}) + assert views[1].mapper_kind == MapperKind.REPLICATED + assert get_unique_layers(views[1]) == {1} + + +def test_v2_builder_splits_mixed_kind_pool_into_per_class_views(): + """A pool coalescing several role classes yields one view per class. + + Miniature of the MiniMax M3 layout at TP degrees where + K == V == INDEX_KEY bytes per block: V2 storage coalesces all three + into one pool and the slot interleaves the sparse layer's INDEX_KEY + between the layers' K/V regions. The builder must split the pool into + a NHD K/V view (non-uniform layer offsets, uniform per-layer size) + and a REPLICATED INDEX_KEY view covering only the sparse layer. + """ + from types import SimpleNamespace + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + (0, Role.INDEX_KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=256, size=128), + (1, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=384, size=128), + (1, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=512, size=128), + } + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.NHD, Role.INDEX_KEY: MapperKind.REPLICATED}, + ) + + views = build_page_table_from_manager(manager).layer_groups[0].pool_views + assert len(views) == 2 + kv_view, idx_view = views + # Both views address the same physical pool. + assert kv_view.pool_idx == idx_view.pool_idx == 0 + + assert kv_view.pool_role == frozenset({"key", "value"}) + assert kv_view.mapper_kind == MapperKind.NHD + assert get_unique_layers(kv_view) == {0, 1} + assert kv_view.bytes_per_layer == 256 + kv_starts, kv_bytes_per_layer = get_layer_byte_ranges(kv_view) + # Layer stride is non-uniform: INDEX_KEY interleaves after layer 0. + assert kv_starts == {0: 0, 1: 384} + assert kv_bytes_per_layer == 256 + + assert idx_view.pool_role == frozenset({"index_key"}) + assert idx_view.mapper_kind == MapperKind.REPLICATED + assert get_unique_layers(idx_view) == {0} + assert idx_view.bytes_per_layer == 128 + idx_starts, _ = get_layer_byte_ranges(idx_view) + assert idx_starts == {0: 256} + + +def test_v2_builder_validates_role_mapper_declaration(): + from types import SimpleNamespace + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + (1, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=256, size=128), + (1, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=384, size=128), + } + + # Default INDEXED declaration keeps the whole-slot HND view. + views = ( + build_page_table_from_manager(_make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED})) + .layer_groups[0] + .pool_views + ) + assert len(views) == 1 + assert views[0].mapper_kind == MapperKind.INDEXED + + # V2 managers must expose the capability method explicitly. + missing_capability = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}) + del missing_capability.get_disagg_role_mapper_kinds + with pytest.raises(AttributeError, match="get_disagg_role_mapper_kinds"): + build_page_table_from_manager(missing_capability) + + with pytest.raises(ValueError, match="must define Role.ALL"): + build_page_table_from_manager( + _make_fake_v2_manager(attrs, {Role.INDEX_KEY: MapperKind.REPLICATED}) + ) + + with pytest.raises(ValueError, match="Invalid disaggregation mapper kind 'HND'"): + build_page_table_from_manager(_make_fake_v2_manager(attrs, {Role.ALL: "HND"})) + + with pytest.raises(ValueError, match="INDEXED is only valid as the Role.ALL mapping"): + build_page_table_from_manager( + _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.NHD, Role.KEY: MapperKind.INDEXED}) + ) + + # INDEXED as the Role.ALL fallback coexists with side-cache roles that + # declare their own kind (the base-manager default for INDEX_KEY); with + # no INDEX_KEY buffers registered the declaration is inert. + views = ( + build_page_table_from_manager( + _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED}, + ) + ) + .layer_groups[0] + .pool_views + ) + assert len(views) == 1 + assert views[0].mapper_kind == MapperKind.INDEXED + + def test_mamba_layer_group_serialization(): from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup, PhysicalPool diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py new file mode 100644 index 000000000000..0fd729d1ae96 --- /dev/null +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -0,0 +1,617 @@ +# 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. + +"""MiniMax M3 heterogeneous-topology KV transfer tests. + +This drives the shared threaded NIXL harness (``kv_transfer_harness``) with +MiniMax M3 cache managers and validates both ordinary K/V and the +replicated INDEX_KEY cache. V2 storage coalesces buffers purely by (life +cycle, size), so at TP degrees where K == V == INDEX_KEY bytes per block the +index-K cache shares the K/V pool and the slot interleaves per layer; at +other degrees it gets its own pool. The disagg page-table builder splits each +pool into per-mapper-kind views, so both layouts (and transfers between them) +are exercised by the topology matrix below. +""" + +from collections.abc import Sequence + +import kv_transfer_harness as transfer_harness +import pytest +import torch + +from tensorrt_llm import Mapping +from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool, get_pool_bytes +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import KvCacheConfig + +NUM_LAYERS = 4 +NUM_KV_HEADS = 2 +HEAD_DIM = 128 +INDEX_DIM = 128 +TOKENS_PER_BLOCK = transfer_harness.TOKENS_PER_BLOCK +SPARSE_LAYERS = [3] + + +class _FakeKVCache: + num_blocks = 3 + + @staticmethod + def get_base_page_indices(_pool_id): + return [4, 5, 6, -1] + + +class _ShortFakeKVCache: + num_blocks = 2 + + @staticmethod + def get_base_page_indices(_pool_id): + return [7, 8, -1, -1] + + +def test_minimax_cache_indices_support_block_count_limit() -> None: + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + manager.kv_cache_map = {7: _FakeKVCache(), 8: _ShortFakeKVCache()} + + assert manager._get_batch_cache_indices_by_pool_id([7]) == [[4, 5, 6, -1]] + assert manager._get_batch_cache_indices_by_pool_id([7], num_blocks_per_seq=[2]) == [[4, 5]] + assert manager._get_batch_cache_indices_by_pool_id([7], num_blocks_per_seq=[99]) == [[4, 5, 6]] + assert manager._get_batch_cache_indices_by_pool_id([7, 8]) == [ + [4, 5, 6, -1], + [7, 8, -1, -1], + ] + assert manager.get_block_ids_per_seq([7]).tolist() == [[4, 5, 6, 0]] + assert manager.get_block_ids_per_seq([7, 8]).tolist() == [ + [4, 5, 6, 0], + [7, 8, 0, 0], + ] + + +def test_v2_disagg_role_mapper_kind_defaults() -> None: + manager = object.__new__(KVCacheManagerV2) + + assert manager.get_disagg_role_mapper_kinds() == { + Role.ALL: MapperKind.INDEXED, + Role.INDEX_KEY: MapperKind.REPLICATED, + } + + +def test_minimax_disagg_role_mapper_kinds() -> None: + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + + role_mapper_kinds = manager.get_disagg_role_mapper_kinds() + + assert role_mapper_kinds == { + Role.ALL: MapperKind.NHD, + Role.INDEX_KEY: MapperKind.REPLICATED, + } + + +def test_minimax_disagg_rejects_unmanaged_index_value(monkeypatch) -> None: + def fake_base_init(self, *args, **kwargs): + self.is_disagg = kwargs.get("is_disagg", False) + self.layer_offsets = {layer_id: layer_id for layer_id in range(kwargs["num_layers"])} + + monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init) + + with pytest.raises(ValueError, match="requires disable_index_value=True"): + MiniMaxM3KVCacheManagerV2( + num_layers=4, + sparse_layer_ids=[3], + disable_index_value_layer_ids=[], + sparse_index_dim=INDEX_DIM, + is_disagg=True, + ) + + +def _create_manager( + mapping: Mapping, dtype: DataType, sparse_layers: list[int] | None = None +) -> MiniMaxM3KVCacheManagerV2: + max_num_tokens = 2048 + kv_cache_dtype = { + DataType.FP8: "fp8", + DataType.NVFP4: "nvfp4", + }.get(dtype, "auto") + return MiniMaxM3KVCacheManagerV2( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_num_tokens, + event_buffer_max_size=0, + dtype=kv_cache_dtype, + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=NUM_LAYERS, + num_kv_heads=NUM_KV_HEADS, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=transfer_harness.MAX_SEQ_LEN, + max_batch_size=transfer_harness.MAX_BATCH_SIZE, + mapping=mapping, + dtype=dtype, + vocab_size=transfer_harness.VOCAB_SIZE, + max_num_tokens=max_num_tokens, + sparse_layer_ids=sparse_layers if sparse_layers is not None else SPARSE_LAYERS, + disable_index_value_layer_ids=sparse_layers if sparse_layers is not None else SPARSE_LAYERS, + sparse_index_dim=INDEX_DIM, + ) + + +def test_minimax_m3_pool_view_scheme_coalesced_vs_separate() -> None: + """Pool composition varies with TP; the view scheme must not. + + TP=2 makes kv_heads_per_rank == 1, so K == V == INDEX_KEY bytes per + block and V2 coalesces all three into one pool whose slot interleaves + the sparse layer's index-K between K/V regions (non-uniform layer + stride). TP=1 doubles K/V, so index-K gets its own pool and strides are + uniform. Both layouts must yield the same two per-class views. + """ + from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + build_page_table_from_manager, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_byte_ranges + + # --- TP=2: coalesced, interleaved slot --- + manager = _create_manager( + Mapping(world_size=2, rank=0, tp_size=2, pp_size=1), DataType.BF16, sparse_layers=[1] + ) + try: + lg = build_page_table_from_manager(manager).layer_groups[0] + assert len(lg.pool_views) == 2 + kv_view = next(pv for pv in lg.pool_views if pv.mapper_kind == MapperKind.NHD) + idx_view = next(pv for pv in lg.pool_views if pv.mapper_kind == MapperKind.REPLICATED) + assert kv_view.pool_role == frozenset({str(Role.KEY), str(Role.VALUE)}) + assert idx_view.pool_role == frozenset({str(Role.INDEX_KEY)}) + # Coalesced: both views address the same physical pool. + assert kv_view.pool_idx == idx_view.pool_idx + starts, bytes_per_layer = get_layer_byte_ranges(kv_view) + unit = bytes_per_layer // 2 # one K or V buffer per block + # Slot: L0K L0V L1K L1V L1IDX L2K L2V L3K L3V — the sparse layer's + # index-K makes the K/V layer stride non-uniform. + assert starts == {0: 0, 1: 2 * unit, 2: 5 * unit, 3: 7 * unit} + idx_starts, idx_bytes_per_layer = get_layer_byte_ranges(idx_view) + assert idx_starts == {1: 4 * unit} + assert idx_bytes_per_layer == unit + finally: + manager.shutdown() + + # --- TP=1: separate pools, uniform strides --- + manager = _create_manager( + Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), DataType.BF16, sparse_layers=[1] + ) + try: + lg = build_page_table_from_manager(manager).layer_groups[0] + assert len(lg.pool_views) == 2 + kv_view = next(pv for pv in lg.pool_views if pv.mapper_kind == MapperKind.NHD) + idx_view = next(pv for pv in lg.pool_views if pv.mapper_kind == MapperKind.REPLICATED) + assert kv_view.pool_idx != idx_view.pool_idx + starts, bytes_per_layer = get_layer_byte_ranges(kv_view) + assert starts == {i: i * bytes_per_layer for i in range(NUM_LAYERS)} + finally: + manager.shutdown() + + +def _create_managers( + tp: int, + pp: int, + enable_dp: bool, + dtype: DataType = DataType.BF16, + sparse_layers: list[int] | None = None, +) -> list[MiniMaxM3KVCacheManagerV2]: + return [ + _create_manager( + Mapping( + world_size=tp * pp, + rank=rank, + tp_size=tp, + pp_size=pp, + enable_attention_dp=enable_dp, + ), + dtype, + sparse_layers, + ) + for rank in range(tp * pp) + ] + + +def _zero_physical_pools(manager: MiniMaxM3KVCacheManagerV2) -> None: + page_table = KVRegionExtractorV1(manager).page_table + unique_pools = {} + for pool_group in page_table.pool_groups: + for pool in pool_group.pools: + unique_pools[pool.base_address] = max( + unique_pools.get(pool.base_address, 0), get_pool_bytes(pool) + ) + for base_address, size in unique_pools.items(): + tensor = convert_to_torch_tensor(TensorWrapper(base_address, DataType.INT8, [size])) + tensor.zero_() + + +def _get_nvfp4_scale_view( + manager: MiniMaxM3KVCacheManagerV2, layer_idx: int +) -> torch.Tensor | None: + if manager.dtype != DataType.NVFP4: + return None + page_table = KVRegionExtractorV1(manager).page_table + local_layer_id = manager.layer_offsets[layer_idx] + scale_roles = frozenset({"key_block_scale", "value_block_scale"}) + for layer_group_id, layer_group in enumerate(page_table.layer_groups): + for pool_view in getattr(layer_group, "pool_views", ()): + # Scale buffers land in their own pool (smaller size class), so + # selecting the pool by role set suffices; every entry for the + # layer inside that pool is a scale buffer. + if not (pool_view.pool_role & scale_roles): + continue + matching_entries = [ + entry + for entry in pool_view.buffer_entries + if int(entry["local_layer_id"]) == local_layer_id + ] + if not matching_entries: + continue + pool = get_physical_pool(page_table, layer_group_id, pool_view.pool_idx) + raw_slots = convert_to_torch_tensor( + TensorWrapper( + pool.base_address, + DataType.INT8, + [pool.num_slots, pool.slot_bytes], + ) + ) + start = min(int(entry["offset"]) for entry in matching_entries) + end = max(int(entry["offset"] + entry["size"]) for entry in matching_entries) + return raw_slots[:, start:end] + raise AssertionError(f"missing NVFP4 scale view for layer {layer_idx}") + + +def _fill_position_dependent( + tensor: torch.Tensor, + *, + layer_idx: int, + first_global_head: int, +) -> None: + """Fill ``[block, role, token, head, dim]`` with exact small integers.""" + block = torch.arange(tensor.shape[0], device=tensor.device)[:, None, None, None, None] + role = torch.arange(tensor.shape[1], device=tensor.device)[None, :, None, None, None] + token = torch.arange(tensor.shape[2], device=tensor.device)[None, None, :, None, None] + head = (first_global_head + torch.arange(tensor.shape[3], device=tensor.device))[ + None, None, None, :, None + ] + dim = torch.arange(tensor.shape[4], device=tensor.device)[None, None, None, None, :] + values = (layer_idx * 17 + block * 11 + role * 13 + token * 3 + head * 19 + dim) % 97 + tensor.copy_(values.to(tensor.dtype)) + + +def _as_nvfp4_scale_tensor( + manager: MiniMaxM3KVCacheManagerV2, + layer_idx: int, +) -> torch.Tensor | None: + scale_view = _get_nvfp4_scale_view(manager, layer_idx) + if scale_view is None: + return None + local_layer_id = manager.layer_offsets[layer_idx] + local_heads = manager.num_kv_heads_per_layer[local_layer_id] + bytes_per_token_head = HEAD_DIM // 16 + return scale_view.view( + scale_view.shape[0], + 2, + TOKENS_PER_BLOCK, + local_heads, + bytes_per_token_head, + ) + + +def _valid_indices( + manager: MiniMaxM3KVCacheManagerV2, + request_id: int, + layer_idx: int, +) -> list[int]: + return [ + index for index in manager.get_batch_cache_indices([request_id], layer_idx)[0] if index >= 0 + ] + + +def _first_global_head(manager: MiniMaxM3KVCacheManagerV2) -> int: + """Map a TP rank to its first logical head, including duplicated heads.""" + if manager.mapping.enable_attention_dp: + return 0 + return manager.mapping.tp_rank * NUM_KV_HEADS // manager.mapping.tp_size + + +def _find_ctx_source( + ctx_managers: list[MiniMaxM3KVCacheManagerV2], + *, + ctx_tp: int, + ctx_enable_dp: bool, + request_idx: int, + layer_idx: int, + global_head: int, +) -> tuple[MiniMaxM3KVCacheManagerV2, int]: + """Return the context manager/local head that owns one logical KV head.""" + for manager in ctx_managers: + if layer_idx not in manager.pp_layers: + continue + tp_rank = manager.mapping.tp_rank + local_layer_id = manager.layer_offsets[layer_idx] + local_heads = manager.num_kv_heads_per_layer[local_layer_id] + if ctx_enable_dp: + if tp_rank == request_idx % ctx_tp: + return manager, global_head + continue + first_global_head = _first_global_head(manager) + if first_global_head <= global_head < first_global_head + local_heads: + return manager, global_head - first_global_head + raise AssertionError( + f"missing context source for request={request_idx}, layer={layer_idx}, " + f"global_head={global_head}" + ) + + +def _initialize_cache( + managers: Sequence[MiniMaxM3KVCacheManagerV2], + _tp: int, + seed_base: int = 0, + fill_random: bool = True, +) -> None: + del seed_base + for manager in managers: + _zero_physical_pools(manager) + if not fill_random: + continue + + for layer_idx in manager.pp_layers: + kv = manager.get_buffers(layer_idx, kv_layout="NHD") + first_global_head = _first_global_head(manager) + _fill_position_dependent( + kv, + layer_idx=layer_idx, + first_global_head=first_global_head, + ) + + index_key = manager.get_index_k_buffer(layer_idx) + if index_key is not None: + index_tensor = index_key.unsqueeze(1) + _fill_position_dependent( + index_tensor, + layer_idx=layer_idx, + first_global_head=0, + ) + + scale_tensor = _as_nvfp4_scale_tensor(manager, layer_idx) + if scale_tensor is not None: + _fill_position_dependent( + scale_tensor, + layer_idx=layer_idx, + first_global_head=first_global_head, + ) + + +def _verify_cache( + request_lengths: list[int], + ctx_managers: Sequence[MiniMaxM3KVCacheManagerV2], + gen_managers: Sequence[MiniMaxM3KVCacheManagerV2], + ctx_tp: int, + ctx_pp: int, + gen_tp: int, + gen_pp: int, + ctx_enable_dp: bool, + gen_enable_dp: bool, + ctx_request_ids: list[int], + gen_request_ids: list[int], +) -> None: + del ctx_pp + + for req_idx, _request_length in enumerate(request_lengths): + gen_request_id = gen_request_ids[req_idx] + for gen_rank, manager in enumerate(gen_managers): + tp_rank = gen_rank % gen_tp + if gen_enable_dp and req_idx % gen_tp != tp_rank: + continue + + for layer_idx in manager.pp_layers: + gen_indices = _valid_indices(manager, gen_request_id, layer_idx) + assert gen_indices + + gen_kv = manager.get_buffers(layer_idx, kv_layout="NHD")[gen_indices] + local_heads = gen_kv.shape[3] + first_global_head = _first_global_head(manager) + for kv_idx in range(2): + for local_head in range(local_heads): + global_head = first_global_head + local_head + ctx_manager, ctx_local_head = _find_ctx_source( + ctx_managers, + ctx_tp=ctx_tp, + ctx_enable_dp=ctx_enable_dp, + request_idx=req_idx, + layer_idx=layer_idx, + global_head=global_head, + ) + ctx_indices = _valid_indices( + ctx_manager, ctx_request_ids[req_idx], layer_idx + ) + ctx_kv = ctx_manager.get_buffers(layer_idx, kv_layout="NHD") + torch.testing.assert_close( + gen_kv[:, kv_idx, :, local_head, :], + ctx_kv[ctx_indices, kv_idx, :, ctx_local_head, :], + rtol=0, + atol=0, + ) + + index_key = manager.get_index_k_buffer(layer_idx) + if index_key is not None: + ctx_manager, _ = _find_ctx_source( + ctx_managers, + ctx_tp=ctx_tp, + ctx_enable_dp=ctx_enable_dp, + request_idx=req_idx, + layer_idx=layer_idx, + global_head=0, + ) + ctx_indices = _valid_indices(ctx_manager, ctx_request_ids[req_idx], layer_idx) + ctx_index_key = ctx_manager.get_index_k_buffer(layer_idx) + assert ctx_index_key is not None + torch.testing.assert_close( + index_key[gen_indices], + ctx_index_key[ctx_indices], + rtol=0, + atol=0, + ) + + gen_scales = _as_nvfp4_scale_tensor(manager, layer_idx) + if gen_scales is not None: + for local_head in range(local_heads): + global_head = first_global_head + local_head + ctx_manager, ctx_local_head = _find_ctx_source( + ctx_managers, + ctx_tp=ctx_tp, + ctx_enable_dp=ctx_enable_dp, + request_idx=req_idx, + layer_idx=layer_idx, + global_head=global_head, + ) + ctx_indices = _valid_indices( + ctx_manager, ctx_request_ids[req_idx], layer_idx + ) + ctx_scales = _as_nvfp4_scale_tensor(ctx_manager, layer_idx) + assert ctx_scales is not None + torch.testing.assert_close( + gen_scales[gen_indices, :, :, local_head, :], + ctx_scales[ctx_indices, :, :, ctx_local_head, :], + rtol=0, + atol=0, + ) + + +# Production is expected to use TEP/DEP context and DEP generation. Bias the +# committed matrix toward head-matched layouts while retaining representative +# head-mismatch, degree-8 duplication, fan-in/fan-out, and PP2 coverage. +HEAD_MATCHED_BF16_TOPOLOGIES = [ + (1, 1, True, 1, 1, True, "dep1_to_dep1"), + (2, 1, True, 2, 1, True, "dep2_to_dep2"), + (4, 1, True, 4, 1, True, "dep4_to_dep4"), + (8, 1, True, 8, 1, True, "dep8_to_dep8"), + (1, 1, True, 8, 1, True, "dep1_to_dep8"), + (8, 1, True, 1, 1, True, "dep8_to_dep1"), + (1, 1, False, 4, 1, True, "tep1_to_dep4"), + (1, 1, False, 8, 1, True, "tep1_to_dep8"), + (1, 2, False, 2, 1, True, "tp1pp2_to_dep2"), +] + +HEAD_MISMATCHED_BF16_TOPOLOGIES = [ + (2, 1, False, 2, 1, True, "tep2_to_dep2"), + (4, 1, False, 4, 1, True, "tep4_to_dep4"), + (8, 1, False, 8, 1, True, "tep8_to_dep8"), + (8, 1, False, 1, 1, True, "tep8_to_dep1"), +] + +BF16_TOPOLOGIES = HEAD_MATCHED_BF16_TOPOLOGIES + HEAD_MISMATCHED_BF16_TOPOLOGIES + +# Smaller quantized-cache sets cover packed K/V geometry across the same +# production directions without repeating the BF16 topology matrix. NVFP4 +# additionally exercises its block-scale pools. +QUANTIZED_TOPOLOGIES = [ + (4, 1, True, 4, 1, True, "dep4_to_dep4"), + (8, 1, True, 1, 1, True, "dep8_to_dep1"), + (1, 1, False, 8, 1, True, "tep1_to_dep8"), + (4, 1, False, 4, 1, True, "tep4_to_dep4"), + (8, 1, False, 1, 1, True, "tep8_to_dep1"), +] + +CACHE_CASES = [(*topology[:6], DataType.BF16, topology[6]) for topology in BF16_TOPOLOGIES] + [ + (*topology[:6], dtype, f"{prefix}_{topology[6]}") + for dtype, prefix in ((DataType.FP8, "fp8"), (DataType.NVFP4, "nvfp4")) + for topology in QUANTIZED_TOPOLOGIES +] + + +@pytest.mark.cuda +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,ctx_pp,ctx_enable_dp,gen_tp,gen_pp,gen_enable_dp,cache_dtype", + [case[:7] for case in CACHE_CASES], + ids=[case[7] for case in CACHE_CASES], +) +@pytest.mark.parametrize( + "update_before_transfer", + [True, False], + ids=["update_before", "update_after"], +) +def test_minimax_m3_kv_transfer( + ctx_tp, + ctx_pp, + ctx_enable_dp, + gen_tp, + gen_pp, + gen_enable_dp, + cache_dtype, + update_before_transfer, +) -> None: + transfer_harness.run_kv_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + update_before_transfer=update_before_transfer, + manager_factory=lambda tp, pp, enable_dp: _create_managers(tp, pp, enable_dp, cache_dtype), + init_fn=_initialize_cache, + verify_fn=_verify_cache, + ) + + +# Multiple sparse layers spread across PP stages: the replicated index-key +# pool overlaps only partially between peers, exercising the layer-strided +# ReplicatedMapper offsets (a single-sparse-layer model always fully +# overlaps and would degenerate to a whole-slot copy). +MULTI_SPARSE_LAYERS = [1, 2, 3] + + +@pytest.mark.cuda +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,ctx_pp,ctx_enable_dp,gen_tp,gen_pp,gen_enable_dp", + [ + (1, 1, True, 1, 2, False), # ctx full index pool -> per-PP-stage subsets + (1, 2, False, 1, 1, True), # per-PP-stage subsets -> full index pool + (2, 1, False, 1, 2, False), # TEP fan-in + PP subset on the same transfer + ], + ids=["dep1_to_tp1pp2", "tp1pp2_to_dep1", "tep2_to_tp1pp2"], +) +def test_minimax_m3_multi_sparse_layer_pp_transfer( + ctx_tp, + ctx_pp, + ctx_enable_dp, + gen_tp, + gen_pp, + gen_enable_dp, +) -> None: + transfer_harness.run_kv_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=ctx_enable_dp, + gen_enable_dp=gen_enable_dp, + update_before_transfer=True, + manager_factory=lambda tp, pp, enable_dp: _create_managers( + tp, pp, enable_dp, DataType.BF16, sparse_layers=MULTI_SPARSE_LAYERS + ), + init_fn=_initialize_cache, + verify_fn=_verify_cache, + ) diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index d2329d74e1aa..80134bde5554 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -1,10 +1,28 @@ +# 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. + import numpy as np import pytest +import tensorrt_llm._torch.disaggregation.native.peer as peer_module +from tensorrt_llm._torch.disaggregation.base.region import MemRegionGroup, SpecRegion from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ( - HeadMatchMapper, - HeadMismatchMapper, - IdentityMapper, + HNDHeadMismatchMapper, + IntactMapper, + NHDHeadMismatchMapper, + ReplicatedMapper, ) from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar @@ -16,6 +34,7 @@ KVCachePageTable, LocalLayer, MambaLayerGroup, + MapperKind, PhysicalPool, PhysicalPoolGroup, PoolView, @@ -31,21 +50,22 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None): if global_layer_ids is None: global_layer_ids = [0, 1] - # Build buffer entries: K + V per local layer - buffer_size = 256 # bytes per buffer entry (arbitrary for tests) - entries = [] - for i in range(len(global_layer_ids)): - base_offset = i * buffer_size * 2 - entries.append((i, base_offset, buffer_size)) - entries.append((i, base_offset + buffer_size, buffer_size)) - buffer_entries = np.array(entries, dtype=BUFFER_ENTRY_DTYPE) - local_layers = [ LocalLayer(local_layer_id=i, global_layer_id=gid) for i, gid in enumerate(global_layer_ids) ] - pool_views = [ - PoolView(pool_idx=pi, buffer_entries=buffer_entries) for pi in range(len(pool_ptrs)) - ] + # Build buffer entries: K + V per local layer, sized so the layers fill + # the pool slot (keeps entry geometry consistent with slot_bytes). + pool_views = [] + for pi, bs in enumerate(block_bytes): + buffer_size = bs // (len(global_layer_ids) * 2) + entries = [] + for i in range(len(global_layer_ids)): + base_offset = i * buffer_size * 2 + entries.append((i, base_offset, buffer_size)) + entries.append((i, base_offset + buffer_size, buffer_size)) + pool_views.append( + PoolView(pool_idx=pi, buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE)) + ) physical_pools = [ PhysicalPool(base_address=ptr, slot_bytes=bs, num_slots=128) for ptr, bs in zip(pool_ptrs, block_bytes) @@ -359,7 +379,15 @@ def test_peer_registrar_get_kv_map_identity(): ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) - assert isinstance(mapper, IdentityMapper) + # Full contiguous overlap: one whole-region fragment per block. + assert isinstance(mapper, IntactMapper) + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=1024)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=1024)), + ) + assert not isinstance(pair, list) + assert pair.src.memory.ptrs.tolist() == [1000] + assert pair.src.memory.bytes_per_region == 1024 def test_peer_registrar_get_kv_map_head_match(): @@ -380,12 +408,23 @@ def test_peer_registrar_get_kv_map_head_match(): ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) - assert isinstance(mapper, HeadMatchMapper) + # Partial layer overlap with matching heads: a single shifted fragment. + assert isinstance(mapper, IntactMapper) + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=1024)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=512)), + ) + # Overlap is global layer 1: self slot holds layers [0, 1] (512B each), + # peer slot holds only layer 1. + assert pair.src.memory.ptrs.tolist() == [1512] + assert pair.dst.memory.ptrs.tolist() == [2000] + assert pair.src.memory.bytes_per_region == 512 def test_peer_registrar_get_kv_map_head_mismatch(): self_rankinfo = make_rankinfo(instance_name="local", page_table=make_page_table()) reg = _make_peer_registrar(self_rankinfo) + # Twice the KV heads per rank -> twice the slot bytes on the peer side. peer_ri = make_rankinfo( instance_name="peer", instance_rank=3, @@ -397,11 +436,156 @@ def test_peer_registrar_get_kv_map_head_mismatch(): tokens_per_block=16, dims_per_head=8, layer_num_per_pp=[2], - page_table=make_page_table(), + page_table=make_page_table(block_bytes=[2048]), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) - assert isinstance(mapper, HeadMismatchMapper) + assert isinstance(mapper, HNDHeadMismatchMapper) + + +def test_nhd_head_mismatch_mapper_slices_each_token(): + self_ri = make_rankinfo( + instance_name="local", + tp_size=2, + tp_rank=0, + dp_size=2, + dp_rank=0, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=2, + element_bytes=2, + enable_attention_dp=True, + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=2, + tp_rank=0, + kv_heads_per_rank=1, + tokens_per_block=2, + dims_per_head=2, + element_bytes=2, + ) + mapper = NHDHeadMismatchMapper( + src_layer_offsets=[0], + dst_layer_offsets=[0], + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=32, + peer_bytes_per_layer=16, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) + + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=32)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=16)), + ) + + # Source NHD has two heads per token: select head 0 at offsets 0 and 8 + # for K, then 16 and 24 for V. Destination has one head per token. + assert pair.src.memory.ptrs.tolist() == [1000, 1008, 1016, 1024] + assert pair.dst.memory.ptrs.tolist() == [2000, 2004, 2008, 2012] + assert pair.src.memory.bytes_per_region == 4 + assert pair.dst.memory.bytes_per_region == 4 + + +def test_nhd_head_mismatch_mapper_uses_scale_pool_geometry(): + self_ri = make_rankinfo( + tp_size=2, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=128, + element_bytes=0.5, + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + kv_heads_per_rank=1, + tokens_per_block=2, + dims_per_head=128, + element_bytes=0.5, + ) + mapper = NHDHeadMismatchMapper( + src_layer_offsets=[0], + dst_layer_offsets=[0], + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=8, + peer_bytes_per_layer=4, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) + + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=8)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=4)), + ) + + assert pair.src.memory.ptrs.tolist() == [1000, 1002, 1004, 1006] + assert pair.dst.memory.ptrs.tolist() == [2000, 2001, 2002, 2003] + assert pair.src.memory.bytes_per_region == 1 + assert pair.dst.memory.bytes_per_region == 1 + + +def test_replicated_mapper_ignores_kv_head_mismatch(): + self_pt = make_page_table(global_layer_ids=[0]) + peer_pt = make_page_table(global_layer_ids=[0]) + for page_table in (self_pt, peer_pt): + view = page_table.layer_groups[0].pool_views[0] + view.pool_role = frozenset({"index_key"}) + view.mapper_kind = MapperKind.REPLICATED + + self_rankinfo = make_rankinfo(instance_name="local", kv_heads_per_rank=1, page_table=self_pt) + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=3, + tp_size=1, + kv_heads_per_rank=8, + layer_num_per_pp=[1], + page_table=peer_pt, + ) + mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + assert isinstance(mapper, ReplicatedMapper) + + +@pytest.mark.parametrize("peer_dp_rank", [0, 1, 3, 7]) +def test_replicated_pool_owner_rotates_by_destination_dp_rank(peer_dp_rank): + """Exactly one owner per fan-in group, rotated by destination DP rank. + + Mirrors the C++ MLACacheFormatter::needSendCache pairing so the + replicated traffic spreads across local ranks for multi-DP generation. + """ + page_table = make_page_table(global_layer_ids=[0]) + view = page_table.layer_groups[0].pool_views[0] + view.mapper_kind = MapperKind.REPLICATED + + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=8, + dp_size=8, + dp_rank=peer_dp_rank, + enable_attention_dp=True, + page_table=page_table, + layer_num_per_pp=[1], + ) + overlap = PeerOverlap() + ownership = [] + for tp_rank in range(8): + self_ri = make_rankinfo( + instance_name="local", + tp_size=8, + tp_rank=tp_rank, + page_table=page_table, + layer_num_per_pp=[1], + ) + reg = _make_peer_registrar(self_ri) + ownership.append(reg.should_send_pool(overlap, peer_ri, 0, 0)) + + # ratio = self_tp(8) / peer_tp_per_dp(1) = 8: exactly one owner, at the + # slot selected by the destination dp rank. + assert sum(ownership) == 1 + assert ownership.index(True) == peer_dp_rank % 8 def test_peer_registrar_tpb_divisible_warns_but_compatible(): @@ -432,3 +616,379 @@ def test_peer_registrar_tpb_not_divisible_raises(): ) with pytest.raises(ValueError): reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + +@pytest.mark.parametrize("mapper_kind", [MapperKind.NHD, MapperKind.REPLICATED]) +def test_peer_registrar_exact_tpb_mapper_rejects_divisible_mismatch(mapper_kind): + self_pt = make_page_table() + peer_pt = make_page_table() + self_pt.layer_groups[0].pool_views[0].mapper_kind = mapper_kind + peer_pt.layer_groups[0].pool_views[0].mapper_kind = mapper_kind + self_ri = make_rankinfo(page_table=self_pt, tokens_per_block=16) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=7, + page_table=peer_pt, + tokens_per_block=32, + ) + reg = _make_peer_registrar(self_ri) + + with pytest.raises(ValueError): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + +def test_intact_mapper_region_size_mismatch_raises(): + with pytest.raises(ValueError, match="cache region size mismatch"): + IntactMapper([0], [0], 256, 128) + + +def test_replicated_mapper_per_layer_size_mismatch_raises(): + with pytest.raises(ValueError, match="Replicated cache region size mismatch"): + ReplicatedMapper( + src_layer_offsets=[0, 128], + dst_layer_offsets=[0, 64], + self_bytes_per_layer=128, + peer_bytes_per_layer=64, + ) + + +def test_replicated_mapper_selects_partial_layer_range(): + """PP mismatch selects the overlap layers via explicit offsets. + + The peer slot holds a superset of layers; only the overlap moves, and + contiguous layers on both sides merge into one fragment. + """ + mapper = ReplicatedMapper( + src_layer_offsets=[0, 128], # self slot: 2 layers x 128B + dst_layer_offsets=[128, 256], # peer slot: 3 layers x 128B, overlap at 1..2 + self_bytes_per_layer=128, + peer_bytes_per_layer=128, + ) + + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=256)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=384)), + ) + + assert pair.src.memory.ptrs.tolist() == [1000] + assert pair.dst.memory.ptrs.tolist() == [2128] + assert pair.src.memory.bytes_per_region == 256 + assert pair.dst.memory.bytes_per_region == 256 + + +def test_intact_mapper_splits_non_contiguous_runs(): + """Interleaved slots (another role class between layers) split runs. + + Source layers sit at non-uniform strides (something interleaves after + layer 0); destination is densely packed. Runs must break where either + side is discontiguous, and each fragment must carry the run's bytes. + """ + mapper = IntactMapper( + src_layer_offsets=[0, 192, 320], # gap after layer 0 (64B interleaved) + dst_layer_offsets=[0, 128, 256], # dense + self_bytes_per_layer=128, + peer_bytes_per_layer=128, + ) + + pairs = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=448)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=384)), + ) + + assert isinstance(pairs, list) and len(pairs) == 2 + # Run 1: layer 0 alone (src gap breaks the run). + assert pairs[0].src.memory.ptrs.tolist() == [1000] + assert pairs[0].dst.memory.ptrs.tolist() == [2000] + assert pairs[0].src.memory.bytes_per_region == 128 + # Run 2: layers 1-2 contiguous on both sides -> merged. + assert pairs[1].src.memory.ptrs.tolist() == [1192] + assert pairs[1].dst.memory.ptrs.tolist() == [2128] + assert pairs[1].src.memory.bytes_per_region == 256 + + +def test_nhd_mapper_rejects_non_divisible_region_geometry(): + self_ri = make_rankinfo(kv_heads_per_rank=2, tokens_per_block=2) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + kv_heads_per_rank=4, + tokens_per_block=2, + ) + + with pytest.raises(ValueError, match="not evenly divisible"): + NHDHeadMismatchMapper( + src_layer_offsets=[0], + dst_layer_offsets=[0], + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=17, + peer_bytes_per_layer=32, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) + + +def test_nhd_mapper_rejects_tokens_per_block_mismatch(): + self_ri = make_rankinfo(tokens_per_block=2) + peer_ri = make_rankinfo(instance_name="peer", tokens_per_block=4) + + with pytest.raises(ValueError, match="requires equal tokens_per_block"): + NHDHeadMismatchMapper( + src_layer_offsets=[0], + dst_layer_offsets=[0], + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=32, + peer_bytes_per_layer=64, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) + + +def test_get_buffers_per_layer_rejects_non_uniform_nhd_entries(): + pool_view = PoolView( + pool_idx=0, + buffer_entries=np.array( + [(0, 0, 16), (0, 16, 16), (1, 32, 16)], + dtype=BUFFER_ENTRY_DTYPE, + ), + mapper_kind=MapperKind.NHD, + ) + + with pytest.raises(ValueError, match="layer_group=3, pool=4"): + PeerRegistrar._get_buffers_per_layer( + pool_view, + layer_group_id=3, + pool_idx=4, + ) + + +def test_non_uniform_view_geometry_rejected_for_all_kinds(): + """Entries-driven views require uniform per-layer regions for every kind. + + Under the unified contract INDEXED views are no longer exempt: a view + whose layers have different region sizes cannot be addressed per layer + and must fail loudly instead of transferring garbage. + """ + entries = np.array( + [(0, 0, 16), (0, 16, 16), (1, 32, 16)], + dtype=BUFFER_ENTRY_DTYPE, + ) + self_pt = make_page_table() + peer_pt = make_page_table() + self_pt.layer_groups[0].pool_views[0].buffer_entries = entries + peer_pt.layer_groups[0].pool_views[0].buffer_entries = entries.copy() + self_ri = make_rankinfo(page_table=self_pt) + peer_ri = make_rankinfo(instance_name="peer", page_table=peer_pt) + reg = _make_peer_registrar(self_ri) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + with pytest.raises(ValueError, match="not uniform"): + reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + + +def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): + """Byte-aligned sub-byte head slicing is allowed on the HND path. + + Per-head bytes are integral here: tpb=16 x dims=8 x 0.5B = 64B. + """ + self_ri = make_rankinfo( + element_bytes=0.5, + kv_heads_per_rank=2, + tp_size=2, + page_table=make_page_table(), + ) + peer_ri = make_rankinfo( + instance_name="peer", + element_bytes=0.5, + kv_heads_per_rank=4, + tp_size=1, + page_table=make_page_table(block_bytes=[2048]), + ) + reg = _make_peer_registrar(self_ri) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + + assert isinstance(mapper, HNDHeadMismatchMapper) + + +def test_peer_registrar_rejects_misaligned_subbyte_head_mismatch(): + """Head slicing that lands mid-byte must fail at registration. + + tpb=1 x dims=1 x 0.5B = 0.5B per head is not byte-aligned. + """ + self_ri = make_rankinfo( + element_bytes=0.5, + kv_heads_per_rank=2, + tokens_per_block=1, + dims_per_head=1, + tp_size=2, + page_table=make_page_table(), + ) + peer_ri = make_rankinfo( + instance_name="peer", + element_bytes=0.5, + kv_heads_per_rank=4, + tokens_per_block=1, + dims_per_head=1, + tp_size=1, + page_table=make_page_table(block_bytes=[2048]), + ) + reg = _make_peer_registrar(self_ri) + + with pytest.raises(ValueError): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + +def test_peer_registrar_dispatches_nhd_mapper(): + self_pt = make_page_table() + peer_pt = make_page_table(block_bytes=[2048]) + self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + self_ri = make_rankinfo( + kv_heads_per_rank=2, + tp_size=2, + page_table=self_pt, + ) + peer_ri = make_rankinfo( + instance_name="peer", + kv_heads_per_rank=4, + tp_size=1, + page_table=peer_pt, + ) + reg = _make_peer_registrar(self_ri) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + + assert isinstance(mapper, NHDHeadMismatchMapper) + + +def test_peer_registrar_warns_for_nhd_head_mismatch(monkeypatch): + self_pt = make_page_table() + peer_pt = make_page_table(block_bytes=[2048]) + self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + self_ri = make_rankinfo(kv_heads_per_rank=2, tp_size=2, page_table=self_pt) + peer_ri = make_rankinfo( + instance_name="peer", + kv_heads_per_rank=4, + tp_size=1, + page_table=peer_pt, + ) + warnings = [] + monkeypatch.setattr( + peer_module.logger, + "warning_once", + lambda *message, key: warnings.append((" ".join(map(str, message)), key)), + ) + + _make_peer_registrar(self_ri).register( + peer_ri.instance_name, + peer_ri.instance_rank, + peer_ri, + ) + + assert len(warnings) == 1 + message, key = warnings[0] + assert "4 NIXL descriptors per transferred token per peer" in message + assert "local_kv_heads=2, peer_kv_heads=4" in message + assert key == "native-nhd-head-mismatch-2-4-4" + + +def test_peer_registrar_nhd_head_match_uses_intact_mapper(): + """NHD + equal heads takes the merged-run fast path, not per-token slicing. + + With a dedicated (non-interleaved) K/V pool the run merge collapses the + whole class region into a single fragment per block — the fragment-count + budget for separate layouts. + """ + self_pt = make_page_table() + peer_pt = make_page_table() + self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD + self_ri = make_rankinfo(page_table=self_pt) + peer_ri = make_rankinfo(instance_name="peer", instance_rank=9, page_table=peer_pt) + reg = _make_peer_registrar(self_ri) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + assert isinstance(mapper, IntactMapper) + assert not isinstance(mapper, NHDHeadMismatchMapper) + + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000, 5000]), bytes_per_region=1024)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000, 6000]), bytes_per_region=1024)), + ) + # Fully contiguous on both sides -> exactly one fragment per block. + assert not isinstance(pair, list) + assert pair.src.memory.ptrs.tolist() == [1000, 5000] + assert pair.dst.memory.ptrs.tolist() == [2000, 6000] + assert pair.src.memory.bytes_per_region == 1024 + + +def test_indexed_head_mismatch_subbyte_geometry_is_entries_derived(): + """HND head-mismatch byte math comes from slot bytes, never element_bytes. + + Emulates an NVFP4-like cache: element_bytes is fractional (0.5), but the + slot size registered by storage is whole bytes, so every derived offset + and fragment size must be an exact integer. + """ + # self: 2 kv heads/rank, per-head bytes = 16 tokens x 8 dims x 0.5B = 64. + self_ri = make_rankinfo( + tp_size=1, + tp_rank=0, + kv_heads_per_rank=2, + tokens_per_block=16, + dims_per_head=8, + element_bytes=0.5, + ) + # peer: twice the TP -> 1 kv head/rank, half the slot bytes. + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=2, + tp_rank=1, + kv_heads_per_rank=1, + tokens_per_block=16, + dims_per_head=8, + element_bytes=0.5, + ) + mapper = HNDHeadMismatchMapper( + src_layer_offsets=[0, 256], # 2 layers x 256B/layer (kv_factor 2 x 128B) + dst_layer_offsets=[0, 128], # 2 layers x 128B/layer (kv_factor 2 x 64B) + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=256, + peer_bytes_per_layer=128, + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) + + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=512)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=256)), + ) + # peer tp_rank=1 selects head 1 inside self's per-rank pair of heads: + # src_head_off = 1 x 64B; fragments = (layer, k/v) x 2 layers. + assert pair.src.memory.ptrs.tolist() == [1064, 1192, 1320, 1448] + assert pair.dst.memory.ptrs.tolist() == [2000, 2064, 2128, 2192] + assert pair.src.memory.bytes_per_region == 64 + assert pair.src.memory.ptrs.dtype == np.int64 + assert pair.dst.memory.ptrs.dtype == np.int64 + + +def test_indexed_head_mismatch_inconsistent_slot_geometry_raises(): + self_ri = make_rankinfo(tp_size=1, kv_heads_per_rank=2) + peer_ri = make_rankinfo(instance_name="peer", tp_size=2, kv_heads_per_rank=1) + with pytest.raises(ValueError, match="HND bytes per head mismatch"): + HNDHeadMismatchMapper( + src_layer_offsets=[0, 256], + dst_layer_offsets=[0, 256], + self_ri=self_ri, + peer_ri=peer_ri, + self_bytes_per_layer=256, + peer_bytes_per_layer=256, # should be 128 for half the heads + self_buffers_per_layer=2, + peer_buffers_per_layer=2, + ) diff --git a/tests/unittest/disaggregated/test_pool_matching.py b/tests/unittest/disaggregated/test_pool_matching.py index 884852c3a539..ba1aeea4d90c 100644 --- a/tests/unittest/disaggregated/test_pool_matching.py +++ b/tests/unittest/disaggregated/test_pool_matching.py @@ -1,3 +1,18 @@ +# 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. + """Golden tests for ``PeerRegistrar.get_pool_mapping``. Locks the current (pre-refactor) behavior of disagg pool matching across the @@ -6,7 +21,7 @@ * single KV pool, full layer overlap (basic MHA) * MLA (kv_factor=1, KEY-only pool) * KV + block-scale pools coexisting in one LG - * KV + FLAT pools coexisting in one LG (FLAT has empty buffer_entries) + * KV + REPLICATED indexer pools coexisting in one LG * PP partial layer overlap (Step-1 LG match by global_layer_id) * Two pools with same role but different layer sets within a peer LG (DSv4 virtual-layer scenario, exercises ``best_overlap``) @@ -33,6 +48,7 @@ PhysicalPoolGroup, PoolView, ) +from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_to_layer_group # --------------------------------------------------------------------------- # Builders @@ -70,13 +86,13 @@ def _pool_view(pool_idx, layer_role_pairs, *, pool_role=None, mapper_kind=Mapper ) -def _empty_pool_view(pool_idx, *, pool_role=frozenset({"indexer_k"})): - """Pool view with no buffer entries — FLAT pool convention.""" - return PoolView( - pool_idx=pool_idx, - buffer_entries=np.array([], dtype=BUFFER_ENTRY_DTYPE), +def _indexer_pool_view(pool_idx, local_layer_ids, *, pool_role=frozenset({"indexer_k"})): + """Replicated indexer pool view: one synthesized entry per local layer.""" + return _pool_view( + pool_idx, + [(lid, "indexer_k") for lid in local_layer_ids], pool_role=pool_role, - mapper_kind=MapperKind.FLAT, + mapper_kind=MapperKind.REPLICATED, ) @@ -229,16 +245,16 @@ def _bq_pool(pool_idx, lids): def test_kv_and_indexer_in_same_lg(): - """KV pool + FLAT pool. FLAT has empty buffer_entries; matches by role.""" + """KV pool + replicated indexer pool match independently by role.""" self_lg = _attn_lg( 0, [(0, 0), (1, 1)], - [_kv_pool_view(0, [0, 1]), _empty_pool_view(1)], + [_kv_pool_view(0, [0, 1]), _indexer_pool_view(1, [0, 1])], ) peer_lg = _attn_lg( 0, [(0, 0), (1, 1)], - [_kv_pool_view(0, [0, 1]), _empty_pool_view(1)], + [_kv_pool_view(0, [0, 1]), _indexer_pool_view(1, [0, 1])], ) self_pt = _page_table([self_lg], pool_specs={0: [(1024, 64, 0x1000), (512, 64, 0x2000)]}) @@ -251,6 +267,58 @@ def test_kv_and_indexer_in_same_lg(): assert mapping == {(0, 0): (0, 0), (0, 1): (0, 1)} +def test_minimax_split_kv_and_replicated_index_pools_match(): + """MiniMax M3 shape: NHD KV pool + REPLICATED index-key pool. + + coalescing_group derivation keeps INDEX_KEY in its own pool on every + topology, so both sides always present the same two role sets and + role-set matching pairs them 1:1 with kinds intact. + """ + + def _lg(): + return _attn_lg( + 0, + [(0, 0), (1, 1)], + [ + _pool_view( + 0, + [(lid, role) for lid in (0, 1) for role in ("key", "value")], + mapper_kind=MapperKind.NHD, + ), + _pool_view( + 1, [(0, "index_key"), (1, "index_key")], mapper_kind=MapperKind.REPLICATED + ), + ], + ) + + self_pt = _page_table([_lg()], pool_specs={0: [(512, 64, 0x1000), (256, 64, 0x2000)]}) + peer_pt = _page_table([_lg()], pool_specs={0: [(512, 64, 0x3000), (256, 64, 0x4000)]}) + + reg = _registrar(self_pt) + peer_ri = _rank_info(name="peer", rank=1, page_table=peer_pt) + + mapping = reg.get_pool_mapping(peer_ri) + assert mapping == {(0, 0): (0, 0), (0, 1): (0, 1)} + + +def test_mismatched_view_kinds_raise(): + """Kind disagreement on one role set must fail loudly. + + Same role set but different kinds on the two sides is a version or + configuration inconsistency. + """ + self_lg = _attn_lg( + 0, [(0, 0)], [_pool_view(0, [(0, "index_key")], mapper_kind=MapperKind.REPLICATED)] + ) + peer_lg = _attn_lg(0, [(0, 0)], [_pool_view(0, [(0, "index_key")], mapper_kind=MapperKind.NHD)]) + + reg = _registrar(_page_table([self_lg])) + peer_ri = _rank_info(name="peer", rank=1, page_table=_page_table([peer_lg])) + + with pytest.raises(ValueError, match="incompatible mapper"): + reg.get_pool_mapping(peer_ri) + + def test_pp_partial_layer_overlap(): """Self covers global layers {10,11}, peer covers {11,12}. Match via overlap on layer 11.""" self_lg = _attn_lg(0, [(0, 10), (1, 11)], [_kv_pool_view(0, [0, 1])]) @@ -265,6 +333,47 @@ def test_pp_partial_layer_overlap(): assert mapping == {(0, 0): (0, 0)} +def test_pool_view_spanning_multiple_peer_lgs_raises(): + """A self pool view spanning two peer LGs must fail loudly. + + Mismatched layer grouping between peers is an unsupported topology; + silently transferring only the first LG's overlap would drop layers. + """ + self_lg = _attn_lg(0, [(0, 10), (1, 11)], [_kv_pool_view(0, [0, 1])]) + # Peer holds the same global layers but split across two layer groups. + peer_lg0 = _attn_lg(0, [(0, 10)], [_kv_pool_view(0, [0])]) + peer_lg1 = _attn_lg(1, [(0, 11)], [_kv_pool_view(0, [0])]) + + reg = _registrar(_page_table([self_lg])) + peer_ri = _rank_info(name="peer", rank=1, page_table=_page_table([peer_lg0, peer_lg1])) + + with pytest.raises(ValueError, match="multiple peer layer groups"): + reg.get_pool_mapping(peer_ri) + + +def test_skewed_buffer_entries_per_layer_raises(): + """A skewed per-layer entry distribution must raise. + + 1 + 3 entries over two layers passes ``total % layers == 0`` (4 % 2) + but is not a uniform per-layer layout — must raise, not return 2. + """ + view = _pool_view(0, [(0, "key"), (1, "key"), (1, "key"), (1, "key")]) + with pytest.raises(ValueError, match="not evenly distributed"): + PeerRegistrar._get_buffers_per_layer(view, layer_group_id=0, pool_idx=0) + + +def test_duplicate_global_layer_id_across_lgs_raises(): + """Layer groups must partition a rank's attention layers. + + The same global_layer_id in two LGs would silently corrupt peer matching. + """ + lg0 = _attn_lg(0, [(0, 10)], [_kv_pool_view(0, [0])]) + lg1 = _attn_lg(1, [(0, 10)], [_kv_pool_view(0, [0])]) + + with pytest.raises(ValueError, match="layer groups must partition"): + get_layer_to_layer_group(_page_table([lg0, lg1])) + + def test_two_pools_distinct_roles_in_same_lg(): """Two pools with distinct pool_role in one LG match by role, not by layer overlap. @@ -274,11 +383,11 @@ def test_two_pools_distinct_roles_in_same_lg(): pool with the same pool_role. """ self_kv = _kv_pool_view(0, [0, 1]) - self_indexer = _empty_pool_view(1) + self_indexer = _indexer_pool_view(1, [0, 1]) self_lg = _attn_lg(0, [(0, 10), (1, 11)], [self_kv, self_indexer]) peer_kv = _kv_pool_view(0, [0, 1]) - peer_indexer = _empty_pool_view(1) + peer_indexer = _indexer_pool_view(1, [0, 1]) peer_lg = _attn_lg(0, [(0, 10), (1, 11)], [peer_kv, peer_indexer]) self_pt = _page_table([self_lg], pool_specs={0: [(1024, 64, 0x1000), (512, 64, 0x2000)]}) @@ -323,20 +432,19 @@ def test_same_role_pools_disambiguated_by_layer_overlap(): @pytest.mark.parametrize( ("self_mapper_kind", "peer_mapper_kind"), [ - (MapperKind.FLAT, MapperKind.INDEXED), - (MapperKind.INDEXED, MapperKind.FLAT), + (MapperKind.REPLICATED, MapperKind.INDEXED), + (MapperKind.INDEXED, MapperKind.REPLICATED), ], ) def test_mixed_mapper_kinds_are_rejected(self_mapper_kind, peer_mapper_kind): """Pool layouts must use the same mapper kind on both peers.""" def _indexer_view(mapper_kind): - if mapper_kind == MapperKind.FLAT: - return _empty_pool_view(0) return _pool_view( 0, [(0, "indexer_k"), (1, "indexer_k")], pool_role=frozenset({"indexer_k"}), + mapper_kind=mapper_kind, ) self_lg = _attn_lg(0, [(0, 10), (1, 11)], [_indexer_view(self_mapper_kind)]) diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index 5c934152efed..f188cfab6ef9 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -1,7 +1,27 @@ +# 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. + +from types import SimpleNamespace + import numpy as np +import pytest +from tensorrt_llm._torch.disaggregation.native import rank_info as rank_info_module from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBufferMeta from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm.bindings import DataType def test_rank_info_construction(): @@ -75,3 +95,41 @@ def test_rank_info_roundtrip_with_aux_meta(): np.testing.assert_array_equal(restored.aux_meta.size, [1024, 2048]) np.testing.assert_array_equal(restored.aux_meta.item_sizes, [64, 128]) assert restored.aux_meta.device == "cpu" + + +@pytest.mark.parametrize( + ("dtype", "expected_element_bytes", "expected_type"), + [(DataType.NVFP4, 0.5, float), (DataType.HALF, 2, int)], +) +def test_rank_info_represents_cache_element_bytes( + monkeypatch, dtype, expected_element_bytes, expected_type +): + monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _manager: None) + manager = SimpleNamespace( + mapping=SimpleNamespace( + rank=0, + tp_size=2, + tp_rank=0, + pp_size=1, + pp_rank=0, + dp_size=1, + cp_size=1, + cp_rank=0, + enable_attention_dp=False, + ), + pp_layers=[0], + num_kv_heads_per_layer=[4], + tokens_per_block=64, + head_dim=128, + dtype=dtype, + kv_factor=2, + ) + + rank_info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) + + assert rank_info.attention.element_bytes == expected_element_bytes + assert isinstance(rank_info.attention.element_bytes, expected_type) + + restored = RankInfo.from_bytes(rank_info.to_bytes()) + assert restored.attention.element_bytes == expected_element_bytes + assert isinstance(restored.attention.element_bytes, expected_type) diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index cfad765968e2..68b38bf896ae 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -20,6 +20,8 @@ from typing import Optional from unittest.mock import Mock +import pytest + from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.transfer import TaskStatus, TxSession from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 @@ -269,6 +271,48 @@ def fake_allgather(payload): assert new_completed == [7] # intersection only (8 is completed on the peer only) +@pytest.mark.skip( + reason="ctx idle fast-path was dropped from this branch. TODO: when the " + "fast-path is reintroduced, its terminal-count reduction must mirror " + "_ctx_consensus()'s communicator scope (TP group, then PP group; TP " + "skipped under attention DP) — a WORLD-scoped allreduce hangs under " + "ADP+PP because independent attention-DP lanes poll on their own " + "schedules. Re-enable this test and add scoped mock coverage for the " + "TP+PP and ADP+PP configurations plus real-collective MP tests." +) +def test_ctx_consensus_fastpath_skips_when_idle(monkeypatch) -> None: + # With the fast-path enabled, an all-zero terminal count (one fixed-size + # allreduce) makes every rank skip the variable-length consensus; a non-zero + # count falls through to the normal consensus path. + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver._CTX_CONSENSUS_FASTPATH", True + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_send_session = True + transceiver._ctx_need_tp_sync = True + transceiver._ctx_need_pp_sync = False + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._dist = Mock() + transceiver._dist.allreduce = Mock(return_value=0) + transceiver._ctx_consensus = Mock(return_value=[]) + transceiver._build_to_process = Mock(return_value=[]) + transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [], [])) + transceiver._transfer_worker = _FakeTransferWorker() + transceiver._close_failed_sessions = Mock() + + completed, failed = transceiver.check_context_transfer_status(at_least_request_num=0) + + assert completed == [] and failed == [] + transceiver._dist.allreduce.assert_called_once() + transceiver._ctx_consensus.assert_not_called() # idle fast-path skipped the consensus + + # Non-zero global terminal count => fast-path does not skip; consensus runs. + transceiver._dist.allreduce = Mock(return_value=2) + transceiver.check_context_transfer_status(at_least_request_num=0) + transceiver._ctx_consensus.assert_called_once() + + def test_tx_session_wait_complete_defaults_to_blocking() -> None: task = _FakeTask(TaskStatus.INIT, wait_result=False) session = _make_tx_session([task]) diff --git a/tests/unittest/scripts/test_build_wheel.py b/tests/unittest/scripts/test_build_wheel.py new file mode 100644 index 000000000000..0b9061efd34d --- /dev/null +++ b/tests/unittest/scripts/test_build_wheel.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +# 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. + +import importlib.util +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +SCRIPT_PATH = REPO_ROOT / "scripts" / "build_wheel.py" +MSA_INTERFACE = Path("python/fmha_sm100/cute/interface.py") + +_SPEC = importlib.util.spec_from_file_location("build_wheel", SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_BUILD_WHEEL = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_BUILD_WHEEL) +stage_msa_package = _BUILD_WHEEL.stage_msa_package + + +def test_stage_msa_package_applies_patch_without_modifying_submodule(tmp_path): + source_interface = REPO_ROOT / "3rdparty" / "MSA" / MSA_INTERFACE + if not source_interface.is_file(): + pytest.skip("3rdparty/MSA is not initialized") + + source_before = source_interface.read_bytes() + staged_package = stage_msa_package(REPO_ROOT, tmp_path) + staged_interface = staged_package / "cute" / "interface.py" + + assert b"def _prepare_paged_hnd_input" not in source_before + assert "def _prepare_paged_hnd_input" in staged_interface.read_text() + assert source_interface.read_bytes() == source_before + + +def test_stage_msa_package_requires_initialized_submodule(tmp_path): + project_dir = tmp_path / "project" + project_dir.mkdir() + + with pytest.raises(FileNotFoundError, match="initialize 3rdparty/MSA"): + stage_msa_package(project_dir, tmp_path / "build")