v1.8: post-LayerNorm encoders (NeuronXLMRobertaModel; post-LN epilogue outputs + no-RoPE switch) and LNC=2 fix for odd head-group counts

Browse files

Files changed (8) hide show

CHANGELOG.md +35 -0
README.md +29 -4
build/torch-neuron/__init__.py +9 -6
build/torch-neuron/encoder.py +7 -0
build/torch-neuron/layers.py +128 -0
build/torch-neuron/nki_kernels/__init__.py +4 -3
build/torch-neuron/nki_kernels/attention_block_l1_ln.py +70 -23
build/torch-neuron/nki_kernels/mlp_res_ln.py +33 -2

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,40 @@
 # Changelog
 ## v1.7 (2026-10)
 * **Fix: `res_layernorm_fold` failed to compile at large batch** (`NCC_IBIR229`, state buffer allocation failed). The

 # Changelog
+## v1.8 (2026-10)
+* **New: post-LayerNorm encoders (BERT / RoBERTa / XLM-R).** `NeuronXLMRobertaModel` is a drop-in forward for
+  `XLMRobertaModel`, `RobertaModel` and `BertModel` (GELU, absolute positions, no RoPE). It runs on the same kernels via
+  two new options:
+  - `attention_block_l1_ln_post` / `mlp_res_ln_post`: the existing fused attention and MLP epilogues also emit the
+    post-LN residual `LN(x_new) * gamma + beta` in fp32, computed from the mean and rstd they already have.
+  - `norope=1` on the attention kernel skips the rotation (and its table loads) for models without RoPE.
+  XLM-RoBERTa-Large, LNC=1, 8 processes (one per logical core) on one trn2.3xlarge, fully padded batches (inferences/s):
+  | seq | stock `torch.compile` | v1.8 `NeuronXLMRobertaModel` (`attention_mask=None`) | speedup | rel_l2 vs fp32 (stock / kernels) |
+  |---|---|---|---|---|
+  | 128 | 2,092.7 | **3,144.8** | 1.50x | 0.031 / 0.0066 |
+  | 256 | 944.4 (BS 16) | **1,869.5** | 1.98x | 0.051 / 0.0147 |
+  | 512 | 397.1 (BS 8) | **935.1** | 2.35x | 0.018 / 0.0143 |
+  Batch per core: 16 / 8 / 4. Passing an all-ones `attention_mask` costs 3-4% (a host-side check that no position is
+  masked: 3,022.7 / 1,789.7 / 902.8). A padded mask falls back to the stock forward (checked: rel_l2 vs fp32 1.2e-6).
+  Measured on the SDK 2.33 DLAMI (`torch_neuronx_2_14_0` venv: torch 2.14.0, torch-neuronx 2.14.3.0, neuronx-cc 2.28.7631,
+  NKI 0.7.0) with transformers 4.57.6. The kernel path needs LNC=1, a sequence length that is a multiple of 128 (up to 2048)
+  and no padding mask (all positions attended). Anything else falls back to the stock forward.
+* **Fix: models with an odd number of head groups failed to compile at LNC=2 on neuronx-cc 2.28** (`NCC_EVRF066`: a
+  one-program attention grid on an LNC=2 build). That covers ViT-S/16, ViT-S+/16 and ViT-H+/16. At LNC=2 the host packing now pads
+  an odd group count with one all-zero head group, so every attention launch uses a two-program grid. The padded heads have
+  zero V and out-projection weights and add exactly 0. Models with an even group count (ViT-B, ViT-L, V-JEPA2) are unchanged.
+  Single core, LNC=2: ViT-S/16 224 BS4 1,148.8 img/s (rel_l2 0.0109), ViT-S+/16 1,097.9 (0.0148), ViT-H+/16 224 BS1/BS4
+  87.8 / 130.9, 448 BS1/BS2 32.6 / 33.9 (rel_l2 0.0059-0.0080). The ViT-H+ LNC=2 path is 1.18-1.34x faster than its LNC=1
+  fallback on the same core.
+* **No change for existing models at existing shapes:** in the 28-case regression (DINOv3 S/S+/B/L/H+ at 224/448 and
+  V-JEPA2 at batch 1/2, LNC=1 and LNC=2), every case v1.7 runs gives output bitwise identical to v1.7, at the same speed
+  (within ±1%). That is 22 cases. The other 6 are the odd-group LNC=2 cases above, which v1.7 cannot compile on this
+  toolchain. Verified on the SDK 2.33 DLAMI.
 ## v1.7 (2026-10)
 * **Fix: `res_layernorm_fold` failed to compile at large batch** (`NCC_IBIR229`, state buffer allocation failed). The

