File tree Expand file tree Collapse file tree
aphrodite/model_executor/layers/fused_moe/runner Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -525,6 +525,8 @@ def forward(
525525
526526 # Record before `_maybe_pad_hidden_states` pads activations to match
527527 # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`
528+ # so routed output can be trimmed before
529+ # shared+routed add / latent up proj if needed.
528530 routed_hidden_dim = hidden_states .shape [- 1 ]
529531 hidden_states , og_hidden_dim = self ._maybe_pad_hidden_states (
530532 shared_experts_input ,
@@ -550,7 +552,7 @@ def forward(
550552
551553 # Extract outputs from result
552554 shared_output , fused_output = _unpack (result )
553- if hidden_dim_was_padded :
555+ if ( shared_output is not None or self . routed_output_transform is not None ) and hidden_dim_was_padded :
554556 fused_output = fused_output [..., :routed_hidden_dim ]
555557
556558 # If combine kernel already reduced fused, reduce shared to match.
You can’t perform that action at this time.
0 commit comments