From ae7fa65d689daf54d86b88f9f0b5c137fd4cbdd3 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 22:12:21 -0700
Subject: [PATCH 01/18] [None][fix] align MTP ADP token metadata

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../_torch/models/modeling_qwen3_next.py      | 82 +++++++++++--------
 .../models/test_qwen3_next_eager_fusion.py    | 74 +++++++++++++++++
 2 files changed, 122 insertions(+), 34 deletions(-)

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/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()

From f5bdfe43b4b8541d1f0bb49c01283b648217322e Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 22:13:51 -0700
Subject: [PATCH 02/18] [None][fix] keep MoE chunk state graph-safe

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../_torch/modules/fused_moe/moe_scheduler.py |  5 +-
 .../_torch/modules/moe/test_moe_scheduler.py  | 71 +++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 tests/unittest/_torch/modules/moe/test_moe_scheduler.py

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/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)

From adebe1f03bc8df9d5320b3e8a905d600de26b98f Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 22:19:20 -0700
Subject: [PATCH 03/18] [None][fix] size FP8 MoE activation backing
 independently

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp    | 18 ++++++++---
 .../kernels/blockScaleMoeActivationTest.cu    | 32 +++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)

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<at::Tensor> 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<at::Tensor> 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/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu
index 4a56dcd8f558..94d08caf875f 100644
--- a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu
+++ b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu
@@ -45,6 +45,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"
@@ -435,4 +436,35 @@ TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNa
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+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<int64_t>(activationCapacity) * intermediateSize * fp8Bits / 8;
+    auto const activationBytesWithGemm1Capacity = static_cast<int64_t>(gemm1Capacity) * intermediateSize * fp8Bits / 8;
+    auto const scaleBytes = static_cast<int64_t>(activationCapacity) * (intermediateSize / kEltsPerSf) * sizeof(float);
+    auto const scaleBytesWithGemm1Capacity
+        = static_cast<int64_t>(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

From f229c59b70e09b7f06a0dedbb669b9f0c2f5cdc6 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 22:48:58 -0700
Subject: [PATCH 04/18] [None][fix] respect explicit DeepGEMM MoE token
 capacity

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 tensorrt_llm/_torch/model_config.py           | 16 +++++++--
 .../modules/fused_moe/fused_moe_deepgemm.py   | 36 ++++++++++++++++---
 .../_torch/modules/moe/test_moe_backend.py    | 30 ++++++++++++++++
 tests/unittest/_torch/test_model_config.py    | 28 +++++++++++++++
 4 files changed, 103 insertions(+), 7 deletions(-)

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/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py
index 7d3951762326..249235670afd 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(
@@ -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,
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/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,

From 56f2a93d9e85bfd10445d443926f237423a182a4 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 6 Aug 2026 22:24:48 -0700
Subject: [PATCH 05/18] [None][fix] keep zero FP8 MoE activation blocks finite

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../blockScaleMoe/DevKernel.cu                | 17 ++++++++----
 .../kernels/blockScaleMoeActivationTest.cu    | 27 ++++++++++---------
 2 files changed, 26 insertions(+), 18 deletions(-)

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<Type>(out[i] / scaleOut);
         }
         *reinterpret_cast<PackedIo*>(params.outPtr + static_cast<int64_t>(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/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu
index 94d08caf875f..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
@@ -53,6 +52,7 @@
 #include <cutlass/numeric_types.h>
 
 #include <algorithm>
+#include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <memory>
@@ -213,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<std::pair<int32_t, int32_t>> zeroedRowBlock = std::nullopt)
     {
@@ -402,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,
@@ -423,14 +422,16 @@ TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNa
     auto const permuted = runOnce(kTileForcePermuted);
 
     auto const sfIdx = static_cast<int64_t>(zeroedRow) + static_cast<int64_t>(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<int64_t>(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<uint8_t>(legacy.bytes[idx]), static_cast<uint8_t>(permuted.bytes[idx]))
-            << "0/0 encoding differs at element " << elt;
+            << "zero-block encoding differs at element " << elt;
     }
 }
 

From 3906e73265e628ee9260f20fe025d3a0b9ae743b Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 7 Aug 2026 02:32:34 -0700
Subject: [PATCH 06/18] [None][perf] avoid unused DeepGEMM permutation
 allocation

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 cpp/tensorrt_llm/thop/moeUtilOp.cpp           | 10 +-
 .../modules/fused_moe/fused_moe_deepgemm.py   | 23 +++--
 .../modules/fused_moe/ops/moe_op_deepgemm.py  | 20 ++--
 .../test_deepgemm_fused_expand_quant.py       | 94 ++++++++++++++++++-
 4 files changed, 125 insertions(+), 22 deletions(-)

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<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Te
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
     auto permuted_token_selected_experts_tensor
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
-    auto permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false));
+    auto permuted_data_tensor = torch::empty({0, hidden_size}, input.options().requires_grad(false));
     auto permuted_token_final_scales_tensor
-        = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+        = torch::empty({0}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+    if (!skip_data_expand)
+    {
+        permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false));
+        permuted_token_final_scales_tensor
+            = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+    }
     auto expert_first_token_offset_tensor = torch::empty(
         {num_experts_per_node + 1}, torch::dtype(torch::kInt64).device(torch::kCUDA).requires_grad(false));
     auto unpermuted_row_to_permuted_row_tensor = torch::empty({static_cast<int64_t>(experts_per_token * num_rows)},
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 249235670afd..53beac5f625b 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py
@@ -644,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
@@ -1008,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/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/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,

From 203211f40eecc27ecbb8cab5bb729d0c8e4d9e25 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Sat, 8 Aug 2026 11:02:34 -0700
Subject: [PATCH 07/18] [None][perf] select cooperative routing for large
 expert tiers

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../blockScaleMoe/routing/RoutingCustom.cu    | 10 ++++----
 .../routing/RoutingCustomPolicy.cuh           |  3 +++
 .../kernels/routing/routingCustomTest.cpp     | 24 +++++++++++++++++++
 3 files changed, 32 insertions(+), 5 deletions(-)

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/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()

From 06a13e2fb0641baff877370669f4e0e6eee71041 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 21:22:15 -0700
Subject: [PATCH 08/18] [None][perf] adapt attention DP balance at low
 occupancy

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 21 +++++++-
 .../_torch/executor/test_py_executor.py       | 53 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

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])

From d8d10ab3540c4341852075a45b6b35bcfa0a23cf Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Sat, 8 Aug 2026 11:02:40 -0700
Subject: [PATCH 09/18] [None][perf] tune cached replay for wide value heads

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../_torch/modules/fla/cached_replay.py       | 40 +++++++-
 .../_torch/modules/mamba/mamba2_metadata.py   | 70 +++++++++++++-
 .../mamba/test_gdn_replay_recurrent.py        | 31 ++++++-
 .../modules/mamba/test_mamba2_metadata.py     | 92 +++++++++++++++++++
 4 files changed, 224 insertions(+), 9 deletions(-)

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/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/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")

From e6b07b8050908a0afa7451664155885ac73c8e51 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 5 Aug 2026 23:29:36 -0700
Subject: [PATCH 10/18] [None][perf] add autotuned low-M BF16 GEMM dispatch

Add an opt-in, offline-cache dispatcher for Blackwell BF16 GEMMs with M <= 32. Route exact graph/eager shapes to FlashInfer direct or split-K tactics, or preserve the existing TensorRT-LLM path, while validating cache and runtime provenance.

Initialize the dispatcher with the model, retain an explicit rollback path, and support debug-only hot-shape collection with focused unit coverage.

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/linear.py         |  30 +
 tensorrt_llm/_torch/modules/low_m_gemm.py     | 785 ++++++++++++++++++
 .../_torch/pyexecutor/model_engine.py         |   7 +
 .../_torch/modules/test_low_m_gemm.py         | 226 +++++
 4 files changed, 1048 insertions(+)
 create mode 100644 tensorrt_llm/_torch/modules/low_m_gemm.py
 create mode 100644 tests/unittest/_torch/modules/test_low_m_gemm.py

diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 7378b332fab3..d29073688c61 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,13 @@ 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 loads decisions produced by an offline
+        # FlashInfer direct/split-K versus cuBLAS sweep. It never benchmarks on
+        # the serving path 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..55e0417d926d
--- /dev/null
+++ b/tensorrt_llm/_torch/modules/low_m_gemm.py
@@ -0,0 +1,785 @@
+# 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 hashlib
+import json
+import os
+import threading
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+import torch
+
+from tensorrt_llm.logger import logger
+from tensorrt_llm.version import __version__ as trtllm_version
+
+from ..flashinfer_utils import get_env_enable_pdl
+
+_BACKEND_ENV = "TRTLLM_LOW_M_GEMM_BACKEND"
+_DISPATCH_CACHE_ENV = "TRTLLM_LOW_M_GEMM_TUNING_CACHE"
+_FLASHINFER_CACHE_ENV = "TRTLLM_FLASHINFER_AUTOTUNER_CACHE"
+_FLASHINFER_COMMIT_ENV = "TRTLLM_FLASHINFER_COMMIT"
+_FLASHINFER_SOURCE_ROOT_ENV = "TRTLLM_FLASHINFER_SOURCE_ROOT"
+_DISABLE_CACHE_ENV = "TRTLLM_DISABLE_GEMM_TUNING_CACHE"
+_SHAPE_LOG_ENV = "TRTLLM_LOW_M_GEMM_SHAPE_LOG"
+_DISPATCH_CACHE_SCHEMA_VERSION = 3
+_SHAPE_LOG_SCHEMA_VERSION = 1
+_FLASHINFER_BACKEND = "cute-dsl"
+_SUPPORTED_SMS = {100, 103}
+_MAX_FLASHINFER_M = 32
+
+
+class LowMGemmBackend(str, enum.Enum):
+    """Runtime choices exposed by the low-M BF16 dispatcher."""
+
+    OFF = "off"
+    AUTO = "auto"
+    FLASHINFER = "flashinfer"
+    CUBLAS = "cublas"
+
+
+@dataclass(frozen=True)
+class GemmDispatchKey:
+    """Properties that can change the best low-M GEMM implementation."""
+
+    sm: int
+    m: int
+    n: int
+    k: int
+    a_type: str = "bf16"
+    b_type: str = "bf16"
+    c_type: str = "bf16"
+    trans_a: bool = False
+    trans_b: bool = True
+    has_bias: bool = False
+    cuda_graph: bool = True
+
+    def cache_key(self) -> str:
+        transpose = ("t" if self.trans_a else "n") + ("t" if self.trans_b else "n")
+        bias = "bias" if self.has_bias else "nobias"
+        execution = "graph" if self.cuda_graph else "eager"
+        return (
+            f"sm{self.sm}:{self.a_type}:{self.m}x{self.n}x{self.k}:{transpose}:{bias}:{execution}"
+        )
+
+
+@dataclass(frozen=True)
+class GemmTuningResult:
+    """One persisted dispatcher decision produced by offline tuning."""
+
+    backend: str
+    algorithm: Optional[str] = None
+    tactic: Optional[dict[str, Any]] = None
+    latency_us: Optional[float] = None
+    baseline_us: Optional[float] = None
+    measurements: Optional[dict[str, Any]] = None
+
+
+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 _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as cache_file:
+        for chunk in iter(lambda: cache_file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+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"))
+
+
+def _flashinfer_commit(flashinfer_module: Any) -> Optional[str]:
+    """Return a verifiable FlashInfer source revision for cache provenance."""
+
+    package_commit = getattr(flashinfer_module, "__git_version__", None)
+    if package_commit == "unknown":
+        package_commit = None
+    declared_commit = os.environ.get(_FLASHINFER_COMMIT_ENV)
+    if package_commit is not None and declared_commit is not None:
+        if package_commit != declared_commit:
+            raise RuntimeError(
+                f"{_FLASHINFER_COMMIT_ENV}={declared_commit} does not match "
+                f"the imported FlashInfer revision {package_commit}."
+            )
+    return package_commit or declared_commit
+
+
+def _current_sm(device: torch.device) -> int:
+    major, minor = torch.cuda.get_device_capability(device)
+    return major * 10 + minor
+
+
+def _configure_flashinfer_source_layout() -> None:
+    source_value = os.environ.get(_FLASHINFER_SOURCE_ROOT_ENV)
+    if not source_value:
+        return
+    source_root = Path(source_value).resolve()
+    source_csrc = source_root / "csrc"
+    source_include = source_root / "include"
+    if not source_csrc.is_dir() or not source_include.is_dir():
+        raise RuntimeError(
+            f"{_FLASHINFER_SOURCE_ROOT_ENV} is not a FlashInfer source checkout: {source_root}"
+        )
+
+    import flashinfer
+    from flashinfer.jit import env as jit_env
+
+    imported_root = Path(flashinfer.__file__).resolve().parents[1]
+    if imported_root != source_root:
+        raise RuntimeError(
+            f"Imported FlashInfer from {imported_root}, but "
+            f"{_FLASHINFER_SOURCE_ROOT_ENV}={source_root}."
+        )
+    # An installed wheel stores build inputs under flashinfer/data. A source
+    # overlay stores them at repository top level. Configure this before any
+    # FlashInfer JIT helper is first invoked.
+    jit_env.FLASHINFER_CSRC_DIR = source_csrc
+    jit_env.FLASHINFER_INCLUDE_DIR = source_include
+
+
+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:
+    """Offline-cache dispatcher for FlashInfer direct/split-K versus cuBLAS."""
+
+    def __init__(self) -> None:
+        self.backend = _configured_backend()
+        if self.backend in (LowMGemmBackend.AUTO, LowMGemmBackend.FLASHINFER):
+            _configure_flashinfer_source_layout()
+        self.cuda_graph_enabled = True
+        self._dispatch_cache: dict[str, GemmTuningResult] = {}
+        self._dispatch_metadata: dict[str, Any] = {}
+        self._prepared = False
+        self._flashinfer_mm = None
+        self._flashinfer_direct_tactic = None
+        self._flashinfer_run_direct = None
+        self._flashinfer_splitk_tactic = None
+        self._flashinfer_run_splitk = None
+        self._engaged: set[tuple[str, int, int, int, bool]] = set()
+        self._cache_misses: set[str] = 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._load_dispatch_cache()
+        needs_flashinfer = self.backend == LowMGemmBackend.FLASHINFER or any(
+            result.backend == LowMGemmBackend.FLASHINFER.value
+            for result in self._dispatch_cache.values()
+        )
+        if needs_flashinfer:
+            self._prepare_flashinfer()
+        self._prepared = True
+
+    def _load_dispatch_cache(self) -> None:
+        if self.backend != LowMGemmBackend.AUTO:
+            return
+        if os.environ.get(_DISABLE_CACHE_ENV, "0") == "1":
+            logger.warning(
+                f"{_DISABLE_CACHE_ENV}=1: low-M GEMM auto mode will use its cache-miss heuristic."
+            )
+            return
+        cache_value = os.environ.get(_DISPATCH_CACHE_ENV)
+        if not cache_value:
+            raise RuntimeError(
+                f"{_BACKEND_ENV}=auto requires {_DISPATCH_CACHE_ENV}; run the "
+                "offline tuner first, use flashinfer for heuristic-only testing, "
+                "or set cublas to roll back."
+            )
+        cache_path = Path(cache_value)
+        if not cache_path.is_file():
+            raise FileNotFoundError(f"Low-M GEMM dispatch cache does not exist: {cache_path}")
+        document = json.loads(cache_path.read_text(encoding="utf-8"))
+        if document.get("schema_version") != _DISPATCH_CACHE_SCHEMA_VERSION:
+            raise RuntimeError(
+                f"Unsupported low-M GEMM cache schema in {cache_path}: "
+                f"{document.get('schema_version')!r}"
+            )
+        self._dispatch_metadata = document.get("metadata", {})
+        if not isinstance(self._dispatch_metadata, dict):
+            raise RuntimeError(f"Low-M GEMM cache {cache_path} has invalid metadata.")
+        expected_pdl = self._dispatch_metadata.get("pdl")
+        if not isinstance(expected_pdl, bool):
+            raise RuntimeError(
+                f"Low-M GEMM cache {cache_path} has no Boolean PDL runtime identity."
+            )
+        actual_pdl = bool(get_env_enable_pdl())
+        if expected_pdl != actual_pdl:
+            raise RuntimeError(
+                f"Low-M GEMM cache was tuned with PDL={expected_pdl}, but runtime PDL={actual_pdl}."
+            )
+        expected_cuda_version = self._dispatch_metadata.get("cuda_version")
+        if not isinstance(expected_cuda_version, str):
+            raise RuntimeError(f"Low-M GEMM cache {cache_path} has no CUDA runtime identity.")
+        actual_cuda_version = torch.version.cuda or "none"
+        if expected_cuda_version != actual_cuda_version:
+            raise RuntimeError(
+                "CUDA version does not match the low-M GEMM dispatch cache: "
+                f"expected {expected_cuda_version}, got {actual_cuda_version}."
+            )
+        entries = document.get("entries")
+        if not isinstance(entries, dict):
+            raise RuntimeError(f"Low-M GEMM cache {cache_path} has no object-valued entries.")
+        self._dispatch_cache = {key: GemmTuningResult(**value) for key, value in entries.items()}
+        logger.info(f"Loaded {len(self._dispatch_cache)} low-M GEMM decisions from {cache_path}.")
+
+    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
+            from flashinfer.gemm.kernels.dense_bf16_gemm_direct import (
+                DirectTactic,
+                run_direct_dense,
+            )
+            from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import (
+                SplitKTactic,
+                run_splitk_dense,
+            )
+        except (ImportError, ModuleNotFoundError) as error:
+            raise RuntimeError(
+                "FlashInfer low-M GEMM requires a FlashInfer build containing "
+                "PR #4266 and nvidia-cutlass-dsl."
+            ) from error
+        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 pinned PR #4266 build is not active."
+            )
+        expected_arch = self._dispatch_metadata.get("gpu_arch")
+        actual_arch = f"sm{_current_sm(torch.device('cuda'))}"
+        if expected_arch is not None and expected_arch != actual_arch:
+            raise RuntimeError(
+                f"Low-M GEMM cache targets {expected_arch}, but this rank uses {actual_arch}."
+            )
+        expected_version = self._dispatch_metadata.get("flashinfer_version")
+        if expected_version is not None and expected_version != flashinfer_module.__version__:
+            raise RuntimeError(
+                "FlashInfer version does not match the low-M GEMM dispatch "
+                f"cache: expected {expected_version}, got "
+                f"{flashinfer_module.__version__}."
+            )
+        expected_trtllm_version = self._dispatch_metadata.get("trtllm_version")
+        if expected_trtllm_version is not None and expected_trtllm_version != trtllm_version:
+            raise RuntimeError(
+                "TensorRT-LLM version does not match the low-M GEMM dispatch "
+                f"cache: expected {expected_trtllm_version}, got "
+                f"{trtllm_version}."
+            )
+        expected_dispatcher_digest = self._dispatch_metadata.get("trtllm_low_m_gemm_sha256")
+        if expected_dispatcher_digest is not None:
+            actual_dispatcher_digest = _sha256(Path(__file__))
+            if actual_dispatcher_digest != expected_dispatcher_digest:
+                raise RuntimeError(
+                    "TensorRT-LLM low-M GEMM source does not match the offline "
+                    "dispatch cache; regenerate the cache with this checkout."
+                )
+        expected_commit = self._dispatch_metadata.get("flashinfer_commit")
+        actual_commit = _flashinfer_commit(flashinfer_module)
+        if (
+            expected_commit is not None
+            and self.backend == LowMGemmBackend.AUTO
+            and actual_commit != expected_commit
+        ):
+            raise RuntimeError(
+                "Imported FlashInfer does not match the cache's pinned "
+                f"commit {expected_commit}; got {actual_commit!r}."
+            )
+        self._flashinfer_mm = mm_bf16
+        self._flashinfer_direct_tactic = DirectTactic
+        self._flashinfer_run_direct = run_direct_dense
+        self._flashinfer_splitk_tactic = SplitKTactic
+        self._flashinfer_run_splitk = run_splitk_dense
+
+        cache_disabled = os.environ.get(_DISABLE_CACHE_ENV, "0") == "1"
+        cache_value = os.environ.get(_FLASHINFER_CACHE_ENV)
+        if cache_disabled or not cache_value:
+            if self.backend == LowMGemmBackend.AUTO and not cache_disabled:
+                raise RuntimeError(
+                    f"{_BACKEND_ENV}=auto selected FlashInfer entries but "
+                    f"{_FLASHINFER_CACHE_ENV} is unset."
+                )
+            logger.warning(
+                "FlashInfer BF16 autotune cache is unavailable for provenance "
+                "validation; forced/cache-miss routing will use FlashInfer's "
+                "heuristic."
+            )
+            return
+
+        cache_path = Path(cache_value)
+        if not cache_path.is_file():
+            raise FileNotFoundError(f"FlashInfer autotune cache does not exist: {cache_path}")
+        expected_digest = self._dispatch_metadata.get("flashinfer_cache_sha256")
+        if expected_digest is not None:
+            actual_digest = _sha256(cache_path)
+            if actual_digest != expected_digest:
+                raise RuntimeError(
+                    "FlashInfer autotune cache checksum does not match the "
+                    "TRT-LLM dispatch cache; regenerate the pair together."
+                )
+        # Do not load this file into FlashInfer's runtime AutoTuner. Its
+        # default mapper aliases M=24 to M=32, whereas this dispatch cache is
+        # exact-M. The selected runner/tactic is launched directly below.
+        logger.info(f"Validated paired FlashInfer BF16 autotune cache {cache_path}.")
+
+    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))
+
+    @staticmethod
+    def _tactic_values(
+        result: GemmTuningResult,
+        expected_fields: tuple[str, ...],
+    ) -> dict[str, int]:
+        tactic = result.tactic
+        if not isinstance(tactic, dict) or set(tactic) != set(expected_fields):
+            raise RuntimeError(
+                f"Invalid {result.algorithm!r} tactic in low-M GEMM cache: "
+                f"expected {expected_fields}, got {tactic!r}."
+            )
+        return {field: int(tactic[field]) for field in expected_fields}
+
+    def _launch_cached_flashinfer(
+        self,
+        input_2d: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        result: GemmTuningResult,
+    ) -> torch.Tensor:
+        output = torch.empty(
+            (input_2d.shape[0], weight.shape[0]),
+            device=input_2d.device,
+            dtype=torch.bfloat16,
+        )
+        weight_t = weight.t()
+        pdl = get_env_enable_pdl()
+        if result.algorithm == "simt":
+            if bias is not None:
+                raise RuntimeError("FlashInfer direct/SIMT GEMM does not support bias.")
+            values = self._tactic_values(
+                result,
+                ("block_size", "outputs_per_block", "rows_per_block"),
+            )
+            tactic = self._flashinfer_direct_tactic(**values)
+            self._flashinfer_run_direct(input_2d, weight_t, output, pdl, tactic)
+            return output
+        if result.algorithm == "splitk":
+            values = self._tactic_values(
+                result,
+                ("mma_m", "mma_n", "split_k", "ab_stages"),
+            )
+            tactic = self._flashinfer_splitk_tactic(**values)
+            self._flashinfer_run_splitk(input_2d, weight_t, bias, output, pdl, tactic)
+            return output
+        raise RuntimeError(
+            f"FlashInfer cache entry has unsupported algorithm {result.algorithm!r}."
+        )
+
+    def _make_key(
+        self,
+        input_tensor: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+    ) -> GemmDispatchKey:
+        input_2d = input_tensor.reshape(-1, input_tensor.shape[-1])
+        return GemmDispatchKey(
+            sm=_current_sm(input_tensor.device),
+            m=int(input_2d.shape[0]),
+            n=int(weight.shape[0]),
+            k=int(weight.shape[1]),
+            has_bias=bias is not None,
+            cuda_graph=self.cuda_graph_enabled,
+        )
+
+    def _select_backend(self, key: GemmDispatchKey) -> LowMGemmBackend:
+        if self.backend != LowMGemmBackend.AUTO:
+            return self.backend
+        if os.environ.get(_DISABLE_CACHE_ENV, "0") != "1":
+            result = self._dispatch_cache.get(key.cache_key())
+            if result is not None:
+                return _normalize_backend(result.backend)
+        # FlashInfer's measured fallback chooses direct for the smallest shapes
+        # and split-K for the rest of M<=32. This is intentionally only a cache
+        # miss policy; offline tuning should cover every hot production shape.
+        if key.cache_key() not in self._cache_misses:
+            self._cache_misses.add(key.cache_key())
+            logger.warning(
+                f"No offline low-M GEMM decision for {key.cache_key()}; using "
+                "the FlashInfer direct/split-K fallback heuristic."
+            )
+        return LowMGemmBackend.FLASHINFER
+
+    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
+
+        key = self._make_key(input_tensor, weight, bias)
+        selected = self._select_backend(key)
+        if selected == LowMGemmBackend.CUBLAS:
+            return None
+        if selected != LowMGemmBackend.FLASHINFER:
+            raise RuntimeError(f"Unsupported cached low-M GEMM backend: {selected.value}")
+        if not self._is_flashinfer_supported(input_tensor.device):
+            return None
+
+        engaged_key = (selected.value, key.m, key.n, key.k, key.has_bias)
+        if engaged_key not in self._engaged:
+            self._engaged.add(engaged_key)
+            result = self._dispatch_cache.get(key.cache_key())
+            detail = ""
+            if result is not None and result.algorithm:
+                detail = f" ({result.algorithm})"
+            logger.info(
+                f"Low-M BF16 GEMM: routing M={key.m} N={key.n} K={key.k} "
+                f"bias={key.has_bias} to FlashInfer {_FLASHINFER_BACKEND}{detail}."
+            )
+
+        input_2d = input_tensor.detach().view(-1, input_tensor.shape[-1])
+        inference_weight = weight.detach()
+        inference_bias = bias.detach() if bias is not None else None
+        result = self._dispatch_cache.get(key.cache_key())
+        if self.backend == LowMGemmBackend.AUTO and result is not None:
+            output = self._launch_cached_flashinfer(
+                input_2d,
+                inference_weight,
+                inference_bias,
+                result,
+            )
+        else:
+            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 load read-only offline tuning caches."""
+
+    _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 tuned 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()
+
+
+def write_dispatch_cache(
+    path: Path, metadata: dict[str, Any], entries: dict[str, GemmTuningResult]
+) -> None:
+    """Write a dispatcher cache atomically for the offline tuning tool."""
+
+    payload = {
+        "schema_version": _DISPATCH_CACHE_SCHEMA_VERSION,
+        "metadata": metadata,
+        "entries": {key: asdict(result) for key, result in sorted(entries.items())},
+    }
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary_path = path.with_suffix(f"{path.suffix}.tmp")
+    temporary_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    temporary_path.replace(path)
+
+
+__all__ = [
+    "GemmDispatchKey",
+    "GemmTuningResult",
+    "LOW_M_GEMM_ACTIVE",
+    "LowMGemmBackend",
+    "apply_low_m_gemm",
+    "flush_low_m_gemm_shape_log",
+    "prepare_low_m_gemm",
+    "write_dispatch_cache",
+]
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 4906ad52f5e5..7a2795a4bb21 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
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..8577d196b561
--- /dev/null
+++ b/tests/unittest/_torch/modules/test_low_m_gemm.py
@@ -0,0 +1,226 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import atexit
+import json
+from types import SimpleNamespace
+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 (
+    GemmDispatchKey,
+    GemmTuningResult,
+    LowMGemmBackend,
+    LowMGemmDispatcher,
+    _configured_backend,
+    _flashinfer_commit,
+    _infer_op_name,
+    _normalize_backend,
+    _ShapeCollector,
+    write_dispatch_cache,
+)
+
+
+@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_flashinfer_commit_uses_package_identity(monkeypatch) -> None:
+    monkeypatch.delenv("TRTLLM_FLASHINFER_COMMIT", raising=False)
+    assert _flashinfer_commit(SimpleNamespace(__git_version__="abc123")) == "abc123"
+
+
+def test_flashinfer_commit_rejects_false_environment_identity(monkeypatch) -> None:
+    monkeypatch.setenv("TRTLLM_FLASHINFER_COMMIT", "declared")
+    with pytest.raises(RuntimeError, match="does not match"):
+        _flashinfer_commit(SimpleNamespace(__git_version__="actual"))
+
+
+def test_dispatch_key_includes_shape_layout_and_graph_mode() -> None:
+    key = GemmDispatchKey(sm=103, m=4, n=2304, k=8192)
+    assert key.cache_key() == "sm103:bf16:4x2304x8192:nt:nobias:graph"
+    assert (
+        GemmDispatchKey(sm=103, m=4, n=2304, k=8192, cuda_graph=False)
+        .cache_key()
+        .endswith(":eager")
+    )
+    assert (
+        GemmDispatchKey(sm=103, m=4, n=2304, k=8192, has_bias=True)
+        .cache_key()
+        .endswith(":bias:graph")
+    )
+
+
+@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
+
+
+def test_write_dispatch_cache(tmp_path) -> None:
+    output = tmp_path / "dispatch.json"
+    key = GemmDispatchKey(sm=103, m=1, n=256, k=8192).cache_key()
+    write_dispatch_cache(
+        output,
+        {"flashinfer_commit": "b195c7a8"},
+        {
+            key: GemmTuningResult(
+                backend="flashinfer",
+                algorithm="direct",
+                latency_us=2.1,
+                baseline_us=3.3,
+            )
+        },
+    )
+    document = json.loads(output.read_text(encoding="utf-8"))
+    assert document["schema_version"] == 3
+    assert document["metadata"]["flashinfer_commit"] == "b195c7a8"
+    assert document["entries"][key]["algorithm"] == "direct"
+
+
+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)
+
+
+def test_auto_cache_can_keep_shape_on_cublas_without_flashinfer(tmp_path, monkeypatch) -> None:
+    output = tmp_path / "dispatch.json"
+    key = GemmDispatchKey(sm=103, m=2, n=2304, k=8192, cuda_graph=False).cache_key()
+    write_dispatch_cache(
+        output,
+        {"pdl": True, "cuda_version": torch.version.cuda or "none"},
+        {key: GemmTuningResult(backend="cublas", algorithm="cublas")},
+    )
+    monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", "auto")
+    monkeypatch.setenv("TRTLLM_LOW_M_GEMM_TUNING_CACHE", str(output))
+    dispatcher = LowMGemmDispatcher()
+    module = torch.nn.Linear(8, 8)
+    dispatcher.prepare(module, cuda_graph_enabled=False)
+    assert (
+        dispatcher._select_backend(GemmDispatchKey(sm=103, m=2, n=2304, k=8192, cuda_graph=False))
+        == LowMGemmBackend.CUBLAS
+    )
+    assert module._low_m_gemm_name == ""
+
+
+def test_cached_tactic_requires_exact_fields() -> None:
+    result = GemmTuningResult(
+        backend="flashinfer",
+        algorithm="simt",
+        tactic={"block_size": 256, "rows_per_block": 4},
+    )
+    with pytest.raises(RuntimeError, match="Invalid 'simt' tactic"):
+        LowMGemmDispatcher._tactic_values(
+            result,
+            ("block_size", "outputs_per_block", "rows_per_block"),
+        )

From 6163a23f0b0817dc476cc11e74f224fd0ecc8deb Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 7 Aug 2026 00:32:40 -0700
Subject: [PATCH 11/18] [None][perf] use cache-free FlashInfer low-M GEMM

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 ATTRIBUTIONS-Python.md                        |   2 +-
 requirements.txt                              |   2 +-
 security_scanning/pyproject.toml              |   2 +-
 .../_torch/modules/fla/flashinfer_chunk.py    |  28 ++
 tensorrt_llm/_torch/modules/linear.py         |   5 +-
 tensorrt_llm/_torch/modules/low_m_gemm.py     | 428 ++----------------
 .../mamba/test_flashinfer_chunk_gdn.py        |  32 ++
 .../_torch/modules/test_low_m_gemm.py         | 145 +++---
 8 files changed, 173 insertions(+), 471 deletions(-)

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 <http://unlicense.org>
   - `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/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/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/linear.py b/tensorrt_llm/_torch/modules/linear.py
index d29073688c61..1b62be18d20f 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -566,9 +566,8 @@ 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 loads decisions produced by an offline
-        # FlashInfer direct/split-K versus cuBLAS sweep. It never benchmarks on
-        # the serving path and returns None for the normal GEMM fallback.
+        # 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:
diff --git a/tensorrt_llm/_torch/modules/low_m_gemm.py b/tensorrt_llm/_torch/modules/low_m_gemm.py
index 55e0417d926d..58a0a9069d52 100644
--- a/tensorrt_llm/_torch/modules/low_m_gemm.py
+++ b/tensorrt_llm/_torch/modules/low_m_gemm.py
@@ -5,32 +5,25 @@
 
 import atexit
 import enum
-import hashlib
 import json
 import os
 import threading
 from collections import defaultdict
-from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Optional
 
 import torch
+from packaging.version import InvalidVersion, Version
 
 from tensorrt_llm.logger import logger
-from tensorrt_llm.version import __version__ as trtllm_version
 
 from ..flashinfer_utils import get_env_enable_pdl
 
 _BACKEND_ENV = "TRTLLM_LOW_M_GEMM_BACKEND"
-_DISPATCH_CACHE_ENV = "TRTLLM_LOW_M_GEMM_TUNING_CACHE"
-_FLASHINFER_CACHE_ENV = "TRTLLM_FLASHINFER_AUTOTUNER_CACHE"
-_FLASHINFER_COMMIT_ENV = "TRTLLM_FLASHINFER_COMMIT"
-_FLASHINFER_SOURCE_ROOT_ENV = "TRTLLM_FLASHINFER_SOURCE_ROOT"
-_DISABLE_CACHE_ENV = "TRTLLM_DISABLE_GEMM_TUNING_CACHE"
 _SHAPE_LOG_ENV = "TRTLLM_LOW_M_GEMM_SHAPE_LOG"
-_DISPATCH_CACHE_SCHEMA_VERSION = 3
 _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
 
@@ -44,43 +37,6 @@ class LowMGemmBackend(str, enum.Enum):
     CUBLAS = "cublas"
 
 
-@dataclass(frozen=True)
-class GemmDispatchKey:
-    """Properties that can change the best low-M GEMM implementation."""
-
-    sm: int
-    m: int
-    n: int
-    k: int
-    a_type: str = "bf16"
-    b_type: str = "bf16"
-    c_type: str = "bf16"
-    trans_a: bool = False
-    trans_b: bool = True
-    has_bias: bool = False
-    cuda_graph: bool = True
-
-    def cache_key(self) -> str:
-        transpose = ("t" if self.trans_a else "n") + ("t" if self.trans_b else "n")
-        bias = "bias" if self.has_bias else "nobias"
-        execution = "graph" if self.cuda_graph else "eager"
-        return (
-            f"sm{self.sm}:{self.a_type}:{self.m}x{self.n}x{self.k}:{transpose}:{bias}:{execution}"
-        )
-
-
-@dataclass(frozen=True)
-class GemmTuningResult:
-    """One persisted dispatcher decision produced by offline tuning."""
-
-    backend: str
-    algorithm: Optional[str] = None
-    tactic: Optional[dict[str, Any]] = None
-    latency_us: Optional[float] = None
-    baseline_us: Optional[float] = None
-    measurements: Optional[dict[str, Any]] = None
-
-
 def _get_rank() -> int:
     for name in ("RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"):
         value = os.environ.get(name)
@@ -107,14 +63,6 @@ def _rank_path(path: str) -> Path:
     return resolved
 
 
-def _sha256(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as cache_file:
-        for chunk in iter(lambda: cache_file.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
 def _normalize_backend(value: str) -> LowMGemmBackend:
     normalized = value.strip().lower().replace("_", "-")
     aliases = {
@@ -139,55 +87,11 @@ def _configured_backend() -> LowMGemmBackend:
     return _normalize_backend(os.environ.get(_BACKEND_ENV, "off"))
 
 
-def _flashinfer_commit(flashinfer_module: Any) -> Optional[str]:
-    """Return a verifiable FlashInfer source revision for cache provenance."""
-
-    package_commit = getattr(flashinfer_module, "__git_version__", None)
-    if package_commit == "unknown":
-        package_commit = None
-    declared_commit = os.environ.get(_FLASHINFER_COMMIT_ENV)
-    if package_commit is not None and declared_commit is not None:
-        if package_commit != declared_commit:
-            raise RuntimeError(
-                f"{_FLASHINFER_COMMIT_ENV}={declared_commit} does not match "
-                f"the imported FlashInfer revision {package_commit}."
-            )
-    return package_commit or declared_commit
-
-
 def _current_sm(device: torch.device) -> int:
     major, minor = torch.cuda.get_device_capability(device)
     return major * 10 + minor
 
 
-def _configure_flashinfer_source_layout() -> None:
-    source_value = os.environ.get(_FLASHINFER_SOURCE_ROOT_ENV)
-    if not source_value:
-        return
-    source_root = Path(source_value).resolve()
-    source_csrc = source_root / "csrc"
-    source_include = source_root / "include"
-    if not source_csrc.is_dir() or not source_include.is_dir():
-        raise RuntimeError(
-            f"{_FLASHINFER_SOURCE_ROOT_ENV} is not a FlashInfer source checkout: {source_root}"
-        )
-
-    import flashinfer
-    from flashinfer.jit import env as jit_env
-
-    imported_root = Path(flashinfer.__file__).resolve().parents[1]
-    if imported_root != source_root:
-        raise RuntimeError(
-            f"Imported FlashInfer from {imported_root}, but "
-            f"{_FLASHINFER_SOURCE_ROOT_ENV}={source_root}."
-        )
-    # An installed wheel stores build inputs under flashinfer/data. A source
-    # overlay stores them at repository top level. Configure this before any
-    # FlashInfer JIT helper is first invoked.
-    jit_env.FLASHINFER_CSRC_DIR = source_csrc
-    jit_env.FLASHINFER_INCLUDE_DIR = source_include
-
-
 class _ShapeCollector:
     """Debug-only hot-shape collector with durable runtime-shape discovery."""
 
@@ -334,23 +238,14 @@ def _infer_op_name(layer: str) -> str:
 
 
 class LowMGemmDispatcher:
-    """Offline-cache dispatcher for FlashInfer direct/split-K versus cuBLAS."""
+    """Dispatch eligible BF16 GEMMs to FlashInfer's low-M CuTe DSL backend."""
 
     def __init__(self) -> None:
         self.backend = _configured_backend()
