From c4bf73708c98df010054c17c921fa9b408cf4eaf Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Mon, 17 Aug 2026 06:37:02 -0700 Subject: [PATCH 1/3] [https://nvbugs/6601578][fix] Avoid MoE multi-GPU rendezvous port race Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - .../_torch/modules/moe/test_moe_module.py | 54 +++++++++++-------- 2 files changed, 33 insertions(+), 22 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e4c505019c5e..daae6cf13d17 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -360,7 +360,6 @@ 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) diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index c418fb2f1a29..a28c9f004481 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/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 From a6ec12098524f19550aaa62bf63121dda1c4e325 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:36:38 -0700 Subject: [PATCH 2/3] [https://nvbugs/6572835][test] Remove stale single-GPU MegaMoE waiver Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index daae6cf13d17..877964dedde5 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -363,7 +363,6 @@ unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu 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) From ec9f88f89cfed08dfb7b99d03de11b6a0b898db7 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:34:27 -0700 Subject: [PATCH 3/3] [https://nvbugs/6572835][fix] Fix MegaMoE CUTLASS DSL 4.6 compatibility Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- .../mega_moe_nvfp4/dynamic_mainloop.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py index b1ce22ccb3c2..834c5a9cc576 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py +++ b/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)