File tree 1 file changed +2
-9
lines changed
tensorrt_llm/_torch/models
1 file changed +2
-9
lines changed Original file line number Diff line number Diff line change @@ -347,8 +347,8 @@ def __init__(
347
347
# self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp(
348
348
# )
349
349
# 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
352
352
353
353
self .self_attn = Llama4Attention (
354
354
model_config ,
@@ -374,9 +374,6 @@ def __init__(
374
374
375
375
# self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp(
376
376
# )
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
-
380
377
else :
381
378
self .feed_forward = Llama4MoE (
382
379
num_experts = config .num_local_experts ,
@@ -388,10 +385,6 @@ def __init__(
388
385
aux_stream = aux_stream ,
389
386
dtype = config .torch_dtype )
390
387
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
-
395
388
# self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp(
396
389
# )
397
390
You can’t perform that action at this time.
0 commit comments