wehub-resource-sync
1038 lines
38 KiB
Python
1038 lines
38 KiB
Python
from itertools import product
|
|
from typing import Tuple # noqa: UP035
|
|
|
|
import ml_dtypes
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.relax.frontend import nn
|
|
from tvm.relax.frontend.nn import spec
|
|
from tvm.s_tir import dlight as dl
|
|
|
|
from mlc_llm.compiler_pass.dispatch_triton_kernel import DispatchTritonKernel
|
|
from mlc_llm.op import batch_matmul, cutlass, moe_matmul, triton
|
|
from mlc_llm.quantization.block_scale_quantization import rowwise_group_quant_fp8
|
|
|
|
# test category "op_correctness"
|
|
pytestmark = [pytest.mark.op_correctness]
|
|
|
|
block_size = (128, 128)
|
|
fp8_dtype = "float8_e4m3fn"
|
|
|
|
torch_fp8_dtype = torch.float8_e4m3fn
|
|
torch_device = torch.device("cuda")
|
|
|
|
torch.set_grad_enabled(False)
|
|
|
|
|
|
def test_fp8_block_matmul_cutlass(M: int, N: int, K: int, dtype: str):
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def cutlass_gemm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor):
|
|
n, k = w.shape
|
|
# assert n % block_size[0] == 0
|
|
assert k % block_size[1] == 0
|
|
assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0]
|
|
assert k // block_size[1] == w_scale.shape[1]
|
|
assert x.shape[1] == k
|
|
x_fp8, x_scale = rowwise_group_quant_fp8(
|
|
x, block_size[1], w.dtype, transpose_scale=True
|
|
)
|
|
assert x_fp8.dtype == w.dtype
|
|
assert x_scale.dtype == "float32"
|
|
o = cutlass.fp8_groupwise_scaled_gemm(x_fp8, x_scale, w, w_scale, block_size, x.dtype)
|
|
return x_fp8, x_scale, o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"cutlass_gemm": {
|
|
"x": spec.Tensor(("m", K), dtype),
|
|
"w": spec.Tensor((N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
x_torch = torch.rand(M, K, dtype=getattr(torch, dtype), device=torch_device) * 2 - 1
|
|
w_full_torch = torch.rand(N, K, dtype=getattr(torch, dtype), device=torch_device) * 2 - 1
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype)
|
|
o_torch = blockwise_matmul(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
x_fp8_tvm, x_scale_tvm, o_tvm = vm["cutlass_gemm"](x_tvm, w_tvm, w_scale_tvm)
|
|
x_fp8_tvm = x_fp8_tvm.numpy()
|
|
x_scale_tvm = x_scale_tvm.numpy()
|
|
o_tvm = o_tvm.numpy()
|
|
|
|
np.testing.assert_allclose(
|
|
x_fp8_tvm,
|
|
x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype),
|
|
atol=1e-1,
|
|
rtol=1e-1,
|
|
)
|
|
np.testing.assert_allclose(x_scale_tvm.T, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5)
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def test_fp8_block_matmul_triton(M: int, N: int, K: int, dtype: str):
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def triton_gemm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor):
|
|
n, k = w.shape
|
|
assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0]
|
|
assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[1]
|
|
assert x.shape[1] == k
|
|
x_fp8, x_scale = rowwise_group_quant_fp8(
|
|
x, block_size[1], w.dtype, transpose_scale=False
|
|
)
|
|
assert x_fp8.dtype == w.dtype
|
|
assert x_scale.dtype == "float32"
|
|
o = triton.fp8_groupwise_scaled_gemm(
|
|
x_fp8,
|
|
x_scale,
|
|
w,
|
|
w_scale,
|
|
block_size,
|
|
x.dtype,
|
|
)
|
|
return x_fp8, x_scale, o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"triton_gemm": {
|
|
"x": spec.Tensor(("m", K), dtype),
|
|
"w": spec.Tensor((N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
mod = DispatchTritonKernel(target)(mod)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
x_torch = torch.randn(M, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_full_torch = torch.randn(N, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype)
|
|
o_torch = blockwise_matmul(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
x_fp8_tvm, x_scale_tvm, o_tvm = vm["triton_gemm"](x_tvm, w_tvm, w_scale_tvm)
|
|
x_fp8_tvm = x_fp8_tvm.numpy()
|
|
x_scale_tvm = x_scale_tvm.numpy()
|
|
o_tvm = o_tvm.numpy()
|
|
np.testing.assert_allclose(
|
|
x_fp8_tvm,
|
|
x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype),
|
|
atol=1e-1,
|
|
rtol=1e-1,
|
|
)
|
|
np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5)
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def test_fp8_block_group_matmul_cutlass(M: int, N: int, K: int, dtype: str):
|
|
num_experts = 256
|
|
top_k = 8
|
|
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def cutlass_group_gemm(
|
|
self,
|
|
x: nn.Tensor,
|
|
w: nn.Tensor,
|
|
w_scale: nn.Tensor,
|
|
indptr: nn.Tensor,
|
|
):
|
|
e, n, k = w.shape
|
|
assert e == num_experts
|
|
assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1]
|
|
assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2]
|
|
assert x.shape[1] == k
|
|
x_fp8, x_scale = rowwise_group_quant_fp8(
|
|
x, block_size[1], w.dtype, transpose_scale=False
|
|
)
|
|
assert x_fp8.dtype == w.dtype
|
|
assert x_scale.dtype == "float32"
|
|
o = cutlass.fp8_groupwise_scaled_group_gemm(
|
|
x_fp8,
|
|
x_scale,
|
|
w,
|
|
w_scale,
|
|
indptr,
|
|
block_size,
|
|
x.dtype,
|
|
)
|
|
return x_fp8, x_scale, o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"cutlass_group_gemm": {
|
|
"x": spec.Tensor(("m", K), dtype),
|
|
"w": spec.Tensor((num_experts, N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
num_experts,
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
"indptr": spec.Tensor((num_experts,), "int64"),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
# Randomly sample `top_k` experts for each token with pytorch
|
|
expert_choices = torch.randint(
|
|
0, num_experts, (M * top_k,), device=torch_device, dtype=torch.int32
|
|
)
|
|
|
|
factor = 1
|
|
# Balance so that the number of tokens for each expert is a multiple of `factor`
|
|
token_balance = 0
|
|
num_tokens_list = [int((expert_choices == i).sum().to("cpu")) for i in range(num_experts)]
|
|
for i in range(num_experts):
|
|
if token_balance > 0:
|
|
diff = min(token_balance, num_tokens_list[i])
|
|
num_tokens_list[i] -= diff
|
|
token_balance -= diff
|
|
if num_tokens_list[i] % factor != 0:
|
|
token_balance += factor - num_tokens_list[i] % factor
|
|
num_tokens_list[i] += factor - num_tokens_list[i] % factor
|
|
assert sum(num_tokens_list) == M * top_k
|
|
|
|
indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int64)
|
|
for i in range(num_experts):
|
|
indptr[i + 1] = indptr[i] + (expert_choices == i).sum()
|
|
token_ids_list = []
|
|
for i in range(num_experts):
|
|
# Get the indices of the tokens that belong to the i-th expert
|
|
token_ids = torch.where(expert_choices == i)[0]
|
|
token_ids_list.append(token_ids)
|
|
|
|
x_torch = torch.randn(M * top_k, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype)
|
|
o_torch = blockwise_group_matmul(
|
|
x_fp8_torch,
|
|
x_scale_torch,
|
|
w_torch,
|
|
w_scale_torch,
|
|
indptr,
|
|
x_torch.dtype,
|
|
)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
indptr_tvm = tvm.runtime.tensor(indptr[1:].cpu().numpy(), device=device)
|
|
x_fp8_tvm, x_scale_tvm, o_tvm = vm["cutlass_group_gemm"](
|
|
x_tvm,
|
|
w_tvm,
|
|
w_scale_tvm,
|
|
indptr_tvm,
|
|
)
|
|
x_fp8_tvm = x_fp8_tvm.numpy()
|
|
x_scale_tvm = x_scale_tvm.numpy()
|
|
o_tvm = o_tvm.numpy()
|
|
np.testing.assert_allclose(
|
|
x_fp8_tvm,
|
|
x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype),
|
|
atol=1e-1,
|
|
rtol=1e-1,
|
|
)
|
|
np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5)
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def test_fp8_block_group_matmul_triton(M: int, N: int, K: int, dtype: str):
|
|
num_experts = 256
|
|
top_k = 8
|
|
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def triton_group_gemm(
|
|
self,
|
|
x: nn.Tensor,
|
|
w: nn.Tensor,
|
|
w_scale: nn.Tensor,
|
|
indptr: nn.Tensor,
|
|
):
|
|
e, n, k = w.shape
|
|
assert e == num_experts
|
|
assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1]
|
|
assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2]
|
|
assert x.shape[1] == k
|
|
x_fp8, x_scale = rowwise_group_quant_fp8(
|
|
x, block_size[1], w.dtype, transpose_scale=False
|
|
)
|
|
assert x_fp8.dtype == w.dtype
|
|
assert x_scale.dtype == "float32"
|
|
o = triton.fp8_groupwise_scaled_group_gemm(
|
|
x_fp8,
|
|
x_scale,
|
|
w,
|
|
w_scale,
|
|
indptr,
|
|
block_size,
|
|
x.dtype,
|
|
)
|
|
return x_fp8, x_scale, o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"triton_group_gemm": {
|
|
"x": spec.Tensor(("m", K), dtype),
|
|
"w": spec.Tensor((num_experts, N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
num_experts,
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
"indptr": spec.Tensor((num_experts + 1,), "int32"),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
mod = DispatchTritonKernel(target)(mod)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
# Randomly sample `top_k` experts for each token with pytorch
|
|
expert_choices = torch.randint(
|
|
0, num_experts, (M * top_k,), device=torch_device, dtype=torch.int32
|
|
)
|
|
|
|
indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int32)
|
|
for i in range(num_experts):
|
|
indptr[i + 1] = indptr[i] + (expert_choices == i).sum()
|
|
token_ids_list = []
|
|
for i in range(num_experts):
|
|
# Get the indices of the tokens that belong to the i-th expert
|
|
token_ids = torch.where(expert_choices == i)[0]
|
|
token_ids_list.append(token_ids)
|
|
|
|
x_torch = torch.randn(M * top_k, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype)
|
|
o_torch = blockwise_group_matmul(
|
|
x_fp8_torch,
|
|
x_scale_torch,
|
|
w_torch,
|
|
w_scale_torch,
|
|
indptr,
|
|
x_torch.dtype,
|
|
)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
indptr_tvm = tvm.runtime.tensor(indptr.cpu().numpy(), device=device)
|
|
x_fp8_tvm, x_scale_tvm, o_tvm = vm["triton_group_gemm"](
|
|
x_tvm,
|
|
w_tvm,
|
|
w_scale_tvm,
|
|
indptr_tvm,
|
|
)
|
|
x_fp8_tvm = x_fp8_tvm.numpy()
|
|
x_scale_tvm = x_scale_tvm.numpy()
|
|
o_tvm = o_tvm.numpy()
|
|
np.testing.assert_allclose(
|
|
x_fp8_tvm,
|
|
x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype),
|
|
atol=1e-1,
|
|
rtol=1e-1,
|
|
)
|
|
np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5)
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def test_fp8_block_bmm_cutlass(M: int, N: int, K: int, H: int, dtype: str):
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def cutlass_bmm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor):
|
|
_, n, k = w.shape
|
|
assert w.shape[0] == x.shape[0] == H
|
|
assert n % block_size[0] == 0
|
|
assert k % block_size[1] == 0
|
|
assert n // block_size[0] == w_scale.shape[1]
|
|
assert k // block_size[1] == w_scale.shape[2]
|
|
assert x.shape[2] == k
|
|
o = batch_matmul.quantized_bmm(x, w, w_scale, block_size)
|
|
return o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"cutlass_bmm": {
|
|
"x": spec.Tensor((H, "m", K), dtype),
|
|
"w": spec.Tensor((H, N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
H,
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
x_torch = torch.randn(H, M, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_full_torch = torch.randn(H, N, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype)
|
|
o_torch = blockwise_bmm(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
o_tvm = vm["cutlass_bmm"](x_tvm, w_tvm, w_scale_tvm)
|
|
o_tvm = o_tvm.numpy()
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def test_fp8_block_gemv_tir(N: int, K: int, up: bool, dtype: str):
|
|
num_experts = 256
|
|
top_k = 8
|
|
M = 1 if up else top_k
|
|
|
|
device = tvm.cuda()
|
|
target = tvm.target.Target.from_device(device)
|
|
|
|
class TestModule(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def tir_moe_gemv(
|
|
self,
|
|
x: nn.Tensor,
|
|
w: nn.Tensor,
|
|
w_scale: nn.Tensor,
|
|
expert_indices: nn.Tensor,
|
|
):
|
|
e, n, k = w.shape
|
|
assert e == num_experts
|
|
assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1]
|
|
assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2]
|
|
assert x.shape[1] == k
|
|
o = moe_matmul.dequantize_block_scale_float8_gemv(
|
|
x, w, w_scale, expert_indices, block_size, x.dtype
|
|
)
|
|
return o
|
|
|
|
mod, _, ext_mods = TestModule().export_tvm(
|
|
spec={
|
|
"tir_moe_gemv": {
|
|
"x": spec.Tensor((M, K), dtype),
|
|
"w": spec.Tensor((num_experts, N, K), fp8_dtype),
|
|
"w_scale": spec.Tensor(
|
|
(
|
|
num_experts,
|
|
(N + block_size[0] - 1) // block_size[0],
|
|
(K + block_size[1] - 1) // block_size[1],
|
|
),
|
|
"float32",
|
|
),
|
|
"expert_indices": spec.Tensor((1, top_k), "int32"),
|
|
},
|
|
},
|
|
allow_extern=True,
|
|
)
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(
|
|
dl.gpu.Matmul(),
|
|
dl.gpu.GEMV(),
|
|
dl.gpu.Reduction(),
|
|
dl.gpu.GeneralReduction(),
|
|
dl.gpu.Fallback(),
|
|
)(mod)
|
|
exec = relax.build(
|
|
mod,
|
|
target=target,
|
|
relax_pipeline=relax.backend.cuda.get_default_pipeline(target),
|
|
)
|
|
vm = relax.VirtualMachine(exec, device)
|
|
|
|
# Randomly sample `top_k` experts for each token with pytorch
|
|
expert_choices = torch.randint(0, num_experts, (top_k,), device=torch_device, dtype=torch.int32)
|
|
indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int32)
|
|
for i in range(num_experts):
|
|
indptr[i + 1] = indptr[i] + (expert_choices == i).sum()
|
|
token_ids_list = []
|
|
for i in range(num_experts):
|
|
# Get the indices of the tokens that belong to the i-th expert
|
|
token_ids = torch.where(expert_choices == i)[0]
|
|
token_ids_list.append(token_ids)
|
|
|
|
x_torch = torch.randn(M, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device)
|
|
w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype)
|
|
x_input_torch = torch.repeat_interleave(x_torch, top_k, dim=0) if up else x_torch
|
|
o_torch = blockwise_group_matmul_unquantized(
|
|
x_input_torch, w_torch, w_scale_torch, expert_choices
|
|
)
|
|
x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device)
|
|
w_tvm = tvm.runtime.tensor(
|
|
w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device)
|
|
expert_choices = tvm.runtime.tensor(
|
|
expert_choices.reshape(1, top_k).cpu().numpy(), device=device
|
|
)
|
|
o_tvm = vm["tir_moe_gemv"](x_tvm, w_tvm, w_scale_tvm, expert_choices)
|
|
o_tvm = o_tvm.numpy()
|
|
atol = 0.5
|
|
rtol = 1e-4
|
|
o_tvm_flat = o_tvm.flatten()
|
|
o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten()
|
|
failed_indices = np.where(
|
|
np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat))
|
|
)[0]
|
|
if len(failed_indices) > 0:
|
|
print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}")
|
|
print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}")
|
|
print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}")
|
|
np.testing.assert_allclose(
|
|
o_tvm,
|
|
o_torch.view(torch.float16).cpu().numpy().view(dtype),
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
|
|
|
|
def blockwise_matmul(
|
|
x_fp8_torch: torch.Tensor,
|
|
x_scale_torch: torch.Tensor,
|
|
w_torch: torch.Tensor,
|
|
w_scale_torch: torch.Tensor,
|
|
dtype,
|
|
):
|
|
o_torch = torch.zeros(
|
|
(x_fp8_torch.shape[0], w_torch.shape[0]), dtype=dtype, device=torch_device
|
|
)
|
|
for j in range(w_scale_torch.shape[0]):
|
|
for k in range(w_scale_torch.shape[1]):
|
|
o_torch[
|
|
:,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[0]),
|
|
] += (
|
|
torch.matmul(
|
|
x_fp8_torch[
|
|
:,
|
|
k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[1]),
|
|
].to(dtype),
|
|
w_torch[
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[0]),
|
|
k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[1]),
|
|
].T.to(dtype),
|
|
)
|
|
* x_scale_torch[:, k : k + 1]
|
|
* w_scale_torch[j, k]
|
|
)
|
|
return o_torch
|
|
|
|
|
|
def blockwise_group_matmul(
|
|
x_fp8_torch: torch.Tensor,
|
|
x_scale_torch: torch.Tensor,
|
|
w_torch: torch.Tensor,
|
|
w_scale_torch: torch.Tensor,
|
|
indptr: torch.Tensor,
|
|
dtype,
|
|
):
|
|
o_torch = torch.zeros(
|
|
(x_fp8_torch.shape[0], w_torch.shape[1]), dtype=dtype, device=torch_device
|
|
)
|
|
for e in range(w_scale_torch.shape[0]):
|
|
if indptr[e + 1] - indptr[e] == 0:
|
|
continue
|
|
indices = slice(indptr[e], indptr[e + 1])
|
|
for j in range(w_scale_torch.shape[1]):
|
|
for k in range(w_scale_torch.shape[2]):
|
|
o_torch[
|
|
indices,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
] += (
|
|
torch.matmul(
|
|
x_fp8_torch.to(dtype)[
|
|
indices,
|
|
k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[1]),
|
|
],
|
|
w_torch[
|
|
e,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]),
|
|
].T.to(dtype),
|
|
)
|
|
* x_scale_torch[indices, k : k + 1]
|
|
* w_scale_torch[e, j, k]
|
|
)
|
|
return o_torch
|
|
|
|
|
|
def blockwise_group_matmul_unquantized(
|
|
x_torch: torch.Tensor,
|
|
w_torch: torch.Tensor,
|
|
w_scale_torch: torch.Tensor,
|
|
expert_choices: torch.Tensor,
|
|
):
|
|
o_torch = torch.zeros(
|
|
(x_torch.shape[0], w_torch.shape[1]), dtype=x_torch.dtype, device=torch_device
|
|
)
|
|
for i, e in enumerate(expert_choices):
|
|
for j in range(w_scale_torch.shape[1]):
|
|
for k in range(w_scale_torch.shape[2]):
|
|
o_torch[
|
|
i,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
] += torch.matmul(
|
|
x_torch[
|
|
i,
|
|
k * block_size[1] : min((k + 1) * block_size[1], x_torch.shape[1]),
|
|
],
|
|
w_torch[
|
|
e,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]),
|
|
].T.to(x_torch.dtype)
|
|
* w_scale_torch[e, j, k].to(x_torch.dtype),
|
|
)
|
|
return o_torch
|
|
|
|
|
|
def blockwise_bmm(
|
|
x_fp8_torch: torch.Tensor,
|
|
x_scale_torch: torch.Tensor,
|
|
w_torch: torch.Tensor,
|
|
w_scale_torch: torch.Tensor,
|
|
dtype,
|
|
):
|
|
o_torch = torch.zeros(
|
|
(x_fp8_torch.shape[0], x_fp8_torch.shape[1], w_torch.shape[1]),
|
|
dtype=dtype,
|
|
device=torch_device,
|
|
)
|
|
for j in range(w_scale_torch.shape[1]):
|
|
for k in range(w_scale_torch.shape[2]):
|
|
o_torch[
|
|
...,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
] += (
|
|
torch.bmm(
|
|
x_fp8_torch[
|
|
...,
|
|
k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[2]),
|
|
].to(dtype),
|
|
w_torch[
|
|
...,
|
|
j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]),
|
|
k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]),
|
|
]
|
|
.transpose(1, 2)
|
|
.to(dtype),
|
|
)
|
|
* x_scale_torch[..., k : k + 1]
|
|
* w_scale_torch[..., j : j + 1, k : k + 1]
|
|
)
|
|
return o_torch
|
|
|
|
|
|
def blockwise_quant_fp8(
|
|
w_full_torch: torch.Tensor,
|
|
block_size: Tuple[int, int], # noqa: UP006
|
|
quant_dtype: torch.dtype,
|
|
):
|
|
w_scale_shape = (
|
|
*w_full_torch.shape[:-2],
|
|
(w_full_torch.shape[-2] + block_size[0] - 1) // block_size[0],
|
|
(w_full_torch.shape[-1] + block_size[1] - 1) // block_size[1],
|
|
)
|
|
# For each (block_size[0], block_size[1]) block, compute the max abs value of `w_full_torch`
|
|
w_max_abs_torch = torch.zeros(w_scale_shape, dtype=torch.float32, device=torch_device)
|
|
for i in range(w_scale_shape[-2]):
|
|
for j in range(w_scale_shape[-1]):
|
|
w_max_abs_torch[..., i, j] = torch.max(
|
|
torch.abs(
|
|
w_full_torch[
|
|
...,
|
|
i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]),
|
|
j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]),
|
|
]
|
|
).flatten(-2, -1),
|
|
dim=-1,
|
|
)[0]
|
|
# Scale is the `w_max_abs_torch` divided by the max value of quant_dtype in ml_dtypes
|
|
fp8_max = float(ml_dtypes.finfo(fp8_dtype).max)
|
|
w_scale_torch = w_max_abs_torch / fp8_max
|
|
# `w_torch` is the `w_full_torch` divided by the `w_scale_torch` (with block awareness),
|
|
# clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
|
|
w_torch = torch.zeros_like(w_full_torch, dtype=quant_dtype, device=torch_device)
|
|
if len(w_scale_shape) == 2:
|
|
for i in range(w_scale_shape[-2]):
|
|
for j in range(w_scale_shape[-1]):
|
|
w_torch[
|
|
i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]),
|
|
j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]),
|
|
] = torch.clamp(
|
|
w_full_torch[
|
|
i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]),
|
|
j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]),
|
|
]
|
|
/ w_scale_torch[..., i, j],
|
|
-fp8_max,
|
|
fp8_max,
|
|
)
|
|
else:
|
|
for e in range(w_scale_shape[0]):
|
|
for i in range(w_scale_shape[-2]):
|
|
for j in range(w_scale_shape[-1]):
|
|
w_torch[
|
|
e,
|
|
i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]),
|
|
j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]),
|
|
] = torch.clamp(
|
|
w_full_torch[
|
|
e,
|
|
i * block_size[0] : min(
|
|
(i + 1) * block_size[0], w_full_torch.shape[-2]
|
|
),
|
|
j * block_size[1] : min(
|
|
(j + 1) * block_size[1], w_full_torch.shape[-1]
|
|
),
|
|
]
|
|
/ w_scale_torch[e, i, j],
|
|
-fp8_max,
|
|
fp8_max,
|
|
)
|
|
|
|
w_scale_torch = (
|
|
torch.rand(w_scale_torch.shape, dtype=torch.float32, device=torch_device) / fp8_max
|
|
)
|
|
return w_torch, w_scale_torch
|
|
|
|
|
|
def rowwise_quant_fp8(
|
|
x_full_torch: torch.Tensor,
|
|
block_size: Tuple[int, int], # noqa: UP006
|
|
quant_dtype: torch.dtype,
|
|
):
|
|
x_scale_shape = (
|
|
*x_full_torch.shape[:-1],
|
|
(x_full_torch.shape[-1] + block_size[1] - 1) // block_size[1],
|
|
)
|
|
# For each (block_size[1]) block, compute the max abs value of `w_full_torch`
|
|
x_max_abs_torch = torch.zeros(x_scale_shape, dtype=torch.float32, device=torch_device)
|
|
for i in range(x_scale_shape[-1]):
|
|
x_max_abs_torch[..., i] = torch.max(
|
|
torch.abs(
|
|
x_full_torch[
|
|
...,
|
|
i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]),
|
|
]
|
|
),
|
|
dim=-1,
|
|
)[0]
|
|
# Scale is the `x_max_abs_torch` divided by the max value of quant_dtype in ml_dtypes
|
|
fp8_max = float(ml_dtypes.finfo(fp8_dtype).max)
|
|
x_scale_torch = x_max_abs_torch / fp8_max
|
|
# `x_torch` is the `x_full_torch` divided by the `x_scale_torch` (with block awareness),
|
|
# clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
|
|
x_torch = torch.zeros_like(x_full_torch, dtype=quant_dtype, device=torch_device)
|
|
for i in range(x_scale_shape[-1]):
|
|
x_torch[
|
|
...,
|
|
i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]),
|
|
] = torch.clamp(
|
|
x_full_torch[
|
|
...,
|
|
i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]),
|
|
]
|
|
/ x_scale_torch[..., i : i + 1],
|
|
-fp8_max,
|
|
fp8_max,
|
|
)
|
|
|
|
x_scale_torch = (
|
|
torch.rand(x_scale_torch.shape, dtype=torch.float32, device=torch_device) / fp8_max
|
|
)
|
|
for i in range(x_scale_shape[-1]):
|
|
x_full_torch[
|
|
...,
|
|
i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]),
|
|
] = (
|
|
x_torch[
|
|
...,
|
|
i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]),
|
|
].to(x_scale_torch.dtype)
|
|
* x_scale_torch[..., i : i + 1]
|
|
)
|
|
return x_full_torch, x_torch, x_scale_torch
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_cutlass_gemm():
|
|
# Cutlass GEMM
|
|
for M, (N, K), dtype in product(
|
|
[4, 128, 256, 1024, 2112],
|
|
[
|
|
(4608, 896),
|
|
(896, 2304),
|
|
(3072, 896),
|
|
(512, 896),
|
|
(3072, 512),
|
|
(4096, 512),
|
|
(896, 2048),
|
|
(129280, 896),
|
|
],
|
|
["bfloat16"],
|
|
):
|
|
print(f"Cutlass, M: {M}, N: {N}, K: {K}, dtype: {dtype}")
|
|
test_fp8_block_matmul_cutlass(M, N, K, dtype)
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_triton_gemm():
|
|
# Triton GEMM
|
|
for M, (N, K), dtype in product(
|
|
[1, 128, 256, 1024, 2111],
|
|
[
|
|
(4608, 896),
|
|
(896, 576),
|
|
(896, 2304),
|
|
],
|
|
["bfloat16"],
|
|
):
|
|
print(f"Triton, M: {M}, N: {N}, K: {K}, dtype: {dtype}")
|
|
test_fp8_block_matmul_triton(M, N, K, dtype)
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_cutlass_group_gemm():
|
|
# Cutlass group GEMM
|
|
for M, (N, K), dtype in product(
|
|
[1, 128, 256, 1024, 2111],
|
|
[
|
|
(512, 896),
|
|
(896, 256),
|
|
],
|
|
["bfloat16"],
|
|
):
|
|
print(f"Cutlass group gemm, M: {M}, N: {N}, K: {K}, dtype: {dtype}")
|
|
test_fp8_block_group_matmul_cutlass(M, N, K, dtype)
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_triton_group_gemm():
|
|
# Triton group GEMM
|
|
for M, (N, K), dtype in product(
|
|
[1, 128, 256, 1024, 2111],
|
|
[
|
|
(512, 896),
|
|
(896, 256),
|
|
],
|
|
["bfloat16"],
|
|
):
|
|
print(f"Triton group gemm, M: {M}, N: {N}, K: {K}, dtype: {dtype}")
|
|
test_fp8_block_group_matmul_triton(M, N, K, dtype)
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_cutlass_bmm():
|
|
# Cutlass BMM
|
|
for M, H, (N, K), dtype in product(
|
|
[4, 128, 256, 1024, 2112],
|
|
[16, 64, 128],
|
|
[
|
|
(512, 128),
|
|
(128, 512),
|
|
],
|
|
["bfloat16"],
|
|
):
|
|
print(f"Cutlass BMM, M: {M}, N: {N}, K: {K}, H: {H}, dtype: {dtype}")
|
|
test_fp8_block_bmm_cutlass(M, N, K, H, dtype)
|
|
|
|
|
|
@pytest.mark.skip(reason="Test requiring SM90a")
|
|
def test_tir_moe_gemv():
|
|
# TIR MoE GEMV
|
|
for (N, K), up, dtype in product(
|
|
[(512, 896), (896, 256)],
|
|
[True, False],
|
|
["bfloat16"],
|
|
):
|
|
print(f"TIR MoE GEMV, N: {N}, K: {K}, up: {up}, dtype: {dtype}")
|
|
test_fp8_block_gemv_tir(N, K, up, dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_cutlass_gemm()
|
|
test_triton_gemm()
|
|
test_cutlass_group_gemm()
|
|
test_triton_group_gemm()
|
|
test_cutlass_bmm()
|
|
test_tir_moe_gemv()
|