v1: V-JEPA2 ViT-L encoder NKI kernels for Trainium2 (4.35x, 44% MFU)

Browse files

Files changed (14) hide show

.gitattributes +4 -2
CHANGELOG.md +26 -0
LICENSE +202 -0
NOTICE +8 -0
README.md +143 -0
build/torch-neuron/__init__.py +27 -0
build/torch-neuron/layers.py +172 -0
build/torch-neuron/metadata.json +1 -0
build/torch-neuron/nki_kernels/__init__.py +19 -0
build/torch-neuron/nki_kernels/attention_block.py +709 -0
build/torch-neuron/nki_kernels/attention_inner.py +1214 -0
build/torch-neuron/nki_kernels/mlp.py +294 -0
build/torch-neuron/nki_kernels/res_layernorm.py +211 -0
build/torch-neuron/nki_kernels/rope_tables.py +47 -0

.gitattributes CHANGED Viewed

@@ -1,3 +1,5 @@
-*.so filter=lfs diff=lfs merge=lfs -text
-*.pyd filter=lfs diff=lfs merge=lfs -text
 *.dylib filter=lfs diff=lfs merge=lfs -text

 *.dylib filter=lfs diff=lfs merge=lfs -text
+*.neff filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

CHANGELOG.md ADDED Viewed

	@@ -0,0 +1,26 @@

+# Changelog
+## v1 (2026-09-25)
+First release. Four NKI kernels replacing the per-layer torch ops of the V-JEPA2 ViT-L video
+encoder on Trainium2, exposed as `NeuronVJEPA2Encoder` for `transformers.KernelConfig`.
+* **4.35x over the stock torch encoder** (104.18 -> 23.94 ms/clip, `torch.compile(backend="neuron")`),
+  **MFU 10.1% -> 44.08%** (neuron-explorer `mfu_estimated_percent`, device window).
+* End-to-end serving 42.8% MFU (24.59 ms/clip) with preprocessing folded into the compiled graph
+  and double-buffered uint8 host-to-device copies.
+* Output **bitwise identical** to the verified reference implementation (`torch.equal`), including
+  through the KernelConfig forward graft under `torch.compile`; rel_l2 4.2e-2 vs the fp32 torch
+  encoder (same as stock -- both run bf16 matmuls).
+* `can_torch_compile = True` set, so the kernels library does not silently fall back to the stock
+  layer under `Mode.TORCH_COMPILE`.
+* The one-time weight fold is `@torch._dynamo.disable`d: compiled folding rounds differently and
+  breaks bitwise parity (observed output norm 4487.45 vs the correct 4486.19).
+* Verified on trn2.3xlarge: neuronx-cc 2.27.2878.0, torch-neuronx 2.11.3, nki 0.6.0,
+  transformers 5.13.1, neuron-explorer 2.32.0.
+Known limitation: 16 frames / 2048 tokens / batch 1 only. 64-frame clips (8192 tokens) are not
+supported -- three of the vendor attention kernel's SBUF buffers scale with keys-per-section
+(347 KiB/lane needed vs a 192 KiB budget, device-confirmed), so 8192 keys require the
+flash-sectioning path whose cross-section combine this integration does not yet reproduce
+correctly. Off-geometry input falls back to the stock encoder.

LICENSE ADDED Viewed

	@@ -0,0 +1,202 @@

+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+   1. Definitions.
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+   END OF TERMS AND CONDITIONS
+   APPENDIX: How to apply the Apache License to your work.
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+   Copyright [yyyy] [name of copyright owner]
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+       http://www.apache.org/licenses/LICENSE-2.0
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE ADDED Viewed

	@@ -0,0 +1,8 @@

+V-JEPA2 encoder NKI kernels for AWS Trainium2
+Copyright 2026 (kernels: res_layernorm.py, attention_block.py, mlp.py, rope_tables.py,
+model glue and tooling -- original work, Apache-2.0)
+kernels/attention_inner.py is a modified copy of the AWS Neuron SDK's nkilib
+core/attention/attention_cte.py, Copyright Amazon.com, Inc. or its affiliates,
+licensed under the Apache License, Version 2.0. See the file header for the
+list of modifications.

README.md ADDED Viewed

	@@ -0,0 +1,143 @@

+---
+license: apache-2.0
+library_name: kernels
+tags:
+  - kernel
+  - neuron
+  - nki
+  - trainium
+  - trn2
+  - vjepa2
+  - video
+  - vision
+  - attention
+  - encoder
+---
+# V-JEPA2 ViT-L Encoder NKI Kernels for AWS Trainium2
+**Version 1 (2026-09-25).** Four NKI kernels replacing the per-layer torch ops of the
+[V-JEPA2](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256) ViT-L video encoder on trn2:
+**4.35x over the stock torch path (104.2 → 23.9 ms/clip), MFU 10.1% → 44.1%** (neuron-explorer
+`mfu_estimated_percent`, device window). Output is **bitwise identical** to the reference
+implementation this was verified against, and rel_l2 4.2e-2 vs the fp32 torch encoder (the
+model's own bf16-matmul sensitivity).
+## What ships
+| Component | Purpose |
+|---|---|
+| `NeuronVJEPA2Encoder` (layer) | KernelConfig-mappable drop-in for `VJEPA2Encoder`: the 4-kernel layer loop, lazy one-time weight fold, stock fallback off-device/off-geometry |
+| `res_layernorm_fold` | K1/K3: residual add + affine-free LayerNorm (fp32 stream in, bf16 normalized out); the LN affines are folded into the neighboring matmuls |
+| `attention_block` | K2: fused QKV projection (LN-gain folded, lane-permuted) + 3D RoPE + flash-attention inner + fused out-projection → [2, N, H] per-core partials (LNC=2) |
+| `mlp` | K4: fc1 + GELU + fc2, bf16 matmuls, fp32 boundary, prologue-prestaged row blocks |
+| `pack_*`, `build_rope_tables` | host-side weight folding (used internally by the layer; exported for direct kernel callers) |
+## Quick start
+```python
+import torch
+from transformers import AutoModel, KernelConfig
+kernel_config = KernelConfig({
+    "VJEPA2Encoder": "jburtoft/vjepa2-neuron-kernels:NeuronVJEPA2Encoder",
+})
+model = AutoModel.from_pretrained(
+    "facebook/vjepa2-vitl-fpc64-256",
+    dtype=torch.float32,
+    kernel_config=kernel_config,
+    device_map="neuron",
+).eval()
+enc = torch.compile(model.encoder, backend="neuron")   # the published numbers are of this graph
+video = your_clip                                       # [1, 16, 3, 256, 256] fp32
+with torch.no_grad():
+    embeddings = enc(video).last_hidden_state           # [1, 2048, 1024]
+```
+> **`torch.compile(backend="neuron")` is part of the configuration.** Eager dispatch adds
+> ~5.5 ms/forward. The layer sets `can_torch_compile`, so the kernels library will not silently
+> swap it out under compile.
+> **Resolving via `get_kernel` / `kernels`:** on a vLLM (0.24) venv where `torch` is a
+> CUDA-tagged build, `import libtorch_neuronx_lite` (or `torch_xla`) **before** `get_kernel(...)`
+> so `torch.neuron` is registered and the `torch-neuron` variant is selected. On a PyTorch Native
+> container (`torch.neuron` already registered) this import is unnecessary.
+## Performance (trn2.3xlarge, verified 2026-09-25)
+| | stock torch encoder | these kernels |
+|---|---|---|
+| throughput, compiled graph | 104.18 ms/clip | **23.94 ms/clip (4.35x)** |
+| device MFU (`mfu_estimated_percent`) | 10.1% | **44.08%** |
+| end-to-end serving, preprocess on device + double-buffered uint8 H2D | — | 24.59 ms/clip = **42.8%** |
+| end-to-end single stream, nothing overlapped | — | 29.86 ms = 35.7% |
+MFU is neuron-explorer's own field: model FLOPs (1.6557e12, equal to the closed-form
+QK+PV+projections+MLP+patch count to 0.000%) over 157.3 TFLOP/s (one logical NeuronCore =
+2 physical cores, bf16) times the window. One logical core of four: a single stream uses ~11%
+of the chip; four concurrent streams reach ~44% chip-wide.
+Where the remaining 2.3x to physics goes (measured, not estimated): the softmax exp runs on the
+Scalar engine at 1 element/lane/cycle (+2.6 ms over physics — unreachable without different
+hardware), instruction issue overhead (~3.8 ms), and the queueing of three engines co-bound
+within 0.5% (~3.7 ms). The exact-math kernel-level candidate classes are enumerated to
+exhaustion; details in the source repository's cycle records.
+## Correctness
+* Full-forward output **bitwise identical** (`torch.equal`) to the verified reference
+  implementation, including through this repo's KernelConfig-style forward graft under
+  `torch.compile`.
+* rel_l2 vs the fp32 torch encoder: 4.2028e-2 — identical to the stock-vs-fp32 difference
+  (both paths run bf16 matmuls; neither is "more approximate" than the other).
+* K1/K3's GpSimd `rsqrt` is measurably *closer* to fp64 than the torch path's (6.9e-8 vs
+  3.9e-5 relative).
+* The vendor attention inner is run-to-run nondeterministic on some layers at a few bf16
+  positions (fp64 agreement unchanged) — an upstream kernel property, reported to the Neuron
+  team. Bitwise comparisons must use a fixed captured sample.
+## Requirements
+* **Hardware**: trn2 (verified on trn2.3xlarge, LNC=2).
+* **SDK**: neuronx-cc 2.27.2878.0; torch-neuronx 2.11.3 (PyTorch Native); nki 0.6.0;
+  nkilib (bundled with the SDK — `ModularAllocator` is imported from it).
+* **Transformers** with `KernelConfig` support; `kernels` library.
+* `NKI_ENABLE_TRACE_CACHE=0` in the environment (the runtime defaults the cross-process trace
+  cache on; a stale NEFF silently captures a different instruction stream).
+## Repository layout
+```
+build/torch-neuron/
+├── __init__.py            # public API re-exports
+├── metadata.json          # kernels-library metadata (backend: neuron)
+├── layers.py              # NeuronVJEPA2Encoder (KernelConfig target)
+└── nki_kernels/
+    ├── res_layernorm.py   # K1/K3
+    ├── attention_block.py # K2 outer (QKV + RoPE + out-proj) and weight packing
+    ├── attention_inner.py # K2 inner: MODIFIED copy of the Neuron SDK's nkilib attention_cte
+    ├── mlp.py             # K4
+    └── rope_tables.py     # host-side 3D RoPE tables
+```
+## Known limitations
+* **16 frames / 2048 tokens / batch 1 / 256x256 only.** The model's native 64-frame clip
+  (8192 tokens) is NOT supported: three of the vendor attention kernel's SBUF buffers scale
+  with keys-per-section (347 KiB/lane needed vs the 192 KiB budget, device-confirmed), so 8192
+  keys require the flash-sectioning path, whose cross-section combine this integration does not
+  yet reproduce correctly. Off-geometry inputs fall back to the stock implementation.
+* **Inference only** — no autograd rule; the layer refuses in training mode.
+* The first device forward performs the one-time weight fold eagerly (~seconds); it is
+  deliberately kept outside `torch.compile` (compiled folding rounds differently and breaks
+  bitwise parity with the verified reference).
+## Attribution and license
+Apache-2.0 (see `LICENSE`, `NOTICE`). `nki_kernels/attention_inner.py` is a **modified copy of
+the AWS Neuron SDK's `nkilib` `core/attention/attention_cte.py`** (Copyright Amazon.com, Inc.
+or its affiliates, Apache-2.0); the file header lists the modifications (four engine-placement
+pins worth ~1.19x end-to-end, and removal of code paths unreachable for this configuration).
+The other three kernels, the layer glue and the packing are original work.

build/torch-neuron/__init__.py ADDED Viewed

	@@ -0,0 +1,27 @@

+"""V-JEPA2 ViT-L encoder NKI kernels for AWS Trainium2 (trn2).
+Public API:
+  NeuronVJEPA2Encoder          -- KernelConfig-mappable layer ("VJEPA2Encoder" -> this)
+  res_layernorm_fold, attention_block, mlp   -- the three kernel entry points
+  pack_* / build_rope_tables / lane_permutation -- host-side weight packing
+"""
+from .layers import NeuronVJEPA2Encoder
+from .nki_kernels import (
+    attention_block,
+    build_rope_tables,
+    lane_permutation,
+    mlp,
+    pack_out_weights,
+    pack_tables,
+    pack_v_weights,
+    pack_weights,
+    res_layernorm_fold,
+)
+__all__ = [
+    "NeuronVJEPA2Encoder",
+    "res_layernorm_fold", "attention_block", "mlp",
+    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
+    "lane_permutation", "build_rope_tables",
+]

build/torch-neuron/layers.py ADDED Viewed

	@@ -0,0 +1,172 @@

