From 9a16d6d2b3ce16be74118f29c471efb63e0e4184 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 03:43:43 +0000
Subject: [PATCH 01/11] [None][feat] MoE / MegaMoE Rubin (SM107) support

Port the MoE feature family from the Rubin development branch that main does
not already carry:

* PDL launch support in the SiTU Triton kernel (`_torch/modules/situ.py`),
  gated on `get_env_enable_pdl()` and SM >= 90.
* `_Pointer.__add__` in `_torch/cute_dsl_kernels/rubin/moe/utils.py`, so a
  Rubin device pointer can be offset by an element count with the alignment
  narrowed to the real `gcd`.
* A `moe_a2a_cft_release` hook on the NVLink one-sided teardown path: the C++
  CFT logical endpoint is destroyed before the last Python reference to its
  MNNVL backing memory is dropped.
* MegaMoE DeepGEMM rendezvous: trust `MASTER_ADDR`/`MASTER_PORT`/`RANK`/
  `WORLD_SIZE` only when they actually describe this MPI world, so a
  disaggregated launcher's outer world cannot be adopted by an inner one.

The NVFP4 SiTU MoE unit test that accompanied this work, its `l0_b200` entry
and the de-duplication of the overlapping tactic-reachability case in
`tests/unittest/_torch/thop/serial/test_moe.py` are not part of this series;
this PR carries the production change only.

Every `cpp/tensorrt_llm/thop/` translation unit this series used to touch is
now left byte-identical to main. The EP-aware padded-token sizing carried for
`fp4BlockScaleMoe.cpp` (`num_experts` -> `local_num_experts` for
`getMaxPermutedPaddedCount` and `getMaxNumCtasInBatchDim`) is under separate
review in #19251: `routingIndicesWarpKernel` in `RoutingLlama4.cu` counts
tokens ungated by `isLocalExpert` and writes a global expert id as the CTA
batch index, so the smaller allocation can be written out of bounds on the
Llama4 routing path. The duplicate-routing-input log-level demote and the
unused `<ATen/native/cuda/Resize.h>` include removals are unrelated to SM107
enablement and are dropped from this series; `th_common` is already built as
C++20 on main, so those includes cost nothing but a line.

Two things this series used to carry are now main's and are not restated
here. `CuteDslFusedMoE.run_moe_nvfp4` already admits `ActivationType.SiTu`
after #19003, so `fused_moe_cute_dsl.py` is left byte-identical to main.
And the MegaMoE NVFP4 CuteDSL kernel tree that lived in
`_torch/cute_dsl_kernels/mega_moe_nvfp4/` was deleted by PR #17956 and
replaced with `_torch/cute_dsl_kernels/cutedsl_megamoe/`, which already
carries the same SM107 work (renamed: `build_sm107_static_idesc_base` ->
`build_static_idesc_base`, `Sm107MegaMoEKernel` ->
`rubin/inference/mega/BlockScaledSwapAbMegaMoeKernel`, and so on). The
`CuteDslFc12FusedMoE` backend and the Rubin fused-FC12 GEMM landed upstream
too.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../cute_dsl_kernels/rubin/moe/utils.py       |  6 ++
 tensorrt_llm/_torch/modules/situ.py           | 14 +++++
 .../communication/nvlink_one_sided.py         | 37 ++++++++++--
 .../fused_moe/mega_moe/mega_moe_deepgemm.py   | 23 ++++++--
 tests/unittest/_torch/moe/test_moe_a2a_cft.py | 57 ++++++++++++++++++-
 5 files changed, 124 insertions(+), 13 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/utils.py
index 8ecdd1f7744c..735880c16f6b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/utils.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/utils.py
@@ -44,6 +44,7 @@
 # This file is copied and modified from cutlass https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/core.py
 
 import ctypes
+import math
 import os
 from typing import Union
 
@@ -100,6 +101,11 @@ def __init__(
             f"pointer must be {self._assumed_align} bytes aligned"
         )
 
+    def __add__(self, offset: int) -> Pointer:  # type: ignore[override]
+        offset_bytes = offset * self._dtype.width // 8
+        assumed_align = math.gcd(offset_bytes, self._assumed_align)
+        return _Pointer(self._pointer + offset_bytes, self._dtype, self._addr_space, assumed_align)
+
     def size_in_bytes(self) -> int:
         return ctypes.sizeof(ctypes.c_void_p(int(self._pointer)))
 
diff --git a/tensorrt_llm/_torch/modules/situ.py b/tensorrt_llm/_torch/modules/situ.py
index 980893271bfe..83e8cadfe719 100644
--- a/tensorrt_llm/_torch/modules/situ.py
+++ b/tensorrt_llm/_torch/modules/situ.py
@@ -12,6 +12,10 @@
 import triton.language.extra.libdevice as tldevice  # type: ignore[import]
 from torch import nn
 
+from tensorrt_llm._utils import get_sm_version
+
+from ..flashinfer_utils import get_env_enable_pdl
+
 
 class SituAndMul(nn.Module):
     """SiTU activation with gate/up multiplicative gating.
@@ -59,6 +63,7 @@ def situ_and_mul_kernel(
     linear_beta,
     BLOCK_SIZE: tl.constexpr,
     HAS_LINEAR_BETA: tl.constexpr,
+    LAUNCH_WITH_PDL: tl.constexpr,
 ) -> None:
     """Fused :class:`SituAndMul` on a packed ``[gate | up]`` row layout."""
     i = tl.program_id(axis=0).to(tl.int64)
@@ -67,6 +72,9 @@ def situ_and_mul_kernel(
     o_row_ptr = o_ptr + o_stride * i
     x_row_ptr = x_ptr + x_stride * i
 
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_wait()
+
     offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offsets < d
 
@@ -80,6 +88,9 @@ def situ_and_mul_kernel(
 
     tl.store(o_row_ptr + offsets, result, mask=mask)
 
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 @torch.library.custom_op("trtllm::situ_and_mul", mutates_args=())
 def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor:
@@ -95,6 +106,7 @@ def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = No
     def grid(meta: Mapping[str, int]) -> tuple[int, int]:
         return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))
 
+    launch_with_pdl = get_env_enable_pdl() and get_sm_version() >= 90
     situ_and_mul_kernel[grid](
         o_ptr=output,
         o_stride=output.stride(0),
@@ -105,6 +117,8 @@ def grid(meta: Mapping[str, int]) -> tuple[int, int]:
         linear_beta=float(linear_beta) if linear_beta is not None else 1.0,
         BLOCK_SIZE=1024,
         HAS_LINEAR_BETA=linear_beta is not None,
+        LAUNCH_WITH_PDL=launch_with_pdl,
+        launch_pdl=launch_with_pdl,
     )
     return output
 
diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
index 3dc1a1f3245c..a59341f5a610 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
@@ -557,6 +557,37 @@ def supports_post_quant_dispatch(self) -> bool:
         """
         return True
 
+    @classmethod
+    def _release_workspace(cls, workspace_key: Tuple[object, ...]) -> None:
+        """Release a shared workspace after its last communicator is destroyed."""
+        workspace_state = cls._WORKSPACES.get(workspace_key)
+        if workspace_state is None:
+            return
+
+        if workspace_state.get("cft_initialized", False):
+            # The C++ manager owns a logical endpoint bound to this allocation.
+            # Destroy it before dropping the final Python references to the
+            # MNNVL memory, otherwise the LE outlives its backing pages.
+            #
+            # The release op is not registered in every build yet (only
+            # moe_a2a_cft_initialize is). Where it is missing the LE is
+            # reclaimed at process exit, which is the pre-existing behavior;
+            # the ordering below is already correct for when it lands.
+            release_cft_manager = getattr(torch.ops.trtllm, "moe_a2a_cft_release", None)
+            if release_cft_manager is None:
+                tllm_logger.warning_once(
+                    "moe_a2a_cft_release is not available in this build; the CFT "
+                    "logical endpoint will only be reclaimed at process exit.",
+                    key="moe_a2a_cft_release_missing",
+                )
+            else:
+                release_cft_manager(workspace_state["workspace"], workspace_state["ep_rank"])
+
+        cls._WORKSPACES.pop(workspace_key)
+        if cls._WORKSPACE is workspace_state:
+            cls._WORKSPACE = None
+        workspace_state.clear()
+
     def destroy(self):
         """Release shared state during explicit, rank-coordinated teardown."""
         if getattr(self, "_destroyed", False):
@@ -580,11 +611,7 @@ def destroy(self):
             NVLinkOneSided._WORKSPACE_REFCOUNTS[workspace_key] = refcount
         else:
             NVLinkOneSided._WORKSPACE_REFCOUNTS.pop(workspace_key, None)
