[Bugfix][MoE] Convert FlashInfer BF16 weights in place (#54699)

Signed-off-by: Yuchen Wang <93700456+yuchenwang3@users.noreply.github.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Co-authored-by: Dakai An <dakaian108@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>

tests/kernels/moe/test_flashinfer.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing
 from dataclasses import dataclass
 
 import pytest
@@ -433,6 +434,263 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
 
     assert w13_converted.shape[0] == num_experts
     assert w2_converted.shape[0] == num_experts
+    assert w13_converted.data_ptr() == w13.data_ptr()
+    assert w2_converted.data_ptr() == w2.data_ptr()
+
+
+@pytest.mark.parametrize("is_gated_act_gemm", [True, False])
+def test_convert_moe_weights_to_flashinfer_trtllm_block_layout_values(
+    is_gated_act_gemm,
+):
+    from flashinfer.fused_moe.core import (
+        _maybe_get_cached_w3_w1_permute_indices,
+        get_w2_permute_indices_with_cache,
+    )
+
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+
+    num_experts, intermediate, hidden = 2, 256, 256
+    w13_multiplier = 2 if is_gated_act_gemm else 1
+    w13 = torch.randn(
+        (num_experts, w13_multiplier * intermediate, hidden),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    w2 = torch.randn(
+        (num_experts, hidden, intermediate),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    def _reference_block_layout(
+        weight: torch.Tensor,
+        is_w13: bool,
+        cache: dict[torch.Size, torch.Tensor],
+    ) -> torch.Tensor:
+        outputs = []
+        for expert in weight:
+            expert_uint8 = expert.view(torch.uint8)
+            if is_w13:
+                indices = _maybe_get_cached_w3_w1_permute_indices(
+                    cache,
+                    expert_uint8,
+                    128,
+                    is_gated_act_gemm=is_gated_act_gemm,
+                )
+                if is_gated_act_gemm:
+                    indices = (indices + expert_uint8.shape[0] // 2) % (
+                        expert_uint8.shape[0]
+                    )
+            else:
+                indices = get_w2_permute_indices_with_cache(
+                    cache,
+                    expert_uint8,
+                    128,
+                )
+            rows, cols = expert_uint8.shape
+            blocks = expert_uint8.view(rows, cols // 128, 128).permute(1, 0, 2)
+            outputs.append(torch.index_select(blocks, 1, indices.to(weight.device)))
+        return torch.stack(outputs).view(torch.bfloat16)
+
+    reference_cache: dict[torch.Size, torch.Tensor] = {}
+    expected_w13 = _reference_block_layout(w13, is_w13=True, cache=reference_cache)
+    expected_w2 = _reference_block_layout(w2, is_w13=False, cache=reference_cache)
+    w13_ptr = w13.data_ptr()
+    w2_ptr = w2.data_ptr()
+
+    actual_w13, actual_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        {},
+        w13,
+        w2,
+        is_gated_act_gemm=is_gated_act_gemm,
+    )
+
+    assert actual_w13.data_ptr() == w13_ptr
+    assert actual_w2.data_ptr() == w2_ptr
+    assert torch.equal(actual_w13, expected_w13)
+    assert torch.equal(actual_w2, expected_w2)
+
+
+def _make_unquantized_flashinfer_test_layer(
+    monkeypatch, intermediate, is_gated, *, device="cuda"
+):
+    from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+        UnquantizedMoeBackend,
+    )
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    moe_config = make_dummy_moe_config(
+        num_experts=2,
+        hidden_dim=256,
+        intermediate_size=intermediate,
+        activation=MoEActivation.SILU if is_gated else MoEActivation.RELU2_NO_MUL,
+    )
+    moe_config.intermediate_size_per_partition = (intermediate + 127) // 128 * 128
+    method = object.__new__(UnquantizedFusedMoEMethod)
+    method.moe = moe_config
+    method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
+    method.moe_kernel = None
+    monkeypatch.setattr(
+        method,
+        "_init_moe_kernel",
+        lambda _: setattr(method, "moe_kernel", object()),
+    )
+
+    layer = torch.nn.Module()
+    layer.moe_config = moe_config
+    with torch.device(device):
+        method.create_weights(
+            layer,
+            num_experts=2,
+            hidden_size=256,
+            intermediate_size_per_partition=moe_config.intermediate_size_per_partition,
+            params_dtype=torch.bfloat16,
+        )
+    return method, layer
+
+
+@pytest.mark.parametrize("intermediate", [192, 256])
+@pytest.mark.parametrize("is_gated", [True, False])
+def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(
+    monkeypatch, intermediate, is_gated
+):
+    """Each raw reload clears padding left dirty by prior in-place conversion."""
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+
+    method, layer = _make_unquantized_flashinfer_test_layer(
+        monkeypatch, intermediate, is_gated
+    )
+    w13_shape, w2_shape = layer.w13_weight.shape, layer.w2_weight.shape
+    padded = w2_shape[-1]
+    w13_ptr = layer.w13_weight.data_ptr()
+    w2_ptr = layer.w2_weight.data_ptr()
+
+    for _ in range(2):
+        reloaded_w13 = torch.randn_like(layer.w13_weight)
+        reloaded_w2 = torch.randn_like(layer.w2_weight)
+        reloaded_w13[:, intermediate:padded].zero_()
+        if is_gated:
+            reloaded_w13[:, padded + intermediate :].zero_()
+        reloaded_w2[:, :, intermediate:].zero_()
+        expected_w13, expected_w2 = (
+            convert_moe_weights_to_flashinfer_trtllm_block_layout(
+                {},
+                reloaded_w13.clone(),
+                reloaded_w2.clone(),
+                is_gated_act_gemm=is_gated,
+            )
+        )
+
+        # The loader writes only logical checkpoint slices, not padding.
+        layer.w13_weight.fill_(float("nan"))
+        layer.w2_weight.fill_(float("nan"))
+        layer.w13_weight[:, :intermediate].copy_(reloaded_w13[:, :intermediate])
+        if is_gated:
+            layer.w13_weight[:, padded : padded + intermediate].copy_(
+                reloaded_w13[:, padded : padded + intermediate]
+            )
+        layer.w2_weight[:, :, :intermediate].copy_(reloaded_w2[:, :, :intermediate])
+        method.process_weights_after_loading(layer)
+
+        assert layer.w13_weight.shape == w13_shape
+        assert layer.w2_weight.shape == w2_shape
+        assert layer.w13_weight.data_ptr() == w13_ptr
+        assert layer.w2_weight.data_ptr() == w2_ptr
+        kernel_w13, kernel_w2 = method._kernel_weights(layer)
+        assert torch.equal(kernel_w13, expected_w13)
+        assert torch.equal(kernel_w2, expected_w2)
+
+
+def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index):
+    from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader
+    from vllm.model_executor.utils import weights_already_processed
+
+    torch.accelerator.set_device_index(device_index)
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        method, layer = _make_unquantized_flashinfer_test_layer(
+            monkeypatch, 192, is_gated, device="meta"
+        )
+        loader = object.__new__(IpcModelLoader)
+        loader.mode = mode
+        loader._apply_entries(layer, entries, {}, device_index)
+        pointers = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr())
+
+        # Like ipc_cache: a fresh method, only tensor metadata, no _setup_kernel.
+        with weights_already_processed():
+            method.process_weights_after_loading(layer)
+        assert method.moe_kernel is not None
+        actual = method._kernel_weights(layer)
+        for weight, reference, pointer in zip(actual, expected, pointers):
+            assert weight.data_ptr() == pointer
+            assert torch.equal(weight.cpu(), reference)
+        # Cross-process addresses need not match; writes prove sharing/isolation.
+        with torch.no_grad():
+            for weight in actual:
+                weight.fill_(1.0)
+        torch.accelerator.synchronize()
+
+
+@pytest.mark.parametrize("is_gated", [True, False])
+@pytest.mark.parametrize("cache_ndim", [3, 4])
+@pytest.mark.parametrize("mode", ["copy", "zero_copy"])
+def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(
+    monkeypatch, is_gated, cache_ndim, mode
+):
+    """A spawned IPC consumer restores views without transient method state."""
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+    from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry
+
+    method, layer = _make_unquantized_flashinfer_test_layer(
+        monkeypatch, 192, is_gated, device="meta"
+    )
+    packed = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        {},
+        torch.randn_like(layer.w13_weight, device="cuda"),
+        torch.randn_like(layer.w2_weight, device="cuda"),
+        is_gated_act_gemm=is_gated,
+    )
+    expected = [weight.cpu() for weight in packed]
+    entries = {}
+    for name, weight in zip(("w13_weight", "w2_weight"), packed):
+        cached = weight.view_as(getattr(layer, name)) if cache_ndim == 3 else weight
+        entries[name] = TensorEntry.from_tensor(cached, kind="param")
+    consumer = multiprocessing.get_context("spawn").Process(
+        target=_check_flashinfer_ipc_weights,
+        args=(
+            entries,
+            expected,
+            is_gated,
+            mode,
+            torch.accelerator.current_device_index(),
+        ),
+    )
+    consumer.start()
+    try:
+        # Keep producer allocations alive until the consumer releases its views.
+        consumer.join(timeout=180)
+        assert not consumer.is_alive(), "IPC consumer timed out"
+        assert consumer.exitcode == 0, "IPC consumer failed; see child traceback"
+        for source, reference in zip(packed, expected):
+            expected_source = (
+                torch.ones_like(reference) if mode == "zero_copy" else reference
+            )
+            assert torch.equal(source.cpu(), expected_source)
+    finally:
+        if consumer.is_alive():
+            consumer.terminate()
+            consumer.join(timeout=10)
+            if consumer.is_alive():
+                consumer.kill()
+                consumer.join(timeout=10)
+        consumer.close()
 
 
 @pytest.mark.parametrize(

tests/kernels/moe/test_moe_weight_loading_padded.py

@@ -12,7 +12,14 @@ correctly handles this mismatch.
 import pytest
 import torch
 
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+    UnquantizedMoeBackend,
+)
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+    UnquantizedFusedMoEMethod,
+)
 
 from .utils import make_dummy_moe_config
 
