[Bugfix][EPLB] Move 1-byte expert tensors as uint8 views

[luoyuctl]

Purpose

Fixes #51567.

When EPLB is enabled, expert tensors whose dtype NCCL does not support make engine startup fail. MXFP4 and MXFP8 experts store block scales as torch.float8_e8m0fnu. ROCm AITER MXFP4 also views expert weights as torch.float4_e2m1fn_x2. Neither dtype exists in NCCL or RCCL.

I hit this with DeepSeek-V4.1-Flash using MXFP4 experts (Humming MoE backend), TP8 + EP + EPLB, on 8x H20. The crash happens in the first rearrange during memory profiling.

With communicator=torch_nccl:

File "vllm/v1/worker/gpu_worker.py", line 572, in determine_available_memory
File "vllm/v1/worker/gpu/eplb_utils.py", line 126, in step
File "vllm/distributed/eplb/eplb_state.py", line 580, in step
...
File "torch/distributed/distributed_c10d.py", line 4280, in all_gather
    work = group.allgather([tensor_list], [tensor], opts)
TypeError: Input tensor data type is not supported for NCCL process group: Float8_e8m0fnu

With communicator=pynccl:

File "vllm/distributed/eplb/eplb_communicator.py", line 756, in create_eplb_communicator
    return _create_pynccl()
File "vllm/distributed/eplb/eplb_communicator.py", line 693, in _create_pynccl
    raise RuntimeError(
RuntimeError: EPLB communicator 'pynccl' requested but expert weights contain unsupported dtypes: (torch.float8_e8m0fnu).

When this triggers

The bug is not tied to a specific model architecture. It depends on the dtype that the MoE backend leaves in the expert parameters after weight post-processing. Startup fails when all of these hold:

1. EPLB is enabled (--enable-eplb, with EP).

2. The experts use MX quantization (MXFP4/MXFP8 with E8M0 block scales).

3. After post-processing, the backend stores an expert tensor in a 1-byte dtype that NCCL/RCCL does not support. For example:

   • Marlin MXFP4 MoE: mxfp4_marlin_process_scales() returns scales as float8_e8m0fnu.
   • ROCm AITER MXFP4: expert weights are viewed as float4_e2m1fn_x2 (oracle/mxfp4.py).
   • Humming MXFP4 (my repro): scales are stored as float8_e8m0fnu.

   Backends that keep these tensors as uint8 are not affected.

4. The tensor is returned by get_expert_weights(), so it is registered for EPLB transport.

Converting at load time is not enough. DeepSeek-V4/V4.1 loaders already view checkpoint E8M0 scales as uint8, but the backend's kernel-format conversion writes float8_e8m0fnu back into the parameter. The fix therefore belongs where the tensors are registered for EPLB, not in individual loaders or backends.

Root cause

EPLB treats expert tensors as opaque payload. It only sends, receives, and all-gathers them, and never reduces them. However, it passes their numeric dtype straight to the communicator. On the first rearrange (the dummy all-gather during determine_available_memory), the communicator has to map that dtype to an NCCL dtype. NCCL only supports int8/uint8/int32/int64/fp16/fp32/fp64/bf16 and FP8 E4M3/E5M2, so the mapping fails:

• torch_nccl raises TypeError inside all_gather.
• pynccl rejects the dtype earlier, when the communicator is created.

Fix

EPLB only copies expert bytes between ranks. It never reduces them, so the numeric dtype does not matter to the transport. In MixtureOfExperts.set_eplb_state(), every 1-byte expert tensor is now registered as w.view(torch.uint8):

• The view keeps the shape, stride, and storage, so a rearrange writes directly into the parameter that the kernels read in the original dtype.
• Applying the view at this registration point covers FusedMoE layers, models with their own get_expert_weights() (for example DeepSeek-V4 MegaMoE), and elastic EP re-registration.
• 1-byte dtypes that NCCL already supports (int8, uint8, float8_e4m3fn) move the same bytes as before.

Test Plan

.venv/bin/python -m pytest tests/distributed/test_eplb_execute.py -k byte_dtypes -v
pre-commit run --files tests/distributed/test_eplb_execute.py vllm/model_executor/models/interfaces.py

Test Result

Unit test: the new test is parametrized over float8_e8m0fnu and float4_e2m1fn_x2. Both cases pass with this change and fail on main. pre-commit (ruff, mypy) passes.

MXFP4 deployment with EPLB (8x H20, SM90): DeepSeek-V4.1-Flash with MXFP4 experts, Humming MoE backend (HummingIndexedExperts), TP8 + EP + EPLB (sync, pynccl, 16 redundant experts, step_interval=1500). The image is vllm/vllm-openai:v0.30.0 with this PR's diff applied to interfaces.py. On this setup, startup previously failed with the tracebacks above.

• Startup: the profile rearrange passes and the server starts.

  Rearranging experts sync mode (profile)...
  Rearranged experts  (profile)  in 2.30 s.
  Application startup complete.
  

• Real rearranges: three ran during the GSM8K run below. The first took 6.64 s and the next two took 0.27 s each. Balancedness was about 0.70 before the first rearrange and 0.82 to 0.87 after.
• Accuracy across rearranges (GSM8K, all 1319 questions, 5-shot chat, concurrency 32): flex-extract 89.92% ± 0.83 (1186/1319). This matches our production deployment of the same model without EPLB (90.1%). All 8 smoke prompts were also answered correctly.
• None of the 8 ranks logged Traceback, Float8_e8m0fnu, unsupported dtypes, or TypeError.

AI assistance

This PR was prepared with AI assistance (Claude). I reviewed every changed line and ran the tests above.

[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.

[luoyuctl]

@ZJY0516 @ilmarkov @SageMoore hi, do you have time to review? thx

[ilmarkov]

Fix looks good to me. Please, add the validation results for mxfp4 model deployment with eplb to PR description.

[luoyuctl]

@ilmarkov i ran a e2e test and full GSM8K eval test , both worked. also added the MXFP4 + EPLB deployment results (DeepSeek-V4.1-Flash, Humming MXFP4 MoE, EP8 with sync EPLB on 8x H20) to the description, you can check it , thanks