File tree Expand file tree Collapse file tree 1 file changed +9
-10
lines changed Expand file tree Collapse file tree 1 file changed +9
-10
lines changed Original file line number Diff line number Diff line change @@ -337,16 +337,13 @@ def forward_deepep(
337
337
topk_weights ,
338
338
forward_mode = forward_mode ,
339
339
)
340
- final_hidden_states = (
341
- self .experts (
342
- hidden_states = hidden_states ,
343
- reorder_topk_ids = reorder_topk_ids ,
344
- seg_indptr = seg_indptr ,
345
- masked_m = masked_m ,
346
- expected_m = expected_m ,
347
- forward_mode = forward_mode ,
348
- )
349
- * self .routed_scaling_factor
340
+ final_hidden_states = self .experts (
341
+ hidden_states = hidden_states ,
342
+ reorder_topk_ids = reorder_topk_ids ,
343
+ seg_indptr = seg_indptr ,
344
+ masked_m = masked_m ,
345
+ expected_m = expected_m ,
346
+ forward_mode = forward_mode ,
350
347
)
351
348
if self .ep_size > 1 :
352
349
final_hidden_states = self .deepep_dispatcher .combine (
@@ -355,6 +352,8 @@ def forward_deepep(
355
352
topk_weights ,
356
353
forward_mode ,
357
354
)
355
+ final_hidden_states *= self .routed_scaling_factor
356
+
358
357
if shared_output is not None :
359
358
final_hidden_states = final_hidden_states + shared_output
360
359
You can’t perform that action at this time.
0 commit comments