@@ -249,6 +256,57 @@ class TestNarrowExpertDataForPadding:
 class TestWeightLoadingWithPaddedHiddenSize:
     """Integration-style tests that simulate padded weight loading."""
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("tp_rank", [0, 1])
+    def test_load_w2_chunks_noncontiguous_cpu_source_to_padded_cuda(
+        self, dtype, tp_rank
+    ):
+        hidden = 4096
+        intermediate = 1024
+        loaded_weight = (
+            torch.arange(hidden * intermediate * 2)
+            .remainder(251)
+            .to(dtype)
+            .reshape(hidden, intermediate * 2)
+        )
+        tp_source = loaded_weight[
+            :, intermediate * tp_rank : intermediate * (tp_rank + 1)
+        ]
+        expert_data_full = torch.zeros(
+            hidden + 8, intermediate + 8, device="cuda", dtype=dtype
+        )
+        destination = expert_data_full[:hidden, :intermediate]
+
+        assert not tp_source.is_contiguous()
+        assert tp_source.nbytes > 1 << 20
+        assert not destination.is_contiguous()
+
+        experts = object.__new__(RoutedExperts)
+        torch.nn.Module.__init__(experts)
+        experts.moe_config = make_dummy_moe_config()
+        experts.moe_config.moe_parallel_config.tp_size = 2
+
+        torch.accelerator.synchronize()
+        allocated_before = torch.accelerator.memory_allocated()
+        torch.accelerator.reset_peak_memory_stats()
+        experts._load_w2(
+            expert_data=expert_data_full,
+            shard_dim=1,
+            loaded_weight=loaded_weight,
+            tp_rank=tp_rank,
+        )
+        torch.accelerator.synchronize()
+        peak_extra = torch.accelerator.max_memory_allocated() - allocated_before
+
+        # A single copy into the strided CUDA view allocates a full-shard
+        # temporary. Correct values alone would not catch that regression.
+        assert peak_extra < tp_source.nbytes
+
+        torch.testing.assert_close(destination.cpu(), tp_source)
+        assert torch.count_nonzero(expert_data_full[hidden:, :]) == 0
+        assert torch.count_nonzero(expert_data_full[:hidden, intermediate:]) == 0
+
     def test_load_w2_with_padding(self):
         """Simulate loading w2 weights when hidden_size is padded."""
         padded_hidden = 3072