README.md CHANGED Viewed

@@ -13,6 +13,8 @@ tags:
   - vjepa2
   - rope
   - swiglu
 ---
 # vit-encoder-neuron-kernels
@@ -22,15 +24,16 @@ Trainium2 (trn2). Each block is 2-3 coarse kernels instead of dozens of small op
 the residual add and LayerNorm, and the MLP / SwiGLU FFN plus the next block's LayerNorm. The
 residual stream stays in fp32 between kernels.
-Drop-in layers are included for two Hugging Face Transformers models:
 | Layer | Replaces | Models |
 |---|---|---|
 | `NeuronDINOv3ViTModel` | `DINOv3ViTModel.forward` | DINOv3 ViT-S/16, S+/16, B/16, L/16, H+/16 |
 | `NeuronVJEPA2Encoder` | `VJEPA2Encoder.forward` | V-JEPA2 ViT-L (video), up to 2048 tokens per clip |
 The building blocks (`BlockWeights` / `pack_block` / `EncoderRunner`) are model-agnostic. Any
-pre-norm ViT with the shape constraints below can use them, but only the two models above
 have been validated.
 ## Performance (trn2.3xlarge, bf16, `torch.compile(backend="neuron")`)
@@ -141,6 +144,25 @@ activation scale with `calibrate_fp8(runner, gelu_amax)`, because trn2's e4m3 sa
 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.
 ## Kernels
 | Kernel | What it does |
@@ -173,13 +195,16 @@ RoPE pairings via `rope=`:
 * Inference only (no autograd). The layers refuse in training mode.
 * ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism.
   The streaming SwiGLU kernel is slower than the compiler's FFN at that width.
-* Validated end to end on DINOv3 ViT-B/L/H+ and V-JEPA2 ViT-L only.
 ## Requirements
 * trn2 (verified on trn2.3xlarge, LNC=1 and LNC=2).
 * PyTorch Native (TorchNeuron) with NKI 0.7 and `nkilib`. Verified on Beta 6 (torch 2.13.0,
-  torch-neuronx 2.13.3, NKI 0.7.0b1, neuronx-cc 2.0.404056).
 * Transformers with DINOv3 / V-JEPA2 (verified on 4.57.6); `kernels` for `get_kernel`.
 ## Repository layout

   - vjepa2
   - rope
   - swiglu
+  - bert
+  - xlm-roberta
 ---
 # vit-encoder-neuron-kernels
 the residual add and LayerNorm, and the MLP / SwiGLU FFN plus the next block's LayerNorm. The
 residual stream stays in fp32 between kernels.
+Drop-in layers are included for these Hugging Face Transformers models:
 | Layer | Replaces | Models |
 |---|---|---|
 | `NeuronDINOv3ViTModel` | `DINOv3ViTModel.forward` | DINOv3 ViT-S/16, S+/16, B/16, L/16, H+/16 |
 | `NeuronVJEPA2Encoder` | `VJEPA2Encoder.forward` | V-JEPA2 ViT-L (video), up to 2048 tokens per clip |
+| `NeuronXLMRobertaModel` (v1.8) | `XLMRobertaModel.forward` (also `RobertaModel` / `BertModel`) | XLM-RoBERTa-Large (validated); post-LN text encoders without RoPE |
 The building blocks (`BlockWeights` / `pack_block` / `EncoderRunner`) are model-agnostic. Any
+pre-norm ViT with the shape constraints below can use them, but only the models above
 have been validated.
 ## Performance (trn2.3xlarge, bf16, `torch.compile(backend="neuron")`)
 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.
+### XLM-RoBERTa / BERT-family text encoders (v1.8)
+```python
+import os
+os.environ.setdefault("NEURON_LOGICAL_NC_CONFIG", "1")         # the post-LN path runs at LNC=1
+import types, torch
+from transformers import XLMRobertaModel
+model = XLMRobertaModel.from_pretrained("FacebookAI/xlm-roberta-large", dtype=torch.float32,
+                                        add_pooling_layer=False).eval().to("neuron")
+model.forward = types.MethodType(k.NeuronXLMRobertaModel.forward, model)   # instance graft (k = get_kernel(...) as above)
+model = torch.compile(model, backend="neuron")
+ids = tokenizer(texts, padding="max_length", max_length=128, return_tensors="pt").input_ids   # seq % 128 == 0
+out = model(input_ids=ids.to("neuron")).last_hidden_state
+```
+The kernel path needs LNC=1, a sequence length that is a multiple of 128 (up to 2048) and no padding mask. Pass
+`attention_mask=None` or all ones; masked batches fall back to the stock forward. Per instance, XLM-RoBERTa-Large at
+seq 128 / 256 / 512 runs at 3,145 / 1,870 / 935 inferences/s, against 2,093 / 944 / 397 stock (see CHANGELOG v1.8).
 ## Kernels
 | Kernel | What it does |
 * Inference only (no autograd). The layers refuse in training mode.
 * ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism.
   The streaming SwiGLU kernel is slower than the compiler's FFN at that width.
