[https://nvbugs/6601578][fix] Avoid MoE multi-GPU rendezvous port race

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py

@@ -296,16 +296,16 @@ def issue_dynamic_block_scaled_mma_tile(
                 }
             else:
                 nvvm_args = {
-                    "kind":
+                    "mma_kind":
                     _nvvm_raw.Tcgen05MMAKind.MXF4NVF4,
                     "cta_group":
                     _nvvm_raw.CTAGroupKind.CTA_2 if mma_tiler_mnk[0] == 256 else
                     _nvvm_raw.CTAGroupKind.CTA_1,
-                    "matrix_d":
+                    "d":
                     operand_d_ptr,
-                    "matrix_a":
+                    "a":
                     operand_a,
-                    "matrix_b":
+                    "b":
                     operand_b,
                     "idesc":
                     idesc.ir_value(),
@@ -315,7 +315,7 @@ def issue_dynamic_block_scaled_mma_tile(
                     operand_sfa_ptr,
                     "scale_b":
                     operand_sfb_ptr,
-                    "block_scale":
-                    _nvvm_raw.Tcgen05MMABlockScale.BLOCK16,
+                    "scale_vec_size":
+                    _nvvm_raw.Tcgen05MMAScaleVecSize.BLOCK16,
                 }
             nvvm.tcgen05_mma_block_scale(**nvvm_args)

tests/integration/test_lists/waives.txt

@@ -360,11 +360,9 @@ unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "MEGAMOE_CU
 unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6602176)
 unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTEDSL" SKIP (https://nvbugs/6644459)
-unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_CUTEDSL" SKIP (https://nvbugs/6601578)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=DeepSeekV3] SKIP (https://nvbugs/6644479)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_CUTEDSL" SKIP (https://nvbugs/6644478)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=IGNORE-e8_k2_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_CUTEDSL-quant=NVFP4-routing=Renormalize-slots=16-eplb=dynamic] SKIP (https://nvbugs/6644477)
-unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_CUTEDSL" SKIP (https://nvbugs/6572835)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169)
 unittest/_torch/modules/test_w4a16_nvfp4_linear.py::test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin SKIP (https://nvbugs/6581071)
 unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124)

tests/unittest/_torch/modules/moe/test_moe_module.py

@@ -30,7 +30,6 @@
 import logging
 import os
 import pickle
-import socket
 import sys
 import tempfile
 import traceback
@@ -109,7 +108,7 @@
     WFP4A16FusedMoEMethod,
     WInt4AFP8FusedMoEMethod,
 )
-from tensorrt_llm._utils import get_sm_version, mpi_rank
+from tensorrt_llm._utils import get_sm_version, mpi_comm, mpi_rank
 from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo
@@ -124,13 +123,6 @@
 )
 
 
-def _get_free_tcp_port() -> int:
-    """Return a local TCP port for MPI-worker torch.distributed rendezvous."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind(("127.0.0.1", 0))
-        return int(sock.getsockname()[1])
-
-
 def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> None:
     """MegaMoE backends resolve an EP ProcessGroup at construction time.
 
@@ -156,13 +148,40 @@ def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> No
             world_size=world_size,
         )
         return
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29561")
+    master_addr = "127.0.0.1"
+    store = None
+    if rank == 0:
+        # Bind before publishing the OS-selected port. wait_for_workers must be
+        # false because clients cannot connect until the port is broadcast.
+        store = dist.TCPStore(
+            host_name=master_addr,
+            port=0,
+            world_size=world_size,
+            is_master=True,
+            wait_for_workers=False,
+        )
+    master_port = mpi_comm().bcast(store.port if store is not None else None, root=0)
+    if store is None:
+        store = dist.TCPStore(
+            host_name=master_addr,
+            port=master_port,
+            world_size=world_size,
+            is_master=False,
+            wait_for_workers=False,
+        )
+
+    os.environ["MASTER_ADDR"] = master_addr
+    os.environ["MASTER_PORT"] = str(master_port)
     os.environ["RANK"] = str(rank)
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["LOCAL_RANK"] = str(rank)
     torch.cuda.set_device(rank)
-    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    dist.init_process_group(
+        backend="nccl",
+        store=store,
+        rank=rank,
+        world_size=world_size,
+    )
 
 
 def _create_mapping_for_parallel_mode(world_size, parallel_mode):
@@ -756,17 +775,11 @@ def run_forward():
 # ---------------------------------------------------------------------------
 
 
-def _moe_init_worker(custom_paths, master_port):
+def _moe_init_worker(custom_paths):
     # Align worker sys.path with the main process for submodule import.
     for custom_path in custom_paths:
         if custom_path.endswith("tests/unittest") and custom_path not in sys.path:
             sys.path.append(custom_path)
-    # Force the local loopback rendezvous: master_port is probed on 127.0.0.1
-    # (see _get_free_tcp_port), so MASTER_ADDR must match it and not an inherited
-    # cluster hostname, otherwise the loopback-probed port may bind a different
-    # interface. All workers share one node here, so loopback is always reachable.
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(master_port)
 
 
 def _reset_moe_comm_state():
@@ -823,10 +836,9 @@ def moe_multi_gpu_executor():
     mpi_pool_executor users, test_moe_a2a / test_autotuner). world_size is 4.
     """
     world_size = 4
-    master_port = _get_free_tcp_port()
     with MPIPoolExecutor(
         initializer=_moe_init_worker,
-        initargs=(sys.path, master_port),
+        initargs=(sys.path,),
         max_workers=world_size,
     ) as executor:
         yield executor