v1.4: 5-bank QK^T PSUM rotation (<=1024 tokens) + experimental opt-in FP8 MLP

QK^T rotates over five PSUM banks across Q groups when a group has <=2 key tiles, so it no longer waits on
the Vector row-max evict: attention device cycles -1.1% to -4.7% vs v1.3; 2048-token layout unchanged.
FP8 e4m3 double_row GELU MLP behind VIT_ENC_MLP_FP8=1 (off by default) with weight and per-layer activation
scaling. Verified on trn2.3xlarge: all 24 DINOv3 cases and V-JEPA2 pass.

Files changed (6) hide show

CHANGELOG.md +20 -0
README.md +5 -0
build/torch-neuron/__init__.py +2 -2
build/torch-neuron/encoder.py +23 -1
build/torch-neuron/nki_kernels/attention_inner.py +13 -1
build/torch-neuron/nki_kernels/mlp_fp8.py +225 -0

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,25 @@
 # Changelog
 ## v1.3 (2026-10)
 Attention-path schedule improvements, plus a SwiGLU fallback fix. Output of every attention kernel is

 # Changelog
+## v1.4 (2026-10)
+* **QK^T PSUM banks:** for 1024 or fewer tokens per image (two key tiles or fewer), the QK^T matmul now rotates over
+  five PSUM banks across Q groups instead of reusing four, so it no longer waits for the Vector row-max evict of the
+  tile four back. Attention kernel device cycles vs v1.3: -2.1% (N=256 and N=1024, LNC=2), -1.1% (N=1024, LNC=1),
+  -4.7% (all-group kernel, N=256, LNC=1). At 2048 tokens the old layout is kept (the rotation collided with other banks
+  and measured +7%). `VIT_ENC_MM1_WIDE=0` restores the v1.3 layout.
+  End to end (single core, real HF weights) vs v1.3: DINOv3 ViT-L/16 -0.5% to +3.7%, ViT-H+/16 0.0% to +2.6%.
+  V-JEPA2 (2048 tokens): 23.80 ms/clip vs 23.81, output identical to v1.3. All 24 DINOv3 cases pass (rel_l2 vs fp32
+  0.56-1.75e-2).
+* **EXPERIMENTAL FP8 MLP** (`VIT_ENC_MLP_FP8=1`, GELU-MLP models only, off by default): fc1 and fc2 run as e4m3
+  `double_row` matmuls, with power-of-two weight scaling and a per-layer GELU-output scale (`calibrate_fp8`). Measured
+  on DINOv3 ViT-L/16, single core: +7.3% (224 px, LNC=2) to +10.5% (448 px, LNC=1) end to end. On imagenette
+  (real weights), k-NN and linear-probe predictions were identical to fp32. But the CLS feature error rises to
+  rel_l2 5.8e-2 (bf16 path: 0.84e-2). **Calibrate before use**: trn2's e4m3 saturates at 240, and DINOv3's first
+  block produces GELU outputs near 2,000. Without calibration the default scale (1/32) is used, which is safe for
+  DINOv3 but coarse.
+* Outputs are deterministic per compiled graph. Across separate compilations they can differ at bf16-ulp level
+  (rel ~1e-5), as in v1.3.
 ## v1.3 (2026-10)
 Attention-path schedule improvements, plus a SwiGLU fallback fix. Output of every attention kernel is

README.md CHANGED Viewed

@@ -128,6 +128,11 @@ Notes:
 * On a vLLM-Neuron venv (CUDA-tagged torch), `import libtorch_neuronx_lite` before `get_kernel`
   so the `torch-neuron` variant is selected. On PyTorch Native this is not needed.
 Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per
 instance. **LNC=2** (4 larger cores) gives the lowest latency per request.

 * On a vLLM-Neuron venv (CUDA-tagged torch), `import libtorch_neuronx_lite` before `get_kernel`
   so the `torch-neuron` variant is selected. On PyTorch Native this is not needed.
+**Experimental FP8 MLP** (v1.4, GELU-MLP models only): set `VIT_ENC_MLP_FP8=1` before importing to run the MLP
+matmuls in FP8 e4m3 (`double_row`). That is +7-10% end to end on DINOv3 ViT-L/16, at ~7x the feature error of the
+default bf16 path (CLS rel_l2 5.8e-2 vs 0.84e-2; imagenette k-NN predictions unchanged). Calibrate the per-layer
+activation scale with `calibrate_fp8(runner, gelu_amax)`, because trn2's e4m3 saturates at 240. Off by default.
 Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per
 instance. **LNC=2** (4 larger cores) gives the lowest latency per request.

build/torch-neuron/__init__.py CHANGED Viewed