+* Validated end to end on DINOv3 ViT-B/L/H+, V-JEPA2 ViT-L and (v1.8) XLM-RoBERTa-Large.
+* Post-LN text encoders (v1.8, `NeuronXLMRobertaModel`): LNC=1 only, no padding mask, no token-type / relative
+  position variants beyond the standard BERT embeddings, GELU FFN.
 ## Requirements
 * trn2 (verified on trn2.3xlarge, LNC=1 and LNC=2).
 * PyTorch Native (TorchNeuron) with NKI 0.7 and `nkilib`. Verified on Beta 6 (torch 2.13.0,
+  torch-neuronx 2.13.3, NKI 0.7.0b1, neuronx-cc 2.0.404056) and (v1.8) on the **SDK 2.33 DLAMI**
+  `torch_neuronx_2_14_0` venv (torch 2.14.0, torch-neuronx 2.14.3.0, NKI 0.7.0, neuronx-cc 2.28.7631).
 * Transformers with DINOv3 / V-JEPA2 (verified on 4.57.6); `kernels` for `get_kernel`.
 ## Repository layout

build/torch-neuron/__init__.py CHANGED Viewed

@@ -3,22 +3,25 @@
 Model layers (drop-in forwards for Hugging Face Transformers):
   NeuronDINOv3ViTModel   -- DINOv3ViTModel (S/S+/B/L/H+)
   NeuronVJEPA2Encoder    -- VJEPA2Encoder
 Model-agnostic building blocks:
   BlockWeights, pack_block, rope_tables, EncoderRunner   -- pre-norm encoder loop on the kernels
-  kernels: attention_block, attention_block_l1, attention_block_l1_ln, res_layernorm_fold,
-           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,
     attention_block_l1,
     attention_block_l1_ln,
     lane_permutation,
     mlp,
     mlp_res_ln,
     pack_out_weights,
     pack_tables,
     pack_v_weights,
@@ -29,9 +32,9 @@ from .nki_kernels import (
 )
 __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",
 ]

 Model layers (drop-in forwards for Hugging Face Transformers):
   NeuronDINOv3ViTModel   -- DINOv3ViTModel (S/S+/B/L/H+)
   NeuronVJEPA2Encoder    -- VJEPA2Encoder
+  NeuronXLMRobertaModel  -- XLMRobertaModel / RobertaModel / BertModel (post-LN, no RoPE; v1.8)
 Model-agnostic building blocks:
   BlockWeights, pack_block, rope_tables, EncoderRunner   -- pre-norm encoder loop on the kernels
