fabolous005
d1ccc5c294
Measured on the live 16k agent setup (A/B with runtime flags, greedy answers vs reference): decode 3.49 -> 5.24 tok/s, first token 4.3 -> 3.3 s, 2.3k-token prompt 27.5 -> 20.6 s, same answers. Dropped after measuring: route stats on the CPU, a Triton router GEMV, pregate prefetch from disk (all within noise or slower). The GEMV keeps one config instead of autotuning: Triton benchmarks with a 256 MiB buffer that torch keeps cached, which left 8 MiB of VRAM and made lazy kernel loads fail (hipErrorNoBinaryForGpu). prain also returns torch cached blocks after materializing and whenever the driver has < 128 MiB free. quality.py score --greedy-only skips prompt logprobs, which OOM the 16k config.
112 lines
6.0 KiB
Python
112 lines
6.0 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 Kernels(unittest.TestCase):
|
|
def setUp(self):
|
|
torch.manual_seed(0)
|
|
|
|
def test_unpack_writes_what_to_gpu_int4_computes(self):
|
|
from prain.kernels import unpack_int4
|
|
from prain.quant import to_gpu_int4
|
|
|
|
for n, half in ((256, 128), (1408, 1408), (2816, 352)):
|
|
cpu_layout = torch.randint(0, 256, (n, half), dtype=torch.uint8, device="cuda")
|
|
dst = torch.zeros_like(cpu_layout)
|
|
unpack_int4(cpu_layout.flatten(), dst)
|
|
self.assertTrue(torch.equal(dst, to_gpu_int4(cpu_layout)), (n, half))
|
|
|
|
def test_decode_kernel_computes_what_fused_experts_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.kernels import moe_int4_decode
|
|
from prain.quant import dequantize, moe_quant_config, quantize
|
|
from prain.tiered import cpu_experts, slot_table
|
|
|
|
hidden, inter, experts, topk = 256, 128, 6, 4
|
|
w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
|
|
w2 = torch.randn(experts, hidden, inter) * 0.1
|
|
parts = [quantize(w, "int4") for w in (w13, w2)] # gpu layout
|
|
held = [1, 3, 4] # the round's experts, in slots 0..2
|
|
slots = [[q[held].cuda(), s[held].cuda()] for q, s in parts]
|
|
table = slot_table({e: i for i, e in enumerate(held)}, experts, "cuda")
|
|
for tokens in (1, 2, 3, 4):
|
|
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
|
ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).int()
|
|
ids[0, -1] = experts # padding route: the table's last row
|
|
w = torch.rand(tokens, topk)
|
|
got = moe_int4_decode(x.cuda(), w.cuda(), ids.cuda(), table, *slots[0], *slots[1])
|
|
with set_current_vllm_config(VllmConfig()):
|
|
vllm = fused_experts(x.cuda(), slots[0][0], slots[1][0], w.cuda(), ids.cuda(),
|
|
activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=table.numel(), expert_map=table,
|
|
quant_config=moe_quant_config("int4", *[s[1] for s in slots]))
|
|
exact = cpu_experts(x, w, torch.where(ids < experts, ids, -1), set(held), lambda e: tuple(
|
|
dequantize(q, s, "int4")[e] for q, s in parts))
|
|
self.assertEqual((got.shape, got.dtype), ((tokens, hidden), torch.bfloat16))
|
|
torch.testing.assert_close(got.float().cpu(), exact, rtol=3e-2, atol=3e-2)
|
|
torch.testing.assert_close(got.float(), vllm.float(), rtol=3e-2, atol=3e-2)
|
|
|
|
def test_staged_copy_fills_a_slot_like_the_repack_path(self):
|
|
from prain.experts import _Staging, stage_int4
|
|
from prain.quant import quantize, to_cpu_int4, to_gpu_int4
|
|
|
|
w13, w2 = torch.randn(2, 256, 256), torch.randn(2, 256, 128)
|
|
pack = [[to_cpu_int4(q), s] for q, s in (quantize(w, "int4") for w in (w13, w2))]
|
|
slots = [[torch.zeros((3, *t.shape[1:]), dtype=t.dtype, device="cuda") for t in mats]
|
|
for mats in pack]
|
|
ring = _Staging(sum(t[0].nbytes for mats in pack for t in mats), "cuda")
|
|
stream = torch.cuda.current_stream()
|
|
for expert, slot in ((1, 2), (0, 0), (1, 1), (0, 2), (1, 0)): # wraps the ring
|
|
stage_int4([[t[expert] for t in mats] for mats in pack], slots, slot, ring, stream)
|
|
torch.cuda.synchronize()
|
|
for mats, dsts in zip(pack, slots):
|
|
self.assertTrue(torch.equal(dsts[0][slot], to_gpu_int4(mats[0][expert].cuda())))
|
|
self.assertTrue(torch.equal(dsts[1][slot], mats[1][expert].cuda()))
|
|
|
|
|
|
def test_int4_rounds_on_rdna2_match_the_bf16_kernel_even_when_fp16_overflows(self):
|
|
from types import SimpleNamespace
|
|
|
|
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.experts import TieredRoutedExperts
|
|
from prain.quant import moe_quant_config, quantize
|
|
from prain.tiered import slot_table
|
|
|
|
hidden, inter, experts = 256, 128, 4
|
|
parts = [quantize(w, "int4") for w in (torch.randn(experts, 2 * inter, hidden) * 0.1,
|
|
torch.randn(experts, hidden, inter) * 0.1)]
|
|
slots = [[q.cuda(), s.cuda()] for q, s in parts]
|
|
config = moe_quant_config("int4", slots[0][1], slots[1][1])
|
|
layer = SimpleNamespace(_slots=slots, activation=MoEActivation.GELU_TANH, _qconfig=config)
|
|
table = slot_table({e: e for e in range(experts)}, experts, "cuda")
|
|
with set_current_vllm_config(VllmConfig()):
|
|
for tokens, scale in ((2, 1.0), (64, 1.0), (64, 3000.0)): # decode kernel, fp16, overflow
|
|
x = (torch.randn(tokens, hidden) * scale).bfloat16().cuda()
|
|
ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int().cuda()
|
|
w = torch.rand(tokens, 2).cuda()
|
|
got = TieredRoutedExperts._int4_gfx10(layer, x, w, ids, table)
|
|
want = fused_experts(x, slots[0][0], slots[1][0], w, ids, activation=layer.activation,
|
|
global_num_experts=table.numel(), expert_map=table,
|
|
quant_config=config)
|
|
self.assertTrue(torch.isfinite(got).all(), (tokens, scale))
|
|
torch.testing.assert_close(got.float(), want.float(), rtol=3e-2, atol=3e-2 * scale,
|
|
msg=f"{tokens} tokens, x * {scale}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|