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)