+  kernels: attention_block, attention_block_l1, attention_block_l1_ln(_post), res_layernorm_fold,
+           mlp, mlp_res_ln(_post), swiglu_stream, swiglu_stream_ln
 """
 from .encoder import LNC, BlockWeights, EncoderRunner, calibrate_fp8, pack_block, rope_tables
+from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder, NeuronXLMRobertaModel
 from .nki_kernels import (
     attention_block,
     attention_block_l1,
     attention_block_l1_ln,
+    attention_block_l1_ln_post,
     lane_permutation,
     mlp,
     mlp_res_ln,
+    mlp_res_ln_post,
     pack_out_weights,
     pack_tables,
     pack_v_weights,
 )
 __all__ = [
+    "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder", "NeuronXLMRobertaModel",
     "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8",
+    "attention_block", "attention_block_l1", "attention_block_l1_ln", "attention_block_l1_ln_post", "res_layernorm_fold",
+    "mlp", "mlp_res_ln", "mlp_res_ln_post", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
 ]

build/torch-neuron/encoder.py CHANGED Viewed

@@ -105,6 +105,11 @@ def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
     D = F_ // num_heads
     gsz = 512 // D                                   # heads per kernel program
     nh_p = ((num_heads + gsz - 1) // gsz) * gsz      # padded head count
     Fp = nh_p * D
     ngrp = nh_p // gsz
     g1, t1, g2, t2 = _f(bw.ln1_w), _f(bw.ln1_b), _f(bw.ln2_w), _f(bw.ln2_b)
@@ -167,6 +172,8 @@ def rope_tables(cos: torch.Tensor, sin: torch.Tensor, n_prefix: int, n_pad: int,
     S[n_prefix:n_prefix + n_rot] = sin.float().cpu()
     gsz = 512 // D
     nh_p = ((num_heads + gsz - 1) // gsz) * gsz
     t = _ab.pack_tables(C, S, nh_p, D, rope)
     if _osr.environ.get("VIT_ENC_ROPE_BF16", "0") == "1":     # EXPERIMENTAL bf16 RoPE (tables in bf16)
         t = t.to(torch.bfloat16)

     D = F_ // num_heads
     gsz = 512 // D                                   # heads per kernel program
     nh_p = ((num_heads + gsz - 1) // gsz) * gsz      # padded head count
+    if LNC == 2 and (nh_p // gsz) % 2 == 1:
+        # v1.8: at LNC=2 every attention launch must use a 2-program grid (newer neuronx-cc rejects a 1-program
+        # grid on an LNC=2 build: NCC_EVRF066). Pad an odd group count with one all-zero head group; its V and
+        # out-projection weights are zero, so it adds exactly 0 to the output.
+        nh_p += gsz
     Fp = nh_p * D
     ngrp = nh_p // gsz
     g1, t1, g2, t2 = _f(bw.ln1_w), _f(bw.ln1_b), _f(bw.ln2_w), _f(bw.ln2_b)
     S[n_prefix:n_prefix + n_rot] = sin.float().cpu()
     gsz = 512 // D
     nh_p = ((num_heads + gsz - 1) // gsz) * gsz
+    if LNC == 2 and (nh_p // gsz) % 2 == 1:
+        nh_p += gsz                                  # same padding as pack_block (v1.8)
     t = _ab.pack_tables(C, S, nh_p, D, rope)
     if _osr.environ.get("VIT_ENC_ROPE_BF16", "0") == "1":     # EXPERIMENTAL bf16 RoPE (tables in bf16)
         t = t.to(torch.bfloat16)

build/torch-neuron/layers.py CHANGED Viewed

@@ -218,3 +218,131 @@ class NeuronVJEPA2Encoder(torch.nn.Module):
         x = getattr(self, _PACK)(_stack(hs, N), tabs[N], n, B)
         out = self.layernorm(x.reshape(B, N, H)[:, :n])
         return BaseModelOutput(last_hidden_state=out)

         x = getattr(self, _PACK)(_stack(hs, N), tabs[N], n, B)
         out = self.layernorm(x.reshape(B, N, H)[:, :n])
         return BaseModelOutput(last_hidden_state=out)
+# ------------------------------------------------------------------------------------------------
+# XLM-RoBERTa / RoBERTa / BERT (post-LayerNorm encoders, no RoPE)  -- v1.8
+# ------------------------------------------------------------------------------------------------
+_POST = "_vit_enc_neuron_postln"
+def _bert_blocks(model):
+    """BlockWeights per layer for a post-LN BERT-family encoder. Block i's "ln1" is the LayerNorm that produced
+    its input (embeddings.LayerNorm for block 0, else layer i-1's output.LayerNorm) and "ln2" its attention-output
+    LayerNorm, so the kernels' affine-free LN + folded affine reproduce the post-LN math exactly."""
+    out, prev = [], model.embeddings.LayerNorm
+    for lyr in model.encoder.layer:
+        a = lyr.attention
+        out.append(BlockWeights(
+            ln1_w=prev.weight, ln1_b=prev.bias, ln1_eps=prev.eps,
+            wq=a.self.query.weight, bq=a.self.query.bias, wk=a.self.key.weight, bk=a.self.key.bias,
+            wv=a.self.value.weight, bv=a.self.value.bias, wo=a.output.dense.weight, bo=a.output.dense.bias,
+            ln2_w=a.output.LayerNorm.weight, ln2_b=a.output.LayerNorm.bias, ln2_eps=a.output.LayerNorm.eps,
+            fc1_w=lyr.intermediate.dense.weight, fc1_b=lyr.intermediate.dense.bias,
+            fc2_w=lyr.output.dense.weight, fc2_b=lyr.output.dense.bias))
+        prev = lyr.output.LayerNorm
+    return out
+def _bert_ok(model, S):
+    c = model.config
+    return (_supported(c.hidden_size, c.num_attention_heads) and getattr(c, "hidden_act", "gelu") == "gelu"
+            and getattr(c, "position_embedding_type", "absolute") == "absolute" and S <= _MAX_TOKENS
+            and S % (128 * LNC) == 0 and LNC == 1)
+@torch._dynamo.disable
+def _bert_freeze(model):
+    if model.training:
+        raise RuntimeError("NeuronXLMRobertaModel: these kernels have no autograd rule; call model.eval() first.")
+    c = model.config
+    dev = model.embeddings.word_embeddings.weight.device
+    heads, H = c.num_attention_heads, c.hidden_size
+    D = H // heads
+    runner = EncoderRunner([pack_block(bw, heads, "half", dev) for bw in _bert_blocks(model)])
+    f32 = lambda t: t.detach().float().to(dev)
+    lns = [model.embeddings.LayerNorm] + [l.output.LayerNorm for l in model.encoder.layer]
+    st = {"runner": runner, "D": D, "heads": heads, "H": H, "tabs": {},
+          "ln_in": [(f32(m.weight), f32(m.bias), float(m.eps)) for m in lns],
+          "ln_attn": [(f32(l.attention.output.LayerNorm.weight), f32(l.attention.output.LayerNorm.bias),
+                       float(l.attention.output.LayerNorm.eps)) for l in model.encoder.layer]}
+    setattr(model, _POST, st)
+@torch._dynamo.disable
+def _bert_tab(st, S, dev):
+    if S not in st["tabs"]:
+        D = st["D"]                                   # no RoPE: identity tables (the kernel skips them, norope=1)
+        st["tabs"][S] = rope_tables(torch.ones(S, D), torch.zeros(S, D), 0, S, st["heads"], D, "half", dev)
+    return st["tabs"][S]
+@torch._dynamo.disable
+def _bert_masked(attention_mask) -> bool:
+    """True if any position is masked out. Runs outside the compiled graph (one small host check per call)."""
+    return not bool(torch.all(attention_mask == 1))
+class NeuronXLMRobertaModel(torch.nn.Module):
+    """Drop-in forward for ``XLMRobertaModel`` / ``RobertaModel`` / ``BertModel`` (post-LN, GELU, absolute positions).
+    Graft on an instance (``model.forward = types.MethodType(NeuronXLMRobertaModel.forward, model)``); do not use it
+    as a class-level KernelConfig replacement, because the fallback calls ``type(model).forward``.
+    Returns ``last_hidden_state``. Kernel path: LNC=1, sequence length a multiple of 128 and <= 2048, all positions
+    attended (``attention_mask`` all ones / no padding). Anything else (padding mask, LNC=2, other lengths,
+    off-device, training) falls back to the stock forward.
+    """
+    can_torch_compile = True
+    has_backward = False
+    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, **kwargs):
+        import torch.nn.functional as F
+        from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions as Out
+        from .encoder import _ab1ln
+        from .nki_kernels.mlp import nki_mlp_k4_xp
+        from .nki_kernels.mlp_res_ln import nki_mlp_res_ln_post
+        from .nki_kernels.res_layernorm import nki_res_layernorm_fold
+        S = input_ids.shape[1] if input_ids is not None else -1
+        # The mask check runs outside the compiled graph (one host sync per call); pass attention_mask=None to skip it.
+        masked = attention_mask is not None and _bert_masked(attention_mask)
+        if (input_ids is None or input_ids.device.type != "neuron" or self.training or masked or position_ids is not None
+                or kwargs.get("inputs_embeds") is not None or not _bert_ok(self, S)):
+            # stock forward of the grafted instance's own class (XLMRobertaModel / RobertaModel / BertModel). The graft
+            # sets an INSTANCE attribute (types.MethodType), so the class attribute is still the stock forward.
+            return type(self).forward(self, input_ids=input_ids, attention_mask=attention_mask,
+                                      token_type_ids=token_type_ids, position_ids=position_ids, **kwargs)
+        if not hasattr(self, _POST):
+            _bert_freeze(self)
+        st = getattr(self, _POST)
+        R, H, D = st["runner"], st["H"], st["D"]
+        B = input_ids.shape[0]
+        # the kernels need the PRE-LayerNorm embedding sum (the embedding LN is block 0's "ln1"); same as embeddings()
+        e = self.embeddings
+        pos = e.position_ids[:, :S] if getattr(e, "padding_idx", None) is None else None
+        if pos is None:
+            m = input_ids.ne(e.padding_idx).int()
+            pos = (torch.cumsum(m, dim=1) * m).long() + e.padding_idx
+        tt = token_type_ids if token_type_ids is not None else torch.zeros_like(input_ids)
+        x = (e.word_embeddings(input_ids) + e.position_embeddings(pos) + e.token_type_embeddings(tt)).float()
+        x = x.reshape(B * S, H).contiguous()
+        tab = _bert_tab(st, S, x.device)
+        g0, b0, eps0 = st["ln_in"][0]
+        _, z1 = nki_res_layernorm_fold(x, x.new_zeros(1, x.shape[0], H, dtype=torch.bfloat16), eps0)
+        h = F.layer_norm(x, (H,), g0, b0, eps0)
+        nL = len(R.L)
+        from .encoder import _PACK2
+        pk2 = int(_PACK2 and B % 2 == 0 and 2 * S <= 512)
+        for li, p in enumerate(R.L):
+            ga, ba, _ = st["ln_attn"][li]
+            _, z2, h1 = _ab1ln.attention_block_l1_ln_post(
+                z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S, p["eps2"],
+                ga, ba, B, R._res(_ab1ln, S), pk2, 1)
+            go, bo, eo = st["ln_in"][li + 1]
+            if li + 1 < nL:
+                _, z1, h = nki_mlp_res_ln_post(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], h1, eo, go, bo)
+            else:
+                h = F.layer_norm(h1 + nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).float(), (H,), go, bo, eo)
+        out = h.reshape(B, S, H)
+        return Out(last_hidden_state=out, pooler_output=None)

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