-        if self.backend in (LowMGemmBackend.AUTO, LowMGemmBackend.FLASHINFER):
-            _configure_flashinfer_source_layout()
         self.cuda_graph_enabled = True
-        self._dispatch_cache: dict[str, GemmTuningResult] = {}
-        self._dispatch_metadata: dict[str, Any] = {}
         self._prepared = False
         self._flashinfer_mm = None
-        self._flashinfer_direct_tactic = None
-        self._flashinfer_run_direct = None
-        self._flashinfer_splitk_tactic = None
-        self._flashinfer_run_splitk = None
-        self._engaged: set[tuple[str, int, int, int, bool]] = set()
-        self._cache_misses: set[str] = set()
+        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
         )
@@ -369,167 +264,41 @@ def prepare(self, model: torch.nn.Module, cuda_graph_enabled: bool) -> None:
             self._prepared = True
             return
 
-        self._load_dispatch_cache()
-        needs_flashinfer = self.backend == LowMGemmBackend.FLASHINFER or any(
-            result.backend == LowMGemmBackend.FLASHINFER.value
-            for result in self._dispatch_cache.values()
-        )
-        if needs_flashinfer:
-            self._prepare_flashinfer()
+        self._prepare_flashinfer()
         self._prepared = True
 
-    def _load_dispatch_cache(self) -> None:
-        if self.backend != LowMGemmBackend.AUTO:
-            return
-        if os.environ.get(_DISABLE_CACHE_ENV, "0") == "1":
-            logger.warning(
-                f"{_DISABLE_CACHE_ENV}=1: low-M GEMM auto mode will use its cache-miss heuristic."
-            )
-            return
-        cache_value = os.environ.get(_DISPATCH_CACHE_ENV)
-        if not cache_value:
-            raise RuntimeError(
-                f"{_BACKEND_ENV}=auto requires {_DISPATCH_CACHE_ENV}; run the "
-                "offline tuner first, use flashinfer for heuristic-only testing, "
-                "or set cublas to roll back."
-            )
-        cache_path = Path(cache_value)
-        if not cache_path.is_file():
-            raise FileNotFoundError(f"Low-M GEMM dispatch cache does not exist: {cache_path}")
-        document = json.loads(cache_path.read_text(encoding="utf-8"))
-        if document.get("schema_version") != _DISPATCH_CACHE_SCHEMA_VERSION:
-            raise RuntimeError(
-                f"Unsupported low-M GEMM cache schema in {cache_path}: "
-                f"{document.get('schema_version')!r}"
-            )
-        self._dispatch_metadata = document.get("metadata", {})
-        if not isinstance(self._dispatch_metadata, dict):
-            raise RuntimeError(f"Low-M GEMM cache {cache_path} has invalid metadata.")
-        expected_pdl = self._dispatch_metadata.get("pdl")
-        if not isinstance(expected_pdl, bool):
-            raise RuntimeError(
-                f"Low-M GEMM cache {cache_path} has no Boolean PDL runtime identity."
-            )
-        actual_pdl = bool(get_env_enable_pdl())
-        if expected_pdl != actual_pdl:
-            raise RuntimeError(
-                f"Low-M GEMM cache was tuned with PDL={expected_pdl}, but runtime PDL={actual_pdl}."
-            )
-        expected_cuda_version = self._dispatch_metadata.get("cuda_version")
-        if not isinstance(expected_cuda_version, str):
-            raise RuntimeError(f"Low-M GEMM cache {cache_path} has no CUDA runtime identity.")
-        actual_cuda_version = torch.version.cuda or "none"
-        if expected_cuda_version != actual_cuda_version:
-            raise RuntimeError(
-                "CUDA version does not match the low-M GEMM dispatch cache: "
-                f"expected {expected_cuda_version}, got {actual_cuda_version}."
-            )
-        entries = document.get("entries")
-        if not isinstance(entries, dict):
-            raise RuntimeError(f"Low-M GEMM cache {cache_path} has no object-valued entries.")
-        self._dispatch_cache = {key: GemmTuningResult(**value) for key, value in entries.items()}
-        logger.info(f"Loaded {len(self._dispatch_cache)} low-M GEMM decisions from {cache_path}.")
-
     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