@@ -10,7 +10,7 @@ Model-agnostic building blocks:
            mlp, mlp_res_ln, swiglu_stream, swiglu_stream_ln
 """
-from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
 from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder
 from .nki_kernels import (
     attention_block,
@@ -30,7 +30,7 @@ from .nki_kernels import (
 __all__ = [
     "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder",
-    "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC",
     "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
     "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",

            mlp, mlp_res_ln, swiglu_stream, swiglu_stream_ln
 """
+from .encoder import LNC, BlockWeights, EncoderRunner, calibrate_fp8, pack_block, rope_tables
 from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder
 from .nki_kernels import (
     attention_block,
 __all__ = [
     "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder",
+    "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8",
     "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
     "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",

build/torch-neuron/encoder.py CHANGED Viewed

@@ -36,10 +36,14 @@ from .nki_kernels.attention_block import LNC
 from .nki_kernels.mlp import nki_mlp_k4_xp
 from .nki_kernels.mlp import pack_b1 as _mlp_pack_b1
 from .nki_kernels.mlp_res_ln import nki_mlp_res_ln
 from .nki_kernels.res_layernorm import nki_res_layernorm_fold
 from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln
-__all__ = ["BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC"]
 _SBUF_BUDGET = 180 * 1024        # per-partition bytes the LNC=1 resident-weights decision allows
 _SW_MIN_ROWS = 256               # below this many rows per core, SwiGLU runs in torch (measured faster)
@@ -138,6 +142,10 @@ def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
         # Pre-pack fc1 bias to [128, H_mlp/128] so K4/K4+LN load it in one DMA (see mlp.pack_b1).
         p["b1"] = _mlp_pack_b1((_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous())
         p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16)
         p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous()
     return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()}
@@ -160,6 +168,15 @@ def rope_tables(cos: torch.Tensor, sin: torch.Tensor, n_prefix: int, n_pad: int,
     return _ab.pack_tables(C, S, nh_p, D, rope).to(device)
 class EncoderRunner:
     """Runs packed blocks on a stacked ``[B*N, H]`` fp32 residual stream (rows of image b at b*N)."""
@@ -203,6 +220,8 @@ class EncoderRunner:
             g = torch.nn.functional.silu((z2 @ p["s_w1t"]).float() + p["s_b1u"])
             u = (z2 @ p["s_wgt"]).float() + p["s_bgu"]
             return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0)
         return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0)
     def _ffn_fused(self, z2, p, x, eps_next):
@@ -211,6 +230,9 @@ class EncoderRunner:
                 return None
             return nki_swiglu_stream_ln(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
                                         p["s_w2t"], p["s_b2"], x, eps_next)
         return nki_mlp_res_ln(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], x, eps_next)
     def __call__(self, x: torch.Tensor, tab: torch.Tensor, n_valid: int, B: int) -> torch.Tensor:

 from .nki_kernels.mlp import nki_mlp_k4_xp
 from .nki_kernels.mlp import pack_b1 as _mlp_pack_b1
 from .nki_kernels.mlp_res_ln import nki_mlp_res_ln
+import os as _osx
+_MLP_FP8 = _osx.environ.get("VIT_ENC_MLP_FP8", "0") == "1"      # EXPERIMENTAL: FP8 e4m3 GELU MLP (see README)
+if _MLP_FP8:
+    _mf8 = _il.import_module(__package__ + ".nki_kernels.mlp_fp8")
 from .nki_kernels.res_layernorm import nki_res_layernorm_fold
 from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln
+__all__ = ["BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8"]
 _SBUF_BUDGET = 180 * 1024        # per-partition bytes the LNC=1 resident-weights decision allows
 _SW_MIN_ROWS = 256               # below this many rows per core, SwiGLU runs in torch (measured faster)
         # Pre-pack fc1 bias to [128, H_mlp/128] so K4/K4+LN load it in one DMA (see mlp.pack_b1).
         p["b1"] = _mlp_pack_b1((_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous())
         p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16)
+        if _MLP_FP8:
+            p["w1p8"], p["sw1"] = _mf8.pack_fp8_weight((w1 * g2[None]).t().contiguous())
+            p["w2p8"], p["sw2"] = _mf8.pack_fp8_weight((w2 * ls2[:, None]).t().contiguous())
+            p["g8"] = 1.0 / 32           # GELU-output scale; per-layer calibration recommended (calibrate_fp8)
         p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous()
     return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()}
     return _ab.pack_tables(C, S, nh_p, D, rope).to(device)
