iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[MoE] Encapsulate TRT-LLM BF16 weight layout handling (#57405)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Misha Goin
2026-09-17 17:47:30 -04:00
committed by GitHub
parent db7a24c230
commit ac78445bc9
7 changed files with 138 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_flashinfer.py
+21 -9
View File
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.config import (
from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
FlashInferExperts,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
view_as_block_major_k,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
TrtLlmFp8ExpertsModular,
TrtLlmFp8ExpertsMonolithic,
@@ -570,6 +573,10 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(
padded = w2_shape[-1]
w13_ptr = layer.w13_weight.data_ptr()
w2_ptr = layer.w2_weight.data_ptr()
load_numel = (
layer.w13_weight.weight_loader_numel,
layer.w2_weight.weight_loader_numel,
)
for _ in range(2):
reloaded_w13 = torch.randn_like(layer.w13_weight)
@@ -602,19 +609,24 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(
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 (
layer.w13_weight.weight_loader_numel,
layer.w2_weight.weight_loader_numel,
) == load_numel
kernel_w13 = view_as_block_major_k(layer.w13_weight)
kernel_w2 = view_as_block_major_k(layer.w2_weight)
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):
def _check_flashinfer_ipc_weights(entries, expected, 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"
monkeypatch, 192, is_gated=True, device="meta"
)
loader = object.__new__(IpcModelLoader)
loader.mode = mode
@@ -625,7 +637,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde
with weights_already_processed():
method.process_weights_after_loading(layer)
assert method.moe_kernel is not None
actual = method._kernel_weights(layer)
actual = (
view_as_block_major_k(layer.w13_weight),
view_as_block_major_k(layer.w2_weight),
)
for weight, reference, pointer in zip(actual, expected, pointers):
assert weight.data_ptr() == pointer
assert torch.equal(weight.cpu(), reference)
@@ -636,11 +651,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde
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
monkeypatch, cache_ndim, mode
):
"""A spawned IPC consumer restores views without transient method state."""
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
@@ -649,13 +663,12 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(
from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry
method, layer = _make_unquantized_flashinfer_test_layer(
monkeypatch, 192, is_gated, device="meta"
monkeypatch, 192, is_gated=True, 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 = {}
@@ -667,7 +680,6 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(
args=(
entries,
expected,
is_gated,
mode,
torch.accelerator.current_device_index(),
),
tests/kernels/moe/test_moe_weight_loading_padded.py
-25
View File
@@ -12,7 +12,6 @@ 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,
)
@@ -465,30 +464,6 @@ class TestUnquantizedTrtLlmPrePadding:
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
tests/kernels/moe/test_trtllm_bf16_moe.py
+49
View File
@@ -12,6 +12,7 @@ import pytest
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from tests.kernels.utils import torch_moe
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -27,6 +28,9 @@ from vllm.model_executor.layers.fused_moe.config import (
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
TrtLlmBf16ExpertsModular,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
TrtLlmBf16LoRAExperts,
)
from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
UnquantizedMoeBackend,
convert_to_unquantized_kernel_format,
@@ -150,3 +154,48 @@ def test_trtllm_bf16_moe_modular_no_graph(
atol=1e-1,
rtol=2e-1,
)
@pytest.mark.parametrize("has_lora_delta", [False, True])
@torch.inference_mode()
def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta):
"""LoRA dispatch accepts both 3D parameters and legacy 4D packed views."""
config = make_dummy_moe_config(
num_experts=128, hidden_dim=256, intermediate_size=128
)
experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG)
x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10
w1, w2 = convert_to_unquantized_kernel_format(
UnquantizedMoeBackend.FLASHINFER_TRTLLM,
config,
torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10,
torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10,
)
assert w1.ndim == w2.ndim == 3
topk = (
torch.tensor([[0], [1]], device="cuda", dtype=torch.int32),
torch.ones(2, 1, device="cuda", dtype=torch.float32),
)
delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16)
def invoke(w1, w2):
result = experts.invoke_routed_moe(
hidden_states=x,
w1=w1,
w2=w2,
topk_ids_and_weights=topk,
gemm1_lora_delta=delta if has_lora_delta else None,
global_num_experts=128,
a1q_scale=None,
output=torch.empty_like(x),
)
if has_lora_delta:
# Only compare rows belonging to real tokens, excluding padding.
indices = result[2].flatten().long()
return result[0][indices], result[3][indices]
return result
actual = invoke(w1, w2)
expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64))
for actual_tensor, expected_tensor in zip(actual, expected):
torch.testing.assert_close(actual_tensor, expected_tensor)
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+13 -4
View File
@@ -31,6 +31,15 @@ from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
def view_as_block_major_k(weight: torch.Tensor) -> torch.Tensor:
"""View packed storage, including legacy 4D IPC cache entries."""
if weight.ndim == 4:
return weight
experts, rows, cols = weight.shape
block_k = 128 // weight.element_size()
return weight.view(experts, cols // block_k, rows, block_k)
class TrtLlmBf16ExpertsBase:
"""BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and
monolithic interfaces.
@@ -200,8 +209,8 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)
result = flashinfer.fused_moe.trtllm_bf16_routed_moe(
topk_ids=(topk_ids, topk_weights),
hidden_states=hidden_states,
gemm1_weights=w1,
gemm2_weights=w2,
gemm1_weights=view_as_block_major_k(w1),
gemm2_weights=view_as_block_major_k(w2),
num_experts=global_num_experts,
top_k=topk_ids.size(1),
n_group=None,
@@ -281,8 +290,8 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
routing_logits=router_logits,
routing_bias=e_score_correction_bias,
hidden_states=hidden_states,
gemm1_weights=w1,
gemm2_weights=w2,
gemm1_weights=view_as_block_major_k(w1),
gemm2_weights=view_as_block_major_k(w2),
num_experts=global_num_experts,
top_k=self.topk,
n_group=num_expert_group,
vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
+5 -2
View File
@@ -34,6 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRACon
from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
LoRAExpertsMixin,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
view_as_block_major_k,
)
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceNoOP,
)
@@ -527,8 +530,8 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(
topk_ids=topk_ids_and_weights,
hidden_states=hidden_states,
gemm1_weights=w1,
gemm2_weights=w2,
gemm1_weights=view_as_block_major_k(w1),
gemm2_weights=view_as_block_major_k(w2),
gemm1_lora_delta=gemm1_lora_delta,
num_experts=global_num_experts,
top_k=self.topk,
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+26 -1
View File
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
swap_w13_to_w31,
)
from vllm.platforms import current_platform
from vllm.utils.math_utils import round_up
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
@@ -349,6 +350,17 @@ def select_unquantized_moe_backend(
)
def unquantized_round_up_hidden_size_and_intermediate_size(
backend: UnquantizedMoeBackend,
hidden_size: int,
intermediate_size: int,
) -> tuple[int, int]:
"""Round up dimensions before allocation to satisfy the selected kernel."""
if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
intermediate_size = round_up(intermediate_size, 128)
return hidden_size, intermediate_size
def convert_to_unquantized_kernel_format(
unquantized_backend: UnquantizedMoeBackend,
moe_config: FusedMoEConfig,
@@ -377,13 +389,26 @@ def convert_to_unquantized_kernel_format(
)
moe_config.intermediate_size_per_partition = padded_intermediate
# Reloads only overwrite checkpoint slices. An earlier in-place
# permutation can leave nonzero values in the raw padding slots.
unpadded = moe_config.intermediate_size_per_partition_unpadded
assert unpadded is not None
if padded_intermediate > unpadded:
w13_weight[:, unpadded:padded_intermediate].zero_()
if is_act_and_mul:
w13_weight[:, padded_intermediate + unpadded :].zero_()
w2_weight[:, :, unpadded:].zero_()
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
convert_moe_weights_to_flashinfer_trtllm_block_layout(
_cache_permute_indices,
w13_weight,
w2_weight,
is_gated_act_gemm=is_act_and_mul,
)
# Keep checkpoint-shaped parameters for reload/IPC. The experts create
# BlockMajorK views at dispatch without changing the underlying storage.
return w13_weight, w2_weight
if (
unquantized_backend == UnquantizedMoeBackend.TRITON
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+24 -76
View File
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
convert_to_unquantized_kernel_format,
make_unquantized_moe_kernel,
select_unquantized_moe_backend,
unquantized_round_up_hidden_size_and_intermediate_size,
)
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
SharedExperts,
@@ -34,7 +35,6 @@ 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
@@ -72,9 +72,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
):
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:
@@ -83,22 +80,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
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_data = torch.empty(
num_experts,
w13_up_dim,
hidden_size,
dtype=params_dtype,
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w13_up_dim,
hidden_size,
dtype=params_dtype,
),
requires_grad=False,
)
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
@@ -111,15 +101,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
set_weight_attrs(w13_bias, extra_weight_attrs)
w13_bias.weight_loader_numel = num_experts * unpadded_up
# down_proj (row parallel)
w2_data = torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=params_dtype,
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=params_dtype,
),
requires_grad=False,
)
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 = (
@@ -147,11 +137,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
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
return unquantized_round_up_hidden_size_and_intermediate_size(
self.unquantized_backend, 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
@@ -177,21 +165,6 @@ 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,
@@ -199,10 +172,6 @@ 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
@@ -231,25 +200,6 @@ 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)
@@ -354,11 +304,10 @@ 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=w13,
w2=w2,
w1=layer.w13_weight,
w2=layer.w2_weight,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=layer.activation,
@@ -396,11 +345,10 @@ 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,
w13,
w2,
layer.w13_weight,
layer.w2_weight,
router_logits,
activation=layer.activation,
global_num_experts=layer.global_num_experts,