mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 23:43:44 +08:00
Daniel Koepping
213 lines
10 KiB
Python
213 lines
10 KiB
Python
import functools, math, pathlib
|
|
from tinygrad import Tensor, dtypes
|
|
from tinygrad.helpers import getenv
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
|
|
from tinygrad.renderer import Estimates
|
|
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
|
|
|
|
BLOCK_ROW = 256
|
|
|
|
@functools.cache
|
|
def _router_mfma_fwd(out:UOp, x:UOp, weight:UOp, bias:UOp, *, dname:str) -> UOp:
|
|
*lead, K = x.shape
|
|
M = math.prod(lead)
|
|
E = weight.shape[0]
|
|
threads = UOp.special(256, "lidx0")
|
|
workgroups = UOp.special((M + 63) // 64, "gidx0")
|
|
sink = UOp.sink(out.base, x.base, weight.base, bias.base, threads, workgroups,
|
|
arg=KernelInfo(f"moe_router_mfma_{M}_{K}_{E}", estimates=Estimates(ops=2*M*E*K, mem=(M*K+E*K+E)*2+M*E*4)))
|
|
amd = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
|
|
src = (amd/"moe_router_mfma.cpp").read_text()
|
|
lib = HIPCCCompiler("gfx950", [f"-I{(amd/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
|
|
f"-DROUTER_M={M}", f"-DROUTER_K={K}", f"-DROUTER_E={E}"]).compile_cached(src)
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
|
|
|
|
def _router_mfma_bwd(gradient:UOp, kernel:UOp) -> tuple:
|
|
_, x_u, weight_u, bias_u = kernel.src[1:5]
|
|
x, weight, bias = (Tensor(u, device=u.device) for u in (x_u, weight_u, bias_u))
|
|
reference = x.float() @ weight.float().T + bias.float()
|
|
grad_x, grad_weight, grad_bias = reference.gradient(x, weight, bias, gradient=Tensor(gradient, device=x_u.device))
|
|
return None, grad_x.uop, grad_weight.uop, grad_bias.uop
|
|
|
|
def router_mfma(x:Tensor, weight:Tensor, bias:Tensor) -> Tensor:
|
|
assert x.ndim >= 2 and weight.ndim == 2 and bias.ndim == 1
|
|
K = x.shape[-1]
|
|
E = weight.shape[0]
|
|
assert weight.shape == (E, K) and bias.shape == (E,)
|
|
assert x.dtype == weight.dtype == bias.dtype == dtypes.bfloat16
|
|
assert E == 32 and K % 64 == 0
|
|
if isinstance(x.device, tuple):
|
|
assert x.uop.axis == 0, f"router MFMA requires axis-0 sharding, got axis={x.uop.axis}"
|
|
local_shape = x.uop.shard_shape
|
|
assert local_shape[-1] == K and math.prod(local_shape[:-1]) % 64 == 0, f"unsupported local router shape {local_shape}"
|
|
else:
|
|
assert math.prod(x.shape[:-1]) % 64 == 0
|
|
x, weight, bias = x.contiguous(), weight.contiguous(), bias.contiguous()
|
|
out = _sharded_invalids((*x.shape[:-1], E), dtypes.float32, x.device)
|
|
out, *_ = Tensor.custom_kernel(out, x, weight, bias,
|
|
fxn=functools.partial(_router_mfma_fwd, dname=str(x.device)), grad_fxn=_router_mfma_bwd)
|
|
return out
|
|
|
|
def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
|
|
if isinstance(device, tuple):
|
|
per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
|
|
return Tensor(per.uop.unshard(0), device=device)
|
|
return Tensor.invalids(*shape, dtype=dtype, device=device)
|
|
|
|
def _atomic_add(device:str) -> str:
|
|
return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
|
|
else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
|
|
|
|
def _blk_for(D:int) -> int:
|
|
blk = 64
|
|
while D % blk: blk //= 2
|
|
return blk
|
|
|
|
def _kv_ranges(G, N, D, BLK):
|
|
g = UOp.range(G, 0)
|
|
m = UOp.range(N, 1)
|
|
jo = UOp.range(D // BLK, 2)
|
|
ji = UOp.range(BLK, 3, AxisType.LOCAL)
|
|
return g, m, jo * BLK + ji, jo, ji
|
|
|
|
def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
|
|
G, M, D = out.shape
|
|
g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
|
|
row = idx.index(g, m).cast(dtypes.weakint)
|
|
val = table.index(g, row, j).load()
|
|
return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
|
|
arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
|
|
|
|
def _ggather_zero_kernel(out:UOp) -> UOp:
|
|
i = UOp.range(out.numel(), 0)
|
|
return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
|
|
|
|
def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
|
|
return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
|
|
|
|
def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
|
|
_, table_u, idx_u = kernel.src[1:4]
|
|
dev = table_u.device
|
|
device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
|
|
G, R, D = table_u.shape
|
|
gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
|
|
go = Tensor(gradient, device=dev)
|
|
atomic_str = _atomic_add(device)
|
|
def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
|
|
Gk, M, Dk = gout.shape
|
|
g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
|
|
row = idx.index(g, m).cast(dtypes.weakint)
|
|
val = gout.index(g, m, j).load().cast(dtypes.float32)
|
|
atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
|
|
return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
|
|
grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
|
|
return (None, grad_table.cast(table_u.dtype).uop, None)
|
|
|
|
def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
|
|
G, R, D = table.shape
|
|
M = idx.shape[1]
|
|
out = _sharded_invalids((G, M, D), table.dtype, table.device)
|
|
return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
|
|
|
|
def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
|
|
G, M, D = out.shape
|
|
k = idx.shape[1] // src.shape[1]
|
|
g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
|
|
row = idx.index(g, m).cast(dtypes.weakint)
|
|
val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
|
|
return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
|
|
arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
|
|
|
|
def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
|
|
_, src_u, idx_u = kernel.src[1:4]
|
|
dev = src_u.device
|
|
G, T_l, D = src_u.shape
|
|
k = idx_u.shape[1] // T_l
|
|
if getenv("GGATHER_SUM_HIP", 0):
|
|
from extra.gptoss_kernels.gather_sum import gather_sum
|
|
assert k == 4
|
|
return None, gather_sum(Tensor(gradient), Tensor(idx_u)).uop, None
|
|
sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
|
|
# FP8 dispatch carries BF16 gradients back to the quantizer, not FP8-rounded gradients.
|
|
grad_dtype = dtypes.bfloat16 if src_u.dtype == dtypes.fp8e4m3 else src_u.dtype
|
|
return (None, sel.reshape(G, T_l, k, D).sum(2).cast(grad_dtype).uop, None)
|
|
|
|
def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
|
|
G, T_l, D = src.shape
|
|
zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
|
|
return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
|
|
|
|
def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
|
|
return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
|
|
|
|
class Routing:
|
|
def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int,
|
|
topi:Tensor|None=None, counts:Tensor|None=None):
|
|
self.weights, self.dest_row = weights, dest_row
|
|
self.off = off
|
|
self.topi = topi
|
|
self.counts = counts
|
|
self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
|
|
|
|
@property
|
|
def tile_e(self) -> Tensor:
|
|
G, E = self.off.shape[0], self.off.shape[1] - 1
|
|
tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
|
|
tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
|
|
return ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32).reshape(-1)
|
|
|
|
@property
|
|
def rows_e(self) -> Tensor:
|
|
return self.tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
|
|
|
|
def n_groups_of(t:Tensor) -> int:
|
|
return len(t.device) if isinstance(t.device, tuple) else 1
|
|
|
|
def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
|
|
T, E = logits.shape
|
|
k, G = experts_per_tok, n_groups_of(logits)
|
|
assert T % G == 0, f"tokens {T} must split across {G} devices"
|
|
T_l = T // G
|
|
|
|
if getenv("FUSED_ROUTER_TOPK", 0):
|
|
from extra.gptoss_kernels.router_topk import fused_router_topk
|
|
weights, topi = fused_router_topk(logits.reshape(G, T_l, E))
|
|
else:
|
|
topv, topi = logits.reshape(G, T_l, E).topk(k)
|
|
weights = topv.softmax(-1)
|
|
return route_topk(weights, topi, n_experts)
|
|
|
|
def route_topk(weights:Tensor, topi:Tensor, n_experts:int) -> Routing:
|
|
G, T_l, k = weights.shape
|
|
E, m_l = n_experts, m_max_for(T_l, k, n_experts)
|
|
m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
|
|
|
|
counts = m.sum(1)
|
|
pad = ((counts + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
|
|
off = pad.cumsum(1).pad(((0, 0), (1, 0)))
|
|
dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
|
|
return Routing(weights, dest_row, off, m_l, G, T_l, topi=topi, counts=counts)
|
|
|
|
def dispatch(x:Tensor, r:Routing) -> Tensor:
|
|
G, D = r.n_groups, x.shape[-1]
|
|
return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
|
|
|
|
def dispatch_fp8(x:Tensor|tuple[Tensor, Tensor], r:Routing) -> tuple[Tensor, Tensor]:
|
|
from extra.gptoss_kernels.quantize_mxfp8 import quantize_mxfp8_fused_qe8
|
|
q, e8 = x if isinstance(x, tuple) else quantize_mxfp8_fused_qe8(x)
|
|
if getenv("DISPATCH_GATHER", 0):
|
|
from extra.gptoss_kernels.dispatch import inverse_rows, dispatch_gather
|
|
assert r.counts is not None
|
|
src_row = inverse_rows(r.dest_row, r.counts, r.off, r.m_l)
|
|
return tuple(dispatch_gather(t.reshape(r.n_groups, r.t_local, -1), r.dest_row, src_row) for t in (q, e8))
|
|
return dispatch(q, r), dispatch(e8, r)
|
|
|
|
def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
|
|
G, D, k = r.n_groups, y.shape[-1], experts_per_tok
|
|
if getenv("FUSED_COMBINE", 0):
|
|
from extra.gptoss_kernels.combine import fused_combine
|
|
assert k == 4
|
|
return fused_combine(y.reshape(G, r.m_l, D), r.dest_row, r.weights.reshape(G, r.t_local, k)).reshape(n_tokens, D)
|
|
sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
|
|
return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)
|