diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index acdb9f39fcac..6d0f89248d47 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,33 @@ 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 +@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(dtype) + 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 818c1610500e..bd0a6881c7a9 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1109,8 +1109,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. 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.element_size() == 1 else w + for w in layer.get_expert_weights() + ] + ) layer.set_eplb_state( moe_layer_idx=layer_idx, expert_load_view=expert_load_view,