inference moe expert sp padding

[acisseJZhong]

Stack from ghstack (oldest at bottom):

• fix async tp #4045
• [do not review] inference moe padding #4067
• -> inference moe expert sp padding #4080
• Experimenting MoE new sharding #3996

Summary

Pad MoE inputs to the expert sequence-parallel degree before routing, then
trim the output to its original length. This ensures equal TP shards for
uneven decode batches and removes virtual-padding bookkeeping from token
dispatchers.

Adds a TP=EP=4 vLLM test covering a three-request decode batch and verifying
bitwise decode/prefill parity.

Test Plan

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc-per-node=4     -m pytest -q     torchtitan/experiments/rl/tests/test_generator.py::test_vllm_uneven_decode_tp_padding

[tianyu-l]

The moe part of change looks good to me. But why we would ever enable SP for generator in general? (I know that we are always using expert sp)

[acisseJZhong]

i added this flag just in case we want to turn on SP for inference later. For uneven inference when dense SP on, vllm model runner will pad the global input for us to make sure it's evenly divided. Therefore I need to pass enable_sp in Compilation Config into vllm.

[tianyu-l]

Let's separate it in a different PR if it's needed (with symmetric-memory based AG / RS). Right now it doesn't seem to benefit anyone, so let's split the complexity.

[acisseJZhong]

but it we don't do this, if someone run uneven inference when dense SP on, spmd_types will error out. It's not able to turn R -> S(1) due to unevenness. We need to set enable_sp in Compilation Config for vllm model runner to pad.

[tianyu-l]

if someone run uneven inference when dense SP on

why would they

[acisseJZhong]

you can't control what users would do? why leaving SP + unevenness buggy when we can fix it?

however I admit at this point no user would try this, so doesn't matter.