File size: 12,347 Bytes
3b1bffa | 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 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | # SPDX-License-Identifier: Apache-2.0
# FP8 (e4m3, double_row) variant of the GELU MLP. EXPERIMENTAL, opt-in. See NOTICE.
"""fc1 + GELU + fc2 with FP8 e4m3 ``double_row`` matmuls. EXPERIMENTAL, opt-in (``VIT_ENC_MLP_FP8=1``).
Same math as ``mlp.nki_mlp_k4_xp``: ``gelu(x @ w1.T + b1) @ w2.T + b2``, but both matmuls run in
FP8 ``double_row`` mode (two 128-row halves of a 256-deep contraction per instruction):
* weights are host-packed bf16 ``[K/256, 128, 2, F]`` (:func:`pack_fp8_weight`) and cast to e4m3
in SBUF once per kernel;
* x^T is staged as bf16 ``[128, rows]`` per 128-row K tile (as in ``mlp``), then paired into
``[128, 2, rows]`` e4m3 tiles (rows ``2j`` and ``2j+1`` of the K-tile list);
* the GELU output ``g^T`` tiles are paired the same way for fc2;
* PSUM accumulates in fp32; bias, GELU and the output cast are unchanged.
e4m3 on trn2 is the legacy format (max 240); values are cast with the hardware's saturating
convert. Requires K % 256 == 0 and H % 256 == 0. INFERENCE ONLY.
"""
from __future__ import annotations
import torch
from torch_neuronx import nki_op
import nki
import nki.isa as nisa
import nki.language as nl
import os as _os
LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
_P_MAX = 128
_F_MOV_MAX = 512
_BF16 = nl.bfloat16
_E4 = nl.float8_e4m3
def _div_ceil(n, d):
return (n + d - 1) // d
def _kernel_assert(cond, msg):
assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
def _mlp_fp8_body(x_hbm, w1p_hbm, b1_hbm, w2p_hbm, b2_hbm, inv_sw1, g_scale, out_scale, a_hbm, eps, fuse):
"""x [R,K] bf16; w1p [K/256,128,2,H] bf16 (= pack(fc1.weight.T)); b1 [128,H/128] f32;
w2p [H/256,128,2,K] bf16 (= pack(fc2.weight.T)); b2 [K] f32 -> [R,K] bf16."""
R, K = x_hbm.shape
nk2, P, TWO, H = w1p_hbm.shape
nh2 = w2p_hbm.shape[0]
_kernel_assert(P == _P_MAX and TWO == 2 and nk2 * 256 == K, "w1p must be [K/256,128,2,H]")
_kernel_assert(nh2 * 256 == H and w2p_hbm.shape[3] == K, "w2p must be [H/256,128,2,K]")
_kernel_assert(R % LNC == 0, "R must be divisible by LNC")
rows_per = R // LNC
pid = nl.program_id(0)
r_base = pid * rows_per
if fuse:
xnew_hbm = nl.ndarray((R, K), dtype=nl.float32, buffer=nl.shared_hbm)
z_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
b2_full = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=b2_full, src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
acc_s = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
neg_shift = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
neg_mu = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
var = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
rstd = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
inv_k = 1.0 / float(K)
neg_inv_k = -inv_k
else:
out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
n_h = H // _P_MAX
n_ko = _div_ceil(K, _F_MOV_MAX)
b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
# weights: one DMA per 256-deep chunk (host-packed), cast to e4m3 once
w1_8 = []
for c in range(nk2):
wb = nl.ndarray((_P_MAX, 2, H), dtype=_BF16, buffer=nl.sbuf)
nisa.dma_copy(dst=wb, src=w1p_hbm[c])
w8 = nl.ndarray((_P_MAX, 2, H), dtype=_E4, buffer=nl.sbuf)
nisa.tensor_copy(dst=w8, src=wb)
w1_8.append(w8)
w2_8 = []
for c in range(nh2):
wb = nl.ndarray((_P_MAX, 2, K), dtype=_BF16, buffer=nl.sbuf)
nisa.dma_copy(dst=wb, src=w2p_hbm[c])
w8 = nl.ndarray((_P_MAX, 2, K), dtype=_E4, buffer=nl.sbuf)
nisa.tensor_copy(dst=w8, src=wb)
w2_8.append(w8)
rb_size = _F_MOV_MAX
n_rb = _div_ceil(rows_per, rb_size)
for rb in range(n_rb):
r_local = rb * rb_size
pr = min(rb_size, rows_per - r_local)
n_rs = _div_ceil(pr, _P_MAX)
r_abs = r_base + r_local
# x^T paired into [128, 2, pr] e4m3: pair c holds K rows [c*256, c*256+128) and [c*256+128, c*256+256)
xt8 = []
for c in range(nk2):
xb = nl.ndarray((_P_MAX, 2, rb_size), dtype=_BF16, buffer=nl.sbuf)
for half in range(2):
k0 = c * 256 + half * _P_MAX
nisa.dma_transpose(dst=xb[:, half, 0:pr],
src=x_hbm[r_abs:r_abs + pr, k0:k0 + _P_MAX])
x8 = nl.ndarray((_P_MAX, 2, rb_size), dtype=_E4, buffer=nl.sbuf)
nisa.tensor_copy(dst=x8[:, :, 0:pr], src=xb[:, :, 0:pr])
xt8.append(x8)
# fc1 (H out-rows in 128 tiles), GELU into bf16, then pair two H tiles into one e4m3 [128,2,pr]
g8 = []
for hp in range(nh2):
gb = nl.ndarray((_P_MAX, 2, rb_size), dtype=_BF16, buffer=nl.sbuf)
for half in range(2):
ht = hp * 2 + half
h0 = ht * _P_MAX
ps = nl.ndarray((_P_MAX, rb_size), dtype=nl.float32, buffer=nl.psum)
for c in nl.affine_range(nk2):
nisa.nc_matmul(dst=ps[:, 0:pr], stationary=w1_8[c][:, :, h0:h0 + _P_MAX],
moving=xt8[c][:, :, 0:pr], perf_mode=nisa.matmul_perf_mode.double_row)
# PSUM holds (x @ (w1*sw1)); gelu(scale*psum + b1) with scale = 1/sw1 undoes the weight scale
nisa.activation(dst=gb[:, half, 0:pr], op=nl.gelu, data=ps[:, 0:pr],
bias=b1_sb[:, ht:ht + 1], scale=inv_sw1)
# scale the GELU output into e4m3 range (legacy e4m3 max = 240) before the cast
gs = nl.ndarray((_P_MAX, 2, rb_size), dtype=nl.float32, buffer=nl.sbuf)
nisa.tensor_scalar(dst=gs[:, :, 0:pr], data=gb[:, :, 0:pr], op0=nl.multiply, operand0=g_scale)
gq = nl.ndarray((_P_MAX, 2, rb_size), dtype=_E4, buffer=nl.sbuf)
nisa.tensor_copy(dst=gq[:, :, 0:pr], src=gs[:, :, 0:pr])
g8.append(gq)
# fc2: out[r, k] = sum_h g[r,h] w2[k,h]; stationary = g^T pair [128,2,128 rows], moving = w2 pair
if fuse:
for rs in range(n_rs):
rr0 = rs * _P_MAX
prs = min(_P_MAX, pr - rr0)
r0 = r_abs + rr0
row = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
z_t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
nisa.dma_copy(dst=row[0:prs, 0:K], src=a_hbm[r0:r0 + prs, 0:K])
for ko in range(n_ko):
ko0 = ko * _F_MOV_MAX
nk = min(_F_MOV_MAX, K - ko0)
ps = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
for hp in nl.affine_range(nh2):
nisa.nc_matmul(dst=ps[0:prs, 0:nk], stationary=g8[hp][:, :, rr0:rr0 + prs],
moving=w2_8[hp][:, :, ko0:ko0 + nk], perf_mode=nisa.matmul_perf_mode.double_row)
# row += psum * out_scale (residual accumulated in fp32)
nisa.scalar_tensor_tensor(dst=row[0:prs, ko0:ko0 + nk], data=ps[0:prs, 0:nk], op0=nl.multiply,
operand0=out_scale, op1=nl.add, operand1=row[0:prs, ko0:ko0 + nk])
nisa.tensor_tensor(dst=row[0:prs, 0:K], data1=row[0:prs, 0:K], data2=b2_full[0:prs, 0:K], op=nl.add)
nisa.dma_copy(dst=xnew_hbm[r0:r0 + prs, 0:K], src=row[0:prs, 0:K])
nisa.tensor_scalar(dst=neg_shift[0:prs, 0:1], data=row[0:prs, 0:1], op0=nl.multiply, operand0=-1.0)
nisa.activation_reduce(dst=z_t[0:prs, 0:K], op=nl.copy, data=row[0:prs, 0:K],
reduce_op=nl.add, reduce_res=acc_s[0:prs, 0:1], bias=neg_shift[0:prs, 0:1], scale=1.0)
nisa.tensor_scalar(dst=neg_mu[0:prs, 0:1], data=acc_s[0:prs, 0:1], op0=nl.multiply,
operand0=neg_inv_k, op1=nl.add, operand1=neg_shift[0:prs, 0:1])
nisa.activation_reduce(dst=z_t[0:prs, 0:K], op=nl.square, data=row[0:prs, 0:K],
reduce_op=nl.add, reduce_res=acc_s[0:prs, 0:1], bias=neg_mu[0:prs, 0:1], scale=1.0)
nisa.tensor_scalar(dst=var[0:prs, 0:1], data=acc_s[0:prs, 0:1], op0=nl.multiply,
operand0=inv_k, op1=nl.add, operand1=eps)
nisa.activation(dst=rstd[0:prs, 0:1], op=nl.rsqrt, data=var[0:prs, 0:1])
nisa.tensor_scalar(dst=z_t[0:prs, 0:K], data=row[0:prs, 0:K], op0=nl.add, operand0=neg_mu[0:prs, 0:1],
op1=nl.multiply, operand1=rstd[0:prs, 0:1])
nisa.dma_copy(dst=z_hbm[r0:r0 + prs, 0:K], src=z_t[0:prs, 0:K])
continue
for ko in range(n_ko):
ko0 = ko * _F_MOV_MAX
nk = min(_F_MOV_MAX, K - ko0)
b2_bc = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=b2_bc[:, 0:nk], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, nk]], offset=ko0))
for rs in range(n_rs):
rr0 = rs * _P_MAX
prs = min(_P_MAX, pr - rr0)
ps = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
for hp in nl.affine_range(nh2):
nisa.nc_matmul(dst=ps[0:prs, 0:nk], stationary=g8[hp][:, :, rr0:rr0 + prs],
moving=w2_8[hp][:, :, ko0:ko0 + nk], perf_mode=nisa.matmul_perf_mode.double_row)
o = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
# out = psum * out_scale + b2, out_scale = 1 / (g_scale * sw2)
nisa.scalar_tensor_tensor(dst=o[0:prs, 0:nk], data=ps[0:prs, 0:nk], op0=nl.multiply,
operand0=out_scale, op1=nl.add, operand1=b2_bc[0:prs, 0:nk])
nisa.dma_copy(dst=out_hbm[r_abs + rr0:r_abs + rr0 + prs, ko0:ko0 + nk], src=o[0:prs, 0:nk])
if fuse:
return xnew_hbm, z_hbm
return out_hbm
@nki.jit
def _mlp_fp8_kernel(x_hbm, w1p_hbm, b1_hbm, w2p_hbm, b2_hbm, inv_sw1, g_scale, out_scale):
return _mlp_fp8_body(x_hbm, w1p_hbm, b1_hbm, w2p_hbm, b2_hbm, inv_sw1, g_scale, out_scale, None, 0.0, 0)
@nki.jit
def _mlp_fp8_ln_kernel(x_hbm, w1p_hbm, b1_hbm, w2p_hbm, b2_hbm, a_hbm, inv_sw1, g_scale, out_scale, eps):
return _mlp_fp8_body(x_hbm, w1p_hbm, b1_hbm, w2p_hbm, b2_hbm, inv_sw1, g_scale, out_scale, a_hbm, eps, 1)
_E4_TARGET = 128.0 # scale tensors so |max| lands at 128 (< legacy e4m3 max 240), power of two
def _pow2_scale(amax: float) -> float:
import math
return 2.0 ** math.floor(math.log2(_E4_TARGET / max(amax, 1e-12)))
def pack_fp8_weight(wt: torch.Tensor):
"""``[Kc, F]`` (contraction-major) -> (``[Kc/256, 128, 2, F]`` bf16 scaled, scale). Scale is a power of two."""
Kc, F = wt.shape
assert Kc % 256 == 0, f"contraction {Kc} must be a multiple of 256"
s = _pow2_scale(wt.abs().max().item())
return (wt * s).to(torch.bfloat16).reshape(Kc // 256, 2, 128, F).transpose(1, 2).contiguous(), s
@nki_op("vit_enc::nki_mlp_fp8", mutates_args={})
def nki_mlp_fp8(x: torch.Tensor, w1p: torch.Tensor, b1: torch.Tensor, w2p: torch.Tensor,
b2: torch.Tensor, sw1: float, sw2: float, g_scale: float) -> torch.Tensor:
"""EXPERIMENTAL FP8 GELU MLP. ``w1p, sw1`` / ``w2p, sw2`` from :func:`pack_fp8_weight`; ``g_scale`` = power-of-two
scale for the GELU output (choose so max|gelu| * g_scale <= 128)."""
return _mlp_fp8_kernel[LNC](x, w1p, b1, w2p, b2, 1.0 / sw1, g_scale, 1.0 / (g_scale * sw2))
@nki_op("vit_enc::nki_mlp_fp8_res_ln", mutates_args={})
def nki_mlp_fp8_res_ln(x: torch.Tensor, w1p: torch.Tensor, b1: torch.Tensor, w2p: torch.Tensor, b2: torch.Tensor,
a: torch.Tensor, eps: float, sw1: float, sw2: float, g_scale: float) -> tuple[torch.Tensor, torch.Tensor]:
"""EXPERIMENTAL FP8 GELU MLP + residual + next-layer affine-free LN. Returns (x_new fp32, z bf16)."""
return _mlp_fp8_ln_kernel[LNC](x, w1p, b1, w2p, b2, a, 1.0 / sw1, g_scale, 1.0 / (g_scale * sw2), eps)
|