| """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 |
|
|
| |
|
|
| _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))) |
| 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) |
| 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 |
|
|
|
|
| |
|
|
|
|
| @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 |
|
|
|
|
| |
|
|
|
|
| @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 |
|
|