+def calibrate_fp8(runner, gelu_amax, headroom: float = 2.0):
+    """EXPERIMENTAL (VIT_ENC_MLP_FP8=1). Set each layer's power-of-two GELU-output scale from the observed max |GELU
+    output| per layer (list, one per block): scale = 2^floor(log2(128 / (amax * headroom))). trn2's e4m3 saturates at
+    240; uncalibrated activations above that turn into inf/NaN."""
+    import math
+    for p, a in zip(runner.L, gelu_amax):
+        p["g8"] = 2.0 ** math.floor(math.log2(128.0 / max(a * headroom, 1e-6)))
 class EncoderRunner:
     """Runs packed blocks on a stacked ``[B*N, H]`` fp32 residual stream (rows of image b at b*N)."""
             g = torch.nn.functional.silu((z2 @ p["s_w1t"]).float() + p["s_b1u"])
             u = (z2 @ p["s_wgt"]).float() + p["s_bgu"]
             return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0)
+        if _MLP_FP8:
+            return _mf8.nki_mlp_fp8(z2, p["w1p8"], p["b1"], p["w2p8"], p["b2"], p["sw1"], p["sw2"], p["g8"]).unsqueeze(0)
         return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0)
     def _ffn_fused(self, z2, p, x, eps_next):
                 return None
             return nki_swiglu_stream_ln(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
                                         p["s_w2t"], p["s_b2"], x, eps_next)
+        if _MLP_FP8:
+            return _mf8.nki_mlp_fp8_res_ln(z2, p["w1p8"], p["b1"], p["w2p8"], p["b2"], x, eps_next,
+                                           p["sw1"], p["sw2"], p["g8"])
         return nki_mlp_res_ln(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], x, eps_next)
     def __call__(self, x: torch.Tensor, tab: torch.Tensor, n_valid: int, B: int) -> torch.Tensor:

build/torch-neuron/nki_kernels/attention_inner.py CHANGED Viewed

@@ -109,6 +109,11 @@ _MAX_GLOBAL_CP_DEGREE = 32
 # Tile size / threshold constants used by the LIVE path.
 _Q_GRP_SZ = 128
 _V_TILE_SZ = 128  # V tile size for loading and MM2 operations
 _K_TILE_SZ = 512  # K tile size for loading and MM1+masking operations
 _EXP_TILE_SZ = 512  # Tile size for exp instructions (must equal _K_TILE_SZ)
 _LARGE_TILE_SZ = 2048  # Larger tile size for allocations/pipelining (4 x 512 tiles)
@@ -694,7 +699,14 @@ def _allocate_attention_buffers(
             for k_tile_idx in range(4):
                 mm1_psum_tile = nl.ndarray(
                     (atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype, buffer=nl.psum,
-                    address=(0, (k_tile_idx % 4) * PSUM_BANK_SIZE),
                 )
                 tile_row.append(mm1_psum_tile)
             grp_row.append(tile_row)

 # Tile size / threshold constants used by the LIVE path.
 _Q_GRP_SZ = 128
 _V_TILE_SZ = 128  # V tile size for loading and MM2 operations
+import os as _os
+_MM1_WIDE = _os.environ.get("VIT_ENC_MM1_WIDE", "1") == "1"      # set 0 to restore the v1.3 4-bank MM1 layout
+_MM1_BANKS = (0, 1, 2, 5, 6)   # MM1 rotation; bank 4 = P.V (single large tile), 3/7 = reciprocal transposes
+_MM1_WIDE_MAX_KTILES = 2       # only widen when <= this many key tiles per group: at 4 key tiles (N=2048) the
+                               # 5-bank rotation collides with P.V/producer banks and measured +7% slower
 _K_TILE_SZ = 512  # K tile size for loading and MM1+masking operations
 _EXP_TILE_SZ = 512  # Tile size for exp instructions (must equal _K_TILE_SZ)
 _LARGE_TILE_SZ = 2048  # Larger tile size for allocations/pipelining (4 x 512 tiles)
             for k_tile_idx in range(4):
                 mm1_psum_tile = nl.ndarray(
                     (atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype, buffer=nl.psum,
+                    # v1.4: rotate MM1 over 5 banks (0,1,2,5,6) across consecutive Q groups instead
+                    # of reusing banks 0-3 per group, so a QK matmul does not wait for the Vector row-max evict
+                    # of the tile 4 back. Bank 4 = P.V, 3/7 = reciprocal transposes. Banks 5/6 are also the
+                    # producer's projection banks; the tile framework orders reuse by address, so overlap
+                    # with a later image's producer stays correct (it may serialise).
+                    address=(0, (_MM1_BANKS[(grp_idx * 4 + k_tile_idx) % len(_MM1_BANKS)]
+                                 if (_MM1_WIDE and atp.num_k_tiles_per_section <= _MM1_WIDE_MAX_KTILES)
+                                 else (k_tile_idx % 4)) * PSUM_BANK_SIZE),
                 )
                 tile_row.append(mm1_psum_tile)
             grp_row.append(tile_row)

build/torch-neuron/nki_kernels/mlp_fp8.py ADDED Viewed

	@@ -0,0 +1,225 @@

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