File size: 5,932 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 | """Original SM86 native-packed W4A16 kernels. No native FP4/FP8 MMA or atomics on results."""
from __future__ import annotations
import triton as tr
import triton.language as tl
@tr.jit
def bf16_decode(code, scale):
# E2M1 magnitude as BF16 fields, then exact exponent adjustment. Integer
# construction also handles scale-0 subnormals; no FTZ multiplication here.
mag = code & 7
exp = tl.where(mag < 2, -1, tl.where(mag < 4, 0, tl.where(mag < 6, 1, 2)))
mant = tl.where((mag == 3) | (mag == 5) | (mag == 7), 64, 0)
biased = exp + scale
normal = (biased << 7) | mant
sub = (128 | mant) >> (1 - biased)
bits = tl.where(biased > 0, normal, sub)
bits = tl.where(biased >= 255, 0x7f80, bits)
bits = tl.where(mag == 0, 0, bits) | ((code & 8) << 12)
bits = tl.where(scale == 255, 0x7fc0, bits)
return bits.to(tl.uint16).to(tl.bfloat16, bitcast=True)
@tr.jit
def decode_kernel(C, S, O, COUNT: tl.constexpr, B: tl.constexpr):
i = tl.program_id(0) * B + tl.arange(0, B)
c = tl.load(C+i, i < COUNT, 0).to(tl.int32)
s = tl.load(S+i, i < COUNT, 127).to(tl.int32)
tl.store(O+i, bf16_decode(c, s), i < COUNT)
@tr.jit
def _weights(W, S, e, k, n, N: tl.constexpr, K: tl.constexpr):
p = tl.load(W + (e * (K//2) + k[:, None]//2) * N + n[None, :],
(k[:, None] < K) & (n[None, :] < N), 0).to(tl.int32)
c = (p >> ((k[:, None] & 1) * 4)) & 15
s = tl.load(S + (e * (K//32) + k[:, None]//32) * N + n[None, :],
(k[:, None] < K) & (n[None, :] < N), 127).to(tl.int32)
return bf16_decode(c, s)
@tr.jit
def linear_kernel(X, IDS, W, S, U, US, O, OU, STATUS,
N: tl.constexpr, K: tl.constexpr, E: tl.constexpr,
TOP: tl.constexpr, PER_SLOT: tl.constexpr,
FUSED: tl.constexpr, MMA: tl.constexpr,
BN: tl.constexpr, BK: tl.constexpr):
route = tl.program_id(1)
expert = tl.load(IDS+route)
valid = (expert >= 0) & (expert < E)
if not valid:
tl.atomic_or(STATUS, 1)
expert = tl.minimum(tl.maximum(expert, 0), E-1)
row = route if PER_SLOT else route // TOP
n = tl.program_id(0)*BN + tl.arange(0, BN)
if MMA:
m = tl.arange(0, 16)
acc = tl.full((16, BN), 0, tl.float32)
if FUSED:
up = tl.full((16, BN), 0, tl.float32)
for start in range(tr.cdiv(K, BK)):
k = start*BK + tl.arange(0, BK)
x = tl.load(X + row*K + k[None, :], (m[:, None] == 0) & (k[None, :] < K), 0)
w = _weights(W, S, expert, k, n, N, K)
acc = tl.dot(x, w, acc, out_dtype=tl.float32)
if FUSED:
u = _weights(U, US, expert, k, n, N, K)
up = tl.dot(x, u, up, out_dtype=tl.float32)
value = tl.sum(tl.where(m[:, None] == 0, acc, 0), axis=0)
if FUSED:
uvalue = tl.sum(tl.where(m[:, None] == 0, up, 0), axis=0)
else:
# One/few row SIMT control: FP32 product/reduction, decoded BF16 values.
# Only BK independent partials persist, rather than the whole matrix.
acc = tl.full((BK, BN), 0, tl.float32)
if FUSED:
up = tl.full((BK, BN), 0, tl.float32)
for start in range(tr.cdiv(K, BK)):
k = start*BK + tl.arange(0, BK)
x = tl.load(X+row*K+k, k < K, 0).to(tl.float32)
w = _weights(W, S, expert, k, n, N, K).to(tl.float32)
acc += x[:, None] * w
if FUSED:
u = _weights(U, US, expert, k, n, N, K).to(tl.float32)
up += x[:, None] * u
value = tl.sum(acc, axis=0)
if FUSED:
uvalue = tl.sum(up, axis=0)
tl.store(O+route*N+n, tl.where(valid, value, float('nan')), n < N)
if FUSED:
tl.store(OU+route*N+n, tl.where(valid, uvalue, float('nan')), n < N)
@tr.jit
def activation_kernel(G, U, R, A, COUNT: tl.constexpr, N: tl.constexpr,
LIMIT: tl.constexpr, B: tl.constexpr):
i = tl.program_id(0)*B + tl.arange(0, B)
# Official linear outputs round to BF16 before clamp/SwiGLU; routing is
# applied before the BF16 down input, not moved through its trained QDQ.
gate = tl.minimum(tl.load(G+i, i < COUNT, 0).to(tl.bfloat16).to(tl.float32), LIMIT)
up = tl.minimum(tl.maximum(tl.load(U+i, i < COUNT, 0).to(tl.bfloat16).to(tl.float32), -LIMIT), LIMIT)
route = tl.load(R+i//N, i < COUNT, 0)
act = gate / (1 + tl.exp(-gate)) * up * route
tl.store(A+i, act.to(tl.bfloat16), i < COUNT)
@tr.jit
def combine_kernel(P, R, O, N: tl.constexpr, TOP: tl.constexpr, B: tl.constexpr):
row = tl.program_id(1)
n = tl.program_id(0)*B + tl.arange(0, B)
total = tl.full((B,), 0, tl.float32)
for slot in tl.static_range(TOP):
weight = tl.load(R+row*TOP+slot)
part = tl.load(P+(row*TOP+slot)*N+n, n < N, 0)
total = total + part * weight
tl.store(O+row*N+n, total, n < N)
@tr.jit
def native_combine_kernel(P, IDS, O, N: tl.constexpr, TOP: tl.constexpr, B: tl.constexpr):
row = tl.program_id(1)
n = tl.program_id(0)*B + tl.arange(0, B)
total = tl.full((B,), 0, tl.float32)
# Official expert-loop order. Stable slot order also handles duplicate IDs
# in diagnostic inputs. TP reduction must already be complete per route.
for order in tl.static_range(TOP):
value = tl.full((B,), 0, tl.float32)
for slot in tl.static_range(TOP):
eid = tl.load(IDS+row*TOP+slot)
rank = tl.full((), 0, tl.int32)
for other in tl.static_range(TOP):
other_id = tl.load(IDS+row*TOP+other)
rank += ((other_id < eid) | ((other_id == eid) & (other < slot))).to(tl.int32)
part = tl.load(P+(row*TOP+slot)*N+n, n < N, 0).to(tl.bfloat16).to(tl.float32)
value = tl.where(rank == order, part, value)
total = total + value
tl.store(O+row*N+n, total, n < N)
|