diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 61240b496de7..67baec2cbe94 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/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, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index a4ce0092a0b7..aa5b64d9b652 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 82d1f7d41db9..e263c2cd85ef 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/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() diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index 9dc8f51ac2c8..c99f641e7fea 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/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 """ diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5d582d02fe4f..779bc9784820 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 971f65ce08ab..23f6235540f4 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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,