@@ -10,16 +10,17 @@ from .attention_block import (
     pack_weights,
 )
 from .attention_block_l1 import attention_block_l1
-from .attention_block_l1_ln import attention_block_l1_ln
 from .mlp import nki_mlp_k4_xp as mlp
 from .mlp import pack_b1
 from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
 from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
 from .swiglu_stream import nki_swiglu_stream as swiglu_stream
 from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
 __all__ = [
-    "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
-    "mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
 ]

     pack_weights,
 )
 from .attention_block_l1 import attention_block_l1
+from .attention_block_l1_ln import attention_block_l1_ln, attention_block_l1_ln_post
 from .mlp import nki_mlp_k4_xp as mlp
 from .mlp import pack_b1
 from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
+from .mlp_res_ln import nki_mlp_res_ln_post as mlp_res_ln_post
 from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
 from .swiglu_stream import nki_swiglu_stream as swiglu_stream
 from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
 __all__ = [
+    "attention_block", "attention_block_l1", "attention_block_l1_ln", "attention_block_l1_ln_post", "res_layernorm_fold",
+    "mlp", "mlp_res_ln", "mlp_res_ln_post", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
 ]

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

@@ -87,7 +87,8 @@ def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int:
 @nki.jit
 def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
-                                  wot_hbm, bo_hbm, x_hbm, scale, head_dim, n_valid, eps, nimg=1, resident=0, pack2=0):
     NT, H = h_hbm.shape
     N = NT // nimg
     Fp = wqk_hbm.shape[2]