+"""``NeuronVJEPA2Encoder`` -- the KernelConfig-mappable layer for the V-JEPA2 encoder.
+USAGE (HF Transformers ``KernelConfig``)::
+    from transformers import AutoModel, KernelConfig
+    kernel_config = KernelConfig({
+        "VJEPA2Encoder": "OWNER/vjepa2-neuron-kernels:NeuronVJEPA2Encoder",  # <- set to the published repo id
+    })
+    model = AutoModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256",
+                                      dtype=torch.float32,
+                                      kernel_config=kernel_config,
+                                      device_map="neuron")
+    model.encoder.forward = model.encoder.forward  # (kernelize does the graft)
+WHY THE ENCODER CLASS AND NOT A PER-LAYER MAPPING. The K1/K3 kernel fuses the *previous*
+sublayer's residual add with the *next* sublayer's LayerNorm, so the fusion spans transformer-
+layer boundaries: the state carried between layers is the pair (fp32 residual stream, bf16
+branch partials), not a single summed hidden_states. Replacing the whole encoder forward keeps
+that chain -- and the instruction stream it was verified bitwise against -- intact. A per-layer
+mapping would force a residual re-add at every boundary and change the stream.
+MECHANICS. The kernels library grafts this class's ``forward`` onto the existing
+``VJEPA2Encoder`` instance, so ``self`` inside ``forward`` is the *stock* encoder with its
+stock parameters (checkpoint loading is untouched). On the first device forward, the per-layer
+packed weights (LayerNorm affines folded into QKV/V and MLP-fc1, out-proj bias halved for the
+LNC=2 partials, rope tables) are computed once from those stock parameters and cached on the
+layer objects. Off-device or non-fp32 inputs fall back to the stock implementation.
+``can_torch_compile = True`` is REQUIRED here: the published performance (23.94 ms/forward,
+44% MFU) is of the ``torch.compile(backend="neuron")`` graph, and without this flag the kernels
+library silently swaps the layer back to stock under Mode.TORCH_COMPILE (observed on other
+Neuron kernel repos as a 25-127x silent regression).
+"""
+from __future__ import annotations
+import torch
+from .nki_kernels import (
+    attention_block,
+    build_rope_tables,
+    mlp,
+    pack_out_weights,
+    pack_tables,
+    pack_v_weights,
+    pack_weights,
+    res_layernorm_fold,
+)
+_PACKED_ATTR = "_vjepa2_nki_packed"
+_TAB_ATTR = "_vjepa2_nki_tab"
+def _pack_layer(lyr, tab):
+    """Fold one transformer block's LayerNorm affines into the kernel weight packings.
+    Exact fp32 fold on CPU (the algebra of the verified reference implementation):
+      K1 gamma1/beta1 -> QKV and V packs;  out-proj bias HALVED (two cores' partials sum);
+      K3 gamma2/beta2 -> MLP fc1;  fc2 unfolded.
+    """
+    attn = lyr.attention
+    dev = attn.query.weight.device
+    H = attn.all_head_size
+    D = attn.attention_head_size
+    heads = H // D
+    g1 = lyr.norm1.weight.detach().float().cpu()
+    t1 = lyr.norm1.bias.detach().float().cpu()
+    g2 = lyr.norm2.weight.detach().float().cpu()
+    t2 = lyr.norm2.bias.detach().float().cpu()
+    wq = attn.query.weight.detach().float().cpu(); bq = attn.query.bias.detach().float().cpu()
+    wk = attn.key.weight.detach().float().cpu();   bk = attn.key.bias.detach().float().cpu()
+    wv = attn.value.weight.detach().float().cpu(); bv = attn.value.bias.detach().float().cpu()
+    wqk_f, bqk_f = pack_weights(wq * g1[None, :], wk * g1[None, :],
+                                bq + wq @ t1, bk + wk @ t1, heads, D)
+    wvt_f, bv_f = pack_v_weights(wv * g1[None, :], bv + wv @ t1)
+    wot, bo = pack_out_weights(attn.proj.weight, attn.proj.bias)
+    w1 = lyr.mlp.fc1.weight.detach().float().cpu(); b1 = lyr.mlp.fc1.bias.detach().float().cpu()
+    packed = {
+        "D": D, "heads": heads,
+        "k1_eps": float(lyr.norm1.eps), "k3_eps": float(lyr.norm2.eps),
+        "scale": float(attn.scaling),
+        "wqk_f": wqk_f.to(dev), "bqk_f": bqk_f.to(dev), "tab": tab,
+        "wvt_f": wvt_f.to(dev), "bv_f": bv_f.to(dev), "wot": wot.to(dev), "bo": bo.to(dev),
+        "w1t_f": (w1 * g2[None, :]).t().contiguous().to(torch.bfloat16).to(dev),
+        "b1_f": (b1 + w1 @ t2).contiguous().to(dev),
+        "w2t": lyr.mlp.fc2.weight.detach().t().contiguous().to(torch.bfloat16).to(dev),
+        "b2": lyr.mlp.fc2.bias.detach().float().contiguous().to(dev),
+    }
+    return packed
+def _rope_tables(encoder, tokens):
+    attn0 = encoder.layer[0].attention
+    H = attn0.all_head_size
+    D = attn0.attention_head_size
+    dev = attn0.query.weight.device
+    pos = attn0.get_position_ids(torch.zeros(1, tokens, H, device=dev), masks=None)
+    def to1d(t):
+        return t if t.dim() == 1 else t.reshape(-1, tokens)[0]
+    cos, sinsigned = build_rope_tables(
+        (attn0.d_dim, attn0.h_dim, attn0.w_dim),
+        (to1d(pos[0]), to1d(pos[1]), to1d(pos[2])),
+        D, torch.float32, dev)
+    return pack_tables(cos, sinsigned, H // D, D)
+@torch._dynamo.disable
+def _freeze(encoder, tokens):
+    """Pack every layer once; cached on the layer objects. Refuses in training mode.
+    ``@torch._dynamo.disable`` is load-bearing: without it, a user who wraps the model in
+    ``torch.compile`` before the first forward gets the weight-fold arithmetic traced/executed
+    by the compiler instead of eagerly, and the packed weights round DIFFERENTLY (measured:
+    output norm 4487.45 vs the verified 4486.19). The fold must run eagerly, once.
+    """
+    if encoder.training:
+        raise RuntimeError(
+            "NeuronVJEPA2Encoder: encoder.training is True; these kernels have no autograd "
+            "rule. Call model.eval() first.")
+    tab = _rope_tables(encoder, tokens)
+    setattr(encoder, _TAB_ATTR, tab)
+    for lyr in encoder.layer:
+        setattr(lyr, _PACKED_ATTR, _pack_layer(lyr, tab))
+class NeuronVJEPA2Encoder(torch.nn.Module):
+    """Drop-in forward for ``VJEPA2Encoder``: the 4-kernel NKI layer loop.
+    Supported: trn2, 16 frames / 2048 tokens / batch 1, fp32 model. Anything else falls back
+    to the stock implementation (off-device, wrong geometry) or raises (training mode during
+    the first-freeze).
+    """
+    can_torch_compile = True
+    has_backward = False
+    def forward(self, pixel_values_videos=None, **kwargs):
+        hs = self.embeddings(pixel_values_videos)
+        N = hs.shape[1]
+        # Fallback: off-device, or an unsupported geometry (only N=2048 is verified; the vendor
+        # attention kernel cannot hold 8192 keys in SBUF -- see the card's Known limitations).
+        if hs.device.type != "neuron" or N != 2048:
+            import transformers.models.vjepa2.modeling_vjepa2 as mv
+            return mv.VJEPA2Encoder.forward(self, pixel_values_videos, **kwargs)
+        if not hasattr(self.layer[0], _PACKED_ATTR):
+            _freeze(self, N)
+        x = hs[0]                                        # [N, H] fp32 residual stream
+        H = x.shape[1]
+        branch = torch.zeros(1, N, H, dtype=torch.bfloat16, device=x.device)
+        for lyr in self.layer:
+            p = getattr(lyr, _PACKED_ATTR)
+            x, z1 = res_layernorm_fold(x, branch, p["k1_eps"])
+            attn_partials = attention_block(
+                z1, p["wqk_f"], p["bqk_f"], p["tab"], p["wvt_f"], p["bv_f"],
+                p["wot"], p["bo"], p["scale"], p["D"])
+            x, z2 = res_layernorm_fold(x, attn_partials, p["k3_eps"])
+            branch = mlp(z2, p["w1t_f"], p["b1_f"], p["w2t"], p["b2"]).unsqueeze(0)
+        x = x + branch[0].to(torch.float32)
+        out = self.layernorm(x.unsqueeze(0))
+        import transformers.models.vjepa2.modeling_vjepa2 as mv
+        return mv.BaseModelOutput(last_hidden_state=out)

build/torch-neuron/metadata.json ADDED Viewed

	@@ -0,0 +1 @@


1	+ { "name": "vjepa2-neuron-kernels", "id": "vjepa2_neuron_kernels", "version": 1, "license": "Apache-2.0", "backend": {"type": "neuron", "archs": []}, "python-depends": [] }

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

	@@ -0,0 +1,19 @@

+"""The four NKI kernels of the V-JEPA2 encoder plus their host-side weight packing."""
+from .attention_block import (
+    attention_block,
+    lane_permutation,
+    pack_out_weights,
+    pack_tables,
+    pack_v_weights,
+    pack_weights,
+)
+from .mlp import mlp
+from .res_layernorm import res_layernorm_fold
+from .rope_tables import build_rope_tables
+__all__ = [
+    "res_layernorm_fold", "attention_block", "mlp",
+    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
+    "lane_permutation", "build_rope_tables",
+]

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

	@@ -0,0 +1,709 @@

+"""K2 outer: fused QKV projection + 3D RoPE + vendor attention + fused out-projection, in ONE
+top-level NKI kernel, FLATTENED from ``nki_attn_block_k2v_h10`` (which forks
+``nki_attn_block_k2v_dp4``). Self-contained: imports only ``nki``, ``nkilib``, ``torch``, and
+``_attention_cte`` from the sibling ``attention_inner`` -- NO imports from ``experiment/`` or
+any project-local modules.
+WHAT THIS COMPUTES (per core; LNC=2, this core owns 8 of the 16 heads)
+======================================================================
+``result[LNC, N, H] bf16`` = out-projection PARTIALS of
+``attention(rope(h@Wq), rope(h@Wk), h@Wv)``. Each core computes a partial sum over its 8 heads'
+512 out-projection contraction rows; K3's residual add sums the two cores' partials.
+Phase 1 (producer): dma_transpose h -> x_t; project Q/K (lane-permuted) and V; apply the six
+rope ``tensor_tensor`` ops; scatter K into ``k_tiles`` (SBUF), Q into ``q_priv`` (private_hbm),
+V into ``v_tiles`` (SBUF).
+Phase 2 (attention + out-proj): reset the allocator to reclaim the dead producer buffers, call
+``_attention_cte`` per local head into feature-major SBUF O slices, then contract the fused
+out-projection (tokens on partitions) and DMA token-major to ``result``.
+PSUM BANK PLAN (producer): V uses bank 7 (``(0, 7*_PSUM_BANK)``); Q/K parities use banks 5/6
+(``(5+par)*_PSUM_BANK``); the fused out-projection uses bank 0 (``(0, 0)``). The vendor
+attention (``attention_inner``) pins its own low banks {0-3} for MM1 and {4-7} for MM2 during
+its phase, which runs after the producer's high-water mark, so there is no live overlap.
+THE TWO ENGINE PINS (both to Scalar), and WHY (from PERF_INVENTORY)
+-------------------------------------------------------------------
+  * ``:276`` V-scatter ``tensor_copy`` -> SCALAR (H9's inherited mechanism): eight copies read
+    one ``vout`` before the next vsub overwrites it (a WAR chain of 8); Scalar (240 ns) beats
+    GpSimd (375 ns) for that chain and sheds the copy off Vector, K2's binding engine.
+  * ``:458`` QK bias-add-and-evict ``tensor_scalar`` -> SCALAR (H10's mechanism): the largest
+    remaining Vector-resident relocatable site (768/core); ``data=ps`` is PSUM so GpSimd is
+    illegal, and Scalar's 612 ns fee is faster than the 744 ns Vector cell it replaces, so the
+    zero-lead rope chain that consumes ``proj`` cannot lengthen.
+The ``:224`` bias-cast ``tensor_copy`` and ``:483`` K-scatter ``tensor_copy`` stay UNPINNED
+(compiler default) -- pinning them was dropped (:224 is 0.09% of the window; :483 has no
+affordable landing engine), and they MUST stay unpinned to match the shipping stream.
+HOST-SIDE FOLD ALGEBRA (packing helpers below, copied flat)
+-----------------------------------------------------------
+Affine-free LN1 folds gamma1/beta1 into the K2 weights at freeze time:
+``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, ...)`` and ``pack_v_weights(wv*g1, bv+wv@t1)``.
+The lane permutation puts each quad-of-heads' even/odd RoPE lanes in their own 128-partition
+group so RoPE is six full-width ``tensor_tensor`` ops with no transposes; it cancels in
+attention's Q/K contraction (a dot product is invariant under a shared permutation of the
+contraction axis), so V is natural and nothing downstream un-permutes.
+LNC2 BIAS HALVING: ``pack_out_weights`` halves the out-projection bias (``bo*0.5``) so the two
+cores' partials sum to exactly one bias -- halving is bitwise exact in bf16 (it only decrements
+the exponent, and summing two identical bf16 values is exact).
+INFERENCE ONLY: no autograd; the packed weights are detached caches.
+SBUF BUDGET (bytes/lane; budget = 192 KiB = 196608 B/lane), TWO GEOMETRIES
+==========================================================================
+N=2048 (16-frame clips) is the SHIPPING path: nothing below is gated on, so buffer shapes,
+byte layout, DMA target addresses and fill VALUES are BIT-FOR-BIT the reference version (verified:
+torch.equal(perf_out.pt) == True, max abs diff 0.0). N=8192 (64-frame clips) is a SEPARATE,
+gated path (``if N > 2048``) that the N<=2048 forward never enters.
+                                   |  N=2048 (16f)  |  N=8192 (64f, gated)
+    ------------------------------ + -------------- + --------------------------------------
+    k_tiles (8 heads resident)     |     32768      |  0  (streamed to private_hbm)
+    v_tiles (8 heads resident)     |     16384      |  0  (streamed to private_hbm)
+    k_stage / v_stage (1 head)     |       -        |  16384 + 8192 = 24576 (one head at a time)
+    x_t   (hidden^T ring)          |     32768      |  32768 (depth-4 token-block ring, was 131072)
+    tab   (rope tables ring)       |     32768      |  32768 (depth-4 token-block ring, was 131072)
+    producer working set           |    ~25000      |  ~25000
+    ------------------------------ + -------------- + --------------------------------------
+    PHASE-1 PEAK                   |    ~124 KiB    |  ~106 KiB   (FITS)
+    ------------------------------ + -------------- + --------------------------------------
+    phase2_base (persistent K/V)   |     49152      |  24576 (only k_stage/v_stage persist)
+    wo_sb + bo_bc + o_feat + out   |     29696      |  78848
+    vendor inner (per head)        |    ~45 KiB     |  ~251 KiB (single 8192 section, n_large=4:
+                                   |                |   exp_tp_sb 131072 + mm1_masked 65536 +
+                                   |                |   exp_sb 32768 balloon -- see below)
+    ------------------------------ + -------------- + --------------------------------------
+    PHASE-2 PEAK                   |    ~125 KiB    |  ~354912 B/lane  (DOES NOT FIT -- measured
+                                   |                |   on device [NCC_INLA001] 128x354912, 1.81x
+                                   |                |   the 196608 budget => N>2048 is REFUSED)
+WHY N>2048 IS REFUSED. At N=8192 the shipping design has TWO independent SBUF overflows
+(both confirmed on device as ``[NCC_INLA001] ... out of bound 128x526944``):
+  (1) PHASE 1: the four N-scaled residents k_tiles(131072)+v_tiles(65536)+x_t(131072)+
+      tab(131072) = 458752 B/lane coexist -- 2.3x the whole budget.
+  (2) PHASE 2: eight heads' K/V resident (196608 B/lane) is the ENTIRE budget.
+The phase-1 fixes work (x_t/tab depth-4 token-block rings, bit-for-bit at N<=2048, saving
+196608 B/lane; K/V streamed through private_hbm one head at a time) => phase-1 peak ~106 KiB.
+But PHASE 2 does NOT fit: handing one head all 8192 keys makes the single-section flatten run
+with num_large_tiles_per_section=ceil(8192/2048)=4, ballooning three inner buffers (exp_tp_sb
+131072 + mm1_masked 65536 + exp_sb 32768 = 229376 B/lane) -> phase-2 peak measured 354912 B/lane,
+1.81x the budget. The shrinkable phase-2 region (o_feat/out/wo/bo = 78848) is smaller than the
+158304 overflow, so streaming O or shrinking o_feat cannot close the gap; re-enabling flash
+sectioning fits but is numerically wrong (rel_l2 0.90, ``perf/SBUF_8192_STOPPED.md``), and
+editing the inner would perturb the bit-for-bit 16-frame path. So the host wrapper
+``attention_block`` raises ``NotImplementedError`` for N>2048 rather than emit a non-fitting or
+wrong kernel (see ``perf/SBUF_8192_STOPPED2.md``). The N<=2048 path is bit-for-bit unchanged.
+"""
+from __future__ import annotations
+import torch
+from nki.framework.kernel import wrap_nki_to_torch_hop
+from torch_neuronx import nki_op
+import nki
+import nki.isa as nisa
+import nki.language as nl
+from nki.isa import engine
+from nkilib.core.utils.modular_allocator import ModularAllocator
+# The vendor attention, flattened to this model's reached path with the h8/h10 pins inlined.
+# The N<=2048 path calls THIS single-section flatten; the N>2048 path is REFUSED (see below).
+# seqlen_k=8192 < the vendor's _FLASH_ATTENTION_THRESHOLD (10240), so num_sections==1 -- the
+# SAME single-section, non-flash code path the verified 16-frame configuration uses (no
+# multi-section online-softmax combine, which is numerically wrong here: SBUF_8192_STOPPED.md).
+# BUT single-section at 8192 does NOT fit: the inner's mm1_masked / exp_sb / exp_tp_sb are
+# allocated with num_free_tiles carrying num_large_tiles_per_section = ceil(8192/2048) = 4
+# (attention_inner.py:658-701), so those three buffers BALLOON with total_seqlen_k -- they are
+# NOT reused across the large-tile loop. exp_tp_sb 131072 + mm1_masked 65536 + exp_sb 32768 =
+# 229376 B/lane, driving phase-2 peak to ~354912 B/lane (1.81x the 196608 budget; measured on
+# device [NCC_INLA001] 128x354912). See the SBUF table and refusal rationale below.
+from .attention_inner import _attention_cte
+LNC = 2
+_P_MAX = 128
+_PSUM_FMAX = 512
+_PSUM_BANK = 2048          # bytes per partition per bank
+_K_TILE = 512              # vendor k_cache_sbuf tiles are (d, 512)
+_V_TILE = 128              # vendor v_cache_sbuf tiles are (128, d)
+def _kernel_assert(cond: bool, msg: str):
+    assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
+# ==========================================================================================
+# Host-side packing helpers (copied flat from nki_qk_fused_2core / nki_attn_block_k2v_dp4).
+# All are frozen-weight work, cached once; INFERENCE ONLY.
+# ==========================================================================================
+def lane_permutation(heads: int, D: int) -> torch.Tensor:
+    """Column order that puts each quad-of-heads' evens and odds in their own 128 group.
+    FOLD ALGEBRA: RoPE pairs adjacent lanes ``out[2j]=x[2j]*cos+x[2j+1]*sinsigned`` etc. A
+    projection produces lanes on the partition axis, so a pair is a stride-2 partition access.
+    Reorder the projection's output columns so group ``2a`` holds the EVEN lanes of heads
+    ``4a..4a+3`` and group ``2a+1`` their ODD lanes (``quad*half = 4*(D//2) = 128`` rows when
+    D==64). Then one matmul yields a full [128,512] all-evens tile and another all-odds, and
+    RoPE is six full-width ``tensor_tensor`` ops -- no transposes, no partition waste. The
+    permutation cancels in attention's Q/K contraction, so nothing downstream un-permutes.
+    Returns ``perm`` of length ``heads*D`` s.t. ``w.t()[:, perm]`` is the reordered weight:
+    ``perm[p]`` is the natural column that lands at permuted row ``p % 128`` of group ``p // 128``.
+    """
+    half = D // 2
+    quad = _P_MAX // half
+    _kernel_assert(quad * half == _P_MAX, f"D={D} gives half={half}, needs quad*half=128")
+    _kernel_assert(heads % quad == 0, f"heads={heads} must be divisible by quad={quad}")
+    perm = torch.empty(heads * D, dtype=torch.long)
+    for hd in range(heads):
+        a, b = hd // quad, hd % quad
+        for j in range(half):
+            perm[(2 * a) * _P_MAX + b * half + j] = hd * D + 2 * j          # even lane
+            perm[(2 * a + 1) * _P_MAX + b * half + j] = hd * D + 2 * j + 1  # odd lane
+    return perm
+def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
+                 bq: torch.Tensor | None, bk: torch.Tensor | None,
+                 heads: int, D: int):
+    """``(w_packed [2,H,H] bf16, b_packed [2,H,1] bf16)`` -- transposed and column-permuted.
+    FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
+    ``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, heads, D)``. ``wq``/``wk`` are torch
+    ``Linear.weight`` ``[H_out, H_in]``; the kernel wants ``[H_in, H_out]`` with ``H_out``
+    permuted, i.e. ``w.t()[:, perm]``. The bias is shaped ``[2,H,1]`` (not ``[2,H]``) so the
+    kernel can slice a ``[128,1]`` tile straight out of HBM.
+    """
+    H = wq.shape[0]
+    perm = lane_permutation(heads, D).to(wq.device)
+    w = torch.stack([wq.detach().t()[:, perm].contiguous(),
+                     wk.detach().t()[:, perm].contiguous()], 0)
+    zero = torch.zeros(H, dtype=wq.dtype, device=wq.device)
+    b = torch.stack([(zero if bq is None else bq.detach())[perm].contiguous(),
+                     (zero if bk is None else bk.detach())[perm].contiguous()], 0)
+    b = b.unsqueeze(-1)
+    return w.to(torch.bfloat16).contiguous(), b.to(torch.bfloat16).contiguous()
+def pack_tables(cos: torch.Tensor, sinsigned: torch.Tensor, heads: int, D: int):
+    """``tab [4, 128, N] fp32`` = (cos_even, cos_odd, sin_even, sin_odd), quad-tiled.
+    FOLD ALGEBRA: ``cos``/``sinsigned`` are ``[N, D]`` shared by every head. A group's 128 rows
+    are one head's 32 even (or odd) lanes tiled ``quad`` (=4) times to match the lane_permutation
+    layout, so each RoPE operand is a full 128-partition tile with a matching partition base.
+    Built once via the rope-table build and shared across all layers.
+    """
+    half = D // 2
+    quad = _P_MAX // half
+    parts = [cos[:, 0::2], cos[:, 1::2], sinsigned[:, 0::2], sinsigned[:, 1::2]]
+    tab = torch.stack([p.t().contiguous().repeat(quad, 1) for p in parts], 0)
+    return tab.to(torch.float32).contiguous()
+def pack_v_weights(wv: torch.Tensor, bv: torch.Tensor | None):
+    """``(wv_t [H, H] bf16, bv [1, H] bf16)``. V is NOT lane-permuted (natural order).
+    FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
+    ``pack_v_weights(wv*g1, bv+wv@t1)``.
+    """
+    H = wv.shape[0]
+    zero = torch.zeros(H, dtype=wv.dtype, device=wv.device)
+    return (wv.detach().t().contiguous().to(torch.bfloat16),
+            (bv.detach() if bv is not None else zero).reshape(1, H).contiguous().to(torch.bfloat16))
+def pack_out_weights(wo, bo):
+    """``(wo_t [H, H] bf16, bo/2 [1, H] bf16)`` for the fused out projection.
+    ``wo`` is ``attn.proj.weight`` [H, H] (attention features -> hidden). The kernel contracts
+    over attention features, so it wants those on partitions: ``wo.T`` = [features, hidden].
+    LNC2 BIAS HALVING. Both cores add ``bo/2`` and K3's partial sum restores ``bo``. Adding the
+    full bias on both would double it; adding it on one core only would need ``program_id`` as an
+    arithmetic operand (it is symbolic -- neither a compile-time constant nor a [P,1] tile).
+    Halving is bitwise exact in bf16: it only decrements the exponent, and summing two identical
+    bf16 values is exact.
+    """
+    import torch as _t
+    Hh = wo.shape[0]
+    zero = _t.zeros(Hh, dtype=wo.dtype, device=wo.device)
+    half = (bo.detach() if bo is not None else zero) * 0.5
+    return (wo.detach().t().contiguous().to(_t.bfloat16),
+            half.reshape(1, Hh).contiguous().to(_t.bfloat16))
+# ==========================================================================================
+# THE KERNEL. Flattened from _attn_block_k2v_h10_dp_kernel (verbatim dp4 body with the two
+# engine= keywords hardcoded to engine.scalar at :276 and :458; :224 and :483 left unpinned).
+# ==========================================================================================
+@nki.jit
+def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
+                            wot_hbm, bo_hbm, scale, head_dim):
+    """out[16, N, D] bf16 = attention(rope(h@Wq), rope(h@Wk), h@Wv), one kernel (per-core partials)."""
+    N, H = h_hbm.shape
+    D = head_dim
+    half = D // 2
+    quad = _P_MAX // half                 # heads per 128-partition parity group = 4
+    heads = H // D                        # 16
+    local_heads = heads // LNC            # 8
+    local_quads = local_heads // quad     # 2
+    n_chunks = H // _P_MAX                # 8 contraction chunks
+    n_tt = N // _PSUM_FMAX                # 4 token tiles of 512
+    n_ktiles = N // _K_TILE               # 4 K tiles / head
+    n_vtiles = N // _V_TILE               # 16 V tiles / head
+    n_ochunks = (local_heads * D) // _P_MAX   # feature chunks of the out projection
+    v_per_tt = _PSUM_FMAX // _V_TILE      # 4 V tiles per 512-token tile
+    # N>2048 gate (see the phase-1/phase-2 note below). Computed early so the K/V residents can
+    # be allocated as HBM (streamed) for the big path.
+    _big = N > _PSUM_FMAX * 4
+    # N>2048 IS NOT SUPPORTED (guarded at the host wrapper `attention_block` below, which raises
+    # NotImplementedError before this kernel is ever traced -- NKI kernels cannot `raise`). The
+    # phase-1 memory fixes (x_t/tab depth-4 rings, K/V streamed through private_hbm one head at a
+    # time) DO fit -- phase-1 peak ~106 KiB/lane. But phase 2 does NOT: handing one head all 8192
+    # keys makes the flattened single-section _attention_cte run with num_large_tiles_per_section
+    # =ceil(8192/2048)=4, which balloons three N-scaled inner buffers -- exp_tp_sb (8*4*4*512*2 =
+    # 131072), mm1_masked (2*4*2048*4 = 65536), exp_sb (2*4*2048*2 = 32768) = 229376 B/lane for
+    # three buffers alone. Measured on device (2026-09-25, neuronx-cc 2.25): [NCC_INLA001] out of
+    # bound 128x354912 -> phase-2 peak 354912 B/lane, 1.81x the 196608 budget. The shrinkable
+    # phase-2 region (o_feat 65536 + out 1024 + wo/bo 12288 = 78848) is smaller than the 158304
+    # overflow, so streaming O or shrinking o_feat CANNOT close the gap -- the inner's N-scaled
+    # buffers are the wall. Re-enabling flash sectioning (2048-key sections) shrinks the inner but
+    # activates the vendor's multi-section online-softmax combine, a path the shipping config never
+    # exercises and which produced rel_l2 0.90 (wrong). See perf/SBUF_8192_STOPPED2.md.
+    _kernel_assert(wqk_hbm.shape == (2, H, H), f"wqk must be [2,{H},{H}]")
+    _kernel_assert(bqk_hbm.shape == (2, H, 1), f"bqk must be [2,{H},1]")
+    _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), f"tab must be [4,{_P_MAX},{N}]")
+    _kernel_assert(wvt_hbm.shape == (H, H) and bv_hbm.shape == (1, H), "bad V weights")
+    _kernel_assert(quad * half == _P_MAX, f"D={D} must give 4 heads per group")
+    _kernel_assert(heads % (quad * LNC) == 0, "heads must split into per-core quads")
+    _kernel_assert(local_heads * D == _PSUM_FMAX,
+                   "inline V projection wants all local heads in one 512-wide stream")
+    _kernel_assert(N % _PSUM_FMAX == 0, f"N={N} must tile by 512")
+    _kernel_assert(str(h_hbm.dtype) == str(nl.bfloat16), "h must be bf16")
+    # [LNC, N, H] PARTIALS. The out projection contracts over all H features == all 16 heads,
+    # and this core owns 8, so what it computes is a partial sum; K3 adds the two.
+    result = nl.ndarray((LNC, N, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    pid = nl.program_id(0)                # SYMBOLIC: index arithmetic only
+    head0 = pid * local_heads
+    col0 = pid * (local_heads * D)        # first permuted qk weight column
+    vcol0 = pid * (local_heads * D)       # first natural v weight column (same offset)
+    allocator = ModularAllocator(initial_address=0)
+    # ---- persistent consumer tiles (written by producer, read by attention) ---------
+    # N<=2048: all 8 heads' K/V resident in SBUF (bit-for-bit shipping path). N>2048: K/V for
+    # all 8 heads (196608 B/lane) is the ENTIRE 192 KiB budget, so they cannot stay resident
+    # alongside x_t/tab/working set. STREAM them: the producer scatters each head's K/V to
+    # private_hbm, and phase-2 loads ONE head's K/V into small resident staging rings just
+    # before its attention call (attention runs per head and writes O to HBM, so only one
+    # head's K/V need be live at a time).
+    k_tiles = []
+    v_tiles = []
+    k_hbm = []
+    v_hbm = []
+    if _big:
+        for _lh in range(local_heads):
+            k_hbm.append(nl.ndarray((n_ktiles, D, _K_TILE), dtype=nl.bfloat16,
+                                    buffer=nl.private_hbm))
+            v_hbm.append(nl.ndarray((n_vtiles, _V_TILE, D), dtype=nl.bfloat16,
+                                    buffer=nl.private_hbm))
+        # ONE head's resident staging rings (reused across heads in phase 2).
+        k_stage = allocator.alloc_sbuf_tensor(
+            shape=(D, _K_TILE), dtype=nl.bfloat16,
+            block_dim=[n_ktiles], num_free_tiles=[n_ktiles], align_to=32)
+        v_stage = allocator.alloc_sbuf_tensor(
+            shape=(_V_TILE, D), dtype=nl.bfloat16,
+            block_dim=[n_vtiles], num_free_tiles=[n_vtiles])
+    else:
+        for _lh in range(local_heads):
+            kt_ = allocator.alloc_sbuf_tensor(
+                shape=(D, _K_TILE), dtype=nl.bfloat16,
+                block_dim=[n_ktiles], num_free_tiles=[n_ktiles], align_to=32)
+            k_tiles.append(kt_)
+            vt_ = allocator.alloc_sbuf_tensor(
+                shape=(_V_TILE, D), dtype=nl.bfloat16,
+                block_dim=[n_vtiles], num_free_tiles=[n_vtiles])
+            v_tiles.append(vt_)
+    # HIGH WATER MARK for phase 2: k_tiles/v_tiles are the only phase-1 products attention
+    # needs, so every buffer above this address is dead once the producer loop finishes.
+    phase2_base = allocator.get_current_address()
+    # ------------------------------------------------------------------------------------
+    # N>2048 GATE (64-frame clips, N=8192). At N<=2048 every ring below keeps its full-N
+    # depth, so the buffer shape, byte layout, DMA target addresses and fill VALUES are
+    # bit-for-bit identical to the shipping 16-frame path (see the SBUF budget table in the
+    # module docstring), and the attention call goes to the pinned single-section flatten
+    # ``_attention_cte`` (all engine pins intact). At N>2048 two independent overflows appear:
+    #   (1) PHASE 1: the four N-scaled residents (k_tiles+v_tiles+x_t+tab = 458752 B/lane)
+    #       coexist. FIX: x_t and tab become depth-XT_RING token-block rings (a rolling 4-tile
+    #       window filled inside the per-token-tile loop; a pure copy, identical values).
+    #   (2) PHASE 2: eight heads' K/V resident (196608 B/lane) is the whole budget. FIX: stream
+    #       K/V through private_hbm and stage ONE head at a time into k_stage/v_stage. The
+    #       attention call is the SAME single-section flatten ``_attention_cte`` the N<=2048 path
+    #       uses (seqlen_k=8192 < 10240 threshold => num_sections==1, no multi-section flash
+    #       combine); its large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers.
+    # The N<=2048 path enters NEITHER branch and is bit-for-bit unchanged.
+    _big = N > _PSUM_FMAX * 4                 # N>2048  -> the 64-frame geometry
+    XT_RING = min(n_tt, 4) if _big else n_tt  # token-block ring depth for x_t / tab
+    # ---- resident producer inputs ----------------------------------------------------
+    # x^T: hidden on partitions, tokens on the free axis, as a per-512-token-block ring.
+    # At N<=2048 XT_RING==n_tt so the ring holds every block (== the old (128,N) buffer,
+    # same bytes, same addresses); at N=8192 it holds a rolling window of 4 blocks. The fill
+    # is a pure dma_transpose (no arithmetic), so any fill ORDER yields identical values.
+    x_t = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
+        block_dim=[n_chunks, n_tt], num_free_tiles=[n_chunks, XT_RING], align_to=32)
+    if not _big:
+        # N<=2048: fill the whole buffer up front, exactly as the shipping path did (pure
+        # dma_transpose; block-indexed AP lands at the same addresses as the old (128,N) buffer).
+        for s in range(n_tt):
+            for c in range(n_chunks):
+                nisa.dma_transpose(
+                    dst=x_t[c][s].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1], [1, _PSUM_FMAX]],
+                                     offset=0),
+                    src=h_hbm.ap([[H, _PSUM_FMAX], [1, 1], [1, 1], [1, _P_MAX]],
+                                 offset=(s * _PSUM_FMAX) * H + c * _P_MAX),
+                )
+    # this core's Q/K weight columns (lane-permuted) and V columns (natural)
+    w_sb = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, local_quads * 2 * _P_MAX), dtype=nl.bfloat16,
+        block_dim=[2, n_chunks], num_free_tiles=[2, n_chunks])
+    for qk in range(2):
+        for c in range(n_chunks):
+            nisa.dma_copy(dst=w_sb[qk][c],
+                          src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX,
+                                      col0:col0 + local_quads * 2 * _P_MAX])
+    wv_sb = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, local_heads * D), dtype=nl.bfloat16,
+        block_dim=[n_chunks], num_free_tiles=[n_chunks])
+    for c in range(n_chunks):
+        nisa.dma_copy(dst=wv_sb[c],
+                      src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX,
+                                  vcol0:vcol0 + local_heads * D])
+    # biases: qk as (128,1) per group; v broadcast to (128, 8*64) once
+    b_sb = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, 1), dtype=nl.float32,
+        block_dim=[2, local_quads * 2], num_free_tiles=[2, local_quads * 2], align_to=4)
+    b_stage = allocator.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.bfloat16)
+    for qk in range(2):
+        for g in range(local_quads * 2):
+            nisa.dma_copy(dst=b_stage,
+                          src=bqk_hbm[qk, col0 + g * _P_MAX:col0 + (g + 1) * _P_MAX, 0:1])
+            # SITE :224 -- bias CAST, left UNPINNED on purpose (0.09% of the window, already at
+            # the Vector fixed fee; nothing to move).
+            nisa.tensor_copy(dst=b_sb[qk][g], src=b_stage)
+    bv_bc = allocator.alloc_sbuf_tensor(
+        shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
+    nisa.dma_copy(dst=bv_bc,
+                  src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]],
+                                offset=vcol0))
+    # rope tables, resident (quad-tiled on the host: identical rows for every quad). Token-block
+    # ring, same gating as x_t: at N<=2048 TAB_RING==n_tt holds every block (== old (128,N)
+    # buffer, bit-for-bit); at N=8192 a rolling 4-block window filled inside the tt loop. Pure
+    # dma_copy -> identical values regardless of fill order.
+    tab = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+        block_dim=[4, n_tt], num_free_tiles=[4, XT_RING], align_to=4)
+    if not _big:
+        for s in range(n_tt):
+            for t_i in range(4):
+                nisa.dma_copy(dst=tab[t_i][s],
+                              src=tab_hbm[t_i, 0:_P_MAX,
+                                          s * _PSUM_FMAX:(s + 1) * _PSUM_FMAX])
+    # ---- per-tile working set (reused across iterations; same tensor objects) --------
+    proj = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, block_dim=[2],
+        num_free_tiles=[2])
+    acc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                      align_to=4)
+    tmp = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                      align_to=4)
+    roped = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, block_dim=[2],
+        num_free_tiles=[2])
+    vout = allocator.alloc_sbuf_tensor(
+        shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
+    # Q staging via private_hbm, one per LOCAL head (vendor segmented-kernel pattern)
+    q_priv = []
+    for _lh in range(local_heads):
+        q_priv.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
+    # ---- producer: per 512-token tile, straight through the pipeline ----------------
+    for tt in range(n_tt):
+        t0 = tt * _PSUM_FMAX
+        # N>2048: fill this token-block's x_t / tab slice NOW (rolling ring window). At
+        # N<=2048 the whole buffer was already filled up front (bit-for-bit), so skip.
+        if _big:
+            for c in range(n_chunks):
+                nisa.dma_transpose(
+                    dst=x_t[c][tt].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1], [1, _PSUM_FMAX]],
+                                      offset=0),
+                    src=h_hbm.ap([[H, _PSUM_FMAX], [1, 1], [1, 1], [1, _P_MAX]],
+                                 offset=t0 * H + c * _P_MAX),
+                )
+            for t_i in range(4):
+                nisa.dma_copy(dst=tab[t_i][tt],
+                              src=tab_hbm[t_i, 0:_P_MAX, t0:t0 + _PSUM_FMAX])
+        # V for all 8 local heads of this token tile.
+        for vsub in range(v_per_tt):
+            ps_v = nl.ndarray((_V_TILE, local_heads * D), dtype=nl.float32,
+                              buffer=nl.psum, address=(0, 7 * _PSUM_BANK))
+            for c in range(n_chunks):
+                nisa.nc_matmul(
+                    dst=ps_v,
+                    stationary=x_t[c][tt][:, vsub * _V_TILE:(vsub + 1) * _V_TILE],
+                    moving=wv_sb[c],
+                    accumulate=(c > 0),
+                )
+            nisa.tensor_tensor(dst=vout, data1=ps_v, data2=bv_bc, op=nl.add)
+            for lh in range(local_heads):
+                if _big:
+                    # N>2048: stream V to HBM (SBUF->HBM dma from vout directly). The pinned
+                    # :276 tensor_copy is the N<=2048 shipping site and is left untouched below.
+                    nisa.dma_copy(
+                        dst=v_hbm[lh][tt * v_per_tt + vsub],
+                        src=vout[:, lh * D:(lh + 1) * D])
+                    continue
+                # ======== SITE :276 -- pinned to Scalar. WAR chain of 8 copies read one vout
+                # before the next vsub overwrites it; Scalar (240 ns) beats GpSimd (375 ns) and
+                # sheds the copy off Vector, the binding engine. ========
+                nisa.tensor_copy(
+                    dst=v_tiles[lh][tt * v_per_tt + vsub],
+                    src=vout[:, lh * D:(lh + 1) * D],
+                    engine=engine.scalar)
+        # Q and K for both local quads, both parities, this token tile.
+        for qk in range(2):
+            for j in range(local_quads):
+                for par in range(2):
+                    g = j * 2 + par
+                    cols = g * _P_MAX
+                    ps = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                    buffer=nl.psum,
+                                    address=(0, (5 + par) * _PSUM_BANK))
+                    for c in range(n_chunks):
+                        nisa.nc_matmul(
+                            dst=ps,
+                            stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
+                            moving=x_t[c][tt][:, 0:_PSUM_FMAX],
+                            accumulate=(c > 0),
+                        )
+                    # ======== SITE :458 -- pinned to Scalar. Largest remaining Vector-resident
+                    # relocatable site (768/core); data=ps is PSUM so GpSimd is illegal, and
+                    # Scalar's 612 ns fee is faster than the 744 ns Vector cell, so the
+                    # zero-lead rope chain consuming proj cannot lengthen. ========
+                    nisa.tensor_scalar(dst=proj[par], data=ps,
+                                       op0=nl.add, operand0=b_sb[qk][g],
+                                       engine=engine.scalar)
+                # rope: the six ops of nki_rope, same order, tables sliced per tile
+                nisa.tensor_tensor(dst=acc, data1=proj[0],
+                                   data2=tab[0][tt][:, 0:_PSUM_FMAX], op=nl.multiply)
+                nisa.tensor_tensor(dst=tmp, data1=proj[1],
+                                   data2=tab[2][tt][:, 0:_PSUM_FMAX], op=nl.multiply)
+                nisa.tensor_tensor(dst=roped[0], data1=acc, data2=tmp, op=nl.add)
+                nisa.tensor_tensor(dst=acc, data1=proj[1],
+                                   data2=tab[1][tt][:, 0:_PSUM_FMAX], op=nl.multiply)
+                nisa.tensor_tensor(dst=tmp, data1=proj[0],
+                                   data2=tab[3][tt][:, 0:_PSUM_FMAX], op=nl.multiply)
+                nisa.tensor_tensor(dst=roped[1], data1=acc, data2=tmp, op=nl.add)
+                # scatter: head b of the quad owns rows [b*32,(b+1)*32)
+                for b in range(quad):
+                    lh = j * quad + b
+                    r0 = b * half
+                    if qk == 1:
+                        for par2 in range(2):
+                            if _big:
+                                # N>2048: stream K to HBM (SBUF->HBM dma). The :483 tensor_copy
+                                # below is the N<=2048 shipping site, left untouched.
+                                nisa.dma_copy(
+                                    dst=k_hbm[lh][tt][par2 * half:(par2 + 1) * half,
+                                                      0:_K_TILE],
+                                    src=roped[par2][r0:r0 + half, 0:_PSUM_FMAX])
+                                continue
+                            # SITE :483 -- K-scatter COPY, left UNPINNED (no affordable landing
+                            # engine: it is a WAR chain against the next iteration's rope write).
+                            nisa.tensor_copy(
+                                dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half,
+                                                    0:_K_TILE],
+                                src=roped[par2][r0:r0 + half, 0:_PSUM_FMAX])
+                    else:
+                        for par2 in range(2):
+                            nisa.dma_copy(
+                                dst=q_priv[lh][0, par2 * half:(par2 + 1) * half,
+                                               t0:t0 + _PSUM_FMAX],
+                                src=roped[par2][r0:r0 + half, 0:_PSUM_FMAX])
+    # ---- PHASE 2: reclaim the dead producer buffers. NO ALIASING. -------------------
+    allocator.set_current_address(phase2_base)
+    # ---- out-projection weights: THIS core's feature rows, all H output columns -------
+    wo_sb = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, H), dtype=nl.bfloat16,
+        block_dim=[n_ochunks], num_free_tiles=[n_ochunks])
+    for oc in range(n_ochunks):
+        nisa.dma_copy(dst=wo_sb[oc],
+                      src=wot_hbm[col0 + oc * _P_MAX:col0 + (oc + 1) * _P_MAX, 0:H])
+    # bias already HALVED on the host (pack_out_weights) so the two cores' partials sum to
+    # exactly one bias. Broadcast across partitions (output columns are on the free axis).
+    bo_bc = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+    nisa.dma_copy(dst=bo_bc,
+                  src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
+    # ---- O STAYS IN SBUF, ALLOCATED EXACTLY ONCE. tp_out=True gives [D, 128 tokens] per Q
+    # group (feature-major), which the projection contracts over, so O never transposes.
+    # The list handed to the vendor is SLICES into o_feat, so _write_back_o_impl copies
+    # straight into final position. ------------------------------------------------------
+    n_grp = N // _P_MAX
+    o_feat = allocator.alloc_sbuf_tensor(
+        shape=(_P_MAX, N), dtype=nl.bfloat16,
+        block_dim=[n_ochunks], num_free_tiles=[n_ochunks], align_to=32)
+    o_sb = []
+    for _lh in range(local_heads):
+        oc = (_lh * D) // _P_MAX
+        rr = (_lh * D) % _P_MAX
+        # explicit loop, NOT a comprehension: the NKI tracer rejects comprehensions in a kernel.
+        grp = []
+        for _g in range(n_grp):
+            grp.append(o_feat[oc][rr:rr + D, _g * _P_MAX:(_g + 1) * _P_MAX])
+        o_sb.append(grp)
+    attn_base = allocator.get_current_address()
+    for lh in range(local_heads):
+        if _big:
+            # N>2048: load THIS head's K/V from HBM into the resident staging rings (only one
+            # head live at a time -> K/V cost is one head's 24576 B/lane, not eight heads'
+            # 196608). Then run the SAME single-section flatten the N<=2048 path uses: seqlen_k
+            # =8192 < 10240 threshold => num_sections==1, no multi-section flash combine. The
+            # inner's large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers; O is written
+            # straight into this head's o_feat SBUF slices (o_sb[lh]), exactly as N<=2048.
+            for _i in range(n_ktiles):
+                nisa.dma_copy(dst=k_stage[_i], src=k_hbm[lh][_i])
+            for _i in range(n_vtiles):
+                nisa.dma_copy(dst=v_stage[_i], src=v_hbm[lh][_i])
+            # 1D python-list views [seq_tile] (the flatten indexes k_sb[k_tile_idx_in_section]
+            # and v_sb[v_tile_idx] as flat lists -- same convention as k_tiles[lh]/v_tiles[lh]).
+            k1d = []
+            for _i in range(n_ktiles):
+                k1d.append(k_stage[_i])
+            v1d = []
+            for _i in range(n_vtiles):
+                v1d.append(v_stage[_i])
+            _attention_cte(
+                q_priv[lh], None, None,
+                scale=scale, causal_mask=False,
+                tp_q=False, tp_k=False, tp_out=True,
+                cache_softmax=False,
+                k_cache_sbuf=k1d,
+                v_cache_sbuf=v1d,
+                out_o_hbm=o_sb[lh],
+                init_sbuf_addr=attn_base,
+            )
+        else:
+            # N<=2048: the pinned single-section flatten (h8/h10 engine pins intact, bit-for-bit).
+            _attention_cte(
+                q_priv[lh], None, None,
+                scale=scale, causal_mask=False,
+                tp_q=False, tp_k=False, tp_out=True,
+                cache_softmax=False,
+                k_cache_sbuf=k_tiles[lh],
+                v_cache_sbuf=v_tiles[lh],
+                out_o_hbm=o_sb[lh],
+                init_sbuf_addr=attn_base,
+            )
+        allocator.set_current_address(attn_base)
+    # ---- PARTIAL out projection: tokens on partitions, output columns on the free axis
+    # (chosen by the HBM write: result is token-major, so a row of consecutive output columns
+    # is a 1024-byte contiguous run). 128 calls/layer/core (16 token blocks x 2 col blocks x
+    # 4 chunks). --------------------------------------------------------------------------
+    out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX),
+                                         dtype=nl.bfloat16, align_to=32)
+    for tb in range(N // _P_MAX):
+        for ob in range(H // _PSUM_FMAX):
+            ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
+                                address=(0, 0))
+            for oc in range(n_ochunks):
+                nisa.nc_matmul(
+                    dst=ps_out,
+                    stationary=o_feat[oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
+                    moving=wo_sb[oc][:, ob * _PSUM_FMAX:(ob + 1) * _PSUM_FMAX],
+                    accumulate=(oc > 0),
+                )
+            # bias is per output column (free axis) -> tensor_tensor against the pre-broadcast
+            # tile, not tensor_scalar. Same single eviction pass.
+            nisa.tensor_tensor(dst=out_sb, data1=ps_out,
+                               data2=bo_bc[:, ob * _PSUM_FMAX:(ob + 1) * _PSUM_FMAX],
+                               op=nl.add)
+            # [128 tokens, 512 out cols] -> token-major HBM: contiguous along the free axis.
+            nisa.dma_copy(
+                dst=result.ap(pattern=[[H, _P_MAX], [1, _PSUM_FMAX]],
+                              offset=pid * N * H + tb * _P_MAX * H + ob * _PSUM_FMAX),
+                src=out_sb)
+    return result
+_ATTN_BLOCK = wrap_nki_to_torch_hop(_attention_block_kernel[LNC])
+def _no_grad(*tensors) -> None:
+    for t in tensors:
+        if t is not None and t.requires_grad:
+            raise RuntimeError(
+                "attention_block registers no autograd rule and its packed weights are "
+                "detached caches. Use under torch.no_grad()/eval().")
+@nki_op("vjepa2::attention_block", mutates_args={})
+def attention_block(
+    h: torch.Tensor,
+    wqk: torch.Tensor,
+    bqk: torch.Tensor,
+    tab: torch.Tensor,
+    wvt: torch.Tensor,
+    bv: torch.Tensor,
+    wot: torch.Tensor,
+    bo: torch.Tensor,
+    scale: float,
+    head_dim: int,
+) -> torch.Tensor:
+    """Fused norm1'd-hidden -> QKV+rope -> attention -> fused out-projection PARTIALS.
+    Returns ``[LNC=2, N, H]`` bf16 per-core out-projection partials; K3's residual add sums the
+    two cores. The two Scalar engine pins (V-scatter :276, QK bias-evict :458) are hardcoded.
+    """
+    _no_grad(h, wqk, bqk, tab, wvt, bv, wot, bo)
+    # N>2048 (e.g. 64-frame clips) is NOT supported. The phase-1 memory fixes fit (~106 KiB/lane),
+    # but the single-section vendor inner overflows phase 2 (measured 354912 B/lane vs the 196608
+    # budget: exp_tp_sb/mm1_masked/exp_sb scale with num_large_tiles_per_section=ceil(N/2048)=4).
+    # The shrinkable phase-2 region (78848 B/lane) is smaller than the 158304 overflow, so
+    # streaming O / shrinking o_feat cannot close the gap; re-enabling flash sectioning fits but
+    # is numerically wrong (rel_l2 0.90). Refuse here rather than emit a non-fitting or wrong
+    # kernel. See perf/SBUF_8192_STOPPED2.md. NKI kernels cannot `raise`, so the guard lives here.
+    N = h.shape[0]
+    if N > 2048:
+        raise NotImplementedError(
+            f"attention_block N={N} (>2048, e.g. 64-frame clips) is not supported: the "
+            f"single-section vendor inner does not fit SBUF (phase-2 peak measured 354912 B/lane "
+            f"vs the 196608 budget). Streaming O or shrinking o_feat cannot close the 158304 "
+            f"B/lane gap, and flash sectioning is numerically wrong (rel_l2 0.90). See "
+            f"perf/SBUF_8192_STOPPED2.md. Use N<=2048 (up to 16-frame clips) -- the bit-for-bit "
+            f"shipping path."
+        )
+    return _ATTN_BLOCK(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim)

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

	@@ -0,0 +1,1214 @@

+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#
+# MODIFIED from the AWS Neuron SDK's nkilib ``core/attention/attention_cte.py``
+# (Apache-2.0, Copyright Amazon.com, Inc. or its affiliates). Modifications:
+#   1. Four engine-placement edits (two memsets -> GpSimd, the tp_out eviction
+#      tensor_copy -> Scalar, the list-branch write-back tensor_copy -> GpSimd),
+#      measured worth ~1.19x end-to-end in the V-JEPA2 encoder.
+#   2. Code paths unreachable for this model's configuration removed (causal/SWA
+#      masks, prefix caching, flash sectioning at seq<=10240, sinks, fp8,
+#      sequence packing, CP) -- the executed instruction stream is unchanged.
+# The original supports the general case; use nkilib for anything but this
+# non-causal, single-section, SBUF-KV configuration.
+"""K2 inner attention: the vendor ``_attention_cte`` context-encoding kernel, FLATTENED to
+exactly the branches vjepa2's 4-kernel encoder reaches, with the four h8/h10 engine pins
+applied in place. Self-contained: imports only ``nki``, ``nkilib``, ``torch`` (none here) --
+Self-contained: depends only on nki, nkilib and torch.
+WHAT THIS COMPUTES
+==================
+``Output = softmax(scale * Q @ K^T) @ V`` for ONE head, called 8x per core by the K2 outer
+kernel (``attention_block.py``). The pipeline, in order, per 128-token Q group:
+    load_q  ->  QK^T (MM1) + row-max  ->  update running max  ->  exp (softmax numerator)
+            ->  P@V (MM2)  ->  write-back (reciprocal-normalise, evict to the caller's SBUF O).
+Group software-pipelining overlaps three groups at a time (grp_i does PV+writeback, grp_i+1
+does exp, grp_i+2 does load-Q + QK/max), which is the vendor's Level-4 loop structure.
+PSUM BANK PLAN (this config, tp_out=True, single 2048 section)
+--------------------------------------------------------------
+  * MM1 (QK^T) tiles: banks {0,1,2,3}   (``address=(0, (k_tile_idx % 4) * PSUM_BANK_SIZE)``).
+  * MM2 (P@V) tiles:  banks {4,5,6,7}   (``address=(0, (4 + large_tile_idx % 4) * ...)``).
+  * tp reciprocal / correction transposes: bank {3 or 7} (``((grp%2)*4 + 3) * ...``).
+The producer (K2 outer) owns banks 5/6/7 during its phase; the attention region runs after the
+producer's high-water mark, so there is no live overlap.
+RING / MULTI-BUFFER DEPTHS (from the dp/dp2/dp4 fork lineage; kept verbatim)
+---------------------------------------------------------------------------
+  * ``exp_sb``   group-axis depth 2 (dp2): exp(g+1) can run while transpose(g) still drains.
+  * ``exp_tp_sb`` group-axis depth 8 (dp4): P^T transposes issue as soon as their exp output
+    exists, spreading DMA traffic out of the P.V collision window.
+THE FOUR ENGINE PINS (h8/h10), and WHY each exists (one line each, from PERF_INVENTORY)
+---------------------------------------------------------------------------------------
+  1. ``_scale_reciprocal_write_back_impl`` tp_out PSUM->SBUF evict ``tensor_copy`` -> SCALAR
+     (h8 site 1): the un-pinned copy defaults to Vector, K2's binding engine; Scalar/Vector are
+     the only legal engines (PSUM source excludes GpSimd) and Scalar is the exercised majority.
+  2. ``_write_back_o_impl`` tp_out list-branch ``tensor_copy`` -> GPSIMD (h8 site 3): SBUF->SBUF
+     fp32->bf16 CAST, all three engines legal, GpSimd is the idle engine so the copy sheds off
+     Vector.
+  3. ``_qk_and_max_impl`` ``memset`` -> GPSIMD (h10 :2876): initialises mm1_partial_max, first
+     consumer is several ops downstream, so a slower idle engine is safe and sheds Vector.
+  4. ``_fused_qkmax_and_pv_impl`` ``memset`` -> GPSIMD (h10 :3294): same rationale, 2688/core.
+The vendor ``_exp_impl`` memset (dp4 :2993) stays UNPINNED (engine.unknown / Vector) on purpose
+-- pinning it would fork ``_exp_impl`` and re-home the softmax-exp source strings; it MUST stay
+unpinned to reproduce the shipping instruction stream.
+DEAD PATHS REMOVED (§4 of PERF_INVENTORY) -- guard that made each unreachable for this call:
+``causal_mask=False, use_sbuf_kv=True, seqlen_k=2048<10240, bs=1, tp_out=True, fuse_v_proj=False,
+no prior/APC/SWA/CP/seqpack/sink, cache_softmax=False``. See per-function docstrings.
+The allocator usage (``nkilib`` ``ModularAllocator``) is kept EXACTLY as the vendor uses it.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, List, Optional
+import nki
+import nki.isa as nisa
+import nki.language as nl
+from nki.isa import engine, reduce_cmd
+from nki.language.opcode import maximum as _maximum
+from nkilib.core.utils.allocator import align_to
+from nkilib.core.utils.kernel_assert import assert_shape, kernel_assert
+from nkilib.core.utils.kernel_helpers import (
+    PSUM_BANK_SIZE,
+    div_ceil,
+    get_verified_program_sharding_info,
+    is_sbuf_tensor,
+)
+from nkilib.core.utils.logging import get_logger
+from nkilib.core.utils.modular_allocator import ModularAllocator
+from nkilib.core.utils.stream_shuffle_broadcast import stream_shuffle_broadcast
+logger = get_logger("attention_inner")
+_FLOAT32_MIN = -3.4028235e38  # used for initialization and masking
+# Tested-range constraints (checked in _attention_cte).
+_MAX_BS = 512
+_MAX_SEQLEN = 131072
+_MAX_BS_TIMES_SEQLEN_QK = 32.0 * 36864 * 36864
+_MAX_HEAD_DIM = 128
+_MIN_GLOBAL_CP_DEGREE = 1
+_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)
+_FLASH_ATTENTION_THRESHOLD = 10 * 1024  # Use flash attention above this K/V length
+_FLASH_ATTENTION_SECTION_LENGTH = 8 * 1024  # Section size when using flash attention
+# ==========================================================================================
+# Buffer / config dataclasses (kept verbatim -- the four required by the reached path).
+# ==========================================================================================
+@dataclass
+class AttnConfig(nl.NKIObject):
+    """High-level attention configuration set at kernel entry point."""
+    seqlen_q: int = None
+    seqlen_k_active: int = None
+    seqlen_k_prior: int = None
+    d: int = None
+    tp_q: bool = None
+    tp_k: bool = None
+    tp_out: bool = None
+    is_prefix_caching: bool = None
+    causal_mask: bool = None
+    use_swa: bool = None
+    sliding_window: int = None
+    use_cp: bool = None
+    global_cp_deg: int = None
+    cp_strided_q_slicing: bool = None
+    cp_striped_input: bool = None
+    scale: float = None
+    cache_softmax: bool = None
+    skip_output_normalization: bool = None
+    dtype: Any = None
+    softmax_dtype: Any = None
+    mm_out_dtype: Any = None
+    is_sequence_packed: bool = None
+@dataclass
+class AttnInternalBuffers(nl.NKIObject):
+    """Container for all SBUF and PSUM tensor buffers used in attention computation."""
+    q_sb = None
+    k_sb = None
+    v_sb = None
+    k_sb_prior = None
+    v_sb_prior = None
+    k_scale_sb = None
+    mm1_psum = None
+    mm1_copy_sb = None
+    mm1_affine_select_output = None
+    mm1_masked = None
+    mm1_partial_max = None
+    mm1_section_max = None
+    mm1_running_max = None
+    prev_mm1_running_max = None
+    flash_attn_correction_factor = None
+    exp_sb = None
+    exp_partial_sum = None
+    exp_section_sum = None
+    exp_tp_sb = None
+    exp_running_sum = None
+    prev_exp_running_sum = None
+    exp_sum_reciprocal = None
+    mm2_psum = None
+    mm2_sb = None
+    mm2_prev_output = None
+    mm2_accum_flash_attn = None
+    mm2_final = None
+    tp_flash_attn_correction_factor_psum = None
+    tp_flash_attn_correction_factor_sb = None
+    tp_exp_sum_reciprocal_psum = None
+    tp_exp_sum_reciprocal_sb = None
+    mm2_prev_output_scaled = None
+    zero_bias_tensor = None
+    sink_sb = None
+    range_sel_lbs = None
+    range_sel_ubs = None
+    range_sel_lbs_prior = None
+    range_sel_ubs_prior = None
+    k_offset_sb_u32 = None
+@dataclass
+class SectionParams(nl.NKIObject):
+    section_idx = None
+    section_offset = None
+    section_offset_active = None
+    next_section_offset_active = None
+    section_contains_prefix = None
+    next_section_contains_prefix = None
+    kv_section_idx = None
+@dataclass
+class AttnTileParams(nl.NKIObject):
+    """Tile and buffer sizing parameters computed during implementation."""
+    seqlen_k_active_updated: int = None
+    seqlen_k_prior_padded: int = None
+    is_causal: bool = None
+    sb_p: int = None
+    num_grps: int = None
+    num_q_grps_per_load: int = None
+    can_pack_q_load: bool = None
+    num_large_tiles_per_section: int = None
+    num_k_tiles_per_section: int = None
+    num_v_tiles_per_section: int = None
+    exp_inst_elems: int = None
+    num_exp_insts_per_large_tile: int = None
+    num_tps_in_mm2_grp: int = None
+    mm2_grp_sz: int = None
+    use_swa_optimized_allocation: bool = None
+    dynamic_sel_mask: bool = None
+    section_len: int = None
+    num_sections: int = None
+# ==========================================================================================
+# Entry: _attention_cte. Dead paths removed: prefix caching / APC / sink / seqpack / CP / SWA /
+# fused-v-proj / seqlen-q sharding remainder / cache_softmax. This call ALWAYS takes
+# use_sbuf_kv=True, causal_mask=False, tp_out=True, single-shard (num_shard=1), bs=1.
+# ==========================================================================================
+def _attention_cte(
+    q: nl.ndarray,
+    k: Optional[nl.ndarray] = None,
+    v: Optional[nl.ndarray] = None,
+    scale: float = 1.0,
+    causal_mask: bool = True,
+    tp_q: bool = True,
+    tp_k: bool = False,
+    tp_out: bool = False,
+    cache_softmax: bool = False,
+    softmax_dtype=nl.float32,
+    mm_out_dtype=nl.float32,
+    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
+    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
+    out_o_hbm=None,
+    init_sbuf_addr: int = 0,
+):
+    """One head of context-encoding attention (SBUF-KV path).
+    LIVE contract for the K2 encoder: ``k=v=None`` (K/V handed in via ``k_cache_sbuf`` /
+    ``v_cache_sbuf``), ``causal_mask=False``, ``tp_q=False, tp_k=False, tp_out=True``,
+    ``cache_softmax=False``, ``out_o_hbm`` is a LIST of SBUF slices into the caller's O tile.
+    DEAD (removed): prefix caching (k_prior/v_prior/prior_used_len), sink, sliding_window, CP
+    (cp_offset/global_cp_deg/striped/strided), sequence packing (bound_min/bound_max),
+    skip_output_normalization, fp8 k_scale, and the fused v-projection (h_hbm/wv_t/bv).
+    """
+    is_prefix_caching = False
+    use_sbuf_kv = (k_cache_sbuf is not None and len(k_cache_sbuf) > 0) or (
+        v_cache_sbuf is not None and len(v_cache_sbuf) > 0
+    )
+    kernel_assert(
+        k_cache_sbuf is not None and len(k_cache_sbuf) > 0, "k_cache_sbuf must be provided when using sbuf kv"
+    )
+    kernel_assert(
+        v_cache_sbuf is not None and len(v_cache_sbuf) > 0, "v_cache_sbuf must be provided when using sbuf kv"
+    )
+    seqlen_q, seqlen_k_active, seqlen_k_prior, d, out_shape, softmax_shape = _check_input_and_return_shape(
+        q,
+        k,
+        v,
+        is_prefix_caching,
+        None,
+        None,
+        None,
+        tp_q,
+        tp_k,
+        tp_out,
+        cache_softmax,
+        k_cache_sbuf,
+        v_cache_sbuf,
+        None,
+        None,
+        skip_v_check=False,
+    )
+    # Use out_o_hbm if provided (LIVE: it is a list of SBUF slices).
+    if out_o_hbm is not None:
+        result = out_o_hbm
+    else:
+        result = nl.ndarray(shape=out_shape, dtype=q.dtype, buffer=nl.shared_hbm)
+    bs = q.shape[0]
+    bs_kv = 1  # sbuf-kv: batch_size_kv is per-batch tiles
+    kernel_assert(bs > 0, f"Batch size must be positive, got {bs}")
+    kernel_assert(bs <= _MAX_BS, f"attention_cte kernel is not tested for batch size above {_MAX_BS}, got {bs}.")
+    kernel_assert(bs_kv > 0, f"Batch size must be positive, got {bs_kv}")
+    kernel_assert(bs % bs_kv == 0, f"Q batch size must be a multiple of KV batch size, got {bs=}, {bs_kv=}")
+    seqlen_k_total = seqlen_k_active
+    kernel_assert(seqlen_q <= _MAX_SEQLEN, f"seqlen above {_MAX_SEQLEN}, got {seqlen_q=}.")
+    kernel_assert(seqlen_k_total <= _MAX_SEQLEN, f"seqlen above {_MAX_SEQLEN}, got {seqlen_k_total=}.")
+    bs_seqlen_qk_product = float(bs * seqlen_q) * seqlen_k_total
+    if bs_seqlen_qk_product > _MAX_BS_TIMES_SEQLEN_QK:
+        logger.warn(
+            f"attention_cte kernel is not tested for batch size x seqlen_q x seqlen_k above "
+            f"{_MAX_BS_TIMES_SEQLEN_QK}, got {bs_seqlen_qk_product=}.",
+        )
+    kernel_assert(d > 0, f"d must be > 0, got {d=}.")
+    kernel_assert(d <= _MAX_HEAD_DIM, f"we do not support head_dim > {_MAX_HEAD_DIM}, got head dim {d}")
+    kernel_assert(
+        (str(mm_out_dtype) == str(nl.float32))
+        or (str(mm_out_dtype) == str(nl.bfloat16) and (nisa.get_nc_version() >= nisa.nc_version.gen4)),
+        f"mm_out_dtype (psum) should be in [float32, bfloat16], got {mm_out_dtype} in {nisa.get_nc_version()}.",
+    )
+    ac = AttnConfig(
+        seqlen_q=seqlen_q,
+        seqlen_k_active=seqlen_k_active,
+        seqlen_k_prior=seqlen_k_prior,
+        d=d,
+        tp_q=tp_q,
+        tp_k=tp_k,
+        tp_out=tp_out,
+        is_prefix_caching=is_prefix_caching,
+        causal_mask=causal_mask,
+        use_swa=False,
+        sliding_window=0,
+        use_cp=False,
+        global_cp_deg=None,
+        cp_strided_q_slicing=False,
+        cp_striped_input=False,
+        scale=scale,
+        cache_softmax=cache_softmax,
+        skip_output_normalization=False,
+        dtype=q.dtype,
+        softmax_dtype=softmax_dtype,
+        mm_out_dtype=mm_out_dtype,
+        is_sequence_packed=False,
+    )
+    # Sharding info. use_sbuf_kv disables internal sharding (caller already sharded), so
+    # num_shard=1, shard_id=0, num_bs_per_shard=bs, has_remainder=False.
+    get_verified_program_sharding_info("attention_cte", max_sharding=2)
+    num_shard = 1
+    shard_id = 0
+    num_bs_per_shard = bs // num_shard
+    bs_offset = shard_id * num_bs_per_shard
+    for batch_idx in range(num_bs_per_shard):
+        # use_sbuf_kv: bs is LOCAL batch size and b is the LOCAL index.
+        batch_id = batch_idx
+        kv_batch_id = _q_to_kv_batch_id(batch_id, bs, bs_kv)
+        _attention_cte_impl(
+            q,
+            k,
+            v,
+            None,
+            None,
+            None,
+            result,
+            batch_id,
+            kv_batch_id,
+            ac,
+            k_cache_sbuf=k_cache_sbuf,
+            v_cache_sbuf=v_cache_sbuf,
+            init_sbuf_addr=init_sbuf_addr,
+        )
+    # has_remainder is False (num_shard==1); the seqlen-q sharding path is DEAD.
+    return result
+def _attention_cte_impl(
+    q,
+    k_active,
+    v_active,
+    k_prior,
+    v_prior,
+    prior_used_len,
+    o,
+    batch_id,
+    batch_id_kv,
+    ac: AttnConfig,
+    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
+    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
+    init_sbuf_addr: int = 0,
+):
+    """Single-batch attention body. LIVE path only: single section (num_sections==1), pipelined
+    group loop (num_grps=16 > 1), SBUF-KV (no K/V HBM load), tp_out=True.
+    DEAD (removed): seqlen-q sharding, sink, SWA memset, sections>1 flash accumulation, K/V HBM
+    load (_load_k_tile / _load_v_tile), SBUF-prior storage, cache_softmax intermediate writes,
+    the non-pipelined (shard_seqlen_q_length<=1) branch.
+    """
+    # Not seqlen-sharded: start at group 0, cover all groups.
+    atp = _compute_tile_parameters(ac, is_seqlen_sharded=False)
+    shard_seqlen_q_start = 0
+    shard_seqlen_q_length = atp.num_grps
+    allocator = ModularAllocator(initial_address=init_sbuf_addr)
+    bufs = AttnInternalBuffers()
+    bufs.k_scale_sb = None
+    bufs.zero_bias_tensor = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, 1), dtype=nl.float32)
+    nisa.memset(bufs.zero_bias_tensor, 0.0)
+    # _setup_range_select_bounds is a NO-OP here: dynamic_sel_mask=False and not prefix caching.
+    _setup_range_select_bounds(ac, atp, bufs, allocator, None, None, None, None, batch_id=batch_id)
+    bufs.mm1_running_max = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)
+    bufs.exp_running_sum = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)
+    bufs.exp_sum_reciprocal = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)
+    sbuf_addr_outer = allocator.get_current_address()
+    # Single section (num_sections == 1).
+    for section_idx in range(atp.num_sections):
+        section_offset = atp.section_len * section_idx
+        section_offset_active, section_contains_prefix = _compute_section_offset_active(
+            section_offset, ac.is_prefix_caching, atp.seqlen_k_prior_padded
+        )
+        next_section_offset = atp.section_len * (section_idx + 1)
+        next_section_offset_active, next_section_contains_prefix = _compute_section_offset_active(
+            next_section_offset, ac.is_prefix_caching, atp.seqlen_k_prior_padded
+        )
+        sp = SectionParams(
+            section_idx=section_idx,
+            section_offset=section_offset,
+            section_offset_active=section_offset_active,
+            next_section_offset_active=next_section_offset_active,
+            section_contains_prefix=section_contains_prefix,
+            next_section_contains_prefix=next_section_contains_prefix,
+        )
+        allocator.set_current_address(sbuf_addr_outer)
+        _allocate_attention_buffers(allocator, ac, atp, bufs, sink=None,
+                                    k_cache_sbuf=k_cache_sbuf, v_cache_sbuf=v_cache_sbuf)
+        sbuf_addr = allocator.get_current_address()
+        # SBUF-KV: K and V are already resident, so the HBM K/V load block is skipped entirely.
+        # Software pipelining: num_grps=16 > 1, so the pipelined branch is the only reached one.
+        _load_q_impl(shard_seqlen_q_start, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
+        _qk_and_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, batch_id)
+        _update_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)
+        _exp_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)
+        _load_q_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
+        _qk_and_max_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, batch_id)
+        _update_max_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, None)
+        shard_seqlen_q_end = shard_seqlen_q_start + shard_seqlen_q_length
+        for grp_i in range(shard_seqlen_q_start, shard_seqlen_q_end - 2):
+            # grp_i: PV, write_back ; grp_i+1: EXP ; grp_i+2: load Q, QK+Max.
+            _load_q_impl(grp_i + 2, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
+            _exp_impl(grp_i + 1, ac, atp, sp, bufs, None)
+            _fused_qkmax_and_pv_impl(grp_i, ac, atp, sp, bufs, batch_id)
+            _write_back_impl(grp_i, ac, atp, sp, bufs, o, batch_id)
+            _update_max_impl(grp_i + 2, ac, atp, sp, bufs, None)
+        _pv_impl(shard_seqlen_q_end - 2, ac, atp, sp, bufs)
+        _write_back_impl(shard_seqlen_q_end - 2, ac, atp, sp, bufs, o, batch_id)
+        _exp_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs, None)
+        _pv_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs)
+        _write_back_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs, o, batch_id)
+    # cache_softmax intermediate writes (out_neg_max / out_sum_recip) are DEAD (cache_softmax=False).
+def _compute_tile_parameters(ac: AttnConfig, is_seqlen_sharded: bool) -> AttnTileParams:
+    """Tile parameters for the LIVE config. Dead: SWA/prefix/CP scale checks and seqlen updates,
+    flash sectioning (total_seqlen_k=2048 <= 10240 -> single section)."""
+    atp = AttnTileParams()
+    # LIVE: causal_mask=False, no CP/SWA/seqpack -> is_causal=False, dynamic_sel_mask=False.
+    atp.is_causal = ac.causal_mask
+    atp.dynamic_sel_mask = False
+    atp.seqlen_k_active_updated = ac.seqlen_k_active
+    atp.use_swa_optimized_allocation = False
+    atp.sb_p = nl.tile_size.pmax
+    kernel_assert(_Q_GRP_SZ == atp.sb_p, f"expect Q group size == pmax, got {_Q_GRP_SZ=}, {atp.sb_p=}")
+    kernel_assert(_V_TILE_SZ == atp.sb_p, f"expect V tile size == pmax, got {_V_TILE_SZ=}, {atp.sb_p=}")
+    atp.num_grps = div_ceil(ac.seqlen_q, atp.sb_p)
+    atp.can_pack_q_load = not is_seqlen_sharded
+    num_q_grps_per_load_dtype = 4 if ac.dtype == nl.float32 else 8
+    atp.num_q_grps_per_load = min(num_q_grps_per_load_dtype if atp.can_pack_q_load else 1, atp.num_grps)
+    kernel_assert(atp.num_q_grps_per_load > 0, f"num_q_grps_per_load must be positive, got {atp.num_q_grps_per_load}.")
+    atp.seqlen_k_prior_padded = None  # not prefix caching
+    total_seqlen_k = atp.seqlen_k_active_updated
+    # total_seqlen_k = 2048 <= 10240 -> no flash attention, single section.
+    use_flash_attn = total_seqlen_k > _FLASH_ATTENTION_THRESHOLD
+    if use_flash_attn:
+        atp.section_len = min(total_seqlen_k, _FLASH_ATTENTION_SECTION_LENGTH)
+    else:
+        atp.section_len = total_seqlen_k
+    kernel_assert(atp.section_len > 0, f"section_len must be positive, got {atp.section_len}")
+    atp.num_sections = div_ceil(total_seqlen_k, atp.section_len)
+    if not use_flash_attn:
+        kernel_assert(atp.num_sections == 1, "Logic fault, must only have 1 section if not using flash_attn")
+    atp.num_large_tiles_per_section = div_ceil(atp.section_len, _LARGE_TILE_SZ)
+    atp.num_k_tiles_per_section = div_ceil(atp.section_len, _K_TILE_SZ)
+    atp.num_v_tiles_per_section = div_ceil(atp.section_len, _V_TILE_SZ)
+    atp.exp_inst_elems = _EXP_TILE_SZ
+    atp.num_exp_insts_per_large_tile = _LARGE_TILE_SZ // atp.exp_inst_elems
+    atp.num_tps_in_mm2_grp = _K_TILE_SZ // atp.sb_p
+    atp.mm2_grp_sz = _K_TILE_SZ
+    return atp
+def _setup_range_select_bounds(
+    ac: AttnConfig,
+    atp: AttnTileParams,
+    bufs: AttnInternalBuffers,
+    allocator: ModularAllocator,
+    cp_offset: Any,
+    prior_used_len: Any,
+    bound_min: Any,
+    bound_max: Any,
+    batch_id: int = 0,
+) -> tuple:
+    """NO-OP for this config: only CP / SWA / prefix-caching / sequence-packing populate range
+    bounds, and all are off (``dynamic_sel_mask=False``, ``is_prefix_caching=False``). The entire
+    vendor body was gated by ``if atp.dynamic_sel_mask:`` / ``if ac.is_prefix_caching:`` / SWA,
+    so it emits no instructions here and is removed."""
+    return
+def _allocate_attention_buffers(
+    allocator: ModularAllocator,
+    ac: AttnConfig,
+    atp: AttnTileParams,
+    bufs: AttnInternalBuffers,
+    sink: Any,
+    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
+    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
+):
+    """Allocate SBUF/PSUM buffers. LIVE: tp_out=True (so all tp_out branches taken), K/V come
+    from k_cache_sbuf/v_cache_sbuf, sink=None, dynamic_sel_mask=False (so the static-mask
+    mm1_copy_sb / mm1_affine_select_output buffers ARE allocated -- kept). The tp_out=False
+    branches and the fused-v-proj branch are DEAD and removed."""
+    mm2_p, mm2_n = atp.sb_p, ac.d
+    # K/V provided in SBUF by the caller.
+    bufs.k_sb = k_cache_sbuf
+    bufs.v_sb = v_cache_sbuf
+    bufs.q_sb = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p * atp.num_q_grps_per_load),
+        dtype=nl.bfloat16,
+        block_dim=[div_ceil(atp.num_grps, atp.num_q_grps_per_load)],
+        num_free_tiles=[2],
+        align_to=32,
+    )
+    bufs.flash_attn_correction_factor = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    mm1_partial_max_n_elts = atp.num_k_tiles_per_section  # sink is None
+    bufs.mm1_partial_max = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, mm1_partial_max_n_elts), dtype=nl.float32,
+        block_dim=[atp.num_grps], num_free_tiles=[2], align_to=4,
+    )
+    bufs.mm1_section_max = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    n_final_reduce_sum_elts = div_ceil(atp.section_len, atp.exp_inst_elems)  # sink is None
+    bufs.exp_partial_sum = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, n_final_reduce_sum_elts), dtype=nl.float32,
+        block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.exp_section_sum = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.prev_mm1_running_max = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.prev_exp_running_sum = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    # tp_out=True everywhere below.
+    bufs.mm2_prev_output = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.mm2_accum_flash_attn = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.mm2_prev_output_scaled = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.tp_flash_attn_correction_factor_psum = []
+    bufs.tp_exp_sum_reciprocal_psum = []
+    for grp_idx in range(atp.num_grps):
+        tp_flash_attn_correction_factor_psum_tile = nl.ndarray(
+            (ac.d, atp.sb_p), dtype=nl.float32, buffer=nl.psum,
+            address=(0, ((grp_idx % 2) * 4 + 3) * PSUM_BANK_SIZE),
+        )
+        bufs.tp_flash_attn_correction_factor_psum.append(tp_flash_attn_correction_factor_psum_tile)
+        tp_exp_sum_reciprocal_psum_tile = nl.ndarray(
+            (ac.d, atp.sb_p), dtype=nl.float32, buffer=nl.psum,
+            address=(0, ((grp_idx % 2) * 4 + 3) * PSUM_BANK_SIZE),
+        )
+        bufs.tp_exp_sum_reciprocal_psum.append(tp_exp_sum_reciprocal_psum_tile)
+    bufs.tp_flash_attn_correction_factor_sb = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[4],
+    )
+    bufs.tp_exp_sum_reciprocal_sb = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[4],
+    )
+    bufs.mm2_final = allocator.alloc_sbuf_tensor(
+        shape=(ac.d, atp.sb_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.mm2_sb = allocator.alloc_sbuf_tensor(
+        shape=(mm2_n, mm2_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
+    )
+    bufs.mm1_masked = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, _LARGE_TILE_SZ), dtype=nl.float32,
+        block_dim=[atp.num_grps, atp.num_large_tiles_per_section],
+        num_free_tiles=[2, atp.num_large_tiles_per_section],
+    )
+    bufs.exp_sb = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, _LARGE_TILE_SZ), dtype=nl.bfloat16,
+        block_dim=[atp.num_grps, atp.num_large_tiles_per_section],
+        # DP2: group-axis depth 2 so exp(g+1) runs while transpose(g) still drains exp_sb(g).
+        num_free_tiles=[2, atp.num_large_tiles_per_section],
+    )
+    # MM1 PSUM tiles: banks {0,1,2,3}.
+    bufs.mm1_psum = []
+    for grp_idx in range(atp.num_grps):
+        grp_row = []
+        for large_tile_idx in range(atp.num_large_tiles_per_section):
+            tile_row = []
+            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)
+        bufs.mm1_psum.append(grp_row)
+    # dynamic_sel_mask=False -> the static-mask staging buffers are allocated (no-mask branch
+    # does not read them, but the buffers exist in the vendor's non-dynamic allocation).
+    bufs.mm1_copy_sb = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype,
+        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, 4], num_free_tiles=[1, 1, 2],
+    )
+    bufs.mm1_affine_select_output = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype,
+        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, 4], num_free_tiles=[1, 1, 2],
+    )
+    bufs.exp_tp_sb = allocator.alloc_sbuf_tensor(
+        shape=(atp.sb_p, atp.mm2_grp_sz), dtype=nl.bfloat16,
+        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, atp.num_tps_in_mm2_grp],
+        # DP4: 8 group-blocks resident so P^T transposes issue as soon as their exp exists.
+        num_free_tiles=[8, atp.num_large_tiles_per_section, atp.num_tps_in_mm2_grp],
+        align_to=32,
+    )
+    # MM2 PSUM tiles: banks {4,5,6,7}. tp_out=True shape (mm2_n, mm2_p).
+    bufs.mm2_psum = []
+    for grp_idx in range(atp.num_grps):
+        grp_row = []
+        for large_tile_idx in range(atp.num_large_tiles_per_section):
+            mm2_psum_tile = nl.ndarray(
+                (mm2_n, mm2_p), dtype=ac.mm_out_dtype, buffer=nl.psum,
+                address=(0, ((4 + (large_tile_idx % 4)) * PSUM_BANK_SIZE)),
+            )
+            grp_row.append(mm2_psum_tile)
+        bufs.mm2_psum.append(grp_row)
+def _q_to_kv_batch_id(batch_id: int, bs: int, bs_kv: int) -> int:
+    """Map Q batch id to KV batch id for native GQA support. bs_kv=1 here -> always 0."""
+    return batch_id // (bs // bs_kv)
+def _check_input_and_return_shape(
+    q,
+    k,
+    v,
+    is_prefix_caching,
+    k_prior,
+    v_prior,
+    prior_used_len,
+    tp_q,
+    tp_k,
+    tp_out,
+    cache_softmax,
+    k_cache_sbuf=None,
+    v_cache_sbuf=None,
+    k_prior_sbuf=None,
+    v_prior_sbuf=None,
+    skip_v_check=False,
+) -> tuple:
+    """Validate shapes and return (seqlen_q, seqlen_k, seqlen_k_prior, d, out_shape,
+    cache_softmax_shape). LIVE: use_sbuf_kv=True, tp_q=False, tp_out=True, not prefix caching,
+    cache_softmax=False."""
+    # SBUF-KV path.
+    kernel_assert(k_cache_sbuf is not None and len(k_cache_sbuf) > 0, "k_cache_sbuf must be provided")
+    kernel_assert(v_cache_sbuf is not None and len(v_cache_sbuf) > 0, "v_cache_sbuf must be provided")
+    if tp_q:
+        batch_size, seqlen_q, d = q.shape
+    else:
+        batch_size, d, seqlen_q = q.shape
+    d_from_k = k_cache_sbuf[0].shape[0]
+    kernel_assert(d == d_from_k, f"Head dimension mismatch: q has d={d}, but k_cache_sbuf has d={d_from_k}")
+    seqlen_k = len(k_cache_sbuf) * _K_TILE_SZ
+    batch_size_kv = 1
+    seqlen_k_prior = None  # not prefix caching
+    out_seqlen = seqlen_q
+    if tp_out:
+        out_shape = (batch_size, d, out_seqlen)
+    else:
+        out_shape = (batch_size, out_seqlen, d)
+    cache_softmax_tile_size = 128
+    padded_seq_grps = div_ceil(out_seqlen, cache_softmax_tile_size)
+    cache_softmax_shape = [batch_size, cache_softmax_tile_size, padded_seq_grps]
+    return (seqlen_q, seqlen_k, seqlen_k_prior, d, out_shape, cache_softmax_shape)
+def _load_q_tile(q, out, tp_q, batch_id, grp_i, seqlen_offset, load_dtype, sbuf_addr, grps_per_load=1) -> int:
+    """Load a Q tile HBM->SBUF. LIVE: tp_q=False (q is [bs, d, seqlen]), no transpose; the whole
+    ``if tp_q:`` transpose sub-tree (dma_transpose / nc_transpose paths) is DEAD and removed."""
+    kernel_assert(str(load_dtype) == str(out[0].dtype), "Conflicting dtype")
+    _, d, seqlen = q.shape
+    num_f = min(seqlen - seqlen_offset, _Q_GRP_SZ * grps_per_load)
+    out_dst_pat = out[grp_i // grps_per_load].ap(pattern=[[_Q_GRP_SZ * grps_per_load, d], [1, num_f]], offset=0)
+    q_src_pat = q.ap(pattern=[[seqlen, d], [1, num_f]], offset=batch_id * d * seqlen + seqlen_offset)
+    nisa.dma_copy(dst=out_dst_pat, src=q_src_pat)
+def _get_kv_tile_apc(
+    is_prefix_caching,
+    k_active,
+    k_prior,
+    seqlen_active,
+    seqlen_prior,
+    seqlen_offset,
+    load_offset_active,
+) -> tuple:
+    """Determine which KV tensor (active or prior) to use. LIVE: is_prefix_caching=False, so this
+    returns the trivial active-tile mapping; the prefix-caching branch is DEAD and removed."""
+    return k_active, seqlen_active, seqlen_offset, load_offset_active
+def _compute_section_offset_active(k_section_offset, is_prefix_caching, seqlen_k_prior_padded):
+    """Active-K offset for a section. LIVE: is_prefix_caching=False -> pass-through."""
+    section_contains_prefix = False
+    k_section_offset_active = k_section_offset
+    return k_section_offset_active, section_contains_prefix
+def _load_q_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+                 bufs: AttnInternalBuffers, q, batch_id, sbuf_addr):
+    """Load a Q group if compute is needed. LIVE: is_causal=False -> has_any_compute is always
+    True; k_scale_sb is None (fp8 dequant path DEAD)."""
+    if grp_i % atp.num_q_grps_per_load == 0:
+        # is_causal=False -> has_any_compute_pred is True.
+        q_seqlen_offset = grp_i * _Q_GRP_SZ
+        _load_q_tile(
+            q, bufs.q_sb, ac.tp_q, batch_id, grp_i, q_seqlen_offset,
+            bufs.q_sb[0].dtype, sbuf_addr, grps_per_load=atp.num_q_grps_per_load,
+        )
+        # k_scale_sb is None -> the delayed-fp8 Q scaling is DEAD.
+def _qk_and_max_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+                     bufs: AttnInternalBuffers, batch_id: int = 0):
+    """MM1 (QK^T) + row-max for a Q group. LIVE: is_causal=False -> always compute.
+    h10 EDIT (site :2876): the ``mm1_partial_max`` init memset is pinned to GpSimd -- it sheds
+    Vector (K2's binding engine) and its first consumer is several ops downstream, so a slower
+    idle engine is safe."""
+    # is_causal=False -> has_any_compute_pred is True.
+    nisa.memset(bufs.mm1_partial_max[grp_i], value=_FLOAT32_MIN, engine=engine.gpsimd)
+    for large_tile_idx in range(atp.num_large_tiles_per_section):
+        _qk_and_max_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs, batch_id)
+def _update_max_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+                     bufs: AttnInternalBuffers, sink):
+    """Update running max. LIVE: is_causal=False (always compute), sink=None, num_sections==1 so
+    the reduce writes straight into mm1_running_max (max-direct promotion) and the flash
+    multi-section correction is DEAD."""
+    # is_causal=False -> has_any_compute_pred is True.
+    # num_sections==1: the reduce writes straight into mm1_running_max (max-direct); the [128,1]
+    # copy the vendor once had here is gone (3,072 Vector instructions removed upstream).
+    _max_dst = bufs.mm1_running_max[:, grp_i]
+    nisa.tensor_reduce(_max_dst, nl.maximum, bufs.mm1_partial_max[grp_i], 1, negate=True)
+    # num_sections==1: no running-max copy / flash correction factor. sink block DEAD (sink=None).
+def _exp_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+              bufs: AttnInternalBuffers, sink):
+    """Softmax exp (Scalar ACTIVATE) + P^T dma_transposes. LIVE: is_causal=False, use_swa=False,
+    not prefix caching -> every k-tile is emitted (emit_idxs is full), wide_exp=False, sink=None.
+    The ``memset`` on ``exp_partial_sum`` stays UNPINNED (engine.unknown / Vector) ON PURPOSE
+    (vendor :2993): pinning it would fork this 280-line function and re-home the softmax-exp and
+    dma_transpose source strings, which MUST keep their ``attention_inner`` sites stable and
+    match the shipping stream. The ACCUM FORK (reset on first tile, accumulate, drain once) and
+    the reduce_cmd discipline are preserved verbatim."""
+    # is_causal=False -> has_any_compute_pred is True.
+    q_seqlen_offset = grp_i * atp.sb_p
+    nisa.memset(bufs.exp_partial_sum[grp_i][...], value=0.0)  # UNPINNED on purpose (:2993).
+    for large_tile_idx in range(atp.num_large_tiles_per_section):
+        kernel_assert(atp.exp_inst_elems == 512, "Internal validation failed.")
+        # ACCUM FORK pre-scan: is_causal=False and not prefix caching -> every tile emits, so
+        # emit_idxs is [0, 1, 2, 3]. reset on the first, accumulate, drain once.
+        emit_idxs = []
+        for scan_idx in range(atp.num_exp_insts_per_large_tile):
+            scan_prior, scan_seqlen_k, scan_k_start, _ = _get_kv_tile_apc(
+                ac.is_prefix_caching,
+                False,
+                True,
+                atp.seqlen_k_active_updated,
+                ac.seqlen_k_prior,
+                sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + scan_idx * atp.exp_inst_elems,
+                None,
+            )
+            scan_mask = True  # is_causal=False, use_swa=False
+            if scan_mask and scan_seqlen_k > scan_k_start:
+                emit_idxs.append(scan_idx)
+        first_emit = emit_idxs[0] if emit_idxs else -1
+        last_emit = emit_idxs[-1] if emit_idxs else -1
+        for exp_tile_idx in range(atp.num_exp_insts_per_large_tile):
+            is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
+                ac.is_prefix_caching,
+                False,
+                True,
+                atp.seqlen_k_active_updated,
+                ac.seqlen_k_prior,
+                sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + exp_tile_idx * atp.exp_inst_elems,
+                None,
+            )
+            num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+            num_f = min(seqlen_k - k_start_pos, atp.exp_inst_elems)
+            exp_sel_mask = True  # is_causal=False, use_swa=False
+            if exp_sel_mask and seqlen_k > k_start_pos:
+                # Step 1: exp. reduce_cmd is ALWAYS explicit (see the vendor note: an idle cmd
+                # with reduce_op/reduce_res would write a stale register).
+                if exp_tile_idx == first_emit:
+                    exp_reduce_cmd = nisa.reduce_cmd.reset_reduce
+                else:
+                    exp_reduce_cmd = nisa.reduce_cmd.reduce
+                if exp_tile_idx == last_emit:
+                    exp_reduce_res = bufs.exp_partial_sum[grp_i][
+                        :num_p,
+                        large_tile_idx * atp.num_exp_insts_per_large_tile + exp_tile_idx,
+                    ]
+                else:
+                    exp_reduce_res = None
+                nisa.activation(
+                    bufs.exp_sb[grp_i][large_tile_idx][
+                        :num_p, nl.ds(exp_tile_idx * atp.exp_inst_elems, num_f)
+                    ],
+                    op=nl.exp,
+                    data=bufs.mm1_masked[grp_i][large_tile_idx][
+                        :num_p, nl.ds(exp_tile_idx * atp.exp_inst_elems, num_f)
+                    ],
+                    reduce_op=nl.add,
+                    reduce_cmd=exp_reduce_cmd,
+                    reduce_res=exp_reduce_res,
+                    bias=bufs.mm1_running_max[:num_p, grp_i],
+                )
+                # Step 2: P^T dma_transpose into exp_tp_sb ([128, 4, 128]).
+                num_f_outer = num_f // atp.sb_p
+                num_f_inner = num_f % atp.sb_p
+                if num_f_outer >= 1:
+                    nisa.dma_transpose(
+                        dst=bufs.exp_tp_sb[grp_i][large_tile_idx][exp_tile_idx].ap(
+                            [
+                                [atp.mm2_grp_sz, atp.sb_p],
+                                [1, 1],
+                                [atp.sb_p, num_f_outer],
+                                [1, num_p],
+                            ]
+                        ),
+                        src=bufs.exp_sb[grp_i][large_tile_idx].ap(
+                            [
+                                [_LARGE_TILE_SZ, num_p],
+                                [1, 1],
+                                [atp.sb_p, num_f_outer],
+                                [1, atp.sb_p],
+                            ],
+                            offset=exp_tile_idx * atp.mm2_grp_sz,
+                        ),
+                    )
+                if num_f_inner > 0:
+                    nisa.dma_transpose(
+                        dst=bufs.exp_tp_sb[grp_i][large_tile_idx][exp_tile_idx].ap(
+                            [
+                                [atp.mm2_grp_sz, num_f_inner],
+                                [1, 1],
+                                [atp.sb_p, 1],
+                                [1, num_p],
+                            ],
+                            offset=num_f_outer * atp.sb_p,
+                        ),
+                        src=bufs.exp_sb[grp_i][large_tile_idx].ap(
+                            [
+                                [_LARGE_TILE_SZ, num_p],
+                                [1, 1],
+                                [atp.sb_p, 1],
+                                [1, num_f_inner],
+                            ],
+                            offset=exp_tile_idx * atp.mm2_grp_sz + num_f_outer * atp.sb_p,
+                        ),
+                    )
+    # sink block DEAD (sink=None).
+def _pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams, bufs: AttnInternalBuffers):
+    """MM2 (P@V) for a Q group. LIVE: is_causal=False -> always compute."""
+    # is_causal=False -> has_any_compute_pred is True.
+    nisa.memset(bufs.mm2_sb[grp_i][...], value=0.0)
+    for large_tile_idx in range(atp.num_large_tiles_per_section):
+        _pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)
+def _fused_qkmax_and_pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+                             bufs: AttnInternalBuffers, batch_id: int = 0):
+    """Pipelined MM2 for grp_i while doing QK+max for grp_i+2. LIVE: is_causal=False -> both
+    predicates True.
+    h10 EDIT (site :3294): the ``mm1_partial_max[qkmax_grp]`` init memset is pinned to GpSimd
+    (same rationale as :2876; 2,688/core, the larger of the two)."""
+    qkmax_grp = grp_i + 2
+    # is_causal=False -> both has_any_compute predicates are True.
+    nisa.memset(bufs.mm1_partial_max[qkmax_grp][...], value=_FLOAT32_MIN, engine=engine.gpsimd)
+    for large_tile_idx in range(atp.num_large_tiles_per_section):
+        _pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)
+        _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac, atp, sp, bufs, batch_id)
+def _write_back_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
+                     bufs: AttnInternalBuffers, o, batch_id):
+    """Finalise output. LIVE: is_causal=False (always compute), num_sections==1,
+    skip_output_normalization=False -> single-section reciprocal + scale_reciprocal_write_back.
+    The flash multi-section accumulation and the skip-normalization branches are DEAD."""
+    # is_causal=False -> has_any_compute_pred is True.
+    q_seqlen_offset = grp_i * atp.sb_p
+    nisa.tensor_reduce(bufs.exp_section_sum[grp_i][...], nl.add, bufs.exp_partial_sum[grp_i], axis=1)
+    # num_sections==1, not skip_output_normalization.
+    nisa.reciprocal(bufs.exp_sum_reciprocal[:, grp_i], bufs.exp_section_sum[grp_i])
+    num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+    num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+    # Single section: scale by reciprocal and write current output.
+    _scale_reciprocal_write_back_impl(bufs.mm2_sb[grp_i], grp_i, ac, atp, bufs, o, batch_id, num_p, num_f)
+def _scale_reciprocal_write_back_impl(src_buf, grp_i, ac: AttnConfig, atp: AttnTileParams,
+                                      bufs: AttnInternalBuffers, o, batch_id, num_p, num_f):
+    """Reciprocal-normalise and write back (final section). LIVE: tp_out=True.
+    h8 EDIT (site 1): the tp_out PSUM->SBUF evict ``tensor_copy`` is pinned to Scalar -- the
+    un-pinned copy defaults to Vector (K2's binding engine); GpSimd is illegal (PSUM source), so
+    Scalar is the only shedding target and is the exercised majority path. The tp_out=False
+    ``tensor_scalar(..., engine=engine.vector)`` branch is DEAD (tp_out always True) and removed."""
+    # tp_out=True.
+    nisa.nc_transpose(
+        bufs.tp_exp_sum_reciprocal_psum[grp_i].ap(pattern=[[atp.sb_p, ac.d], [1, atp.sb_p]], offset=0),
+        bufs.exp_sum_reciprocal.ap(pattern=[[atp.num_grps, atp.sb_p], [0, ac.d]], offset=grp_i),
+    )
+    # ---- h8 SITE 1: PSUM->SBUF evict pinned to Scalar. ----
+    nisa.tensor_copy(
+        bufs.tp_exp_sum_reciprocal_sb[grp_i][: ac.d, :num_f],
+        bufs.tp_exp_sum_reciprocal_psum[grp_i][: ac.d, :num_f],
+        engine=engine.scalar,
+    )
+    nisa.tensor_tensor(
+        bufs.mm2_final[grp_i][: ac.d, :num_f],
+        src_buf[: ac.d, :num_f],
+        bufs.tp_exp_sum_reciprocal_sb[grp_i][: ac.d, :num_f],
+        nl.multiply,
+    )
+    _write_back_o_impl(bufs.mm2_final[grp_i], grp_i, ac, atp, o, batch_id, num_p, num_f)
+def _write_back_o_impl(src_buf, grp_i, ac: AttnConfig, atp: AttnTileParams, o, batch_id, num_p, num_f):
+    """Evict finished O. LIVE: ``o`` is a LIST of caller-owned SBUF slices, tp_out=True, so the
+    output is copied into the caller's O tile (kept on-chip for the fused out-projection) rather
+    than DMA'd to HBM. The HBM-sink branch is DEAD (o is always a list) and removed.
+    h8 EDIT (site 3): the tp_out list-branch ``tensor_copy`` (SBUF->SBUF fp32->bf16 CAST) is
+    pinned to GpSimd -- all three engines are legal and GpSimd is the idle engine, so it sheds
+    off Vector."""
+    # o is a list (caller-owned SBUF slices), tp_out=True.
+    # ---- h8 SITE 3: SBUF->SBUF CAST pinned to GpSimd. ----
+    nisa.tensor_copy(dst=o[grp_i][: ac.d, :num_f], src=src_buf[: ac.d, :num_f], engine=engine.gpsimd)
+def _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
+                                sp: SectionParams, bufs: AttnInternalBuffers, batch_id: int = 0):
+    """MM1 (QK^T) + masking for a Q group and a 2048 K tile. LIVE: is_causal=False,
+    dynamic_sel_mask=False, is_prior_tile=False -> the FINAL ``else: # no masking``
+    ``tensor_scalar_reduce`` branch. The diagonal (affine_select) and range_select branches are
+    DEAD and removed; SBUF-prefix-caching K-tile selection is DEAD (k_sb_prior is None)."""
+    q_seqlen_offset = qkmax_grp * atp.sb_p
+    num_k_tiles_in_large_tile = _LARGE_TILE_SZ // _K_TILE_SZ
+    for k_tile_idx in range(num_k_tiles_in_large_tile):
+        mm1_psum_tile = bufs.mm1_psum[qkmax_grp][large_tile_idx][k_tile_idx]
+        # dynamic_sel_mask=False, but the no-mask branch below reads neither mm1_copy_sb nor
+        # mm1_affine_select_output, so their per-tile locals are not needed.
+        mm1_masked_tile = bufs.mm1_masked[qkmax_grp][large_tile_idx]
+        mm1_partial_max_tile = bufs.mm1_partial_max[qkmax_grp]
+        k_tile_idx_in_section = large_tile_idx * num_k_tiles_in_large_tile + k_tile_idx
+        _kv_sec_idx = sp.kv_section_idx if sp.kv_section_idx is not None else sp.section_idx
+        k_tile_idx_global = atp.num_k_tiles_per_section * _kv_sec_idx + k_tile_idx_in_section
+        is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
+            ac.is_prefix_caching,
+            False,
+            True,
+            atp.seqlen_k_active_updated,
+            ac.seqlen_k_prior,
+            k_tile_idx_global * _K_TILE_SZ,
+            None,
+        )
+        matmul_selection = True  # is_causal=False
+        if q_seqlen_offset >= ac.seqlen_q or k_start_pos >= seqlen_k:
+            matmul_selection = False
+        if matmul_selection and k_tile_idx_in_section < atp.num_k_tiles_per_section:
+            num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)
+            num_q_free = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+            # Step 1: MM1 matmul. Non-APC path -> use k_tile_idx_in_section.
+            k_tile_to_use = bufs.k_sb[k_tile_idx_in_section]
+            nisa.nc_matmul(
+                mm1_psum_tile[:num_q_free, :num_f],
+                bufs.q_sb[qkmax_grp // atp.num_q_grps_per_load][
+                    : ac.d,
+                    nl.ds((qkmax_grp % atp.num_q_grps_per_load) * _Q_GRP_SZ, num_q_free),
+                ],
+                k_tile_to_use[:, :num_f],
+            )
+            # Step 2: no masking (is_causal=False, dynamic_sel_mask=False).
+            num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+            num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)
+            nisa.tensor_scalar_reduce(
+                mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
+                data=mm1_psum_tile[:num_p, :num_f],
+                op0=nl.multiply,
+                operand0=ac.scale,
+                reduce_op=nl.maximum,
+                reduce_res=mm1_partial_max_tile[:num_p, k_tile_idx_in_section],
+            )
+def _pv_large_tile_impl(pv_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
+                        sp: SectionParams, bufs: AttnInternalBuffers):
+    """MM2 (P@V) for a Q group and 2048 V tile. LIVE: is_causal=False (all tiles selected),
+    tp_out=True, single large tile so ``large_tile_idx==0`` -> the ``tensor_copy`` PSUM->SBUF
+    reduction (never the accumulate ``tensor_tensor`` branch). SBUF-prefix V-tile selection is
+    DEAD (v_sb_prior is None); the tp_out=False MM2 orientation is DEAD."""
+    q_seqlen_offset = pv_grp * atp.sb_p
+    num_mm2_grps_in_large_tile = _LARGE_TILE_SZ // atp.mm2_grp_sz
+    mm2_psum_set = False
+    mm2_psum_tile = bufs.mm2_psum[pv_grp][large_tile_idx]
+    # Step 1: matmul + accumulate in PSUM per group.
+    for mm2_grp_i in range(num_mm2_grps_in_large_tile):
+        num_mm2_per_grp = atp.mm2_grp_sz // _V_TILE_SZ
+        num_mm2_per_large_tile = num_mm2_per_grp * num_mm2_grps_in_large_tile
+        exp_tp_sb_tile = bufs.exp_tp_sb[pv_grp][large_tile_idx][mm2_grp_i]
+        is_prior_tile, seqlen_k, k_start_pos_512_tile, _ = _get_kv_tile_apc(
+            ac.is_prefix_caching,
+            False,
+            True,
+            atp.seqlen_k_active_updated,
+            ac.seqlen_k_prior,
+            sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + mm2_grp_i * atp.mm2_grp_sz,
+            None,
+        )
+        for mm2_i in range(num_mm2_per_grp):
+            v_tile_idx = large_tile_idx * num_mm2_per_large_tile + mm2_grp_i * num_mm2_per_grp + mm2_i
+            k_start_pos = k_start_pos_512_tile + mm2_i * _V_TILE_SZ
+            num_p = min(seqlen_k - k_start_pos, _V_TILE_SZ)
+            num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+            mm2_sel_mask = True  # is_causal=False
+            if mm2_sel_mask and v_tile_idx < atp.num_v_tiles_per_section and num_p > 0 and num_f > 0:
+                mm2_psum_set = True
+                # Non-APC path -> use v_tile_idx.
+                v_tile_to_use = bufs.v_sb[v_tile_idx]
+                # tp_out=True orientation.
+                nisa.nc_matmul(
+                    mm2_psum_tile[: ac.d, :num_f],
+                    v_tile_to_use[:num_p, : ac.d],
+                    exp_tp_sb_tile[:num_p, nl.ds(mm2_i * _V_TILE_SZ, num_f)],
+                )
+    # Step 2: reduce MM2 groups into the large tile. is_prior_tile=False, is_causal=False ->
+    # loop_reduce_sel True; single large tile -> large_tile_idx==0 -> tensor_copy (never the
+    # accumulate tensor_tensor branch).
+    is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
+        ac.is_prefix_caching,
+        False,
+        True,
+        atp.seqlen_k_active_updated,
+        ac.seqlen_k_prior,
+        sp.section_offset + large_tile_idx * _LARGE_TILE_SZ,
+        None,
+    )
+    k_seqlen_mask = k_start_pos < seqlen_k
+    loop_reduce_sel = True  # is_causal=False
+    if loop_reduce_sel and k_seqlen_mask and mm2_psum_set:
+        # tp_out=True.
+        num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
+        if large_tile_idx == 0:
+            nisa.tensor_copy(bufs.mm2_sb[pv_grp][:, :num_f], mm2_psum_tile[:, :num_f])
+        else:
+            nisa.tensor_tensor(
+                bufs.mm2_sb[pv_grp][:, :num_f],
+                bufs.mm2_sb[pv_grp][:, :num_f],
+                mm2_psum_tile[:, :num_f],
+                nl.add,
+            )

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

	@@ -0,0 +1,294 @@

+"""K4 -- fused MLP (fc1 + gelu + fc2), x-prestage ("XP") in prologue mode.
+WHAT IT COMPUTES
+----------------
+``out = gelu(x @ w1.T + b1) @ w2.T + b2``, with the rows sharded across LNC cores and
+the weights pre-transposed to bf16. Returns ``[R, K]`` bf16.
+WHY IT IS SHAPED THIS WAY
+-------------------------
+This is the effective shipping kernel: it is the ``b1pre`` MLP (fc1 bias staged in the
+prologue so every gelu's bias is resident from t=0) with the x-prestage applied. Two
+design choices carry the performance:
+* b1 first: the 32 per-block ``b1`` bias tiles are DMA'd before the 16 MB of weight
+  staging, so they land within ~1 us of kernel start and no gelu ever waits on its bias.
+* x prestaged in the prologue: BOTH row blocks' ``x`` load + dma_transpose staging is
+  issued before the ``rb`` loop, in row-block order. rb=0's staging already hid in the
+  pre-K4 seam (where the Tensor engine is dark for ~148 us because K1/K3 emit no Tensor
+  instructions); prestaging rb=1 gives it the same cover, filling the ~26.6 us Tensor
+  idle window that used to sit at the rb0->rb1 boundary. Appending rb=1's stages after
+  rb=0's, in order, means rb=0's readiness cannot regress -- the first matmul cannot be
+  delayed.
+The x staging is emitted through a single module-level function (``_stage_xt``) rather
+than inlined per row block, so the ``x`` load and the dma_transpose each keep exactly one
+source site. The two-row-block staging is preserved exactly: ``xt_sets`` holds one
+transposed-tile set per row block, and the loop consumes ``xt_sets[rb]``.
+PSUM: fc1 uses 1 bank, fc2 uses up to 4 banks; the ``y1_psum`` / ``out_psums`` tiles are
+declared inside their loops so the ring never accumulates onto a stale PSUM value. The
+b2 bias broadcast is kept (removing it regressed the sibling kernel by 23-30%).
+There are no ``engine=`` pins anywhere in this kernel. INFERENCE ONLY: no autograd rule.
+"""
+from __future__ import annotations
+import math
+import torch
+from torch_neuronx import nki_op
+import nki
+import nki.isa as nisa
+import nki.language as nl
+#: PE array partition limit, and the moving-operand free-axis limit.
+_P_MAX = 128
+_F_MOV_MAX = 512
+_BF16 = nl.bfloat16
+#: Logical-to-physical NeuronCore degree. Must match the shipping shard's.
+LNC = 2
+_INV_SQRT2 = 1.0 / math.sqrt(2.0)
+_INV_SQRT_2PI = 1.0 / math.sqrt(2.0 * math.pi)
+def _div_ceil(n: int, d: int) -> int:
+    """Ceiling division (never write (n + d - 1) // d inline)."""
+    return (n + d - 1) // d
+def _kernel_assert(condition: bool, error_text: str):
+    """Assert with an NKI-formatted error message."""
+    assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"
+def _stage_xt(x_hbm, r_abs, pr, n_k, K, rb_size):
+    """Emit one row block's x staging and return its ``n_k`` transposed bf16 tiles.
+    ``n_rs`` loads of ``[<=128, K]`` bf16, each followed by ``n_k`` DMA transposes into
+    the row block's ``xt_tiles``. Kept as a module-level emitter so the ``x`` load and
+    the dma_transpose each map to a single source site regardless of where it is called.
+    Args:
+        x_hbm:   [R, K] @ HBM bf16, the MLP input.
+        r_abs:   absolute starting row of this row block.
+        pr:      number of rows in this row block (<= rb_size).
+        n_k:     number of K tiles (fc1 contraction dimension tiles).
+        K:       hidden-in dimension.
+        rb_size: row-block width (free axis of the transposed tiles).
+    Returns:
+        list of ``n_k`` transposed ``[128, rb_size]`` bf16 SBUF tiles.
+    """
+    n_rs = _div_ceil(pr, _P_MAX)
+    xt_tiles = []
+    for _ in range(n_k):
+        xt_tiles.append(nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf))
+    for rs in range(n_rs):
+        rr0 = rs * _P_MAX
+        prs = min(_P_MAX, pr - rr0)
+        x_b16 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+        nisa.dma_copy(
+            dst=x_b16[0:prs, 0:K],
+            src=x_hbm[r_abs + rr0:r_abs + rr0 + prs, 0:K],
+        )
+        for kt in range(n_k):
+            k0 = kt * _P_MAX
+            pk = min(_P_MAX, K - k0)
+            nisa.dma_transpose(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs],
+                               src=x_b16[0:prs, k0:k0 + pk])
+    return xt_tiles
+@nki.jit
+def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
+    """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
+    Args:
+        x_hbm   (nl.ndarray): [R, K] @ HBM, **bfloat16** (kernel 1 emits bf16).
+        w1t_hbm (nl.ndarray): [K, H] @ HBM, **bf16**, == fc1.weight.T.
+        b1_hbm  (nl.ndarray): [H]    @ HBM, float32.
+        w2t_hbm (nl.ndarray): [H, K] @ HBM, **bf16**, == fc2.weight.T.
+        b2_hbm  (nl.ndarray): [K]    @ HBM, float32.
+    Returns:
+        nl.ndarray: [R, K] @ HBM, **bfloat16**.
+    """
+    R, K = x_hbm.shape
+    K_w1t, H = w1t_hbm.shape
+    H_w2t, K_w2t = w2t_hbm.shape
+    _kernel_assert(R >= 1 and K >= 1 and H >= 1, f"R({R}) K({K}) H({H}) must be >= 1")
+    _kernel_assert(K_w1t == K, f"w1t in-dim {K_w1t} != K {K}; expected w1t = w1.T of shape [K, H]")
+    _kernel_assert(K_w2t == K, f"w2t out-dim {K_w2t} != K {K}; expected w2t = w2.T of shape [H, K]")
+    _kernel_assert(H_w2t == H, f"w2t in-dim {H_w2t} != H {H}")
+    _kernel_assert(
+        R % LNC == 0,
+        f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. "
+        f"program_id is a symbolic value, so a partial final row block cannot be "
+        f"clamped; use the one-core pre-transposed kernel for uneven R",
+    )
+    rows_per = R // LNC
+    pid = nl.program_id(0)
+    r_base = pid * rows_per
+    out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    n_k = _div_ceil(K, _P_MAX)   # K tiles (fc1 contraction)
+    n_h = _div_ceil(H, _P_MAX)   # H tiles (fc1 out-rows / fc2 contraction)
+    # b1 FIRST, before the 16 MB of weight DMA: 32 [128,1] tiles issue at t=0 and land in
+    # ~1 us, so every gelu's bias is resident long before the first matmul retires.
+    b1_tiles = []
+    for ht in range(n_h):
+        h0 = ht * _P_MAX
+        ph = min(_P_MAX, H - h0)
+        t = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+        nisa.dma_copy(dst=t[0:ph, 0:1], src=b1_hbm[h0:h0 + ph])
+        b1_tiles.append(t)
+    # Weight staging (pre-transposed bf16, one DMA per tile).
+    w1t_tiles = []
+    for kt in range(n_k):
+        k0 = kt * _P_MAX
+        pk = min(_P_MAX, K - k0)
+        t = nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf)
+        nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :])
+        w1t_tiles.append(t)
+    w2t_tiles = []
+    for ht in range(n_h):
+        h0 = ht * _P_MAX
+        ph = min(_P_MAX, H - h0)
+        t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+        nisa.dma_copy(dst=t[0:ph, :], src=w2t_hbm[h0:h0 + ph, :])
+        w2t_tiles.append(t)
+    # Main fused loop.
+    rb_size = _F_MOV_MAX
+    n_rb = _div_ceil(rows_per, rb_size)
+    n_ko = _div_ceil(K, _F_MOV_MAX)
+    # THE X-PRESTAGE: stage EVERY row block's x here, in rb order, so rb=1's serialised
+    # DMA chain runs inside the pre-K4 seam (Tensor dark) instead of inside the idle
+    # window at the rb0->rb1 boundary. Appending rb=1's stages after rb=0's leaves rb=0's
+    # stages exactly where they are on the Sync engine's sub-stream.
+    xt_sets = []
+    for rb_pre in range(n_rb):
+        xt_sets.append(_stage_xt(
+            x_hbm, r_base + rb_pre * rb_size,
+            min(rb_size, rows_per - rb_pre * rb_size), n_k, K, 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 bf16 tiles: consume the prologue-staged set for this row block.
+        xt_tiles = xt_sets[rb]
+        # ---- fc1: gt_tiles[ht] = gelu(w1 @ x.T + b1), bf16, SBUF-resident ----
+        gt_tiles = []
+        for ht in range(n_h):
+            h0 = ht * _P_MAX
+            ph = min(_P_MAX, H - h0)
+            # Declared INSIDE the ht loop, so the ring never reuses a written PSUM tile.
+            # nc_matmul's default accumulate=None auto-detects "previously written" and
+            # would accumulate onto a stale value.
+            y1_psum = nl.ndarray((_P_MAX, rb_size), dtype=nl.float32, buffer=nl.psum)
+            for kt in nl.affine_range(n_k):
+                k0 = kt * _P_MAX
+                pk = min(_P_MAX, K - k0)
+                nisa.nc_matmul(
+                    dst=y1_psum[0:ph, 0:pr],
+                    stationary=w1t_tiles[kt][0:pk, h0:h0 + ph],
+                    moving=xt_tiles[kt][0:pk, 0:pr],
+                )
+            # b1 comes from the prologue-resident tile.
+            gt = nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf)
+            nisa.activation(dst=gt[0:ph, 0:pr], op=nl.gelu,
+                            data=y1_psum[0:ph, 0:pr], bias=b1_tiles[ht][0:ph, 0:1])
+            gt_tiles.append(gt)
+        # ---- fc2: out[r,k] = sum_h g[r,h] * w2[k,h] + b2[k] ----
+        for ko in range(n_ko):
+            ko0 = ko * _F_MOV_MAX
+            nk = min(_F_MOV_MAX, K - ko0)
+            out_psums = []
+            for _ in range(n_rs):
+                out_psums.append(
+                    nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
+                )
+            for ht in nl.affine_range(n_h):
+                h0 = ht * _P_MAX
+                ph = min(_P_MAX, H - h0)
+                for rs in range(n_rs):
+                    rr0 = rs * _P_MAX
+                    prs = min(_P_MAX, pr - rr0)
+                    nisa.nc_matmul(
+                        dst=out_psums[rs][0:prs, 0:nk],
+                        stationary=gt_tiles[ht][0:ph, rr0:rr0 + prs],
+                        moving=w2t_tiles[ht][0:ph, ko0:ko0 + nk],
+                    )
+            # b2 broadcast, hoisted out of the rs loop (removing the broadcast itself
+            # cost 23% and 30% in two separate attempts).
+            b2_bc = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
+            nisa.dma_copy(dst=b2_bc[0:_P_MAX, 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)
+                out_sb = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
+                nisa.tensor_tensor(dst=out_sb[0:prs, 0:nk],
+                                   data1=out_psums[rs][0:prs, 0:nk],
+                                   data2=b2_bc[0:prs, 0:nk],
+                                   op=nl.add)
+                nisa.dma_copy(
+                    dst=out_hbm[r_abs + rr0:r_abs + rr0 + prs, ko0:ko0 + nk],
+                    src=out_sb[0:prs, 0:nk],
+                )
+    return out_hbm
+def _no_grad(*tensors) -> None:
+    """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
+    for t in tensors:
+        if t is not None and t.requires_grad:
+            raise RuntimeError(
+                "nki_mlp_k4_xp registers no autograd rule (INFERENCE ONLY). "
+                "Use under torch.no_grad()/eval()."
+            )
+@nki_op("vjepa2::nki_mlp_k4_xp", mutates_args={})
+def nki_mlp_k4_xp(
+    x: torch.Tensor,
+    w1t: torch.Tensor,
+    b1: torch.Tensor,
+    w2t: torch.Tensor,
+    b2: torch.Tensor,
+) -> torch.Tensor:
+    """Fused MLP on LNC=2, pre-transposed bf16 weights, x prestaged in the prologue.
+    Args:
+        x   : [R, K] **bf16** on 'neuron'. ``R`` must be divisible by :data:`LNC`.
+        w1t : [K, H] bf16 == fc1.weight.T.   b1 : [H] float32.
+        w2t : [H, K] bf16 == fc2.weight.T.   b2 : [K] float32.
+    Returns:
+        [R, K] bf16.
+    """
+    _no_grad(x, w1t, b1, w2t, b2)
+    return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2)
+#: Export alias. Same object as the nki_op above, so the runtime
+#: op-name string ("vjepa2::nki_mlp_k4_xp") is unchanged whichever name is used.
+mlp = nki_mlp_k4_xp

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

	@@ -0,0 +1,211 @@

+"""K1/K3 -- fused residual add + affine-free LayerNorm (the "fold" kernel).
+WHAT IT COMPUTES
+----------------
+Given a residual stream ``a`` [R, H] fp32 and a stack of 1..4 branch tensors
+``b`` [B, R, H] bf16, per token row this kernel computes::
+    x_new = a + sum_j b[j]                     # fp32, the untouched residual stream
+    z     = (x_new - mean(x_new)) * rstd       # bf16, LayerNorm WITHOUT affine
+and returns ``(x_new, z)``. gamma/beta are NOT applied here -- they have been folded
+into the downstream consumers' packed weights at freeze time (K1's z feeds K2's Q/K/V
+projections; K3's z feeds K4's fc1). The affine fold is exact:
+``W @ (z*g + bt) + c == (W * g) @ z + (W @ bt + c)``, so scaling the consumer's input
+columns by gamma and adding ``W @ bt`` to its bias reproduces the same math with the
+affine removed from the norm kernel.
+WHY IT IS SHAPED THIS WAY
+-------------------------
+* LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
+  selected by ``program_id(0)``, so the two cores never touch the same tokens.
+* Ring of depth 4 over the token tiles: four rotating sets of SBUF tiles let the DMA
+  loads of tile ``t+1`` overlap the compute of tile ``t`` without a write-after-read
+  hazard on the buffers.
+* fp32 residual / bf16 branch: ``a`` (the residual stream) stays fp32 so ``x_new`` is
+  bitwise-invariant math and can be carried forward losslessly; the branch inputs and
+  the normalized output ``z`` are bf16, which is the precision boundary the downstream
+  kernels consume.
+* Thinned stat chain: mean and variance are computed with a shift-and-reduce chain
+  (neg_shift -> copy-reduce -> neg_mu -> square-reduce -> var -> rsqrt) that keeps the
+  per-lane statistics in fp32; the final normalize is a single fused ``tensor_scalar``
+  (``z = (x_new + neg_mu) * rstd``) whose two per-lane operands cost the same as one op
+  and whose bf16 cast rides the write.
+There are no ``engine=`` pins anywhere in this kernel; nothing engine-related to
+preserve. INFERENCE ONLY: no autograd rule is registered.
+"""
+from __future__ import annotations
+import torch
+from torch_neuronx import nki_op
+import nki
+import nki.isa as nisa
+import nki.language as nl
+#: SBUF partition count, and the token-tile height (rows processed per instruction).
+_P_MAX = 128
+#: The bf16 dtype used for branch inputs and the normalized output.
+_BF16 = nl.bfloat16
+#: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
+LNC = 2
+#: Ring depth: number of rotating SBUF tile sets that overlap load with compute.
+_RING = 4
+def _kernel_assert(condition: bool, error_text: str):
+    """Assert with an NKI-formatted error message."""
+    assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"
+@nki.jit
+def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
+    """Emit the fold kernel body: ``(a + sum(b), LN_no_affine(a + sum(b)))``, token-major.
+    gamma/beta live in the consumers' packed weights, so this kernel takes no affine
+    parameters. Rows are sharded across ``LNC`` cores; each token tile is loaded through
+    a depth-4 ring, summed in fp32, normalized with the thinned stat chain, and the
+    normalized result is cast to bf16 on write.
+    Args:
+        a_hbm (nl.ndarray): [R, H] @ HBM float32 -- the residual stream.
+        b_hbm (nl.ndarray): [B, R, H] @ HBM bf16 -- B branch tensors (1 <= B <= 4), summed in.
+        eps   (float):      LayerNorm epsilon.
+    Returns:
+        tuple: ``(x_new [R, H] float32, z [R, H] bf16)`` both in shared HBM.
+    """
+    R, H = a_hbm.shape
+    n_b, Rb, Hb = b_hbm.shape
+    _kernel_assert(R >= 1 and H >= 1, f"R({R}) H({H}) must be >= 1")
+    _kernel_assert(Rb == R and Hb == H, f"b {(Rb, Hb)} must match a {(R, H)}")
+    _kernel_assert(1 <= n_b <= 4, f"b must stack 1..4 branches, got {n_b}")
+    _kernel_assert(
+        H <= 16384,
+        f"H={H} is one instruction's free axis here; no chunking is implemented.",
+    )
+    _kernel_assert(
+        R % (LNC * _P_MAX) == 0,
+        f"need R divisible by LNC*{_P_MAX} = {LNC * _P_MAX}, got R={R}",
+    )
+    rows_per = R // LNC
+    n_t = rows_per // _P_MAX
+    pid = nl.program_id(0)
+    r_base = pid * rows_per
+    x_new_hbm = nl.ndarray((R, H), dtype=nl.float32, buffer=nl.shared_hbm)
+    z_hbm = nl.ndarray((R, H), dtype=_BF16, buffer=nl.shared_hbm)
+    inv_h = 1.0 / float(H)
+    neg_inv_h = -inv_h
+    # Ring buffers: _RING rotating sets of (a tile, B branch tiles, sum tile, z tile).
+    a_sb = []
+    b_sb = []
+    s_sb = []
+    z_sb = []
+    for _i in range(_RING):
+        a_sb.append(nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf))
+        _slot = []
+        for _j in range(n_b):
+            _slot.append(nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf))
+        b_sb.append(_slot)
+        s_sb.append(nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf))
+        z_sb.append(nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf))
+    # Scratch destination for the reduce activations (its full tile is discarded; only
+    # the reduce result is consumed).
+    sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)
+    # Per-lane [128, 1] statistics.
+    acc = 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)
+    for t in range(n_t):
+        k = t % _RING
+        r0 = r_base + t * _P_MAX
+        # ---- load residual + branches for this tile ----
+        nisa.dma_copy(dst=a_sb[k], src=a_hbm[r0:r0 + _P_MAX, 0:H])
+        for j in range(n_b):
+            nisa.dma_copy(dst=b_sb[k][j], src=b_hbm[j, r0:r0 + _P_MAX, 0:H])
+        # ---- x_new = a + sum_j b[j] in fp32, then write the residual stream out ----
+        nisa.tensor_tensor(dst=s_sb[k], data1=a_sb[k], data2=b_sb[k][0], op=nl.add)
+        for j in range(1, n_b):
+            nisa.tensor_tensor(dst=s_sb[k], data1=s_sb[k], data2=b_sb[k][j], op=nl.add)
+        nisa.dma_copy(dst=x_new_hbm[r0:r0 + _P_MAX, 0:H], src=s_sb[k])
+        # ---- thinned stat chain: mean via shift-and-reduce, variance, rstd ----
+        nisa.tensor_scalar(dst=neg_shift, data=s_sb[k][0:_P_MAX, 0:1],
+                           op0=nl.multiply, operand0=-1.0)
+        nisa.activation_reduce(dst=sink, op=nl.copy, data=s_sb[k],
+                               reduce_op=nl.add, reduce_res=acc,
+                               bias=neg_shift, scale=1.0)
+        nisa.tensor_scalar(dst=neg_mu, data=acc, op0=nl.multiply, operand0=neg_inv_h,
+                           op1=nl.add, operand1=neg_shift)
+        nisa.activation_reduce(dst=sink, op=nl.square, data=s_sb[k],
+                               reduce_op=nl.add, reduce_res=acc, bias=neg_mu, scale=1.0)
+        nisa.tensor_scalar(dst=var, data=acc, op0=nl.multiply, operand0=inv_h,
+                           op1=nl.add, operand1=eps)
+        nisa.activation(dst=rstd, op=nl.rsqrt, data=var)
+        # ---- one-op affine-free apply: z = (x_new + neg_mu) * rstd, bf16 cast on write.
+        # Both operands are per-lane [128, 1] vectors; the fused two-op tensor_scalar
+        # costs the same as one op.
+        nisa.tensor_scalar(dst=z_sb[k], data=s_sb[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_sb[k])
+    return x_new_hbm, z_hbm
+def _no_grad(*tensors: torch.Tensor) -> None:
+    """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
+    for t in tensors:
+        if t.requires_grad:
+            raise RuntimeError(
+                "nki_res_layernorm_fold registers no autograd rule. Use under "
+                "torch.no_grad()/eval()."
+            )
+@nki_op("vjepa2::nki_res_layernorm_fold", mutates_args={})
+def nki_res_layernorm_fold(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused residual add + affine-free LayerNorm on LNC=2 (gamma/beta folded away).
+    The consumers MUST have been packed with folded weights; feeding ``z`` to unfolded
+    weights silently drops the affine.
+    Args:
+        a:   [R, H] float32 on 'neuron'. ``R`` divisible by ``LNC * 128``.
+        b:   [B, R, H] bf16 on 'neuron'. B branch tensors (1 <= B <= 4), summed in.
+        eps: LayerNorm epsilon.
+    Returns:
+        ``(x_new [R, H] float32, z [R, H] bf16)``.
+    """
+    _no_grad(a, b)
+    return _res_layernorm_fold_kernel[LNC](a, b, eps)
+#: Export alias. Same object as the nki_op above, so the runtime
+#: op-name string ("vjepa2::nki_res_layernorm_fold") is unchanged whichever name is used.
+res_layernorm_fold = nki_res_layernorm_fold

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

	@@ -0,0 +1,47 @@

+"""Host-side 3D RoPE cos / signed-sin tables for the attention block.
+Pure torch, computed once at install time per model (positions depend only on the token grid).
+The attention block consumes these through ``pack_tables`` (see attention_block.py).
+"""
+from __future__ import annotations
+import torch
+__all__ = ["build_rope_tables"]
+def build_rope_tables(sub_dims, positions, head_dim, dtype, device):
+    """Build faithful cos + signed-sin tables [N, head_dim] for 3D RoPE.
+    Args:
+        sub_dims  : (d_dim, h_dim, w_dim) sub-block widths (each even).
+        positions : (pos_d, pos_h, pos_w) each [N] int/long token positions.
+        head_dim  : full head dim D (>= sum(sub_dims)); remainder is passthrough.
+        dtype, device: output dtype/device.
+    Returns:
+        (cos, sinsigned): each [N, head_dim] float32:
+            cos[n,k]=cos(pos_axis*omega[k%half]) for rotated dims, else 1.
+            sinsigned[n,k]=g[k]*sin(...) for rotated dims, else 0;  g=[-1,+1,...].
+    """
+    (pos_d, pos_h, pos_w) = positions
+    N = pos_d.shape[0]
+    cos = torch.ones(N, head_dim, dtype=torch.float32, device=device)
+    sin = torch.zeros(N, head_dim, dtype=torch.float32, device=device)
+    s = 0
+    for sub, pos in zip(sub_dims, (pos_d, pos_h, pos_w)):
+        half = sub // 2
+        omega = torch.arange(half, dtype=torch.float32, device=device) / float(half)
+        omega = 1.0 / (10000.0 ** omega)            # [half]
+        freq = pos.to(device=device, dtype=torch.float32).unsqueeze(-1) * omega   # [N, half]
+        c, sn = freq.cos(), freq.sin()
+        cos[:, s:s + half] = c
+        cos[:, s + half:s + sub] = c
+        sin[:, s:s + half] = sn
+        sin[:, s + half:s + sub] = sn
+        s += sub
+    g = torch.ones(head_dim, dtype=torch.float32, device=device)
+    g[0::2] = -1.0
+    sinsigned = sin * g
+    return cos.to(dtype), sinsigned.to(dtype)