[https://nvbugs/5784543][fix] Setup dist before using autotuner. (#10491)

Signed-off-by: Yuxian Qiu <[email protected]>

Author: yuxianq

tests/unittest/_torch/modules/test_fused_moe.py

@@ -24,6 +24,7 @@
                         skip_pre_hopper)
 
 from tensorrt_llm._torch.autotuner import AutoTuner, autotune
+from tensorrt_llm._torch.distributed import MPIDist, TorchDist
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import \
     CuteDslFusedMoE
@@ -44,7 +45,7 @@
 from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import \
     IS_TRITON_KERNELS_AVAILABLE
 from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
-from tensorrt_llm._utils import get_sm_version, mpi_rank
+from tensorrt_llm._utils import get_sm_version, mpi_disabled, mpi_rank
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 
@@ -104,6 +105,12 @@ def test_fused_moe(moe_backend,
 
     mapping = mapping or Mapping()
     mapping.rank = mpi_rank()
+    if mpi_disabled():
+        dist = TorchDist(mapping=mapping)
+    else:
+        dist = MPIDist(mapping=mapping)
+
+    AutoTuner.get().setup_distributed_state(mapping, dist)
 
     torch.cuda.set_device(mapping.rank)