| """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): |
| |
| |
| 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: |
| |
| |
| 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) |
| |
| |
| 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) |
| |
| |
| 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) |
|
|