-            workspace_state = NVLinkOneSided._WORKSPACES.pop(workspace_key, None)
-            if NVLinkOneSided._WORKSPACE is workspace_state:
-                NVLinkOneSided._WORKSPACE = None
-            if workspace_state is not None:
-                workspace_state.clear()
+            NVLinkOneSided._release_workspace(workspace_key)
 
         self.mnnvl_mem = None
         self.workspace = None
diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
index fb601445fb2f..2acbf202aaba 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
@@ -592,13 +592,24 @@ def _pick_rendezvous():
                 return (host, port)
             return None
 
-        # Respect pre-set launcher env vars (Slurm, Ray, torchrun).
-        if not all(os.environ.get(k) for k in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")):
+        # Respect launcher env vars only when they describe this MPI world.
+        # A disaggregated launcher may export its outer context/generation
+        # world into each server's independent model MPI world.
+        launcher_env_matches_mpi = False
+        if all(os.environ.get(k) for k in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")):
+            try:
+                launcher_env_matches_mpi = (
+                    int(os.environ["RANK"]) == rank and int(os.environ["WORLD_SIZE"]) == world_size
+                )
+            except ValueError:
+                pass
+
+        if not launcher_env_matches_mpi:
             host, port = comm.bcast(_pick_rendezvous(), root=0)
-            os.environ.setdefault("MASTER_ADDR", host)
-            os.environ.setdefault("MASTER_PORT", str(port))
-            os.environ.setdefault("RANK", str(rank))
-            os.environ.setdefault("WORLD_SIZE", str(world_size))
+            os.environ["MASTER_ADDR"] = host
+            os.environ["MASTER_PORT"] = str(port)
+            os.environ["RANK"] = str(rank)
+            os.environ["WORLD_SIZE"] = str(world_size)
         else:
             host = os.environ["MASTER_ADDR"]
             port = int(os.environ["MASTER_PORT"])
diff --git a/tests/unittest/_torch/moe/test_moe_a2a_cft.py b/tests/unittest/_torch/moe/test_moe_a2a_cft.py
index 0d9bf38dbac3..7ca336070abf 100644
--- a/tests/unittest/_torch/moe/test_moe_a2a_cft.py
+++ b/tests/unittest/_torch/moe/test_moe_a2a_cft.py
@@ -13,7 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import MagicMock
+
 import pytest
+import torch
 
 from tensorrt_llm._torch.moe.fused_moe.communication.moe_alltoall import (
     get_force_cft as get_force_cft_standalone,
@@ -23,12 +26,14 @@
 )
 from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import (
     FORCE_CFT_ENV,
+    NVLinkOneSided,
     get_force_cft,
     should_use_cft,
 )
 
-# Environment-variable parsing only. The marker is also what makes the file
-# reachable: the CPU stage collects only files that carry it.
+# Every check in this file is pure mock/env logic with no GPU work, so it can
+# ride the CPU-only CI stage. tests/unittest/conftest.py drops any file without
+# this marker from CPU stages.
 pytestmark = pytest.mark.cpu_only
 
 
@@ -76,3 +81,51 @@ def test_should_use_cft(
         should_use_cft_standalone(can_use_cft, force_cft, 128, runtime_max_tokens_per_rank)
         is expected
     )
+
+
+def test_destroy_releases_cft_manager_before_workspace_allocation(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    workspace_key = ("cft-workspace",)
+    workspace = object()
+    mnnvl_mem = object()
+    workspace_state = {
+        "cft_initialized": True,
+        "workspace": workspace,
+        "ep_rank": 3,
+        "mnnvl_mem": mnnvl_mem,
+    }
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state})
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {workspace_key: 1})
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+
+    release_cft_manager = MagicMock()
+
+    def verify_workspace_is_alive(workspace_arg: object, ep_rank: int) -> None:
+        assert workspace_arg is workspace
+        assert ep_rank == 3
+        assert NVLinkOneSided._WORKSPACES[workspace_key]["mnnvl_mem"] is mnnvl_mem
+
+    release_cft_manager.side_effect = verify_workspace_is_alive
+    monkeypatch.setattr(torch.ops.trtllm, "moe_a2a_cft_release", release_cft_manager, raising=False)
+
+    comm = NVLinkOneSided.__new__(NVLinkOneSided)
+    comm._destroyed = False
+    comm._workspace_key = workspace_key
+    comm._workspace_state = workspace_state
+    comm._workspace_lifecycle = None
+    comm._workspace_registered = True
+    comm.mnnvl_mem = mnnvl_mem
+    comm.workspace = workspace
+    comm._dispatch_state = {"phase": "idle"}
+
+    comm.destroy()
+
+    release_cft_manager.assert_called_once_with(workspace, 3)
+    assert workspace_key not in NVLinkOneSided._WORKSPACES
+    assert workspace_key not in NVLinkOneSided._WORKSPACE_REFCOUNTS
+    assert NVLinkOneSided._WORKSPACE is None
+    assert workspace_state == {}
+    assert comm.mnnvl_mem is None
+    assert comm.workspace is None

From c056fb25715b260a8d5f0a5e51c4447aabd18af1 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 04:56:02 +0000
Subject: [PATCH 02/11] [None][fix] Restore the moe_a2a_cft_release op

NVLinkOneSided teardown calls trtllm::moe_a2a_cft_release to destroy the CFT
logical endpoint while its workspace is still alive, but the op itself was
never registered: the caller survived a rebase that dropped the definition.
Any teardown of a CFT-initialized workspace therefore failed, and the endpoint
could outlive the MNNVL virtual address it is bound to.

Add moeA2ACftReleaseOp to moeAlltoAllOp.cpp, adapted to the current
single-manager layout: it is idempotent, and resets g_cft_manager only when
the manager is actually bound to the rank region of the workspace being torn
down, so a manager belonging to another allocation survives. ~CftLeManager
already unbinds and destroys the endpoints.

With the op registered, drop the hasattr fallback in _release_workspace and
call it directly.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp   | 51 ++++++++++++++++++-
 .../communication/nvlink_one_sided.py         | 20 ++------
 2 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp
index be1b71cc0f10..fe2d63c8f582 100644
--- a/cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp
+++ b/cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp
@@ -299,7 +299,15 @@ torch::Tensor moeA2AInitializeOp(torch::Tensor const& workspace, int64_t epRank,
 // CFT Handle-Based Counted Writes Initialization
 // ============================================================================
 
-// Static CftLeManager — lives for the process lifetime (like workspace).
+// Static CftLeManager — one per process, bound to a single workspace.
+//
+// The manager owns a logical endpoint bound to the workspace's MNNVL
+// allocation, so it must be destroyed before that allocation is freed and its
+// virtual address is recycled; otherwise a later binding could resolve to an
+// endpoint left over from a dead allocation. The Python NVLinkOneSided
+// teardown calls moe_a2a_cft_release while the workspace is still alive.
+// MoeAlltoAll otherwise retains its shared workspaces for the process
+// lifetime, so the manager normally survives until static destruction.
 static std::unique_ptr<tensorrt_llm::kernels::moe_comm::CftLeManager> g_cft_manager;
 
 // Initialize CFT Logical Endpoints by binding the LE to the MNNVL workspace.
@@ -396,6 +404,45 @@ void moeA2ACftInitializeOp(torch::Tensor const& workspace, int64_t workspaceMemH
     }
 }
 
+// Release the CFT logical endpoint before its backing workspace is freed.
+//
+// Idempotent, and a no-op unless the manager is actually bound to this
+// workspace's rank region: the caller passes the workspace it is tearing down,
+// and a manager bound to some other allocation must outlive that teardown.
+// Destroying the manager here — rather than at static destruction — keeps the
+// endpoint from outliving the virtual address it is bound to.
+void moeA2ACftReleaseOp(torch::Tensor const& workspace, int64_t epRank)
+{
+    CHECK_TH_CUDA(workspace);
+    CHECK_TYPE(workspace, torch::kUInt8);
+    TORCH_CHECK(workspace.dim() == 2, "workspace must be a 2D tensor of shape [epSize, sizePerRank]");
+    TORCH_CHECK(epRank >= 0 && epRank < workspace.size(0), "epRank must be in the range [0, epSize)");
+
+    if (!g_cft_manager)
+    {
+        return;
+    }
+
+    // An uninitialized manager holds no endpoint (initialization threw part
+    // way through); drop it unconditionally so a retry starts clean.
+    if (!g_cft_manager->isInitialized())
+    {
+        g_cft_manager.reset();
+        return;
+    }
+
+    CUdeviceptr workspaceRankPtr
+        = reinterpret_cast<CUdeviceptr>(workspace.data_ptr<uint8_t>() + epRank * workspace.stride(0));
+    if (g_cft_manager->getLocalBackingPtr() != workspaceRankPtr)
+    {
+        return;
+    }
+
+    // ~CftLeManager runs destroy(): unbind, destroy the local and imported
+    // endpoints, and release the reserved LE id block.
+    g_cft_manager.reset();
+}
+
 // MoE All-to-All Dispatch Operation
 // This operation dispatches tokens and their associated payloads to different expert ranks.
 //