@@ -117,6 +118,9 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
     _kernel_assert(NT == nimg * N, "rows must be nimg x N")
     _kernel_assert(x_hbm.shape == (NT, H), "x must be [nimg*N, H] fp32")
     xnew_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
     z_hbm = nl.ndarray((NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
     inv_h = 1.0 / float(H)
@@ -158,13 +162,15 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
             r_wv.append(wv_sb_)
             r_b.append(b_sb_)
             r_bv.append(bv_bc_)
-        r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
-                                     block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
-        for s in range(n_tt):
-            for t_i in range(4):
-                _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
-                nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw],
-                              src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
         r_wo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
                                     block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
         for oc in range(G * n_ochunks):
@@ -208,7 +214,9 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
                 )
         # rope tables: shared by every group.
-        if resident and PK == 1:
             tab = r_tab
         elif PK == 2:
             # pack2: the same per-image tables, side by side (one 512-wide tile)
@@ -316,20 +324,26 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
                                 nisa.nc_matmul(dst=ps[:, 0:tw],
                                                stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
                                                moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
-                            nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
-                                               op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
-                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
-                                           data2=tab[0][tt][:, 0:tw], op=nl.multiply)
-                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
-                                           data2=tab[2][tt][:, 0:tw], op=nl.multiply)
-                        nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw],
-                                           data2=tmp[:, 0:tw], op=nl.add)
-                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
-                                           data2=tab[1][tt][:, 0:tw], op=nl.multiply)
-                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
-                                           data2=tab[3][tt][:, 0:tw], op=nl.multiply)
-                        nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw],
-                                           data2=tmp[:, 0:tw], op=nl.add)
                         for b in range(quad):
                             lh = j * quad + b
                             r0 = b * half
@@ -413,6 +427,11 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
             st.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.float32,
                                                   block_dim=[5], num_free_tiles=[5], align_to=4))
         sink = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
         n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
         for tb in range(n_grp):
             k = tb % _RING