@@ -364,6 +422,74 @@ class TestWeightLoadingWithPaddedHiddenSize:
         )
 
 
+class TestUnquantizedTrtLlmPrePadding:
+    @staticmethod
+    def _make_method(
+        intermediate: int,
+        backend: UnquantizedMoeBackend = UnquantizedMoeBackend.FLASHINFER_TRTLLM,
+    ) -> UnquantizedFusedMoEMethod:
+        moe_config = make_dummy_moe_config(
+            num_experts=2,
+            hidden_dim=64,
+            intermediate_size=intermediate,
+        )
+        method = object.__new__(UnquantizedFusedMoEMethod)
+        method.moe = moe_config
+        method.unquantized_backend = backend
+        method.moe_kernel = None
+        return method
+
+    @pytest.mark.parametrize(
+        "backend,original,expected",
+        [
+            (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1344, 1408),
+            (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1408, 1408),
+            (UnquantizedMoeBackend.TRITON, 1344, 1344),
+        ],
+    )
+    def test_rounds_intermediate_before_weight_allocation(
+        self,
+        backend: UnquantizedMoeBackend,
+        original: int,
+        expected: int,
+    ):
+        method = self._make_method(original, backend)
+
+        hidden, intermediate = method.maybe_roundup_sizes(
+            hidden_size=64,
+            intermediate_size_per_partition=original,
+            act_dtype=torch.bfloat16,
+            moe_parallel_config=method.moe.moe_parallel_config,
+        )
+
+        assert hidden == 64
+        assert intermediate == expected
+
+    @pytest.mark.parametrize("is_act_and_mul", [True, False])
+    def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool):
+        method = self._make_method(1344)
+        if not is_act_and_mul:
+            method.moe.activation = MoEActivation.RELU2_NO_MUL
+        method.moe.intermediate_size_per_partition = 1408
+        layer = torch.nn.Module()
+
+        method.create_weights(
+            layer=layer,
+            num_experts=2,
+            hidden_size=64,
+            intermediate_size_per_partition=1408,
+            params_dtype=torch.bfloat16,
+        )
+
+        up_mult = 2 if is_act_and_mul else 1
+        assert layer.w13_weight.shape == (2, up_mult * 1408, 64)
+        assert layer.w2_weight.shape == (2, 64, 1408)
+        assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0
+        if is_act_and_mul:
+            assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0
+        assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0
+
+
 class TestLoadWeightsExpertBias:
     """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize
     all-zero per-expert `.bias` tensors for models whose experts have no bias

tests/kernels/moe/test_trtllm_bf16_moe.py

@@ -75,6 +75,10 @@ def test_trtllm_bf16_moe_modular_no_graph(
         a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
         w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
         w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+        # The FlashInfer conversion may rewrite unpadded input storage in place.
+        # Preserve the original layout for the independent torch reference.
+        reference_w1 = w1.clone()
+        reference_w2 = w2.clone()
         score = torch.randn((m, e), device="cuda", dtype=dtype)
         scores = torch.softmax(score, dim=-1, dtype=torch.float32)
         topk_weights, topk_ids = torch.topk(scores, topk)
@@ -133,8 +137,8 @@ def test_trtllm_bf16_moe_modular_no_graph(
 
         torch_output = torch_moe(
             a,
-            w1,
-            w2,
+            reference_w1,
+            reference_w2,
             score,
             topk,
             activation=MoEActivation.SILU,

tests/model_executor/model_loader/test_reload.py

@@ -5,6 +5,7 @@ import importlib.machinery
 import inspect
 import sys
 import types
+from types import SimpleNamespace
 from unittest.mock import Mock
 from weakref import WeakKeyDictionary, ref
 
@@ -467,6 +468,110 @@ def test_model_cleanup(dist_init, default_vllm_config):
     assert len(mock_info_dict) == 0
 
 
+@pytest.mark.parametrize("is_gated", [False, True])
+@pytest.mark.parametrize("has_bias", [False, True])
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("padded", [False, True])
+def test_padded_moe_reload_releases_each_layer(
+    monkeypatch, is_gated, has_bias, tp_rank, padded
+):
+    """Checkpoint-sized copies finish each layer without global finalization."""
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    hidden, intermediate, experts = 4, 3, 2
+    stored_hidden, stored_intermediate = (8, 8) if padded else (hidden, intermediate)
+    config = SimpleNamespace(
+        hidden_dim_unpadded=hidden,
+        intermediate_size_per_partition_unpadded=intermediate,
+        is_act_and_mul=is_gated,
+        has_bias=has_bias,
+        tp_rank=tp_rank,
+        moe_parallel_config=SimpleNamespace(tp_size=2),
+    )
+    model = torch.nn.ModuleList()
+    processed: list[torch.nn.Module] = []
+    for _ in range(2):
+        method = object.__new__(UnquantizedFusedMoEMethod)
+        torch.nn.Module.__init__(method)
+        method.moe = config
+        # The regression concerns streaming reload, not kernel conversion.
+        monkeypatch.setattr(method, "process_weights_after_loading", processed.append)
+        layer = object.__new__(RoutedExperts)
+        torch.nn.Module.__init__(layer)
+        layer.moe_config = config
+        layer.quant_config = None
+        layer.quant_method = method
+        layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i)
+        layer._loaded_expert_biases = set()
+        method.create_weights(
+            layer,
+            experts,
+            stored_hidden,
+            stored_intermediate,
+            torch.float32,
+            weight_loader=layer.weight_loader,
+        )
+        model.append(layer)
+
+    record_metadata_for_reloading(model)
+    original_params = [dict(layer.named_parameters()) for layer in model]
+    shards = ["w1", "w3", "w2"] if is_gated else ["w1", "w2"]
+    for cycle in range(2):
+        initialize_layerwise_reload(model)
+        for layer_index, layer in enumerate(model):
+            info = reload_layerwise.get_layerwise_info(layer)
+            inputs = []
+            expected = {
+                name: torch.full_like(p, float("nan"))
+                for name, p in original_params[layer_index].items()
+            }
+            params = dict(layer.named_parameters())
+            calls = [
+                (e, s, b)
+                for e in range(experts)
+                for s in shards
+                for b in ([False, True] if has_bias else [False])
+            ]
+            for call_index, (expert, shard, bias) in enumerate(calls):
+                name = ("w2" if shard == "w2" else "w13") + (
+                    "_bias" if bias else "_weight"
+                )
+                shape = (
+                    ((hidden,) if bias else (hidden, 2 * intermediate))
+                    if shard == "w2"
+                    else ((2 * intermediate,) if bias else (2 * intermediate, hidden))
+                )
+                weight = torch.arange(
+                    torch.Size(shape).numel(), dtype=torch.float32
+                ).reshape(shape)
+                weight = weight + 100 * (1 + call_index + cycle)
+                inputs.append(ref(weight))
+                # Direct checkpoint loading is the reference for deferred reload.
+                layer.weight_loader(expected[name], weight, name, shard, expert)
+                param = params[name]
+                param.weight_loader(param, weight, name, shard, expert)
+                del weight
+                if call_index != len(calls) - 1:
+                    assert info.can_load(), (
+                        "Layer processed before its final checkpoint shard"
+                    )
+
+            assert not info.can_load(), (
+                "Padding must not defer the layer until finalization"
+            )
+            assert not info.loaded_weights
+            assert len(processed) == cycle * len(model) + layer_index + 1
+            assert all(source() is None for source in inputs)
+            for name, original in original_params[layer_index].items():
+                assert getattr(layer, name) is original
+                # Kernel-specific tests cover padding, which is not checkpoint data.
+                mask = torch.isfinite(expected[name])
+                assert torch.equal(original[mask], expected[name][mask])
+
+
 def test_get_numel_loaded():
     param = torch.empty(10, device="meta")
     loaded_weight = torch.empty(10)

vllm/model_executor/layers/fused_moe/routed_experts.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     resolve_quant_method,
 )
+from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
@@ -566,6 +567,23 @@ class RoutedExperts(PluggableLayer):
             hidden_dim=hidden_dim,
             shard_dim=shard_dim,
         )
+        if (
+            loaded_weight.device.type == "cpu"
+            and expert_data.device.type == "cuda"
+            and not expert_data.is_contiguous()
+            and expert_data.ndim == 2
+        ):
+            # A strided CPU-to-CUDA copy can allocate a full-shard CUDA
+            # temporary. Make the TP slice contiguous on CPU and copy rows
+            # in chunks to limit the temporary used for the padded view.
+            loaded_weight = loaded_weight.contiguous()
+            num_chunks = max(cdiv(loaded_weight.nbytes, 1 << 20), 1)
+            for dst, src in zip(
+                expert_data.chunk(num_chunks, dim=0),
+                loaded_weight.chunk(num_chunks, dim=0),
+            ):
+                dst.copy_(src)
+            return
         expert_data.copy_(loaded_weight)
 
     def _load_single_value(

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

@@ -11,6 +11,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
+    FusedMoEParallelConfig,
     FusedMoEQuantConfig,
     biased_moe_quant_config,
 )
@@ -33,6 +34,7 @@ from vllm.model_executor.utils import (
     set_weight_attrs,
 )
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import round_up
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
@@ -68,22 +70,38 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
+        unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded
+        assert unpadded_intermediate is not None
+        has_intermediate_padding = (
+            intermediate_size_per_partition > unpadded_intermediate
+        )
         if self.moe.is_act_and_mul:
             w13_up_dim = 2 * intermediate_size_per_partition
         else:
             w13_up_dim = intermediate_size_per_partition
+        unpadded_hidden = self.moe.hidden_dim_unpadded
+        assert unpadded_hidden is not None
+        unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1)
         # Fused gate_up_proj (column parallel)
-        w13_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                w13_up_dim,
-                hidden_size,
-                dtype=params_dtype,
-            ),
-            requires_grad=False,
+        w13_data = torch.empty(
+            num_experts,
+            w13_up_dim,
+            hidden_size,
+            dtype=params_dtype,
         )
+        if has_intermediate_padding:
+            # The checkpoint loader only writes the unpadded slices. Zero the
+            # tails now so no full-size padding copy is needed after loading.
+            w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_()
+            if self.moe.is_act_and_mul:
+                w13_data[
+                    :,
+                    intermediate_size_per_partition + unpadded_intermediate :,
+                ].zero_()
+        w13_weight = torch.nn.Parameter(w13_data, requires_grad=False)
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
+        w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden
         if self.moe.has_bias:
             w13_bias = torch.nn.Parameter(
                 torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
@@ -91,18 +109,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
             )
             layer.register_parameter("w13_bias", w13_bias)
             set_weight_attrs(w13_bias, extra_weight_attrs)
+            w13_bias.weight_loader_numel = num_experts * unpadded_up
         # down_proj (row parallel)
-        w2_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                hidden_size,
-                intermediate_size_per_partition,
-                dtype=params_dtype,
-            ),
-            requires_grad=False,
+        w2_data = torch.empty(
+            num_experts,
+            hidden_size,
+            intermediate_size_per_partition,
+            dtype=params_dtype,
         )
+        if has_intermediate_padding:
+            w2_data[:, :, unpadded_intermediate:].zero_()
+        w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
+        w2_weight.weight_loader_numel = (
+            num_experts * unpadded_hidden * unpadded_intermediate
+        )
         if self.moe.has_bias:
             w2_bias = torch.nn.Parameter(
                 torch.zeros(num_experts, hidden_size, dtype=params_dtype),
@@ -110,6 +132,26 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
             )
             layer.register_parameter("w2_bias", w2_bias)
             set_weight_attrs(w2_bias, extra_weight_attrs)
+            w2_bias.weight_loader_numel = num_experts * unpadded_hidden
+
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
+            hidden_size,
+            intermediate_size_per_partition,
+            act_dtype,
+            moe_parallel_config,
+        )
+        if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
+            intermediate_size_per_partition = round_up(
+                intermediate_size_per_partition, 128
+            )
+        return hidden_size, intermediate_size_per_partition
 
     def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
         # Pad the weight tensor. This is an optimization on ROCm platform, which
@@ -135,6 +177,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         w13: torch.Tensor,
         w2: torch.Tensor,
     ) -> None:
+        w13_parameter_shape = w13.shape
+        w2_parameter_shape = w2.shape
+
+        if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
+            unpadded = self.moe.intermediate_size_per_partition_unpadded
+            assert unpadded is not None
+            intermediate = w2.shape[-1]
+            # Reloads overwrite only checkpoint slices; the previous in-place
+            # permutation can leave nonzero values in the raw padding slots.
+            if intermediate > unpadded:
+                w13[:, unpadded:intermediate].zero_()
+                if self.moe.is_act_and_mul:
+                    w13[:, intermediate + unpadded :].zero_()
+                w2[:, :, unpadded:].zero_()
+
         # Shuffle weights to runtime format.
         w13_new, w2_new = convert_to_unquantized_kernel_format(
             self.unquantized_backend,
@@ -142,6 +199,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
             w13_weight=w13,
             w2_weight=w2,
         )
+        if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
+            w13_new = w13_new.view(w13_parameter_shape)
+            w2_new = w2_new.view(w2_parameter_shape)
+
         # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;
         # On the first call we replace the parameter normally. On subsequent
         # calls (e.g. RL weight updates that re-trigger
@@ -170,6 +231,25 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
                 assert self.moe_kernel is not None
                 self.moe_kernel.fused_experts.process_weights_after_loading(layer)
 
+    def _kernel_weights(
+        self, layer: "RoutedExperts"
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM:
+            return layer.w13_weight, layer.w2_weight
+
+        # IPC caches restore tensor shapes/storage, not method attributes.
+        # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements.
+        def block_view(weight: torch.Tensor) -> torch.Tensor:
+            if weight.ndim == 4:
+                return weight
+            experts, rows, cols = weight.shape
+            return weight.view(experts, cols // 64, rows, 64)
+
+        return (
+            block_view(layer.w13_weight),
+            block_view(layer.w2_weight),
+        )
+
     def _init_moe_kernel(self, layer: "RoutedExperts") -> None:
         """Build the MoE kernel from the layer's current (shuffled) weights."""
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
@@ -274,10 +354,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
         assert self.moe_kernel is not None
+        w13, w2 = self._kernel_weights(layer)
         return self.moe_kernel.apply(
             hidden_states=x,
-            w1=layer.w13_weight,
-            w2=layer.w2_weight,
+            w1=w13,
+            w2=w2,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             activation=layer.activation,
@@ -315,10 +396,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
     ) -> torch.Tensor | UnfinalizedMoEOutput:
         assert self.is_monolithic
         assert self.moe_kernel is not None
+        w13, w2 = self._kernel_weights(layer)
         return self.moe_kernel.apply_monolithic(
             x,
-            layer.w13_weight,
-            layer.w2_weight,
+            w13,
+            w2,
             router_logits,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

@@ -112,8 +112,9 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Convert expert weights to FlashInfer's block layout.
 
-    This reorders W13 and W2 into the expected epilogue-tiled block layout and
-    returns the shuffled weight tensors.
+    This reorders W13 and W2 in place into the expected epilogue-tiled block
+    layout and returns views of the shuffled weight tensors. Using one expert
+    as scratch space avoids allocating another full copy of both weights.
     """
     if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
         raise ValueError(
@@ -147,51 +148,50 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
             out=out,
         )
 
-    w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape
-    w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape
-    w13_weights_shuffled_tensor = torch.empty(
-        (num_experts, w13_cols // block_k, w13_rows, block_k),
-        dtype=torch.uint8,
-        device=w13_weight.device,
-    )
-    w2_weights_shuffled_tensor = torch.empty(
-        (num_experts, w2_cols // block_k, w2_rows, block_k),
-        dtype=torch.uint8,
-        device=w2_weight.device,
-    )
-
-    for i in range(num_experts):
-        w13_expert_uint8 = w13_weight[i].view(torch.uint8)
-
-        permute_indices = _maybe_get_cached_w3_w1_permute_indices(
-            cache_permute_indices,
-            w13_expert_uint8,
-            epilogue_tile_m,
-            is_gated_act_gemm=is_gated_act_gemm,
-        )
-        if is_gated_act_gemm:
-            rows = w13_expert_uint8.shape[0]
-            permute_indices = (permute_indices + rows // 2) % rows
-        _copy_permuted_expert_to_block_layout(
-            w13_weights_shuffled_tensor[i],
-            w13_expert_uint8,
-            permute_indices,
+    def _convert_weight_in_place(
+        weight: torch.Tensor,
+        is_w13: bool,
+    ) -> torch.Tensor:
+        rows, cols = weight[0].view(torch.uint8).shape
+        block_layout_shape = (num_experts, cols // block_k, rows, block_k)
+        expert_scratch = torch.empty(
+            block_layout_shape[1:],
+            dtype=torch.uint8,
+            device=weight.device,
         )
 
-        permute_indices = get_w2_permute_indices_with_cache(
-            cache_permute_indices,
-            w2_weight[i].view(torch.uint8),
-            epilogue_tile_m,
-        )
-        _copy_permuted_expert_to_block_layout(
-            w2_weights_shuffled_tensor[i],
-            w2_weight[i].view(torch.uint8),
-            permute_indices,
-        )
+        for i in range(num_experts):
+            expert_uint8 = weight[i].view(torch.uint8)
+            if is_w13:
+                permute_indices = _maybe_get_cached_w3_w1_permute_indices(
+                    cache_permute_indices,
+                    expert_uint8,
+                    epilogue_tile_m,
+                    is_gated_act_gemm=is_gated_act_gemm,
+                )
+                if is_gated_act_gemm:
+                    permute_indices = (
+                        permute_indices + expert_uint8.shape[0] // 2
+                    ) % expert_uint8.shape[0]
+            else:
+                permute_indices = get_w2_permute_indices_with_cache(
+                    cache_permute_indices,
+                    expert_uint8,
+                    epilogue_tile_m,
+                )
+
+            _copy_permuted_expert_to_block_layout(
+                expert_scratch,
+                expert_uint8,
+                permute_indices,
+            )
+            expert_uint8.view(-1).copy_(expert_scratch.view(-1))
+
+        return weight.view(torch.uint8).view(block_layout_shape).view(torch.bfloat16)
 
     return (
-        w13_weights_shuffled_tensor.view(torch.bfloat16),
-        w2_weights_shuffled_tensor.view(torch.bfloat16),
+        _convert_weight_in_place(w13_weight, is_w13=True),
+        _convert_weight_in_place(w2_weight, is_w13=False),
     )
 
 

vllm/model_executor/model_loader/reload/meta.py

@@ -9,7 +9,7 @@ from torch.utils._python_dispatch import TorchDispatchMode
 
 from .sanitize import restore_layer_refs, sanitize_layer_refs
 from .types import LayerReloadingInfo, LayerTensors
-from .utils import get_layer_params_buffers, get_layer_tensors
+from .utils import get_layer_params_buffers, get_layer_tensors, get_tensor_load_numel
 
 __all__ = [
     "to_meta_tensor",
@@ -206,5 +206,5 @@ def get_numel_loaded(
     numel = counter.copied_numel
     param = args.arguments.get("param", None)
     if isinstance(param, torch.Tensor):
-        numel = min(numel, param.numel())
+        numel = min(numel, get_tensor_load_numel(param))
     return numel, return_value

vllm/model_executor/model_loader/reload/utils.py

@@ -11,6 +11,7 @@ __all__ = [
     "get_layer_tensors",
     "get_layer_params_buffers",
     "get_layer_size",
+    "get_tensor_load_numel",
     "has_device_tensors",
     "get_info_size",
 ]
@@ -30,6 +31,11 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
     )
 
 
+def get_tensor_load_numel(tensor: torch.Tensor) -> int:
+    """Count checkpoint elements, excluding padding declared by the weight creator."""
+    return getattr(tensor, "weight_loader_numel", tensor.numel())
+
+
 def get_layer_size(layer: torch.nn.Module) -> int:
     """Calculate total number of elements across loadable tensors in a layer.
 
@@ -39,7 +45,7 @@ def get_layer_size(layer: torch.nn.Module) -> int:
     from .meta import SKIP_LOAD_TENSORS
 
     return sum(
-        tensor.numel()
+        get_tensor_load_numel(tensor)
         for name, tensor in get_layer_tensors(layer).items()
         if name not in SKIP_LOAD_TENSORS
     )