| """Fused Triton kernels for the draft step's elementwise chains (SM86). Each one has a plain-torch twin in ops.py and a |
| test (tests/test_fused.py) that holds it to that twin: bit-identical where the math allows, else <= 1 bf16 step. |
| |
| Unfused, the step ran ~1,400 kernels (the 20-iteration Sinkhorn alone ~480): ~2 ms of launch-bound work per round. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import triton |
| import triton.language as tl |
| from triton.language.extra import libdevice |
|
|
| from .ops import _markov_bias_kernel |
|
|
| |
|
|
|
|
| @triton.jit |
| def _hc_mix_partial_kernel(x_ptr, fn_ptr, part_ptr, K, CH, NS, BK: tl.constexpr): |
| """Split-K partial dot products of one token's flattened stream with the 24 mix rows, plus its sum of squares.""" |
| t = tl.program_id(0) |
| sp = tl.program_id(1) |
| r32 = tl.arange(0, 32) |
| rmask = r32 < 24 |
| acc = tl.zeros([32], tl.float32) |
| ss = tl.zeros([BK], tl.float32) |
| for k0 in range(sp * CH, sp * CH + CH, BK): |
| cols = k0 + tl.arange(0, BK) |
| x = tl.load(x_ptr + t.to(tl.int64) * K + cols).to(tl.float32) |
| ss += x * x |
| w = tl.load(fn_ptr + r32[:, None] * K + cols[None, :], mask=rmask[:, None], other=0.0) |
| acc += tl.sum(w * x[None, :], axis=1) |
| acc = tl.where(r32 == 24, tl.sum(ss, axis=0), acc) |
| tl.store(part_ptr + (t * NS + sp) * 32 + r32, acc) |
|
|
|
|
| @triton.jit |
| def _hc_mix_finish_kernel(part_ptr, scale_ptr, base_ptr, pre_ptr, post_ptr, comb_ptr, K, NS, rms_eps, hc_eps, |
| ITERS: tl.constexpr, NSP: tl.constexpr): |
| t = tl.program_id(0) |
| sidx = tl.arange(0, NSP) |
| smask = sidx < NS |
| r4 = tl.arange(0, 4) |
| r16 = tl.arange(0, 16) |
| rows = part_ptr + (t * NS + sidx) * 32 |
| acc_pre = tl.sum(tl.load(rows[:, None] + r4[None, :], mask=smask[:, None], other=0.0), axis=0) |
| acc_post = tl.sum(tl.load(rows[:, None] + 4 + r4[None, :], mask=smask[:, None], other=0.0), axis=0) |
| acc_comb = tl.sum(tl.load(rows[:, None] + 8 + r16[None, :], mask=smask[:, None], other=0.0), axis=0) |
| ss = tl.sum(tl.load(rows + 24, mask=smask, other=0.0), axis=0) |
| r = 1.0 / tl.sqrt(ss / K + rms_eps) |
| s0 = tl.load(scale_ptr) |
| s1 = tl.load(scale_ptr + 1) |
| s2 = tl.load(scale_ptr + 2) |
| pre = 1.0 / (1.0 + tl.exp(-(acc_pre * r * s0 + tl.load(base_ptr + r4)))) + hc_eps |
| post = 2.0 / (1.0 + tl.exp(-(acc_post * r * s1 + tl.load(base_ptr + 4 + r4)))) |
| c = tl.reshape(acc_comb * r * s2 + tl.load(base_ptr + 8 + r16), (4, 4)) |
| c = tl.exp(c - tl.max(c, axis=1)[:, None]) |
| c = c / tl.sum(c, axis=1)[:, None] + hc_eps |
| c = c / (tl.sum(c, axis=0)[None, :] + hc_eps) |
| for _ in range(ITERS - 1): |
| c = c / (tl.sum(c, axis=1)[:, None] + hc_eps) |
| c = c / (tl.sum(c, axis=0)[None, :] + hc_eps) |
| tl.store(pre_ptr + t * 4 + r4, pre) |
| tl.store(post_ptr + t * 4 + r4, post) |
| tl.store(comb_ptr + t * 16 + r16, tl.reshape(c, (16,))) |
|
|
|
|
| _HC_SPLITS = 40 |
|
|
|
|
| def hc_mixes(stream: torch.Tensor, fn: torch.Tensor, scale: torch.Tensor, base: torch.Tensor, rms_eps: float, |
| hc_eps: float, iters: int): |
| """ops.hc_mixes: split-K GEMV of the flattened stream with the 24 mix rows (fixed-order sums, deterministic), |
| then sigmoid / softmax / Sinkhorn per token.""" |
| T, hc, H = stream.shape |
| K = hc * H |
| assert hc == 4 and stream.is_contiguous() and fn.is_contiguous() and fn.shape == (24, K) |
| ns = _HC_SPLITS |
| assert K % (ns * 128) == 0 |
| part = torch.empty(T, ns, 32, dtype=torch.float32, device=stream.device) |
| _hc_mix_partial_kernel[(T, ns)](stream, fn, part, K, K // ns, ns, BK=128, num_warps=4) |
| pre = torch.empty(T, 4, dtype=torch.float32, device=stream.device) |
| post = torch.empty(T, 4, dtype=torch.float32, device=stream.device) |
| comb = torch.empty(T, 4, 4, dtype=torch.float32, device=stream.device) |
| _hc_mix_finish_kernel[(T,)](part, scale, base, pre, post, comb, K, ns, rms_eps, hc_eps, ITERS=iters, |
| NSP=triton.next_power_of_2(ns), num_warps=1) |
| return pre, post, comb |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _collapse_norm_kernel(s_ptr, pre_ptr, w_ptr, out_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) |
| tl.store(out_ptr + t.to(tl.int64) * H + h, (w * (yf * r)).to(tl.bfloat16)) |
|
|
|
|
| def collapse_norm(stream: torch.Tensor, pre: torch.Tensor | None, w: torch.Tensor, eps: float) -> torch.Tensor: |
| """rms_norm(hc_collapse(stream, pre), w): collapse with pre-mix (or copy 0), round to bf16, RMSNorm.""" |
| T, hc, H = stream.shape |
| assert hc == 4 and stream.is_contiguous() |
| out = torch.empty(T, H, dtype=torch.bfloat16, device=stream.device) |
| _collapse_norm_kernel[(T,)](stream, pre if pre is not None else stream, w, out, H, eps, |
| HAS_PRE=pre is not None, BH=1024, num_warps=4, enable_fp_fusion=False) |
| return out |
|
|
|
|
| @triton.jit |
| def _collapse_kernel(s_ptr, pre_ptr, out_ptr, H, BH: tl.constexpr): |
| t = tl.program_id(0) |
| h = tl.program_id(1) * BH + tl.arange(0, BH) |
| m = h < H |
| base = s_ptr + t.to(tl.int64) * 4 * H |
| y = (tl.load(pre_ptr + t * 4) * tl.load(base + h, mask=m).to(tl.float32) |
| + tl.load(pre_ptr + t * 4 + 1) * tl.load(base + H + h, mask=m).to(tl.float32) |
| + tl.load(pre_ptr + t * 4 + 2) * tl.load(base + 2 * H + h, mask=m).to(tl.float32) |
| + tl.load(pre_ptr + t * 4 + 3) * tl.load(base + 3 * H + h, mask=m).to(tl.float32)) |
| tl.store(out_ptr + t.to(tl.int64) * H + h, y.to(tl.bfloat16), mask=m) |
|
|
|
|
| def collapse(stream: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: |
| T, hc, H = stream.shape |
| out = torch.empty(T, H, dtype=torch.bfloat16, device=stream.device) |
| _collapse_kernel[(T, triton.cdiv(H, 1024))](stream, pre, out, H, BH=1024, num_warps=4, enable_fp_fusion=False) |
| return out |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _hc_post_kernel(y_ptr, s_ptr, post_ptr, comb_ptr, out_ptr, H, BH: tl.constexpr): |
| t = tl.program_id(0) |
| h = tl.program_id(1) * BH + tl.arange(0, BH) |
| m = h < H |
| base = t.to(tl.int64) * 4 * H |
| y = tl.load(y_ptr + t.to(tl.int64) * H + h, mask=m).to(tl.float32) |
| s0 = tl.load(s_ptr + base + h, mask=m).to(tl.float32) |
| s1 = tl.load(s_ptr + base + H + h, mask=m).to(tl.float32) |
| s2 = tl.load(s_ptr + base + 2 * H + h, mask=m).to(tl.float32) |
| s3 = tl.load(s_ptr + base + 3 * H + h, mask=m).to(tl.float32) |
| for j in tl.static_range(4): |
| c0 = tl.load(comb_ptr + t * 16 + j) |
| c1 = tl.load(comb_ptr + t * 16 + 4 + j) |
| c2 = tl.load(comb_ptr + t * 16 + 8 + j) |
| c3 = tl.load(comb_ptr + t * 16 + 12 + j) |
| v = (c0 * s0 + c1 * s1 + c2 * s2 + c3 * s3) + tl.load(post_ptr + t * 4 + j) * y |
| tl.store(out_ptr + base + j * H + h, v.to(tl.bfloat16), mask=m) |
|
|
|
|
| def hc_post(y: torch.Tensor, stream: torch.Tensor, post: torch.Tensor, comb: torch.Tensor) -> torch.Tensor: |
| T, hc, H = stream.shape |
| out = torch.empty_like(stream) |
| _hc_post_kernel[(T, triton.cdiv(H, 1024))](y, stream, post, comb, out, H, BH=1024, num_warps=4, |
| enable_fp_fusion=False) |
| return out |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _rms_norm_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.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.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), mask=m) |
|
|
|
|
| def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: |
| """x [N, D] bf16 (or fp32: rounded to bf16 first, as ops.rms_norm(x.to(bf16), ...)).""" |
| if x.dtype != torch.bfloat16: |
| x = x.to(torch.bfloat16) |
| x = x.contiguous() |
| N, D = x.shape |
| out = torch.empty_like(x) |
| _rms_norm_kernel[(N,)](x, w, out, D, eps, BD=1024 if D >= 1024 else 512, num_warps=4, enable_fp_fusion=False) |
| return out |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _rope_kernel(x_ptr, out_ptr, pos_ptr, cs_ptr, ROWS_PER_TOKEN, D: tl.constexpr, RD: tl.constexpr, INV: tl.constexpr, |
| QDQ: tl.constexpr): |
| row = tl.program_id(0).to(tl.int64) |
| tok = row // ROWS_PER_TOKEN |
| p = tl.load(pos_ptr + tok) |
| NOPE: tl.constexpr = D - RD |
| i = tl.arange(0, RD // 2) |
| c = tl.load(cs_ptr + p * RD + i) |
| s = tl.load(cs_ptr + p * RD + RD // 2 + i) |
| if INV: |
| s = -s |
| x0 = tl.load(x_ptr + row * D + NOPE + 2 * i).to(tl.float32) |
| x1 = tl.load(x_ptr + row * D + NOPE + 2 * i + 1).to(tl.float32) |
| tl.store(out_ptr + row * D + NOPE + 2 * i, (x0 * c - x1 * s).to(tl.bfloat16)) |
| tl.store(out_ptr + row * D + NOPE + 2 * i + 1, (x1 * c + x0 * s).to(tl.bfloat16)) |
| g = tl.arange(0, 64) |
| for b in tl.static_range(NOPE // 64): |
| v = tl.load(x_ptr + row * D + b * 64 + g) |
| if QDQ: |
| vf = v.to(tl.float32) |
| amax = tl.maximum(tl.max(tl.abs(vf), axis=0) * (1.0 / 448.0), 2.0 ** -126) |
| bits = amax.to(tl.uint32, bitcast=True) |
| e = ((bits >> 23) & 0xFF).to(tl.int32) - 127 + tl.where((bits & 0x7FFFFF) != 0, 1, 0) |
| sc = ((e + 127) << 23).to(tl.uint32).to(tl.float32, bitcast=True) |
| q = tl.minimum(tl.maximum(vf / sc, -448.0), 448.0) |
| |
| qb = q.to(tl.uint32, bitcast=True) |
| qe = tl.maximum(((qb >> 23) & 0xFF).to(tl.int32) - 127, -6) |
| quant = ((qe - 3 + 127) << 23).to(tl.uint32).to(tl.float32, bitcast=True) |
| q = libdevice.rint(q / quant) * quant |
| v = (q * sc).to(tl.bfloat16) |
| tl.store(out_ptr + row * D + b * 64 + g, v) |
|
|
|
|
| def rope(x: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor, inverse: bool = False) -> torch.Tensor: |
| """ops.apply_rope for bf16 x [N, D] or [N, heads, D] (rotates the last 64 channels).""" |
| x = x.contiguous() |
| D = x.shape[-1] |
| rows = x.numel() // D |
| per = rows // x.shape[0] |
| out = torch.empty_like(x) |
| _rope_kernel[(rows,)](x, out, pos, cos_sin, per, D=D, RD=cos_sin.shape[-1], INV=inverse, QDQ=False, |
| num_warps=1, enable_fp_fusion=False) |
| return out |
|
|
|
|
| def kv_rope_qdq(kv: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor) -> torch.Tensor: |
| """ops.kv_record_qdq(ops.apply_rope(kv, pos)) for kv [N, 512] bf16: the V4 record values.""" |
| kv = kv.contiguous() |
| out = torch.empty_like(kv) |
| _rope_kernel[(kv.shape[0],)](kv, out, pos, cos_sin, 1, D=kv.shape[-1], RD=cos_sin.shape[-1], INV=False, |
| QDQ=True, num_warps=1, enable_fp_fusion=False) |
| return out |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _markov_argmax_kernel(ids_ptr, slot_ptr, cache_ptr, emb_ptr, w2_ptr, base_ptr, vals_ptr, idx_ptr, V, |
| cache_stride, w2_stride, emb_stride, base_stride, NBLK, ids_stride, R: tl.constexpr, |
| BLOCK_V: tl.constexpr, EMB_BY_TOKEN: tl.constexpr): |
| b = tl.program_id(0) |
| blk = tl.program_id(1) |
| cols = blk * BLOCK_V + tl.arange(0, BLOCK_V) |
| mask = cols < V |
| tok = tl.load(ids_ptr + b * ids_stride).to(tl.int64) |
| slot = tl.load(slot_ptr + tok) |
| if slot >= 0: |
| bias = tl.load(cache_ptr + slot.to(tl.int64) * cache_stride + cols, mask=mask, other=0.0) |
| else: |
| r = tl.arange(0, R) |
| if EMB_BY_TOKEN: |
| e = tl.load(emb_ptr + tok * emb_stride + r).to(tl.float32) |
| else: |
| e = tl.load(emb_ptr + b * emb_stride + r).to(tl.float32) |
| w = tl.load(w2_ptr + cols[:, None].to(tl.int64) * w2_stride + r[None, :], mask=mask[:, None], other=0.0) |
| bias = tl.sum(w.to(tl.float32) * e[None, :], axis=1).to(tl.bfloat16) |
| lg = tl.load(base_ptr + b.to(tl.int64) * base_stride + cols, mask=mask, other=0.0) |
| li = (lg.to(tl.float32) + bias.to(tl.float32)).to(tl.bfloat16).to(tl.float32) |
| li = tl.where(mask, li, float("-inf")) |
| v, i = tl.max(li, axis=0, return_indices=True, return_indices_tie_break_left=True) |
| tl.store(vals_ptr + b * NBLK + blk, v) |
| tl.store(idx_ptr + b * NBLK + blk, (blk * BLOCK_V + i).to(tl.int64)) |
|
|
|
|
| def markov_argmax(base: torch.Tensor, prev: torch.Tensor, emb: torch.Tensor, mb) -> torch.Tensor: |
| """argmax_v bf16(base[b, v] + bias(prev_b)[v]) per row, ties -> lowest id. base [B, V] bf16 (any row stride).""" |
| B, V = base.shape |
| assert base.stride(1) == 1 |
| BLOCK_V = 64 |
| nblk = triton.cdiv(V, BLOCK_V) |
| vals = torch.empty(B, nblk, dtype=torch.float32, device=base.device) |
| idx = torch.empty(B, nblk, dtype=torch.int64, device=base.device) |
| emb = emb.contiguous() |
| _markov_argmax_kernel[(B, nblk)](prev, mb.slot_of_token, mb.cache, emb, mb.w2, base, vals, idx, V, |
| mb.cache.stride(0), mb.w2.stride(0), emb.stride(0), base.stride(0), nblk, |
| prev.stride(0), R=emb.shape[1], BLOCK_V=BLOCK_V, EMB_BY_TOKEN=False, |
| num_warps=4) |
| best = vals.argmax(dim=-1, keepdim=True) |
| return idx.gather(1, best).squeeze(1) |
|
|
|
|
| @triton.jit |
| def _markov_cached_kernel(ids_ptr, slot_ptr, cache_ptr, base_ptr, vals_ptr, idx_ptr, V, cache_stride, base_stride, |
| NBLK, ids_stride, BLOCK_V: tl.constexpr): |
| """Cached previous token: block max of bf16(base + cache row) (first maximum = lowest id). Exits otherwise.""" |
| b = tl.program_id(0) |
| blk = tl.program_id(1) |
| tok = tl.load(ids_ptr + b * ids_stride).to(tl.int64) |
| slot = tl.load(slot_ptr + tok) |
| if slot >= 0: |
| cols = blk * BLOCK_V + tl.arange(0, BLOCK_V) |
| mask = cols < V |
| bias = tl.load(cache_ptr + slot.to(tl.int64) * cache_stride + cols, mask=mask, other=0.0) |
| lg = tl.load(base_ptr + b.to(tl.int64) * base_stride + cols, mask=mask, other=0.0) |
| li = (lg.to(tl.float32) + bias.to(tl.float32)).to(tl.bfloat16).to(tl.float32) |
| li = tl.where(mask, li, float("-inf")) |
| v, i = tl.max(li, axis=0, return_indices=True, return_indices_tie_break_left=True) |
| tl.store(vals_ptr + b * NBLK + blk, v) |
| tl.store(idx_ptr + b * NBLK + blk, (blk * BLOCK_V + i).to(tl.int64)) |
|
|
|
|
| @triton.jit |
| def _markov_uncached_kernel(ids_ptr, slot_ptr, emb_ptr, w2_ptr, base_ptr, vals_ptr, idx_ptr, V, w2_stride, |
| emb_stride, base_stride, NBLK, ids_stride, R: tl.constexpr, BLOCK_V: tl.constexpr, |
| SUB: tl.constexpr): |
| """Uncached previous token: the serve's bias arithmetic (_markov_bias_kernel: [BLOCK_V, R] tile, fp32 products, |
| one bf16 rounding) over SUB consecutive BLOCK_V-column blocks, running max with strict > (lowest id on ties). |
| Exits for cached tokens.""" |
| b = tl.program_id(0) |
| pb = tl.program_id(1) |
| tok = tl.load(ids_ptr + b * ids_stride).to(tl.int64) |
| slot = tl.load(slot_ptr + tok) |
| if slot < 0: |
| r = tl.arange(0, R) |
| e = tl.load(emb_ptr + tok * emb_stride + r).to(tl.float32) |
| best_v = float("-inf") |
| best_i = (pb * SUB * BLOCK_V).to(tl.int64) |
| for sb in range(SUB): |
| cols = (pb * SUB + sb) * BLOCK_V + tl.arange(0, BLOCK_V) |
| mask = cols < V |
| w = tl.load(w2_ptr + cols[:, None].to(tl.int64) * w2_stride + r[None, :], mask=mask[:, None], other=0.0) |
| bias = tl.sum(w.to(tl.float32) * e[None, :], axis=1).to(tl.bfloat16) |
| lg = tl.load(base_ptr + b.to(tl.int64) * base_stride + cols, mask=mask, other=0.0) |
| li = (lg.to(tl.float32) + bias.to(tl.float32)).to(tl.bfloat16).to(tl.float32) |
| li = tl.where(mask, li, float("-inf")) |
| v, i = tl.max(li, axis=0, return_indices=True, return_indices_tie_break_left=True) |
| take = v > best_v |
| best_i = tl.where(take, ((pb * SUB + sb) * BLOCK_V + i).to(tl.int64), best_i) |
| best_v = tl.where(take, v, best_v) |
| tl.store(vals_ptr + b * NBLK + pb, best_v) |
| tl.store(idx_ptr + b * NBLK + pb, best_i) |
|
|
|
|
| @triton.jit |
| def _argmax_finish_kernel(ids_ptr, slot_ptr, vals_c_ptr, idx_c_ptr, vals_u_ptr, idx_u_ptr, out_ptr, NBLK_C, NBLK_U, |
| ids_stride, out_stride, col, BLOCK: tl.constexpr): |
| b = tl.program_id(0) |
| tok = tl.load(ids_ptr + b * ids_stride).to(tl.int64) |
| cached = tl.load(slot_ptr + tok) >= 0 |
| j = tl.arange(0, BLOCK) |
| vc = tl.load(vals_c_ptr + b * NBLK_C + j, mask=j < NBLK_C, other=float("-inf")) |
| vu = tl.load(vals_u_ptr + b * NBLK_U + j, mask=j < NBLK_U, other=float("-inf")) |
| v = tl.where(cached, vc, vu) |
| _, jm = tl.max(v, axis=0, return_indices=True, return_indices_tie_break_left=True) |
| if cached: |
| d = tl.load(idx_c_ptr + b * NBLK_C + jm) |
| else: |
| d = tl.load(idx_u_ptr + b * NBLK_U + jm) |
| tl.store(out_ptr + b * out_stride + col, d) |
|
|
|
|
| def markov_argmax_into(base: torch.Tensor, prev: torch.Tensor, mb, out: torch.Tensor, col: int) -> None: |
| """out[:, col] = markov_argmax(base, prev, mb.embed(prev), mb) (same values; ties -> lowest id). Three launches: |
| cached previous tokens take a wide-block kernel (2048 columns per program), uncached ones the serve's bias tile |
| looped over 8 blocks per program (each kernel exits for the other kind), and a finish kernel reduces the block |
| maxima straight into the output column. The Markov embedding row is read by token. prev may be a strided column |
| (out[:, col - 1]).""" |
| B, V = base.shape |
| assert base.stride(1) == 1 and out.dtype == torch.int64 |
| BC, BU, SUB = 2048, 64, 8 |
| nc, nu = triton.cdiv(V, BC), triton.cdiv(V, BU * SUB) |
| vals_c = torch.empty(B, nc, dtype=torch.float32, device=base.device) |
| idx_c = torch.empty(B, nc, dtype=torch.int64, device=base.device) |
| vals_u = torch.empty(B, nu, dtype=torch.float32, device=base.device) |
| idx_u = torch.empty(B, nu, dtype=torch.int64, device=base.device) |
| _markov_cached_kernel[(B, nc)](prev, mb.slot_of_token, mb.cache, base, vals_c, idx_c, V, mb.cache.stride(0), |
| base.stride(0), nc, prev.stride(0), BLOCK_V=BC, num_warps=8) |
| _markov_uncached_kernel[(B, nu)](prev, mb.slot_of_token, mb.w1, mb.w2, base, vals_u, idx_u, V, mb.w2.stride(0), |
| mb.w1.stride(0), base.stride(0), nu, prev.stride(0), R=mb.w1.shape[1], |
| BLOCK_V=BU, SUB=SUB, num_warps=4) |
| _argmax_finish_kernel[(B,)](prev, mb.slot_of_token, vals_c, idx_c, vals_u, idx_u, out, nc, nu, prev.stride(0), |
| out.stride(0), col, BLOCK=triton.next_power_of_2(max(nc, nu)), num_warps=4) |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _block_attn_kernel(q_ptr, ring_ptr, ring_pos_ptr, slots_ptr, lo_ptr, last_ptr, blk_ptr, sink_ptr, out_ptr, |
| scale, k, ring_slot_stride, ring_pos_stride, W: tl.constexpr, D: tl.constexpr, HH: tl.constexpr, |
| BM: tl.constexpr, KB: tl.constexpr, BD: tl.constexpr, DO: tl.constexpr): |
| """Program (b, m, n): BM query rows (row = t * HH + h) of request b, output dims [n * DO, (n + 1) * DO). |
| Scores over the W window ring rows of b's slot (masked to positions in [lo, last]) and its k block rows (padded |
| to KB), sink in the denominator, single-pass softmax, probabilities rounded to bf16 for P @ V. Keys = values.""" |
| b = tl.program_id(0) |
| rows = tl.program_id(1) * BM + tl.arange(0, BM) |
| d_out = tl.program_id(2) * DO + tl.arange(0, DO) |
| nrows = k * HH |
| rmask = rows < nrows |
| slot = tl.load(slots_ptr + b).to(tl.int64) |
| lo = tl.load(lo_ptr + b) |
| hi = tl.load(last_ptr + b) |
| jc = tl.arange(0, W) |
| jb = tl.arange(0, KB) |
| bmask = jb < k |
| cpos = tl.load(ring_pos_ptr + slot * ring_pos_stride + jc) |
| cvalid = (cpos >= lo) & (cpos <= hi) |
| ring_base = ring_ptr + slot * ring_slot_stride |
| blk_base = blk_ptr + b.to(tl.int64) * k * D |
| q_base = q_ptr + b.to(tl.int64) * nrows * D |
| sc = tl.zeros([BM, W], tl.float32) |
| sb = tl.zeros([BM, KB], tl.float32) |
| for d0 in tl.static_range(0, D, BD): |
| d = d0 + tl.arange(0, BD) |
| qc = tl.load(q_base + rows[:, None] * D + d[None, :], mask=rmask[:, None], other=0.0) |
| kc = tl.load(ring_base + jc[:, None] * D + d[None, :]) |
| kb = tl.load(blk_base + jb[:, None] * D + d[None, :], mask=bmask[:, None], other=0.0) |
| sc += tl.dot(qc, tl.trans(kc)) |
| sb += tl.dot(qc, tl.trans(kb)) |
| sc = tl.where(cvalid[None, :], sc * scale, float("-inf")) |
| sb = tl.where(bmask[None, :], sb * scale, float("-inf")) |
| sk = tl.load(sink_ptr + rows % HH, mask=rmask, other=0.0) |
| mx = tl.maximum(tl.maximum(tl.max(sc, axis=1), tl.max(sb, axis=1)), sk) |
| pc = tl.exp(sc - mx[:, None]) |
| pb = tl.exp(sb - mx[:, None]) |
| denom = tl.sum(pc, axis=1) + tl.sum(pb, axis=1) + tl.exp(sk - mx) |
| vc = tl.load(ring_base + jc[:, None] * D + d_out[None, :]) |
| vb = tl.load(blk_base + jb[:, None] * D + d_out[None, :], mask=bmask[:, None], other=0.0) |
| acc = tl.dot(pc.to(tl.bfloat16), vc) + tl.dot(pb.to(tl.bfloat16), vb) |
| o_base = out_ptr + b.to(tl.int64) * nrows * D |
| tl.store(o_base + rows[:, None] * D + d_out[None, :], (acc / denom[:, None]).to(tl.bfloat16), mask=rmask[:, None]) |
|
|
|
|
| def block_attention_ring(q: torch.Tensor, ring_stage: torch.Tensor, ring_pos: torch.Tensor, slots: torch.Tensor, |
| ctx_lo: torch.Tensor, last_pos: torch.Tensor, blk_kv: torch.Tensor, sink: torch.Tensor, |
| scale: float, BM: int | None = None, DO: int = 256, BD: int = 64, |
| num_warps: int | None = None) -> torch.Tensor: |
| """ops.block_attention with the context read straight from the window ring of each request's slot. |
| |
| q [B, k, Hh, D] bf16; ring_stage [S, W, D] (one stage's ring, any slot stride); ring_pos [S, W] int64; |
| slots / ctx_lo / last_pos int64 [B]; blk_kv [B, k, D] bf16; sink [Hh] fp32. Returns bf16 [B, k, Hh, D]. |
| Default tiles (3090, profiler kernel time per stage): <= 640 query rows BM 16 / 8 warps (14.7 us at 1x5x64), |
| else BM 32 / 4 warps (21.0 us at 4x5x64); the gathered torch path took 70-155 us.""" |
| B, k, Hh, D = q.shape |
| if BM is None: |
| BM, num_warps = (16, 8) if B * k * Hh <= 640 else (32, 4) |
| W = ring_stage.shape[1] |
| KB = 16 |
| assert k <= KB and ring_stage.stride(2) == 1 and ring_stage.stride(1) == D and ring_pos.stride(1) == 1 |
| assert W == triton.next_power_of_2(W) and D % DO == 0 and D % BD == 0 |
| q = q.contiguous() |
| blk_kv = blk_kv.contiguous() |
| out = torch.empty_like(q) |
| _block_attn_kernel[(B, triton.cdiv(k * Hh, BM), D // DO)]( |
| q, ring_stage, ring_pos, slots, ctx_lo, last_pos, blk_kv, sink, out, scale, k, ring_stage.stride(0), |
| ring_pos.stride(0), W=W, D=D, HH=Hh, BM=BM, KB=KB, BD=BD, DO=DO, num_warps=num_warps, num_stages=1) |
| return out |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _amax_kernel(x_ptr, amax_ptr, N, BLOCK: tl.constexpr): |
| offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) |
| v = tl.load(x_ptr + offs, mask=offs < N, other=0.0).to(tl.float32) |
| tl.atomic_max(amax_ptr, tl.max(tl.abs(v), axis=0)) |
|
|
|
|
| @triton.jit |
| def _prescale_kernel(x_ptr, amax_ptr, out_ptr, sc_ptr, N, BLOCK: tl.constexpr): |
| |
| |
| a = tl.maximum(tl.load(amax_ptr), 1.0) / 16384.0 |
| sc = libdevice.exp2(tl.maximum(libdevice.ceil(libdevice.log2(a)), 0.0)) |
| pid = tl.program_id(0) |
| offs = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK) |
| m = offs < N |
| v = tl.load(x_ptr + offs, mask=m, other=0.0).to(tl.float32) |
| tl.store(out_ptr + offs, (v / sc).to(tl.float16), mask=m) |
| if pid == 0: |
| tl.store(sc_ptr, sc) |
|
|
|
|
| def prescale_fp16(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """main_proj's input scaling (vllm_exl3_moe dense.py): x (any float) -> (fp16 x / sc, fp32 sc [1]), sc the power |
| of two >= 1 that brings max |x| under 2^14. Two launches (+ a memset) instead of the ~12 torch kernels.""" |
| x = x.contiguous() |
| N = x.numel() |
| amax = torch.zeros(1, dtype=torch.float32, device=x.device) |
| out = torch.empty(x.shape, dtype=torch.float16, device=x.device) |
| sc = torch.empty(1, dtype=torch.float32, device=x.device) |
| BLOCK = 4096 |
| grid = (triton.cdiv(N, BLOCK),) |
| _amax_kernel[grid](x, amax, N, BLOCK=BLOCK, num_warps=8) |
| _prescale_kernel[grid](x, amax, out, sc, N, BLOCK=BLOCK, num_warps=8, enable_fp_fusion=False) |
| return out, sc |
|
|
|
|
| @triton.jit |
| def _scaled_rms_norm_kernel(y_ptr, sc_ptr, w_ptr, out_ptr, D, eps, BD: tl.constexpr): |
| r_ = tl.program_id(0).to(tl.int64) |
| sc = tl.load(sc_ptr) |
| 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(y_ptr + r_ * D + d, mask=m, other=0.0) * sc).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(y_ptr + r_ * D + d, mask=m, other=0.0) * sc).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), mask=m) |
|
|
|
|
| def scaled_rms_norm(y: torch.Tensor, sc: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: |
| """rms_norm(y * sc, w) for fp32 y [N, D] and a device scalar sc: the product rounded to bf16 first, as |
| rms_norm(x.to(bf16)) after the plugin's y * sc.""" |
| y = y.contiguous() |
| N, D = y.shape |
| out = torch.empty(N, D, dtype=torch.bfloat16, device=y.device) |
| _scaled_rms_norm_kernel[(N,)](y, sc, w, out, D, eps, BD=1024 if D >= 1024 else 512, num_warps=4, |
| enable_fp_fusion=False) |
| return out |
|
|
|
|
| @triton.jit |
| def _kv_store_kernel(y_ptr, w_ptr, normed_ptr, pos_ptr, cs_ptr, slot_ptr, ring_ptr, ring_pos_ptr, eps, |
| ring_slot_stride, ring_pos_stride, W, D: tl.constexpr, RD: tl.constexpr, |
| STORE_POS: tl.constexpr): |
| """Row r: kv_norm (as _rms_norm_kernel) -> RoPE + V4 record (as _rope_kernel QDQ) -> ring[slot, pos % W].""" |
| row = tl.program_id(0).to(tl.int64) |
| d = tl.arange(0, D) |
| x = tl.load(y_ptr + row * D + d).to(tl.bfloat16).to(tl.float32) |
| ss = x * x |
| r = 1.0 / tl.sqrt(tl.sum(ss, axis=0) / D + eps) |
| w = tl.load(w_ptr + d).to(tl.float32) |
| tl.store(normed_ptr + row * D + d, (w * (x * r)).to(tl.bfloat16)) |
| tl.debug_barrier() |
| p = tl.load(pos_ptr + row) |
| slot = tl.load(slot_ptr + row).to(tl.int64) |
| dst = ring_ptr + slot * ring_slot_stride + (p % W) * D |
| NOPE: tl.constexpr = D - RD |
| i = tl.arange(0, RD // 2) |
| c = tl.load(cs_ptr + p * RD + i) |
| s = tl.load(cs_ptr + p * RD + RD // 2 + i) |
| x0 = tl.load(normed_ptr + row * D + NOPE + 2 * i).to(tl.float32) |
| x1 = tl.load(normed_ptr + row * D + NOPE + 2 * i + 1).to(tl.float32) |
| tl.store(dst + NOPE + 2 * i, (x0 * c - x1 * s).to(tl.bfloat16)) |
| tl.store(dst + NOPE + 2 * i + 1, (x1 * c + x0 * s).to(tl.bfloat16)) |
| g = tl.arange(0, 64) |
| for b in tl.static_range(NOPE // 64): |
| vf = tl.load(normed_ptr + row * D + b * 64 + g).to(tl.float32) |
| amax = tl.maximum(tl.max(tl.abs(vf), axis=0) * (1.0 / 448.0), 2.0 ** -126) |
| bits = amax.to(tl.uint32, bitcast=True) |
| e = ((bits >> 23) & 0xFF).to(tl.int32) - 127 + tl.where((bits & 0x7FFFFF) != 0, 1, 0) |
| sc = ((e + 127) << 23).to(tl.uint32).to(tl.float32, bitcast=True) |
| q = tl.minimum(tl.maximum(vf / sc, -448.0), 448.0) |
| qb = q.to(tl.uint32, bitcast=True) |
| qe = tl.maximum(((qb >> 23) & 0xFF).to(tl.int32) - 127, -6) |
| quant = ((qe - 3 + 127) << 23).to(tl.uint32).to(tl.float32, bitcast=True) |
| q = libdevice.rint(q / quant) * quant |
| tl.store(dst + b * 64 + g, (q * sc).to(tl.bfloat16)) |
| if STORE_POS: |
| tl.store(ring_pos_ptr + slot * ring_pos_stride + p % W, p) |
|
|
|
|
| def kv_store(y: torch.Tensor, w: torch.Tensor, eps: float, pos: torch.Tensor, cos_sin: torch.Tensor, |
| ring_stage: torch.Tensor, slots: torch.Tensor, ring_pos: torch.Tensor | None = None) -> None: |
| """ring_stage[slots, pos % W] = kv_rope_qdq(rms_norm(y, w), pos) for fp32 y [N, D] (wkv output); also |
| ring_pos[slots, pos % W] = pos when ring_pos is given. Callers pass distinct (slot, pos % W) per call.""" |
| y = y.contiguous() |
| N, D = y.shape |
| W = ring_stage.shape[1] |
| assert ring_stage.stride(2) == 1 and ring_stage.stride(1) == D and (ring_pos is None or ring_pos.stride(1) == 1) |
| normed = torch.empty(N, D, dtype=torch.bfloat16, device=y.device) |
| _kv_store_kernel[(N,)](y, w, normed, pos, cos_sin, slots, ring_stage, ring_pos if ring_pos is not None else pos, |
| eps, ring_stage.stride(0), ring_pos.stride(0) if ring_pos is not None else 0, W, D=D, |
| RD=cos_sin.shape[-1], STORE_POS=ring_pos is not None, num_warps=4, enable_fp_fusion=False) |
|
|
|
|
| |
|
|
|
|
| @triton.jit |
| def _wo_a_fp8_kernel(o_ptr, w_ptr, sc_ptr, out_ptr, T, o_stride, out_stride, C: tl.constexpr, R: tl.constexpr, |
| BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): |
| """out[t, g*R + r] = sum_c o[t, g*C + c] * (w[g*R + r, c] * sc[(g*R + r) // 32, c // 32]); the dequantized weight is |
| exact in bf16, products accumulate in fp32 (bf16 MMA), bf16 out. BN <= 32 and BN | 32: one scale row per tile. |
| e4m3 -> value: the byte's sign and 7 magnitude bits placed as fp16 bits (exponent field = e4m3 exponent, mantissa |
| top bits = e4m3 mantissa) are the value / 2^8 exactly, subnormals included; x (2^8 * scale) is exact.""" |
| g = tl.program_id(0) |
| r0 = g * R + tl.program_id(1) * BN |
| t = tl.arange(0, BM) |
| r = r0 + tl.arange(0, BN) |
| tmask = t < T |
| acc = tl.zeros([BM, BN], tl.float32) |
| for k0 in range(0, C, BK): |
| c = k0 + tl.arange(0, BK) |
| a = tl.load(o_ptr + t[:, None].to(tl.int64) * o_stride + g * C + c[None, :], mask=tmask[:, None], other=0.0) |
| b16 = tl.load(w_ptr + r[:, None].to(tl.int64) * C + c[None, :]).to(tl.int16) |
| h = (((b16 & 0x80) << 8) | ((b16 & 0x7F) << 7)).to(tl.float16, bitcast=True) |
| scv = tl.load(sc_ptr + (r0 // 32) * (C // 32) + c // 32) * 256.0 |
| w = (h.to(tl.float32) * scv[None, :]).to(tl.bfloat16) |
| acc += tl.dot(a, tl.trans(w)) |
| tl.store(out_ptr + t[:, None].to(tl.int64) * out_stride + r[None, :], acc.to(tl.bfloat16), mask=tmask[:, None]) |
|
|
|
|
| def wo_a_fp8(o: torch.Tensor, w8: torch.Tensor, sc: torch.Tensor, groups: int, BN: int = 32, BK: int = 128, |
| num_warps: int = 2, num_stages: int | None = None) -> torch.Tensor: |
| """o [T, groups * C] bf16 (rows may be strided), w8 [groups * R, C] uint8 (e4m3fn bytes, no NaN), |
| sc [groups*R/32, C/32] fp32 (UE8M0 values) -> [T, groups * R] bf16: the engine's W8A16 wo_a |
| (Sm12xMarlinMxfp8Bmm) on SM86, reading the FP8 bytes (32 MiB per stage instead of 64 MiB of expanded bf16).""" |
| T = o.shape[0] |
| GR, C = w8.shape |
| R = GR // groups |
| assert o.shape[1] == groups * C and o.stride(1) == 1 and w8.is_contiguous() and sc.is_contiguous() |
| assert 32 % BN == 0 and R % 32 == 0 and C % BK == 0 |
| out = torch.empty(T, GR, dtype=torch.bfloat16, device=o.device) |
| BM = max(16, triton.next_power_of_2(T)) |
| if num_stages is None: |
| num_stages = 3 if BM == 16 else 2 |
| _wo_a_fp8_kernel[(groups, R // BN)](o, w8, sc, out, T, o.stride(0), out.stride(0), C=C, R=R, BM=BM, BN=BN, |
| BK=BK, num_warps=num_warps, num_stages=num_stages) |
| return out |
|
|
