# 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)