mirror of https://github.com/vllm-project/vllm.git
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
import torch
|
|
import triton.language as tl
|
|
|
|
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
|
|
invoke_moe_batched_triton_kernel)
|
|
|
|
|
|
@dataclass
|
|
class BatchedMMConfig:
|
|
dtype: torch.dtype
|
|
num_experts: int
|
|
max_tokens_per_expert: int
|
|
K: int
|
|
N: int
|
|
|
|
|
|
@dataclass
|
|
class BatchedMMTensors:
|
|
A: torch.Tensor # [E, max_tokens, K]
|
|
B: torch.Tensor # [E, K, N] - column major
|
|
C: torch.Tensor # [E, max_tokens, N]
|
|
num_expert_tokens: torch.Tensor # [E]
|
|
|
|
@staticmethod
|
|
def make_tensors(config: BatchedMMConfig):
|
|
A = torch.randn(
|
|
(config.num_experts, config.max_tokens_per_expert, config.K),
|
|
device="cuda",
|
|
dtype=config.dtype) / 10
|
|
B = torch.randn((config.num_experts, config.N, config.K),
|
|
device="cuda",
|
|
dtype=config.dtype)
|
|
C = torch.zeros(
|
|
(config.num_experts, config.max_tokens_per_expert, config.N),
|
|
device="cuda",
|
|
dtype=config.dtype)
|
|
num_expert_tokens = torch.randint(low=0,
|
|
high=config.max_tokens_per_expert,
|
|
size=(config.num_experts, ),
|
|
device="cuda",
|
|
dtype=torch.int32)
|
|
return BatchedMMTensors(A, B, C, num_expert_tokens)
|
|
|
|
|
|
def ref_impl(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
|
|
num_expert_tokens: torch.Tensor) -> torch.Tensor:
|
|
|
|
num_expert_tokens_cpu = num_expert_tokens.clone()
|
|
num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu")
|
|
num_experts = num_expert_tokens.size(0)
|
|
|
|
for e in range(num_experts):
|
|
num_tokens = num_expert_tokens_cpu[e]
|
|
C[e, :num_tokens, :] = A[e, :num_tokens, :] @ B[e].transpose(0, 1)
|
|
|
|
return C
|
|
|
|
|
|
@pytest.mark.parametrize("num_experts", [16, 32])
|
|
@pytest.mark.parametrize("max_tokens_per_expert",
|
|
[32, 64, 128, 192, 224, 256, 512])
|
|
@pytest.mark.parametrize("K", [128, 256, 1024])
|
|
@pytest.mark.parametrize("N", [128, 256, 512, 1024])
|
|
@pytest.mark.parametrize("dtype",
|
|
[torch.float32, torch.float16, torch.bfloat16])
|
|
def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
|
|
N: int, dtype: torch.dtype):
|
|
|
|
config = BatchedMMConfig(dtype, num_experts, max_tokens_per_expert, K, N)
|
|
tensors = BatchedMMTensors.make_tensors(config)
|
|
|
|
test_output = tensors.C
|
|
ref_output = test_output.clone()
|
|
|
|
compute_tl_dtype = {
|
|
torch.float16: tl.float16,
|
|
torch.bfloat16: tl.bfloat16,
|
|
torch.float32: tl.float32
|
|
}[test_output.dtype]
|
|
invoke_moe_batched_triton_kernel(
|
|
tensors.A,
|
|
tensors.B,
|
|
test_output,
|
|
tensors.num_expert_tokens,
|
|
compute_tl_dtype,
|
|
# Quantization data
|
|
None,
|
|
None,
|
|
None,
|
|
# Quantization schemes
|
|
False,
|
|
False,
|
|
False,
|
|
config={
|
|
"BLOCK_SIZE_M": 16,
|
|
"BLOCK_SIZE_N": 16,
|
|
"BLOCK_SIZE_K": 16
|
|
})
|
|
|
|
ref_output = ref_impl(tensors.A, tensors.B, ref_output,
|
|
tensors.num_expert_tokens)
|
|
|
|
rtol, atol = {
|
|
torch.float16: (6e-2, 6e-2),
|
|
torch.bfloat16: (6e-2, 6e-2),
|
|
torch.float32: (1e-2, 1e-2),
|
|
}[test_output.dtype]
|
|
|
|
torch.testing.assert_close(test_output, ref_output, atol=atol, rtol=rtol)
|