diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index 4d2111f72cdd..2bb292f1cbc4 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -5261,7 +5261,7 @@ For more information, please refer to - `Tracker`: https://github.com/tox-dev/py-filelock/issues -## flashinfer-python (0.6.16) +## flashinfer-python (0.6.17.dev20260806) ### Licenses License: `Apache-2.0` diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu index 43a1a6846cdc..1b91784037af 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu @@ -265,12 +265,14 @@ constexpr int DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA = 128; // and strides over the row space. This visits the per-expert tile padding that // the expanded-space kernel skips (~4% extra rows at 32 local experts); those // rows are dropped by the finalize kernel. The arithmetic below deliberately -// preserves the legacy kernel's 0/0 -> NaN behavior for an all-zero block. +// matches activationDeepSeekKernel bit for bit, including its finite all-zero +// block handling. constexpr int kDsActWarpSize = 32; constexpr int kDsActEltsPerSf = 128; constexpr int kDsActEltsPerThread = kDsActEltsPerSf / kDsActWarpSize; constexpr int kDsActWarpsPerCta = 4; constexpr int kDsActPermutedNumThreadsPerCta = kDsActWarpSize * kDsActWarpsPerCta; +constexpr float kDsActAmaxEpsilon = 1.0e-10F; constexpr bool shouldUsePermutedActivation(int innerDim, int numTokens, int topK, int numExperts, int tileTokensDim) { @@ -352,7 +354,11 @@ __global__ void activationDeepSeekPermutedKernel(KernelParams params) aMax = fmaxf(aMax, __shfl_xor_sync(0xffffffffu, aMax, offset)); } - float const scaleOut = aMax / kE4m3MaxVal; + // Keep an all-zero activation block finite. Without the floor, scaleOut + // is zero and quantizing the zero values evaluates 0 / 0, producing FP8 + // NaNs that poison FC2 and all following layers. This matches the + // epsilon used by the DeepGEMM FP8 activation quantizer. + float const scaleOut = fmaxf(aMax, kDsActAmaxEpsilon) / kE4m3MaxVal; if (lane == 0) { @@ -367,7 +373,7 @@ __global__ void activationDeepSeekPermutedKernel(KernelParams params) // Divide; do NOT hoist a reciprocal. `x / s` and `x * (1/s)` round // differently, and an equivalence run showed that single ulp flip a // greedy-decoded token. This must match activationDeepSeekKernel - // bit for bit, including 0/0 -> NaN on an all-zero scale block. + // bit for bit. outElts[i] = static_cast(out[i] / scaleOut); } *reinterpret_cast(params.outPtr + static_cast(permutedIdx) * outputDim + hiddenBase) @@ -504,10 +510,11 @@ __global__ void activationDeepSeekKernel(KernelParams params) { continue; } - s_scaleOutArr[tokenInCtaIdx] = aMaxArr[tokenInCtaIdx] / E4m3MaxVal; + float const scaleOut = fmaxf(aMaxArr[tokenInCtaIdx], kDsActAmaxEpsilon) / E4m3MaxVal; + s_scaleOutArr[tokenInCtaIdx] = scaleOut; int const scaleOut_idx = permutedIdxArr[tokenInCtaIdx] + totalNumPaddedTokens * (hiddenIdx / 128); - params.outDqSfsPtr[scaleOut_idx] = aMaxArr[tokenInCtaIdx] / E4m3MaxVal; + params.outDqSfsPtr[scaleOut_idx] = scaleOut; } } __syncthreads(); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu index 6a0ba120d4cc..9badec330872 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu @@ -1629,10 +1629,10 @@ void run(Data const& data, void* stream) bool const useStaticBlock = data.mNumTokens <= BlockKernelMaxNumTokens; int32_t const dispatchedMaxExperts = queryDispatchedMaxExperts(data); - // Cooperative block kernel: fastest path for tiny batches. Requires an elementwise - // preprocess (any but softmax-over-experts) and one CUDA block's worth of experts. - // Critical for large expert counts, where the classic one-warp-per-token TopK spills - // registers under the 1024-thread launch bounds (e.g. 896 experts / topK 16 at decode). + // Cooperative block kernel: fastest path for tiny batches in the large-expert tiers. + // It requires an elementwise preprocess (any but softmax-over-experts) and one CUDA + // block's worth of experts. The classic one-warp-per-token TopK is faster through the + // 512-expert tier, but spills registers in larger tiers (e.g. 896 experts / topK 16). bool const preprocessIsElementwise = data.mPreprocessType == RoutingPreprocessType::None || data.mPreprocessType == RoutingPreprocessType::Sigmoid || data.mPreprocessType == RoutingPreprocessType::SigmoidBias; @@ -1643,7 +1643,7 @@ void run(Data const& data, void* stream) return env != nullptr && env[0] == '1'; }(); bool const useCoopBlock = !disableCoopBlock && useStaticBlock && preprocessIsElementwise - && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts; + && dispatchedMaxExperts >= CoopBlockKernelMinNumExperts && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts; bool const useDynBlock = !useStaticBlock && data.mNumTokens <= DynBlockKernelMaxNumTokens && dispatchedMaxExperts <= DynBlockKernelMaxNumExperts; bool const useSingleBlock = useStaticBlock || useDynBlock; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh index ad0163cf9294..7853ab999a95 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh @@ -391,6 +391,9 @@ static constexpr int MaxNumTokensSingleClusterScores = NumBlocksPerCluster * Num static constexpr int BlockKernelMaxNumTokens = 4; static constexpr int DynBlockKernelMaxNumTokens = 16; static constexpr int DynBlockKernelMaxNumExperts = 256; +// The classic block kernel is faster through the 512-expert tier. The cooperative +// kernel avoids register spilling for the larger tiers. +static constexpr int CoopBlockKernelMinNumExperts = 576; // Cooperative block kernel: one thread per expert, so at most 1024 experts (1 CUDA block). static constexpr int CoopBlockKernelMaxNumExperts = 1024; diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index c3007343bf68..595eccbae0b3 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -212,6 +212,15 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit int32_t max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( max_num_padded_tokens, 2 * args.intermediate_size, btg::dtypeGetNumBits(args.mDtypeElt)); + // FC2 reads activation_output through TRTLLM-Gen's TMA-OOB path. Blackwell TMA requires at + // least 128 KiB of backing memory when the descriptor's addressable extent is at least + // 128 KiB, even when every logical access is in bounds. A gated activation is half as wide as + // gemm1_output, so reusing max_num_padded_tokens_gemm1 can leave its backing allocation below + // that contract for small decode batches. Compute the capacity from its own row width instead; + // this also grows the associated FP32 scale allocation from 2 KiB to the required 4 KiB minimum. + int32_t max_num_padded_tokens_activation + = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + max_num_padded_tokens, args.intermediate_size, btg::dtypeGetNumBits(args.mDtypeElt)); int32_t max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( max_num_padded_tokens, args.hidden_size, btg::dtypeGetNumBits(args.mDtypeOut)); @@ -254,10 +263,11 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); at::Tensor gemm1_output_scale = at::detail::empty_cuda({2 * intermediate_size / 128, max_num_padded_tokens_gemm1}, at::ScalarType::Float, routing_device, std::nullopt); - at::Tensor activation_output = at::detail::empty_cuda( - {max_num_padded_tokens_gemm1, intermediate_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); - at::Tensor activation_output_scale = at::detail::empty_cuda( - {intermediate_size / 128, max_num_padded_tokens_gemm1}, at::ScalarType::Float, routing_device, std::nullopt); + at::Tensor activation_output = at::detail::empty_cuda({max_num_padded_tokens_activation, intermediate_size}, + at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); + at::Tensor activation_output_scale + = at::detail::empty_cuda({intermediate_size / 128, max_num_padded_tokens_activation}, at::ScalarType::Float, + routing_device, std::nullopt); at::Tensor gemm2_output = at::detail::empty_cuda( {max_num_padded_tokens_gemm2, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); diff --git a/cpp/tensorrt_llm/thop/moeUtilOp.cpp b/cpp/tensorrt_llm/thop/moeUtilOp.cpp index e5496a89cdb0..340451fb73fe 100644 --- a/cpp/tensorrt_llm/thop/moeUtilOp.cpp +++ b/cpp/tensorrt_llm/thop/moeUtilOp.cpp @@ -138,9 +138,15 @@ std::tuple(experts_per_token * num_rows)}, diff --git a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu index 4a56dcd8f558..190449b7d71f 100644 --- a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu +++ b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu @@ -22,12 +22,11 @@ // * `activationDeepSeekPermutedKernel` - grids directly over the permuted row // space with one warp per (row, 128-element scale block), // via `shouldUsePermutedActivation()`. Both must produce *identical bits* for -// every row that carries a real token: DevKernel.cu documents that the permuted +// every row that carries a real token. DevKernel.cu documents that the permuted // kernel must not hoist a reciprocal out of `out / scaleOut`, because `x / s` // and `x * (1/s)` round differently and one ulp was enough to flip a // greedy-decoded token. An `isClose`-style comparison would not catch that -// regression, so everything below compares raw bit patterns (which also makes -// the NaN cases comparable). +// regression, so everything below compares raw bit patterns. // // Note on coverage: fp8 e4m3 carries three mantissa bits, so most 1-ulp fp32 // differences vanish when the result is rounded back down to fp8 -- only values @@ -45,6 +44,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" @@ -52,6 +52,7 @@ #include #include +#include #include #include #include @@ -212,8 +213,8 @@ protected: // Allocates device buffers, fills the inputs deterministically and uploads // them. `zeroedRowBlock`, when set, forces one (row, scale block) pair of - // the input to all-zero so both kernels take the aMax == 0 -> 0/0 -> NaN - // path on exactly the same element. + // the input to all-zero so both kernels exercise the finite aMax floor on + // exactly the same element. void setUp(ActivationEquivParam const& param, PermutedLayout const& layout, std::optional> zeroedRowBlock = std::nullopt) { @@ -401,12 +402,11 @@ INSTANTIATE_TEST_SUITE_P(BlockScaleMoeActivation, BlockScaleMoeActivationEquival //////////////////////////////////////////////////////////////////////////////////////////////////// -// An all-zero scale block yields aMax == 0, so the quantization does 0 / 0. The -// resulting NaN encoding is unspecified, but both kernels evaluate the same -// expression and must therefore land on the same bits -- which is exactly what -// would break if one of them replaced the division with a multiply by the -// reciprocal. -TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNaNs) +// An all-zero scale block must remain finite. A zero dequantization scale would +// make quantization evaluate 0 / 0 and emit FP8 NaNs, which then poison FC2 and +// all following layers. Both kernels floor aMax with the same epsilon and must +// emit identical zero bytes and a finite, positive scale. +TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesFiniteZeros) { ActivationEquivParam const param{"zero_block", /*numTokens=*/64, /*topK=*/4, /*numExperts=*/32, /*numLocalExperts=*/8, /*intermediateSize=*/256, /*paddingTile=*/8, /*hasSwigluLimit=*/false, @@ -422,17 +422,50 @@ TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNa auto const permuted = runOnce(kTileForcePermuted); auto const sfIdx = static_cast(zeroedRow) + static_cast(mTotalRows) * zeroedBlock; - EXPECT_EQ(floatBits(legacy.scales[sfIdx]), 0U) << "an all-zero block must give scaleOut == +0"; + EXPECT_TRUE(std::isfinite(legacy.scales[sfIdx])); + EXPECT_GT(legacy.scales[sfIdx], 0.F); EXPECT_EQ(floatBits(legacy.scales[sfIdx]), floatBits(permuted.scales[sfIdx])); for (int32_t elt = 0; elt < kEltsPerSf; ++elt) { auto const idx = static_cast(zeroedRow) * mOutputDim + zeroedBlock * kEltsPerSf + elt; + EXPECT_EQ(legacy.bytes[idx], toFp8Byte(0.F)) << "zero block emitted non-zero FP8 at element " << elt; ASSERT_EQ(static_cast(legacy.bytes[idx]), static_cast(permuted.bytes[idx])) - << "0/0 encoding differs at element " << elt; + << "zero-block encoding differs at element " << elt; } } //////////////////////////////////////////////////////////////////////////////////////////////////// +TEST(BlockScaleMoeActivationBackingTest, PadsActivationUsingItsOwnRowWidth) +{ + // A single-token Qwen-style decode can have only 32 padded rows. FC1 writes + // 2 * intermediateSize elements per row, while the gated activation read by + // FC2 is half as wide. Reusing FC1's capacity would therefore allocate only + // about half of the backing required by Blackwell's TMA-OOB contract. + constexpr int32_t maxNumPaddedTokens = 32; + constexpr int32_t intermediateSize = 2304; + constexpr int64_t minActivationBytes = 128 * 1024; + constexpr int64_t minScaleBytes = 4 * 1024; + auto const fp8Bits = tg::dtypeGetNumBits(tg::Dtype::E4m3); + + auto const gemm1Capacity = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + maxNumPaddedTokens, 2 * intermediateSize, fp8Bits); + auto const activationCapacity = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + maxNumPaddedTokens, intermediateSize, fp8Bits); + + auto const activationBytes = static_cast(activationCapacity) * intermediateSize * fp8Bits / 8; + auto const activationBytesWithGemm1Capacity = static_cast(gemm1Capacity) * intermediateSize * fp8Bits / 8; + auto const scaleBytes = static_cast(activationCapacity) * (intermediateSize / kEltsPerSf) * sizeof(float); + auto const scaleBytesWithGemm1Capacity + = static_cast(gemm1Capacity) * (intermediateSize / kEltsPerSf) * sizeof(float); + + EXPECT_GE(activationBytes, minActivationBytes); + EXPECT_GE(scaleBytes, minScaleBytes); + EXPECT_LT(activationBytesWithGemm1Capacity, minActivationBytes); + EXPECT_LT(scaleBytesWithGemm1Capacity, minScaleBytes); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace tensorrt_llm::tests::kernels::blockscalemoe diff --git a/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp index 4d1120420486..d77c3a198ded 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp @@ -295,6 +295,30 @@ TYPED_TEST(RoutingCustomKernelTest, BlockLevelParallelizationWithExpertParalleli this->runTest(param); }; +TYPED_TEST(RoutingCustomKernelTest, BlockLevelClassicBoundaryE512K16) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(4) + .withNumExperts(512) + .withTopK(16) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + +TYPED_TEST(RoutingCustomKernelTest, BlockLevelCooperativeBoundaryE576K8) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(4) + .withNumExperts(576) + .withTopK(8) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + TYPED_TEST(RoutingCustomKernelTest, BlockLevelParallelizationWithInvalidTopKInput) { auto param = RoutingKernelTestParam() diff --git a/requirements.txt b/requirements.txt index d63bc1a93be2..37c207b0e7c2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -58,7 +58,7 @@ ordered-set peft>=0.18.1,<0.19.0 patchelf einops -flashinfer-python==0.6.16 +flashinfer-python @ https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.17-20260806/flashinfer_python-0.6.17.dev20260806-py3-none-any.whl#sha256=4ed64b717bf979d268fd7249dfa354d100effe36bc9d340efac8dcebd354611d xgrammar==0.1.32 llguidance==0.7.29 jsonschema diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index b1bf50620b34..dbade31f5f35 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -55,7 +55,7 @@ dependencies = [ "peft (>=0.18.1,<0.19.0)", "patchelf (>=0.19.1.0,<0.20.0.0)", "einops (>=0.8.2,<0.9.0)", - "flashinfer-python (==0.6.16)", + "flashinfer-python @ https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.17-20260806/flashinfer_python-0.6.17.dev20260806-py3-none-any.whl#sha256=4ed64b717bf979d268fd7249dfa354d100effe36bc9d340efac8dcebd354611d", "xgrammar (==0.1.32)", "llguidance (==0.7.29)", "jsonschema (>=4.26.0,<5.0.0)", diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index c1f10ee89158..6917583d2885 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -177,6 +177,12 @@ class ModelConfig(Generic[TConfig]): max_seq_len: Optional[int] = None moe_max_num_tokens: Optional[int] = None + # Preserve whether the value came from the user across dataclasses.replace(). + # DeepGEMM uses this metadata to apply its conservative default only when + # max_num_tokens was not explicitly configured. + _moe_max_num_tokens_is_default: Optional[bool] = field(default=None, + repr=False, + compare=False) moe_load_balancer: Optional[MoeLoadBalancerConfig] = None attn_backend: str = 'TRTLLM' @@ -277,10 +283,16 @@ def get_all_reduce_strategy(strategy: str = "AUTO"): self.allreduce_strategy = get_all_reduce_strategy( self.allreduce_strategy) - # Set default moe_max_num_tokens if not specified - # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled + # Set default moe_max_num_tokens if not specified. The maximum number + # of tokens in MoE is multiplied by DP size when attention DP is + # enabled. + if self._moe_max_num_tokens_is_default is None: + self._moe_max_num_tokens_is_default = (self.moe_max_num_tokens + is None) if self.moe_max_num_tokens is None: self.moe_max_num_tokens = self.max_num_tokens * self.mapping.dp_size + if self.moe_max_num_tokens <= 0: + raise ValueError("moe_max_num_tokens must be a positive integer") @property def torch_dtype(self) -> torch.dtype: diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py index 9ee2a8a55e17..253d7e4b6bb0 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py @@ -3,6 +3,7 @@ import threading from abc import ABC, abstractmethod +from bisect import bisect_left, bisect_right from typing import Any, Dict, Iterator, Tuple, Union from tensorrt_llm.mapping import Mapping @@ -25,12 +26,15 @@ class ConsumableWeightsDict: def __init__(self, weights: Dict[str, Any]): self._weights = weights self._lock = threading.Lock() + self._key_index: list[str] | None = None def __getitem__(self, key: str) -> Any: return self._weights[key] def __setitem__(self, key: str, value: Any) -> None: with self._lock: + if key not in self._weights: + self._key_index = None self._weights[key] = value def __delitem__(self, key: str) -> None: @@ -68,6 +72,8 @@ def get(self, key: str, default: Any = None) -> Any: def update(self, other: Dict[str, Any]) -> None: with self._lock: + if any(key not in self._weights for key in other): + self._key_index = None self._weights.update(other) def clear(self) -> None: @@ -79,6 +85,33 @@ def clear(self) -> None: """ with self._lock: self._weights.clear() + self._key_index = [] + + def filter_prefix(self, prefix: str) -> Dict[str, Any]: + """Return weights below ``prefix`` without scanning the full mapping.""" + with self._lock: + if not prefix: + return dict(self._weights) + prefix_with_separator = prefix + "." + return { + key[len(prefix_with_separator):]: self._weights[key] + for key in self._matching_keys_locked(prefix_with_separator) + } + + def _matching_keys_locked(self, prefix_with_separator: str) -> list[str]: + if self._key_index is None: + self._key_index = sorted(self._weights) + begin = bisect_left(self._key_index, prefix_with_separator) + end = bisect_right(self._key_index, + prefix_with_separator + chr(0x10FFFF)) + return [ + key for key in self._key_index[begin:end] if key in self._weights + ] + + def _delete_keys_locked(self, keys: list[str]) -> int: + for key in keys: + del self._weights[key] + return len(keys) def mark_consumed_keys(self, keys) -> int: """Delete an exact set of keys to free memory. @@ -86,13 +119,11 @@ def mark_consumed_keys(self, keys) -> int: Use instead of :meth:`mark_consumed` when a module consumed specific tensors rather than a whole ``name.*`` subtree. """ - deleted = 0 with self._lock: - for key in keys: - if key in self._weights: - del self._weights[key] - deleted += 1 - return deleted + keys_to_delete = [ + key for key in dict.fromkeys(keys) if key in self._weights + ] + return self._delete_keys_locked(keys_to_delete) def mark_consumed(self, prefix: str) -> int: """ @@ -107,12 +138,8 @@ def mark_consumed(self, prefix: str) -> int: Thread-safe: uses a lock to prevent concurrent modification issues. """ with self._lock: - keys_to_delete = [ - k for k in self._weights.keys() if k.startswith(prefix + ".") - ] - for key in keys_to_delete: - del self._weights[key] - return len(keys_to_delete) + keys_to_delete = self._matching_keys_locked(prefix + ".") + return self._delete_keys_locked(keys_to_delete) class BaseWeightLoader(ABC): diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index a77c7f74609d..d620429c07e9 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -1,9 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from abc import ABC, abstractmethod from typing import Callable, List, Union from torch import nn from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + ConsumableWeightsDict from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM @@ -166,6 +171,8 @@ def should_skip_module(self, module_name: str) -> bool: for skip_module in self._skip_modules) def filter_weights(self, prefix: str, weights: dict) -> dict: + if isinstance(weights, ConsumableWeightsDict): + return weights.filter_prefix(prefix) result = {} for k, v in weights.items(): if k.startswith(prefix): diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index b71fe8495928..622f91b0c6f9 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -36,6 +36,7 @@ ) from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types from ..utils import is_nvfp4_marlin_supported_sm +from .checkpoints.base_weight_loader import ConsumableWeightsDict from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from .modeling_qwen3_next import Qwen3NextForCausalLM @@ -80,6 +81,20 @@ def _get_qwen35_moe_model_defaults(llm_args: "TorchLlmArgs") -> dict: return defaults +def _filter_language_model_weights(weights: Dict[str, torch.Tensor]): + """Drop vision weights without disabling incremental weight consumption.""" + filtered_weights = { + key: value for key, value in weights.items() if not key.startswith("model.visual.") + } + if isinstance(weights, ConsumableWeightsDict): + # The filtered mapping aliases the source tensors. Transfer ownership + # to a fresh consumable wrapper so the inner loader can release routed + # experts layer by layer instead of pinning the whole checkpoint. + weights.clear() + return ConsumableWeightsDict(filtered_weights) + return filtered_weights + + def _translate_mtp_pattern(name, n_hidden_layers): """Translate an HF ``mtp.*`` exclude pattern to a TRT-LLM module path. @@ -758,7 +773,7 @@ def load_weights( ) if weight_mapper.model is not self.llm: weight_mapper.init_model_and_config(self.llm, self.llm.model_config) - filtered_weights = {k: v for k, v in weights.items() if not k.startswith("model.visual.")} + filtered_weights = _filter_language_model_weights(weights) params_map = { r"^model\.language_model\.(.*)$": r"model.\1", } diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index aa2626bf1da7..05a6a11873dc 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -881,44 +881,58 @@ def forward( spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: - del all_rank_num_tokens - - def norm_embeds(): - return self.pre_fc_norm_embedding(embed_tokens(input_ids)) - - def norm_hidden(): - return self.pre_fc_norm_hidden(hidden_states) - - inputs_embeds, hidden_states = maybe_execute_in_parallel( - norm_embeds, - norm_hidden, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - self.aux_stream, - disable_on_compile=True, - ) - hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) + # Draft iterations can have a different per-rank token distribution + # from the target forward stored in attn_metadata. In particular, MTP + # Eagle processes one token per sequence after step 0. MoE chunking + # consumes this metadata, so keep it aligned with the draft input and + # restore the target metadata before returning to the caller. + previous_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + + try: + + def norm_embeds(): + return self.pre_fc_norm_embedding(embed_tokens(input_ids)) + + def norm_hidden(): + return self.pre_fc_norm_hidden(hidden_states) + + inputs_embeds, hidden_states = maybe_execute_in_parallel( + norm_embeds, + norm_hidden, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + disable_on_compile=True, + ) + hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) - tp_size = self.model_config.mapping.tp_size - tp_rank = self.model_config.mapping.tp_rank - if tp_size > 1 and not self.model_config.mapping.enable_attention_dp: - hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank] + tp_size = self.model_config.mapping.tp_size + tp_rank = self.model_config.mapping.tp_rank + if tp_size > 1 and not self.model_config.mapping.enable_attention_dp: + hidden_states = torch.chunk(hidden_states, tp_size, + dim=-1)[tp_rank] - hidden_states = self.fc(hidden_states) + hidden_states = self.fc(hidden_states) - hidden_states, residual = super().forward( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=None, - spec_metadata=spec_metadata, - **kwargs, - ) - hidden_states, _ = self.shared_head.norm(hidden_states, residual) - if spec_metadata is not None: - spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + hidden_states, residual = super().forward( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=None, + spec_metadata=spec_metadata, + **kwargs, + ) + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states( + 0, hidden_states, None) - return hidden_states + return hidden_states + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous_all_rank_num_tokens ALL_DECODER_LAYER_TYPES = { diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index bb81db4cf27a..849d5f55b7ff 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -1155,13 +1155,21 @@ def rename_weights_with_regex(pattern_mapping: Dict[str, str], weights: Dict): if key not in matched_keys: renamed_weights[key] = weights[key] - # Preserve ConsumableWeightsDict type if that's what was passed in + # The renamed mapping aliases every tensor in the input. Transfer + # ownership so mark_consumed() can actually release the last source-side + # reference while modules are loaded. if is_consumable: + weights.clear() return ConsumableWeightsDict(renamed_weights) return renamed_weights def filter_weights(prefix, weights: Dict): + from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + ConsumableWeightsDict + + if isinstance(weights, ConsumableWeightsDict): + return weights.filter_prefix(prefix) result = {} for k, v in weights.items(): if k.startswith(prefix): @@ -1170,6 +1178,42 @@ def filter_weights(prefix, weights: Dict): return result +def _pageout_safetensors_after_moe_load() -> None: + """Bound safetensors file-cache growth after loading each MoE layer. + + Do not gate this on CUDA's ``is_integrated`` property: GB300 reports false + even though four ranks share a host-memory cgroup and can exhaust it with + resident file pages. ``MADV_DONTNEED`` only discards pages already faulted + in, so untouched future weights do not incur extra I/O. + """ + from tensorrt_llm._torch.mmap_utils import pageout_file_backed_regions + + logger.info_once( + "Releasing resident safetensors pages after each MoE layer load", + key="weight_load_safetensors_pageout", + ) + torch.cuda.synchronize() + pageout_file_backed_regions(".safetensors", mode="dontneed") + + +def _get_load_weights_num_workers() -> Optional[int]: + """Return the configured bound for concurrent module weight loading.""" + env_name = "TRT_LLM_LOAD_WEIGHTS_NUM_WORKERS" + value = os.environ.get(env_name) + if value is None or not value.strip(): + return None + + try: + num_workers = int(value) + except ValueError as error: + raise ValueError( + f"{env_name} must be a positive integer, got {value!r}") from error + if num_workers <= 0: + raise ValueError( + f"{env_name} must be a positive integer, got {value!r}") + return num_workers + + def run_concurrently(func, args_list, reduce_func=None, @@ -1267,7 +1311,8 @@ def load_single_module(name, module): # and weights loading is done in the backend, so module name includes '.backend'. # We need to use parent module name (without .backend) to match saved weight names. # After MoE refactoring is fully complete, all paths will follow this branch. - if names[-1] == "backend" and isinstance(module, MoE): + is_moe_backend = names[-1] == "backend" and isinstance(module, MoE) + if is_moe_backend: name = '.'.join(names[:-1]) names = name.split('.') @@ -1337,6 +1382,9 @@ def load_single_module(name, module): weights.mark_consumed_keys( f'{name}.{n}' for n in loaded_own_params) + if is_moe_backend: + _pageout_safetensors_after_moe_load() + if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "False") in ["True", "true", "1", "yes", "y"]: for name, module in tqdm(list( @@ -1368,7 +1416,15 @@ def load_single_module(name, module): for name, module in model.named_modules(remove_duplicate=False) if name not in serial_load_modules ] - run_concurrently(load_single_module, args_list, pbar=pbar) + num_workers = _get_load_weights_num_workers() + if num_workers is not None: + logger.info( + f"Limiting concurrent module weight loading to {num_workers} workers" + ) + run_concurrently(load_single_module, + args_list, + pbar=pbar, + num_workers=num_workers) def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM], @@ -1400,7 +1456,8 @@ def load_single_module(name, module): # and weights loading is done in the backend, so module name includes '.backend'. # We need to use parent module name (without .backend) to match saved weight names. # After MoE refactoring is fully complete, all paths will follow this branch. - if names[-1] == "backend" and isinstance(module, MoE): + is_moe_backend = names[-1] == "backend" and isinstance(module, MoE) + if is_moe_backend: name = '.'.join(names[:-1]) names = name.split('.') @@ -1462,6 +1519,9 @@ def load_single_module(name, module): weights.mark_consumed_keys( f'{name}.{n}' for n in loaded_own_params) + if is_moe_backend: + _pageout_safetensors_after_moe_load() + if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "False") in ["True", "true", "1", "yes", "y"]: for name, module in tqdm(list( @@ -1493,4 +1553,12 @@ def load_single_module(name, module): for name, module in model.named_modules(remove_duplicate=False) if name not in serial_load_modules ] - run_concurrently(load_single_module, args_list, pbar=pbar) + num_workers = _get_load_weights_num_workers() + if num_workers is not None: + logger.info( + f"Limiting concurrent module weight loading to {num_workers} workers" + ) + run_concurrently(load_single_module, + args_list, + pbar=pbar, + num_workers=num_workers) diff --git a/tensorrt_llm/_torch/modules/fla/cached_replay.py b/tensorrt_llm/_torch/modules/fla/cached_replay.py index e279e1d72348..30b4effda5ee 100644 --- a/tensorrt_llm/_torch/modules/fla/cached_replay.py +++ b/tensorrt_llm/_torch/modules/fla/cached_replay.py @@ -16,6 +16,8 @@ _BV16_MAX_HEAD_TILES = 64 _BV32_MAX_HEAD_TILES = 128 _RATIO2_FINE_MAPPING_MAX_HEAD_TILES = 256 +_RATIO8_FINE_MAPPING_MAX_HEAD_TILES = 512 +_RATIO8_ROLLOVER_FINE_HEAD_TILES = 32 _EIGHT_WARP_COMMIT_HEAD_TILES = 1024 _PIPELINED_COMMIT_HEAD_TILES = 2048 _TWO_STAGE_REPLAY_HEAD_TILES = 4096 @@ -51,6 +53,17 @@ def _supports_fine_grained_replay_tiling( ) +def _supports_ratio8_replay_tiling( + num_key_heads: int, + num_value_heads: int, + head_tiles: int, +) -> bool: + """Return whether the measured B300 8:1 mapping applies.""" + return ( + num_value_heads == 8 * num_key_heads and head_tiles <= _RATIO8_FINE_MAPPING_MAX_HEAD_TILES + ) + + @triton.jit def _gdc_wait_with_memory_clobber(): tl.inline_asm_elementwise( @@ -612,6 +625,13 @@ def commit_gdn_cached_replay_history_layers( and V == 128 and ssm_states.dtype == torch.bfloat16 ) + use_ratio8_bf16_mapping = ( + history_size <= 16 + and HV == 8 * H + and K == 128 + and V == 128 + and ssm_states.dtype == torch.bfloat16 + ) use_small_grid_mapping = ( use_tuned_bf16_mapping and per_layer_head_tiles <= _SMALL_GRID_HEAD_TILES ) @@ -637,6 +657,8 @@ def commit_gdn_cached_replay_history_layers( 5 if use_tuned_bf16_mapping and per_layer_head_tiles >= _PIPELINED_COMMIT_HEAD_TILES else 5 + if use_ratio8_bf16_mapping and N >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + else 5 if not use_tuned_bf16_mapping and use_large_workload_mapping else 1 ) @@ -770,15 +792,29 @@ def fused_recurrent_gated_delta_rule_cached_replay_update( ) head_tiles = N * HV use_tuned_bf16_mapping = use_production_bf16_shape and HV == 4 * H + use_ratio8_mapping = use_production_bf16_shape and _supports_ratio8_replay_tiling( + H, HV, head_tiles + ) use_fine_grained_mapping = use_production_bf16_shape and _supports_fine_grained_replay_tiling( H, HV, head_tiles ) use_small_grid_mapping = use_fine_grained_mapping and head_tiles <= _SMALL_GRID_HEAD_TILES if block_v is None: - block_v = _default_cached_replay_block_v(use_fine_grained_mapping, head_tiles, V) + # The 8:1 B300 shape needs finer tiling at very small N to keep its + # periodic checkpoint-commit path occupied. The wider tile wins from + # N=16 and remains tied on ordinary steps below that boundary. + block_v = ( + (16 if head_tiles <= _RATIO8_ROLLOVER_FINE_HEAD_TILES else 64) + if use_ratio8_mapping + else _default_cached_replay_block_v(use_fine_grained_mapping, head_tiles, V) + ) if num_warps is None: num_warps = ( - 2 if use_fine_grained_mapping and (use_small_grid_mapping or launch_with_pdl) else 4 + (2 if head_tiles <= _RATIO8_ROLLOVER_FINE_HEAD_TILES else 4) + if use_ratio8_mapping + else ( + 2 if use_fine_grained_mapping and (use_small_grid_mapping or launch_with_pdl) else 4 + ) ) BV = block_v use_large_workload_mapping = ( diff --git a/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py b/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py index e819f1690728..e7aa3ca5aeb2 100644 --- a/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py +++ b/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py @@ -28,6 +28,7 @@ at process start; do not import it lazily inside hot paths. """ +import inspect from typing import Optional, Tuple import torch @@ -40,6 +41,30 @@ from tensorrt_llm._utils import is_sm_100f +def _enable_cutlass_dsl_45_pipeline_compatibility(pipeline) -> None: + """Accept FlashInfer's explicit DSL 4.6 default on CUTLASS DSL 4.5. + + FlashInfer 0.6.17.dev20260806 permits CUTLASS DSL 4.5 but passes the new + ``enable_multicast_signaling=False`` spelling. DSL 4.5 already implements + that non-multicast behavior; its Python signature simply predates the + keyword. Only the equivalent ``False`` case is adapted here. + """ + create = pipeline.PipelineTmaAsync.create + if "enable_multicast_signaling" in inspect.signature(create).parameters: + return + if getattr(create, "_trtllm_cutlass_dsl_45_compatibility", False): + return + + def create_compatible(*args, **kwargs): + enable_multicast_signaling = kwargs.pop("enable_multicast_signaling", False) + if enable_multicast_signaling: + raise ValueError("enable_multicast_signaling=True requires nvidia-cutlass-dsl>=4.6") + return create(*args, **kwargs) + + create_compatible._trtllm_cutlass_dsl_45_compatibility = True + pipeline.PipelineTmaAsync.create = staticmethod(create_compatible) + + # Mirror the @torch.compiler.disable on the legacy Triton wrapper # (chunk.py:119): Dynamo must not trace this wrapper because it imports # `flashinfer` lazily and calls into FI's CuTe-DSL kernels, neither of @@ -65,6 +90,9 @@ def chunk_gated_delta_rule( # FlashInfer is imported lazily so importing this module on a non-FlashInfer # build does not error until the function is actually called. import flashinfer + from cutlass import pipeline + + _enable_cutlass_dsl_45_pipeline_compatibility(pipeline) # --- Step 1: pre-flight asserts -------------------------------------- assert head_first is False, "head_first=True is not supported by this wrapper" diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index 4ab178c75239..6c728b54e321 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -149,7 +149,7 @@ def create_strategy( force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD", communication_method) if force_method is not None: - return CommunicationFactory._create_forced_method( + strategy = CommunicationFactory._create_forced_method( force_method, model_config, num_experts, @@ -161,6 +161,11 @@ def create_strategy( use_flashinfer, hidden_size=hidden_size, ) + logger.info( + f"Selected communication strategy: {strategy.__class__.__name__} " + f"(forced by TRTLLM_FORCE_COMM_METHOD={force_method})" + ) + return strategy # Auto-selection: Try strategies in priority order using try-catch # Priority: NVLinkOneSided > NVLinkTwoSided > NcclEP > DeepEP > DeepEPLowLatency > AllGather diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index 7d3951762326..53beac5f625b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -22,6 +22,7 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm import deep_gemm from tensorrt_llm._utils import get_sm_version, nvtx_range +from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo from ...memory_buffer_utils import get_memory_buffers @@ -34,6 +35,32 @@ MoEWeightLoadingMode, UnquantizedFusedMoEMethod) from .routing import BaseMoeRoutingMethod +_DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS = 18688 + + +def _configure_deepgemm_moe_max_num_tokens(model_config: ModelConfig) -> None: + moe_max_num_tokens = model_config.moe_max_num_tokens + if moe_max_num_tokens is None: + raise ValueError("moe_max_num_tokens must be set before creating MoE") + + if not model_config._moe_max_num_tokens_is_default: + if moe_max_num_tokens > _DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS: + logger.warning_once( + "DeepGEMM moe_max_num_tokens is explicitly set to " + f"{moe_max_num_tokens}, above the conservative default " + f"{_DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS}; this may increase " + "GPU memory usage.", + key="deepgemm_explicit_moe_max_num_tokens") + return + + if moe_max_num_tokens <= _DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS: + return + + was_frozen = model_config._frozen + model_config._frozen = False + model_config.moe_max_num_tokens = (_DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS) + model_config._frozen = was_frozen + @triton.jit def _masked_index_copy_group_quant_fp8( @@ -617,7 +644,7 @@ def preprocess_after_permute(expert_first_token_offset_tensor, Only the number of permuted (expanded) tokens is needed here, not the permuted activations themselves. Callers that run moe_permute_op with - skip_data_expand=True leave permuted_data_tensor uninitialized, so the count + skip_data_expand=True return an empty permuted_data_tensor, so the count must come from a populated tensor (e.g. permuted_row_to_unpermuted_row_tensor.shape[0]). """ total_tokens = num_permuted_tokens @@ -828,11 +855,10 @@ def __init__( # max_num_tokens = ((mtp+1)*max_batch_size+max_isl+128+63)//64*64 = 9344 # moe_max_num_tokens = max_num_tokens * 2 = 18688 # It can avoid OOM for 8k/1k cases. - default_moe_max_num_tokens = 18688 - if model_config.moe_max_num_tokens > default_moe_max_num_tokens: - model_config._frozen = False - model_config.moe_max_num_tokens = default_moe_max_num_tokens - model_config._frozen = True + # Preserve an explicit deployment value. Only clamp the derived + # default, which keeps the existing OOM-safe behavior when the user + # does not size the DeepGEMM workspace deliberately. + _configure_deepgemm_moe_max_num_tokens(model_config) super().__init__( routing_method=routing_method, @@ -982,23 +1008,26 @@ def run_moe( assert token_selected_experts is not None assert token_final_scales is not None - # Permutation. - # skip_data_expand=True computes the permutation maps but skips the - # data-copy step (expandInputRowsKernel), so permuted_data_tensor and - # permuted_token_final_scales_tensor are returned with UNINITIALIZED - # contents (still full-size, just never written). The fused expand+quant - # kernel re-derives the activations from x via + # Permutation. With skip_data_expand=True, the operator needs the input + # row count and dtype but never reads its hidden dimension. Use a + # zero-width view so older operator libraries, which still size an + # unused expanded return from input.shape[1], allocate no activation + # storage. The fused expand+quant kernel reads the original x below. + # Newer operator libraries return the unused activation and scale + # tensors empty as well. + # + # The fused expand+quant kernel re-derives the activations from x via # permuted_row_to_unpermuted_row_tensor instead, so all unused outputs are # discarded with `_`. ( permuted_row_to_unpermuted_row_tensor, _, # permuted_token_selected_experts_tensor (unused) - _, # permuted_data_tensor (uninitialized under skip_data_expand) + _, # permuted_data_tensor (unused and allocation-free) expert_first_token_offset_tensor, - _, # permuted_token_final_scales_tensor (uninitialized under skip_data_expand) + _, # permuted_token_final_scales_tensor (unused) unpermuted_row_to_permuted_row_tensor, ) = torch.ops.trtllm.moe_permute_op( - x, + x[:, :0], token_selected_experts, token_final_scales, None, # w3_w1_weight.view(weight_dtype), diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 85e2689e9905..a83fc754b4b7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -819,9 +819,12 @@ def register_to_fix_weight_fn(self, weight_name: str): assert isinstance( weight_tensor, torch.Tensor), f'weight {weight_name} should be a tensor' - assert weight_tensor.is_contiguous(), ( - f'weight {weight_name} should be contiguous, ' - f'shape={weight_tensor.shape}, strides={weight_tensor.stride()}') + # DeepGemm FP8 block scales are column-major TMA views. The backing + # storage is complete and starts at offset zero, so migrating the + # allocation preserves the view's shape and strides. + assert weight_tensor.storage_offset() == 0, ( + f'weight {weight_name} should start at storage offset zero, ' + f'offset={weight_tensor.storage_offset()}') assert weight_tensor.numel() * weight_tensor.element_size( ) == weight_tensor.untyped_storage().size(), ( f'weight {weight_name} shape={weight_tensor.shape} ' @@ -861,9 +864,11 @@ def register_all_parameter_slot_and_to_fix_weight_fns( self.layer_load_balancer.host_tensor_sharer.share_host_tensor_with_shape( expert_id, weight_name, weight_tensor[local_slot_id]) else: + transfer_tensor = self.layer_load_balancer.host_tensor_sharer.get_transfer_view( + weight_tensor[0]) self.layer_load_balancer.host_tensor_sharer.pre_register_host_tensor_with_shape( - expert_id, weight_name, weight_tensor.dtype, - weight_tensor[0].shape) + expert_id, weight_name, transfer_tensor.dtype, + transfer_tensor.shape) def _register_layer(self, model_config: ModelConfig): self.register_to_config = False diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py index 115891c860c0..a55601542b16 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py @@ -1,4 +1,21 @@ +# 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 gc +import json +import mmap import os import threading from contextlib import nullcontext @@ -20,6 +37,74 @@ from ..multi_stream_utils import do_multi_stream +class _FileBackedSharedMemory: + """Shared, mmap-backed storage with the ``SharedMemory`` interface.""" + + def __init__(self, + name: str, + directory: str, + create: bool = False, + size: int = 0): + if os.path.basename(name) != name: + raise ValueError(f"Invalid EPLB shared-memory name: {name}") + os.makedirs(directory, exist_ok=True) + self._path = os.path.join(directory, name) + self._unlinked = False + flags = os.O_RDWR + if create: + flags |= os.O_CREAT | os.O_EXCL + fd = os.open(self._path, flags, 0o600) + try: + if create: + if size <= 0: + raise ValueError( + f"EPLB shared-memory size must be positive, got {size}") + os.posix_fallocate(fd, 0, size) + else: + size = os.fstat(fd).st_size + self._mmap = mmap.mmap(fd, size, access=mmap.ACCESS_WRITE) + finally: + os.close(fd) + + @property + def buf(self): + return self._mmap + + def flush(self): + """Write dirty expert pages back before loading the next layer.""" + self._mmap.flush() + + def close(self): + self._mmap.close() + + def unlink(self): + if self._unlinked: + return + try: + os.unlink(self._path) + except FileNotFoundError: + pass + self._unlinked = True + + +def _get_moe_weight_transfer_view(t: torch.Tensor) -> torch.Tensor: + """Return a row-major view compatible with ``MoeWeight`` 2D copies. + + Most expert tensors are already row-major. DeepGemm FP8 block scales use + a TMA column-major layout, represented by a 2D tensor whose first stride + is one. Transposing that tensor exposes the same bytes as a row-major + pitched matrix without changing the layout consumed by the kernel. + """ + if t.dim() != 2: + return t + if t.stride(1) == 1 and t.stride(0) >= t.size(1): + return t + if t.stride(0) == 1 and t.stride(1) >= t.size(0): + return t.transpose(0, 1) + raise ValueError(f"Unsupported MoE weight layout: shape={tuple(t.shape)}, " + f"strides={t.stride()}") + + def _tensor_to_weight(t: torch.Tensor) -> _tbr.MoeWeight: """ Convert a tensor to a MoeWeight object. @@ -28,6 +113,7 @@ def _tensor_to_weight(t: torch.Tensor) -> _tbr.MoeWeight: t: The tensor to convert """ assert t.dim() <= 2, "t.dim() should be less than or equal to 2" + t = _get_moe_weight_transfer_view(t) shape = [1, 1] pitch = 1 elt_size = torch.tensor([], dtype=t.dtype).element_size() @@ -84,6 +170,13 @@ def __init__(self, layer_id: int, expert_count: int, self.loaded_shared_weights = [] + self.file_backed_directory = os.getenv("TRTLLM_EPLB_SHM_DIR") + if self.file_backed_directory is not None: + logger.info_once( + "Using file-backed storage for online EPLB host weights at " + f"{self.file_backed_directory}", + key="eplb_file_backed_storage") + def set_shared_memory_base_name(self, shared_memory_base_name): """ Set the shared memory base name for the layer. @@ -93,6 +186,11 @@ def set_shared_memory_base_name(self, shared_memory_base_name): """ self.shared_memory_base_name = shared_memory_base_name + @staticmethod + def get_transfer_view(t: torch.Tensor) -> torch.Tensor: + """Return the row-major view used by EPLB weight migration.""" + return _get_moe_weight_transfer_view(t) + def get_shared_memory_name(self, rank: Optional[int] = None): """ Get the shared memory name for the layer. @@ -108,6 +206,17 @@ def get_shared_memory_name(self, rank: Optional[int] = None): shared_memory_name = f"{self.shared_memory_base_name}_l{self.layer_id}_lr{rank}_all" return shared_memory_name + def _open_shared_memory(self, + name: str, + create: bool = False, + size: int = 0): + if self.file_backed_directory is not None: + return _FileBackedSharedMemory(name, + self.file_backed_directory, + create=create, + size=size) + return shared_memory.SharedMemory(name=name, create=create, size=size) + def pre_register_host_tensor_with_shape(self, expert_id: int, name: str, dtype, tensor_shape): """ @@ -148,7 +257,7 @@ def share_host_tensor_with_shape(self, expert_id: int, name: str, """ assert len( t.shape) <= 2, "tensor_shape dim must be less than or equal to 2" - assert t.is_contiguous() == True, "t.is_contiguous() must be True" + t = _get_moe_weight_transfer_view(t).contiguous() assert (expert_id, name) not in self.shared_tensors.keys() assert self.expert_start <= expert_id < self.expert_end self.shared_tensors[(expert_id, name)] = t @@ -192,19 +301,19 @@ def finalize_layer_weights(self): shm_name = self.get_shared_memory_name() try: - shm = shared_memory.SharedMemory(name=shm_name, - create=True, - size=total_size) + shm = self._open_shared_memory(shm_name, + create=True, + size=total_size) except FileExistsError: tensorrt_llm.logger.warning( f'Found exist EPLB shared memory name: {shm_name}, unlinking...' ) - existing_shm = shared_memory.SharedMemory(name=shm_name) + existing_shm = self._open_shared_memory(shm_name) existing_shm.close() existing_shm.unlink() - shm = shared_memory.SharedMemory(name=shm_name, - create=True, - size=total_size) + shm = self._open_shared_memory(shm_name, + create=True, + size=total_size) self.own_shm = shm offset = 0 @@ -227,6 +336,12 @@ def finalize_layer_weights(self): ), f"key={key} already exists" self.host_weights[key] = st offset += aligned_size + # Online EPLB can stage hundreds of GiB per node. Explicitly make + # completed file-backed layers writeback-eligible before materializing + # the next layer, otherwise dirty mmap pages can exhaust the Slurm + # memory cgroup during checkpoint loading. + if isinstance(shm, _FileBackedSharedMemory): + shm.flush() self.shared_tensors = {} for raw_weight in self.loaded_shared_weights: @@ -245,7 +360,7 @@ def finalize_host_tensor_sharing(self, add_host_weight_fn: Callable = None): continue shm_name = self.get_shared_memory_name(rank) - shm = shared_memory.SharedMemory(name=shm_name) + shm = self._open_shared_memory(shm_name) self.imported_shms.append(shm) rank_expert_start = rank * self.expert_count // self.local_size @@ -281,7 +396,8 @@ def pre_shutdown_cleanup(self): """ for shm in self.imported_shms: shm.close() - resource_tracker.unregister(shm._name, "shared_memory") + if isinstance(shm, shared_memory.SharedMemory): + resource_tracker.unregister(shm._name, "shared_memory") self.imported_shms = None if self.own_shm: self.own_shm.close() @@ -295,6 +411,12 @@ def post_shutdown_cleanup(self): self.own_shm = None +class _MoeLoadBalancerRuntimeState: + + def __init__(self) -> None: + self.active = False + + class SingleLayerMoeLoadBalancer: """ A class representing a single layer of the Mixture of Experts (MoE) load balancer. @@ -308,7 +430,8 @@ def __init__( expert_count: int, updates_enabled: bool = True, repeated_count=1, - aux_stream: Optional[torch.cuda.Stream] = None): + aux_stream: Optional[torch.cuda.Stream] = None, + runtime_state: Optional[_MoeLoadBalancerRuntimeState] = None): """ Initialize a SingleLayerMoeLoadBalancer instance. @@ -318,6 +441,7 @@ def __init__( expert_count: total number of experts updates_enabled: whether to enable weight updates repeated_count: the repeated count of current layer, used when forward is repeated more than once like MTP. + runtime_state: shared state indicating whether this forward runs online EPLB synchronization """ self.single_layer_load_balancer_impl = single_layer_load_balancer_impl self.single_layer_load_balancer_ptr = single_layer_load_balancer_impl.get_pointer( @@ -325,6 +449,7 @@ def __init__( self.expert_count = expert_count self.updates_enabled = updates_enabled self.repeated_count = repeated_count + self.runtime_state = runtime_state layer_id = self.single_layer_load_balancer_impl.get_layer_id() self.host_tensor_sharer = HostMoeTensorSharer( layer_id, expert_count, @@ -389,6 +514,10 @@ def is_static_routing(self): def is_dynamic_routing(self): return self.updates_enabled + def is_runtime_active(self): + return self.updates_enabled and (self.runtime_state is None + or self.runtime_state.active) + def need_load_shared_weights(self): return self.is_dynamic_routing() @@ -519,6 +648,8 @@ def start_wait_gpu_stage(self): """ Start to wait for the GPU stage to complete. """ + if not self.is_runtime_active(): + return assert self.func_called_count["start_wait_gpu_stage"] == 0 self.func_called_count["start_wait_gpu_stage"] += 1 if self.updates_enabled: @@ -537,6 +668,8 @@ def done_wait_gpu_stage(self): """ Done waiting for the GPU stage to complete. """ + if not self.is_runtime_active(): + return assert self.func_called_count["start_wait_gpu_stage"] == 1 assert self.func_called_count["done_wait_gpu_stage"] == 0 self.func_called_count["done_wait_gpu_stage"] += 1 @@ -548,6 +681,8 @@ def start_set_cpu_stage(self): """ Start to set the CPU stage. """ + if not self.is_runtime_active(): + return assert self.func_called_count["done_wait_gpu_stage"] == 1 assert self.func_called_count["start_set_cpu_stage"] == 0 self.func_called_count["start_set_cpu_stage"] += 1 @@ -567,6 +702,8 @@ def done_set_cpu_stage(self): """ Done setting the CPU stage. """ + if not self.is_runtime_active(): + return assert self.func_called_count["start_set_cpu_stage"] == 1 for name in self.func_called_count: self.func_called_count[name] = 0 @@ -585,6 +722,8 @@ def update_local_statistic(self, local_raw_expert_ids: torch.Tensor, is_first_stage: Whether this is the first stage is_last_stage: Whether this is the last stage """ + if not self.is_runtime_active(): + return assert self.func_called_count["done_wait_gpu_stage"] == 1 assert self.func_called_count["update_statistic_with_global_ids"] == 0 self.func_called_count["update_local_statistic"] += 1 @@ -616,6 +755,8 @@ def get_local_statistic_tensor(self) -> Optional[torch.Tensor]: Returns: The local statistic tensor if using statistic else None """ + if not self.is_runtime_active(): + return None assert self.func_called_count["update_local_statistic"] > 0 self.func_called_count["get_local_statistic_tensor"] += 1 if self.updates_enabled: @@ -634,6 +775,8 @@ def update_statistic_with_gathered_statistic( Args: gathered_local_statistic_tensor: gathered local statistics info, should have shape (world_size, self.expert_count) """ + if not self.is_runtime_active(): + return assert self.func_called_count["get_local_statistic_tensor"] > 0 assert self.func_called_count["update_statistic_with_local_ids"] == 0 assert self.func_called_count["update_statistic_with_global_ids"] == 0 @@ -670,6 +813,8 @@ def update_statistic_with_local_ids(self, is_last_stage: Whether this is the last stage allreduce: The allreduce object """ + if not self.is_runtime_active(): + return assert self.func_called_count["done_wait_gpu_stage"] == 1 assert self.func_called_count[ "update_statistic_with_gathered_statistic"] == 0 @@ -704,6 +849,8 @@ def update_statistic_with_global_ids(self, is_first_stage: Whether this is the first stage is_last_stage: Whether this is the last stage """ + if not self.is_runtime_active(): + return assert self.func_called_count["done_wait_gpu_stage"] == 1 assert self.func_called_count[ "update_statistic_with_gathered_statistic"] == 0 @@ -737,7 +884,7 @@ def route(self, Returns: A tensor of routed slot IDs """ - if self.is_dynamic_routing(): + if self.is_runtime_active(): assert self.func_called_count["done_wait_gpu_stage"] == 1 self.func_called_count["route"] += 1 return torch.ops.trtllm.moe_load_balance_routing( @@ -773,7 +920,8 @@ def __init__(self, ep_rank: int, ep_size: int, layer_updates_per_iter: int, - shared_memory_base_name: Optional[str] = None): + shared_memory_base_name: Optional[str] = None, + iteration_interval: Optional[int] = None): """ Initialize a MoeLoadBalancer instance. @@ -782,11 +930,51 @@ def __init__(self, ep_size: The total number of processes in expert parallelism layer_updates_per_iter: The number of layers to update per iteration shared_memory_base_name: Shared memory base name, will use 'moe_shared' if None + iteration_interval: Number of model forwards between online EPLB statistic/update iterations """ self.is_shutdown = True + if iteration_interval is None: + iteration_interval = int( + os.getenv('TRTLLM_EPLB_ITERATION_INTERVAL', '1')) + if iteration_interval < 1: + raise ValueError( + f"iteration_interval must be positive, got {iteration_interval}" + ) + self.statistics_per_cycle = int( + os.getenv('TRTLLM_EPLB_STATISTICS_PER_CYCLE', '0')) + self.cycle_start_iter = int( + os.getenv('TRTLLM_EPLB_CYCLE_START_ITER', '0')) + self.statistic_interval = int( + os.getenv('TRTLLM_EPLB_STATISTIC_INTERVAL', '1')) + self.update_interval = int(os.getenv('TRTLLM_EPLB_UPDATE_INTERVAL', + '2')) + self.cycle_interval = int(os.getenv('TRTLLM_EPLB_CYCLE_INTERVAL', '0')) + if self.statistics_per_cycle < 0: + raise ValueError( + "TRTLLM_EPLB_STATISTICS_PER_CYCLE must be non-negative, " + f"got {self.statistics_per_cycle}") + if self.cycle_start_iter < 0: + raise ValueError( + "TRTLLM_EPLB_CYCLE_START_ITER must be non-negative, " + f"got {self.cycle_start_iter}") + if self.statistics_per_cycle > 0: + if self.statistic_interval < 1: + raise ValueError( + "TRTLLM_EPLB_STATISTIC_INTERVAL must be positive, " + f"got {self.statistic_interval}") + if self.update_interval < 2: + raise ValueError( + "TRTLLM_EPLB_UPDATE_INTERVAL must be at least two so " + "each migration has a drain forward, " + f"got {self.update_interval}") + if self.cycle_interval < 0: + raise ValueError( + "TRTLLM_EPLB_CYCLE_INTERVAL must be non-negative, " + f"got {self.cycle_interval}") self.ep_rank = ep_rank self.ep_size = ep_size self.layer_updates_per_iter = layer_updates_per_iter + self.iteration_interval = iteration_interval self.load_balancer_impl = _tbr.MoeLoadBalancer(ep_rank, ep_size, layer_updates_per_iter) self._previous_balancer = None @@ -797,13 +985,25 @@ def __init__(self, self.is_shutdown = False self.iter_id = 0 + self.forward_iter_id = 0 self.in_iter = False + self.runtime_state = _MoeLoadBalancerRuntimeState() self.enable_statistic = False self.enable_update_weights = False self.next_layer_repeated_count = None + self.telemetry_path = os.getenv('TRTLLM_EPLB_TELEMETRY_PATH') + self._telemetry_previous_mode = None + self._telemetry_previous_placement = None + self._telemetry_counts = { + 'sample': 0, + 'update': 0, + 'drain': 0, + 'skip': 0, + } + def __del__(self): if not self.is_shutdown: self.shutdown() @@ -870,7 +1070,8 @@ def add_layer( expert_count, updates_enabled=updates_enabled, repeated_count=repeat_count, - aux_stream=aux_stream) + aux_stream=aux_stream, + runtime_state=self.runtime_state) single_layer_load_balancer.set_shared_memory_base_name( self.shared_memory_base_name) self.single_layer_load_balancers.append(single_layer_load_balancer) @@ -896,6 +1097,93 @@ def finalize_model(self): single_layer_load_balancer.py_finalize_model() self.load_balancer_impl.finalize_model() torch.cuda.empty_cache() + self._telemetry_previous_placement = self._placement_snapshot() + self._write_telemetry('initialized', + placement=self._telemetry_previous_placement) + + def _placement_snapshot(self): + """Return all layer placements from the EPLB CPU control plane.""" + if self.ep_rank != 0 or self.telemetry_path is None: + return None + return [{ + 'layer_id': layer.get_layer_idx(), + 'rank_expert_ids': layer.get_old_rank_expert_ids(), + } for layer in self.single_layer_load_balancers] + + @staticmethod + def _count_changed_slots(previous, current): + if previous is None or current is None: + return 0 + changed = 0 + for previous_layer, current_layer in zip(previous, current): + for previous_rank, current_rank in zip( + previous_layer['rank_expert_ids'], + current_layer['rank_expert_ids']): + changed += sum(previous_expert != current_expert + for previous_expert, current_expert in zip( + previous_rank, current_rank)) + return changed + + def _expert_bytes_by_layer(self): + result = [] + for layer in self.single_layer_load_balancers: + sharer = layer.host_tensor_sharer + expert_bytes = 0 + if sharer is not None: + for dtype, shape in sharer.name_info.values(): + element_size = torch.empty((), dtype=dtype).element_size() + expert_bytes += int(np.prod(shape)) * element_size + result.append(expert_bytes) + return result + + def _write_telemetry(self, event: str, **fields) -> None: + if self.ep_rank != 0 or self.telemetry_path is None: + return + directory = os.path.dirname(self.telemetry_path) + if directory: + os.makedirs(directory, exist_ok=True) + record = { + 'event': event, + 'forward_iter_id': self.forward_iter_id, + 'iter_id': self.iter_id, + **fields, + } + with open(self.telemetry_path, 'a', encoding='utf-8') as telemetry_file: + telemetry_file.write( + json.dumps(record, separators=(',', ':')) + '\n') + + def record_iter_mode(self, mode: str) -> None: + """Record batched Online-EPLB lifecycle and post-drain placement.""" + if self.ep_rank != 0 or self.telemetry_path is None: + return + self._telemetry_counts[mode] += 1 + if mode == 'skip': + if self._telemetry_previous_mode != 'skip': + self._write_telemetry('stable', counts=self._telemetry_counts) + elif mode == 'drain': + placement = self._placement_snapshot() + previous = self._telemetry_previous_placement + changed_slots = self._count_changed_slots(previous, placement) + changed_bytes = 0 + if previous is not None and placement is not None: + for previous_layer, current_layer, expert_bytes in zip( + previous, placement, self._expert_bytes_by_layer()): + for previous_rank, current_rank in zip( + previous_layer['rank_expert_ids'], + current_layer['rank_expert_ids']): + changed_bytes += expert_bytes * sum( + previous_expert != current_expert + for previous_expert, current_expert in zip( + previous_rank, current_rank)) + self._telemetry_previous_placement = placement + self._write_telemetry('drain_complete', + counts=self._telemetry_counts, + changed_slots=changed_slots, + changed_bytes=changed_bytes, + placement=placement) + else: + self._write_telemetry(mode, counts=self._telemetry_counts) + self._telemetry_previous_mode = mode def set_warm_up_iter_count(self, iter_count: int): """ @@ -913,6 +1201,73 @@ def set_iter_info(self, enable_statistic: Optional[bool], if enable_update_weights is not None: self.enable_update_weights = enable_update_weights + def uses_batched_update_cycle(self) -> bool: + """Whether statistics and migrations use separate serving forwards.""" + return self.statistics_per_cycle > 0 + + def _get_batched_update_cycle_mode(self) -> str: + """Return the mode for the configured statistics/migration cycle.""" + if self.forward_iter_id < self.cycle_start_iter: + return 'skip' + + layer_count = len(self.single_layer_load_balancers) + if layer_count == 0: + return 'skip' + update_count = (layer_count + self.layer_updates_per_iter - + 1) // self.layer_updates_per_iter + # The C++ update plan appends an empty group when every group has the + # same size. Count it here so a cycle consumes the complete plan and + # the next cycle starts again at the first real layer group. + if layer_count % update_count == 0: + update_count += 1 + last_statistic_offset = ((self.statistics_per_cycle - 1) * + self.statistic_interval) + first_update_offset = last_statistic_offset + 1 + last_update_offset = (first_update_offset + + (update_count - 1) * self.update_interval) + active_span = last_update_offset + 2 + period = active_span + self.cycle_interval + offset = (self.forward_iter_id - self.cycle_start_iter) % period + + if offset <= last_statistic_offset: + if offset % self.statistic_interval == 0: + return 'sample' if self.enable_statistic else 'skip' + return 'skip' + if offset > last_update_offset + 1: + return 'skip' + migration_offset = offset - first_update_offset + if migration_offset % self.update_interval == 0: + return 'update' if self.enable_update_weights else 'skip' + if migration_offset % self.update_interval == 1: + return 'drain' if self.enable_update_weights else 'skip' + return 'skip' + + def get_next_iter_mode(self) -> str: + """Return the online EPLB action for the next model forward.""" + if self.uses_batched_update_cycle(): + if not self.enable_statistic and not self.enable_update_weights: + return 'skip' + return self._get_batched_update_cycle_mode() + if self.iteration_interval == 1: + return 'sample' + if not self.enable_statistic and not self.enable_update_weights: + return 'skip' + phase = self.forward_iter_id % self.iteration_interval + if phase == 0: + return 'sample' + if phase == 1: + return 'drain' + return 'skip' + + def requires_eager_forward(self) -> bool: + """Whether the next forward must run eager for EPLB synchronization.""" + return self.is_dynamic_routing() and self.get_next_iter_mode() != 'skip' + + def advance_forward_iter(self) -> None: + if (self.uses_batched_update_cycle() or self.iteration_interval + > 1) and (self.enable_statistic or self.enable_update_weights): + self.forward_iter_id += 1 + def reconfigure_mask_only(self, dead_ranks: list[int]) -> None: """ Reconfigure EPLB routing so slots on dead EP ranks are unreachable. @@ -926,14 +1281,19 @@ def reconfigure_mask_only(self, dead_ranks: list[int]) -> None: "Cannot reconfigure EPLB mask while an iteration is active") self.load_balancer_impl.reconfigure_mask_only(list(dead_ranks)) - def start_iter(self): + def start_iter(self, + enable_statistic: Optional[bool] = None, + enable_update_weights: Optional[bool] = None): """ Start a new iteration. """ assert self.in_iter == False, "already in forward" self.in_iter = True - self.load_balancer_impl.start_iter(self.iter_id, self.enable_statistic, - self.enable_update_weights) + self.runtime_state.active = True + statistic = self.enable_statistic if enable_statistic is None else enable_statistic + update_weights = self.enable_update_weights if enable_update_weights is None else enable_update_weights + self.load_balancer_impl.start_iter(self.iter_id, statistic, + update_weights) def end_iter(self): """ @@ -945,6 +1305,7 @@ def end_iter(self): assert self.in_iter, "not in forward, cannot end_iter" self.load_balancer_impl.end_iter(self.iter_id) self.in_iter = False + self.runtime_state.active = False self.iter_id += 1 def shutdown(self): @@ -968,7 +1329,15 @@ def __repr__(self): Returns: A string representation of the load balancer """ - return f"MoeLoadBalancer(ep_rank={self.ep_rank}, ep_size={self.ep_size}, layer_updates_per_iter={self.layer_updates_per_iter})" + return ( + f"MoeLoadBalancer(ep_rank={self.ep_rank}, ep_size={self.ep_size}, " + f"layer_updates_per_iter={self.layer_updates_per_iter}, " + f"iteration_interval={self.iteration_interval}, " + f"statistics_per_cycle={self.statistics_per_cycle}, " + f"cycle_start_iter={self.cycle_start_iter}, " + f"statistic_interval={self.statistic_interval}, " + f"update_interval={self.update_interval}, " + f"cycle_interval={self.cycle_interval})") def __enter__(self): """ @@ -1057,8 +1426,8 @@ def maybe_create_moe_load_balancer( layer_updates_per_iter=model_config.moe_load_balancer. layer_updates_per_iter) logger.info( - f"Created MoE LoadBalancer, layer_updates_per_iter={model_config.moe_load_balancer.layer_updates_per_iter}..." - ) + f"Created MoE LoadBalancer, layer_updates_per_iter={model_config.moe_load_balancer.layer_updates_per_iter}, " + f"iteration_interval={moe_load_balancer.iteration_interval}...") return moe_load_balancer @@ -1071,6 +1440,7 @@ def __init__(self, self.moe_load_balancer = moe_load_balancer self.enable_statistic = enable_statistic self.enable_updates = enable_updates + self.iter_mode = 'skip' def __enter__(self): """ @@ -1083,7 +1453,19 @@ def __enter__(self): ): self.moe_load_balancer.set_iter_info(self.enable_statistic, self.enable_updates) - self.moe_load_balancer.start_iter() + self.iter_mode = self.moe_load_balancer.get_next_iter_mode() + if self.iter_mode == 'sample': + if self.moe_load_balancer.uses_batched_update_cycle(): + self.moe_load_balancer.start_iter( + enable_statistic=True, enable_update_weights=False) + else: + self.moe_load_balancer.start_iter() + elif self.iter_mode == 'update': + self.moe_load_balancer.start_iter(enable_statistic=False, + enable_update_weights=True) + elif self.iter_mode == 'drain': + self.moe_load_balancer.start_iter(enable_statistic=False, + enable_update_weights=False) return self def __exit__(self, exc_type, exc_val, exc_tb): @@ -1100,7 +1482,10 @@ def __exit__(self, exc_type, exc_val, exc_tb): """ if self.moe_load_balancer is not None and not self.moe_load_balancer.is_static_routing( ): - self.moe_load_balancer.end_iter() + if self.iter_mode != 'skip': + self.moe_load_balancer.end_iter() + self.moe_load_balancer.record_iter_mode(self.iter_mode) + self.moe_load_balancer.advance_forward_iter() return False diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index f86ecf887aee..19343fa25ae0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -638,7 +638,10 @@ def _forward_multiple_chunks( ) # ========== Empty-chunk substitution (DP only) ========== - chunked_used = torch.ones(num_chunks, dtype=torch.bool) + # This is host-only bookkeeping. Keep it as Python state so checking a + # chunk below does not dispatch Tensor.__bool__ while a CUDA Graph is + # being captured. + chunked_used = [True] * num_chunks if moe.use_dp: # The split heuristic guarantees chunk 0 has >= 1 token, so it can # stand in for any empty chunk on this rank. Without substitution, diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py index 9b86c02673f2..0dc8a8ed019e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py @@ -169,20 +169,21 @@ def compute_moe( intermediate_size = module.intermediate_size hidden_size = x.shape[1] - # Permute the data for expert-parallel processing. - # Unlike DeepGemmFusedMoE (which fuses gather+finalize and never touches - # permuted_data_tensor), this op reuses permuted_data_tensor as a - # write-before-read scratch buffer in the gather+finalize tail below, so - # it is kept; only the genuinely unused outputs are discarded with `_`. + # Permute the data for expert-parallel processing. The skipped + # permutation uses only the row count and dtype, so a zero-width view + # prevents older operator libraries from allocating an unused expanded + # activation return. This path still needs a write-before-read scratch + # buffer in the gather+finalize tail, so it allocates that buffer at the + # point of use below. ( permuted_row_to_unpermuted_row_tensor, _, # permuted_token_selected_experts_tensor (unused) - permuted_data_tensor, + _, # permuted_data_tensor (empty under skip_data_expand) expert_first_token_offset_tensor, - _, # permuted_token_final_scales_tensor (uninitialized under skip_data_expand) + _, # permuted_token_final_scales_tensor (empty under skip_data_expand) unpermuted_row_to_permuted_row_tensor, ) = torch.ops.trtllm.moe_permute_op( - x, + x[:, :0], token_selected_slots, token_final_scales, None, # w3_w1_weight @@ -290,6 +291,9 @@ def compute_moe( ) # Gather results back to original token order + permuted_data_tensor = torch.empty((num_permuted_tokens, hidden_size), + dtype=x.dtype, + device=x.device) triton_masked_index_gather(permuted_data_tensor, h3, expert_first_token_offset_tensor, token_to_expert_map) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 8a0f63d397d6..4f203c82d2d9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -278,7 +278,8 @@ def _online_eplb_not_supported(cls, module): @classmethod def _online_eplb_not_verified(cls, module): - if cls.need_load_shared_weights(module): + if (cls.eplb_support_status == EplbSupportStatus.NOT_VERIFIED + and cls.need_load_shared_weights(module)): logger.warning(f'{cls.__name__} online EPLB is not verified yet') def create_weights( @@ -1314,6 +1315,7 @@ def resmooth_and_transform_fp8_scale( class DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm( DeepSeekFP8BlockScalesFusedMoEMethod): + eplb_support_status = EplbSupportStatus.SUPPORTED def _needs_e8m0_resmooth(self): return is_sm_100f() or get_sm_version() == 120 diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 7378b332fab3..1b62be18d20f 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import enum @@ -33,6 +36,26 @@ from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_nvfp4_marlin_supported_sm, replace_parameter_and_save_metadata, unswizzle_sf) +from .low_m_gemm import LOW_M_GEMM_ACTIVE, apply_low_m_gemm + +_LOW_M_GEMM_MAX_M = 32 +_LOW_M_GEMM_SHAPE_COLLECTION_ACTIVE = bool( + os.environ.get("TRTLLM_LOW_M_GEMM_SHAPE_LOG")) + + +def _should_apply_low_m_gemm(input: torch.Tensor) -> bool: + """Skip dispatcher overhead outside FlashInfer's supported M domain.""" + + if not LOW_M_GEMM_ACTIVE: + return False + # The debug collector intentionally inventories every BF16 Linear shape, + # including normal-path M values above the FlashInfer kernel domain. + if _LOW_M_GEMM_SHAPE_COLLECTION_ACTIVE: + return True + if input.ndim < 1: + return False + k = int(input.shape[-1]) + return k > 0 and input.numel() <= _LOW_M_GEMM_MAX_M * k class WeightMode(str, enum.Enum): @@ -543,6 +566,12 @@ def create_weights(self, module: Linear, in_features: int, def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + # The opt-in low-M dispatcher uses FlashInfer's packaged direct/split-K + # heuristic and returns None for the normal GEMM fallback. + if _should_apply_low_m_gemm(input): + output = apply_low_m_gemm(module, input, module.weight, bias) + if output is not None: + return output # CuTe DSL BF16 GEMM path for Blackwell if (module.use_cute_dsl_bf16_gemm and is_sm_100f() and module.weight.dtype == torch.bfloat16): diff --git a/tensorrt_llm/_torch/modules/low_m_gemm.py b/tensorrt_llm/_torch/modules/low_m_gemm.py new file mode 100644 index 000000000000..a4af2dfc986e --- /dev/null +++ b/tensorrt_llm/_torch/modules/low_m_gemm.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import atexit +import enum +import functools +import json +import os +import threading +from collections import defaultdict +from pathlib import Path +from typing import Any, Optional + +import torch +from packaging.version import InvalidVersion, Version + +from tensorrt_llm.logger import logger + +from ..flashinfer_utils import get_env_enable_pdl + +_BACKEND_ENV = "TRTLLM_LOW_M_GEMM_BACKEND" +_SHAPE_LOG_ENV = "TRTLLM_LOW_M_GEMM_SHAPE_LOG" +_SHAPE_LOG_SCHEMA_VERSION = 1 +_FLASHINFER_BACKEND = "cute-dsl" +_MIN_FLASHINFER_VERSION = Version("0.6.17.dev20260806") +_SUPPORTED_SMS = {100, 103} +_MAX_FLASHINFER_M = 32 +_CUBLAS_8192X2048_M = frozenset({5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 17, 18, 19, 21, 24}) + + +class LowMGemmBackend(str, enum.Enum): + """Runtime choices exposed by the low-M BF16 dispatcher.""" + + OFF = "off" + AUTO = "auto" + FLASHINFER = "flashinfer" + CUBLAS = "cublas" + + +def _get_rank() -> int: + for name in ("RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"): + value = os.environ.get(name) + if value is not None: + return int(value) + return 0 + + +def _get_world_size() -> int: + for name in ("WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "SLURM_NTASKS"): + value = os.environ.get(name) + if value is not None: + return int(value) + return 1 + + +def _rank_path(path: str) -> Path: + rank = _get_rank() + if "{rank}" in path: + return Path(path.format(rank=rank)) + resolved = Path(path) + if _get_world_size() > 1: + return resolved.with_name(f"{resolved.name}.rank{rank}") + return resolved + + +def _normalize_backend(value: str) -> LowMGemmBackend: + normalized = value.strip().lower().replace("_", "-") + aliases = { + "": LowMGemmBackend.OFF, + "0": LowMGemmBackend.OFF, + "none": LowMGemmBackend.OFF, + "off": LowMGemmBackend.OFF, + "auto": LowMGemmBackend.AUTO, + "flashinfer": LowMGemmBackend.FLASHINFER, + "flashinfer-cute-dsl": LowMGemmBackend.FLASHINFER, + "cute-dsl": LowMGemmBackend.FLASHINFER, + "cublas": LowMGemmBackend.CUBLAS, + "cublaslt": LowMGemmBackend.CUBLAS, + } + if normalized not in aliases: + choices = ", ".join(backend.value for backend in LowMGemmBackend) + raise ValueError(f"Invalid {_BACKEND_ENV}={value!r}; expected one of {choices}.") + return aliases[normalized] + + +def _configured_backend() -> LowMGemmBackend: + return _normalize_backend(os.environ.get(_BACKEND_ENV, "off")) + + +@functools.cache +def _device_sm(device_index: int) -> int: + major, minor = torch.cuda.get_device_capability(device_index) + return major * 10 + minor + + +def _current_sm(device: torch.device) -> int: + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + return _device_sm(device_index) + + +def _prefer_cublas_for_auto(m: int, n: int, k: int, sm: int) -> bool: + """Return whether cuBLAS wins for a measured low-M Blackwell shape. + + FlashInfer selects its direct versus split-K tactic internally. This + narrow crossover handles shapes where the packaged CuTe DSL kernels are + slower than the normal cuBLAS Linear path in the 2026-08-06 nightly. + Explicit ``flashinfer`` selection remains an unconditional override. + """ + + if sm != 103: + return False + if (n, k) == (8192, 128): + return m >= 8 + if (n, k) == (8192, 2048): + return m in _CUBLAS_8192X2048_M + if (n, k) == (15520, 8192): + return m >= 15 + return (m, n, k) == (16, 2304, 8192) + + +class _ShapeCollector: + """Debug-only hot-shape collector with durable runtime-shape discovery.""" + + def __init__(self, output_path: str): + self.output_path = _rank_path(output_path) + self._counts: dict[tuple[Any, ...], int] = defaultdict(int) + self._lock = threading.Lock() + self._flush_lock = threading.Lock() + self._persisted_problem_keys: set[tuple[Any, ...]] = set() + self._pending_problem_keys: set[tuple[Any, ...]] = set() + self._warmup_complete = False + atexit.register(self.flush) + + def record( + self, + module: torch.nn.Module, + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + cuda_graph: bool, + ) -> None: + input_2d = input_tensor.reshape(-1, input_tensor.shape[-1]) + layer = getattr(module, "_low_m_gemm_name", module.__class__.__name__) + op_name = _infer_op_name(layer) + batch_shape = tuple(int(dim) for dim in input_tensor.shape[:-1]) + key = ( + int(input_2d.shape[0]), + int(weight.shape[0]), + int(weight.shape[1]), + str(input_tensor.dtype).removeprefix("torch."), + False, + True, + cuda_graph, + bias is not None, + batch_shape, + layer, + op_name, + ) + problem_key = key[:8] + with self._lock: + self._counts[key] += 1 + flush_new_runtime_shape = ( + self._warmup_complete + and problem_key not in self._persisted_problem_keys + and problem_key not in self._pending_problem_keys + ) + if flush_new_runtime_shape: + self._pending_problem_keys.add(problem_key) + if flush_new_runtime_shape: + # SRT terminates serving workers after a benchmark, so their atexit + # handlers are not reliable. Persist each shape first discovered + # after warmup immediately. This collector is debug-only and is + # never enabled for a performance measurement. + try: + self.flush() + finally: + with self._lock: + self._pending_problem_keys.discard(problem_key) + + def flush(self, *, mark_warmup_complete: bool = False) -> None: + with self._flush_lock: + with self._lock: + if mark_warmup_complete: + self._warmup_complete = True + counts = dict(self._counts) + if not counts: + return + rows = [] + for key, count in counts.items(): + ( + m, + n, + k, + dtype, + trans_a, + trans_b, + cuda_graph, + has_bias, + batch_shape, + layer, + op_name, + ) = key + rows.append( + { + "m": m, + "n": n, + "k": k, + "dtype": dtype, + "trans_a": trans_a, + "trans_b": trans_b, + "cuda_graph": cuda_graph, + "has_bias": has_bias, + "batch_shape": list(batch_shape), + "layer": layer, + "op_name": op_name, + "call_count": count, + } + ) + rows.sort( + key=lambda row: ( + -row["call_count"], + row["m"], + row["n"], + row["k"], + row["layer"], + ) + ) + payload = { + "schema_version": _SHAPE_LOG_SCHEMA_VERSION, + "rank": _get_rank(), + "world_size": _get_world_size(), + "shapes": rows, + } + self.output_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self.output_path.with_suffix(f"{self.output_path.suffix}.tmp") + temporary_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + temporary_path.replace(self.output_path) + with self._lock: + self._persisted_problem_keys.update(key[:8] for key in counts) + + +def _infer_op_name(layer: str) -> str: + lowered = layer.lower() + patterns = ( + ("mtp", "mtp"), + ("lm_head", "lm_head"), + ("lmhead", "lm_head"), + ("qkv", "attention_qkv"), + ("q_proj", "attention_qkv"), + ("k_proj", "attention_qkv"), + ("v_proj", "attention_qkv"), + ("o_proj", "attention_o_proj"), + ("out_proj", "out_proj"), + ("in_proj", "in_proj"), + ("shared_expert", "shared_expert"), + ) + for needle, op_name in patterns: + if needle in lowered: + return op_name + # Appended MTP fusion projections may be named simply ``fc``. + if lowered.endswith(".fc"): + return "mtp_fusion_fc" + return "linear" + + +class LowMGemmDispatcher: + """Dispatch eligible BF16 GEMMs to FlashInfer's low-M CuTe DSL backend.""" + + def __init__(self) -> None: + self.backend = _configured_backend() + self.cuda_graph_enabled = True + self._prepared = False + self._flashinfer_mm = None + self._engaged: set[tuple[int, int, int, bool]] = set() + self._collector = ( + _ShapeCollector(os.environ[_SHAPE_LOG_ENV]) if os.environ.get(_SHAPE_LOG_ENV) else None + ) + + @property + def enabled(self) -> bool: + return self.backend not in (LowMGemmBackend.OFF, LowMGemmBackend.CUBLAS) + + def prepare(self, model: torch.nn.Module, cuda_graph_enabled: bool) -> None: + self.cuda_graph_enabled = cuda_graph_enabled + for name, module in model.named_modules(): + if module.__class__.__name__.endswith("Linear"): + setattr(module, "_low_m_gemm_name", name) + + if self._prepared or not self.enabled: + self._prepared = True + return + + self._prepare_flashinfer() + self._prepared = True + + def _prepare_flashinfer(self) -> None: + try: + import flashinfer as flashinfer_module + from flashinfer import mm_bf16 + from flashinfer.cute_dsl.utils import is_cute_dsl_available + except (ImportError, ModuleNotFoundError) as error: + raise RuntimeError( + "FlashInfer low-M GEMM requires flashinfer-python " + f">={_MIN_FLASHINFER_VERSION} and nvidia-cutlass-dsl." + ) from error + version_value = getattr(flashinfer_module, "__version__", "unknown") + try: + flashinfer_version = Version(version_value) + except InvalidVersion as error: + raise RuntimeError(f"FlashInfer has an invalid version: {version_value!r}.") from error + if flashinfer_version < _MIN_FLASHINFER_VERSION: + raise RuntimeError( + "FlashInfer low-M GEMM requires flashinfer-python " + f">={_MIN_FLASHINFER_VERSION}; found {flashinfer_version}." + ) + if not is_cute_dsl_available(): + raise RuntimeError("FlashInfer low-M GEMM requires nvidia-cutlass-dsl.") + if not callable(getattr(mm_bf16, "is_backend_supported", None)): + raise RuntimeError( + "Installed FlashInfer mm_bf16 has no backend capability API; " + "the required PR #4266 implementation is not active." + ) + self._flashinfer_mm = mm_bf16 + logger.info( + "FlashInfer low-M BF16 GEMM is using the packaged cache-free " + f"direct/split-K heuristic from flashinfer-python {flashinfer_version}." + ) + + def _is_candidate_shape( + self, input_tensor: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] + ) -> bool: + # FlashInfer's DLPack bridge is inference-only and intentionally does + # not build an autograd graph. Preserve normal Linear semantics for + # training or any grad-enabled caller. + if ( + torch.is_grad_enabled() + or input_tensor.ndim < 1 + or weight.ndim != 2 + or not input_tensor.is_cuda + or not weight.is_cuda + or input_tensor.device != weight.device + or input_tensor.dtype != torch.bfloat16 + or weight.dtype != torch.bfloat16 + or not input_tensor.is_contiguous() + or not weight.is_contiguous() + ): + return False + k = int(input_tensor.shape[-1]) + if k <= 0 or int(weight.shape[1]) != k or k % 128: + return False + m = input_tensor.numel() // k + if not 1 <= m <= _MAX_FLASHINFER_M: + return False + if input_tensor.data_ptr() % 32 or weight.data_ptr() % 32: + return False + if bias is not None and ( + not bias.is_cuda + or bias.device != input_tensor.device + or bias.dtype != torch.bfloat16 + or bias.shape != (weight.shape[0],) + or not bias.is_contiguous() + ): + return False + sm = _current_sm(input_tensor.device) + if sm not in _SUPPORTED_SMS: + return False + return True + + def _is_flashinfer_supported(self, device: torch.device) -> bool: + if self._flashinfer_mm is None: + self._prepare_flashinfer() + sm = _current_sm(device) + return bool(self._flashinfer_mm.is_backend_supported(_FLASHINFER_BACKEND, sm)) + + def apply( + self, + module: torch.nn.Module, + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + if ( + self._collector is not None + and input_tensor.is_cuda + and input_tensor.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + ): + self._collector.record(module, input_tensor, weight, bias, self.cuda_graph_enabled) + if not self.enabled: + return None + if not self._prepared: + # Unit tests and direct module users do not construct ModelEngine. + self.prepare(module, cuda_graph_enabled=False) + if not self._is_candidate_shape(input_tensor, weight, bias): + return None + + input_2d = input_tensor.detach().view(-1, input_tensor.shape[-1]) + m = int(input_2d.shape[0]) + n = int(weight.shape[0]) + k = int(weight.shape[1]) + sm = _current_sm(input_tensor.device) + if self.backend == LowMGemmBackend.AUTO and _prefer_cublas_for_auto(m, n, k, sm): + return None + + if not self._is_flashinfer_supported(input_tensor.device): + return None + + engaged_key = (m, n, k, bias is not None) + if engaged_key not in self._engaged: + self._engaged.add(engaged_key) + logger.info( + f"Low-M BF16 GEMM: routing M={m} N={n} K={k} " + f"bias={bias is not None} to FlashInfer {_FLASHINFER_BACKEND}'s " + "packaged direct/split-K heuristic." + ) + + inference_weight = weight.detach() + inference_bias = bias.detach() if bias is not None else None + output = self._flashinfer_mm( + input_2d, + inference_weight.t(), + bias=inference_bias, + pdl=get_env_enable_pdl(), + out_dtype=torch.bfloat16, + backend=_FLASHINFER_BACKEND, + ) + return output.view(*input_tensor.shape[:-1], weight.shape[0]) + + def flush_shape_log(self) -> None: + if self._collector is not None: + self._collector.flush(mark_warmup_complete=True) + + +_DISPATCHER = LowMGemmDispatcher() +LOW_M_GEMM_ACTIVE = _DISPATCHER.enabled or _DISPATCHER._collector is not None + + +def prepare_low_m_gemm(model: torch.nn.Module, cuda_graph_enabled: bool) -> None: + """Label Linear modules and initialize the cache-free FlashInfer path.""" + + _DISPATCHER.prepare(model, cuda_graph_enabled) + + +def apply_low_m_gemm( + module: torch.nn.Module, + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + """Return a low-M BF16 result, or ``None`` for the normal path.""" + + return _DISPATCHER.apply(module, input_tensor, weight, bias) + + +def flush_low_m_gemm_shape_log() -> None: + """Persist the debug shape inventory at a deterministic warmup boundary.""" + + _DISPATCHER.flush_shape_log() + + +__all__ = [ + "LOW_M_GEMM_ACTIVE", + "LowMGemmBackend", + "apply_low_m_gemm", + "flush_low_m_gemm_shape_log", + "prepare_low_m_gemm", +] diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 96ef5223b7bb..146fda2cc9c7 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -31,6 +31,46 @@ REPLAY_WORK_PNAT = 2 REPLAY_WORK_CACHE_BUF_IDX = 3 REPLAY_WORK_ITEM_WIDTH = 4 +_FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE = 256 + + +@triton.jit +def _prepare_gdn_replay_work_items_kernel( + state_indices, + prev_num_accepted_tokens, + cache_buf_idx, + work_items, + n_writes_output, + num_decodes, + replay_step_width: tl.constexpr, + replay_history_size: tl.constexpr, + work_item_width: tl.constexpr, + position_field: tl.constexpr, + cache_slot_field: tl.constexpr, + pnat_field: tl.constexpr, + cache_buf_idx_field: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Build the write-first GDN replay partition in one launch.""" + offsets = tl.arange(0, BLOCK_SIZE) + active = offsets < num_decodes + slots = tl.load(state_indices + offsets, mask=active, other=0) + pnat = tl.load(prev_num_accepted_tokens + slots, mask=active, other=0) + active_buffer = tl.load(cache_buf_idx + slots, mask=active, other=0) + writes = active & (pnat + replay_step_width > replay_history_size) + writes_i32 = writes.to(tl.int32) + inclusive_write_offsets = tl.cumsum(writes_i32, axis=0) + write_offsets = inclusive_write_offsets - writes_i32 + n_writes = tl.sum(writes_i32, axis=0) + no_write_offsets = offsets - write_offsets + output_offsets = tl.where(writes, write_offsets, + n_writes + no_write_offsets) + output_base = work_items + output_offsets * work_item_width + tl.store(output_base + position_field, offsets, mask=active) + tl.store(output_base + cache_slot_field, slots, mask=active) + tl.store(output_base + pnat_field, pnat, mask=active) + tl.store(output_base + cache_buf_idx_field, active_buffer, mask=active) + tl.store(n_writes_output, n_writes) @triton.jit @@ -299,8 +339,9 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, self.replay_num_decodes = num_decodes if num_decodes == 0: return - if getattr(kv_cache_manager, 'use_gdn_cached_replay_all_layer_commit', - False): + use_gdn_all_layer_commit = getattr( + kv_cache_manager, "use_gdn_cached_replay_all_layer_commit", False) + if use_gdn_all_layer_commit: from tensorrt_llm._torch.modules.fla.cached_replay import \ CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE @@ -310,7 +351,6 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: return - self.replay_n_writes.zero_() if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): raise RuntimeError( "Replay state update is enabled, but the KV cache manager " @@ -327,6 +367,30 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, replay_step_width = replay_metadata.replay_step_width replay_history_size = replay_metadata.replay_history_size + if (use_gdn_all_layer_commit + and num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE): + block_size = triton.next_power_of_2(num_decodes) + _prepare_gdn_replay_work_items_kernel[(1, )]( + self.state_indices[num_contexts:batch_size], + prev_num_accepted_tokens, + cache_buf_idx, + self.replay_work_items, + self.replay_n_writes, + num_decodes, + replay_step_width=replay_step_width, + replay_history_size=replay_history_size, + work_item_width=REPLAY_WORK_ITEM_WIDTH, + position_field=REPLAY_WORK_POSITION_IN_DECODE_BATCH, + cache_slot_field=REPLAY_WORK_CACHE_SLOT, + pnat_field=REPLAY_WORK_PNAT, + cache_buf_idx_field=REPLAY_WORK_CACHE_BUF_IDX, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return + + self.replay_n_writes.zero_() + position_in_decode_batch = torch.arange( num_decodes, dtype=torch.int32, device=self.state_indices.device) cache_slot = self.state_indices[num_contexts:batch_size] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4906ad52f5e5..1fd88af71b14 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -493,6 +493,11 @@ def __init__( self.model_is_wrapped = True else: self.model_is_wrapped = False + from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm + if LOW_M_GEMM_ACTIVE: + prepare_low_m_gemm(self.model, + cuda_graph_enabled=llm_args.cuda_graph_config + is not None) self.sparse_attention_config = self.model.model_config.sparse_attention_config # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): @@ -1367,6 +1372,8 @@ def warmup(self, resource_manager: ResourceManager) -> None: # fails every step and padded batches silently run eager. self.cuda_graph_runner.preallocate_padding_dummies(resource_manager) log_mem_snapshot("warmup/after_preallocate_padding_dummies") + from ..modules.low_m_gemm import flush_low_m_gemm_shape_log + flush_low_m_gemm_shape_log() def _warmup_dg_paged_mqa_logits_metadata(self) -> None: """Pre-compile DeepGEMM's `get_paged_mqa_logits_metadata` helper for @@ -6904,8 +6911,12 @@ def encoder_forward(self, inputs: Dict[str, Any], attn_metadata = self._set_up_attn_metadata( kv_cache_manager=None ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata - graph_attn_metadata, key = self.encoder_cuda_graph_runner.maybe_get_cuda_graph( - padded_inputs, attn_metadata) + if moe_load_balancer is not None and moe_load_balancer.requires_eager_forward( + ): + graph_attn_metadata, key = None, None + else: + graph_attn_metadata, key = self.encoder_cuda_graph_runner.maybe_get_cuda_graph( + padded_inputs, attn_metadata) # Unpad seq_lens when fallback to eager path. if key is None: padded_inputs['seq_lens'] = padded_inputs[ @@ -7114,18 +7125,22 @@ def forward(self, ResourceManagerType.PEFT_CACHE_MANAGER) peft_cache_data_type = peft_cache_manager.data_type - maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( - padded_graph_requests, - enable_spec_decode=self.enable_spec_decode, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_tokens_cuda=self.draft_tokens_cuda - if self.is_spec_decode else None, - new_tensors_device=new_tensors_device, - spec_resource_manager=spec_resource_manager, - promoted_context_request_ids=promoted_context_request_ids, - peft_cache_data_type=peft_cache_data_type, - ) + if moe_load_balancer is not None and moe_load_balancer.requires_eager_forward( + ): + maybe_attn_metadata, maybe_spec_metadata, key = None, None, None + else: + maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( + padded_graph_requests, + enable_spec_decode=self.enable_spec_decode, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_tokens_cuda=self.draft_tokens_cuda + if self.is_spec_decode else None, + new_tensors_device=new_tensors_device, + spec_resource_manager=spec_resource_manager, + promoted_context_request_ids=promoted_context_request_ids, + peft_cache_data_type=peft_cache_data_type, + ) can_run_graph = key is not None if can_run_graph: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 60d801d904c7..630d5f402f26 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -619,6 +619,21 @@ def __init__( if self.attention_dp_enable_balance: self.attention_dp_time_out_iters = self.llm_args.attention_dp_config.timeout_iters self.attention_dp_batching_wait_iters = self.llm_args.attention_dp_config.batching_wait_iters + self.attention_dp_low_occupancy_timeout_iters = int( + os.environ.get( + "TLLM_ADP_BALANCE_LOW_OCCUPANCY_TIMEOUT_ITERS", + self.attention_dp_time_out_iters, + )) + self.attention_dp_min_generation_requests = int( + os.environ.get("TLLM_ADP_BALANCE_MIN_GENERATION_REQUESTS", 0)) + if self.attention_dp_low_occupancy_timeout_iters < 0: + raise ValueError( + "TLLM_ADP_BALANCE_LOW_OCCUPANCY_TIMEOUT_ITERS must be " + "greater than or equal to 0") + if not 0 <= self.attention_dp_min_generation_requests <= max_batch_size: + raise ValueError( + "TLLM_ADP_BALANCE_MIN_GENERATION_REQUESTS must be between " + f"0 and max_batch_size ({max_batch_size})") self.batch_wait_timeout_ms = self.llm_args.batch_wait_timeout_ms self.batch_wait_timeout_iters = self.llm_args.batch_wait_timeout_iters self.batch_wait_max_tokens_ratio = self.llm_args.batch_wait_max_tokens_ratio @@ -5409,7 +5424,11 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest], else: self.adp_ctx_waiting_iters_count += 1 balanced_context_requests = [] - timeout_reached = self.adp_ctx_waiting_iters_count >= self.attention_dp_time_out_iters + timeout_iters = self.attention_dp_time_out_iters + if (min(all_ranks_num_scheduled_generation_requests) + < self.attention_dp_min_generation_requests): + timeout_iters = self.attention_dp_low_occupancy_timeout_iters + timeout_reached = self.adp_ctx_waiting_iters_count >= timeout_iters if timeout_reached or not all_ranks_have_gen_requests: self.adp_ctx_waiting_iters_count = 0 balanced_context_requests = context_requests diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 161ac02eb995..6524ab9ff9d9 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -2582,3 +2582,56 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self): assert disagg_gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS # Context requests are not generation requests and must be left alone. assert ctx.num_draft_tokens == 0 + + +class TestAttentionDpLowOccupancyBalance: + @staticmethod + def _make_executor(all_ranks_requests, min_generation_requests=48): + executor = object.__new__(PyExecutor) + executor.dist = Mock() + executor.dist.tp_allgather.return_value = all_ranks_requests + executor.max_batch_size = 64 + executor.attention_dp_enable_balance = True + executor.attention_dp_time_out_iters = 60 + executor.attention_dp_low_occupancy_timeout_iters = 0 + executor.attention_dp_min_generation_requests = min_generation_requests + executor.attention_dp_batching_wait_iters = 0 + executor.adp_ctx_waiting_iters_count = 0 + executor.adp_ctx_batching_wait_iters_count = 0 + return executor + + @staticmethod + def _make_context_request(): + request = Mock() + request.get_tokens.return_value = [0] * 2048 + return request + + def test_saturated_imbalance_uses_configured_timeout(self): + executor = self._make_executor([[1, 63, 2111], [0, 63, 63]]) + context_request = self._make_context_request() + + balanced = executor._balance_adp_requests([context_request], [Mock()] * 63) + + assert balanced == [] + assert executor.adp_ctx_waiting_iters_count == 1 + executor.dist.tp_allgather.assert_called_once_with([1, 63, 2111]) + + def test_low_occupancy_on_any_rank_releases_context(self): + executor = self._make_executor([[1, 63, 2111], [0, 20, 20]]) + context_request = self._make_context_request() + + balanced = executor._balance_adp_requests([context_request], [Mock()] * 63) + + assert balanced == [context_request] + assert executor.adp_ctx_waiting_iters_count == 0 + executor.dist.tp_allgather.assert_called_once_with([1, 63, 2111]) + + def test_default_threshold_preserves_existing_behavior(self): + executor = self._make_executor([[1, 20, 2068], [0, 20, 20]], min_generation_requests=0) + context_request = self._make_context_request() + + balanced = executor._balance_adp_requests([context_request], [Mock()] * 20) + + assert balanced == [] + assert executor.adp_ctx_waiting_iters_count == 1 + executor.dist.tp_allgather.assert_called_once_with([1, 20, 2068]) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py index 36c9d8b0c9f7..753bc5b7613a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py @@ -18,9 +18,13 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import Qwen3_5MoeForCausalLM, Qwen3_5MoeVLModel from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM -from tensorrt_llm._torch.models.modeling_qwen3_5 import _normalize_qwen35_moe_vl_config +from tensorrt_llm._torch.models.modeling_qwen3_5 import ( + _filter_language_model_weights, + _normalize_qwen35_moe_vl_config, +) from tensorrt_llm._torch.pyexecutor.config_utils import ( extract_mamba_kv_cache_params, load_pretrained_config, @@ -198,6 +202,23 @@ def test_qwen35_moe_model_defaults( assert llm_args.nvfp4_gemm_config.allowed_backends == expected_gemm_backends +def test_qwen35_vl_filter_preserves_consumable_weights() -> None: + language_weight = torch.tensor([1.0]) + weights = ConsumableWeightsDict( + { + "model.language_model.layers.0.weight": language_weight, + "model.visual.patch_embed.weight": torch.tensor([2.0]), + } + ) + + filtered_weights = _filter_language_model_weights(weights) + + assert isinstance(filtered_weights, ConsumableWeightsDict) + assert len(weights) == 0 + assert len(filtered_weights) == 1 + assert filtered_weights["model.language_model.layers.0.weight"] is language_weight + + def test_qwen35_moe_vl_placeholder_metadata_registered() -> None: metadata = MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata("qwen3_5_moe") diff --git a/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py b/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py new file mode 100644 index 000000000000..637c56ad71d5 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest import mock + +import pytest + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_utils import ( + _get_load_weights_num_workers, + _pageout_safetensors_after_moe_load, + rename_weights_with_regex, +) + + +def test_filter_prefix_matches_only_subtree_and_strips_prefix(): + weights = ConsumableWeightsDict( + { + "model.layers.1.weight": 1, + "model.layers.1.bias": 2, + "model.layers.10.weight": 3, + } + ) + + assert weights.filter_prefix("model.layers.1") == { + "weight": 1, + "bias": 2, + } + + +def test_prefix_index_tracks_new_and_consumed_keys(): + weights = ConsumableWeightsDict({"a.weight": 1}) + + assert weights.filter_prefix("a") == {"weight": 1} + weights["a.bias"] = 2 + weights.update({"a.scale": 3, "b.weight": 4}) + assert weights.filter_prefix("a") == { + "weight": 1, + "bias": 2, + "scale": 3, + } + assert weights.mark_consumed("a") == 3 + assert weights.filter_prefix("a") == {} + assert weights.filter_prefix("b") == {"weight": 4} + + +def test_prefix_index_handles_replacement_deletion_and_clear(): + weights = ConsumableWeightsDict({"a.weight": 1, "b.weight": 2}) + + assert weights.filter_prefix("a") == {"weight": 1} + weights["a.weight"] = 3 + del weights["a.weight"] + assert weights.filter_prefix("a") == {} + assert weights.filter_prefix("") == {"b.weight": 2} + weights.clear() + assert weights.filter_prefix("") == {} + assert weights.mark_consumed("b") == 0 + + +def test_mark_consumed_keys_accepts_missing_and_duplicate_keys(): + weights = ConsumableWeightsDict({"a.weight": 1, "a.bias": 2}) + + assert weights.mark_consumed_keys(["a.weight", "missing", "a.weight"]) == 1 + assert weights.filter_prefix("a") == {"bias": 2} + + +def test_regex_rename_transfers_consumable_weight_ownership(): + tensor = object() + weights = ConsumableWeightsDict( + { + "model.language_model.layer.weight": tensor, + "model.visual.weight": object(), + } + ) + + renamed = rename_weights_with_regex({r"^model\.language_model\.(.*)$": r"model.\1"}, weights) + + assert isinstance(renamed, ConsumableWeightsDict) + assert len(weights) == 0 + assert renamed["model.layer.weight"] is tensor + assert "model.visual.weight" in renamed + + +def test_pageout_safetensors_after_moe_load(): + with ( + mock.patch("torch.cuda.synchronize") as synchronize, + mock.patch( + "tensorrt_llm._torch.mmap_utils.pageout_file_backed_regions" + ) as pageout_file_backed_regions, + ): + _pageout_safetensors_after_moe_load() + + synchronize.assert_called_once_with() + pageout_file_backed_regions.assert_called_once_with(".safetensors", mode="dontneed") + + +@pytest.mark.parametrize( + "value, expected", [(None, None), ("", None), (" ", None), ("1", 1), ("4", 4)] +) +def test_get_load_weights_num_workers(monkeypatch, value, expected): + env_name = "TRT_LLM_LOAD_WEIGHTS_NUM_WORKERS" + if value is None: + monkeypatch.delenv(env_name, raising=False) + else: + monkeypatch.setenv(env_name, value) + + assert _get_load_weights_num_workers() == expected + + +@pytest.mark.parametrize("value", ["0", "-1", "1.5", "workers"]) +def test_get_load_weights_num_workers_rejects_invalid_values(monkeypatch, value): + monkeypatch.setenv("TRT_LLM_LOAD_WEIGHTS_NUM_WORKERS", value) + + with pytest.raises(ValueError, match="must be a positive integer"): + _get_load_weights_num_workers() diff --git a/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py index e0a73d9c1f05..f00cb74fc320 100644 --- a/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py +++ b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py @@ -16,16 +16,20 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest import torch from torch import nn from tensorrt_llm._torch.distributed import AllReduceFusionOp from tensorrt_llm._torch.models.modeling_qwen3_next import ( Qwen3NextForCausalLM, + Qwen3NextFullAttentionDecoderLayer, Qwen3NextLinearDecoderLayer, + Qwen3NextMTP, _eager_fusion_enabled, ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.utils import EventType def _new_causal_lm() -> Qwen3NextForCausalLM: @@ -34,6 +38,76 @@ def _new_causal_lm() -> Qwen3NextForCausalLM: return model +def _new_mtp() -> Qwen3NextMTP: + model = Qwen3NextMTP.__new__(Qwen3NextMTP) + nn.Module.__init__(model) + model.pre_fc_norm_embedding = nn.Identity() + model.pre_fc_norm_hidden = nn.Identity() + model.fc = nn.Identity() + model.shared_head = nn.Module() + model.shared_head.norm = MagicMock( + side_effect=lambda hidden_states, residual: (hidden_states, None) + ) + model.event_dict = {EventType.Main: None, EventType.MoeShared: None} + model.aux_stream = None + model.model_config = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, tp_rank=0, enable_attention_dp=True) + ) + return model + + +@torch.no_grad() +def test_mtp_forward_uses_and_restores_draft_rank_token_counts(monkeypatch) -> None: + model = _new_mtp() + target_rank_tokens = [1024, 1024] + draft_rank_tokens = [1, 1] + attn_metadata = SimpleNamespace(all_rank_num_tokens=target_rank_tokens) + + def decoder_forward(self, **kwargs): + assert kwargs["attn_metadata"].all_rank_num_tokens is draft_rank_tokens + return kwargs["hidden_states"], None + + monkeypatch.setattr(Qwen3NextFullAttentionDecoderLayer, "forward", decoder_forward) + + hidden_states = model( + input_ids=torch.tensor([0]), + position_ids=torch.tensor([0]), + hidden_states=torch.ones(1, 2), + embed_tokens=nn.Embedding(1, 2), + attn_metadata=attn_metadata, + all_rank_num_tokens=draft_rank_tokens, + ) + + assert hidden_states.shape == (1, 4) + assert attn_metadata.all_rank_num_tokens is target_rank_tokens + + +@torch.no_grad() +def test_mtp_forward_restores_rank_token_counts_after_failure(monkeypatch) -> None: + model = _new_mtp() + target_rank_tokens = [1024, 1024] + draft_rank_tokens = [1, 1] + attn_metadata = SimpleNamespace(all_rank_num_tokens=target_rank_tokens) + + def decoder_forward(self, **kwargs): + assert kwargs["attn_metadata"].all_rank_num_tokens is draft_rank_tokens + raise RuntimeError("draft forward failed") + + monkeypatch.setattr(Qwen3NextFullAttentionDecoderLayer, "forward", decoder_forward) + + with pytest.raises(RuntimeError, match="draft forward failed"): + model( + input_ids=torch.tensor([0]), + position_ids=torch.tensor([0]), + hidden_states=torch.ones(1, 2), + embed_tokens=nn.Embedding(1, 2), + attn_metadata=attn_metadata, + all_rank_num_tokens=draft_rank_tokens, + ) + + assert attn_metadata.all_rank_num_tokens is target_rank_tokens + + @torch.no_grad() def test_setup_aliases_does_not_read_meta_weights() -> None: model = _new_causal_lm() diff --git a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py index e1558d0aab4a..c026d158a3b6 100644 --- a/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py +++ b/tests/unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py @@ -69,6 +69,9 @@ class ExpandQuantShape: SHAPES = [ ExpandQuantShape("anchor", num_source_tokens=32, hidden=4096, num_experts=128, top_k=8), ExpandQuantShape("anchor_h7168", num_source_tokens=32, hidden=7168, num_experts=128, top_k=8), + ExpandQuantShape( + "large_h8192_topk10", num_source_tokens=32, hidden=8192, num_experts=512, top_k=10 + ), ExpandQuantShape("small", num_source_tokens=1, hidden=512, num_experts=8, top_k=4), ExpandQuantShape("medium_batch", num_source_tokens=64, hidden=7168, num_experts=128, top_k=8), ExpandQuantShape("topk4", num_source_tokens=16, hidden=4096, num_experts=64, top_k=4), @@ -96,6 +99,45 @@ def _alloc_outputs(num_experts: int, m_max: int, hidden: int, *, device: str): return output_q, output_s +@skip_unsupported +def test_skip_data_expand_omits_unused_outputs() -> None: + num_rows, hidden, num_experts, top_k = 4, 512, 8, 4 + x = torch.randn((num_rows, hidden), device="cuda", dtype=torch.float32) + token_selected_experts = torch.arange(top_k, device="cuda", dtype=torch.int32).repeat( + num_rows, 1 + ) + token_final_scales = torch.full( + (num_rows, top_k), 1.0 / top_k, device="cuda", dtype=torch.float32 + ) + + outputs = torch.ops.trtllm.moe_permute_op( + x, + token_selected_experts, + token_final_scales, + None, + None, + None, + input_sf=None, + num_experts_on_rank=num_experts, + tp_size=1, + tp_rank=0, + ep_size=1, + ep_rank=0, + cluster_size=1, + cluster_rank=0, + min_latency_mode=False, + use_fp8_block_scaling=False, + skip_data_expand=True, + ) + + permuted_row_to_unpermuted_row_tensor = outputs[0] + permuted_data_tensor = outputs[2] + permuted_token_final_scales_tensor = outputs[4] + assert permuted_row_to_unpermuted_row_tensor.numel() == num_rows * top_k + assert permuted_data_tensor.shape == (0, hidden) + assert permuted_token_final_scales_tensor.shape == (0,) + + @skip_unsupported @pytest.mark.parametrize("shape", SHAPES, ids=lambda s: s.name) def test_fused_expand_quant_matches_unfused(shape: ExpandQuantShape) -> None: @@ -156,6 +198,54 @@ def test_fused_expand_quant_matches_unfused(shape: ExpandQuantShape) -> None: # hold; moe_permute_op produces it in fp32. assert permuted_data_tensor.dtype == torch.float32 + # The fused path needs only the permutation maps. Verify that skipping the + # expand copy preserves those maps without retaining the top-k-materialized + # activation and scale buffers. + ( + fused_permuted_row_to_unpermuted_row_tensor, + _fused_permuted_token_selected_experts_tensor, + skipped_permuted_data_tensor, + fused_expert_first_token_offset_tensor, + skipped_permuted_token_final_scales_tensor, + _fused_unpermuted_row_to_permuted_row_tensor, + ) = torch.ops.trtllm.moe_permute_op( + x[:, :0], + token_selected_experts, + token_final_scales, + None, + None, + None, + input_sf=None, + num_experts_on_rank=num_experts_per_node, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + cluster_size=cluster_size, + cluster_rank=cluster_rank, + min_latency_mode=False, + use_fp8_block_scaling=False, + skip_data_expand=True, + ) + + assert skipped_permuted_data_tensor.numel() == 0 + # Newer operator libraries omit the unused scales too. Older libraries may + # retain this small 1-D return, but the top-k activation storage must be + # empty in either case. + assert skipped_permuted_token_final_scales_tensor.dim() == 1 + torch.testing.assert_close( + fused_permuted_row_to_unpermuted_row_tensor, + permuted_row_to_unpermuted_row_tensor, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + fused_expert_first_token_offset_tensor, + expert_first_token_offset_tensor, + rtol=0, + atol=0, + ) + _masked_m, token_to_expert_map = preprocess_after_permute( expert_first_token_offset_tensor, permuted_data_tensor.shape[0] ) @@ -179,8 +269,8 @@ def test_fused_expand_quant_matches_unfused(shape: ExpandQuantShape) -> None: out_q_new, out_s_new, x, - permuted_row_to_unpermuted_row_tensor, - expert_first_token_offset_tensor, + fused_permuted_row_to_unpermuted_row_tensor, + fused_expert_first_token_offset_tensor, token_to_expert_map, experts_per_token=top_k, group_size=GROUP_SIZE, diff --git a/tests/unittest/_torch/modules/mamba/test_flashinfer_chunk_gdn.py b/tests/unittest/_torch/modules/mamba/test_flashinfer_chunk_gdn.py index 834b76ea5ba1..8815cbd767ca 100644 --- a/tests/unittest/_torch/modules/mamba/test_flashinfer_chunk_gdn.py +++ b/tests/unittest/_torch/modules/mamba/test_flashinfer_chunk_gdn.py @@ -99,6 +99,38 @@ def test_wrapper_module_importable(): ) +def test_cutlass_dsl_45_pipeline_compatibility(): + """The FlashInfer nightly's explicit False keyword is a no-op on DSL 4.5.""" + from tensorrt_llm._torch.modules.fla.flashinfer_chunk import ( + _enable_cutlass_dsl_45_pipeline_compatibility, + ) + + calls = [] + + class PipelineTmaAsync: + @staticmethod + def create(*, num_stages): + calls.append(num_stages) + return num_stages + + class Pipeline: + pass + + Pipeline.PipelineTmaAsync = PipelineTmaAsync + + _enable_cutlass_dsl_45_pipeline_compatibility(Pipeline) + assert Pipeline.PipelineTmaAsync.create(num_stages=3, enable_multicast_signaling=False) == 3 + assert calls == [3] + + # Repeated inference calls must not wrap the global CUTLASS class again. + first_compatible_create = Pipeline.PipelineTmaAsync.create + _enable_cutlass_dsl_45_pipeline_compatibility(Pipeline) + assert Pipeline.PipelineTmaAsync.create is first_compatible_create + + with pytest.raises(ValueError, match="nvidia-cutlass-dsl>=4.6"): + Pipeline.PipelineTmaAsync.create(num_stages=3, enable_multicast_signaling=True) + + # Parity tests against the Triton reference ------------------------------- diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py b/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py index 4dc17d9bda1b..5e31ad0c28eb 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py @@ -26,6 +26,7 @@ import torch from tensorrt_llm._torch.modules.fla.cached_replay import ( + _supports_ratio8_replay_tiling, fused_recurrent_gated_delta_rule_cached_replay_update, ) from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule_update @@ -61,6 +62,19 @@ def _seq_ref_step(S, q, k, v, g, beta, scale): return torch.stack(outs), states +@pytest.mark.parametrize( + "H,HV,head_tiles,expected", + [ + (1, 8, 8, True), + (1, 8, 512, True), + (1, 8, 520, False), + (2, 8, 512, False), + ], +) +def test_gdn_replay_ratio8_tiling_boundary(H, HV, head_tiles, expected): + assert _supports_ratio8_replay_tiling(H, HV, head_tiles) is expected + + @pytest.mark.parametrize("fused_gating", [False, True], ids=["pre_gated", "fused_gating"]) @pytest.mark.parametrize( "pool_dtype", [torch.bfloat16, torch.float32], ids=["bf16_pool", "fp32_pool"] @@ -72,8 +86,13 @@ def _seq_ref_step(S, q, k, v, g, beta, scale): ) @pytest.mark.parametrize( "H,HV,K,V", - [(4, 8, 128, 128), (2, 4, 64, 64), (4, 16, 128, 128)], - ids=["qwen3_like", "small", "qwen3_5_like"], + [ + (4, 8, 128, 128), + (2, 4, 64, 64), + (4, 16, 128, 128), + (1, 8, 128, 128), + ], + ids=["qwen3_like", "small", "qwen3_5_like", "ratio8"], ) def test_gdn_replay_vs_legacy_and_ref(H, HV, K, V, T, HIST, iters, pool_dtype, fused_gating): if not torch.cuda.is_available(): @@ -431,8 +450,13 @@ def test_gdn_cached_replay_all_layer_commit_matches_reference(state_layout): torch.testing.assert_close(actual_states.float(), expected_states.float(), rtol=2e-2, atol=2e-2) +@pytest.mark.parametrize( + "H,HV,K,V", + [(2, 4, 64, 64), (1, 8, 128, 128)], + ids=["ratio2", "ratio8"], +) @pytest.mark.parametrize("batch_size", [8, 16], ids=["small_fused", "large_all_layer"]) -def test_gdn_cached_replay_dispatch_cuda_graph_matches_eager(batch_size): +def test_gdn_cached_replay_dispatch_cuda_graph_matches_eager(batch_size, H, HV, K, V): """Both sides of the BS16 dispatch must be safe under CUDA graphs.""" if not torch.cuda.is_available(): pytest.skip("CUDA required") @@ -446,7 +470,6 @@ def test_gdn_cached_replay_dispatch_cuda_graph_matches_eager(batch_size): device = "cuda" dtype = torch.bfloat16 num_layers, T, HIST = 2, 4, 16 - H, HV, K, V = 2, 4, 64, 64 num_slots = batch_size + 1 use_all_layer_commit = batch_size >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 8aa2a847802b..970df255aba1 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -39,6 +39,44 @@ ) +class _GdnReplayCacheManager: + use_replay_state_update = True + use_gdn_cached_replay_all_layer_commit = True + + def __init__(self, prev_num_accepted_tokens, cache_buf_idx): + self.prev_num_accepted_tokens = prev_num_accepted_tokens + self.cache_buf_idx = cache_buf_idx + + def get_replay_state_update_metadata(self): + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=6, + replay_history_size=MIN_REPLAY_HISTORY_SIZE, + ) + + +def _reference_gdn_replay_work_items(state_indices, prev_num_accepted_tokens, cache_buf_idx): + positions = torch.arange(state_indices.numel(), dtype=torch.int32, device="cuda") + cache_slots = state_indices.to(torch.int32) + slot_indices = cache_slots.to(torch.long) + pnat = prev_num_accepted_tokens[slot_indices].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[slot_indices].to(torch.int32) + writes = pnat + 6 > MIN_REPLAY_HISTORY_SIZE + writes_i32 = writes.to(torch.int32) + write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 + n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) + output_offsets = torch.where(writes, write_offsets, n_writes + positions - write_offsets).to( + torch.long + ) + output = torch.empty(state_indices.numel(), 4, dtype=torch.int32, device="cuda") + output[output_offsets, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = positions + output[output_offsets, REPLAY_WORK_CACHE_SLOT] = cache_slots + output[output_offsets, REPLAY_WORK_PNAT] = pnat + output[output_offsets, REPLAY_WORK_CACHE_BUF_IDX] = active_cache_buf_idx + return output, n_writes + + @skip_no_cuda class TestCuSeqlensToChunkIndicesOffsets: """Tests for cu_seqlens_to_chunk_indices_offsets_triton function.""" @@ -155,6 +193,60 @@ def get_replay_state_update_metadata(self): assert actual[0, REPLAY_WORK_PNAT] == 11 assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 + @pytest.mark.parametrize("num_decodes", [16, 40, 256, 257]) + def test_prepare_gdn_replay_work_items_matches_reference(self, num_decodes): + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + state_indices = torch.randperm(num_slots, device="cuda")[:num_decodes].to(torch.int32) + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) + metadata.state_indices.copy_(state_indices) + + metadata._prepare_replay_work_items(manager, num_decodes, 0) + expected_items, expected_n_writes = _reference_gdn_replay_work_items( + state_indices, prev_num_accepted_tokens, cache_buf_idx + ) + + torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) + torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + + def test_prepare_gdn_replay_work_items_cuda_graph_replay(self): + num_decodes = 40 + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) + metadata.state_indices.copy_(torch.arange(num_decodes, dtype=torch.int32, device="cuda")) + + metadata._prepare_replay_work_items(manager, num_decodes, 0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + metadata._prepare_replay_work_items(manager, num_decodes, 0) + + updated_state_indices = torch.arange( + num_slots - 1, + num_slots - num_decodes - 1, + -1, + dtype=torch.int32, + device="cuda", + ) + metadata.state_indices.copy_(updated_state_indices) + prev_num_accepted_tokens.copy_( + (torch.arange(num_slots, dtype=torch.int32, device="cuda") * 7) % 21 + ) + cache_buf_idx.bitwise_xor_(1) + graph.replay() + torch.cuda.synchronize() + + expected_items, expected_n_writes = _reference_gdn_replay_work_items( + updated_state_indices, prev_num_accepted_tokens, cache_buf_idx + ) + torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) + torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + def test_single_sequence_unaligned(self): """Test with a single sequence that doesn't align with chunk size.""" cu_seqlens = torch.tensor([0, 10], dtype=torch.int, device="cuda") diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 3919ea233af4..ea2c3f542b98 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -198,6 +198,35 @@ def test_forced_nccl_ep_allows_missing_moe_max_num_tokens( assert strategy.moe_max_num_tokens is None +def test_forced_strategy_logs_selected_implementation(monkeypatch: pytest.MonkeyPatch): + class _FakeForcedStrategy: + pass + + messages = [] + monkeypatch.setenv("TRTLLM_FORCE_COMM_METHOD", "NVLINK_ONE_SIDED") + monkeypatch.setattr( + communication_factory.CommunicationFactory, + "_create_forced_method", + lambda *args, **kwargs: _FakeForcedStrategy(), + ) + monkeypatch.setattr(communication_factory.logger, "info", messages.append) + + strategy = communication_factory.CommunicationFactory.create_strategy( + _make_model_config(), + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + hidden_size=4096, + ) + + assert isinstance(strategy, _FakeForcedStrategy) + assert messages == [ + "Selected communication strategy: _FakeForcedStrategy " + "(forced by TRTLLM_FORCE_COMM_METHOD=NVLINK_ONE_SIDED)" + ] + + def test_auto_selection_uses_nccl_ep_with_missing_moe_max_num_tokens( monkeypatch: pytest.MonkeyPatch, ): diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 4b8764c548a6..697bd54f60d5 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -60,6 +60,10 @@ ) from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import ( + _DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS, + _configure_deepgemm_moe_max_num_tokens, +) from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoECommPlan, MoERunContext from tensorrt_llm._torch.modules.fused_moe.interface import ( @@ -89,6 +93,32 @@ } +@pytest.mark.parametrize( + ("max_num_tokens", "expected_value"), + [ + pytest.param(8192, 8192, id="below-conservative-default"), + pytest.param(65536, _DEFAULT_DEEPGEMM_MOE_MAX_NUM_TOKENS, id="above-conservative-default"), + ], +) +def test_deepgemm_clamps_only_derived_moe_max_num_tokens(max_num_tokens, expected_value): + model_config = ModelConfig(max_num_tokens=max_num_tokens) + model_config._frozen = True + + _configure_deepgemm_moe_max_num_tokens(model_config) + + assert model_config.moe_max_num_tokens == expected_value + assert model_config._frozen is True + + +@pytest.mark.parametrize("configured_value", [32768, 65536]) +def test_deepgemm_preserves_explicit_moe_max_num_tokens(configured_value): + model_config = ModelConfig(moe_max_num_tokens=configured_value) + + _configure_deepgemm_moe_max_num_tokens(model_config) + + assert model_config.moe_max_num_tokens == configured_value + + def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): valid_tactics = [ [8, 0], diff --git a/tests/unittest/_torch/modules/moe/test_moe_scheduler.py b/tests/unittest/_torch/modules/moe/test_moe_scheduler.py new file mode 100644 index 000000000000..eb731c95df6f --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_moe_scheduler.py @@ -0,0 +1,71 @@ +# 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. + +"""Focused tests for MoE scheduler chunk execution.""" + +from types import MethodType + +import pytest +import torch + +from tensorrt_llm._torch.modules.fused_moe.moe_scheduler import ExternalCommMoEScheduler + + +class _DummyMoe: + use_dp = False + enable_alltoall = False + aux_stream = None + backend = object() + repeat_idx = 0 + repeat_count = 1 + + @staticmethod + def split_chunk(num_tokens: int, num_chunks: int) -> list[int]: + quotient, remainder = divmod(num_tokens, num_chunks) + return [quotient + (chunk_idx < remainder) for chunk_idx in range(num_chunks)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_multichunk_forward_is_cuda_graph_capturable() -> None: + """Host-only chunk bookkeeping must not become a CUDA tensor.""" + scheduler = ExternalCommMoEScheduler(_DummyMoe()) + + def forward_chunk(_self, x_chunk, _router_logits, *_args, **_kwargs): + return x_chunk + 1 + + scheduler._forward_chunk_impl = MethodType(forward_chunk, scheduler) + x = torch.arange(8, dtype=torch.float32, device="cuda").reshape(4, 2) + router_logits = torch.zeros((4, 2), dtype=torch.float32, device="cuda") + + def forward() -> torch.Tensor: + return scheduler._forward_multiple_chunks( + x, + router_logits, + num_chunks=2, + output_dtype=None, + all_rank_num_tokens=[4], + use_dp_padding=False, + ) + + with torch.device("cuda"): + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = forward() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(output, x + 1) diff --git a/tests/unittest/_torch/modules/test_low_m_gemm.py b/tests/unittest/_torch/modules/test_low_m_gemm.py new file mode 100644 index 000000000000..1250c638f2ee --- /dev/null +++ b/tests/unittest/_torch/modules/test_low_m_gemm.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import atexit +import json +import sys +from types import ModuleType +from unittest.mock import MagicMock + +import pytest +import torch + +from tensorrt_llm._torch.modules import linear as linear_module +from tensorrt_llm._torch.modules.low_m_gemm import ( + LowMGemmBackend, + LowMGemmDispatcher, + _configured_backend, + _device_sm, + _infer_op_name, + _normalize_backend, + _prefer_cublas_for_auto, + _ShapeCollector, +) + + +@pytest.mark.parametrize( + "value,expected", + [ + ("off", LowMGemmBackend.OFF), + ("auto", LowMGemmBackend.AUTO), + ("cute-dsl", LowMGemmBackend.FLASHINFER), + ("flashinfer_cute_dsl", LowMGemmBackend.FLASHINFER), + ("cublaslt", LowMGemmBackend.CUBLAS), + ], +) +def test_normalize_backend(value: str, expected: LowMGemmBackend) -> None: + assert _normalize_backend(value) == expected + + +def test_normalize_backend_rejects_unknown_value() -> None: + with pytest.raises(ValueError, match="TRTLLM_LOW_M_GEMM_BACKEND"): + _normalize_backend("split-all-shapes") + + +def test_legacy_exact_m_backend_environment_is_ignored(monkeypatch) -> None: + monkeypatch.delenv("TRTLLM_LOW_M_GEMM_BACKEND", raising=False) + monkeypatch.setenv("TLLM_BF16_GEMM_BACKEND", "heuristic") + assert _configured_backend() == LowMGemmBackend.OFF + + +def test_device_sm_capability_is_cached(monkeypatch) -> None: + get_device_capability = MagicMock(return_value=(10, 3)) + monkeypatch.setattr(torch.cuda, "get_device_capability", get_device_capability) + _device_sm.cache_clear() + try: + assert _device_sm(7) == 103 + assert _device_sm(7) == 103 + get_device_capability.assert_called_once_with(7) + finally: + _device_sm.cache_clear() + + +def _install_fake_flashinfer( + monkeypatch, + *, + version: str = "0.6.17.dev20260806", + cute_dsl_available: bool = True, +) -> MagicMock: + mm_bf16 = MagicMock() + mm_bf16.is_backend_supported = MagicMock(return_value=True) + + flashinfer_module = ModuleType("flashinfer") + flashinfer_module.__path__ = [] + flashinfer_module.__version__ = version + flashinfer_module.mm_bf16 = mm_bf16 + cute_dsl_module = ModuleType("flashinfer.cute_dsl") + cute_dsl_module.__path__ = [] + cute_dsl_utils_module = ModuleType("flashinfer.cute_dsl.utils") + cute_dsl_utils_module.is_cute_dsl_available = lambda: cute_dsl_available + + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer_module) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl", cute_dsl_module) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl.utils", cute_dsl_utils_module) + return mm_bf16 + + +def test_auto_prepare_uses_packaged_flashinfer_without_tuning_cache(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", "auto") + monkeypatch.setenv("TRTLLM_LOW_M_GEMM_TUNING_CACHE", "/missing/dispatch.json") + monkeypatch.setenv("TRTLLM_FLASHINFER_AUTOTUNER_CACHE", "/missing/flashinfer.json") + mm_bf16 = _install_fake_flashinfer(monkeypatch) + module = torch.nn.Linear(8, 8) + + dispatcher = LowMGemmDispatcher() + dispatcher.prepare(module, cuda_graph_enabled=True) + + assert dispatcher._flashinfer_mm is mm_bf16 + assert dispatcher._prepared + assert module._low_m_gemm_name == "" + + +def test_prepare_rejects_flashinfer_before_split_k_nightly(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", "auto") + _install_fake_flashinfer(monkeypatch, version="0.6.15") + + with pytest.raises(RuntimeError, match="0.6.17.dev20260806"): + LowMGemmDispatcher().prepare(torch.nn.Linear(8, 8), cuda_graph_enabled=False) + + +@pytest.mark.parametrize(("m", "expected"), ((32, True), (33, False))) +def test_flashinfer_candidate_shape_stops_at_m32(m: int, expected: bool, monkeypatch) -> None: + device = torch.device("cuda") + input_tensor = MagicMock( + ndim=2, + is_cuda=True, + device=device, + dtype=torch.bfloat16, + shape=(m, 128), + ) + input_tensor.is_contiguous.return_value = True + input_tensor.numel.return_value = m * 128 + input_tensor.data_ptr.return_value = 32 + weight = MagicMock( + ndim=2, + is_cuda=True, + device=device, + dtype=torch.bfloat16, + shape=(256, 128), + ) + weight.is_contiguous.return_value = True + weight.data_ptr.return_value = 32 + monkeypatch.setattr("tensorrt_llm._torch.modules.low_m_gemm._current_sm", lambda unused: 103) + + with torch.inference_mode(): + assert LowMGemmDispatcher()._is_candidate_shape(input_tensor, weight, None) is expected + + +def test_linear_fast_rejects_m_above_flashinfer_domain(monkeypatch) -> None: + monkeypatch.setattr(linear_module, "LOW_M_GEMM_ACTIVE", True) + monkeypatch.setattr(linear_module, "_LOW_M_GEMM_SHAPE_COLLECTION_ACTIVE", False) + + assert linear_module._should_apply_low_m_gemm(torch.empty((32, 128))) + assert not linear_module._should_apply_low_m_gemm(torch.empty((33, 128))) + + +def test_linear_fast_reject_preserves_full_shape_collection(monkeypatch) -> None: + monkeypatch.setattr(linear_module, "LOW_M_GEMM_ACTIVE", True) + monkeypatch.setattr(linear_module, "_LOW_M_GEMM_SHAPE_COLLECTION_ACTIVE", True) + + assert linear_module._should_apply_low_m_gemm(torch.empty((64, 128))) + + +@pytest.mark.parametrize( + "layer,expected", + [ + ("LMHead", "lm_head"), + ("model.layers.92.fc", "mtp_fusion_fc"), + ("model.layers.0.self_attn.qkv_proj", "attention_qkv"), + ], +) +def test_infer_op_name_for_hot_modules(layer: str, expected: str) -> None: + assert _infer_op_name(layer) == expected + + +@pytest.mark.parametrize( + "m,n,k,expected", + [ + (7, 8192, 128, False), + (8, 8192, 128, True), + (4, 8192, 2048, False), + (5, 8192, 2048, True), + (15, 8192, 2048, True), + (16, 8192, 2048, False), + (17, 8192, 2048, True), + (20, 8192, 2048, False), + (21, 8192, 2048, True), + (24, 8192, 2048, True), + (25, 8192, 2048, False), + (14, 15520, 8192, False), + (15, 15520, 8192, True), + (15, 2304, 8192, False), + (16, 2304, 8192, True), + (17, 2304, 8192, False), + (32, 8192, 1024, False), + ], +) +def test_auto_cublas_crossover_for_blackwell(m: int, n: int, k: int, expected: bool) -> None: + assert _prefer_cublas_for_auto(m, n, k, sm=103) is expected + assert not _prefer_cublas_for_auto(m, n, k, sm=100) + + +def test_apply_uses_public_flashinfer_heuristic(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", "auto") + monkeypatch.setenv("TRTLLM_ENABLE_PDL", "0") + dispatcher = LowMGemmDispatcher() + dispatcher._prepared = True + dispatcher._flashinfer_mm = MagicMock(return_value=torch.empty((4, 256))) + monkeypatch.setattr(dispatcher, "_is_candidate_shape", lambda *unused: True) + monkeypatch.setattr(dispatcher, "_is_flashinfer_supported", lambda unused: True) + monkeypatch.setattr("tensorrt_llm._torch.modules.low_m_gemm._current_sm", lambda unused: 103) + + input_tensor = torch.empty((2, 2, 128), dtype=torch.bfloat16) + weight = torch.empty((256, 128), dtype=torch.bfloat16) + bias = torch.empty((256,), dtype=torch.bfloat16) + with torch.inference_mode(): + output = dispatcher.apply(torch.nn.Linear(1, 1), input_tensor, weight, bias) + + assert output.shape == (2, 2, 256) + args, kwargs = dispatcher._flashinfer_mm.call_args + assert args[0].shape == (4, 128) + assert args[1].shape == (128, 256) + assert kwargs["bias"].data_ptr() == bias.data_ptr() + assert kwargs["bias"].shape == bias.shape + assert kwargs["bias"].dtype == bias.dtype + assert kwargs["pdl"] is False + assert kwargs["out_dtype"] == torch.bfloat16 + assert kwargs["backend"] == "cute-dsl" + + +@pytest.mark.parametrize( + "backend,expected_flashinfer_calls", + [("auto", 0), ("flashinfer", 1)], +) +def test_explicit_flashinfer_overrides_auto_cublas_crossover( + monkeypatch, backend: str, expected_flashinfer_calls: int +) -> None: + monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", backend) + dispatcher = LowMGemmDispatcher() + dispatcher._prepared = True + dispatcher._flashinfer_mm = MagicMock(return_value=torch.empty((8, 8192))) + monkeypatch.setattr(dispatcher, "_is_candidate_shape", lambda *unused: True) + monkeypatch.setattr(dispatcher, "_is_flashinfer_supported", lambda unused: True) + monkeypatch.setattr("tensorrt_llm._torch.modules.low_m_gemm._current_sm", lambda unused: 103) + + input_tensor = torch.empty((8, 128), dtype=torch.bfloat16) + weight = torch.empty((8192, 128), dtype=torch.bfloat16) + with torch.inference_mode(): + output = dispatcher.apply(torch.nn.Linear(1, 1), input_tensor, weight, None) + + assert dispatcher._flashinfer_mm.call_count == expected_flashinfer_calls + assert (output is not None) is bool(expected_flashinfer_calls) + + +def test_shape_collector_persists_new_runtime_shape_after_warmup(tmp_path) -> None: + output = tmp_path / "shapes.json" + collector = _ShapeCollector(str(output)) + module = torch.nn.Linear(8, 4, bias=False) + module._low_m_gemm_name = "model.layers.0.self_attn.qkv_proj" + weight = torch.empty((4, 8)) + + try: + collector.record(module, torch.empty((1, 8)), weight, None, cuda_graph=True) + assert not output.exists() + + collector.flush(mark_warmup_complete=True) + warmup_document = json.loads(output.read_text(encoding="utf-8")) + assert [(row["m"], row["call_count"]) for row in warmup_document["shapes"]] == [(1, 1)] + + collector.record(module, torch.empty((1, 8)), weight, None, cuda_graph=True) + collector.record(module, torch.empty((2, 8)), weight, None, cuda_graph=True) + runtime_document = json.loads(output.read_text(encoding="utf-8")) + assert {(row["m"], row["call_count"]) for row in runtime_document["shapes"]} == { + (1, 2), + (2, 1), + } + + another_module = torch.nn.Linear(8, 4, bias=False) + another_module._low_m_gemm_name = "model.layers.1.self_attn.qkv_proj" + collector.record(another_module, torch.empty((2, 8)), weight, None, cuda_graph=True) + unchanged_document = json.loads(output.read_text(encoding="utf-8")) + assert unchanged_document == runtime_document + + collector.flush() + final_document = json.loads(output.read_text(encoding="utf-8")) + assert { + (row["m"], row["layer"], row["call_count"]) for row in final_document["shapes"] + } == { + (1, "model.layers.0.self_attn.qkv_proj", 2), + (2, "model.layers.0.self_attn.qkv_proj", 1), + (2, "model.layers.1.self_attn.qkv_proj", 1), + } + finally: + atexit.unregister(collector.flush) diff --git a/tests/unittest/_torch/modules/test_moe_host_sharer.py b/tests/unittest/_torch/modules/test_moe_host_sharer.py index 0ed0ee609bb8..901a77b4f350 100644 --- a/tests/unittest/_torch/modules/test_moe_host_sharer.py +++ b/tests/unittest/_torch/modules/test_moe_host_sharer.py @@ -1,11 +1,29 @@ +# 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 os import unittest +from tempfile import TemporaryDirectory +from unittest.mock import patch import numpy as np import torch from mpi4py import MPI -from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import \ - HostMoeTensorSharer +from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import ( + HostMoeTensorSharer, _FileBackedSharedMemory, _tensor_to_weight) class TestHostMoeTensorSharer(unittest.TestCase): @@ -70,6 +88,80 @@ def generate_tensor_data(self, expert_id, tensor_shape): tensor_data[i, j] = expert_id * 1000 + i * 100 + j return tensor_data + def test_column_major_transfer_view(self): + """Column-major TMA tensors use a transposed EPLB copy descriptor.""" + row_major = torch.arange(24, dtype=torch.int32).reshape(4, 6) + column_major = row_major.transpose(0, 1) + + self.assertFalse(column_major.is_contiguous()) + transfer_view = HostMoeTensorSharer.get_transfer_view(column_major) + self.assertTrue(transfer_view.is_contiguous()) + torch.testing.assert_close(transfer_view, row_major) + + moe_weight = _tensor_to_weight(column_major) + self.assertEqual(moe_weight.height, row_major.shape[0]) + self.assertEqual(moe_weight.width, + row_major.shape[1] * row_major.element_size()) + self.assertEqual(moe_weight.pitch, + row_major.stride(0) * row_major.element_size()) + + def test_file_backed_storage(self): + """File-backed EPLB staging preserves tensors outside /dev/shm.""" + comm = MPI.COMM_SELF + tensor_shape = (16, 32) + tensor_data = self.generate_tensor_data(0, tensor_shape) + flush_calls = [] + original_flush = _FileBackedSharedMemory.flush + + def tracked_flush(storage): + flush_calls.append(storage) + return original_flush(storage) + + with TemporaryDirectory() as directory, patch.dict( + "os.environ", { + "TRTLLM_EPLB_SHM_DIR": directory, + "TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL": "True", + }), patch.object(_FileBackedSharedMemory, "flush", + tracked_flush): + sharer = HostMoeTensorSharer(0, 1, comm) + sharer.set_shared_memory_base_name("test_file_backed_sharer") + sharer.share_host_tensor_with_shape(0, "weight", tensor_data) + sharer.finalize_layer_weights() + self.assertEqual(len(flush_calls), 1) + retrieved = {} + sharer.finalize_host_tensor_sharing( + lambda expert_id, name, tensor: retrieved.setdefault( + (expert_id, name), tensor.clone())) + + torch.testing.assert_close(retrieved[(0, "weight")], tensor_data) + backing_path = os.path.join(directory, + sharer.get_shared_memory_name()) + self.assertTrue(os.path.exists(backing_path)) + + sharer.pre_shutdown_cleanup() + sharer.post_shutdown_cleanup() + self.assertFalse(os.path.exists(backing_path)) + + def test_file_backed_mappings_are_coherent_after_flush(self): + """Completed file-backed layers are flushed before further loading.""" + with TemporaryDirectory() as directory: + creator = _FileBackedSharedMemory("test_eplb_map", + directory, + create=True, + size=4096) + peer = _FileBackedSharedMemory("test_eplb_map", directory) + try: + creator.buf[:8] = b"EPLBTEST" + creator.flush() + self.assertEqual(peer.buf[:8], b"EPLBTEST") + peer.buf[:8] = b"PEERVIEW" + peer.flush() + self.assertEqual(creator.buf[:8], b"PEERVIEW") + finally: + peer.close() + creator.close() + creator.unlink() + def test_host_tensor_sharing_basic(self): """Basic test for host tensor sharing""" # Get MPI communication information diff --git a/tests/unittest/_torch/modules/test_moe_load_balancer.py b/tests/unittest/_torch/modules/test_moe_load_balancer.py index 816c26a0d88d..9fc35703f4c1 100644 --- a/tests/unittest/_torch/modules/test_moe_load_balancer.py +++ b/tests/unittest/_torch/modules/test_moe_load_balancer.py @@ -1,4 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json import os +import tempfile import unittest from unittest.mock import MagicMock, patch @@ -311,6 +316,184 @@ def test_moe_load_balancer_lifecycle_methods(self, mock_load_balancer_impl): balancer.shutdown() mock_load_balancer_impl.return_value.shutdown.assert_called_once() + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_online_iteration_cadence(self, mock_load_balancer_impl): + """Test that cadence samples, drains pending copies, then skips.""" + + torch.cuda.set_device(0) + balancer = MoeLoadBalancer(0, 4, 2, iteration_interval=4) + balancer.set_iter_info(True, True) + + self.assertTrue(balancer.requires_eager_forward()) + with MoeLoadBalancerIterContext(balancer): + self.assertTrue(balancer.runtime_state.active) + + self.assertTrue(balancer.requires_eager_forward()) + with MoeLoadBalancerIterContext(balancer): + self.assertTrue(balancer.runtime_state.active) + + self.assertFalse(balancer.requires_eager_forward()) + with MoeLoadBalancerIterContext(balancer): + self.assertFalse(balancer.runtime_state.active) + with MoeLoadBalancerIterContext(balancer): + self.assertFalse(balancer.runtime_state.active) + + self.assertTrue(balancer.requires_eager_forward()) + with MoeLoadBalancerIterContext(balancer): + self.assertTrue(balancer.runtime_state.active) + + self.assertEqual( + mock_load_balancer_impl.return_value.start_iter.call_args_list, [ + unittest.mock.call(0, True, True), + unittest.mock.call(1, False, False), + unittest.mock.call(2, True, True), + ]) + self.assertEqual( + mock_load_balancer_impl.return_value.end_iter.call_args_list, [ + unittest.mock.call(0), + unittest.mock.call(1), + unittest.mock.call(2), + ]) + + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_invalid_online_iteration_interval(self, mock_load_balancer_impl): + """Test that a non-positive cadence interval is rejected.""" + + torch.cuda.set_device(0) + with self.assertRaisesRegex(ValueError, "must be positive"): + MoeLoadBalancer(0, 4, 2, iteration_interval=0) + mock_load_balancer_impl.assert_not_called() + + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_batched_online_statistics_and_updates(self, + mock_load_balancer_impl): + """Test that a cycle observes before it migrates any layer.""" + + torch.cuda.set_device(0) + environment = { + 'TRTLLM_EPLB_STATISTICS_PER_CYCLE': '3', + 'TRTLLM_EPLB_CYCLE_START_ITER': '2', + 'TRTLLM_EPLB_STATISTIC_INTERVAL': '2', + 'TRTLLM_EPLB_UPDATE_INTERVAL': '3', + 'TRTLLM_EPLB_CYCLE_INTERVAL': '4', + } + with patch.dict(os.environ, environment): + balancer = MoeLoadBalancer(0, 4, 2) + balancer.single_layer_load_balancers = [MagicMock() for _ in range(5)] + balancer.set_iter_info(True, True) + + expected_modes = [ + 'skip', 'skip', 'sample', 'skip', 'sample', 'skip', 'sample', + 'update', 'drain', 'skip', 'update', 'drain', 'skip', 'update', + 'drain', 'skip', 'skip', 'skip', 'skip', 'sample' + ] + observed_modes = [] + for _ in expected_modes: + observed_modes.append(balancer.get_next_iter_mode()) + with MoeLoadBalancerIterContext(balancer): + pass + + self.assertEqual(observed_modes, expected_modes) + self.assertEqual( + mock_load_balancer_impl.return_value.start_iter.call_args_list, [ + unittest.mock.call(0, True, False), + unittest.mock.call(1, True, False), + unittest.mock.call(2, True, False), + unittest.mock.call(3, False, True), + unittest.mock.call(4, False, False), + unittest.mock.call(5, False, True), + unittest.mock.call(6, False, False), + unittest.mock.call(7, False, True), + unittest.mock.call(8, False, False), + unittest.mock.call(9, True, False), + ]) + + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_batched_online_update_interval_requires_drain( + self, mock_load_balancer_impl): + """Test that adjacent migrations cannot omit their drain forward.""" + + torch.cuda.set_device(0) + environment = { + 'TRTLLM_EPLB_STATISTICS_PER_CYCLE': '1', + 'TRTLLM_EPLB_UPDATE_INTERVAL': '1', + } + with patch.dict(os.environ, environment): + with self.assertRaisesRegex(ValueError, "at least two"): + MoeLoadBalancer(0, 4, 2) + mock_load_balancer_impl.assert_not_called() + + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_batched_online_cycle_consumes_cpp_empty_update_group( + self, mock_load_balancer_impl): + """Test cadence stays aligned with a divisible C++ update plan.""" + + torch.cuda.set_device(0) + environment = { + 'TRTLLM_EPLB_STATISTICS_PER_CYCLE': '1', + 'TRTLLM_EPLB_UPDATE_INTERVAL': '2', + } + with patch.dict(os.environ, environment): + balancer = MoeLoadBalancer(0, 4, 4) + # Six layers split into two round-robin groups of three. The configured + # maximum is four, so this also proves the empty-group condition follows + # the actual C++ plan rather than layer_count % layer_updates_per_iter. + balancer.single_layer_load_balancers = [MagicMock() for _ in range(6)] + balancer.set_iter_info(True, True) + + observed_modes = [] + for _ in range(8): + observed_modes.append(balancer.get_next_iter_mode()) + with MoeLoadBalancerIterContext(balancer): + pass + + self.assertEqual(observed_modes, [ + 'sample', 'update', 'drain', 'update', 'drain', 'update', 'drain', + 'sample' + ]) + self.assertEqual( + mock_load_balancer_impl.return_value.start_iter.call_args_list, [ + unittest.mock.call(0, True, False), + unittest.mock.call(1, False, True), + unittest.mock.call(2, False, False), + unittest.mock.call(3, False, True), + unittest.mock.call(4, False, False), + unittest.mock.call(5, False, True), + unittest.mock.call(6, False, False), + unittest.mock.call(7, True, False), + ]) + + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') + def test_inactive_online_iteration_routes_without_statistics( + self, mock_load_balancer_impl): + """Test that skipped iterations retain routing without EPLB sync.""" + + torch.cuda.set_device(0) + balancer = MoeLoadBalancer(0, 4, 2, iteration_interval=4) + mock_single_layer_impl = MagicMock() + selected_experts = torch.tensor([[0, 1]], dtype=torch.int32) + routed_slots = torch.tensor([[2, 3]], dtype=torch.int32) + layer = SingleLayerMoeLoadBalancer(mock_single_layer_impl, + MPI.COMM_WORLD, + expert_count=4, + runtime_state=balancer.runtime_state) + + with patch('torch.ops.trtllm.moe_load_balance_wait_gpu_stage') as mock_wait, \ + patch('torch.ops.trtllm.moe_load_balance_set_cpu_stage') as mock_set_cpu, \ + patch('torch.ops.trtllm.moe_load_balance_routing', return_value=routed_slots) as mock_route: + layer.start_wait_gpu_stage() + layer.done_wait_gpu_stage() + layer.update_local_statistic(selected_experts, True, True) + self.assertIsNone(layer.get_local_statistic_tensor()) + result = layer.route(selected_experts) + layer.start_set_cpu_stage() + layer.done_set_cpu_stage() + + mock_wait.assert_not_called() + mock_set_cpu.assert_not_called() + mock_route.assert_called_once() + self.assertTrue(torch.equal(result, routed_slots)) + @patch('tensorrt_llm.bindings.internal.runtime.MoeLoadBalancer') def test_reconfigure_mask_only_rejects_active_iteration( self, mock_load_balancer_impl): @@ -327,6 +510,58 @@ def test_reconfigure_mask_only_rejects_active_iteration( mock_load_balancer_impl.return_value.reconfigure_mask_only.assert_not_called( ) + def test_eplb_telemetry_records_lifecycle_and_placement(self): + """Test optional telemetry records phases and migrated bytes.""" + + previous_placement = [{ + 'layer_id': 7, + 'rank_expert_ids': [[0, 1], [2, 3]], + }] + current_placement = [[0, 4], [2, 5]] + layer = MagicMock() + layer.get_layer_idx.return_value = 7 + layer.get_old_rank_expert_ids.return_value = current_placement + layer.host_tensor_sharer.name_info = { + 'weight': (torch.float16, (2, 3)), + } + + with tempfile.TemporaryDirectory() as temp_dir: + balancer = object.__new__(MoeLoadBalancer) + balancer.is_shutdown = True + balancer.ep_rank = 0 + balancer.telemetry_path = os.path.join(temp_dir, 'lifecycle.jsonl') + balancer.forward_iter_id = 10 + balancer.iter_id = 3 + balancer.single_layer_load_balancers = [layer] + balancer._telemetry_previous_mode = None + balancer._telemetry_previous_placement = previous_placement + balancer._telemetry_counts = { + 'sample': 0, + 'update': 0, + 'drain': 0, + 'skip': 0, + } + + for mode in ('sample', 'update', 'drain', 'skip'): + balancer.record_iter_mode(mode) + + with open(balancer.telemetry_path, + encoding='utf-8') as telemetry_file: + records = [json.loads(line) for line in telemetry_file] + + self.assertEqual([record['event'] for record in records], + ['sample', 'update', 'drain_complete', 'stable']) + self.assertEqual(records[2]['changed_slots'], 2) + self.assertEqual(records[2]['changed_bytes'], 24) + self.assertEqual(records[2]['placement'][0]['rank_expert_ids'], + current_placement) + self.assertEqual(records[3]['counts'], { + 'sample': 1, + 'update': 1, + 'drain': 1, + 'skip': 1, + }) + def test_real_statistic_kernel(self): """Test the real statistic kernel functionality.""" diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 6e3daea5e695..9d4658352608 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,6 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import json import struct import types +from dataclasses import replace import pytest import torch @@ -16,6 +20,30 @@ pytestmark = pytest.mark.cpu_only +@pytest.mark.parametrize( + ("configured_value", "expected_value", "is_default"), + [ + pytest.param(None, 32768, True, id="derived-default"), + pytest.param(32768, 32768, False, id="explicit-32768"), + pytest.param(65536, 65536, False, id="explicit-65536"), + ], +) +def test_moe_max_num_tokens_tracks_explicit_configuration( + configured_value, expected_value, is_default +): + model_config = ModelConfig(max_num_tokens=32768, moe_max_num_tokens=configured_value) + + assert model_config.moe_max_num_tokens == expected_value + assert model_config._moe_max_num_tokens_is_default is is_default + assert replace(model_config)._moe_max_num_tokens_is_default is is_default + + +@pytest.mark.parametrize("value", [0, -1]) +def test_moe_max_num_tokens_rejects_nonpositive_values(value): + with pytest.raises(ValueError, match="must be a positive integer"): + ModelConfig(moe_max_num_tokens=value) + + def make_pretrained_config( *, num_attention_heads: int = 16,