vllm-project / vllm-project/vllm

[Bug]: gpt-oss-120b MXFP4 + --enable-expert-parallel crashes at startup: modular MoE finalize allocates padded (3072) output for an unpadded (2880) expert result

Open
#52,051 0 comments 0 reactions 0 assignees View on GitHub
quantization
Dominant language
Python
Stars
92k
Forks
22.3k
PR merge metrics
PR metrics pending

Description

**Before submitting a new issue...**

- [x] I have searched for existing and past issues, at four breadths: the symptom
(padded/unpadded shape mismatch in MoE finalize), the subsystem (`gpt-oss` MXFP4
expert-parallel MoE), the faulting file/class (`modular_kernel.py`,
`FusedMoEPrepareAndFinalize`), and the exact symbols (`has_unpadded_output`,
`hidden_dim_unpadded`). The closest existing item is PR #44452, which restructures
this same call site but never mentions MXFP4 padding or DP+EP; #42118 is the same
file/class with a different defect (unswizzled mxfp8 scales).

---

### Your current environment

The output of python collect_env.py

```text
`collect_env.py` output is not available: this was found by an automated CI
triage agent from a batch job log, and the container is not retained after the
run, so the script could not be executed retroactively. The relevant values
were taken verbatim from the engine's own startup log instead:

vLLM version : 0.26.1rc1.dev602+g65b7662d3
PyTorch : 2.x + CUDA 13.0 wheel (from the vLLM nightly container built on
this commit; exact torch build string not printed in the log)
Python : 3.12
CUDA : 13.0
NCCL : 2.30.7 in the plain nightly container; 2.29.7 in the container
that additionally installs FlashInfer nightlies (pynccl prints
"vLLM is using nccl==2.30.7" and "...==2.29.7" respectively).
The crash is byte-identical at both NCCL versions.
FlashInfer : 0.6.18.dev20260811 (flashinfer-python / -cubin / -jit-cache+cu130
nightlies)
GPU : 8 x NVIDIA B200 (sm100), single node
Model : openai/gpt-oss-120b, MXFP4 checkpoint (hidden_size = 2880)
```

Selected engine config echoed by vLLM at startup:

```text
dtype=torch.bfloat16, max_seq_len=4096, tensor_parallel_size=1,
pipeline_parallel_size=1, data_parallel_size=8, quantization=gpt_oss_mxfp4,
enforce_eager=False, load_format=runai_streamer
```

Two log lines matter for backend selection:

```text
INFO [mxfp4.py:528] Using 'FLASHINFER_TRTLLM_MXFP4_BF16' Mxfp4 MoE backend.
INFO [mxfp4.py:1819] Using MoEPrepareAndFinalizeNaiveDPEPModular
INFO [mxfp4.py:1820] Using TrtLlmMxfp4ExpertsModular
```

The naive DP/EP all-to-all path is selected under these plain default flags:

```text
INFO [cuda_communicator.py:203] Using AgRsAll2AllManager all2all manager.
```

Note on DeepEP: one of the two containers additionally logs
`DeepEP v2 requires NCCL >= 2.30.4 but found 2.29.7`, but **that is not what puts
this configuration on the naive path.** The other container ships NCCL 2.30.7,
prints no DeepEP-unavailable message at all, and still selects
`MoEPrepareAndFinalizeNaiveDPEPModular` and still crashes identically. The naive
DP/EP `prepare_finalize` is simply what `--enable-expert-parallel` resolves to
here under default flags, so raising NCCL does not avoid the bug.

### 🐛 Describe the bug

`gpt-oss-120b` (MXFP4) fails deterministically during V1 startup memory
profiling when expert parallelism is enabled on Blackwell, before the server ever
becomes ready. The same build serves the model fine without
`--enable-expert-parallel`.

#### Reproduce

On a single 8xB200 node, with the `FLASHINFER_TRTLLM_MXFP4_BF16` MoE backend
selected automatically (no MoE-backend or activation override):

```console
vllm serve openai/gpt-oss-120b \
--tensor-parallel-size 1 \
--data-parallel-size 8 \
--enable-expert-parallel \
--max-model-len 4096
```

The failing ingredients are: `hidden_size` that is **not** a multiple of 256
(gpt-oss-120b is 2880), the `FLASHINFER_TRTLLM_MXFP4_BF16` backend, and a
modular `prepare_finalize` that copies into a caller-provided buffer — here
`MoEPrepareAndFinalizeNaiveDPEPModular`, which is what `--enable-expert-parallel`
resolves to here under default flags (it is reached at NCCL 2.30.7 as well as
2.29.7, so it is not a consequence of DeepEP being unavailable). Removing
`--enable-expert-parallel` avoids it; so does
forcing `--quantization-config.moe.activation mxfp8`, but that changes activation
precision (see "Why the mxfp8 workaround is not a fix" below).

#### Actual behavior

Every DP rank aborts in `determine_available_memory` → `profile_run` →
`_dummy_run` → `gpt_oss.py forward` → `torch.ops.vllm.moe_forward`:

```
RuntimeError: The size of tensor a (3072) must match the size of tensor b (2880)
at non-singleton dimension 1
File ".../vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py",
line 207, in finalize
output.copy_(get_ep_group().combine(...))
```

