mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
Daniel Koepping
49 lines
2.8 KiB
Python
49 lines
2.8 KiB
Python
import functools, pathlib
|
|
from tinygrad import Tensor, dtypes
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
|
|
from extra.llama_kernels import alloc_like, compile_hip
|
|
|
|
@functools.cache
|
|
def _inverse_rows(out:UOp, dest_row:UOp, counts:UOp, off:UOp) -> UOp:
|
|
groups, rows = out.shape
|
|
entries, experts = dest_row.shape[1], counts.shape[1]
|
|
assert rows >= entries and counts.shape == (groups, experts) and off.shape == (groups, experts+1)
|
|
sink = UOp.sink(out.base, dest_row.base, counts.base, off.base, UOp.special(256, "lidx0"),
|
|
UOp.special(groups*((rows+255)//256), "gidx0"), arg=KernelInfo(f"dispatch_inverse_{groups}_{rows}_{entries}_{experts}"))
|
|
src = (pathlib.Path(__file__).parent/"inverse.cpp").read_text()
|
|
lib = compile_hip(src, [f"-DROWS={rows}", f"-DENTRIES={entries}", f"-DEXPERTS={experts}"])
|
|
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
|
|
|
|
def inverse_rows(dest_row:Tensor, counts:Tensor, off:Tensor, rows:int) -> Tensor:
|
|
out = alloc_like((dest_row.shape[0], rows), dtypes.int32, dest_row.device, dest_row.uop.axis)
|
|
return Tensor.custom_kernel(out, dest_row, counts, off, fxn=_inverse_rows)[0]
|
|
|
|
@functools.cache
|
|
def _dispatch_gather(out:UOp, x:UOp, dest_row:UOp, src_row:UOp) -> UOp:
|
|
from extra.gemm.moe_routing import _blk_for
|
|
groups, tokens, hidden = x.shape
|
|
rows = src_row.shape[1]
|
|
assert out.shape == (groups*rows, hidden) and dest_row.shape == (groups, tokens*4)
|
|
assert x.dtype == out.dtype and x.dtype.itemsize == 1
|
|
block = _blk_for(hidden)
|
|
g, row = UOp.range(groups, 0, AxisType.GLOBAL), UOp.range(rows, 1, AxisType.GLOBAL)
|
|
outer, inner = UOp.range(hidden//block, 2), UOp.range(block, 3, AxisType.LOCAL)
|
|
col = outer*block + inner
|
|
source = src_row.index(g, row).load().cast(dtypes.weakint)
|
|
anchor = dest_row.index(g, 0).load() >= 0
|
|
valid = (source >= 0) & anchor
|
|
token = valid.where(source, source.const_like(0)) // 4
|
|
value = valid.where(x.index(g, token, col).load(), UOp.const(0, x.dtype))
|
|
return out.reshape(groups, rows, hidden).index(g, row, col).store(value).end(g, row, outer, inner).sink(
|
|
arg=KernelInfo(f"dispatch_gather_{groups}_{rows}_{tokens}_{hidden}", opts_to_apply=()))
|
|
|
|
def _dispatch_gather_bwd(gradient:UOp, kernel:UOp) -> tuple:
|
|
from extra.gemm.moe_routing import _gscatter_bwd
|
|
# The first three arguments match scatter's backward; src_row is forward-only.
|
|
x = kernel.src[2]
|
|
return (*_gscatter_bwd(gradient.reshape(x.shape[0], -1, x.shape[-1]), kernel), None)
|
|
|
|
def dispatch_gather(x:Tensor, dest_row:Tensor, src_row:Tensor) -> Tensor:
|
|
out = alloc_like((x.shape[0]*src_row.shape[1], x.shape[-1]), x.dtype, x.device, x.uop.axis)
|
|
return Tensor.custom_kernel(out, x, dest_row, src_row, fxn=_dispatch_gather, grad_fxn=_dispatch_gather_bwd)[0]
|