@@ -1062,6 +1109,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module)
     module.def(
         "moe_a2a_cft_initialize(Tensor(a!) workspace, int workspace_mem_handle, "
         "int workspace_size_per_rank, int ep_rank, int ep_size) -> ()");
+    module.def("moe_a2a_cft_release(Tensor(a!) workspace, int ep_rank) -> ()");
     module.def(
         "moe_a2a_initialize(Tensor(a!) workspace, int ep_rank, int ep_size, int max_num_tokens_per_rank, "
         "int? eplb_stats_num_experts=None, bool can_use_cft_counted_writes=False) -> Tensor");
@@ -1088,4 +1136,5 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, module)
     module.impl(
         "moe_a2a_get_combine_payload_tensor", &tensorrt_llm::torch_ext::moe_comm::moeA2AGetCombinePayloadTensorOp);
     module.impl("moe_a2a_cft_initialize", &tensorrt_llm::torch_ext::moe_comm::moeA2ACftInitializeOp);
+    module.impl("moe_a2a_cft_release", &tensorrt_llm::torch_ext::moe_comm::moeA2ACftReleaseOp);
 }
diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
index a59341f5a610..9fdbf3646f35 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
@@ -567,21 +567,11 @@ def _release_workspace(cls, workspace_key: Tuple[object, ...]) -> None:
         if workspace_state.get("cft_initialized", False):
             # The C++ manager owns a logical endpoint bound to this allocation.
             # Destroy it before dropping the final Python references to the
-            # MNNVL memory, otherwise the LE outlives its backing pages.
-            #
-            # The release op is not registered in every build yet (only
-            # moe_a2a_cft_initialize is). Where it is missing the LE is
-            # reclaimed at process exit, which is the pre-existing behavior;
-            # the ordering below is already correct for when it lands.
-            release_cft_manager = getattr(torch.ops.trtllm, "moe_a2a_cft_release", None)
-            if release_cft_manager is None:
-                tllm_logger.warning_once(
-                    "moe_a2a_cft_release is not available in this build; the CFT "
-                    "logical endpoint will only be reclaimed at process exit.",
-                    key="moe_a2a_cft_release_missing",
-                )
-            else:
-                release_cft_manager(workspace_state["workspace"], workspace_state["ep_rank"])
+            # MNNVL memory, otherwise the endpoint outlives its backing pages
+            # and a recycled virtual address could resolve to it.
+            torch.ops.trtllm.moe_a2a_cft_release(
+                workspace_state["workspace"], workspace_state["ep_rank"]
+            )
 
         cls._WORKSPACES.pop(workspace_key)
         if cls._WORKSPACE is workspace_state:

From 1a9e3f27a906a37ab520283371d2cbc81e73781c Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 07:16:02 +0000
Subject: [PATCH 03/11] [None][fix] bind MegaMoE ranks to their local CUDA
 device

The CuteDSL MegaMoE bootstrap selected a device from the global MPI rank, which
is wrong on any multi-node launch: rank 8 on a 8-GPU-per-node job would index
past the local device count or collide with another node's rank. Derive the
device from local_mpi_comm().Get_rank() modulo the visible device count, set it
before the process group is created, and pass device_id= to
init_process_group so NCCL binds to the same device.

Ported from the Rubin branch, where this landed as 70bf7e4fd1 and was then lost
when a later replay commit (4f32abfc01) removed the device-binding half.
Restored on the internal branch as f7c74778b4. The other half of that original
commit, _setup_cutlass_dsl_compatibility, survived and is already in the core
PR of this series -- it is not duplicated here.

NOT RUN: no GPU and no multi-rank launcher here; the bootstrap path needs a real
MPI run to validate.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../moe/fused_moe/mega_moe/mega_moe_cute_dsl.py  | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py
index 49c52f1965e0..7e9e1ffb208f 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py
@@ -818,11 +818,12 @@ def _maybe_init_torch_dist_under_mpi(self):
         """
         if not dist.is_available() or dist.is_initialized():
             return
-        from tensorrt_llm._utils import mpi_comm, mpi_rank, mpi_world_size
+        from tensorrt_llm._utils import local_mpi_comm, mpi_comm, mpi_rank, mpi_world_size
 
         try:
             world = mpi_world_size()
             rank = mpi_rank()
+            local_rank = local_mpi_comm().Get_rank()
         except Exception as e:  # not under MPI either -> leave uninitialized
             logger.debug(
                 f"[MegaMoECuteDsl] MPI rank query failed ({e!r}); "
@@ -867,13 +868,24 @@ def _pick_rendezvous():
         host, port = mpi_comm().bcast(_pick_rendezvous() if rank == 0 else None, root=0)
         os.environ["MASTER_ADDR"] = str(host)
         os.environ["MASTER_PORT"] = str(port)
+        device_id = None
+        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
+            device_index = local_rank % torch.cuda.device_count()
+            torch.cuda.set_device(device_index)
+            device_id = torch.device("cuda", device_index)
         logger.info(
             f"[MegaMoECuteDsl] torch.distributed not initialized under MPI; "
             f"bootstrapping NCCL WORLD group (rank={rank}/{world}, "
+            f"local_rank={local_rank}, device={device_id}, "
             f"{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}) for the "
             f"EP rendezvous."
         )
-        dist.init_process_group(backend="cuda:nccl,cpu:gloo", rank=rank, world_size=world)
+        dist.init_process_group(
+            backend="cuda:nccl,cpu:gloo",
+            rank=rank,
+            world_size=world,
+            device_id=device_id,
+        )
 
     def _resolve_ep_pg(self):
         """Return the torch.distributed ProcessGroup for the EP sub-world.

From c1142adc177adc799f83b20d6f8778a49eb2b735 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 09:26:00 +0000
Subject: [PATCH 04/11] [None][fix] make MoE teardown and MASTER_PORT reuse
 fail cleanly

`_release_workspace` called `moe_a2a_cft_release` before removing the
workspace from `_WORKSPACES`. If the op raised, the entry survived --
and by then `destroy` had already decremented the refcount and
unregistered the lifecycle, so nothing would call it again. The next
communicator on that key would adopt an allocation whose endpoint state
is unknown. Drop the workspace in a `finally` instead; the exception
still propagates, it just no longer leaves a reusable workspace behind.

Tests: the existing destroy test now calls `destroy()` twice, since
rank-coordinated teardown can reach it more than once and a second CFT
release would destroy an endpoint this communicator no longer owns. A
new test covers the release-failure path.

Separately, the MegaMoE DeepGEMM rendezvous accepted the launcher
environment on RANK/WORLD_SIZE alone and only converted MASTER_PORT
afterwards, so a non-numeric port raised an uncaught ValueError instead
of falling back to the MPI rendezvous. Parse it inside the same guard.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../communication/nvlink_one_sided.py         | 35 +++++++-----
 .../fused_moe/mega_moe/mega_moe_deepgemm.py   | 10 +++-
 tests/unittest/_torch/moe/test_moe_a2a_cft.py | 53 +++++++++++++++++++
 3 files changed, 83 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
index 9fdbf3646f35..d9e33c78958a 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py
@@ -564,19 +564,28 @@ def _release_workspace(cls, workspace_key: Tuple[object, ...]) -> None:
         if workspace_state is None:
             return
 
-        if workspace_state.get("cft_initialized", False):
-            # The C++ manager owns a logical endpoint bound to this allocation.
-            # Destroy it before dropping the final Python references to the
-            # MNNVL memory, otherwise the endpoint outlives its backing pages
-            # and a recycled virtual address could resolve to it.
-            torch.ops.trtllm.moe_a2a_cft_release(
-                workspace_state["workspace"], workspace_state["ep_rank"]
-            )
-
-        cls._WORKSPACES.pop(workspace_key)
-        if cls._WORKSPACE is workspace_state:
-            cls._WORKSPACE = None
-        workspace_state.clear()
+        try:
+            if workspace_state.get("cft_initialized", False):
+                # The C++ manager owns a logical endpoint bound to this
+                # allocation. Destroy it before dropping the final Python
+                # references to the MNNVL memory, otherwise the endpoint
+                # outlives its backing pages and a recycled virtual address
+                # could resolve to it.
+                torch.ops.trtllm.moe_a2a_cft_release(
+                    workspace_state["workspace"], workspace_state["ep_rank"]
+                )
+        finally:
+            # Drop the workspace whether or not the release succeeded. By the
+            # time we get here ``destroy`` has already decremented the refcount
+            # and unregistered the lifecycle, so nothing will call this again:
+            # leaving the entry behind would hand a later communicator an
+            # allocation whose endpoint state is unknown. The exception still
+            # propagates -- the caller learns the release failed, but not by
+            # inheriting a reusable half-released workspace.
+            cls._WORKSPACES.pop(workspace_key, None)
+            if cls._WORKSPACE is workspace_state:
+                cls._WORKSPACE = None
+            workspace_state.clear()
 
     def destroy(self):
         """Release shared state during explicit, rank-coordinated teardown."""
diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
index 2acbf202aaba..bb97994af9c4 100644
--- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py
@@ -596,13 +596,19 @@ def _pick_rendezvous():
         # A disaggregated launcher may export its outer context/generation
         # world into each server's independent model MPI world.
         launcher_env_matches_mpi = False
+        launcher_port = None
         if all(os.environ.get(k) for k in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")):
             try:
+                # Parse the port here too: accepting the launcher env on
+                # RANK/WORLD_SIZE alone and only converting MASTER_PORT later
+                # turns a non-numeric port into an uncaught ValueError instead
+                # of falling back to the MPI rendezvous.
+                launcher_port = int(os.environ["MASTER_PORT"])
                 launcher_env_matches_mpi = (
                     int(os.environ["RANK"]) == rank and int(os.environ["WORLD_SIZE"]) == world_size
                 )
             except ValueError:
-                pass
+                launcher_env_matches_mpi = False
 
         if not launcher_env_matches_mpi:
             host, port = comm.bcast(_pick_rendezvous(), root=0)
@@ -612,7 +618,7 @@ def _pick_rendezvous():
             os.environ["WORLD_SIZE"] = str(world_size)
         else:
             host = os.environ["MASTER_ADDR"]
-            port = int(os.environ["MASTER_PORT"])
+            port = launcher_port
 
         device_id = None
         if torch.cuda.is_available() and torch.cuda.device_count() > 0:
diff --git a/tests/unittest/_torch/moe/test_moe_a2a_cft.py b/tests/unittest/_torch/moe/test_moe_a2a_cft.py
index 7ca336070abf..3f2589e7e997 100644
--- a/tests/unittest/_torch/moe/test_moe_a2a_cft.py
+++ b/tests/unittest/_torch/moe/test_moe_a2a_cft.py
@@ -120,6 +120,11 @@ def verify_workspace_is_alive(workspace_arg: object, ep_rank: int) -> None:
     comm.workspace = workspace
     comm._dispatch_state = {"phase": "idle"}
 
+    comm.destroy()
+    # Teardown is rank-coordinated and may be reached twice (explicit destroy
+    # plus a later sweep). The second call must be inert: releasing the CFT
+    # endpoint twice would destroy an endpoint this communicator no longer
+    # owns, and the workspace state has already been cleared.
     comm.destroy()
 
     release_cft_manager.assert_called_once_with(workspace, 3)
@@ -129,3 +134,51 @@ def verify_workspace_is_alive(workspace_arg: object, ep_rank: int) -> None:
     assert workspace_state == {}
     assert comm.mnnvl_mem is None
     assert comm.workspace is None
+
+
+def test_destroy_drops_workspace_when_cft_release_fails(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """A failed CFT release must not leave a reusable workspace behind.
+
+    ``destroy`` decrements the refcount and unregisters the lifecycle before
+    calling ``_release_workspace``, so nothing retries. If the release raised
+    and the entry survived in ``_WORKSPACES``, the next communicator built on
+    the same key would adopt an allocation whose endpoint state is unknown.
+    """
+    workspace_key = ("cft-workspace-failing",)
+    workspace = object()
+    mnnvl_mem = object()
+    workspace_state = {
+        "cft_initialized": True,
+        "workspace": workspace,
+        "ep_rank": 3,
+        "mnnvl_mem": mnnvl_mem,
+    }
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state})
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {workspace_key: 1})
+    monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+
+    release_cft_manager = MagicMock(side_effect=RuntimeError("cft release failed"))
+    monkeypatch.setattr(torch.ops.trtllm, "moe_a2a_cft_release", release_cft_manager, raising=False)
+
+    comm = NVLinkOneSided.__new__(NVLinkOneSided)
+    comm._destroyed = False
+    comm._workspace_key = workspace_key
+    comm._workspace_state = workspace_state
+    comm._workspace_lifecycle = None
+    comm._workspace_registered = True
+    comm.mnnvl_mem = mnnvl_mem
+    comm.workspace = workspace
+    comm._dispatch_state = {"phase": "idle"}
+
+    # The failure is reported rather than swallowed.
+    with pytest.raises(RuntimeError, match="cft release failed"):
+        comm.destroy()
+
+    release_cft_manager.assert_called_once_with(workspace, 3)
+    # ...but the workspace is gone either way, so nothing can adopt it.
+    assert workspace_key not in NVLinkOneSided._WORKSPACES
+    assert NVLinkOneSided._WORKSPACE is None
+    assert workspace_state == {}

From 8b8169bd90fae2f634d4b9c0013ba5285f35e5a7 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 09:30:59 +0000
Subject: [PATCH 05/11] [None][fix] mask the Rubin manual-MMA scale-factor
 descriptor fields

``rubin/moe/manual_mma_128dp.compute_idesc`` OR'd the shifted SF-ID fields
in without masking. ``Int32`` is signed, so ``>>`` is an arithmetic shift: a
TMEM address with bit 31 set sign-extends and leaves stray high bits. Bit 31
is ``_BIT_K_SIZE_LO``, so for umma_k=128 that turns k_size from 2 into 3 and
the MMA runs the wrong K.

``compute_idesc`` in
cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py -- the
reference this encoding was copied from -- already masks both fields; this
copy had dropped it.

``build_static_idesc_base`` also accepted umma_m=64, which is not encodable
here: ``m_dim = umma_m >> 4`` at bit 24 only lines up with the real m_dim
field (bits [27,29), holding M >> 7) for 128 and 256. For 64 it sets bit 26,
which is ``_BIT_SFA_LAYOUT``, silently selecting SFA_128dp_Unique and
discarding the caller's sfa_layout. Only 128 and 256 ever reach this copy,
so the assertion now says so.

