File size: 9,399 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """Lane E fusions of the draft step's small replicated chains (used by tp2.py at tp = 1 and 2; O2's fused.py is
untouched). Each has a plain-torch twin it is tested against (tests/test_fused_e.py):
route_fused gate logits (split-K fp32 dot, fixed order) -> sqrt(softplus) -> + bias -> top-k -> renormalized
weights x route_scale -> logical ids, physical slots (perm), fp16 weights: 2 launches instead of
~10 per stage (cuBLAS sgemm + split-K reduce, softplus, sqrt, add, sort, topk, gather, sum, div,
mul, index, cast). The logits differ from cuBLAS in rounding only (a different summation order),
so a routing decision can differ only at a 3rd / 4th near-tie (the class O2 already measured
between the 3090 and the engine); ties of equal scores go to the lower expert id.
swiglu_fp16 silu(min(g, lim)) * clamp(u, -lim, lim) in fp32 -> fp16 (torch's silu formula x / (1 + exp(-x))):
1 launch instead of 5
collapse_norm2 fused.collapse_norm that also stores the fp16 conversion of its bf16 output (the next GEMM's input)
rms_norm_fp16 fused.rms_norm storing fp16(bf16 result) (wq_a -> q_norm -> wq_b)
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
from triton.language.extra import libdevice
# --------------------------------------------------------------------------------------------- routing
_ROUTE_SPLITS = 80
@triton.jit
def _route_partial_kernel(x_ptr, w_ptr, part_ptr, T, H, CH, NS, E: tl.constexpr, BK: tl.constexpr):
"""Split sp: partial logits of all T tokens over columns [sp * CH, (sp + 1) * CH) (CH == BK): the gate tile is
loaded once (coalesced rows) and kept in registers for every token; fp32 products, fixed-order sums."""
sp = tl.program_id(0)
e = tl.arange(0, E)
cols = sp * CH + tl.arange(0, BK)
w = tl.load(w_ptr + e[:, None].to(tl.int64) * H + cols[None, :])
for t in range(T):
x = tl.load(x_ptr + t * H + cols).to(tl.float32)
tl.store(part_ptr + (t * NS + sp) * E + e, tl.sum(w * x[None, :], axis=1))
@triton.jit
def _route_finish_kernel(part_ptr, bias_ptr, perm_ptr, ids_ptr, sel_ptr, w32_ptr, w16_ptr, NS, scale,
E: tl.constexpr, TOPK: tl.constexpr, NSP: tl.constexpr, KP: tl.constexpr):
t = tl.program_id(0)
e = tl.arange(0, E)
sidx = tl.arange(0, NSP)
p = tl.load(part_ptr + (t * NS + sidx[:, None]) * E + e[None, :], mask=sidx[:, None] < NS, other=0.0)
logits = tl.sum(p, axis=0)
sp = tl.where(logits > 20.0, logits, libdevice.log1p(libdevice.exp(logits))) # torch softplus, threshold 20
scores = tl.sqrt_rn(sp)
sel_s = scores + tl.load(bias_ptr + e)
k = tl.arange(0, KP)
ids = tl.zeros([KP], tl.int32)
ws = tl.zeros([KP], tl.float32)
for j in tl.static_range(TOPK):
m = tl.max(sel_s, axis=0)
i = tl.min(tl.where(sel_s == m, e, E), axis=0) # lowest id among equal scores
ids = tl.where(k == j, i, ids)
ws = tl.where(k == j, tl.sum(tl.where(e == i, scores, 0.0), axis=0), ws)
sel_s = tl.where(e == i, float("-inf"), sel_s)
kmask = k < TOPK
w0 = tl.sum(tl.where(k == 0, ws, 0.0), axis=0)
w1 = tl.sum(tl.where(k == 1, ws, 0.0), axis=0)
w2 = tl.sum(tl.where(k == 2, ws, 0.0), axis=0)
denom = (w0 + w1) + w2 + 1e-20
wn = ws / denom * scale
phys = tl.load(perm_ptr + ids, mask=kmask, other=0)
tl.store(ids_ptr + t * TOPK + k, ids.to(tl.int64), mask=kmask)
tl.store(sel_ptr + t * TOPK + k, phys, mask=kmask)
tl.store(w32_ptr + t * TOPK + k, wn, mask=kmask)
tl.store(w16_ptr + t * TOPK + k, wn.to(tl.float16), mask=kmask)
def route_fused(x: torch.Tensor, gate_w: torch.Tensor, bias: torch.Tensor, perm: torch.Tensor, top_k: int,
route_scale: float):
"""x bf16 [T, H], gate_w fp32 [E, H], bias fp32 [E], perm int64 [E] (logical -> physical slot).
Returns (ids int64 [T, k], sel int64 [T, k] = perm[ids], w fp32 [T, k], w fp16 [T, k])."""
T, H = x.shape
E = gate_w.shape[0]
assert top_k == 3 and x.is_contiguous() and gate_w.is_contiguous() and E == triton.next_power_of_2(E)
ns = _ROUTE_SPLITS
assert H % ns == 0 and (H // ns) == triton.next_power_of_2(H // ns)
dev = x.device
part = torch.empty(T, ns, E, dtype=torch.float32, device=dev)
_route_partial_kernel[(ns,)](x, gate_w, part, T, H, H // ns, ns, E=E, BK=H // ns, num_warps=4)
ids = torch.empty(T, top_k, dtype=torch.int64, device=dev)
sel = torch.empty(T, top_k, dtype=torch.int64, device=dev)
w32 = torch.empty(T, top_k, dtype=torch.float32, device=dev)
w16 = torch.empty(T, top_k, dtype=torch.float16, device=dev)
_route_finish_kernel[(T,)](part, bias, perm, ids, sel, w32, w16, ns, float(route_scale), E=E, TOPK=top_k,
NSP=triton.next_power_of_2(ns), KP=4, num_warps=4, enable_fp_fusion=False)
return ids, sel, w32, w16
# --------------------------------------------------------------------------------------------- SwiGLU
@triton.jit
def _swiglu_kernel(g_ptr, u_ptr, out_ptr, N, lim, BLOCK: tl.constexpr):
offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
m = offs < N
g = tl.minimum(tl.load(g_ptr + offs, mask=m, other=0.0), lim)
u = tl.minimum(tl.maximum(tl.load(u_ptr + offs, mask=m, other=0.0), -lim), lim)
a = g / (1.0 + libdevice.exp(-g)) * u
tl.store(out_ptr + offs, a.to(tl.float16), mask=m)
def swiglu_fp16(g: torch.Tensor, u: torch.Tensor, lim: float) -> torch.Tensor:
"""fp16(silu(g.clamp(max=lim)) * u.clamp(-lim, lim)) for fp32 g, u (contiguous, same shape)."""
g, u = g.contiguous(), u.contiguous()
out = torch.empty(g.shape, dtype=torch.float16, device=g.device)
N = g.numel()
_swiglu_kernel[(triton.cdiv(N, 1024),)](g, u, out, N, float(lim), BLOCK=1024, num_warps=4,
enable_fp_fusion=False)
return out
# --------------------------------------------------------------------------------------------- norms with fp16 out
@triton.jit
def _collapse_norm2_kernel(s_ptr, pre_ptr, w_ptr, out_ptr, out16_ptr, H, eps, HAS_PRE: tl.constexpr,
BH: tl.constexpr):
t = tl.program_id(0)
base = s_ptr + t.to(tl.int64) * 4 * H
ss = tl.zeros([BH], tl.float32)
if HAS_PRE:
p0 = tl.load(pre_ptr + t * 4)
p1 = tl.load(pre_ptr + t * 4 + 1)
p2 = tl.load(pre_ptr + t * 4 + 2)
p3 = tl.load(pre_ptr + t * 4 + 3)
for h0 in range(0, H, BH):
h = h0 + tl.arange(0, BH)
if HAS_PRE:
y = (p0 * tl.load(base + h).to(tl.float32) + p1 * tl.load(base + H + h).to(tl.float32)
+ p2 * tl.load(base + 2 * H + h).to(tl.float32) + p3 * tl.load(base + 3 * H + h).to(tl.float32))
yb = y.to(tl.bfloat16)
else:
yb = tl.load(base + h)
yf = yb.to(tl.float32)
ss += yf * yf
tl.store(out_ptr + t.to(tl.int64) * H + h, yb)
r = 1.0 / tl.sqrt(tl.sum(ss, axis=0) / H + eps)
for h0 in range(0, H, BH):
h = h0 + tl.arange(0, BH)
yf = tl.load(out_ptr + t.to(tl.int64) * H + h).to(tl.float32)
w = tl.load(w_ptr + h).to(tl.float32)
ob = (w * (yf * r)).to(tl.bfloat16)
tl.store(out_ptr + t.to(tl.int64) * H + h, ob)
tl.store(out16_ptr + t.to(tl.int64) * H + h, ob.to(tl.float32).to(tl.float16))
def collapse_norm2(stream: torch.Tensor, pre: torch.Tensor | None, w: torch.Tensor, eps: float):
"""(fused.collapse_norm(stream, pre, w, eps), its fp16 conversion): bf16 [T, H] and fp16 [T, H]."""
T, hc, H = stream.shape
assert hc == 4 and stream.is_contiguous()
out = torch.empty(T, H, dtype=torch.bfloat16, device=stream.device)
out16 = torch.empty(T, H, dtype=torch.float16, device=stream.device)
_collapse_norm2_kernel[(T,)](stream, pre if pre is not None else stream, w, out, out16, H, eps,
HAS_PRE=pre is not None, BH=1024, num_warps=4, enable_fp_fusion=False)
return out, out16
@triton.jit
def _rms_norm_fp16_kernel(x_ptr, w_ptr, out_ptr, D, eps, BD: tl.constexpr):
r_ = tl.program_id(0).to(tl.int64)
ss = tl.zeros([BD], tl.float32)
for d0 in range(0, D, BD):
d = d0 + tl.arange(0, BD)
m = d < D
x = tl.load(x_ptr + r_ * D + d, mask=m, other=0.0).to(tl.bfloat16).to(tl.float32)
ss += x * x
r = 1.0 / tl.sqrt(tl.sum(ss, axis=0) / D + eps)
for d0 in range(0, D, BD):
d = d0 + tl.arange(0, BD)
m = d < D
x = tl.load(x_ptr + r_ * D + d, mask=m, other=0.0).to(tl.bfloat16).to(tl.float32)
w = tl.load(w_ptr + d, mask=m, other=0.0).to(tl.float32)
tl.store(out_ptr + r_ * D + d, (w * (x * r)).to(tl.bfloat16).to(tl.float32).to(tl.float16), mask=m)
def rms_norm_fp16(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
"""fused.rms_norm(x, w, eps).to(float16) for fp32 or bf16 x [N, D] (rounded to bf16 first, like rms_norm)."""
x = x.contiguous()
N, D = x.shape
out = torch.empty(N, D, dtype=torch.float16, device=x.device)
_rms_norm_fp16_kernel[(N,)](x, w, out, D, eps, BD=1024 if D >= 1024 else 512, num_warps=4,
enable_fp_fusion=False)
return out
|