From d31a1eb4ff777f0833a6c43eb3642e2fea753d7f Mon Sep 17 00:00:00 2001 From: zack <51604064+luoyuctl@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:40:00 +0800 Subject: [PATCH 1/2] [Bugfix][EPLB] Register E8M0 expert scales as uint8 views NCCL has no Float8_e8m0fnu dtype, so EPLB startup fails for MXFP4/MXFP8 experts. Register those tensors as same-storage uint8 views so EPLB moves the raw scale bytes while kernels keep reading the E8M0 parameter. Fixes #51567 Co-authored-by: Claude Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com> --- tests/distributed/test_eplb_execute.py | 22 ++++++++++++++++++++++ vllm/model_executor/models/interfaces.py | 10 ++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index acdb9f39fcac..6e54172ca1eb 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random +from types import SimpleNamespace import pytest import torch @@ -23,11 +24,32 @@ ensure_model_parallel_initialized, get_tp_group, ) +from vllm.model_executor.models.interfaces import MixtureOfExperts from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device +def test_set_eplb_state_registers_e8m0_as_uint8_view(): + """NCCL rejects E8M0, so EPLB must see the scale bytes as uint8 views + sharing storage with the parameter the kernel reads.""" + weight = torch.zeros(2, 8, dtype=torch.int32) + scale = torch.zeros(2, 4, dtype=torch.uint8).view(torch.float8_e8m0fnu) + layer = SimpleNamespace( + get_expert_weights=lambda: [weight, scale], + set_eplb_state=lambda **_: None, + ) + model = SimpleNamespace(moe_layers=[layer]) + empty = torch.empty(0) + MixtureOfExperts.set_eplb_state(model, empty, empty, empty) + + registered_weight, registered_scale = model.expert_weights[0] + assert registered_weight is weight + assert registered_scale.dtype == torch.uint8 + registered_scale.fill_(127) + assert torch.equal(scale.view(torch.uint8), registered_scale) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch): """Gloo receives must reach the GPU without an implicit pageable copy.""" diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 5f51ece0b322..039fb93942bf 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1067,8 +1067,14 @@ def set_eplb_state( """ self.expert_weights = [] for layer_idx, layer in enumerate(self.moe_layers): - # Register the expert weights. - self.expert_weights.append(list(layer.get_expert_weights())) + # Register the expert weights. NCCL has no E8M0 dtype, so + # E8M0 scales are moved as same-storage uint8 views. + self.expert_weights.append( + [ + w.view(torch.uint8) if w.dtype == torch.float8_e8m0fnu else w + for w in layer.get_expert_weights() + ] + ) layer.set_eplb_state( moe_layer_idx=layer_idx, expert_load_view=expert_load_view, From fd73d9660db1b2d60d9f189f0dad228ea63d653e Mon Sep 17 00:00:00 2001 From: zack <51604064+luoyuctl@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:30:08 +0800 Subject: [PATCH 2/2] [Bugfix][EPLB] Generalize to all 1-byte expert dtypes EPLB only copies bytes, so view every 1-byte expert tensor as uint8 instead of special-casing E8M0. This also covers float4_e2m1fn_x2 (ROCm AITER MXFP4), which RCCL does not support either. Co-authored-by: Claude Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com> --- tests/distributed/test_eplb_execute.py | 7 ++++--- vllm/model_executor/models/interfaces.py | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 6e54172ca1eb..6d0f89248d47 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -30,11 +30,12 @@ from .eplb_utils import distributed_run, set_env_vars_and_device -def test_set_eplb_state_registers_e8m0_as_uint8_view(): - """NCCL rejects E8M0, so EPLB must see the scale bytes as uint8 views +@pytest.mark.parametrize("dtype", [torch.float8_e8m0fnu, torch.float4_e2m1fn_x2]) +def test_set_eplb_state_registers_byte_dtypes_as_uint8_views(dtype): + """NCCL lacks these dtypes, so EPLB must see their bytes as uint8 views sharing storage with the parameter the kernel reads.""" weight = torch.zeros(2, 8, dtype=torch.int32) - scale = torch.zeros(2, 4, dtype=torch.uint8).view(torch.float8_e8m0fnu) + scale = torch.zeros(2, 4, dtype=torch.uint8).view(dtype) layer = SimpleNamespace( get_expert_weights=lambda: [weight, scale], set_eplb_state=lambda **_: None, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 039fb93942bf..c46e3121c0cb 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1067,11 +1067,11 @@ def set_eplb_state( """ self.expert_weights = [] for layer_idx, layer in enumerate(self.moe_layers): - # Register the expert weights. NCCL has no E8M0 dtype, so - # E8M0 scales are moved as same-storage uint8 views. + # Register the expert weights. EPLB only moves bytes, so 1-byte + # dtypes NCCL lacks (e.g. E8M0, FP4x2) go as uint8 views. self.expert_weights.append( [ - w.view(torch.uint8) if w.dtype == torch.float8_e8m0fnu else w + w.view(torch.uint8) if w.element_size() == 1 else w for w in layer.get_expert_weights() ] )