95 lines
5.3 KiB
Python
95 lines
5.3 KiB
Python
import importlib.util
|
|
import unittest
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
torch = None
|
|
|
|
GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
|
|
|
|
|
|
@unittest.skipUnless(GPU, "needs vLLM and a GPU")
|
|
class KernelMatch(unittest.TestCase):
|
|
def test_cpu_experts_compute_what_the_vllm_kernel_computes(self):
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
|
|
|
from prain.tiered import cpu_experts
|
|
|
|
torch.manual_seed(0)
|
|
tokens, hidden, inter, experts = 8, 256, 128, 4
|
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
|
w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
|
|
w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
|
|
topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
|
|
topk_weights = torch.rand(tokens, 2)
|
|
with set_current_vllm_config(VllmConfig()):
|
|
gpu = fused_experts(x.cuda(), w13.cuda(), w2.cuda(), topk_weights.cuda(),
|
|
topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=experts).float().cpu()
|
|
cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
|
|
torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
|
|
|
|
def test_a_round_computes_only_the_experts_its_slot_table_maps(self):
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
|
|
|
from prain.tiered import cpu_experts, slot_table
|
|
|
|
torch.manual_seed(0)
|
|
tokens, hidden, inter, experts = 64, 256, 128, 6
|
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
|
w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
|
|
w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
|
|
topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
|
|
topk_ids[0, 1] = experts # a padding route
|
|
topk_weights = torch.rand(tokens, 2)
|
|
table = slot_table({1: 0, 4: 1}, experts, "cuda")
|
|
with set_current_vllm_config(VllmConfig()):
|
|
gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(),
|
|
topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=table.numel(), expert_map=table).float().cpu()
|
|
cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e]))
|
|
torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
|
|
|
|
def test_quantized_slots_compute_what_the_cpu_path_computes(self):
|
|
"""Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that
|
|
holds experts 1 and 3, against the CPU kernels on the pack layout and against the
|
|
dequantized weights."""
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
|
|
|
from prain.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4,
|
|
to_gpu_int4)
|
|
from prain.tiered import cpu_experts, slot_table
|
|
|
|
torch.manual_seed(0)
|
|
tokens, hidden, inter, experts = 8, 256, 128, 4
|
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
|
w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
|
|
w2 = torch.randn(experts, hidden, inter) * 0.1
|
|
topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
|
|
topk_weights = torch.rand(tokens, 2)
|
|
table = slot_table({1: 0, 3: 1}, experts, "cuda")
|
|
for quant in ("int8", "int4"):
|
|
parts = [quantize(w, quant) for w in (w13, w2)] # [q, scale], gpu layout
|
|
pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts]
|
|
slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(),
|
|
s[[1, 3]].cuda()] for q, s in pack]
|
|
with set_current_vllm_config(VllmConfig()):
|
|
gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(),
|
|
topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=table.numel(), expert_map=table,
|
|
quant_config=moe_quant_config(quant, slots[0][1], slots[1][1]))
|
|
cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
|
|
cpu_matmul(quant, q[e], s[e]) for q, s in pack))
|
|
exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
|
|
dequantize(q, s, quant)[e] for q, s in parts))
|
|
self.assertGreater(exact.norm().item(), 1.0)
|
|
torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant)
|
|
torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant)
|