diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2e022ede36..53d760c339 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/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( diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index d3716a196b..423f181132 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/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 diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index 9e50b4fc3b..f7466c63c8 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/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, diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b4a544d2fc..844149ce1b 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 42ea15dcda..e3a287ddc2 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/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( diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 523d1df0d8..a1b2471f5d 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 044285ac85..1bfd7e61f9 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/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), ) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 4d08d49165..9874ebd67b 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/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 diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py index b7dc361503..203aa82af3 100644 --- a/vllm/model_executor/model_loader/reload/utils.py +++ b/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 )