3072 is `round_up(2880, 256)`; 2880 is the model's true `hidden_size`.

#### Expected behavior

The engine starts and serves, as it does with EP disabled.

#### Root cause

The modular fused-MoE kernel allocates the finalize destination at the **padded**
hidden dim, but this backend produces an **unpadded** expert output:

1. `moe_runner._maybe_pad_hidden_states` pads hidden states 2880 → 3072, because
`GptOssMxfp4MoEMethod.skip_forward_padding` is `True` only for
`FLASHINFER_TRTLLM_MXFP4_MXFP8` — so `..._MXFP4_BF16` gets padding.
2. `FusedMoEKernelModularImpl.apply` allocates
`output = torch.empty_like(hidden_states)` → width **3072**.
3. `GptOssMxfp4MoEMethod.has_unpadded_output` is `True` for `..._MXFP4_BF16`, so
`TrtLlmMxfp4ExpertsModular.workspace_shapes` returns
`(M, hidden_dim_unpadded)` → `fused_out` is **2880** wide.
`AgRsAll2AllManager.combine` reduce-scatters on dim 0 only, so the trailing dim
stays 2880.
4. `finalize` is therefore asked to copy 2880 into 3072 → the `RuntimeError`.

`FLASHINFER_TRTLLM_MXFP4_BF16` appears to be the only backend that sets padding
**on** and unpadded output **on** at the same time, which is why nothing else
trips this. `..._MXFP4_MXFP8` sets both `skip_forward_padding` and
`has_unpadded_output`, so it never mismatches. The monolithic
prepare/finalize sibling is immune because it *returns* its own tensor rather
than copying into a caller buffer.

Notably, `has_unpadded_output` / `hidden_dim_unpadded` do not appear anywhere in
`modular_kernel.py` or under `fused_moe/prepare_finalize/`: the runner learned the
padded/unpadded distinction, but the modular kernel's output allocation did not.
There is even an existing `# TODO(bnell): move to MK/expert_class?` directly above
the `has_unpadded_output` property in `mxfp4.py`.

#### The copy is not the real defect — the allocation is

Making only the `finalize` copy width-tolerant fixes the `RuntimeError` and then
exposes a second, deeper failure in the compiled path:

```
File ".../torch/_inductor/output_code.py", line 763, in __call__
return self.current_callable(inputs)
File ".../inductor_cache/fi/cfi3kg....py", line 994, in call
assert_size_stride(buf8, (s72, 2880), (2880, 1),
'torch.ops.vllm.moe_forward.default')
AssertionError: expected size 8192==8192, stride 3072==2880 at dim=0;
expected size 3072==2880, stride 1==1 at dim=1
Error in op: torch.ops.vllm.moe_forward.default
This error most often comes from a incorrect fake (aka meta) kernel for a custom op.
```

`_finalize` ends with `return output`, so the 3072-wide buffer **is** the custom
op's return value, while `moe_runner._moe_forward_fake` declares the op returns
`hidden_dim_unpadded` (2880). `_maybe_reduce_final_output` does strip the padding,
but it runs *outside* the op, so the op boundary itself must already be 2880.
The buffer must be allocated at the unpadded width, not merely copied into
tolerantly.

#### Suggested fix

Recover the output width from the expert output once it is known, in
`FusedMoEKernelModularImpl.apply` (before `self._finalize(...)`):

```python
if fused_out.shape[-1] != output.shape[-1]:
output = torch.empty(
(*output.shape[:-1], fused_out.shape[-1]),
dtype=output.dtype,
device=output.device,
)
```

This appears safe: `output` reaches the experts only as `output_alias`, and
`_fused_experts` ignores that alias unless `output_alias.shape == fused_out.shape`
— which is false in exactly the mismatching case, so nothing has been written into
the discarded buffer. It also fixes every `prepare_finalize` implementation at
once rather than one `finalize()`.

I have validated this change by patching it into the library in a nightly
container and re-running the failing 8xB200 EP configuration: startup memory
profiling completes and the server serves normally, while the non-EP
configuration stays green. Happy to open a PR if maintainers prefer this shape of
fix.

PR #44452 ("Delegate finalize output-buffer allocation to prepare_finalize")
restructures this exact site and recovers the width from
`fused_expert_output.shape[-1]`, so it would likely fix this incidentally — but it
is a broad refactor, currently open and needing a rebase, and does not mention
MXFP4/padding/DP+EP. A targeted fix may be worth landing independently.

#### Why the mxfp8 workaround is not a fix

Setting `--quantization-config.moe.activation mxfp8` switches to
`FLASHINFER_TRTLLM_MXFP4_MXFP8`, which sets `skip_forward_padding = True` and so
never hits the mismatch. That masks the bug and changes activation precision;
MXFP4 with BF16 activations under DP+EP remains broken.

---

_This issue was drafted with assistance from the `opus` AI model._

Contributor guide

Open the contributing guide

Research direction

Start in FusedMoEKernelModularImpl.apply and trace the output allocation through naive_dp_ep.py finalize, mxfp4.py, and moe_runner.py. Reproduce the 8xB200 gpt-oss-120b configuration, then verify that the modular custom op returns the unpadded width and that startup profiling and serving succeed with BF16 MXFP4 expert parallelism.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
63/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.