-            from flashinfer.gemm.kernels.dense_bf16_gemm_direct import (
-                DirectTactic,
-                run_direct_dense,
-            )
-            from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import (
-                SplitKTactic,
-                run_splitk_dense,
-            )
         except (ImportError, ModuleNotFoundError) as error:
             raise RuntimeError(
-                "FlashInfer low-M GEMM requires a FlashInfer build containing "
-                "PR #4266 and nvidia-cutlass-dsl."
+                "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 pinned PR #4266 build is not active."
-            )
-        expected_arch = self._dispatch_metadata.get("gpu_arch")
-        actual_arch = f"sm{_current_sm(torch.device('cuda'))}"
-        if expected_arch is not None and expected_arch != actual_arch:
-            raise RuntimeError(
-                f"Low-M GEMM cache targets {expected_arch}, but this rank uses {actual_arch}."
-            )
-        expected_version = self._dispatch_metadata.get("flashinfer_version")
-        if expected_version is not None and expected_version != flashinfer_module.__version__:
-            raise RuntimeError(
-                "FlashInfer version does not match the low-M GEMM dispatch "
-                f"cache: expected {expected_version}, got "
-                f"{flashinfer_module.__version__}."
-            )
-        expected_trtllm_version = self._dispatch_metadata.get("trtllm_version")
-        if expected_trtllm_version is not None and expected_trtllm_version != trtllm_version:
-            raise RuntimeError(
-                "TensorRT-LLM version does not match the low-M GEMM dispatch "
-                f"cache: expected {expected_trtllm_version}, got "
-                f"{trtllm_version}."
-            )
-        expected_dispatcher_digest = self._dispatch_metadata.get("trtllm_low_m_gemm_sha256")
-        if expected_dispatcher_digest is not None:
-            actual_dispatcher_digest = _sha256(Path(__file__))
-            if actual_dispatcher_digest != expected_dispatcher_digest:
-                raise RuntimeError(
-                    "TensorRT-LLM low-M GEMM source does not match the offline "
-                    "dispatch cache; regenerate the cache with this checkout."
-                )
-        expected_commit = self._dispatch_metadata.get("flashinfer_commit")
-        actual_commit = _flashinfer_commit(flashinfer_module)
-        if (
-            expected_commit is not None
-            and self.backend == LowMGemmBackend.AUTO
-            and actual_commit != expected_commit
-        ):
-            raise RuntimeError(
-                "Imported FlashInfer does not match the cache's pinned "
-                f"commit {expected_commit}; got {actual_commit!r}."
+                "the required PR #4266 implementation is not active."
             )
         self._flashinfer_mm = mm_bf16