@@ -455,10 +474,21 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
             nisa.tensor_scalar(dst=z_t[k], data=row[k], op0=nl.add, operand0=neg_mu,
                                op1=nl.multiply, operand1=rstd)
             nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_t[k])
     return xnew_hbm, z_hbm
 _ATTN_L1_LN = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 @nki_op("vit_enc::attention_block_l1_ln", mutates_args={})
@@ -489,3 +519,20 @@ def attention_block_l1_ln(
     if bqk.dim() == 3 and bqk.shape[1] != 128:
         bqk = _pack_qk_bias(bqk)
     return _ATTN_L1_LN(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, pack2)

 @nki.jit
 def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
+                                  wot_hbm, bo_hbm, x_hbm, scale, head_dim, n_valid, eps, nimg=1, resident=0, pack2=0,
+                                  g_hbm=None, bt_hbm=None, norope=0):
     NT, H = h_hbm.shape
     N = NT // nimg
     Fp = wqk_hbm.shape[2]
     _kernel_assert(NT == nimg * N, "rows must be nimg x N")
     _kernel_assert(x_hbm.shape == (NT, H), "x must be [nimg*N, H] fp32")
+    post = g_hbm is not None                       # post-LN mode: also emit h = LN(x_new) * g + bt (fp32)
+    if post:
+        hout_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
     xnew_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
     z_hbm = nl.ndarray((NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
     inv_h = 1.0 / float(H)
             r_wv.append(wv_sb_)
             r_b.append(b_sb_)
             r_bv.append(bv_bc_)
+        r_tab = None
+        if not norope:
+            r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
+                                         block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+            for s in range(n_tt):
+                for t_i in range(4):
+                    _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                    nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw],
+                                  src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
         r_wo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
                                     block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
         for oc in range(G * n_ochunks):
                 )
         # rope tables: shared by every group.
+        if norope:
+            tab = None
+        elif resident and PK == 1:
             tab = r_tab
         elif PK == 2:
             # pack2: the same per-image tables, side by side (one 512-wide tile)
                                 nisa.nc_matmul(dst=ps[:, 0:tw],
                                                stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
                                                moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
+                            if norope:
+                                # no positional rotation (e.g. BERT/XLM-R): bias-add straight into the roped buffers
+                                nisa.tensor_scalar(dst=roped[par][:, 0:tw], data=ps[:, 0:tw],
+                                                   op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
+                            else:
+                                nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
+                                                   op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
+                        if not norope:
+                            nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
+                                               data2=tab[0][tt][:, 0:tw], op=nl.multiply)
+                            nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
+                                               data2=tab[2][tt][:, 0:tw], op=nl.multiply)
+                            nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw],
+                                               data2=tmp[:, 0:tw], op=nl.add)
+                            nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
+                                               data2=tab[1][tt][:, 0:tw], op=nl.multiply)
+                            nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
+                                               data2=tab[3][tt][:, 0:tw], op=nl.multiply)
+                            nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw],
+                                               data2=tmp[:, 0:tw], op=nl.add)
                         for b in range(quad):
                             lh = j * quad + b
                             r0 = b * half
             st.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.float32,
                                                   block_dim=[5], num_free_tiles=[5], align_to=4))
         sink = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+        if post:
