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

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

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