-        self._flashinfer_direct_tactic = DirectTactic
-        self._flashinfer_run_direct = run_direct_dense
-        self._flashinfer_splitk_tactic = SplitKTactic
-        self._flashinfer_run_splitk = run_splitk_dense
-
-        cache_disabled = os.environ.get(_DISABLE_CACHE_ENV, "0") == "1"
-        cache_value = os.environ.get(_FLASHINFER_CACHE_ENV)
-        if cache_disabled or not cache_value:
-            if self.backend == LowMGemmBackend.AUTO and not cache_disabled:
-                raise RuntimeError(
-                    f"{_BACKEND_ENV}=auto selected FlashInfer entries but "
-                    f"{_FLASHINFER_CACHE_ENV} is unset."
-                )
-            logger.warning(
-                "FlashInfer BF16 autotune cache is unavailable for provenance "
-                "validation; forced/cache-miss routing will use FlashInfer's "
-                "heuristic."
-            )
-            return
-
-        cache_path = Path(cache_value)
-        if not cache_path.is_file():
-            raise FileNotFoundError(f"FlashInfer autotune cache does not exist: {cache_path}")
-        expected_digest = self._dispatch_metadata.get("flashinfer_cache_sha256")
-        if expected_digest is not None:
-            actual_digest = _sha256(cache_path)
-            if actual_digest != expected_digest:
-                raise RuntimeError(
-                    "FlashInfer autotune cache checksum does not match the "
-                    "TRT-LLM dispatch cache; regenerate the pair together."
-                )
-        # Do not load this file into FlashInfer's runtime AutoTuner. Its
-        # default mapper aliases M=24 to M=32, whereas this dispatch cache is
-        # exact-M. The selected runner/tactic is launched directly below.
-        logger.info(f"Validated paired FlashInfer BF16 autotune cache {cache_path}.")
+        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]
@@ -577,89 +346,6 @@ def _is_flashinfer_supported(self, device: torch.device) -> bool:
         sm = _current_sm(device)
         return bool(self._flashinfer_mm.is_backend_supported(_FLASHINFER_BACKEND, sm))
 
