[None][chore] Make use_low_precision_moe_combine as a llm arg

examples/llm-api/quickstart_advanced.py

@@ -73,6 +73,11 @@ def add_llm_args(parser):
     parser.add_argument('--moe_ep_size', type=int, default=-1)
     parser.add_argument('--moe_tp_size', type=int, default=-1)
     parser.add_argument('--moe_cluster_size', type=int, default=-1)
+    parser.add_argument(
+        '--use_low_precision_moe_combine',
+        default=False,
+        action='store_true',
+        help='Use low precision combine in MoE (only for NVFP4 quantization)')
 
     # KV cache
     parser.add_argument('--kv_cache_dtype', type=str, default='auto')
@@ -229,7 +234,7 @@ def setup_llm(args, **kwargs):
             enable_piecewise_cuda_graph= \
                 args.use_piecewise_cuda_graph)
         if args.use_torch_compile else None,
-        moe_config=MoeConfig(backend=args.moe_backend),
+        moe_config=MoeConfig(backend=args.moe_backend, use_low_precision_moe_combine=args.use_low_precision_moe_combine),
         sampler_type=args.sampler_type,
         max_seq_len=args.max_seq_len,
         max_batch_size=args.max_batch_size,

tensorrt_llm/_torch/model_config.py

@@ -117,6 +117,8 @@ class ModelConfig(Generic[TConfig]):
     moe_backend: str = 'CUTLASS'  # options can be CUTLASS, TRTLLM
     # IF true, disables FC2+finalize fusion in CUTLASS MoE backend
     moe_disable_finalize_fusion: bool = False
+    # If true, use low precision combine in MoE operations (only for NVFP4 quantization)
+    use_low_precision_moe_combine: bool = False
 
     allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO
 

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

@@ -195,9 +195,8 @@ def __init__(
             self.use_postquant_alltoall = (os.environ.get(
                 "TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1")
                                            == "1") and qm.has_nvfp4()
-            self.use_low_precision_combine = (os.environ.get(
-                "TRTLLM_MOE_USE_LOW_PRECISION_COMBINE", "0")
-                                              == "1") and qm.has_nvfp4()
+            self.use_low_precision_combine = model_config.use_low_precision_moe_combine and qm.has_nvfp4(
+            )
 
             if self.alltoall_method_type == AlltoallMethodType.MNNVL:
                 MnnvlMemory.initialize()

tensorrt_llm/_torch/pyexecutor/config.py

@@ -60,6 +60,7 @@ class PyTorchConfig:
     moe_backend: str = 'CUTLASS'
 
     moe_disable_finalize_fusion: bool = False
+    use_low_precision_moe_combine: bool = False
 
     enable_mixed_sampler: bool = False
     """

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -307,6 +307,8 @@ def __init__(
             moe_backend=pytorch_backend_config.moe_backend,
             moe_disable_finalize_fusion=pytorch_backend_config.
             moe_disable_finalize_fusion,
+            use_low_precision_moe_combine=pytorch_backend_config.
+            use_low_precision_moe_combine,
             load_format=pytorch_backend_config.load_format,
             max_num_tokens=max_num_tokens,
             moe_max_num_tokens=pytorch_backend_config.moe_max_num_tokens,

tensorrt_llm/llmapi/llm_args.py

@@ -191,6 +191,12 @@ class MoeConfig(StrictBaseModel):
         "Disable FC2+finalize kernel fusion in CUTLASS MoE backend. Setting this to True recovers deterministic numerical behavior with top-k > 2."
     )
 
+    use_low_precision_moe_combine: bool = Field(
+        default=False,
+        description=
+        "Use low precision combine in MoE operations (only for NVFP4 quantization). When enabled, uses lower precision for combining expert outputs to improve performance."
+    )
+
     @classmethod
     def from_dict(cls, data: dict):
         return cls(**data)
@@ -2586,6 +2592,8 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig":
             moe_load_balancer=self.moe_config.load_balancer,
             attn_backend=self.attn_backend,
             moe_backend=self.moe_config.backend,
+            use_low_precision_moe_combine=self.moe_config.
+            use_low_precision_moe_combine,
             enable_mixed_sampler=self.enable_mixed_sampler,
             sampler_type=self.sampler_type,
             kv_cache_dtype=self.kv_cache_config.dtype,