Skip to content

Commit d7cc508

Browse files
committed
Revert "enable moe/mlp fusion"
This reverts commit 6e175ce047c10302bd435d25af20f2892cb14bee.
1 parent 31f25cb commit d7cc508

File tree

1 file changed

+2
-9
lines changed

1 file changed

+2
-9
lines changed

tensorrt_llm/_torch/models/modeling_llama.py

+2-9
Original file line numberDiff line numberDiff line change
@@ -347,8 +347,8 @@ def __init__(
347347
# self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp(
348348
# )
349349
# TODO: re-enable these fusions
350-
# self.fusion_config.PRE_MOE_FUSION = False
351-
# self.fusion_config.POST_MLP_FUSION = False
350+
self.fusion_config.PRE_MOE_FUSION = False
351+
self.fusion_config.POST_MLP_FUSION = False
352352

353353
self.self_attn = Llama4Attention(
354354
model_config,
@@ -374,9 +374,6 @@ def __init__(
374374

375375
# self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp(
376376
# )
377-
self.fusion_config.PRE_MLP_FUSION = model_config.mapping.has_tp()
378-
self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp()
379-
380377
else:
381378
self.feed_forward = Llama4MoE(
382379
num_experts=config.num_local_experts,
@@ -388,10 +385,6 @@ def __init__(
388385
aux_stream=aux_stream,
389386
dtype=config.torch_dtype)
390387

391-
self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp()
392-
self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp()
393-
394-
395388
# self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp(
396389
# )
397390

0 commit comments

Comments
 (0)