Skip to content

Commit d02fadf

Browse files
Propagate disable_finalize_fusion config flag in WIDEEP MoE backend
Signed-off-by: Sergey Klevtsov <[email protected]>
1 parent 88ea2c4 commit d02fadf

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,8 @@ def __init__(
226226
f"Not available alltoall method type: {self.alltoall_method_type!r}"
227227
)
228228

229+
self.use_fused_finalize = not model_config.moe_disable_finalize_fusion
230+
229231
self._weights_created = False
230232
if not model_config.skip_create_weights_in_init:
231233
self.create_weights()
@@ -724,7 +726,7 @@ def forward_chunk(
724726
input_sf=x_sf,
725727
swizzled_input_sf=False,
726728
min_latency_mode=False,
727-
use_fused_finalize=True,
729+
use_fused_finalize=self.use_fused_finalize,
728730
tuner_num_tokens=tuner_num_tokens,
729731
tuner_top_k=tuner_top_k,
730732
)

0 commit comments

Comments
 (0)