112 lines
5.4 KiB
Python
112 lines
5.4 KiB
Python
"""Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors.
|
|
|
|
usage: python scripts/tune_moe.py OUT_DIR [FORMAT...] (scripts/env.sh points
|
|
VLLM_TUNED_CONFIG_FOLDER at build/moe-configs)
|
|
Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map;
|
|
up to 32 tokens route into the slots like decode, more tokens spread like prefill),
|
|
Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM
|
|
config files for every slot count E (the best config hardly depends on E). Every config in
|
|
the files compiled, ran and matched the default config's output on this GPU.
|
|
"""
|
|
|
|
import itertools
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import torch
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe import override_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config,
|
|
int8_w8a16_moe_quant_config)
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name,
|
|
get_default_config)
|
|
|
|
from prain.quant import GROUP, quantize
|
|
|
|
H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16
|
|
SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048)
|
|
FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])}
|
|
|
|
|
|
def candidates(tokens):
|
|
# ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per
|
|
# format); widen it when a GPU compiles faster.
|
|
if tokens <= 32:
|
|
space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,))
|
|
else:
|
|
space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,))
|
|
for bm, bn, bk, group, warps, stages in space:
|
|
yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group,
|
|
"SPLIT_K": 1, "num_warps": warps, "num_stages": stages} # int4 needs SPLIT_K
|
|
|
|
|
|
def timed(fn, repeat):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
for _ in range(repeat):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
return (time.perf_counter() - start) * 1000 / repeat
|
|
|
|
|
|
def main(out_dir, formats):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
torch.manual_seed(0)
|
|
w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16()
|
|
w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16()
|
|
table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda")
|
|
table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda")
|
|
for name in formats:
|
|
dtype, block_shape = FORMATS[name]
|
|
if name == "bf16":
|
|
a, b, quant_config = w13, w2, None
|
|
else:
|
|
(a, s13), (b, s2) = quantize(w13, name), quantize(w2, name)
|
|
make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config
|
|
quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None,
|
|
block_shape=block_shape)
|
|
best = {}
|
|
for tokens in SIZES:
|
|
x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16()
|
|
pool = SLOTS if tokens <= 32 else EXPERTS
|
|
ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK]
|
|
for _ in range(tokens)]).to(torch.int32)
|
|
weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1)
|
|
|
|
def run():
|
|
return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH,
|
|
global_num_experts=EXPERTS + 1, expert_map=table,
|
|
quant_config=quant_config)
|
|
|
|
default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape)
|
|
with override_config(default): # vLLM's own choice when no config file exists
|
|
default_ms = timed(run, 10 if tokens <= 32 else 3)
|
|
reference = run()
|
|
results = [(default_ms, json.dumps(default, sort_keys=True))] # never worse than it
|
|
for config in candidates(tokens):
|
|
# catch inside: vLLM's override_config skips its restore when the body raises
|
|
with override_config(config):
|
|
try:
|
|
ms = timed(run, 10 if tokens <= 32 else 3)
|
|
out = run()
|
|
except Exception: # noqa: BLE001 - configs that do not compile or fit
|
|
continue
|
|
if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6:
|
|
results.append((ms, json.dumps(config, sort_keys=True)))
|
|
ms, config = min(results)
|
|
best[str(tokens)] = json.loads(config)
|
|
print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms "
|
|
f"({len(results)} configs) {config}", flush=True)
|
|
for e in range(1, EXPERTS + 1): # any slot count a placement can give a layer
|
|
with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f:
|
|
json.dump(best, f, indent=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
with set_current_vllm_config(VllmConfig()):
|
|
main(sys.argv[1], sys.argv[2:] or list(FORMATS))
|