diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 7126756032cc..68e73963bb1b 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') @@ -236,7 +241,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 0703a4420072..f8c84781b3cb 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -133,6 +133,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 fd018023e4dc..f5f7b02fd52b 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 @@ -193,8 +193,7 @@ def __init__( if self.enable_alltoall: self.use_postquant_alltoall = (os.environ.get( "TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1") - self.use_low_precision_combine = (os.environ.get( - "TRTLLM_MOE_USE_LOW_PRECISION_COMBINE", "0") == "1") + self.use_low_precision_combine = model_config.use_low_precision_moe_combine 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 25edf9dc5d5b..1060a6bff461 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -62,6 +62,7 @@ class PyTorchConfig: moe_backend: str = 'CUTLASS' moe_disable_finalize_fusion: bool = False + use_low_precision_moe_combine: bool = False sampler_type: SamplerType = SamplerType.auto """ diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index eb3618dabb41..e3d9cfc54107 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -303,7 +303,9 @@ def _load_and_validate_config( attn_backend=self.pytorch_backend_config.attn_backend, moe_backend=self.pytorch_backend_config.moe_backend, moe_disable_finalize_fusion=self.pytorch_backend_config. - moe_disable_finalize_fusion) + moe_disable_finalize_fusion, + use_low_precision_moe_combine=self.pytorch_backend_config. + use_low_precision_moe_combine) validate_and_set_kv_cache_quant( config, self.pytorch_backend_config.kv_cache_dtype) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5a05ee741f3e..9a7faede4cee 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -192,6 +192,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) @@ -2614,6 +2620,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, sampler_type=self.sampler_type, kv_cache_dtype=self.kv_cache_config.dtype, mamba_ssm_cache_dtype=self.kv_cache_config.mamba_ssm_cache_dtype,