iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
78300cdabffbaaffe7625ed34412ec2aec75535c
vLLM/tests/kernels/moe/test_moe.py
T
Misha GoinandOpenAI Codex 82b7d49a6e [MoE] Generalize masked activation for padded layouts (#51217) 
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-09-01 11:34:13 -04:00

2248 lines
71 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the MOE layers.
Run `pytest tests/kernels/test_moe.py`.
"""
import functools
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
import pytest
import torch
from torch.nn import Parameter
from torch.nn import functional as F
import vllm.model_executor.layers.fused_moe # noqa
import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module
from tests.kernels.moe.utils import (
fused_moe,
make_dummy_moe_config,
modular_triton_fused_moe,
)
from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import (
MoEActivation,
fused_topk,
)
from vllm.model_executor.layers.fused_moe.activation import (
ApplyMoEActivationConfig,
apply_moe_activation,
apply_moe_activation_supported,
)
from vllm.model_executor.layers.fused_moe.config import (
FUSED_MOE_UNQUANTIZED_CONFIG,
int4_w4a16_moe_quant_config,
int8_w8a16_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
batched_fused_marlin_moe,
fused_marlin_moe,
)
from vllm.model_executor.layers.fused_moe.utils import (
moe_use_td_hw_supported,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
marlin_permute_bias,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
rand_marlin_weight_mxfp4_like,
rand_marlin_weight_nvfp4_like,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
marlin_quant_fp8_torch,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
awq_marlin_quantize,
marlin_quantize,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
from vllm.platforms import current_platform
from vllm.scalar_type import ScalarType, scalar_types
from vllm.triton_utils import tl
from vllm.utils.math_utils import next_power_of_2
from vllm.utils.torch_utils import set_random_seed
DEVICE_TYPE = current_platform.device_type
def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None:
captured: dict[str, torch.Tensor] = {}
class FakeKernel:
def __getitem__(self, grid):
def launch(*args, **kwargs) -> None:
captured["a_scale"] = args[4]
return launch
monkeypatch.setattr(fused_moe_module, "fused_moe_kernel", FakeKernel())
a_scale = torch.tensor(0.5)
fused_moe_module.invoke_fused_moe_triton_kernel(
A=torch.ones((1, 1)),
B=torch.ones((1, 1, 1)),
C=torch.empty((1, 1, 1)),
A_scale=a_scale,
B_scale=torch.ones(1),
topk_weights=torch.ones((1, 1)),
sorted_token_ids=None,
expert_ids=torch.zeros(1, dtype=torch.int32),
num_tokens_post_padded=torch.ones(1, dtype=torch.int32),
mul_routed_weight=True,
top_k=1,
config={"BLOCK_SIZE_M": 1, "BLOCK_SIZE_N": 1, "BLOCK_SIZE_K": 1},
compute_type=tl.float32,
use_fp8_w8a8=True,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
)
captured_scale = captured["a_scale"]
assert a_scale.ndim == 0
assert captured_scale.shape == (1,)
assert captured_scale.data_ptr() == a_scale.data_ptr()
def iterative_moe(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
gating_output: torch.Tensor,
topk: int,
global_num_experts: int,
expert_map: torch.Tensor = None,
renormalize: bool = False,
) -> torch.Tensor:
"""
Baseline implementation of fused moe.
Args:
hidden_states: [*, hidden_size]
w1: [num_experts, intermediate_size * 2, hidden_size]
w2: [num_experts, hidden_size, intermediate_size]
gating_output: [*, num_experts]
expert_map: [num_experts]
"""
orig_shape = hidden_states.shape
hidden_size = hidden_states.shape[-1]
num_tokens = hidden_states.shape[:-1].numel()
num_experts = w1.shape[0]
intermediate_size = w2.shape[-1]
dtype = hidden_states.dtype
hidden_states = hidden_states.view(num_tokens, hidden_size)
gating_output = gating_output.view(num_tokens, global_num_experts)
topk_weights = gating_output.softmax(dim=-1, dtype=torch.float)
topk_weights, selected_experts = topk_weights.topk(topk, dim=-1)
if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
topk_weights = topk_weights.to(dtype)
if expert_map is not None:
selected_experts = expert_map[selected_experts]
final_hidden_states = None
for expert_idx in range(num_experts):
expert_w1 = w1[expert_idx]
expert_w2 = w2[expert_idx]
expert_mask = selected_experts == expert_idx
expert_weights = (topk_weights * expert_mask).sum(dim=-1, keepdim=True)
x = F.linear(hidden_states, expert_w1)
gate = F.silu(x[:, :intermediate_size])
x = x[:, intermediate_size:] * gate
x = F.linear(x, expert_w2)
current_hidden_states = x * expert_weights
if final_hidden_states is None:
final_hidden_states = current_hidden_states
else:
final_hidden_states = final_hidden_states + current_hidden_states
return final_hidden_states.view(orig_shape) # type: ignore
NUM_EXPERTS = [8, 64, 192]
NUM_EXPERTS_LARGE = [128, 256]
EP_SIZE = [1, 4]
TOP_KS = [2, 6]
TOP_KS_SMALL = [1, 2]
MOE_MARLIN_QUANT_TEST_CONFIGS = [
# AWQ-INT4
{"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]},
# GPTQ-INT4
{
"b_type": scalar_types.uint4b8,
"support_act_order": True,
"group_blocks": [-1, 2, 4, 8],
},
# GPTQ-INT8
{
"b_type": scalar_types.uint8b128,
"support_act_order": True,
"group_blocks": [-1, 2, 4, 8],
},
# FP8
{"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]},
# NVFP4
{"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]},
# MXFP4
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float4_e2m1f,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# MXFP8
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float8_e4m3fn,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# AWQ-INT4 with INT8 activation
{
"a_type": [scalar_types.int8],
"b_type": scalar_types.uint4,
"group_blocks": [-1, 2, 4, 8],
},
# GPTQ-INT4 with INT8 activation
{
"a_type": [scalar_types.int8],
"b_type": scalar_types.uint4b8,
"group_blocks": [-1, 2, 4, 8],
},
# GPTQ-INT4 with FP8 activation
{
"a_type": [scalar_types.float8_e4m3fn],
"b_type": scalar_types.uint4b8,
"group_blocks": [-1, 2, 4, 8],
},
# AWQ-INT4 with FP8 activation
{
"a_type": [scalar_types.float8_e4m3fn],
"b_type": scalar_types.uint4,
"group_blocks": [-1, 2, 4, 8],
},
# MXFP4 with FP8 activation
{
"a_type": [scalar_types.float8_e4m3fn],
"b_type": scalar_types.float4_e2m1f,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
]
FUSED_MOE_MNK_FACTORS = [
(1, 128, 128),
(1, 2048, 128),
(33, 2048, 128),
(32768, 2048, 511),
(40000, 1024, 1024),
]
FUSED_MOE_MNK_FACTORS_SMALL_M = [
(1, 128, 128),
(1, 2048, 128),
(2, 2048, 128),
(2, 2048, 511),
]
FUSED_MOE_WN16_MNK_FACTORS = [
(1, 128, 128),
(1, 1024, 1024),
(32, 2048, 128),
(222, 2048, 1024),
]
vllm_config = VllmConfig()
def run_moe_test(
baseline: Callable | torch.Tensor,
moe_fn: Callable,
a: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
score: torch.Tensor,
topk: int,
global_num_experts: int = -1,
expert_map: torch.Tensor | None = None,
padding: bool = False,
use_compile: bool = False,
use_cudagraph: bool = False,
atol: float = 2e-2,
rtol: float = 0,
) -> torch.Tensor:
if isinstance(baseline, torch.Tensor):
baseline_output = baseline
else:
baseline_output = baseline(
a,
w1,
w2,
score,
topk,
global_num_experts=global_num_experts,
expert_map=expert_map,
)
# Pad the weight if moe padding is enabled
if padding:
w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
if use_compile:
moe_fn = torch.compile(moe_fn, backend="inductor", fullgraph=True)
torch._dynamo.mark_dynamic(a, 0)
torch._dynamo.mark_dynamic(score, 0)
test_output = moe_fn(
a,
w1,
w2,
score,
topk,
global_num_experts=global_num_experts,
expert_map=expert_map,
)
if use_cudagraph:
test_output.fill_(0)
stream = torch.cuda.Stream()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=stream):
test_output = moe_fn(
a,
w1,
w2,
score,
topk,
global_num_experts=global_num_experts,
expert_map=expert_map,
)
torch.accelerator.synchronize()
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(test_output, baseline_output, atol=atol, rtol=rtol)
return baseline_output
@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS)
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
@pytest.mark.parametrize("ep_size", EP_SIZE)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("padding", [True, False])
@pytest.mark.parametrize("use_td", [False, True])
def test_fused_moe(
m: int,
n: int,
k: int,
e: int,
topk: int,
ep_size: int,
dtype: torch.dtype,
padding: bool,
use_td: bool,
monkeypatch,
workspace_init,
):
if use_td and not hasattr(tl, "make_tensor_descriptor"):
pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor")
if use_td and not moe_use_td_hw_supported():
pytest.skip(
"tensor_descriptor.gather requires XPU or NVIDIA Blackwell "
"(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas "
"rejects on Hopper (sm90) and earlier"
)
monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0")
set_random_seed(7)
#
# Setup test data
#
#
# Setup test data
#
a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype)
if ep_size > 1:
local_e = e // ep_size
e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32)
e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32)
e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32)
w1 = w1[e_ids]
w2 = w2[e_ids]
else:
e_map = None
#
# Setup test functions
#
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
def m_fused_moe(
a: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
score: torch.Tensor,
topk: int,
global_num_experts: int = -1,
expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
return m_fused_moe_fn.apply(
a,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
global_num_experts=global_num_experts,
expert_map=expert_map,
apply_router_weight_on_input=False,
)
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
#
# Run tests
#
runner = functools.partial(
run_moe_test,
a=a,
w1=w1,
w2=w2,
score=score,
topk=topk,
global_num_experts=e,
expert_map=e_map,
padding=padding,
)
# Note: for now use_compile will error out if the problem size is
# large enough to trigger chunking. I'm leaving the flag and
# setup code in case we are able to revisit this later.
use_compile = False
use_cudagraph = n >= 1024 and k >= 1024 and current_platform.is_cuda_alike()
with set_current_vllm_config(vllm_config):
baseline_output = runner(torch_moe, iterative_moe)
runner(
baseline_output,
fused_moe_fn,
use_compile=use_compile,
use_cudagraph=use_cudagraph,
)
runner(
baseline_output,
m_fused_moe,
use_compile=use_compile,
use_cudagraph=use_cudagraph,
)
def test_fused_shared_expert_alignment(workspace_init):
set_random_seed(7)
m, n, k = 4, 64, 128
routed_experts = 8
physical_experts = routed_experts + 1
dtype = torch.bfloat16
a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
w1 = torch.randn((physical_experts, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
w2 = torch.randn((physical_experts, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
topk_ids = torch.tensor(
[[0, 8], [1, 8], [2, 8], [3, 8]], device=DEVICE_TYPE, dtype=torch.int32
)
topk_weights = torch.tensor(
[[0.5, 1.0]] * m, device=DEVICE_TYPE, dtype=torch.float32
)
moe_config = make_dummy_moe_config(
num_experts=physical_experts,
experts_per_token=2,
hidden_dim=k,
intermediate_size=n,
in_dtype=dtype,
max_num_tokens=m,
)
modular_moe = modular_triton_fused_moe(moe_config, FUSED_MOE_UNQUANTIZED_CONFIG)
with set_current_vllm_config(vllm_config):
expected = torch_experts(a, w1, w2, topk_weights, topk_ids)
actual = modular_moe.apply(
hidden_states=a,
w1=w1,
w2=w2,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SILU,
global_num_experts=routed_experts,
expert_map=None,
apply_router_weight_on_input=False,
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
def test_fused_moe_int64_overflow(workspace_init):
"""Regression test for int32 overflow in stride*offset products.
With large M, stride_cm * offs_token can exceed int32 max. Verifies
the offs_token int64 cast (fix for #34413) prevents overflow and
produces correct results.
Reproduces the scenario from PR #34279.
"""
# ~12 GB GPU memory needed for intermediate caches
free_mem = torch.accelerator.get_memory_info()[0]
if free_mem < 12 * 1024**3:
pytest.skip("Insufficient GPU memory for overflow test")
set_random_seed(7)
m, n, k, e, topk = 100000, 2048, 1024, 8, 6
dtype = torch.bfloat16
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
score = torch.randn((m, e), device="cuda", dtype=dtype)
# Verify the test exercises the overflow condition:
# C has shape (M, topk, N) where N = w1.size(1) = 2*n
# stride_cm = C.stride(1) = N, max offs_token = M * topk
# Product must exceed int32 max for this test to be meaningful
N = w1.size(1)
assert N * m * topk > 2**31 - 1, "Test params don't trigger int32 overflow"
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
with set_current_vllm_config(vllm_config):
run_moe_test(
torch_moe,
fused_moe_fn,
a=a,
w1=w1,
w2=w2,
score=score,
topk=topk,
global_num_experts=e,
)
@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS_SMALL_M)
@pytest.mark.parametrize("e", NUM_EXPERTS_LARGE)
@pytest.mark.parametrize("topk", TOP_KS_SMALL)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("padding", [True, False])
def test_naive_block_assignment_moe(
m: int,
n: int,
k: int,
e: int,
topk: int,
dtype: torch.dtype,
padding: bool,
monkeypatch,
workspace_init,
):
set_random_seed(7)
#
# Setup test data
#
#
# Setup test data
#
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
score = torch.randn((m, e), device="cuda", dtype=dtype)
e_map = None
#
# Setup test functions
#
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
def m_fused_moe(
a: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
score: torch.Tensor,
topk: int,
global_num_experts: int = -1,
expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
return m_fused_moe_fn.apply(
a,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
global_num_experts=global_num_experts,
expert_map=expert_map,
apply_router_weight_on_input=False,
)
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
#
# Run tests
#
runner = functools.partial(
run_moe_test,
a=a,
w1=w1,
w2=w2,
score=score,
topk=topk,
global_num_experts=e,
expert_map=e_map,
padding=padding,
)
# Note: for now use_compile will error out if the problem size is
# large enough to trigger chunking. I'm leaving the flag and
# setup code in case we are able to revisit this later.
use_compile = False
use_cudagraph = n >= 1024 and k >= 1024 and current_platform.is_cuda_alike()
with set_current_vllm_config(vllm_config):
baseline_output = runner(torch_moe, iterative_moe)
runner(
baseline_output,
fused_moe_fn,
use_compile=use_compile,
use_cudagraph=use_cudagraph,
)
runner(
baseline_output,
m_fused_moe,
use_compile=use_compile,
use_cudagraph=use_cudagraph,
)
@pytest.mark.parametrize("m,n,k", FUSED_MOE_WN16_MNK_FACTORS)
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
@pytest.mark.parametrize("ep_size", EP_SIZE)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.parametrize("has_zp", [True, False])
@pytest.mark.parametrize("weight_bits", [4, 8])
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="MoE WNA16 Triton kernels require CUDA/ROCm or XPU.",
)
def test_fused_moe_wn16(
m: int,
n: int,
k: int,
e: int,
topk: int,
ep_size: int,
dtype: torch.dtype,
group_size: int,
has_zp: bool,
weight_bits: int,
):
a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype)
if weight_bits == 4:
pack_factor = 2
quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
elif weight_bits == 8:
pack_factor = 1
quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128
w1_ref = w1.clone()
w2_ref = w2.clone()
w1_qweight = torch.empty(
(e, 2 * n, k // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8
)
w2_qweight = torch.empty(
(e, k, n // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8
)
w1_scales = torch.empty(
(e, 2 * n, k // group_size), device=DEVICE_TYPE, dtype=dtype
)
w2_scales = torch.empty((e, k, n // group_size), device=DEVICE_TYPE, dtype=dtype)
w1_qzeros = torch.empty(
(e, 2 * n // pack_factor, k // group_size),
device=DEVICE_TYPE,
dtype=torch.uint8,
)
w2_qzeros = torch.empty(
(e, k // pack_factor, n // group_size), device=DEVICE_TYPE, dtype=torch.uint8
)
for i in range(e * 2):
expert_id = i % e
if i // e == 0:
w, w_ref, w_qweight, w_scales, w_qzeros = (
w1,
w1_ref,
w1_qweight,
w1_scales,
w1_qzeros,
)
else:
w, w_ref, w_qweight, w_scales, w_qzeros = (
w2,
w2_ref,
w2_qweight,
w2_scales,
w2_qzeros,
)
weight, qweight, scales, qzeros = quantize_weights(
w[expert_id].T, quant_type, group_size, has_zp, False
)
weight = weight.T
qweight = qweight.T.contiguous().to(torch.uint8)
scales = scales.T
if has_zp:
qzeros = qzeros.T.contiguous().to(torch.uint8)
if weight_bits == 4:
qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]
if has_zp:
qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]
w_ref[expert_id] = weight
w_qweight[expert_id] = qweight
w_scales[expert_id] = scales
if has_zp:
w_qzeros[expert_id] = qzeros
if ep_size > 1:
local_e = e // ep_size
e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32)
e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32)
e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32)
w1_ref = w1_ref[e_ids]
w2_ref = w2_ref[e_ids]
w1_qweight = w1_qweight[e_ids]
w2_qweight = w2_qweight[e_ids]
w1_scales = w1_scales[e_ids]
w2_scales = w2_scales[e_ids]
w1_qzeros = w1_qzeros[e_ids]
w2_qzeros = w2_qzeros[e_ids]
else:
e_map = None
if weight_bits == 4:
quant_config_builder = int4_w4a16_moe_quant_config
else:
assert weight_bits == 8
quant_config_builder = int8_w8a16_moe_quant_config
quant_config = quant_config_builder(
w1_scale=w1_scales,
w2_scale=w2_scales,
w1_zp=w1_qzeros if has_zp else None,
w2_zp=w2_qzeros if has_zp else None,
block_shape=[0, group_size],
)
with set_current_vllm_config(vllm_config):
triton_output = fused_moe(
a,
w1_qweight,
w2_qweight,
score,
topk,
renormalize=False,
global_num_experts=e,
expert_map=e_map,
quant_config=quant_config,
)
torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, expert_map=e_map)
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
MARLIN_MOE_SCENARIOS = [
# (m, n, k, e, topk, ep_size, act_order, is_k_full)
# No act_order: is_k_full=True matches usual case (marlin_is_k_full).
# N>=256 required for Marlin kernel thread config for MXFP8.
# Single token, small matrices
(1, 128, 256, 5, 2, 1, False, True),
# Single token, large matrices
(1, 1024, 2048, 5, 2, 1, False, True),
# Unaligned m, small matrices
(133, 256, 256, 5, 2, 1, False, True),
# Unaligned m, large matrices
(133, 1024, 2048, 12, 3, 1, False, True),
# Aligned batch, small matrices
(128, 256, 256, 5, 2, 1, False, True),
# Aligned batch, large matrices
(128, 1024, 2048, 12, 3, 1, False, True),
# Expert parallelism
(64, 1024, 2048, 12, 3, 4, False, True),
# Act order with is_k_full=True (no tensor parallelism)
(1, 1024, 2048, 5, 2, 1, True, True),
# Act order with is_k_full=False (tensor parallelism)
(133, 256, 256, 5, 2, 1, True, False),
]
def marlin_moe_generate_valid_test_cases():
import itertools
def is_valid(
a_type,
b_type,
c_type,
group_blocks,
m,
n,
k,
e,
topk,
ep_size,
act_order,
is_k_full,
):
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
if group_size > 0 and k % group_size != 0:
return False
if act_order and group_size in [-1, k, n]:
return False
if group_size in [k, n]:
return False
if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
return False
return a_type.size_bits < 16 or a_type is c_type
cases = []
for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
f16_types = [scalar_types.float16]
inner_combinations = list(
itertools.product(
quant_test_config.get("a_type", f16_types),
[quant_test_config["b_type"]],
quant_test_config.get("c_type", f16_types),
quant_test_config["group_blocks"],
)
)
supports_act_order = quant_test_config.get("support_act_order", False)
for sub_case in inner_combinations:
if (
sub_case[0] == scalar_types.float8_e4m3fn
and not current_platform.is_device_capability(89)
and not current_platform.is_device_capability_family(120)
):
continue
for scenario in MARLIN_MOE_SCENARIOS:
m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
if act_order and not supports_act_order:
continue
args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
if is_valid(*args):
cases.append(args)
return cases
@dataclass
class MarlinMoEWeightData:
w_ref: torch.Tensor
qweight: torch.Tensor
scales: torch.Tensor
global_scale: torch.Tensor | None
a_scales_factor: torch.Tensor | None
g_idx: torch.Tensor | None
zeros: torch.Tensor | None
sort_indices: torch.Tensor | None
marlin_bias: torch.Tensor | None
@staticmethod
def make(
w: torch.Tensor,
quant_type: ScalarType,
group_size: int,
act_order: bool | None = None,
bias: torch.Tensor | None = None,
input_type: ScalarType = None,
) -> "MarlinMoEWeightData":
assert w.ndim == 3
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
k = w.shape[-1]
if input_type == scalar_types.int8:
input_dtype = torch.int8
elif input_type == scalar_types.float8_e4m3fn:
input_dtype = torch.float8_e4m3fn
else:
input_dtype = w.dtype
w_ref_l: list[torch.Tensor] = []
qweight_l: list[torch.Tensor] = []
scales_l: list[torch.Tensor] = []
global_scale_l: list[torch.Tensor] = []
zeros_l: list[torch.Tensor] = []
g_idx_l: list[torch.Tensor] = []
sort_indices_l: list[torch.Tensor] = []
bias_l: list[torch.Tensor] = []
for i in range(w.shape[0]):
if quant_type == scalar_types.float4_e2m1f:
if group_size == 16:
w_ref, qweight, scales, global_scale = (
rand_marlin_weight_nvfp4_like(
w[i], group_size, input_dtype=input_dtype
)
)
else:
w_ref, qweight, scales = rand_marlin_weight_mxfp4_like(
w[i], group_size, input_dtype=input_dtype
)
global_scale = None
w_ref_l.append(w_ref.T)
qweight_l.append(qweight)
scales_l.append(scales)
if global_scale is not None:
global_scale_l.append(global_scale)
elif quant_type == scalar_types.float8_e4m3fn:
w_ref, qweight, scales = marlin_quant_fp8_torch(
w[i], group_size, input_dtype=input_dtype
)
w_ref_l.append(w_ref.T)
qweight_l.append(qweight)
scales_l.append(scales)
elif has_zp:
w_ref, qweight, scales, zeros = awq_marlin_quantize(
w[i].transpose(1, 0),
quant_type,
group_size,
input_dtype=input_dtype,
)
w_ref_l.append(w_ref.T)
qweight_l.append(qweight)
scales_l.append(scales)
zeros_l.append(zeros)
else:
test_perm = torch.randperm(k)
w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
w[i].transpose(1, 0),
quant_type,
group_size,
act_order,
test_perm,
input_dtype=input_dtype,
)
w_ref_l.append(w_ref.T)
qweight_l.append(qweight)
scales_l.append(scales)
g_idx_l.append(g_idx)
sort_indices_l.append(sort_indices)
if bias is not None:
bias_l.append(marlin_permute_bias(bias[i]))
w_ref = stack_and_dev(w_ref_l)
qweight = stack_and_dev(qweight_l).contiguous()
scales = stack_and_dev(scales_l)
global_scale = stack_and_dev(global_scale_l) if global_scale_l else None
g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
zeros = stack_and_dev(zeros_l) if zeros_l else None
sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
marlin_bias = stack_and_dev(bias_l) if bias_l else None
a_scales_factor = None
if input_type == scalar_types.int8 and group_size != -1:
a_scales_factor = 1 / 4096 * scales.max().float()
scales = scales / scales.max() * 4096
scales = scales.round().to(torch.int16).view(w.dtype)
return MarlinMoEWeightData(
w_ref=w_ref,
qweight=qweight,
scales=scales,
global_scale=global_scale,
a_scales_factor=a_scales_factor,
g_idx=g_idx,
zeros=zeros,
sort_indices=sort_indices,
marlin_bias=marlin_bias,
)
@pytest.mark.flaky(reruns=2)
@pytest.mark.parametrize(
(
"a_type, b_type, c_type, group_blocks,"
"m, n, k, e, topk, ep_size, act_order, is_k_full"
),
marlin_moe_generate_valid_test_cases(),
)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
@pytest.mark.usefixtures("default_vllm_config")
def test_fused_marlin_moe(
a_type: ScalarType,
b_type: ScalarType,
c_type: ScalarType,
group_blocks: int,
m: int,
n: int,
k: int,
e: int,
topk: int,
ep_size: int,
act_order: bool,
is_k_full: bool,
):
set_random_seed(1)
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
if c_type == scalar_types.float16:
dtype = torch.float16
elif c_type == scalar_types.bfloat16:
dtype = torch.bfloat16
else:
raise RuntimeError("unsupported c_type")
if a_type == scalar_types.int8:
a_dtype = torch.int8
elif a_type == scalar_types.float8_e4m3fn:
a_dtype = torch.float8_e4m3fn
else:
a_dtype = dtype
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
if ep_size > 1:
local_e = e // ep_size
e_ids = torch.randperm(e, device="cuda", dtype=torch.int32)[:local_e]
e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32)
e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
w1 = w1[e_ids]
w2 = w2[e_ids]
else:
e_map = None
w1_data = MarlinMoEWeightData.make(
w=w1,
quant_type=b_type,
group_size=group_size,
act_order=act_order,
input_type=a_type,
)
w2_data = MarlinMoEWeightData.make(
w=w2,
quant_type=b_type,
group_size=group_size,
act_order=act_order,
input_type=a_type,
)
score = torch.randn((m, e), device="cuda", dtype=dtype)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
with set_current_vllm_config(vllm_config):
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
torch_output = torch_experts(
a,
w1_data.w_ref,
w2_data.w_ref,
topk_weight=topk_weight,
topk_ids=topk_ids,
global_num_experts=e,
expert_map=e_map,
quant_dtype=a_dtype,
per_act_token_quant=True,
)
def instance_activation(
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
*,
topk_ids: torch.Tensor | None = None,
expert_map: torch.Tensor | None = None,
) -> None:
apply_moe_activation(
activation,
output,
input,
activation_config=ApplyMoEActivationConfig(),
topk_ids=topk_ids,
expert_map=expert_map,
)
marlin_output = fused_marlin_moe(
a,
w1_data.qweight,
w2_data.qweight,
None,
None,
w1_data.scales,
w2_data.scales,
topk_weights,
topk_ids,
global_num_experts=e,
expert_map=e_map,
global_scale1=w1_data.global_scale,
global_scale2=w2_data.global_scale,
g_idx1=w1_data.g_idx,
g_idx2=w2_data.g_idx,
input_global_scale1=w1_data.a_scales_factor,
input_global_scale2=w2_data.a_scales_factor,
sort_indices1=w1_data.sort_indices,
sort_indices2=w2_data.sort_indices,
w1_zeros=w1_data.zeros,
w2_zeros=w2_data.zeros,
input_dtype=a_dtype,
quant_type_id=b_type.id,
is_k_full=is_k_full,
activation_func=instance_activation,
)
torch.testing.assert_close(marlin_output, torch_output, atol=4e-2, rtol=0)
@pytest.mark.flaky(reruns=2)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
@pytest.mark.usefixtures("default_vllm_config")
@pytest.mark.parametrize("m", [1, 256])
def test_fused_marlin_moe_with_bias(m):
set_random_seed(0)
e, topk = 32, 4
n, k = 2048, 2048
group_size = 128
act_order = False
is_k_full = True
quant_type = scalar_types.uint4b8
dtype = torch.half
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
b_bias1 = torch.randn((e, 2 * n), device="cuda", dtype=dtype) / 10
b_bias2 = torch.randn((e, k), device="cuda", dtype=dtype) / 10
w1_data = MarlinMoEWeightData.make(
w=w1,
quant_type=quant_type,
group_size=group_size,
act_order=act_order,
bias=b_bias1,
)
w2_data = MarlinMoEWeightData.make(
w=w2,
quant_type=quant_type,
group_size=group_size,
act_order=act_order,
bias=b_bias2,
)
score = torch.randn((m, e), device="cuda", dtype=dtype)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
with set_current_vllm_config(vllm_config):
torch_output = torch_moe(
a, w1_data.w_ref, w2_data.w_ref, score, topk, b_bias1, b_bias2
)
marlin_output = fused_marlin_moe(
a,
w1_data.qweight,
w2_data.qweight,
w1_data.marlin_bias,
w2_data.marlin_bias,
w1_data.scales,
w2_data.scales,
topk_weights,
topk_ids,
global_num_experts=e,
expert_map=None,
global_scale1=w1_data.global_scale,
global_scale2=w2_data.global_scale,
g_idx1=w1_data.g_idx,
g_idx2=w2_data.g_idx,
sort_indices1=w1_data.sort_indices,
sort_indices2=w2_data.sort_indices,
w1_zeros=w1_data.zeros,
w2_zeros=w2_data.zeros,
quant_type_id=quant_type.id,
is_k_full=is_k_full,
)
torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
@pytest.mark.flaky(reruns=2)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
@pytest.mark.usefixtures("default_vllm_config")
@pytest.mark.parametrize("m", [1, 64, 256])
@pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])
@pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])
@pytest.mark.parametrize("activation", [MoEActivation.RELU2_NO_MUL])
def test_fused_marlin_moe_non_gated(
m: int, n: int, k: int, e: int, topk: int, activation: MoEActivation
):
"""Test Marlin MoE with non-gated activation (relu2_no_mul).
Non-gated activations like relu2 don't have the gate-up projection pattern,
so w1 has shape (e, n, k) instead of (e, 2*n, k).
"""
set_random_seed(42)
group_size = 16 # NVFP4 group size
is_k_full = True
quant_type = scalar_types.float4_e2m1f
dtype = torch.bfloat16
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
# Non-gated: w1 shape is (e, n, k) not (e, 2*n, k)
w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
w1_data = MarlinMoEWeightData.make(
w=w1,
quant_type=quant_type,
group_size=group_size,
act_order=False,
)
w2_data = MarlinMoEWeightData.make(
w=w2,
quant_type=quant_type,
group_size=group_size,
act_order=False,
)
score = torch.randn((m, e), device="cuda", dtype=dtype)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
with set_current_vllm_config(vllm_config):
torch_output = torch_moe(
a,
w1_data.w_ref,
w2_data.w_ref,
score,
topk,
activation=activation,
)
marlin_output = fused_marlin_moe(
a,
w1_data.qweight,
w2_data.qweight,
None, # bias1
None, # bias2
w1_data.scales,
w2_data.scales,
topk_weights,
topk_ids,
global_num_experts=e,
expert_map=None,
global_scale1=w1_data.global_scale,
global_scale2=w2_data.global_scale,
g_idx1=w1_data.g_idx,
g_idx2=w2_data.g_idx,
sort_indices1=w1_data.sort_indices,
sort_indices2=w2_data.sort_indices,
w1_zeros=w1_data.zeros,
w2_zeros=w2_data.zeros,
quant_type_id=quant_type.id,
is_k_full=is_k_full,
activation=activation,
)
torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)
def _make_humming_indexed_experts(activation: MoEActivation):
pytest.importorskip("humming")
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
HummingIndexedExperts,
)
from vllm.model_executor.layers.quantization.utils import humming_utils
from vllm.utils import humming
top_k, num_experts = 6, 12
hidden_size, intermediate_size = 2688, 1856
gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size
num_w13_stacks = 2 if activation.is_gated else 1
moe_config = make_dummy_moe_config(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
activation=activation,
)
layer = torch.nn.Module()
layer.moe_config = moe_config
layer.params_dtype = torch.bfloat16
weight_schema = humming.ModeloptNvfp4WeightSchema()
for sublayer_name, shape_n, shape_k, stack_size in (
("w13", gate_up_size, hidden_size, num_w13_stacks),
("w2", hidden_size, intermediate_size, 1),
):
tensor_attrs = weight_schema.get_tensors_attrs(
shape_n=shape_n,
shape_k=shape_k,
param_dtype=layer.params_dtype,
num_experts=num_experts,
stack_size=stack_size,
)
for tensor_name, attrs in tensor_attrs.items():
layer.register_parameter(
f"{sublayer_name}_{tensor_name}",
Parameter(
torch.ones(
attrs["shape"],
dtype=attrs["dtype"],
device="cuda",
),
requires_grad=False,
),
)
humming_utils.convert_to_humming_moe_kernel_format(
layer,
weight_schema=weight_schema,
input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16),
)
layer.local_num_experts = layer.global_num_experts = num_experts
layer.hidden_size = hidden_size
layer.intermediate_size_per_partition = intermediate_size
quant_config = humming_utils.get_humming_moe_quant_config(layer)
experts = HummingIndexedExperts(
moe_config=moe_config,
quant_config=quant_config,
)
return experts, layer
@pytest.mark.parametrize("activation", list(MoEActivation))
def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation):
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
HummingExpertsBase,
)
assert HummingExpertsBase._supports_activation(
activation
) == apply_moe_activation_supported(activation)
@pytest.mark.parametrize(
("env_override", "use_ep", "expected"),
[
(None, False, "indexed"),
(None, True, "grouped_contiguous"),
("auto", True, "grouped_contiguous"),
("indexed", True, "indexed"),
("grouped", False, "grouped_contiguous"),
],
)
def test_humming_selects_gemm_from_parallelism_and_override(
monkeypatch: pytest.MonkeyPatch,
env_override: str | None,
use_ep: bool,
expected: str,
):
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
get_humming_moe_gemm_type,
)
if env_override is None:
monkeypatch.delenv("VLLM_HUMMING_MOE_GEMM_TYPE", raising=False)
else:
monkeypatch.setenv("VLLM_HUMMING_MOE_GEMM_TYPE", env_override)
moe_config = make_dummy_moe_config()
moe_config.moe_parallel_config.use_ep = use_ep
assert get_humming_moe_gemm_type(moe_config) == expected
@pytest.mark.parametrize(
("dp_token_counts", "topk_shape", "expected"),
[
pytest.param(None, (7, 3), 21, id="runtime-layout"),
pytest.param([3, 5], (48, 1), 48, id="expanded-dp-layout"),
],
)
def test_humming_global_valid_shape_m(
monkeypatch: pytest.MonkeyPatch,
dp_token_counts: list[int] | None,
topk_shape: tuple[int, int],
expected: int,
):
from types import SimpleNamespace
import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
dp_metadata = (
None
if dp_token_counts is None
else SimpleNamespace(
num_tokens_across_dp_cpu=torch.tensor(dp_token_counts, dtype=torch.int32)
)
)
monkeypatch.setattr(
humming,
"get_forward_context",
lambda: SimpleNamespace(dp_metadata=dp_metadata),
)
experts = SimpleNamespace(moe_config=make_dummy_moe_config(experts_per_token=6))
topk_ids = torch.empty(topk_shape, dtype=torch.int64)
result = humming.HummingExpertsBase.get_global_valid_shape_m(experts, topk_ids)
assert result == expected
def test_humming_permute_scratch_is_keyed_by_runtime_topk(
monkeypatch: pytest.MonkeyPatch,
):
from types import SimpleNamespace
from unittest.mock import Mock
import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
scratch_topk6 = Mock()
scratch_topk1 = Mock()
scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])
monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)
monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type)
moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)
moe_config.moe_parallel_config.dp_size = 2
experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config)
assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1
assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
assert scratch_type.call_count == 2
first_call, second_call = scratch_type.call_args_list
assert first_call.kwargs["max_num_tokens"] == 1024
assert first_call.kwargs["topk"] == 6
assert second_call.kwargs["max_num_tokens"] == 6144
assert second_call.kwargs["topk"] == 1
def test_humming_delegates_to_instance_activation():
from types import SimpleNamespace
from unittest.mock import Mock
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
HummingExpertsBase,
)
activation_func = Mock()
activation_config = ApplyMoEActivationConfig(
clamp_limit=7.0,
alpha=1.5,
beta=0.25,
activation_situ_beta=2.0,
activation_situ_linear_beta=3.0,
)
experts = SimpleNamespace(
activation=activation_func,
activation_config=activation_config,
)
input = torch.empty(1, 2)
output = torch.empty(1, 1)
HummingExpertsBase.apply_activation(
experts, MoEActivation.SWIGLUOAI_UNINTERLEAVE, output, input
)
activation_func.assert_called_once_with(
activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
input=input,
output=output,
)
def test_humming_grouped_apply_forwards_valid_prefix(
monkeypatch: pytest.MonkeyPatch,
):
from types import MethodType, SimpleNamespace
from unittest.mock import Mock
import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
HummingExpertsBase,
HummingGroupedExperts,
)
activation_func = Mock()
buffers = {
"gate_up_output": torch.empty(6, 4),
"activation_output": torch.empty(6, 2),
"down_output": torch.empty(6, 4),
}
experts = SimpleNamespace(
activation=activation_func,
num_experts=2,
estimate_local_valid_shape_m=lambda _: 6,
prepare_buffers=lambda *_: buffers,
_get_permute_scratch=lambda _: None,
quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: (
inputs,
input_scale,
),
humming_forward=Mock(),
compute_config_str="",
w13_tuning_config_str="",
w2_tuning_config_str="",
)
experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts)
hidden_states = torch.empty(3, 4)
output = torch.empty_like(hidden_states)
topk_ids = torch.zeros((3, 2), dtype=torch.int64)
topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
expected_counts = torch.tensor([4], dtype=torch.int32)
expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64)
monkeypatch.setattr(
humming,
"moe_permute",
lambda **_: (
torch.empty(6, 4),
None,
expert_offsets,
torch.arange(6),
None,
),
)
monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None)
HummingGroupedExperts.apply(
experts,
output=output,
hidden_states=hidden_states,
w1=torch.empty(0),
w2=torch.empty(0),
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SITU,
global_num_experts=2,
expert_map=None,
a1q_scale=None,
a2_scale=None,
workspace13=torch.empty(0),
workspace2=torch.empty(0),
expert_tokens_meta=None,
apply_router_weight_on_input=False,
)
activation_func.assert_called_once()
call_kwargs = activation_func.call_args.kwargs
assert call_kwargs["activation"] == MoEActivation.SITU
assert call_kwargs["input"].shape == (6, 4)
assert call_kwargs["output"].shape == (6, 2)
torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts)
def test_batched_marlin_activation_uses_expert_token_counts(
monkeypatch: pytest.MonkeyPatch,
):
from types import SimpleNamespace
from unittest.mock import Mock
import vllm.model_executor.layers.fused_moe.experts.marlin_moe as marlin
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
BatchedMarlinExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import (
ExpertTokensMetadata,
)
activation_func = Mock()
experts = SimpleNamespace(
activation=activation_func,
w1_bias=None,
w2_bias=None,
w1_scale=None,
w2_scale=None,
quant_type_id=0,
a1_gscale=None,
a2_gscale=None,
w13_g_idx=None,
w2_g_idx=None,
w13_g_idx_sort_indices=None,
w2_g_idx_sort_indices=None,
w1_zp=None,
w2_zp=None,
input_dtype=None,
is_k_full=True,
activation_config=ApplyMoEActivationConfig(
clamp_limit=3.0, alpha=1.3, beta=0.5
),
)
hidden_states = torch.empty(2, 3, 4)
expert_num_tokens = torch.tensor([1, 2], dtype=torch.int32)
topk_ids = torch.zeros((2, 3), dtype=torch.int64)
expert_map = torch.arange(2, dtype=torch.int32)
def fake_batched_fused_marlin_moe(**kwargs):
kwargs["activation_func"](
kwargs["activation"],
torch.empty(6, 4),
torch.empty(6, 8),
topk_ids=topk_ids,
expert_map=expert_map,
)
return kwargs["output"]
monkeypatch.setattr(
marlin, "batched_fused_marlin_moe", fake_batched_fused_marlin_moe
)
output = torch.empty_like(hidden_states)
BatchedMarlinExperts.apply(
experts,
output=output,
hidden_states=hidden_states,
w1=torch.empty(0),
w2=torch.empty(0),
topk_weights=torch.ones_like(topk_ids, dtype=torch.float32),
topk_ids=topk_ids,
activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
global_num_experts=2,
expert_map=expert_map,
a1q_scale=None,
a2_scale=None,
workspace13=torch.empty(0),
workspace2=torch.empty(0),
expert_tokens_meta=ExpertTokensMetadata(expert_num_tokens, None),
apply_router_weight_on_input=False,
)
activation_func.assert_called_once()
call = activation_func.call_args
assert call.args[0] == MoEActivation.SWIGLUOAI_UNINTERLEAVE
assert call.args[1].shape == (2, 3, 4)
assert call.args[2].shape == (2, 3, 8)
torch.testing.assert_close(call.kwargs["valid_token_counts"], expert_num_tokens)
@pytest.mark.parametrize(
"activation",
[
MoEActivation.SILU,
MoEActivation.RELU2_NO_MUL,
],
ids=["gated", "non_gated"],
)
def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
from vllm.utils import humming
experts, layer = _make_humming_indexed_experts(activation)
moe_config = experts.moe_config
top_k = moe_config.experts_per_token
num_experts = moe_config.num_experts
hidden_size = moe_config.hidden_dim
intermediate_size = moe_config.intermediate_size
gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size
w13_meta, w2_meta = (experts.humming_configs[name] for name in ("w13", "w2"))
for meta in (w13_meta, w2_meta):
assert meta.a_dtype == humming.dtypes.bfloat16
assert meta.b_dtype == humming.dtypes.float4e2m1
assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size
assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size
buffer_metas, _ = experts.get_buffer_metas(
M=1,
topk=top_k,
activation=moe_config.activation,
)
assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size
assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size
assert experts.moe_problem_size(
a1=torch.empty(1, hidden_size),
w1=torch.empty(num_experts, 1),
w2=torch.empty(num_experts, 1),
topk_ids=torch.empty(1, top_k, dtype=torch.long),
) == (num_experts, 1, intermediate_size, hidden_size, top_k)
def test_humming_indexed_writes_supplied_output_buffer():
from vllm.forward_context import set_forward_context
activation = MoEActivation.SILU
experts, layer = _make_humming_indexed_experts(activation)
moe_config = experts.moe_config
num_tokens = 1
top_k = moe_config.experts_per_token
hidden_size = moe_config.hidden_dim
num_experts = moe_config.num_experts
workspace13_shape, workspace2_shape, _ = experts.workspace_shapes(
M=num_tokens,
N=moe_config.intermediate_size,
K=hidden_size,
topk=top_k,
global_num_experts=num_experts,
local_num_experts=num_experts,
expert_tokens_meta=None,
activation=activation,
)
device = torch.device("cuda")
dtype = layer.params_dtype
workspace13 = torch.empty(workspace13_shape, dtype=dtype, device=device)
workspace2 = torch.empty(workspace2_shape, dtype=dtype, device=device)
hidden_states = torch.ones((num_tokens, hidden_size), dtype=dtype, device=device)
output = torch.full_like(hidden_states, torch.nan)
topk_weights = torch.full(
(num_tokens, top_k),
1 / top_k,
dtype=dtype,
device=device,
)
topk_ids = torch.arange(top_k, dtype=torch.int32, device=device).unsqueeze(0)
with set_forward_context(None, vllm_config, num_tokens=num_tokens):
experts.apply(
output=output,
hidden_states=hidden_states,
w1=layer.w13_weight,
w2=layer.w2_weight,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=activation,
global_num_experts=num_experts,
expert_map=None,
a1q_scale=None,
a2_scale=None,
workspace13=workspace13,
workspace2=workspace2,
expert_tokens_meta=None,
apply_router_weight_on_input=False,
)
assert torch.isfinite(output).all()
@pytest.mark.parametrize("ep_size", [1, 2])
def test_moe_align_block_size_opcheck(ep_size):
num_experts = 4
block_size = 4
expert_map = None
if ep_size != 1:
local_num_experts = num_experts // ep_size
expert_ids = torch.randint(
0, num_experts, (local_num_experts,), device="cuda", dtype=torch.int32
)
expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
expert_map[expert_ids] = torch.arange(
local_num_experts, device="cuda", dtype=torch.int32
)
topk_ids = torch.randint(0, num_experts, (3, 4), dtype=torch.int32, device="cuda")
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
sorted_ids = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids.fill_(topk_ids.numel())
max_num_m_blocks = max_num_tokens_padded // block_size
expert_ids = torch.empty(
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
opcheck(
torch.ops._moe_C.moe_align_block_size,
(
topk_ids,
num_experts,
block_size,
sorted_ids,
expert_ids,
num_tokens_post_pad,
expert_map,
),
)
def test_batched_moe_align_block_size_opcheck():
max_tokens_per_batch = 512
num_experts = 4
block_size = 16
expert_num_tokens = torch.randint(
low=0,
high=max_tokens_per_batch,
size=(num_experts,),
dtype=torch.int32,
device="cuda",
)
max_num_tokens_padded = num_experts * max(max_tokens_per_batch, block_size)
sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
assert max_num_tokens_padded % block_size == 0
max_num_m_blocks = max_num_tokens_padded // block_size
expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda")
num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device="cuda")
opcheck(
torch.ops._moe_C.batched_moe_align_block_size,
(
max_tokens_per_batch,
block_size,
expert_num_tokens,
sorted_ids,
expert_ids,
num_tokens_post_pad,
),
)
# topk=8 covers topk > 4; k=511 covers the non-vectorized scalar path. The
# layouts exercise contiguous input plus the two non-contiguous cases: a
# transpose (strided hidden -> scalar gather) and a topk-slice (hidden still
# contiguous -> vectorized).
@pytest.mark.parametrize("m", [1, 33, 222])
@pytest.mark.parametrize("topk", [*TOP_KS, 8])
@pytest.mark.parametrize("k", [128, 511, 1024])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("layout", ["contig", "transpose", "slice"])
def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype, layout: str):
if layout == "transpose":
input = torch.randn((m, k, topk), device="cuda", dtype=dtype).transpose(1, 2)
elif layout == "slice":
input = torch.randn((m, 2 * topk, k), device="cuda", dtype=dtype)[:, ::2, :]
else:
input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
assert input.is_contiguous() == (layout == "contig")
actual = torch.empty((m, k), device="cuda", dtype=dtype)
# Reduction accumulates in fp32.
expected = input.float().sum(dim=1).to(dtype)
torch.ops._moe_C.moe_sum(input, actual)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
opcheck(torch.ops._moe_C.moe_sum, (input, actual))
@pytest.mark.parametrize("topk", [2, 8])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("topk_ids_dtype", [torch.int32, torch.int64])
def test_moe_sum_pad_aware(topk: int, dtype: torch.dtype, topk_ids_dtype: torch.dtype):
m, k = 17, 128
input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
topk_ids = torch.randint(0, 4, (m, topk), device="cuda", dtype=topk_ids_dtype)
# Force some slots unrouted, independent of expert parallelism.
topk_ids[0, :] = -1
topk_ids[1, 0] = -1
expert_map = torch.tensor([0, -1, 1, -1], device="cuda", dtype=torch.int32)
actual = torch.empty((m, k), device="cuda", dtype=dtype)
torch.ops._moe_C.moe_sum(input, actual, topk_ids, expert_map)
routed = expert_map[topk_ids.clamp(min=0).long()] >= 0
routed &= topk_ids >= 0
expected = (input.float() * routed.unsqueeze(-1)).sum(dim=1).to(dtype)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
opcheck(torch.ops._moe_C.moe_sum, (input, actual, topk_ids, expert_map))
def _batched_fused_marlin_moe_cases() -> list[Any]:
cases = [
pytest.param(
m,
128,
128,
e,
topk,
max_tokens_per_batch,
torch.bfloat16,
scalar_types.float4_e2m1f,
None,
1e-3,
id=(
f"m{m}-n128-k128-e{e}-topk{topk}-max_tokens{max_tokens_per_batch}-mxfp4"
),
)
for m in [16, 32, 64]
for e in [8, 12, 16, 32]
for topk in [2, 4]
for max_tokens_per_batch in [16, 32, 64]
]
cases.append(
pytest.param(
32,
128,
128,
8,
2,
64,
torch.float16,
scalar_types.uint4,
scalar_types.int8,
4e-2,
id="awq-int8-activation-metadata",
)
)
return cases
@pytest.mark.parametrize(
("m,n,k,e,topk,max_tokens_per_batch,dtype,quant_dtype,input_type,atol"),
_batched_fused_marlin_moe_cases(),
)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
def test_batched_fused_marlin_moe(
m: int,
n: int,
k: int,
e: int,
topk: int,
max_tokens_per_batch: int,
dtype: torch.dtype,
quant_dtype: ScalarType,
input_type: ScalarType | None,
atol: float,
):
print(
f"testing m={m}, n={n}, k={k}, e={e}, "
f"topk={topk}, "
f"max_tokens_per_batch={max_tokens_per_batch}, "
f"dtype={dtype}, quant_dtype={quant_dtype}, input_type={input_type}"
)
set_random_seed(0)
group_size = 32
if input_type == scalar_types.int8:
input_dtype = torch.int8
elif input_type == scalar_types.float8_e4m3fn:
input_dtype = torch.float8_e4m3fn
else:
input_dtype = None
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
w1_data = MarlinMoEWeightData.make(
w=w1,
quant_type=quant_dtype,
group_size=group_size,
act_order=None,
input_type=input_type,
)
w2_data = MarlinMoEWeightData.make(
w=w2,
quant_type=quant_dtype,
group_size=group_size,
act_order=None,
input_type=input_type,
)
score = torch.randn((m, e), device="cuda", dtype=dtype)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
class BatchedRun:
@staticmethod
def _make_expert_num_tokens_cpu(
e: int, # num_experts
topk_ids_cpu: torch.Tensor,
) -> torch.Tensor:
expert_num_tokens_cpu = torch.zeros((e,), dtype=torch.int32, device="cpu")
for topk_id in torch.flatten(topk_ids_cpu):
expert_num_tokens_cpu[topk_id] += 1
return expert_num_tokens_cpu
def __init__(
self,
max_tokens_per_batch: int,
num_experts: int,
_topk_ids: torch.Tensor,
_topk_weights: torch.Tensor,
):
self.max_tokens_per_batch = max_tokens_per_batch
self.e = num_experts
self.topk_ids_cpu = _topk_ids.to("cpu")
self.topk_weights_cpu = _topk_weights.to("cpu")
self.expert_num_tokens_cpu = self._make_expert_num_tokens_cpu(
self.e, self.topk_ids_cpu
)
def is_valid(self):
"""
Return True only if the input can be represented in a Batched
format.
"""
return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch)
def _scatter(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states_cpu = hidden_states.to("cpu")
K = hidden_states_cpu.size(1)
batched_hidden_states_cpu = torch.empty(
(e, max_tokens_per_batch, K),
dtype=hidden_states_cpu.dtype,
device="cpu",
)
counter_cpu = torch.zeros_like(self.expert_num_tokens_cpu)
for t_idx, token in enumerate(hidden_states_cpu):
for topk_id in self.topk_ids_cpu[t_idx]:
pos_in_batch = counter_cpu[topk_id]
batched_hidden_states_cpu[topk_id, pos_in_batch] = token
counter_cpu[topk_id] += 1
assert torch.allclose(counter_cpu, self.expert_num_tokens_cpu)
return batched_hidden_states_cpu.to("cuda")
def _gather(
self, batched_outputs: torch.Tensor, gather_outputs: torch.Tensor
) -> torch.Tensor:
batched_outputs_cpu = batched_outputs.to("cpu")
gather_outputs_cpu = torch.zeros_like(gather_outputs)
counter_cpu = torch.zeros((e,), device="cpu", dtype=torch.int32)
md = gather_outputs_cpu.size(0)
for t_idx in range(md):
token = None
for topk_id, topk_weight in zip(
self.topk_ids_cpu[t_idx], self.topk_weights_cpu[t_idx]
):
pos_in_batch = counter_cpu[topk_id]
t = batched_outputs_cpu[topk_id, pos_in_batch] * topk_weight
if token is None:
token = t
else:
token += t
counter_cpu[topk_id] += 1
assert token is not None
gather_outputs_cpu[t_idx] = token
gather_outputs.copy_(gather_outputs_cpu)
return gather_outputs
def run(
self, hidden_states: torch.Tensor, fused_marlin_moe_kwargs: dict[Any, Any]
) -> torch.Tensor:
assert hidden_states.ndim == 2
assert self.is_valid()
batched_hidden_states = self._scatter(hidden_states)
kwargs = fused_marlin_moe_kwargs | {
"hidden_states": batched_hidden_states,
"expert_num_tokens": self.expert_num_tokens_cpu.to("cuda"),
}
batched_outputs = batched_fused_marlin_moe(**kwargs)
output = torch.zeros_like(hidden_states)
output = self._gather(batched_outputs, output)
return output
kwargs = {
"w1": w1_data.qweight,
"w2": w2_data.qweight,
"bias1": None,
"bias2": None,
"w1_scale": w1_data.scales,
"w2_scale": w2_data.scales,
"global_num_experts": e,
"expert_map": None,
"global_scale1": w1_data.global_scale,
"global_scale2": w2_data.global_scale,
"g_idx1": w1_data.g_idx,
"g_idx2": w2_data.g_idx,
"sort_indices1": w1_data.sort_indices,
"sort_indices2": w2_data.sort_indices,
"w1_zeros": w1_data.zeros,
"w2_zeros": w2_data.zeros,
"quant_type_id": quant_dtype.id,
"is_k_full": True,
}
if input_dtype is not None:
kwargs["input_dtype"] = input_dtype
if w1_data.a_scales_factor is not None:
kwargs["input_global_scale1"] = w1_data.a_scales_factor
if w2_data.a_scales_factor is not None:
kwargs["input_global_scale2"] = w2_data.a_scales_factor
# Reference
fused_marlin_moe_kwargs = kwargs | {
"hidden_states": a,
"topk_ids": topk_ids,
"topk_weights": topk_weights,
}
ref_marlin_output = fused_marlin_moe(**fused_marlin_moe_kwargs)
# Batched
br = BatchedRun(max_tokens_per_batch, e, topk_ids, topk_weights)
if not br.is_valid():
pytest.skip("Cannot represent data in Batched Format.")
marlin_output = br.run(a, kwargs)
torch.testing.assert_close(marlin_output, ref_marlin_output, atol=atol, rtol=0)
@pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)])
@pytest.mark.parametrize("e,topk", [(8, 2)])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="TRTLLM backend test only runs on Blackwell GPUs (SM10x).",
)
def test_unquantized_bf16_flashinfer_trtllm_backend(
m: int,
n: int,
k: int,
e: int,
topk: int,
dtype: torch.dtype,
workspace_init,
):
"""
Test BF16 unquantized MoE with FlashInfer TRTLLM backend.
"""
set_random_seed(7)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
UnquantizedMoeBackend,
)
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
UnquantizedFusedMoEMethod,
)
# Setup test data
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
router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
moe_config = FusedMoEConfig(
num_experts=e,
experts_per_token=topk,
hidden_dim=k,
intermediate_size=n,
num_local_experts=e,
num_logical_experts=e,
activation=MoEActivation.SILU,
device="cuda",
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
in_dtype=dtype,
routing_method=RoutingMethodType.Renormalize,
max_num_tokens=next_power_of_2(m),
moe_backend="flashinfer_trtllm",
)
with set_current_vllm_config(vllm_config):
quant_method = UnquantizedFusedMoEMethod(moe_config)
# Verify TRTLLM backend was selected
assert (
quant_method.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
), f"Expected FLASHINFER_TRTLLM backend, got {quant_method.unquantized_backend}"
# Verify it's using monolithic path
assert quant_method.is_monolithic, (
"FLASHINFER_TRTLLM backend should use monolithic forward"
)
layer = torch.nn.Module()
layer.w13_weight = Parameter(w1.clone(), requires_grad=False)
layer.w2_weight = Parameter(w2.clone(), requires_grad=False)
layer.global_num_experts = e
layer.local_num_experts = e
layer.top_k = topk
layer.num_expert_group = 1
layer.topk_group = 1
layer.intermediate_size_per_partition = n
layer.ep_rank = 0
layer.activation = MoEActivation.SILU
layer.e_score_correction_bias = None
layer.routing_method_type = RoutingMethodType.Renormalize
layer.expert_map = None
layer.apply_router_weight_on_input = False
layer.routed_scaling_factor = 2.446
layer.shared_experts = None
layer._expert_routing_tables = lambda: None
quant_method.process_weights_after_loading(layer)
assert quant_method.moe_kernel is not None, (
"moe_kernel should be set after process_weights_after_loading"
)
assert quant_method.supports_internal_mk, (
"supports_internal_mk should be True after setup"
)
trtllm_output = quant_method.apply_monolithic(
layer=layer,
x=a,
router_logits=router_logits,
)
# Compute torch baseline
w1_original = w1.clone()
w2_original = w2.clone()
baseline_output = (
torch_moe(a, w1_original, w2_original, router_logits, topk)
* layer.routed_scaling_factor
)
close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
assert close.float().mean() > 0.925
Reference in New Issue View Git Blame Copy Permalink