The FC12 cluster/TMA guards that shipped alongside this fix targeted
``cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py``, a tree PR #17956 removed
and replaced with ``cute_dsl_kernels/cutedsl_megamoe/``; the replacement
already carries the equivalent guards, so only the manual-MMA half remains.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../rubin/moe/manual_mma_128dp.py              | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/manual_mma_128dp.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/manual_mma_128dp.py
index cf80927d2f30..ba5448d00b6b 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/manual_mma_128dp.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/manual_mma_128dp.py
@@ -73,7 +73,12 @@ def build_static_idesc_base(
     OR's them in from the SF TMEM addresses.  n_dim is static (non-swapAB) and
     folded in here.
     """
-    assert umma_m in (64, 128, 256), f"Unsupported UMMA_M={umma_m}"
+    # 64 is not encodable here. ``m_dim = umma_m >> 4`` placed at bit 24 only
+    # coincides with the real m_dim field (bits [27,29), holding M >> 7) for
+    # 128 -> bit 27 and 256 -> bit 28. For 64 it yields 4 << 24, i.e. bit 26,
+    # which is _BIT_SFA_LAYOUT: the descriptor would silently select
+    # SFA_128dp_Unique and discard the caller's sfa_layout argument.
+    assert umma_m in (128, 256), f"Unsupported UMMA_M={umma_m}"
     assert umma_k in _K_SIZE_FIELD, f"Unsupported UMMA_K={umma_k}"
     assert 0 <= sfa_layout < 2
 
@@ -108,8 +113,15 @@ def compute_idesc(
     idesc = Int32(static_base)
     sfa_top = Int32(sfa_tmem_addr_i32) & Int32(0xC0000000)
     sfb_top = Int32(sfb_tmem_addr_i32) & Int32(0xC0000000)
-    idesc = idesc | (sfa_top >> Int32(30 - _BIT_A_SF_ID))
-    idesc = idesc | (sfb_top >> Int32(30 - _BIT_B_SF_ID))
+    # Mask after shifting. ``Int32`` is signed, so ``>>`` is an arithmetic
+    # shift: a TMEM address with bit 31 set sign-extends and would leave stray
+    # high bits set. Bit 31 is _BIT_K_SIZE_LO, so for umma_k=128 (k_size 2)
+    # that flips k_size to 3 and the instruction runs the wrong K. Matches
+    # ``compute_idesc`` in
+    # cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py,
+    # which masks both fields; this copy had dropped it.
+    idesc = idesc | ((sfa_top >> Int32(30 - _BIT_A_SF_ID)) & Int32(0x3 << _BIT_A_SF_ID))
+    idesc = idesc | ((sfb_top >> Int32(30 - _BIT_B_SF_ID)) & Int32(0x3 << _BIT_B_SF_ID))
     return idesc
 
 

From 4458a875edcfff8779b95a72316222ba947f2d26 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Wed, 16 Sep 2026 02:47:25 +0000
Subject: [PATCH 06/11] [None][fix] cap MoE tactic profiling at the reachable
 token count

tune_max_num_tokens was bounded by moe_max_num_tokens and a per-expert
heuristic, but not by the largest batch a forward can present. moe_max_num_tokens
is a chunking ceiling and may be configured well above that bound, and the
profiler workspace scales with it (maxM * top_k expanded tokens) and is
allocated with a raw cudaMalloc, so an inflated ceiling asks the driver for a
workspace for a shape that can never occur.

Add default_moe_max_num_tokens (max_num_tokens * dp_size) to the existing
min(). FusedMoeTRTLLMGen already caches model_config.max_num_tokens for the
same purpose; this brings the Cutlass backend in line. Tactic selection is
unaffected because the shapes removed are ones the engine cannot produce.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py
index 8e6d295912d0..2eccd16e2b39 100755
--- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py
+++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py
@@ -370,6 +370,8 @@ def __init__(
         self.tune_max_num_tokens = min(
             self.moe_max_num_tokens,
             16384 * self.num_slots // routing_method.get_experts_per_token(),
+            # A forward can never present more tokens than this.
+            default_moe_max_num_tokens,
         )
         self.has_been_profiled = False
         self.has_been_profiled_min_latency = False

From e05415a31770d53cc1bb3000d95bdab62b15f092 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 15 Sep 2026 03:56:14 +0000
Subject: [PATCH 07/11] [None][chore] Core: deepgemm JIT warmup, postproc
 metrics, cleanup

Ports the non-Rubin-specific core changes from the internal Rubin branch.

Autotune / JIT warmup:
- Add deep_gemm_jit_warmup_buckets() and use it for the three DeepGemm
  runners whose tuning call exists to drive JIT warmup. A step-16 M grid
  over the whole range is required: the SM100 layout heuristic selects on
  ceil_div(m, block_m) with every candidate block_m a multiple of 16, and
  the last-wave-utilization tie-break keeps oscillating at high M, so a
  coarse high-M band silently skips layouts that then compile
  mid-inference (nvcc fork under the GIL, stalling every attention-DP
  rank). Fp8BlockScalingGemmRunner also gets exclude_from_cache so a warm
  disk cache cannot short-circuit the warmup.

Executor:
- Propagate decoding_iter, avg_decoded_tokens_per_iter and cached_tokens
  through PostprocWorker.Output, so a result served by a postproc worker
  reports the same metrics as the in-process path.
- When TRTLLM_WORKER_DISABLE_GC=1, also disable dynamo's post-compile
  gc.collect(1): with automatic GC off it walks every object allocated
  since the previous compile, costing seconds per recompile.

Models:
- dwdp: resolve the MoE wrapper model-agnostically, so Kimi K3's
  block_sparse_moe/routed_experts spelling is handled alongside
  DeepSeek's mlp/experts.

Misc:
- CuTe DSL compatibility shim for legacy cute.core.ThrCopy/ThrMma and
  cute.make_fragment, needed by QuACK and Transformer Engine against the
  pinned CUTLASS DSL.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tensorrt_llm/__init__.py                      | 23 +++++++++++++
 .../_torch/custom_ops/torch_custom_ops.py     | 20 +++++++----
 tensorrt_llm/_torch/modules/dwdp/setup.py     | 33 +++++++++++++++++--
 tensorrt_llm/_torch/utils.py                  | 32 ++++++++++++++++++
 tensorrt_llm/executor/postproc_worker.py      | 15 +++++++--
 tensorrt_llm/executor/result.py               |  4 +++
 tensorrt_llm/executor/worker.py               |  5 +++
 7 files changed, 121 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py
index 516e31ad2c52..814aa9874244 100644
--- a/tensorrt_llm/__init__.py
+++ b/tensorrt_llm/__init__.py
@@ -41,6 +41,29 @@
 # ImportError: libc10.so: cannot open shared object file: No such file or directory
 import torch  # noqa
 
+
+def _setup_cutlass_dsl_compatibility():
+    """Expose legacy CuTe APIs required by TensorRT-LLM and its dependencies."""
+    try:
+        import cutlass.cute as cute
+    except ImportError:
+        return
+
+    # The pinned CUTLASS DSL exposes these types at cute.*, while QuACK and
+    # Transformer Engine still resolve their annotations from cute.core.
+    # Keep this list explicit: copying the full namespace also replaces
+    # cute.core.tuple with the cutlass.cute.tuple module.
+    for name in ("ThrCopy", "ThrMma"):
+        if hasattr(cute, name) and not hasattr(cute.core, name):
+            setattr(cute.core, name, getattr(cute, name))
+
+    # CUTLASS DSL renamed make_fragment to make_rmem_tensor.
+    if hasattr(cute, "make_rmem_tensor") and not hasattr(cute, "make_fragment"):
+        cute.make_fragment = cute.make_rmem_tensor
+
+
+_setup_cutlass_dsl_compatibility()
+
 from .logger import logger
 from .version import __version__
 
diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index cf7fc63ef8b3..730eecdf945a 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -49,7 +49,7 @@
 
 from ..modules.multi_stream_utils import do_multi_stream
 from ..modules.swiglu import silu_and_mul_kernel
-from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets,
+from ..utils import (ActivationType, deep_gemm_jit_warmup_buckets,
                      fp4_scale_infer_shape,
                      get_last_power_of_2_num_tokens_buckets,
                      get_power_of_2_num_tokens_buckets,
@@ -2074,7 +2074,7 @@ def _(
     return input.new_empty((M, N), dtype=output_dtype)
 
 
-# deep_gemm_gen_tuning_buckets is imported from ..utils
+# deep_gemm_jit_warmup_buckets is imported from ..utils
 
 _USE_FUSED_FP8_QUANT_PACK = os.environ.get("TRTLLM_FUSED_FP8_QUANT_PACK",
                                            "1") == "1"
@@ -2147,7 +2147,7 @@ class fp8SwapABGemmRunner(TunableRunner):
     # every process startup.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         exclude_from_cache=True,
     )
 
@@ -2195,9 +2195,12 @@ def forward(
 class Fp8PrequantizedSwapABGemmRunner(TunableRunner):
     """Runs DeepGemm with pre-quantized FP8 activations and packed scales."""
 
+    # The same step-16 grid as the other two DeepGemm runners: a layout no
+    # bucket selects is compiled mid-inference instead, and DeepGemm forks
+    # nvcc while holding the GIL.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         constraint_specs=(ConstraintSpec(
             1, 0, lambda input_shapes: input_shapes[0][0]), ),
         exclude_from_cache=True,
@@ -2340,12 +2343,17 @@ def _(
     return input.new_empty((input.size(0), weight.size(0)), dtype=output_dtype)
 
 
-# The runner is used to trigger deepgemm jit during autotune.
+# The runner is used to trigger deepgemm jit during autotune. Only Hopper has
+# work to do: on SM100 this GEMM dispatches to TrtllmGenGemmRunner's prebuilt
+# cubins and compiles nothing.
 class Fp8BlockScalingGemmRunner(TunableRunner):
+    # Without exclude_from_cache, a warm disk cache short-circuits tuning and
+    # the JIT warmup never runs.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         tune_max_num_tokens=4096,
+        exclude_from_cache=True,
     )
 
     def get_valid_tactics(
diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py
index e8a1cb740a7a..b15f4330e69f 100644
--- a/tensorrt_llm/_torch/modules/dwdp/setup.py
+++ b/tensorrt_llm/_torch/modules/dwdp/setup.py
@@ -689,7 +689,11 @@ def fixup_moe_backends(
         # ConfigurableMoE has its own ep_size, slot_start, etc. that are used
         # in its forward path.  The backend is the inner module that holds
         # weight parameters.
-        configurable_moe = getattr(layer.mlp, "experts", None)
+        # ``moe_module`` is what _get_moe_and_experts() just resolved, so this
+        # is model-agnostic: on DeepSeek it is layer.mlp and this stays exactly
+        # equivalent to the old getattr(layer.mlp, "experts", None); on K3 it
+        # is layer.block_sparse_moe, which has no ``.mlp`` at all.
+        configurable_moe = _get_configurable_moe(moe_module)
         targets = [experts_module]
         if configurable_moe is not None and configurable_moe is not experts_module:
             targets.insert(0, configurable_moe)
@@ -1110,6 +1114,20 @@ def _get_decoder_model(model: nn.Module) -> nn.Module:
     )
 
 
+def _get_configurable_moe(moe_module: Optional[nn.Module]) -> Optional[nn.Module]:
+    """The ConfigurableMoE wrapper of an MoE module, if the model uses one.
+
+    DeepSeek calls it ``experts``; Kimi K3's ``KimiK3MoERuntime`` calls the
+    same thing ``routed_experts``.  Returns None when the module has neither.
+    """
+    if moe_module is None:
+        return None
+    experts = getattr(moe_module, "experts", None)
+    if experts is None:
+        experts = getattr(moe_module, "routed_experts", None)
+    return experts
+
+
 def _get_moe_and_experts(
     layer: nn.Module,
 ) -> Tuple[Optional[nn.Module], Optional[nn.Module]]:
@@ -1118,13 +1136,22 @@ def _get_moe_and_experts(
     The standard path for DeepSeek is:
         layer.mlp (Deepseekv3MoE) -> .experts (MoE backend)
 
+    Kimi K3 spells the same shape differently:
+        layer.block_sparse_moe (KimiK3MoERuntime) -> .routed_experts (MoE backend)
+
     Returns:
         Tuple of (moe_module, experts_module) where moe_module is the wrapper
         (e.g. Deepseekv3MoE) and experts_module is the backend (e.g.
         CutlassFusedMoE, ConfigurableMoE, etc.). Both may be None if the
         layer is not an MoE layer.
     """
+    # K3's *dense* layers do carry an ``mlp``, but this function is only ever
+    # reached for layer indices that registered themselves from
+    # ConfigurableMoE.__init__, so a dense layer never gets here and the
+    # ``mlp``-first order stays safe.
     mlp = getattr(layer, "mlp", None)
+    if mlp is None:
+        mlp = getattr(layer, "block_sparse_moe", None)
     if mlp is None:
         return None, None
 
@@ -1132,8 +1159,8 @@ def _get_moe_and_experts(
     if hasattr(mlp, "w3_w1_weight"):
         return mlp, mlp
 
-    # Standard path: mlp.experts
-    experts = getattr(mlp, "experts", None)
+    # Standard path: mlp.experts (K3: block_sparse_moe.routed_experts)
+    experts = _get_configurable_moe(mlp)
     if experts is not None:
         # Prefer the inner backend (ConfigurableMoE wraps it)
         backend = getattr(experts, "backend", None)
diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py
index 8404828a4276..f06c717d001e 100644
--- a/tensorrt_llm/_torch/utils.py
+++ b/tensorrt_llm/_torch/utils.py
@@ -422,6 +422,38 @@ def deep_gemm_gen_tuning_buckets(x: int):
     return buckets
 
 
+def deep_gemm_jit_warmup_buckets(max_m: int):
+    """M grid for the DeepGemm runners that exist only to drive JIT warmup.
+
+    DeepGemm picks its tile layout from a heuristic over M and compiles one
+    kernel per selected layout. A layout no bucket selects gets compiled
+    mid-inference instead, and DeepGemm forks nvcc while holding the GIL, so
+    that compile stalls every rank of the attention-DP group.
+
+    Step 16 is exactly the right spacing, and it is needed over the whole
+    range. In ``deepgemm/csrc/jit_kernels/heuristics/sm100.hpp`` the selection
+    depends on M only through ``ceil_div(m, block_m)``, and every candidate
+    ``block_m`` is a multiple of 16, so the choice is constant on each window
+    ``[16j + 1, 16j + 16]``: one sample per window misses nothing, and anything
+    coarser skips whole windows.
+
+    A coarse high-M band is *not* safe -- ``compare`` tie-breaks on
+    ``last_wave_util = num_blocks % num_sms``, which keeps oscillating. At
+    148 SMs and ``n=128, k=512``, ``M in [2305, 2368]`` selects a layout of its
+    own (``block_m=16``: one wave, best last-wave utilization) that a step-128
+    grid steps over, sampling 2304 and 2432.
+    """
+    # A worker whose M never leaves the low band -- a disagg GEN worker runs at
+    # batch x MTP tokens -- must not be pulled up to the 4096 floor. Measured
+    # cost of doing so: +283 s of GEN autotune, +1096 tuning-cache entries.
+    if max_m < 128:
+        return tuple(range(8, 128, 8))
+    max_m = max(min(max_m, 8192), 4096)
+    low = range(8, 128, 8)
+    dense = range(128, max_m, 16)
+    return tuple(low) + tuple(dense) + (max_m, )
+
+
 def fp4_scale_infer_shape(input_shapes: List[List[int]]) -> int:
     """Calculate the swizzled scale size for a packed FP4 input tensor."""
     unpacked_shape = list(input_shapes[0])
diff --git a/tensorrt_llm/executor/postproc_worker.py b/tensorrt_llm/executor/postproc_worker.py
index 184b9f924c8e..b8943aac24ab 100644
--- a/tensorrt_llm/executor/postproc_worker.py
+++ b/tensorrt_llm/executor/postproc_worker.py
@@ -88,6 +88,9 @@ class Output(NamedTuple):
         should_abort: bool = False
         finish_reason: Optional[str] = None
         num_generated_tokens: Optional[int] = None
+        decoding_iter: int = 0
+        avg_decoded_tokens_per_iter: Optional[float] = None
+        cached_tokens: int = 0
 
     def __init__(
         self,
@@ -237,8 +240,9 @@ async def handle_single_input(inp: PostprocWorker.Input,
                 self._records.pop(client_id, None)
                 return
             try:
-                is_final = inp.rsp.result.is_final if is_llm_response(
-                    inp.rsp) else True
+                response_result = inp.rsp.result if is_llm_response(
+                    inp.rsp) else None
+                is_final = response_result.is_final if response_result else True
                 res, metrics, perf_metrics, disaggregated_params = await self._handle_input(
                     inp)
                 record = self._records.get(client_id)
@@ -264,6 +268,13 @@ async def handle_single_input(inp: PostprocWorker.Input,
                         should_abort=should_abort,
                         finish_reason=finish_reason,
                         num_generated_tokens=num_generated_tokens,
+                        decoding_iter=getattr(response_result, "decoding_iter",
+                                              0),
+                        avg_decoded_tokens_per_iter=getattr(
+                            response_result, "avg_decoded_tokens_per_iter",
+                            None),
+                        cached_tokens=getattr(response_result, "cached_tokens",
+                                              0),
                     ))
                 if is_final:
                     self._records.pop(client_id, None)
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index 9ccc7808a99e..2143a1735c81 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -564,6 +564,10 @@ def _handle_response(self,
 
         if isinstance(response, PostprocWorker.Output):
             self._done = response.is_final
+            self.decoding_iter = response.decoding_iter
+            self.avg_decoded_tokens_per_iter = (
+                response.avg_decoded_tokens_per_iter)
+            self.cached_tokens = response.cached_tokens
             if isinstance(response.res, CompletionOutput):
                 # in streaming mode
                 self._outputs[0] = response.res
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 40cee5b65b52..a9fa5d7a6c88 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -422,6 +422,11 @@ def notify_proxy_threads_to_quit():
     # Optionally disable GC (default: not disabled)
     if os.getenv("TRTLLM_WORKER_DISABLE_GC", "0") == "1":
         gc.disable()
+        # With automatic GC off, dynamo's post-compile gc.collect(1) walks every
+        # object allocated since the previous compile (seconds per recompile).
+        if "TORCH_DYNAMO_RUN_GC_AFTER_COMPILE" not in os.environ:
+            import torch._dynamo.config
+            torch._dynamo.config.run_gc_after_compile = False
 
     with worker:
         try:

From ace591d8fab416204f0afb6494484dd11a0ea986 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 22 Sep 2026 02:15:21 +0000
Subject: [PATCH 08/11] [None][test] Port the disagg and KV-cache-v2 regression
 tests from the Rubin branch

Merge-back of the portable subset of the Rubin runtime work that main does
not already carry. PR #19040 ("MLA-backboned standalone DSpark drafter")
landed the production half of this port from the same lineage, so the DFlash
weight-load checks, the DSpark `norm_dim` kernel knob, `FUSED_MODULE_COMPONENTS`,
the external drafter `max_seq_len` bound and the DFlash position-id clamp are
all already on main, in an equal or better form. What is left is test
coverage plus one explanatory comment.

* kv_cache_manager_v2/_block_radix_tree: document why the stale-tail prune
  requires every life cycle to be pageless, mirroring
  `Block::clearStaleBlocksAfterPageUnlink` in the C++ implementation.
  Comment only, no behavior change.

* New `test_block_radix_tree_stale_prune` covers that predicate on the Python
  radix tree, including the negative control that dead tails still get pruned.

* `test_disagg_index_mapper_early_release` gains a case asserting that
  `release_index_slot` detaches every page-index view before the slot is
  reused. `is_draft` is stubbed because the guard main added at the top of
  `release_index_slot` reads it before anything else.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 .../kv_cache_manager_v2/_block_radix_tree.py  |   5 +
 .../test_disagg_index_mapper_early_release.py |  45 +++++-
 .../test_block_radix_tree_stale_prune.py      | 151 ++++++++++++++++++
 3 files changed, 200 insertions(+), 1 deletion(-)
 create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py

diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
index d1e3ec662cc7..6a0f2e7ec42c 100644
--- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
+++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
@@ -620,6 +620,11 @@ def clear_stale_blocks_after_page_unlink(
         # check if consecutive available blocks is sufficient for window_size. (TRTLLM-8802)
         # But for simplicity, we leave it for now.
         curr = start
+        # Only detach blocks with no live page in ANY life cycle: a childless tip
+        # that lost this life cycle's page may still hold live pages of other life
+        # cycles; detaching it would orphan the committed chain of an in-flight
+        # sequence. Mirrors Block::clearStaleBlocksAfterPageUnlink in
+        # cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp.
         while (
             (
                 isinstance(curr, Block)
diff --git a/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py
index 7807ef566f48..e24209dfdc91 100644
--- a/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py
+++ b/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py
@@ -19,7 +19,7 @@
 iterations of requests accumulate in-flight.
 """
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import pytest
 import torch
@@ -258,6 +258,49 @@ def test_gather_k_block_offsets_matches_beam_zero_index_select(self):
         assert torch.count_nonzero(destination[:, :, 1] != -1) == 0
         assert torch.count_nonzero(destination[:, 4, 0] != -1) == 0
 
+    def test_release_detaches_page_index_views_before_slot_reuse(self):
+        from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+        from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import (
+            IndexMapper,
+        )
+
+        manager = object.__new__(KVCacheManagerV2)
+        # release_index_slot reads is_draft before anything else; a target-side
+        # manager is the case that must detach the views.
+        manager.is_draft = False
+        manager.max_beam_width = 2
+        manager.num_pools = 3
+        index_mapper = IndexMapper(max_batch_size=1, max_beam_width=2)
+        index_mapper.add_new_sequence(1)
+        events = []
+        manager.index_mapper = MagicMock(wraps=index_mapper)
+        manager.index_mapper.remove_sequence.side_effect = lambda request_id: (
+            events.append(("remove", request_id)),
+            index_mapper.remove_sequence(request_id),
+        )[1]
+        manager._early_freed_index_requests = set()
+        kv_cache = MagicMock()
+        kv_cache.set_base_page_index_buf.side_effect = (
+            lambda beam_idx, pool_idx, value: events.append(("detach", beam_idx, pool_idx, value))
+        )
+        manager.kv_cache_map = {1: kv_cache}
+
+        manager.release_index_slot(1)
+
+        expected_calls = [
+            call(beam_idx, pool_idx, None)
+            for beam_idx in range(manager.max_beam_width)
+            for pool_idx in range(manager.num_pools)
+        ]
+        assert kv_cache.set_base_page_index_buf.call_args_list == expected_calls
+        assert events == [
+            ("detach", beam_idx, pool_idx, None)
+            for beam_idx in range(manager.max_beam_width)
+            for pool_idx in range(manager.num_pools)
+        ] + [("remove", 1)]
+        assert not _has_sequence(index_mapper, 1)
+        assert manager._early_freed_index_requests == {1}
+
 
 class TestFreeResourcesDoubleReleaseSafety:
     """Test that free_resources handles already-released IndexMapper slots."""
diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py b/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py
new file mode 100644
index 000000000000..6517e7e1b760
--- /dev/null
+++ b/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py
@@ -0,0 +1,151 @@
+# 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.
+"""Pure unit tests for the stale-block tail prune in ``clear_stale_blocks_after_page_unlink``.
+
+Regression guard for a multi-life-cycle tree: unlinking one life cycle's page
+from a childless tip must not detach the block while another life cycle still
+holds a live page there. Hybrid models (Kimi K3: MLA attention + KDA/SSM) are
+the ones that hit this, because they are the ones with more than one life cycle.
+
+The C++ mirror was fixed upstream in PR #17323
+(``Block::clearStaleBlocksAfterPageUnlink`` in
+``cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp``);
+this covers the Python implementation, which is the one selected by
+``TLLM_KV_CACHE_MANAGER_V2_BACKEND=python``.
+"""
+
+import unittest
+from collections.abc import Iterator
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, cast
+
+if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None:
+    from kv_cache_manager_v2 import TokenId
+    from kv_cache_manager_v2._block_radix_tree import Block, BlockRadixTree, ReuseScope
+    from kv_cache_manager_v2._life_cycle_registry import (
+        AttnLifeCycle,
+        LifeCycleId,
+        LifeCycleRegistry,
+    )
+else:
+    from tensorrt_llm.runtime.kv_cache_manager_v2 import TokenId
+    from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import (
+        Block,
+        BlockRadixTree,
+        ReuseScope,
+    )
+    from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import (
+        AttnLifeCycle,
+        LifeCycleId,
+        LifeCycleRegistry,
+    )
+
+
+class _TwoLifeCycles:
+    """Minimal ``LifeCycleRegistry`` stand-in: only ``size`` is reached here.
+
+    ``Block.storage`` is sized from ``num_life_cycles``, which is all the prune
+    predicate needs. Two life cycles is the smallest tree that can express
+    "this block lost life cycle 0's page but still holds life cycle 1's".
+    """
+
+    size = LifeCycleId(2)
+
+    @property
+    def ssm_life_cycle_id(self) -> None:
+        return None
+
+    def attention_life_cycles(self) -> Iterator[tuple[object, object]]:
+        return iter(())
+
+
+class _DeadPageRef:
+    """Stand-in for ``rawref.ref[CommittedPage]``: occupied slot, dead referent.
+
+    The predicate under test only compares the slot against ``None``. Returning
+    ``None`` from ``__call__`` keeps ``Block._release_pages`` on its
+    already-collected path, so teardown never dereferences a fake page.
+    """
+
+    def __call__(self) -> None:
+        return None
+
+
+# Windowed, no sink blocks: keeps ``clear_stale_blocks_after_page_unlink`` off
+# the ``remove_subtree`` branch (that branch fires for full attention or sink
+# blocks and would drop the subtree regardless of the tail-prune predicate),
+# so the test isolates the prune loop.
+_WINDOWED_ATTN = AttnLifeCycle(window_size=64, num_sink_blocks=0)
+
+_LC_UNLINKED = LifeCycleId(0)
+_LC_OTHER = LifeCycleId(1)
+
+
+class TestStaleTailPrune(unittest.TestCase):
+    def _build_chain(self) -> "tuple[BlockRadixTree, object, Block, Block]":
+        """Root -> first -> tip, two life cycles, tokens_per_block=2."""
+        tree = BlockRadixTree(cast(LifeCycleRegistry, _TwoLifeCycles()), tokens_per_block=2)
+        root = tree.add_or_get_existing(ReuseScope())
+        first = Block([TokenId(1), TokenId(2)], root)
+        tip = Block([TokenId(3), TokenId(4)], first)
+        self.assertEqual(len(tip.storage), 2)
+        return tree, root, first, tip
+
+    def test_tip_with_live_page_in_another_life_cycle_is_kept(self) -> None:
+        tree, root, first, tip = self._build_chain()
+        # Life cycle 0's page was just unlinked; life cycle 1 still holds one.
+        tip.storage[_LC_UNLINKED] = None
+        tip.storage[_LC_OTHER] = cast(object, _DeadPageRef())
+
+        Block.clear_stale_blocks_after_page_unlink(tip, _LC_UNLINKED, _WINDOWED_ATTN)
+
+        # Detaching the tip here would orphan the committed chain of an
+        # in-flight sequence that is still using life cycle 1.
+        self.assertIn(tip.key, first.next)
+        self.assertIs(first.next[tip.key], tip)
+        self.assertIsNotNone(tip._prev())
+        self.assertIn(first.key, root.next)
+
+    def test_tip_with_no_live_page_anywhere_is_pruned(self) -> None:
+        # Negative control: the fix must not stop the prune it is supposed to
+        # allow, otherwise dead tails accumulate forever.
+        tree, root, first, tip = self._build_chain()
+        tip.storage[_LC_UNLINKED] = None
+        tip.storage[_LC_OTHER] = None
+
+        Block.clear_stale_blocks_after_page_unlink(tip, _LC_UNLINKED, _WINDOWED_ATTN)
+
+        # tip is detached, and the walk continues up through `first`, which is
+        # now itself a childless tip with no pages; the emptied root is then
+        # dropped from the tree.
+        self.assertNotIn(tip.key, first.next)
+        self.assertNotIn(first.key, root.next)
+        self.assertEqual(tree.next, {})
+
+    def test_tip_keeps_when_only_the_unlinked_life_cycle_is_empty(self) -> None:
+        # Same as the first case but with the roles of the two life cycles
+        # swapped, so the test cannot pass by hard-coding an index.
+        tree, root, first, tip = self._build_chain()
+        tip.storage[_LC_OTHER] = None
+        tip.storage[_LC_UNLINKED] = cast(object, _DeadPageRef())
+
+        Block.clear_stale_blocks_after_page_unlink(tip, _LC_OTHER, _WINDOWED_ATTN)
+
+        self.assertIn(tip.key, first.next)
+        self.assertIs(first.next[tip.key], tip)
+
+
+if __name__ == "__main__":
+    unittest.main()

From 27cd6cf0813a655b5bb331df740adec512029d2a Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 22 Sep 2026 02:16:41 +0000
Subject: [PATCH 09/11] [None][fix] restore the V2 draft-KV max_tokens
 derivation dropped in the rebase

On user/lizhiz/rubin-advance, configure_kv_cache_capacity re-derives
max_tokens from the final estimated budget once estimation finishes whenever
the V2 KV cache manager is in use. The rebase onto main dropped that block;
main never had it either.

Without it, a one-model speculative-decoding draft KV cache that shares the
config reads the same max_gpu_total_bytes as the target. That cap does not
scale with the draft's much smaller per-token footprint, which depends on
num_local_layers, so both managers claim the whole budget and the draft
OOMs. Deriving max_tokens restores V1's behaviour: V2's quota becomes
min(max_gpu_total_bytes, max_tokens * bytes_per_token), which picks the
layer-scaled draft budget when the draft manager reads the shared config.

build_managers already splits max_gpu_total_bytes per manager where it can,
but that split is skipped during KV cache estimation and bails out whenever
_get_target_and_draft_cache_costs cannot model the per-manager costs, so
this remains the backstop for those paths.

Adapted to main's structure: the max_tokens block already restores an
explicit user-provided value for V2, so this only fills in the derivation
when none was given, and it stays after the max_gpu_total_bytes clamp so the
derivation uses the final budget.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 820a4bcac0d7..d66c36da0e55 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1649,6 +1649,28 @@ def configure_kv_cache_capacity(self,
                 f"max_gpu_total_bytes={self._max_gpu_total_bytes_in / (GB):.2f} GiB is provided. New max memory is {kv_cache_max_memory / (GB):.2f} GiB"
             )
 
+        if self._is_kv_cache_manager_v2 and self._max_kv_tokens_in is None:
+            # The max_tokens block above restores the user's value, which is
+            # None here, so V2 would size purely from max_gpu_total_bytes. That
+            # cap is a byte budget for the TARGET's per-token footprint. A
+            # one-model speculative-decoding draft manager reading the same
+            # config has a much smaller per-token footprint (it scales with
+            # num_local_layers), so the same byte cap lets it claim the whole
+            # budget a second time -> OOM. build_managers splits
+            # max_gpu_total_bytes per manager when it can, but that split is
+            # skipped during estimation and bails out whenever
+            # _get_target_and_draft_cache_costs cannot model the costs, and
+            # those are exactly the paths this backstops.
+            #
+            # Deriving max_tokens from the FINAL budget (after the
+            # max_gpu_total_bytes clamp just above) mirrors V1: V2's quota
+            # becomes min(max_gpu_total_bytes, max_tokens * bytes_per_token),
+            # which is a no-op for the target and picks the layer-scaled budget
+            # for the draft. Hence the placement here rather than in the
+            # max_tokens block, which runs before that clamp.
+            self._kv_cache_config.max_tokens = (self._get_kv_size_per_token(
+            ).tokens_for_budget(kv_cache_max_memory))
+
         logger.info(
             f"Estimated max memory in KV cache : {kv_cache_max_memory / (GB):.2f} GiB"
         )

From ffdf3f363aa8f427f8bb89f71905b27e88d27166 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:17:32 +0000
Subject: [PATCH 10/11] [None][test] restore the MegaMoE local-device bootstrap
 test

The MegaMoE CuteDSL bootstrap in this PR binds the CUDA device from
local_mpi_comm().Get_rank() % device_count rather than the global MPI rank,
so a multi-node run stops driving every rank at device 0. That commit could
not be exercised here -- no GPU and no multi-rank launcher -- which leaves
this monkeypatched unit test as its only verification: it drives
_maybe_init_torch_dist_under_mpi with 8 ranks over 4 visible devices and a
local rank of 5, and asserts set_device(1) and device_id="cuda:1".

It was dropped while pruning unrelated tests from this branch; the fix it
covers is still here, so put it back rather than ship the change untested.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tests/unittest/_torch/moe/test_moe_backend.py | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py
index 9cea31cf4e84..06b2be2363a0 100644
--- a/tests/unittest/_torch/moe/test_moe_backend.py
+++ b/tests/unittest/_torch/moe/test_moe_backend.py
@@ -1093,6 +1093,40 @@ def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks():
     assert moe._weights_transformed is True
 
 
+def test_megamoe_cutedsl_mpi_bootstrap_binds_local_cuda_device(monkeypatch):
+    import tensorrt_llm._utils as utils
+
+    moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl)
+    moe.ep_size = 8
+    mpi_comm = MagicMock()
+    mpi_comm.bcast.return_value = ("127.0.0.1", "29500")
+    local_mpi_comm = MagicMock()
+    local_mpi_comm.Get_rank.return_value = 5
+    init_process_group = MagicMock()
+    set_device = MagicMock()
+
+    monkeypatch.setattr(dist, "is_available", lambda: True)
+    monkeypatch.setattr(dist, "is_initialized", lambda: False)
+    monkeypatch.setattr(dist, "init_process_group", init_process_group)
+    monkeypatch.setattr(utils, "mpi_world_size", lambda: 8)
+    monkeypatch.setattr(utils, "mpi_rank", lambda: 3)
+    monkeypatch.setattr(utils, "mpi_comm", lambda: mpi_comm)
+    monkeypatch.setattr(utils, "local_mpi_comm", lambda: local_mpi_comm)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
+    monkeypatch.setattr(torch.cuda, "set_device", set_device)
+
+    moe._maybe_init_torch_dist_under_mpi()
+
+    set_device.assert_called_once_with(1)
+    init_process_group.assert_called_once_with(
+        backend="cuda:nccl,cpu:gloo",
+        rank=3,
+        world_size=8,
+        device_id=torch.device("cuda", 1),
+    )
+
+
 def test_megamoe_load_weights_invalidates_cached_deepgemm_views():
     method = W4A8MXFP4MXFP8MegaMoEDeepGemmMethod()
     hidden_size = 128

From d8f6fa6a7939a28bf3f8076d83c575be20c434c5 Mon Sep 17 00:00:00 2001
From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Date: Tue, 22 Sep 2026 07:38:05 +0000
Subject: [PATCH 11/11] [None][fix] keep the KDA dtype cast usable under meta
 init

_meta_safe_cast_dtype read param.data before calling _cast. Reading
.data dispatches aten.detach.default, which MetaInitMode rejects, so the
cast raised MetaInitException before _cast could reach its own is_meta
branch -- the guard was there and unreachable.

model_loader catches that exception, logs "Fallback to regular model
init" at INFO level and clears is_meta_init, which skips the
meta-materialisation branch and turns model_loader's model.to("cuda")
from the no-op its comment claims into a transfer of the whole model.
Measured on a GB300 4-node K3 disaggregated GSM8K run: peak allocation
254.34 GiB against a 181.16 GiB steady state, i.e. a ~68 GiB transient,
after which a 336 MiB allocation failed on every rank with under 150 MiB
free of 276.62 GiB.

Backports 6091b98e12 from user/lizhiz/rubin-advance, adapted to this
branch's named_parameters loop. The same line is present on origin/main
at f7aeaef163, so this is a main defect that rubin-advance happens to
carry a fix for rather than a defect of this branch.

Scope, stated because it was measured rather than assumed: removing the
fallback removes this transient. It is not established to remove the
OOM. rubin-advance carries this fix, stays in meta init on all 16 ranks,
and still OOMs on the same preset at a different call site inside weight
loading -- so expert-parallel width looks like the governing pressure
there. At EP8 this preset holds 112 experts per rank per layer against
EP16's 56, and the EP16 variant reports 163.44 GB/rank against EP8's
252.59 GB.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
index a7117413c9ec..c478c42fad46 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
+++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
@@ -71,7 +71,11 @@ def _cast(tensor: torch.Tensor) -> torch.Tensor:
             continue
         child._apply(_cast)
     for name, param in module.named_parameters(recurse=False):
-        param.data = _cast(param.data)
+        # Pass the parameter, not param.data: reading .data detaches, and
+        # MetaInitMode rejects aten.detach on a meta tensor, so casting the
+        # module's own parameters under meta init raised MetaInitException
+        # before _cast ever got to its is_meta branch.
+        param.data = _cast(param)
 
 
 def _stage_state_rows(ssm_pool: torch.Tensor, slot_indices: torch.Tensor) -> torch.Tensor: