@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 78fdf4efb8

[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>

Yuchen Wang committed Sep 17, 2026 at 07:00 UTC 78fdf4efb86dffde1b0f5733207aed038c3530f8
9 files changed +667 -68
tests/kernels/moe/test_flashinfer.py
+258
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import multiprocessing
4 from dataclasses import dataclass
5
6 import pytest
@@ -433,6 +434,263 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
434
435 assert w13_converted.shape[0] == num_experts
436 assert w2_converted.shape[0] == num_experts
437 + assert w13_converted.data_ptr() == w13.data_ptr()
438 + assert w2_converted.data_ptr() == w2.data_ptr()
439 +
440 +
441 +@pytest.mark.parametrize("is_gated_act_gemm", [True, False])
442 +def test_convert_moe_weights_to_flashinfer_trtllm_block_layout_values(
443 + is_gated_act_gemm,
444 +):
445 + from flashinfer.fused_moe.core import (
446 + _maybe_get_cached_w3_w1_permute_indices,
447 + get_w2_permute_indices_with_cache,
448 + )
449 +
450 + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
451 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
452 + )
453 +
454 + num_experts, intermediate, hidden = 2, 256, 256
455 + w13_multiplier = 2 if is_gated_act_gemm else 1
456 + w13 = torch.randn(
457 + (num_experts, w13_multiplier * intermediate, hidden),
458 + dtype=torch.bfloat16,
459 + device="cuda",
460 + )
461 + w2 = torch.randn(
462 + (num_experts, hidden, intermediate),
463 + dtype=torch.bfloat16,
464 + device="cuda",
465 + )
466 +
467 + def _reference_block_layout(
468 + weight: torch.Tensor,
469 + is_w13: bool,
470 + cache: dict[torch.Size, torch.Tensor],
471 + ) -> torch.Tensor:
472 + outputs = []
473 + for expert in weight:
474 + expert_uint8 = expert.view(torch.uint8)
475 + if is_w13:
476 + indices = _maybe_get_cached_w3_w1_permute_indices(
477 + cache,
478 + expert_uint8,
479 + 128,
480 + is_gated_act_gemm=is_gated_act_gemm,
481 + )
482 + if is_gated_act_gemm:
483 + indices = (indices + expert_uint8.shape[0] // 2) % (
484 + expert_uint8.shape[0]
485 + )
486 + else:
487 + indices = get_w2_permute_indices_with_cache(
488 + cache,
489 + expert_uint8,
490 + 128,
491 + )
492 + rows, cols = expert_uint8.shape
493 + blocks = expert_uint8.view(rows, cols // 128, 128).permute(1, 0, 2)
494 + outputs.append(torch.index_select(blocks, 1, indices.to(weight.device)))
495 + return torch.stack(outputs).view(torch.bfloat16)
496 +
497 + reference_cache: dict[torch.Size, torch.Tensor] = {}
498 + expected_w13 = _reference_block_layout(w13, is_w13=True, cache=reference_cache)
499 + expected_w2 = _reference_block_layout(w2, is_w13=False, cache=reference_cache)
500 + w13_ptr = w13.data_ptr()
501 + w2_ptr = w2.data_ptr()
502 +
503 + actual_w13, actual_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout(
504 + {},
505 + w13,
506 + w2,
507 + is_gated_act_gemm=is_gated_act_gemm,
508 + )
509 +
510 + assert actual_w13.data_ptr() == w13_ptr
511 + assert actual_w2.data_ptr() == w2_ptr
512 + assert torch.equal(actual_w13, expected_w13)
513 + assert torch.equal(actual_w2, expected_w2)
514 +
515 +
516 +def _make_unquantized_flashinfer_test_layer(
517 + monkeypatch, intermediate, is_gated, *, device="cuda"
518 +):
519 + from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
520 + UnquantizedMoeBackend,
521 + )
522 + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
523 + UnquantizedFusedMoEMethod,
524 + )
525 +
526 + moe_config = make_dummy_moe_config(
527 + num_experts=2,
528 + hidden_dim=256,
529 + intermediate_size=intermediate,
530 + activation=MoEActivation.SILU if is_gated else MoEActivation.RELU2_NO_MUL,
531 + )
532 + moe_config.intermediate_size_per_partition = (intermediate + 127) // 128 * 128
533 + method = object.__new__(UnquantizedFusedMoEMethod)
534 + method.moe = moe_config
535 + method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
536 + method.moe_kernel = None
537 + monkeypatch.setattr(
538 + method,
539 + "_init_moe_kernel",
540 + lambda _: setattr(method, "moe_kernel", object()),
541 + )
542 +
543 + layer = torch.nn.Module()
544 + layer.moe_config = moe_config
545 + with torch.device(device):
546 + method.create_weights(
547 + layer,
548 + num_experts=2,
549 + hidden_size=256,
550 + intermediate_size_per_partition=moe_config.intermediate_size_per_partition,
551 + params_dtype=torch.bfloat16,
552 + )
553 + return method, layer
554 +
555 +
556 +@pytest.mark.parametrize("intermediate", [192, 256])
557 +@pytest.mark.parametrize("is_gated", [True, False])
558 +def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(
559 + monkeypatch, intermediate, is_gated
560 +):
561 + """Each raw reload clears padding left dirty by prior in-place conversion."""
562 + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
563 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
564 + )
565 +
566 + method, layer = _make_unquantized_flashinfer_test_layer(
567 + monkeypatch, intermediate, is_gated
568 + )
569 + w13_shape, w2_shape = layer.w13_weight.shape, layer.w2_weight.shape
570 + padded = w2_shape[-1]
571 + w13_ptr = layer.w13_weight.data_ptr()
572 + w2_ptr = layer.w2_weight.data_ptr()
573 +
574 + for _ in range(2):
575 + reloaded_w13 = torch.randn_like(layer.w13_weight)
576 + reloaded_w2 = torch.randn_like(layer.w2_weight)
577 + reloaded_w13[:, intermediate:padded].zero_()
578 + if is_gated:
579 + reloaded_w13[:, padded + intermediate :].zero_()
580 + reloaded_w2[:, :, intermediate:].zero_()
581 + expected_w13, expected_w2 = (
582 + convert_moe_weights_to_flashinfer_trtllm_block_layout(
583 + {},
584 + reloaded_w13.clone(),
585 + reloaded_w2.clone(),
586 + is_gated_act_gemm=is_gated,
587 + )
588 + )
589 +
590 + # The loader writes only logical checkpoint slices, not padding.
591 + layer.w13_weight.fill_(float("nan"))
592 + layer.w2_weight.fill_(float("nan"))
593 + layer.w13_weight[:, :intermediate].copy_(reloaded_w13[:, :intermediate])
594 + if is_gated:
595 + layer.w13_weight[:, padded : padded + intermediate].copy_(
596 + reloaded_w13[:, padded : padded + intermediate]
597 + )
598 + layer.w2_weight[:, :, :intermediate].copy_(reloaded_w2[:, :, :intermediate])
599 + method.process_weights_after_loading(layer)
600 +
601 + assert layer.w13_weight.shape == w13_shape
602 + assert layer.w2_weight.shape == w2_shape
603 + assert layer.w13_weight.data_ptr() == w13_ptr
604 + assert layer.w2_weight.data_ptr() == w2_ptr
605 + kernel_w13, kernel_w2 = method._kernel_weights(layer)
606 + assert torch.equal(kernel_w13, expected_w13)
607 + assert torch.equal(kernel_w2, expected_w2)
608 +
609 +
610 +def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index):
611 + from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader
612 + from vllm.model_executor.utils import weights_already_processed
613 +
614 + torch.accelerator.set_device_index(device_index)
615 + with pytest.MonkeyPatch.context() as monkeypatch:
616 + method, layer = _make_unquantized_flashinfer_test_layer(
617 + monkeypatch, 192, is_gated, device="meta"
618 + )
619 + loader = object.__new__(IpcModelLoader)
620 + loader.mode = mode
621 + loader._apply_entries(layer, entries, {}, device_index)
622 + pointers = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr())
623 +
624 + # Like ipc_cache: a fresh method, only tensor metadata, no _setup_kernel.
625 + with weights_already_processed():
626 + method.process_weights_after_loading(layer)
627 + assert method.moe_kernel is not None
628 + actual = method._kernel_weights(layer)
629 + for weight, reference, pointer in zip(actual, expected, pointers):
630 + assert weight.data_ptr() == pointer
631 + assert torch.equal(weight.cpu(), reference)
632 + # Cross-process addresses need not match; writes prove sharing/isolation.
633 + with torch.no_grad():
634 + for weight in actual:
635 + weight.fill_(1.0)
636 + torch.accelerator.synchronize()
637 +
638 +
639 +@pytest.mark.parametrize("is_gated", [True, False])
640 +@pytest.mark.parametrize("cache_ndim", [3, 4])
641 +@pytest.mark.parametrize("mode", ["copy", "zero_copy"])
642 +def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(
643 + monkeypatch, is_gated, cache_ndim, mode
644 +):
645 + """A spawned IPC consumer restores views without transient method state."""
646 + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
647 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
648 + )
649 + from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry
650 +
651 + method, layer = _make_unquantized_flashinfer_test_layer(
652 + monkeypatch, 192, is_gated, device="meta"
653 + )
654 + packed = convert_moe_weights_to_flashinfer_trtllm_block_layout(
655 + {},
656 + torch.randn_like(layer.w13_weight, device="cuda"),
657 + torch.randn_like(layer.w2_weight, device="cuda"),
658 + is_gated_act_gemm=is_gated,
659 + )
660 + expected = [weight.cpu() for weight in packed]
661 + entries = {}
662 + for name, weight in zip(("w13_weight", "w2_weight"), packed):
663 + cached = weight.view_as(getattr(layer, name)) if cache_ndim == 3 else weight
664 + entries[name] = TensorEntry.from_tensor(cached, kind="param")
665 + consumer = multiprocessing.get_context("spawn").Process(
666 + target=_check_flashinfer_ipc_weights,
667 + args=(
668 + entries,
669 + expected,
670 + is_gated,
671 + mode,
672 + torch.accelerator.current_device_index(),
673 + ),
674 + )
675 + consumer.start()
676 + try:
677 + # Keep producer allocations alive until the consumer releases its views.
678 + consumer.join(timeout=180)
679 + assert not consumer.is_alive(), "IPC consumer timed out"
680 + assert consumer.exitcode == 0, "IPC consumer failed; see child traceback"
681 + for source, reference in zip(packed, expected):
682 + expected_source = (
683 + torch.ones_like(reference) if mode == "zero_copy" else reference
684 + )
685 + assert torch.equal(source.cpu(), expected_source)
686 + finally:
687 + if consumer.is_alive():
688 + consumer.terminate()
689 + consumer.join(timeout=10)
690 + if consumer.is_alive():
691 + consumer.kill()
692 + consumer.join(timeout=10)
693 + consumer.close()
694
695
696 @pytest.mark.parametrize(
tests/kernels/moe/test_moe_weight_loading_padded.py
+126
@@ -12,7 +12,14 @@ correctly handles this mismatch.
12 import pytest
13 import torch
14
15 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 +from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
17 + UnquantizedMoeBackend,
18 +)
19 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
20 +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
21 + UnquantizedFusedMoEMethod,
22 +)
23
24 from .utils import make_dummy_moe_config
25
@@ -249,6 +256,57 @@ class TestNarrowExpertDataForPadding:
256 class TestWeightLoadingWithPaddedHiddenSize:
257 """Integration-style tests that simulate padded weight loading."""
258
259 + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
260 + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
261 + @pytest.mark.parametrize("tp_rank", [0, 1])
262 + def test_load_w2_chunks_noncontiguous_cpu_source_to_padded_cuda(
263 + self, dtype, tp_rank
264 + ):
265 + hidden = 4096
266 + intermediate = 1024
267 + loaded_weight = (
268 + torch.arange(hidden * intermediate * 2)
269 + .remainder(251)
270 + .to(dtype)
271 + .reshape(hidden, intermediate * 2)
272 + )
273 + tp_source = loaded_weight[
274 + :, intermediate * tp_rank : intermediate * (tp_rank + 1)
275 + ]
276 + expert_data_full = torch.zeros(
277 + hidden + 8, intermediate + 8, device="cuda", dtype=dtype
278 + )
279 + destination = expert_data_full[:hidden, :intermediate]
280 +
281 + assert not tp_source.is_contiguous()
282 + assert tp_source.nbytes > 1 << 20
283 + assert not destination.is_contiguous()
284 +
285 + experts = object.__new__(RoutedExperts)
286 + torch.nn.Module.__init__(experts)
287 + experts.moe_config = make_dummy_moe_config()
288 + experts.moe_config.moe_parallel_config.tp_size = 2
289 +
290 + torch.accelerator.synchronize()
291 + allocated_before = torch.accelerator.memory_allocated()
292 + torch.accelerator.reset_peak_memory_stats()
293 + experts._load_w2(
294 + expert_data=expert_data_full,
295 + shard_dim=1,
296 + loaded_weight=loaded_weight,
297 + tp_rank=tp_rank,
298 + )
299 + torch.accelerator.synchronize()
300 + peak_extra = torch.accelerator.max_memory_allocated() - allocated_before
301 +
302 + # A single copy into the strided CUDA view allocates a full-shard
303 + # temporary. Correct values alone would not catch that regression.
304 + assert peak_extra < tp_source.nbytes
305 +
306 + torch.testing.assert_close(destination.cpu(), tp_source)
307 + assert torch.count_nonzero(expert_data_full[hidden:, :]) == 0
308 + assert torch.count_nonzero(expert_data_full[:hidden, intermediate:]) == 0
309 +
310 def test_load_w2_with_padding(self):
311 """Simulate loading w2 weights when hidden_size is padded."""
312 padded_hidden = 3072
@@ -364,6 +422,74 @@ class TestWeightLoadingWithPaddedHiddenSize:
422 )
423
424
425 +class TestUnquantizedTrtLlmPrePadding:
426 + @staticmethod
427 + def _make_method(
428 + intermediate: int,
429 + backend: UnquantizedMoeBackend = UnquantizedMoeBackend.FLASHINFER_TRTLLM,
430 + ) -> UnquantizedFusedMoEMethod:
431 + moe_config = make_dummy_moe_config(
432 + num_experts=2,
433 + hidden_dim=64,
434 + intermediate_size=intermediate,
435 + )
436 + method = object.__new__(UnquantizedFusedMoEMethod)
437 + method.moe = moe_config
438 + method.unquantized_backend = backend
439 + method.moe_kernel = None
440 + return method
441 +
442 + @pytest.mark.parametrize(
443 + "backend,original,expected",
444 + [
445 + (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1344, 1408),
446 + (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1408, 1408),
447 + (UnquantizedMoeBackend.TRITON, 1344, 1344),
448 + ],
449 + )
450 + def test_rounds_intermediate_before_weight_allocation(
451 + self,
452 + backend: UnquantizedMoeBackend,
453 + original: int,
454 + expected: int,
455 + ):
456 + method = self._make_method(original, backend)
457 +
458 + hidden, intermediate = method.maybe_roundup_sizes(
459 + hidden_size=64,
460 + intermediate_size_per_partition=original,
461 + act_dtype=torch.bfloat16,
462 + moe_parallel_config=method.moe.moe_parallel_config,
463 + )
464 +
465 + assert hidden == 64
466 + assert intermediate == expected
467 +
468 + @pytest.mark.parametrize("is_act_and_mul", [True, False])
469 + def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool):
470 + method = self._make_method(1344)
471 + if not is_act_and_mul:
472 + method.moe.activation = MoEActivation.RELU2_NO_MUL
473 + method.moe.intermediate_size_per_partition = 1408
474 + layer = torch.nn.Module()
475 +
476 + method.create_weights(
477 + layer=layer,
478 + num_experts=2,
479 + hidden_size=64,
480 + intermediate_size_per_partition=1408,
481 + params_dtype=torch.bfloat16,
482 + )
483 +
484 + up_mult = 2 if is_act_and_mul else 1
485 + assert layer.w13_weight.shape == (2, up_mult * 1408, 64)
486 + assert layer.w2_weight.shape == (2, 64, 1408)
487 + assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0
488 + if is_act_and_mul:
489 + assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0
490 + assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0
491 +
492 +
493 class TestLoadWeightsExpertBias:
494 """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize
495 all-zero per-expert `.bias` tensors for models whose experts have no bias
tests/kernels/moe/test_trtllm_bf16_moe.py
+6 -2
@@ -75,6 +75,10 @@ def test_trtllm_bf16_moe_modular_no_graph(
75 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
76 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
77 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
78 + # The FlashInfer conversion may rewrite unpadded input storage in place.
79 + # Preserve the original layout for the independent torch reference.
80 + reference_w1 = w1.clone()
81 + reference_w2 = w2.clone()
82 score = torch.randn((m, e), device="cuda", dtype=dtype)
83 scores = torch.softmax(score, dim=-1, dtype=torch.float32)
84 topk_weights, topk_ids = torch.topk(scores, topk)
@@ -133,8 +137,8 @@ def test_trtllm_bf16_moe_modular_no_graph(
137
138 torch_output = torch_moe(
139 a,
136 - w1,
137 - w2,
140 + reference_w1,
141 + reference_w2,
142 score,
143 topk,
144 activation=MoEActivation.SILU,
tests/model_executor/model_loader/test_reload.py
+105
@@ -5,6 +5,7 @@ import importlib.machinery
5 import inspect
6 import sys
7 import types
8 +from types import SimpleNamespace
9 from unittest.mock import Mock
10 from weakref import WeakKeyDictionary, ref
11
@@ -467,6 +468,110 @@ def test_model_cleanup(dist_init, default_vllm_config):
468 assert len(mock_info_dict) == 0
469
470
471 +@pytest.mark.parametrize("is_gated", [False, True])
472 +@pytest.mark.parametrize("has_bias", [False, True])
473 +@pytest.mark.parametrize("tp_rank", [0, 1])
474 +@pytest.mark.parametrize("padded", [False, True])
475 +def test_padded_moe_reload_releases_each_layer(
476 + monkeypatch, is_gated, has_bias, tp_rank, padded
477 +):
478 + """Checkpoint-sized copies finish each layer without global finalization."""
479 + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
480 + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
481 + UnquantizedFusedMoEMethod,
482 + )
483 +
484 + hidden, intermediate, experts = 4, 3, 2
485 + stored_hidden, stored_intermediate = (8, 8) if padded else (hidden, intermediate)
486 + config = SimpleNamespace(
487 + hidden_dim_unpadded=hidden,
488 + intermediate_size_per_partition_unpadded=intermediate,
489 + is_act_and_mul=is_gated,
490 + has_bias=has_bias,
491 + tp_rank=tp_rank,
492 + moe_parallel_config=SimpleNamespace(tp_size=2),
493 + )
494 + model = torch.nn.ModuleList()
495 + processed: list[torch.nn.Module] = []
496 + for _ in range(2):
497 + method = object.__new__(UnquantizedFusedMoEMethod)
498 + torch.nn.Module.__init__(method)
499 + method.moe = config
500 + # The regression concerns streaming reload, not kernel conversion.
501 + monkeypatch.setattr(method, "process_weights_after_loading", processed.append)
502 + layer = object.__new__(RoutedExperts)
503 + torch.nn.Module.__init__(layer)
504 + layer.moe_config = config
505 + layer.quant_config = None
506 + layer.quant_method = method
507 + layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i)
508 + layer._loaded_expert_biases = set()
509 + method.create_weights(
510 + layer,
511 + experts,
512 + stored_hidden,
513 + stored_intermediate,
514 + torch.float32,
515 + weight_loader=layer.weight_loader,
516 + )
517 + model.append(layer)
518 +
519 + record_metadata_for_reloading(model)
520 + original_params = [dict(layer.named_parameters()) for layer in model]
521 + shards = ["w1", "w3", "w2"] if is_gated else ["w1", "w2"]
522 + for cycle in range(2):
523 + initialize_layerwise_reload(model)
524 + for layer_index, layer in enumerate(model):
525 + info = reload_layerwise.get_layerwise_info(layer)
526 + inputs = []
527 + expected = {
528 + name: torch.full_like(p, float("nan"))
529 + for name, p in original_params[layer_index].items()
530 + }
531 + params = dict(layer.named_parameters())
532 + calls = [
533 + (e, s, b)
534 + for e in range(experts)
535 + for s in shards
536 + for b in ([False, True] if has_bias else [False])
537 + ]
538 + for call_index, (expert, shard, bias) in enumerate(calls):
539 + name = ("w2" if shard == "w2" else "w13") + (
540 + "_bias" if bias else "_weight"
541 + )
542 + shape = (
543 + ((hidden,) if bias else (hidden, 2 * intermediate))
544 + if shard == "w2"
545 + else ((2 * intermediate,) if bias else (2 * intermediate, hidden))
546 + )
547 + weight = torch.arange(
548 + torch.Size(shape).numel(), dtype=torch.float32
549 + ).reshape(shape)
550 + weight = weight + 100 * (1 + call_index + cycle)
551 + inputs.append(ref(weight))
552 + # Direct checkpoint loading is the reference for deferred reload.
553 + layer.weight_loader(expected[name], weight, name, shard, expert)
554 + param = params[name]
555 + param.weight_loader(param, weight, name, shard, expert)
556 + del weight
557 + if call_index != len(calls) - 1:
558 + assert info.can_load(), (
559 + "Layer processed before its final checkpoint shard"
560 + )
561 +
562 + assert not info.can_load(), (
563 + "Padding must not defer the layer until finalization"
564 + )
565 + assert not info.loaded_weights
566 + assert len(processed) == cycle * len(model) + layer_index + 1
567 + assert all(source() is None for source in inputs)
568 + for name, original in original_params[layer_index].items():
569 + assert getattr(layer, name) is original
570 + # Kernel-specific tests cover padding, which is not checkpoint data.
571 + mask = torch.isfinite(expected[name])
572 + assert torch.equal(original[mask], expected[name][mask])
573 +
574 +
575 def test_get_numel_loaded():
576 param = torch.empty(10, device="meta")
577 loaded_weight = torch.empty(10)
vllm/model_executor/layers/fused_moe/routed_experts.py
+18
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.base_config import (
27 QuantizationConfig,
28 resolve_quant_method,
29 )
30 +from vllm.utils.math_utils import cdiv
31
32 if TYPE_CHECKING:
33 from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
@@ -566,6 +567,23 @@ class RoutedExperts(PluggableLayer):
567 hidden_dim=hidden_dim,
568 shard_dim=shard_dim,
569 )
570 + if (
571 + loaded_weight.device.type == "cpu"
572 + and expert_data.device.type == "cuda"
573 + and not expert_data.is_contiguous()
574 + and expert_data.ndim == 2
575 + ):
576 + # A strided CPU-to-CUDA copy can allocate a full-shard CUDA
577 + # temporary. Make the TP slice contiguous on CPU and copy rows
578 + # in chunks to limit the temporary used for the padded view.
579 + loaded_weight = loaded_weight.contiguous()
580 + num_chunks = max(cdiv(loaded_weight.nbytes, 1 << 20), 1)
581 + for dst, src in zip(
582 + expert_data.chunk(num_chunks, dim=0),
583 + loaded_weight.chunk(num_chunks, dim=0),
584 + ):
585 + dst.copy_(src)
586 + return
587 expert_data.copy_(loaded_weight)
588
589 def _load_single_value(
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+102 -20
@@ -11,6 +11,7 @@ from vllm.logger import init_logger
11 from vllm.model_executor.custom_op import CustomOp
12 from vllm.model_executor.layers.fused_moe.config import (
13 FusedMoEConfig,
14 + FusedMoEParallelConfig,
15 FusedMoEQuantConfig,
16 biased_moe_quant_config,
17 )
@@ -33,6 +34,7 @@ from vllm.model_executor.utils import (
34 set_weight_attrs,
35 )
36 from vllm.platforms import current_platform
37 +from vllm.utils.math_utils import round_up
38
39 if TYPE_CHECKING:
40 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
@@ -68,22 +70,38 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
70 params_dtype: torch.dtype,
71 **extra_weight_attrs,
72 ):
73 + unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded
74 + assert unpadded_intermediate is not None
75 + has_intermediate_padding = (
76 + intermediate_size_per_partition > unpadded_intermediate
77 + )
78 if self.moe.is_act_and_mul:
79 w13_up_dim = 2 * intermediate_size_per_partition
80 else:
81 w13_up_dim = intermediate_size_per_partition
82 + unpadded_hidden = self.moe.hidden_dim_unpadded
83 + assert unpadded_hidden is not None
84 + unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1)
85 # Fused gate_up_proj (column parallel)
76 - w13_weight = torch.nn.Parameter(
77 - torch.empty(
78 - num_experts,
79 - w13_up_dim,
80 - hidden_size,
81 - dtype=params_dtype,
82 - ),
83 - requires_grad=False,
86 + w13_data = torch.empty(
87 + num_experts,
88 + w13_up_dim,
89 + hidden_size,
90 + dtype=params_dtype,
91 )
92 + if has_intermediate_padding:
93 + # The checkpoint loader only writes the unpadded slices. Zero the
94 + # tails now so no full-size padding copy is needed after loading.
95 + w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_()
96 + if self.moe.is_act_and_mul:
97 + w13_data[
98 + :,
99 + intermediate_size_per_partition + unpadded_intermediate :,
100 + ].zero_()
101 + w13_weight = torch.nn.Parameter(w13_data, requires_grad=False)
102 layer.register_parameter("w13_weight", w13_weight)
103 set_weight_attrs(w13_weight, extra_weight_attrs)
104 + w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden
105 if self.moe.has_bias:
106 w13_bias = torch.nn.Parameter(
107 torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
@@ -91,18 +109,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
109 )
110 layer.register_parameter("w13_bias", w13_bias)
111 set_weight_attrs(w13_bias, extra_weight_attrs)
112 + w13_bias.weight_loader_numel = num_experts * unpadded_up
113 # down_proj (row parallel)
95 - w2_weight = torch.nn.Parameter(
96 - torch.empty(
97 - num_experts,
98 - hidden_size,
99 - intermediate_size_per_partition,
100 - dtype=params_dtype,
101 - ),
102 - requires_grad=False,
114 + w2_data = torch.empty(
115 + num_experts,
116 + hidden_size,
117 + intermediate_size_per_partition,
118 + dtype=params_dtype,
119 )
120 + if has_intermediate_padding:
121 + w2_data[:, :, unpadded_intermediate:].zero_()
122 + w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)
123 layer.register_parameter("w2_weight", w2_weight)
124 set_weight_attrs(w2_weight, extra_weight_attrs)
125 + w2_weight.weight_loader_numel = (
126 + num_experts * unpadded_hidden * unpadded_intermediate
127 + )
128 if self.moe.has_bias:
129 w2_bias = torch.nn.Parameter(
130 torch.zeros(num_experts, hidden_size, dtype=params_dtype),
@@ -110,6 +132,26 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
132 )
133 layer.register_parameter("w2_bias", w2_bias)
134 set_weight_attrs(w2_bias, extra_weight_attrs)
135 + w2_bias.weight_loader_numel = num_experts * unpadded_hidden
136 +
137 + def maybe_roundup_sizes(
138 + self,
139 + hidden_size: int,
140 + intermediate_size_per_partition: int,
141 + act_dtype: torch.dtype,
142 + moe_parallel_config: FusedMoEParallelConfig,
143 + ) -> tuple[int, int]:
144 + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
145 + hidden_size,
146 + intermediate_size_per_partition,
147 + act_dtype,
148 + moe_parallel_config,
149 + )
150 + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
151 + intermediate_size_per_partition = round_up(
152 + intermediate_size_per_partition, 128
153 + )
154 + return hidden_size, intermediate_size_per_partition
155
156 def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
157 # Pad the weight tensor. This is an optimization on ROCm platform, which
@@ -135,6 +177,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
177 w13: torch.Tensor,
178 w2: torch.Tensor,
179 ) -> None:
180 + w13_parameter_shape = w13.shape
181 + w2_parameter_shape = w2.shape
182 +
183 + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
184 + unpadded = self.moe.intermediate_size_per_partition_unpadded
185 + assert unpadded is not None
186 + intermediate = w2.shape[-1]
187 + # Reloads overwrite only checkpoint slices; the previous in-place
188 + # permutation can leave nonzero values in the raw padding slots.
189 + if intermediate > unpadded:
190 + w13[:, unpadded:intermediate].zero_()
191 + if self.moe.is_act_and_mul:
192 + w13[:, intermediate + unpadded :].zero_()
193 + w2[:, :, unpadded:].zero_()
194 +
195 # Shuffle weights to runtime format.
196 w13_new, w2_new = convert_to_unquantized_kernel_format(
197 self.unquantized_backend,
@@ -142,6 +199,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
199 w13_weight=w13,
200 w2_weight=w2,
201 )
202 + if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
203 + w13_new = w13_new.view(w13_parameter_shape)
204 + w2_new = w2_new.view(w2_parameter_shape)
205 +
206 # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;
207 # On the first call we replace the parameter normally. On subsequent
208 # calls (e.g. RL weight updates that re-trigger
@@ -170,6 +231,25 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
231 assert self.moe_kernel is not None
232 self.moe_kernel.fused_experts.process_weights_after_loading(layer)
233
234 + def _kernel_weights(
235 + self, layer: "RoutedExperts"
236 + ) -> tuple[torch.Tensor, torch.Tensor]:
237 + if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM:
238 + return layer.w13_weight, layer.w2_weight
239 +
240 + # IPC caches restore tensor shapes/storage, not method attributes.
241 + # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements.
242 + def block_view(weight: torch.Tensor) -> torch.Tensor:
243 + if weight.ndim == 4:
244 + return weight
245 + experts, rows, cols = weight.shape
246 + return weight.view(experts, cols // 64, rows, 64)
247 +
248 + return (
249 + block_view(layer.w13_weight),
250 + block_view(layer.w2_weight),
251 + )
252 +
253 def _init_moe_kernel(self, layer: "RoutedExperts") -> None:
254 """Build the MoE kernel from the layer's current (shuffled) weights."""
255 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
@@ -274,10 +354,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
354 shared_experts_input: torch.Tensor | None,
355 ) -> torch.Tensor:
356 assert self.moe_kernel is not None
357 + w13, w2 = self._kernel_weights(layer)
358 return self.moe_kernel.apply(
359 hidden_states=x,
279 - w1=layer.w13_weight,
280 - w2=layer.w2_weight,
360 + w1=w13,
361 + w2=w2,
362 topk_weights=topk_weights,
363 topk_ids=topk_ids,
364 activation=layer.activation,
@@ -315,10 +396,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
396 ) -> torch.Tensor | UnfinalizedMoEOutput:
397 assert self.is_monolithic
398 assert self.moe_kernel is not None
399 + w13, w2 = self._kernel_weights(layer)
400 return self.moe_kernel.apply_monolithic(
401 x,
320 - layer.w13_weight,
321 - layer.w2_weight,
402 + w13,
403 + w2,
404 router_logits,
405 activation=layer.activation,
406 global_num_experts=layer.global_num_experts,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+43 -43
@@ -112,8 +112,9 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
112 ) -> tuple[torch.Tensor, torch.Tensor]:
113 """Convert expert weights to FlashInfer's block layout.
114
115 - This reorders W13 and W2 into the expected epilogue-tiled block layout and
116 - returns the shuffled weight tensors.
115 + This reorders W13 and W2 in place into the expected epilogue-tiled block
116 + layout and returns views of the shuffled weight tensors. Using one expert
117 + as scratch space avoids allocating another full copy of both weights.
118 """
119 if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
120 raise ValueError(
@@ -147,51 +148,50 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
148 out=out,
149 )
150
150 - w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape
151 - w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape
152 - w13_weights_shuffled_tensor = torch.empty(
153 - (num_experts, w13_cols // block_k, w13_rows, block_k),
154 - dtype=torch.uint8,
155 - device=w13_weight.device,
156 - )
157 - w2_weights_shuffled_tensor = torch.empty(
158 - (num_experts, w2_cols // block_k, w2_rows, block_k),
159 - dtype=torch.uint8,
160 - device=w2_weight.device,
161 - )
162 -
163 - for i in range(num_experts):
164 - w13_expert_uint8 = w13_weight[i].view(torch.uint8)
165 -
166 - permute_indices = _maybe_get_cached_w3_w1_permute_indices(
167 - cache_permute_indices,
168 - w13_expert_uint8,
169 - epilogue_tile_m,
170 - is_gated_act_gemm=is_gated_act_gemm,
171 - )
172 - if is_gated_act_gemm:
173 - rows = w13_expert_uint8.shape[0]
174 - permute_indices = (permute_indices + rows // 2) % rows
175 - _copy_permuted_expert_to_block_layout(
176 - w13_weights_shuffled_tensor[i],
177 - w13_expert_uint8,
178 - permute_indices,
151 + def _convert_weight_in_place(
152 + weight: torch.Tensor,
153 + is_w13: bool,
154 + ) -> torch.Tensor:
155 + rows, cols = weight[0].view(torch.uint8).shape
156 + block_layout_shape = (num_experts, cols // block_k, rows, block_k)
157 + expert_scratch = torch.empty(
158 + block_layout_shape[1:],
159 + dtype=torch.uint8,
160 + device=weight.device,
161 )
162
181 - permute_indices = get_w2_permute_indices_with_cache(
182 - cache_permute_indices,
183 - w2_weight[i].view(torch.uint8),
184 - epilogue_tile_m,
185 - )
186 - _copy_permuted_expert_to_block_layout(
187 - w2_weights_shuffled_tensor[i],
188 - w2_weight[i].view(torch.uint8),
189 - permute_indices,
190 - )
163 + for i in range(num_experts):
164 + expert_uint8 = weight[i].view(torch.uint8)
165 + if is_w13:
166 + permute_indices = _maybe_get_cached_w3_w1_permute_indices(
167 + cache_permute_indices,
168 + expert_uint8,
169 + epilogue_tile_m,
170 + is_gated_act_gemm=is_gated_act_gemm,
171 + )
172 + if is_gated_act_gemm:
173 + permute_indices = (
174 + permute_indices + expert_uint8.shape[0] // 2
175 + ) % expert_uint8.shape[0]
176 + else:
177 + permute_indices = get_w2_permute_indices_with_cache(
178 + cache_permute_indices,
179 + expert_uint8,
180 + epilogue_tile_m,
181 + )
182 +
183 + _copy_permuted_expert_to_block_layout(
184 + expert_scratch,
185 + expert_uint8,
186 + permute_indices,
187 + )
188 + expert_uint8.view(-1).copy_(expert_scratch.view(-1))
189 +
190 + return weight.view(torch.uint8).view(block_layout_shape).view(torch.bfloat16)
191
192 return (
193 - w13_weights_shuffled_tensor.view(torch.bfloat16),
194 - w2_weights_shuffled_tensor.view(torch.bfloat16),
193 + _convert_weight_in_place(w13_weight, is_w13=True),
194 + _convert_weight_in_place(w2_weight, is_w13=False),
195 )
196
197
vllm/model_executor/model_loader/reload/meta.py
+2 -2
@@ -9,7 +9,7 @@ from torch.utils._python_dispatch import TorchDispatchMode
9
10 from .sanitize import restore_layer_refs, sanitize_layer_refs
11 from .types import LayerReloadingInfo, LayerTensors
12 -from .utils import get_layer_params_buffers, get_layer_tensors
12 +from .utils import get_layer_params_buffers, get_layer_tensors, get_tensor_load_numel
13
14 __all__ = [
15 "to_meta_tensor",
@@ -206,5 +206,5 @@ def get_numel_loaded(
206 numel = counter.copied_numel
207 param = args.arguments.get("param", None)
208 if isinstance(param, torch.Tensor):
209 - numel = min(numel, param.numel())
209 + numel = min(numel, get_tensor_load_numel(param))
210 return numel, return_value
vllm/model_executor/model_loader/reload/utils.py
+7 -1
@@ -11,6 +11,7 @@ __all__ = [
11 "get_layer_tensors",
12 "get_layer_params_buffers",
13 "get_layer_size",
14 + "get_tensor_load_numel",
15 "has_device_tensors",
16 "get_info_size",
17 ]
@@ -30,6 +31,11 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
31 )
32
33
34 +def get_tensor_load_numel(tensor: torch.Tensor) -> int:
35 + """Count checkpoint elements, excluding padding declared by the weight creator."""
36 + return getattr(tensor, "weight_loader_numel", tensor.numel())
37 +
38 +
39 def get_layer_size(layer: torch.nn.Module) -> int:
40 """Calculate total number of elements across loadable tensors in a layer.
41
@@ -39,7 +45,7 @@ def get_layer_size(layer: torch.nn.Module) -> int:
45 from .meta import SKIP_LOAD_TENSORS
46
47 return sum(
42 - tensor.numel()
48 + get_tensor_load_numel(tensor)
49 for name, tensor in get_layer_tensors(layer).items()
50 if name not in SKIP_LOAD_TENSORS
51 )