Hugging Face's logo

Kernels:
aws-neuron
/
vit-encoder-neuron-kernels

Kernel card Files
xet
 Community
vit-encoder-neuron-kernels / build /torch-neuron /nki_kernels /mlp_fp8.py
jburtoft's picture
jburtoft
v1.4: 5-bank QK^T PSUM rotation (<=1024 tokens) + experimental opt-in FP8 MLP
3b1bffa
Raw History Blame
12.3 kB
# 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)