-    @staticmethod
-    def _tactic_values(
-        result: GemmTuningResult,
-        expected_fields: tuple[str, ...],
-    ) -> dict[str, int]:
-        tactic = result.tactic
-        if not isinstance(tactic, dict) or set(tactic) != set(expected_fields):
-            raise RuntimeError(
-                f"Invalid {result.algorithm!r} tactic in low-M GEMM cache: "
-                f"expected {expected_fields}, got {tactic!r}."
-            )
-        return {field: int(tactic[field]) for field in expected_fields}
-
-    def _launch_cached_flashinfer(
-        self,
-        input_2d: torch.Tensor,
-        weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
-        result: GemmTuningResult,
-    ) -> torch.Tensor:
-        output = torch.empty(
-            (input_2d.shape[0], weight.shape[0]),
-            device=input_2d.device,
-            dtype=torch.bfloat16,
-        )
-        weight_t = weight.t()
-        pdl = get_env_enable_pdl()
-        if result.algorithm == "simt":
-            if bias is not None:
-                raise RuntimeError("FlashInfer direct/SIMT GEMM does not support bias.")
-            values = self._tactic_values(
-                result,
-                ("block_size", "outputs_per_block", "rows_per_block"),
-            )
-            tactic = self._flashinfer_direct_tactic(**values)
-            self._flashinfer_run_direct(input_2d, weight_t, output, pdl, tactic)
-            return output
-        if result.algorithm == "splitk":
-            values = self._tactic_values(
-                result,
-                ("mma_m", "mma_n", "split_k", "ab_stages"),
-            )
-            tactic = self._flashinfer_splitk_tactic(**values)
-            self._flashinfer_run_splitk(input_2d, weight_t, bias, output, pdl, tactic)
-            return output
-        raise RuntimeError(
-            f"FlashInfer cache entry has unsupported algorithm {result.algorithm!r}."
-        )
-
-    def _make_key(
-        self,
-        input_tensor: torch.Tensor,
-        weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
-    ) -> GemmDispatchKey:
-        input_2d = input_tensor.reshape(-1, input_tensor.shape[-1])
-        return GemmDispatchKey(
-            sm=_current_sm(input_tensor.device),
-            m=int(input_2d.shape[0]),
-            n=int(weight.shape[0]),
-            k=int(weight.shape[1]),
-            has_bias=bias is not None,
-            cuda_graph=self.cuda_graph_enabled,
-        )
-
-    def _select_backend(self, key: GemmDispatchKey) -> LowMGemmBackend:
-        if self.backend != LowMGemmBackend.AUTO:
-            return self.backend
-        if os.environ.get(_DISABLE_CACHE_ENV, "0") != "1":
-            result = self._dispatch_cache.get(key.cache_key())
-            if result is not None:
-                return _normalize_backend(result.backend)
-        # FlashInfer's measured fallback chooses direct for the smallest shapes
-        # and split-K for the rest of M<=32. This is intentionally only a cache
-        # miss policy; offline tuning should cover every hot production shape.
-        if key.cache_key() not in self._cache_misses:
-            self._cache_misses.add(key.cache_key())
-            logger.warning(
-                f"No offline low-M GEMM decision for {key.cache_key()}; using "
-                "the FlashInfer direct/split-K fallback heuristic."
-            )
-        return LowMGemmBackend.FLASHINFER
-
     def apply(
         self,
         module: torch.nn.Module,
@@ -682,47 +368,32 @@ def apply(
         if not self._is_candidate_shape(input_tensor, weight, bias):
             return None
 
-        key = self._make_key(input_tensor, weight, bias)
-        selected = self._select_backend(key)
-        if selected == LowMGemmBackend.CUBLAS:
-            return None
-        if selected != LowMGemmBackend.FLASHINFER:
-            raise RuntimeError(f"Unsupported cached low-M GEMM backend: {selected.value}")
         if not self._is_flashinfer_supported(input_tensor.device):
             return None
 
-        engaged_key = (selected.value, key.m, key.n, key.k, key.has_bias)
+        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])
+        engaged_key = (m, n, k, bias is not None)
         if engaged_key not in self._engaged:
             self._engaged.add(engaged_key)
-            result = self._dispatch_cache.get(key.cache_key())
-            detail = ""
-            if result is not None and result.algorithm:
-                detail = f" ({result.algorithm})"
             logger.info(
-                f"Low-M BF16 GEMM: routing M={key.m} N={key.n} K={key.k} "
-                f"bias={key.has_bias} to FlashInfer {_FLASHINFER_BACKEND}{detail}."
+                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."
             )
 
-        input_2d = input_tensor.detach().view(-1, input_tensor.shape[-1])
         inference_weight = weight.detach()
         inference_bias = bias.detach() if bias is not None else None
-        result = self._dispatch_cache.get(key.cache_key())
-        if self.backend == LowMGemmBackend.AUTO and result is not None:
-            output = self._launch_cached_flashinfer(
-                input_2d,
-                inference_weight,
-                inference_bias,
-                result,
-            )
-        else:
-            output = self._flashinfer_mm(
-                input_2d,
-                inference_weight.t(),
-                bias=inference_bias,
-                pdl=get_env_enable_pdl(),
-                out_dtype=torch.bfloat16,
-                backend=_FLASHINFER_BACKEND,
-            )
+        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:
@@ -735,7 +406,7 @@ def flush_shape_log(self) -> None:
 
 
 def prepare_low_m_gemm(model: torch.nn.Module, cuda_graph_enabled: bool) -> None:
-    """Label Linear modules and load read-only offline tuning caches."""
+    """Label Linear modules and initialize the cache-free FlashInfer path."""
 
     _DISPATCHER.prepare(model, cuda_graph_enabled)
 
@@ -746,7 +417,7 @@ def apply_low_m_gemm(
     weight: torch.Tensor,
     bias: Optional[torch.Tensor],
 ) -> Optional[torch.Tensor]:
-    """Return a tuned low-M BF16 result, or ``None`` for the normal path."""
+    """Return a low-M BF16 result, or ``None`` for the normal path."""
 
     return _DISPATCHER.apply(module, input_tensor, weight, bias)
 
@@ -757,29 +428,10 @@ def flush_low_m_gemm_shape_log() -> None:
     _DISPATCHER.flush_shape_log()
 
 
-def write_dispatch_cache(
-    path: Path, metadata: dict[str, Any], entries: dict[str, GemmTuningResult]
-) -> None:
-    """Write a dispatcher cache atomically for the offline tuning tool."""
-
-    payload = {
-        "schema_version": _DISPATCH_CACHE_SCHEMA_VERSION,
-        "metadata": metadata,
-        "entries": {key: asdict(result) for key, result in sorted(entries.items())},
-    }
-    path.parent.mkdir(parents=True, exist_ok=True)
-    temporary_path = path.with_suffix(f"{path.suffix}.tmp")
-    temporary_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
-    temporary_path.replace(path)
-
-
 __all__ = [
-    "GemmDispatchKey",
-    "GemmTuningResult",
     "LOW_M_GEMM_ACTIVE",
     "LowMGemmBackend",
     "apply_low_m_gemm",
     "flush_low_m_gemm_shape_log",
     "prepare_low_m_gemm",
-    "write_dispatch_cache",
 ]
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/test_low_m_gemm.py b/tests/unittest/_torch/modules/test_low_m_gemm.py
index 8577d196b561..ed35f37e44a5 100644
--- a/tests/unittest/_torch/modules/test_low_m_gemm.py
+++ b/tests/unittest/_torch/modules/test_low_m_gemm.py
@@ -3,7 +3,8 @@
 
 import atexit
 import json
-from types import SimpleNamespace
+import sys
+from types import ModuleType
 from unittest.mock import MagicMock
 
 import pytest
@@ -11,16 +12,12 @@
 
 from tensorrt_llm._torch.modules import linear as linear_module
 from tensorrt_llm._torch.modules.low_m_gemm import (
-    GemmDispatchKey,
-    GemmTuningResult,
     LowMGemmBackend,
     LowMGemmDispatcher,
     _configured_backend,
-    _flashinfer_commit,
     _infer_op_name,
     _normalize_backend,
     _ShapeCollector,
-    write_dispatch_cache,
 )
 
 
@@ -49,30 +46,51 @@ def test_legacy_exact_m_backend_environment_is_ignored(monkeypatch) -> None:
     assert _configured_backend() == LowMGemmBackend.OFF
 
 
-def test_flashinfer_commit_uses_package_identity(monkeypatch) -> None:
-    monkeypatch.delenv("TRTLLM_FLASHINFER_COMMIT", raising=False)
-    assert _flashinfer_commit(SimpleNamespace(__git_version__="abc123")) == "abc123"
+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
 
-def test_flashinfer_commit_rejects_false_environment_identity(monkeypatch) -> None:
-    monkeypatch.setenv("TRTLLM_FLASHINFER_COMMIT", "declared")
-    with pytest.raises(RuntimeError, match="does not match"):
-        _flashinfer_commit(SimpleNamespace(__git_version__="actual"))
+    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_dispatch_key_includes_shape_layout_and_graph_mode() -> None:
-    key = GemmDispatchKey(sm=103, m=4, n=2304, k=8192)
-    assert key.cache_key() == "sm103:bf16:4x2304x8192:nt:nobias:graph"
-    assert (
-        GemmDispatchKey(sm=103, m=4, n=2304, k=8192, cuda_graph=False)
-        .cache_key()
-        .endswith(":eager")
-    )
-    assert (
-        GemmDispatchKey(sm=103, m=4, n=2304, k=8192, has_bias=True)
-        .cache_key()
-        .endswith(":bias:graph")
-    )
+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)))
@@ -130,25 +148,31 @@ def test_infer_op_name_for_hot_modules(layer: str, expected: str) -> None:
     assert _infer_op_name(layer) == expected
 
 
-def test_write_dispatch_cache(tmp_path) -> None:
-    output = tmp_path / "dispatch.json"
-    key = GemmDispatchKey(sm=103, m=1, n=256, k=8192).cache_key()
-    write_dispatch_cache(
-        output,
-        {"flashinfer_commit": "b195c7a8"},
-        {
-            key: GemmTuningResult(
-                backend="flashinfer",
-                algorithm="direct",
-                latency_us=2.1,
-                baseline_us=3.3,
-            )
-        },
-    )
-    document = json.loads(output.read_text(encoding="utf-8"))
-    assert document["schema_version"] == 3
-    assert document["metadata"]["flashinfer_commit"] == "b195c7a8"
-    assert document["entries"][key]["algorithm"] == "direct"
+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)
+
+    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"
 
 
 def test_shape_collector_persists_new_runtime_shape_after_warmup(tmp_path) -> None:
@@ -191,36 +215,3 @@ def test_shape_collector_persists_new_runtime_shape_after_warmup(tmp_path) -> No
         }
     finally:
         atexit.unregister(collector.flush)
-
-
-def test_auto_cache_can_keep_shape_on_cublas_without_flashinfer(tmp_path, monkeypatch) -> None:
-    output = tmp_path / "dispatch.json"
-    key = GemmDispatchKey(sm=103, m=2, n=2304, k=8192, cuda_graph=False).cache_key()
-    write_dispatch_cache(
-        output,
-        {"pdl": True, "cuda_version": torch.version.cuda or "none"},
-        {key: GemmTuningResult(backend="cublas", algorithm="cublas")},
-    )
-    monkeypatch.setenv("TRTLLM_LOW_M_GEMM_BACKEND", "auto")
-    monkeypatch.setenv("TRTLLM_LOW_M_GEMM_TUNING_CACHE", str(output))
-    dispatcher = LowMGemmDispatcher()
-    module = torch.nn.Linear(8, 8)
-    dispatcher.prepare(module, cuda_graph_enabled=False)
-    assert (
-        dispatcher._select_backend(GemmDispatchKey(sm=103, m=2, n=2304, k=8192, cuda_graph=False))
-        == LowMGemmBackend.CUBLAS
-    )
-    assert module._low_m_gemm_name == ""
-
-
-def test_cached_tactic_requires_exact_fields() -> None:
-    result = GemmTuningResult(
-        backend="flashinfer",
-        algorithm="simt",
-        tactic={"block_size": 256, "rows_per_block": 4},
-    )
-    with pytest.raises(RuntimeError, match="Invalid 'simt' tactic"):
-        LowMGemmDispatcher._tactic_values(
-            result,
-            ("block_size", "outputs_per_block", "rows_per_block"),
-        )

From c623a46eca4f891f2d02c20b95782c70851375d2 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Sat, 8 Aug 2026 21:02:17 -0700
Subject: [PATCH 12/18] [None][perf] route measured low-M crossover shapes

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/low_m_gemm.py     | 41 +++++++++++--
 .../_torch/modules/test_low_m_gemm.py         | 57 +++++++++++++++++++
 2 files changed, 93 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/low_m_gemm.py b/tensorrt_llm/_torch/modules/low_m_gemm.py
index 58a0a9069d52..ff3eb0358bb1 100644
--- a/tensorrt_llm/_torch/modules/low_m_gemm.py
+++ b/tensorrt_llm/_torch/modules/low_m_gemm.py
@@ -5,6 +5,7 @@
 
 import atexit
 import enum
+import functools
 import json
 import os
 import threading
@@ -87,11 +88,37 @@ def _configured_backend() -> LowMGemmBackend:
     return _normalize_backend(os.environ.get(_BACKEND_ENV, "off"))
 
 
-def _current_sm(device: torch.device) -> int:
-    major, minor = torch.cuda.get_device_capability(device)
+@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) == (15520, 8192):
+        return m >= 15
+    return (m, n, k) == (16, 2304, 8192)
+
+
 class _ShapeCollector:
     """Debug-only hot-shape collector with durable runtime-shape discovery."""
 
