[Bugfix][MoE] Unpad routed output before shared expert add [Fixes #35949]

[netanel-haber]

Fixes https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4.

FI TRTLLM NVFP4 MoE can pad the routed hidden dim, e.g. 2688 -> 2816, via align_trtllm_fp4_moe_hidden_dim_for_fi.

Before #35949, FusedMoE returned routed and shared outputs separately. The routed output was truncated back to the original hidden dim before model code added the shared expert output, so the world looked like:

routed kernel output: [tokens, 2816] -> truncate -> [tokens, 2688]
shared output:        [tokens, 2688]
add:                  [tokens, 2688] + [tokens, 2688]

#35949 moved the shared/routed add into MoERunner. That changed the order to add first and truncate later:

routed kernel output: [tokens, 2816]
shared output:        [tokens, 2688]
add:                  [tokens, 2816] + [tokens, 2688]

Dynamo catches this during fake tensor tracing as a shape mismatch.

This PR records the routed hidden dim before _maybe_pad_hidden_states() and trims the fused routed output back to that dim before shared expert addition.

DailyOmni is on par for nano-v3-omni before and after this pr.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request introduces logic to handle hidden dimension padding in the Fused MoE runner. It records the original hidden dimension before potential padding and ensures that the fused output is sliced back to its original size if padding was applied. I have no feedback to provide as there are no review comments to evaluate.

[tomeras91]

Thanks! Approved

Cc @robertgshaw2-redhat

BTW @netanel-haber - do you know how this works with latentMoE (regardless of padding)? Is the routed hidden states are added to the shared hidden states only after applying the latent up proj to match hidden dims again?

[netanel-haber]

Thanks! Approved

Cc @robertgshaw2-redhat

BTW @netanel-haber - do you know how this works with latentMoE (regardless of padding)? Is the routed hidden states are added to the shared hidden states only after applying the latent up proj to match hidden dims again?

See image re latent:

[bnellnm]

I think this change might have broken lora/test_gptoss_tp.py::test_gpt_oss_lora_tp2[True-False] again.

[netanel-haber]

Re gptoss test breakages,
See: #40865