|
|
|
|
|
"""Tests for the MOE layers. |
|
|
|
|
|
Run `pytest tests/kernels/test_moe.py`. |
|
|
""" |
|
|
|
|
|
from typing import List |
|
|
|
|
|
import pytest |
|
|
import torch |
|
|
|
|
|
from moe._ops import ops |
|
|
from moe.fused_moe import fused_moe, fused_topk, moe_align_block_size |
|
|
from moe.fused_marlin_moe import fused_marlin_moe |
|
|
from moe.platforms import current_platform |
|
|
from moe.scalar_type import scalar_types |
|
|
from moe.utils.marlin_utils_test import marlin_quantize, quantize_weights |
|
|
|
|
|
from .utils import compute_max_diff, opcheck, torch_moe |
|
|
|
|
|
|
|
|
from torch.nn import Parameter |
|
|
from torch.nn import functional as F |
|
|
|
|
|
def stack_and_dev(tensors: List[torch.Tensor]): |
|
|
dev = tensors[0].device |
|
|
return torch.stack(tensors, dim=0).to(dev) |
|
|
|
|
|
NUM_EXPERTS = [8, 64] |
|
|
EP_SIZE = [1, 4] |
|
|
TOP_KS = [2, 6] |
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("m", [1, 33, 64, 222, 1024 * 128]) |
|
|
@pytest.mark.parametrize("n", [128, 1024, 2048]) |
|
|
@pytest.mark.parametrize("k", [128, 511, 1024]) |
|
|
@pytest.mark.parametrize("e", NUM_EXPERTS) |
|
|
@pytest.mark.parametrize("topk", TOP_KS) |
|
|
@pytest.mark.parametrize("ep_size", EP_SIZE) |
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
|
|
@pytest.mark.parametrize("padding", [True, False]) |
|
|
def test_fused_moe( |
|
|
m: int, |
|
|
n: int, |
|
|
k: int, |
|
|
e: int, |
|
|
topk: int, |
|
|
ep_size: int, |
|
|
dtype: torch.dtype, |
|
|
padding: bool, |
|
|
): |
|
|
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) |
|
|
|
|
|
if ep_size > 1: |
|
|
local_e = e // ep_size |
|
|
e_ids = torch.randint(0, |
|
|
e, (local_e, ), |
|
|
device="cuda", |
|
|
dtype=torch.int32) |
|
|
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 |
|
|
|
|
|
torch_output = torch_moe(a, w1, w2, score, topk, e_map) |
|
|
if padding: |
|
|
w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128] |
|
|
torch.cuda.empty_cache() |
|
|
w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128] |
|
|
torch.cuda.empty_cache() |
|
|
|
|
|
triton_output = fused_moe(a, |
|
|
w1, |
|
|
w2, |
|
|
score, |
|
|
topk, |
|
|
global_num_experts=e, |
|
|
expert_map=e_map, |
|
|
renormalize=False) |
|
|
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) |
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("m", [1, 32, 222]) |
|
|
@pytest.mark.parametrize("n", [128, 1024, 2048]) |
|
|
@pytest.mark.parametrize("k", [128, 1024]) |
|
|
@pytest.mark.parametrize("e", NUM_EXPERTS) |
|
|
@pytest.mark.parametrize("topk", TOP_KS) |
|
|
@pytest.mark.parametrize("ep_size", EP_SIZE) |
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
|
|
@pytest.mark.parametrize("group_size", [64, 128]) |
|
|
@pytest.mark.parametrize("has_zp", [True, False]) |
|
|
@pytest.mark.parametrize("weight_bits", [4, 8]) |
|
|
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): |
|
|
print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits) |
|
|
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) |
|
|
|
|
|
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="cuda", |
|
|
dtype=torch.uint8) |
|
|
w2_qweight = torch.empty((e, k, n // pack_factor), |
|
|
device="cuda", |
|
|
dtype=torch.uint8) |
|
|
w1_scales = torch.empty((e, 2 * n, k // group_size), |
|
|
device="cuda", |
|
|
dtype=dtype) |
|
|
w2_scales = torch.empty((e, k, n // group_size), |
|
|
device="cuda", |
|
|
dtype=dtype) |
|
|
w1_qzeros = torch.empty((e, 2 * n // pack_factor, k // group_size), |
|
|
device="cuda", |
|
|
dtype=torch.uint8) |
|
|
w2_qzeros = torch.empty((e, k // pack_factor, n // group_size), |
|
|
device="cuda", |
|
|
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="cuda", |
|
|
dtype=torch.int32) |
|
|
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_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 |
|
|
|
|
|
triton_output = fused_moe(a, |
|
|
w1_qweight, |
|
|
w2_qweight, |
|
|
score, |
|
|
topk, |
|
|
renormalize=False, |
|
|
use_int4_w4a16=weight_bits == 4, |
|
|
use_int8_w8a16=weight_bits == 8, |
|
|
global_num_experts=e, |
|
|
expert_map=e_map, |
|
|
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]) |
|
|
torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map) |
|
|
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("m", [1, 33, 64, 222]) |
|
|
@pytest.mark.parametrize("n", [128, 2048]) |
|
|
@pytest.mark.parametrize("k", [128, 1024]) |
|
|
@pytest.mark.parametrize("e", NUM_EXPERTS) |
|
|
@pytest.mark.parametrize("topk", TOP_KS) |
|
|
@pytest.mark.parametrize("group_size", [-1, 32, 128]) |
|
|
@pytest.mark.parametrize("act_order", [True, False]) |
|
|
@pytest.mark.parametrize("num_bits", [4, 8]) |
|
|
@pytest.mark.parametrize("is_k_full", [True, False]) |
|
|
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
|
|
def test_fused_marlin_moe( |
|
|
m: int, |
|
|
n: int, |
|
|
k: int, |
|
|
e: int, |
|
|
topk: int, |
|
|
group_size: int, |
|
|
act_order: bool, |
|
|
num_bits: int, |
|
|
is_k_full: bool, |
|
|
): |
|
|
current_platform.seed_everything(7) |
|
|
|
|
|
|
|
|
if act_order: |
|
|
if group_size == -1: |
|
|
return |
|
|
if group_size in (k, n): |
|
|
return |
|
|
else: |
|
|
if not is_k_full: |
|
|
return |
|
|
|
|
|
quant_type = (scalar_types.uint4b8 |
|
|
if num_bits == 4 else scalar_types.uint8b128) |
|
|
dtype = torch.float16 |
|
|
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 |
|
|
|
|
|
w_ref1_l = [] |
|
|
qweight1_l = [] |
|
|
scales1_l = [] |
|
|
g_idx1_l = [] |
|
|
sort_indices1_l = [] |
|
|
|
|
|
for i in range(w1.shape[0]): |
|
|
test_perm = torch.randperm(k) |
|
|
w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = marlin_quantize( |
|
|
w1[i].transpose(1, 0), quant_type, group_size, act_order, |
|
|
test_perm) |
|
|
w_ref1_l.append(w_ref1) |
|
|
qweight1_l.append(qweight1) |
|
|
scales1_l.append(scales1) |
|
|
g_idx1_l.append(g_idx1) |
|
|
sort_indices1_l.append(sort_indices1) |
|
|
|
|
|
w_ref1 = stack_and_dev(w_ref1_l) |
|
|
qweight1 = stack_and_dev(qweight1_l).contiguous() |
|
|
scales1 = stack_and_dev(scales1_l) |
|
|
g_idx1 = stack_and_dev(g_idx1_l) |
|
|
sort_indices1 = stack_and_dev(sort_indices1_l) |
|
|
|
|
|
w_ref2_l = [] |
|
|
qweight2_l = [] |
|
|
scales2_l = [] |
|
|
g_idx2_l = [] |
|
|
sort_indices2_l = [] |
|
|
|
|
|
for i in range(w2.shape[0]): |
|
|
test_perm = torch.randperm(n) |
|
|
w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = marlin_quantize( |
|
|
w2[i].transpose(1, 0), quant_type, group_size, act_order, |
|
|
test_perm) |
|
|
w_ref2_l.append(w_ref2) |
|
|
qweight2_l.append(qweight2) |
|
|
scales2_l.append(scales2) |
|
|
g_idx2_l.append(g_idx2) |
|
|
sort_indices2_l.append(sort_indices2) |
|
|
|
|
|
w_ref2 = stack_and_dev(w_ref2_l) |
|
|
qweight2 = stack_and_dev(qweight2_l).contiguous() |
|
|
scales2 = stack_and_dev(scales2_l) |
|
|
g_idx2 = stack_and_dev(g_idx2_l) |
|
|
sort_indices2 = stack_and_dev(sort_indices2_l) |
|
|
|
|
|
score = torch.randn((m, e), device="cuda", dtype=dtype) |
|
|
|
|
|
topk_weights, topk_ids = fused_topk(a, score, topk, False) |
|
|
|
|
|
triton_output = fused_moe( |
|
|
a, |
|
|
w_ref1.transpose(1, 2).contiguous(), |
|
|
w_ref2.transpose(1, 2).contiguous(), |
|
|
score, |
|
|
topk, |
|
|
renormalize=False, |
|
|
) |
|
|
marlin_output = fused_marlin_moe( |
|
|
a, |
|
|
qweight1, |
|
|
qweight2, |
|
|
scales1, |
|
|
scales2, |
|
|
score, |
|
|
topk_weights, |
|
|
topk_ids, |
|
|
g_idx1=g_idx1, |
|
|
g_idx2=g_idx2, |
|
|
sort_indices1=sort_indices1, |
|
|
sort_indices2=sort_indices2, |
|
|
num_bits=num_bits, |
|
|
is_k_full=is_k_full, |
|
|
) |
|
|
|
|
|
assert compute_max_diff(marlin_output, triton_output) < 4e-2 |
|
|
|
|
|
token_expert_indicies = torch.empty(m, |
|
|
topk, |
|
|
dtype=torch.int32, |
|
|
device=a.device) |
|
|
|
|
|
opcheck(ops.topk_softmax, ( |
|
|
topk_weights, |
|
|
topk_ids, |
|
|
token_expert_indicies, |
|
|
score.float(), |
|
|
)) |
|
|
|
|
|
block_size_m = 4 |
|
|
|
|
|
sorted_token_ids, _, _ = moe_align_block_size(topk_ids, block_size_m, |
|
|
e) |
|
|
|
|
|
max_workspace_size = ((m + 255) // 256) * (max(2 * n, k) // 64) * 16 |
|
|
workspace = torch.zeros(max_workspace_size, |
|
|
dtype=torch.int, |
|
|
device="cuda", |
|
|
requires_grad=False) |
|
|
|
|
|
zp = torch.empty((0, 0), |
|
|
dtype=dtype, |
|
|
device="cuda", |
|
|
requires_grad=False) |
|
|
opcheck(ops.marlin_gemm_moe, |
|
|
(a, qweight1, sorted_token_ids, topk_weights, topk_ids, |
|
|
scales1, zp, g_idx1, sort_indices1, workspace, quant_type.id, |
|
|
m, 2 * n, k, True, e, topk, block_size_m, True, False)) |
|
|
|
|
|
|
|
|
@pytest.mark.skip("This test is here for the sake of debugging, " |
|
|
"don't run it in automated tests.") |
|
|
@pytest.mark.parametrize("m", [64, 512, 222, 33, 1]) |
|
|
@pytest.mark.parametrize("n", [128, 2048, 256, 1024]) |
|
|
@pytest.mark.parametrize("k", [128, 1024, 512]) |
|
|
@pytest.mark.parametrize("e", [8, 64]) |
|
|
@pytest.mark.parametrize("topk", [2, 6]) |
|
|
@pytest.mark.parametrize("group_size", [-1, 32, 64, 128]) |
|
|
@pytest.mark.parametrize("act_order", [True, False]) |
|
|
@pytest.mark.parametrize("num_bits", [4, 8]) |
|
|
@pytest.mark.parametrize("is_k_full", [True, False]) |
|
|
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
|
|
def test_single_marlin_moe_multiply( |
|
|
m: int, |
|
|
n: int, |
|
|
k: int, |
|
|
e: int, |
|
|
topk: int, |
|
|
group_size: int, |
|
|
act_order: bool, |
|
|
num_bits: int, |
|
|
is_k_full: bool, |
|
|
): |
|
|
|
|
|
|
|
|
if act_order: |
|
|
if group_size == -1: |
|
|
return |
|
|
if group_size == k: |
|
|
return |
|
|
else: |
|
|
if not is_k_full: |
|
|
return |
|
|
|
|
|
quant_type = (scalar_types.uint4b8 |
|
|
if num_bits == 4 else scalar_types.uint8b128) |
|
|
dtype = torch.float16 |
|
|
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
|
|
w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10 |
|
|
|
|
|
w_ref_l = [] |
|
|
qweights_l = [] |
|
|
scales_l = [] |
|
|
g_idx_l = [] |
|
|
sort_indices_l = [] |
|
|
|
|
|
for i in range(w.shape[0]): |
|
|
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) |
|
|
w_ref_l.append(w_ref) |
|
|
qweights_l.append(qweight) |
|
|
scales_l.append(scales) |
|
|
g_idx_l.append(g_idx) |
|
|
sort_indices_l.append(sort_indices) |
|
|
|
|
|
w_ref = stack_and_dev(w_ref_l) |
|
|
qweight = stack_and_dev(qweights_l).contiguous() |
|
|
scales = stack_and_dev(scales_l) |
|
|
g_idx = stack_and_dev(g_idx_l) |
|
|
sort_indices = stack_and_dev(sort_indices_l) |
|
|
|
|
|
score = torch.randn((m, e), device="cuda", dtype=dtype) |
|
|
marlin_output = ops.single_marlin_moe( |
|
|
a, |
|
|
qweight, |
|
|
scales, |
|
|
score, |
|
|
topk, |
|
|
renormalize=False, |
|
|
g_idx=g_idx, |
|
|
sort_indices=sort_indices, |
|
|
num_bits=num_bits, |
|
|
is_k_full=is_k_full, |
|
|
) |
|
|
|
|
|
torch_output = torch_moe_single(a, w_ref.transpose(1, 2), score, topk) |
|
|
|
|
|
assert compute_max_diff(marlin_output, torch_output) < 1e-2 |
|
|
|
|
|
|
|
|
def test_moe_align_block_size_opcheck(): |
|
|
num_experts = 4 |
|
|
block_size = 4 |
|
|
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(ops.moe_align_block_size, |
|
|
(topk_ids, num_experts, block_size, sorted_ids, expert_ids, |
|
|
num_tokens_post_pad)) |
|
|
|