@@ -368,13 +395,17 @@ def apply(
         if not self._is_candidate_shape(input_tensor, weight, bias):
             return None
 
-        if not self._is_flashinfer_supported(input_tensor.device):
-            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)
diff --git a/tests/unittest/_torch/modules/test_low_m_gemm.py b/tests/unittest/_torch/modules/test_low_m_gemm.py
index ed35f37e44a5..9ad6b5bb894c 100644
--- a/tests/unittest/_torch/modules/test_low_m_gemm.py
+++ b/tests/unittest/_torch/modules/test_low_m_gemm.py
@@ -15,8 +15,10 @@
     LowMGemmBackend,
     LowMGemmDispatcher,
     _configured_backend,
+    _device_sm,
     _infer_op_name,
     _normalize_backend,
+    _prefer_cublas_for_auto,
     _ShapeCollector,
 )
 
@@ -46,6 +48,18 @@ def test_legacy_exact_m_backend_environment_is_ignored(monkeypatch) -> None:
     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,
     *,
@@ -148,6 +162,24 @@ 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),
+        (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")
@@ -156,6 +188,7 @@ def test_apply_uses_public_flashinfer_heuristic(monkeypatch) -> None:
     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)
@@ -175,6 +208,30 @@ def test_apply_uses_public_flashinfer_heuristic(monkeypatch) -> None:
     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))

From d0afe1b511c1ae2e7a8ee3a8f6920fc80ac51c05 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Sat, 8 Aug 2026 23:09:41 -0700
Subject: [PATCH 13/18] [None][perf] route measured shared projection shapes

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/low_m_gemm.py        | 3 +++
 tests/unittest/_torch/modules/test_low_m_gemm.py | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/tensorrt_llm/_torch/modules/low_m_gemm.py b/tensorrt_llm/_torch/modules/low_m_gemm.py
index ff3eb0358bb1..a4af2dfc986e 100644
--- a/tensorrt_llm/_torch/modules/low_m_gemm.py
+++ b/tensorrt_llm/_torch/modules/low_m_gemm.py
@@ -27,6 +27,7 @@
 _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):
@@ -114,6 +115,8 @@ def _prefer_cublas_for_auto(m: int, n: int, k: int, sm: int) -> bool:
         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)
diff --git a/tests/unittest/_torch/modules/test_low_m_gemm.py b/tests/unittest/_torch/modules/test_low_m_gemm.py
index 9ad6b5bb894c..1250c638f2ee 100644
--- a/tests/unittest/_torch/modules/test_low_m_gemm.py
+++ b/tests/unittest/_torch/modules/test_low_m_gemm.py
@@ -167,6 +167,15 @@ def test_infer_op_name_for_hot_modules(layer: str, expected: str) -> None:
     [
         (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),

From efe186aa0a05f8c12810deb770dff253224086fb Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Tue, 11 Aug 2026 07:31:01 -0700
Subject: [PATCH 14/18] [None][fix] bound checkpoint loading host memory

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../models/checkpoints/base_weight_loader.py  |  51 ++++++--
 .../models/checkpoints/base_weight_mapper.py  |   7 ++
 .../_torch/models/modeling_qwen3_5.py         |  17 ++-
 tensorrt_llm/_torch/models/modeling_utils.py  |  78 +++++++++++-
 .../modeling/test_modeling_qwen3_5_vl_moe.py  |  23 +++-
 .../test_consumable_weights_dict.py           | 115 ++++++++++++++++++
 6 files changed, 272 insertions(+), 19 deletions(-)
 create mode 100644 tests/unittest/_torch/models/checkpoints/test_consumable_weights_dict.py

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_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/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()

From 1b5b4374f102abebdd627ccc1101cd23aed78022 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 7 Aug 2026 00:55:29 -0700
Subject: [PATCH 15/18] [None][fix] repair online EPLB lifecycle

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../_torch/modules/fused_moe/interface.py     |  15 +-
 .../modules/fused_moe/moe_load_balancer.py    | 320 ++++++++++++++++--
 .../_torch/modules/fused_moe/quantization.py  |   4 +-
 .../_torch/pyexecutor/model_engine.py         |  36 +-
 .../_torch/modules/test_moe_host_sharer.py    |  86 ++++-
 .../_torch/modules/test_moe_load_balancer.py  | 181 ++++++++++
 6 files changed, 598 insertions(+), 44 deletions(-)

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..d8dcc3451743 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,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import gc
+import mmap
 import os
 import threading
 from contextlib import nullcontext
@@ -20,6 +36,70 @@
 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 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 +108,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 +165,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 +181,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 +201,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 +252,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 +296,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
@@ -245,7 +349,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 +385,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 +400,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 +419,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 +430,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 +438,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 +503,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 +637,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 +657,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 +670,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 +691,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 +711,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 +744,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 +764,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 +802,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 +838,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 +873,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 +909,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 +919,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,7 +974,9 @@ 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
@@ -870,7 +1049,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)
@@ -913,6 +1093,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 +1173,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 +1197,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 +1221,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 +1318,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 +1332,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 +1345,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 +1374,9 @@ 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.advance_forward_iter()
         return False
 
 
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/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 7a2795a4bb21..1fd88af71b14 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -6911,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[
@@ -7121,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/tests/unittest/_torch/modules/test_moe_host_sharer.py b/tests/unittest/_torch/modules/test_moe_host_sharer.py
index 0ed0ee609bb8..728a270b4093 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,70 @@ 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)
+
+        with TemporaryDirectory() as directory, patch.dict(
+                "os.environ", {
+                    "TRTLLM_EPLB_SHM_DIR": directory,
+                    "TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL": "True",
+                }):
+            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()
+            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_without_flush(self):
+        """Peer mappings see writes without forcing the whole file to disk."""
+        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"
+                self.assertEqual(peer.buf[:8], b"EPLBTEST")
+                peer.buf[:8] = b"PEERVIEW"
+                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..232eb01bf2c2 100644
--- a/tests/unittest/_torch/modules/test_moe_load_balancer.py
+++ b/tests/unittest/_torch/modules/test_moe_load_balancer.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import os
 import unittest
 from unittest.mock import MagicMock, patch
@@ -311,6 +314,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):

From 3ebcb7b912fcd745e0f01d9147bb1d8bf0fa2ee0 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 7 Aug 2026 08:22:11 -0700
Subject: [PATCH 16/18] [None][fix] flush file-backed online EPLB weights

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../modules/fused_moe/moe_load_balancer.py       | 10 ++++++++++
 .../_torch/modules/test_moe_host_sharer.py       | 16 +++++++++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

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 d8dcc3451743..200d40d31308 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
@@ -69,6 +69,10 @@ def __init__(self,
     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()
 
@@ -331,6 +335,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:
diff --git a/tests/unittest/_torch/modules/test_moe_host_sharer.py b/tests/unittest/_torch/modules/test_moe_host_sharer.py
index 728a270b4093..901a77b4f350 100644
--- a/tests/unittest/_torch/modules/test_moe_host_sharer.py
+++ b/tests/unittest/_torch/modules/test_moe_host_sharer.py
@@ -110,16 +110,24 @@ def test_file_backed_storage(self):
         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(
@@ -134,8 +142,8 @@ def test_file_backed_storage(self):
             sharer.post_shutdown_cleanup()
             self.assertFalse(os.path.exists(backing_path))
 
-    def test_file_backed_mappings_are_coherent_without_flush(self):
-        """Peer mappings see writes without forcing the whole file to disk."""
+    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,
@@ -144,8 +152,10 @@ def test_file_backed_mappings_are_coherent_without_flush(self):
             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()

From 7ae0bd47e6b1ba3cb4710036c44085c5cae0a29b Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Sat, 8 Aug 2026 11:02:46 -0700
Subject: [PATCH 17/18] [None][test] make online EPLB transitions auditable

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../modules/fused_moe/moe_load_balancer.py    | 99 +++++++++++++++++++
 .../_torch/modules/test_moe_load_balancer.py  | 54 ++++++++++
 2 files changed, 153 insertions(+)

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 200d40d31308..a55601542b16 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import gc
+import json
 import mmap
 import os
 import threading
@@ -993,6 +994,16 @@ def __init__(self,
 
         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()
@@ -1086,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):
         """
@@ -1386,6 +1484,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         ):
             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/tests/unittest/_torch/modules/test_moe_load_balancer.py b/tests/unittest/_torch/modules/test_moe_load_balancer.py
index 232eb01bf2c2..9fc35703f4c1 100644
--- a/tests/unittest/_torch/modules/test_moe_load_balancer.py
+++ b/tests/unittest/_torch/modules/test_moe_load_balancer.py
@@ -1,7 +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
 
@@ -508,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."""
 

From 9a6889b2a2aba6f6e44483999dd972bc157c297b Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 7 Aug 2026 02:32:43 -0700
Subject: [PATCH 18/18] [None][fix] log forced MoE communication strategy

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
 .../communication/communication_factory.py    |  7 ++++-
 .../modules/moe/test_communication_factory.py | 29 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

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/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,
 ):