+            g_bc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+            bt_bc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+            nisa.dma_copy(dst=g_bc, src=g_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
+            nisa.dma_copy(dst=bt_bc, src=bt_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
         n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
         for tb in range(n_grp):
             k = tb % _RING
             nisa.tensor_scalar(dst=z_t[k], data=row[k], op0=nl.add, operand0=neg_mu,
                                op1=nl.multiply, operand1=rstd)
             nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_t[k])
+            if post:
+                # h = (x_new - mu) * rstd * g + bt, fp32; reuse `sink` (its reduce output is dead here)
+                nisa.tensor_scalar(dst=sink, data=row[k], op0=nl.add, operand0=neg_mu,
+                                   op1=nl.multiply, operand1=rstd)
+                nisa.tensor_tensor(dst=sink, data1=sink, data2=g_bc, op=nl.multiply)
+                nisa.tensor_tensor(dst=sink, data1=sink, data2=bt_bc, op=nl.add)
+                nisa.dma_copy(dst=hout_hbm[r0:r0 + _P_MAX, 0:H], src=sink)
+    if post:
+        return xnew_hbm, z_hbm, hout_hbm
     return xnew_hbm, z_hbm
 _ATTN_L1_LN = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
+_ATTN_L1_LN_POST = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
+_ATTN_L1_LN_POST_NR = _ATTN_L1_LN_POST
 @nki_op("vit_enc::attention_block_l1_ln", mutates_args={})
     if bqk.dim() == 3 and bqk.shape[1] != 128:
         bqk = _pack_qk_bias(bqk)
     return _ATTN_L1_LN(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, pack2)
+@nki_op("vit_enc::attention_block_l1_ln_post", mutates_args={})
+def attention_block_l1_ln_post(
+    h: torch.Tensor, wqk: torch.Tensor, bqk: torch.Tensor, tab: torch.Tensor, wvt: torch.Tensor, bv: torch.Tensor,
+    wot: torch.Tensor, bo: torch.Tensor, x: torch.Tensor, scale: float, head_dim: int, n_valid: int, eps: float,
+    g: torch.Tensor, bt: torch.Tensor, nimg: int = 1, resident: int = 0, pack2: int = 0, norope: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Post-LN variant: (x_new fp32, z2 bf16 = LN0(x_new), h fp32 = LN(x_new) * g + bt). g/bt: [H] fp32."""
+    for t in (h, wqk, bqk, tab, wvt, bv, wot, bo, x, g, bt):
+        if t is not None and t.requires_grad:
+            raise RuntimeError("attention_block_l1_ln_post is inference-only; use torch.no_grad().")
+    if h.shape[0] // nimg > 2048:
+        raise NotImplementedError("attention_block_l1_ln_post supports N<=2048 only")
+    if bqk.dim() == 3 and bqk.shape[1] != 128:
+        bqk = _pack_qk_bias(bqk)
+    return _ATTN_L1_LN_POST(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, pack2, g, bt, norope)

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

@@ -54,7 +54,7 @@ def _kernel_assert(condition: bool, error_text: str):
 @nki.jit
-def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
     Args:
@@ -91,6 +91,9 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     _kernel_assert(tuple(a_hbm.shape) == (R, K), "residual a must be [R, K] fp32")
     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)
     inv_k = 1.0 / float(K)
     neg_inv_k = -inv_k
@@ -141,6 +144,11 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     # b2 broadcast ONCE over all K columns (fp32), plus LN per-lane scratch.
     b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
     nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], 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)
@@ -221,7 +229,18 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
                                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])
     return xnew_hbm, z_hbm
@@ -248,3 +267,15 @@ def nki_mlp_res_ln(
     if b1.dim() == 1:
         b1 = pack_b1(b1)
     return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps)

 @nki.jit
+def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps, g_hbm=None, bt_hbm=None):
     """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
     Args:
     _kernel_assert(tuple(a_hbm.shape) == (R, K), "residual a must be [R, K] fp32")
     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)
+    post = g_hbm is not None                       # post-LN mode: also emit h = LN(x_new) * g + bt (fp32)
+    if post:
+        h_hbm = nl.ndarray((R, K), dtype=nl.float32, buffer=nl.shared_hbm)
     inv_k = 1.0 / float(K)
     neg_inv_k = -inv_k
     # b2 broadcast ONCE over all K columns (fp32), plus LN per-lane scratch.
     b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
     nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
+    if post:
+        g_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+        bt_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+        nisa.dma_copy(dst=g_bc[0:_P_MAX, 0:K], src=g_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
+        nisa.dma_copy(dst=bt_bc[0:_P_MAX, 0:K], src=bt_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)
                                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])
+            if post:
+                # h = (x_new - mu) * rstd * g + bt in fp32 (the post-LN residual for the next sublayer)
+                h_t = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+                nisa.tensor_scalar(dst=h_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.tensor_tensor(dst=h_t[0:prs, 0:K], data1=h_t[0:prs, 0:K], data2=g_bc[0:prs, 0:K], op=nl.multiply)
+                nisa.tensor_tensor(dst=h_t[0:prs, 0:K], data1=h_t[0:prs, 0:K], data2=bt_bc[0:prs, 0:K], op=nl.add)
+                nisa.dma_copy(dst=h_hbm[r0:r0 + prs, 0:K], src=h_t[0:prs, 0:K])
+    if post:
+        return xnew_hbm, z_hbm, h_hbm
     return xnew_hbm, z_hbm
     if b1.dim() == 1:
         b1 = pack_b1(b1)
     return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps)
+@nki_op("vit_enc::nki_mlp_res_ln_post", mutates_args={})
+def nki_mlp_res_ln_post(
+    x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor,
+    a: torch.Tensor, eps: float, g: torch.Tensor, bt: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Post-LN variant: (x_new fp32, z bf16 = LN0(x_new), h fp32 = LN(x_new) * g + bt). g/bt: [K] fp32."""
+    _no_grad(x, w1t, b1, w2t, b2, a, g, bt)
+    if b1.dim() == 1:
+        b1 = pack_b1(b1)
+    return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps, g, bt)