jburtoft commited on 18 days ago

Commit

c855122

1 Parent(s): b821eb9

v1: NKI kernels for RoPE ViT encoders (DINOv3, V-JEPA2) on Trainium2

Browse files

Files changed (19) hide show

.gitattributes +4 -2
CHANGELOG.md +20 -0
LICENSE +202 -0
NOTICE +13 -0
README.md +206 -0
build/torch-neuron/__init__.py +37 -0
build/torch-neuron/encoder.py +237 -0
build/torch-neuron/layers.py +220 -0
build/torch-neuron/metadata.json +1 -0
build/torch-neuron/nki_kernels/__init__.py +23 -0
build/torch-neuron/nki_kernels/attention_block.py +801 -0
build/torch-neuron/nki_kernels/attention_block_l1.py +402 -0
build/torch-neuron/nki_kernels/attention_block_l1_ln.py +451 -0
build/torch-neuron/nki_kernels/attention_inner.py +1220 -0
build/torch-neuron/nki_kernels/mlp.py +296 -0
build/torch-neuron/nki_kernels/mlp_res_ln.py +308 -0
build/torch-neuron/nki_kernels/res_layernorm.py +213 -0
build/torch-neuron/nki_kernels/swiglu_stream.py +291 -0
examples/dinov3_features.py +22 -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,20 @@

+# Changelog
+## v1 (2026-10)
+First release.
+* Drop-in layers for Hugging Face Transformers: `NeuronDINOv3ViTModel` (DINOv3 ViT-S/16, S+/16, B/16,
+  L/16, H+/16) and `NeuronVJEPA2Encoder` (V-JEPA2 ViT-L, any batch size, up to 2048 tokens per clip).
+* Model-agnostic encoder loop: `BlockWeights`, `pack_block`, `rope_tables`, `EncoderRunner`.
+* Kernels: `attention_block` (per head group; batched; padded keys skipped; head_dim 64 or 128;
+  any head count), `attention_block_l1` and `attention_block_l1_ln` (LNC=1, all head groups, optional
+  residual + LayerNorm epilogue), `res_layernorm_fold`, `mlp` and `mlp_res_ln` (GELU MLP, optional
+  residual + next LayerNorm epilogue), `swiglu_stream` and `swiglu_stream_ln` (weight-streaming
+  SwiGLU FFN).
+* RoPE pairing selectable at packing time: `rope="half"` (rotate-half) or `rope="interleaved"`.
+* Verified on trn2.3xlarge, PyTorch Native Beta 6 (torch 2.13.0, torch-neuronx 2.13.3, NKI 0.7.0b1),
+  Transformers 4.57.6, LNC=1 and LNC=2. Every measured case passes rel_l2 <= 5e-2 vs the fp32 CPU model.
+Known limitations: encoder (bidirectional) attention with RoPE only; LayerNorm (not RMSNorm);
+<= 2048 tokens per image; inference only; hidden size 4096 (ViT-7B) not covered by the drop-in layer.

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,13 @@

+vit-encoder-neuron-kernels -- NKI kernels for RoPE ViT encoders on AWS Trainium2
+Licensed under the Apache License, Version 2.0 (see LICENSE).
+nki_kernels/attention_block.py, res_layernorm.py and mlp.py are modified versions of the V-JEPA2
+encoder kernels published at huggingface.co/kernels/jburtoft/vjepa2-neuron-kernels (Apache-2.0),
+developed by Antonio Mena. Modifications: selectable RoPE pairing, any head count, head_dim 128,
+ragged token tiles and padded-key skipping, multi-image launches, resident weights, LNC=1 grid plan.
+nki_kernels/attention_block_l1.py, attention_block_l1_ln.py and mlp_res_ln.py are derived from them.
+nki_kernels/swiglu_stream.py, encoder.py and layers.py are new.
+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, licensed under the
+Apache License, Version 2.0. See the file header for the list of modifications.

README.md ADDED Viewed

	@@ -0,0 +1,206 @@

+---
+license: apache-2.0
+library_name: kernels
+tags:
+  - kernel
+  - neuron
+  - nki
+  - trainium
+  - trn2
+  - vision
+  - vit
+  - dinov3
+  - vjepa2
+  - rope
+  - swiglu
+---
+# vit-encoder-neuron-kernels
+NKI kernels that run the whole transformer stack of **RoPE vision-transformer encoders** on AWS
+Trainium2 (trn2). Each block is 2-3 coarse kernels instead of dozens of small ops: attention plus
+the residual add and LayerNorm, and the MLP / SwiGLU FFN plus the next block's LayerNorm. The
+residual stream stays in fp32 between kernels.
+Drop-in layers are included for two Hugging Face Transformers models:
+| Layer | Replaces | Models |
+|---|---|---|
+| `NeuronDINOv3ViTModel` | `DINOv3ViTModel.forward` | DINOv3 ViT-S/16, S+/16, B/16, L/16, H+/16 |
+| `NeuronVJEPA2Encoder` | `VJEPA2Encoder.forward` | V-JEPA2 ViT-L (video), up to 2048 tokens per clip |
+The building blocks (`BlockWeights` / `pack_block` / `EncoderRunner`) are model-agnostic. Any
+pre-norm ViT with the shape constraints below can use them, but only the two models above
+have been validated.
+## Performance (trn2.3xlarge, bf16, `torch.compile(backend="neuron")`)
+Measured 2026-10 on PyTorch Native Beta 6 (torch 2.13.0, torch-neuronx 2.13.3, NKI 0.7.0b1).
+Real Hugging Face weights. "Stock" is the same model compiled with
+`torch.compile(backend="neuron")` without these kernels: DINOv3 cast to bf16, V-JEPA2 as its fp32
+checkpoint (stock V-JEPA2 is the encoder as shipped, and its accuracy is in the table below).
+**Per instance** (one process per logical core, all cores busy):
+| Model | Resolution | Stock (img/s) | **These kernels (img/s)** | Speedup |
+|---|---|---|---|---|
+| DINOv3 ViT-B/16 | 224 | 2,901.5 | **4,105.2** | 1.41x |
+| DINOv3 ViT-L/16 | 224 | 1,061.2 | **1,767.5** | 1.67x |
+| DINOv3 ViT-L/16 | 448 | 201.1 | **491.9** | 2.45x |
+| DINOv3 ViT-H+/16 | 224 | 400.0 | **626.6** | 1.57x |
+| DINOv3 ViT-H+/16 | 448 | 106.8 | **200.3** | 1.88x |
+| V-JEPA2 ViT-L, 16 frames x 256 px | - | 36.8 clips/s | **157.3 clips/s** | 4.27x |
+Configuration: LNC=1 (8 logical cores) for every row except V-JEPA2 stock, whose best was LNC=2
+(4 cores). DINOv3 kernels run at batch 8 (224 px) or batch 2 (448 px). DINOv3 stock runs at batch 4
+(224 px) or batch 1 (448 px), its best measured. V-JEPA2 runs at batch 1 per core for both.
+**Single core, latency-oriented** (LNC=2, batch 1):
+| Model | Resolution | Stock | **These kernels** | Speedup |
+|---|---|---|---|---|
+| DINOv3 ViT-B/16 | 224 | 251.0 img/s | **510.7 img/s** | 2.03x |
+| DINOv3 ViT-L/16 | 224 | 97.1 img/s | **220.3 img/s** (4.5 ms) | 2.27x |
+| DINOv3 ViT-L/16 | 448 | 25.1 img/s | **83.0 img/s** | 3.31x |
+| DINOv3 ViT-H+/16 | 224 | 21.4 img/s | **88.2 img/s** | 4.12x |
+| DINOv3 ViT-H+/16 | 448 | 17.7 img/s | **32.4 img/s** | 1.83x |
+| V-JEPA2 ViT-L, 16 frames x 256 px | - | 101.5 ms/clip | **25.8 ms/clip** | 3.93x |
+The gain grows with resolution: attention cost grows with token count, and these kernels keep
+attention and its surroundings on chip. At 224 px with large batches the gap narrows. For example,
+stock ViT-L/16 at LNC=2 batch 4 is 252 img/s and the kernels give 321 img/s (1.27x).
+## Accuracy
+Every result above passes a per-image check against the fp32 CPU model (rel_l2 <= 5e-2).
+On DINOv3 the kernels are **closer to fp32 than stock bf16**, because the residual stream and
+LayerNorm statistics stay in fp32:
+| Model | Stock bf16 rel_l2 vs fp32 | These kernels |
+|---|---|---|
+| DINOv3 ViT-B/16 | 1.4-1.8e-2 | 0.9-1.7e-2 |
+| DINOv3 ViT-L/16 | 1.3-1.6e-2 | 0.8-1.0e-2 |
+| DINOv3 ViT-H+/16 | 1.0-1.7e-2 | 0.6-1.0e-2 |
+| V-JEPA2 ViT-L | 1.4e-2 (fp32 model, compiler-chosen precision) | 3.9e-2 (bf16 matmuls) |
+For V-JEPA2 the stock path compiles the fp32 checkpoint and keeps more of it in higher precision,
+so it is closer to fp32 than the kernels (which run every matmul in bf16).
+## Quick start
+```python
+import types, torch
+from kernels import get_kernel
+from transformers import AutoModel
+k = get_kernel("jburtoft/vit-encoder-neuron-kernels", version=1, trust_remote_code=True)
+model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m",
+                                  dtype=torch.float32).eval().to("neuron")
+model.forward = types.MethodType(k.NeuronDINOv3ViTModel.forward, model)
+model = torch.compile(model, backend="neuron")
+with torch.no_grad():
+    out = model(pixel_values.to("neuron"))   # [B, 3, H, W] fp32, any batch size
+feats = out.last_hidden_state                 # [B, 1 + 4 registers + patches, 1024]
+cls = out.pooler_output
+```
+V-JEPA2:
+```python
+model = AutoModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256", dtype=torch.float32).eval().to("neuron")
+model.encoder.forward = types.MethodType(k.NeuronVJEPA2Encoder.forward, model.encoder)
+enc = torch.compile(model.encoder, backend="neuron")
+emb = enc(clip.to("neuron")).last_hidden_state   # clip [B, 16, 3, 256, 256] -> [B, 2048, 1024]
+```
+Notes:
+* Load the model in **fp32**. The layer folds LayerNorm affines and LayerScale into the weights
+  once, on the first device forward, and casts the weights to bf16 itself.
+* **`torch.compile(backend="neuron")` is part of the configuration.** The layers set
+  `can_torch_compile = True`.
+* **Set `NEURON_LOGICAL_NC_CONFIG` before importing**, and to the same value the process runs
+  with: `2` (default) or `1`. The kernels are built for that core configuration. For LNC=1 also pass
+  `NEURON_CC_FLAGS=--lnc=1`.
+* `NKI_ENABLE_TRACE_CACHE=0` is recommended when you edit kernels. The cross-process trace cache
+  can otherwise reuse a stale NEFF.
+* On a vLLM-Neuron venv (CUDA-tagged torch), `import libtorch_neuronx_lite` before `get_kernel`
+  so the `torch-neuron` variant is selected. On PyTorch Native this is not needed.
+Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per
+instance. **LNC=2** (4 larger cores) gives the lowest latency per request.
+## Kernels
+| Kernel | What it does |
+|---|---|
+| `attention_block` | LN'd hidden -> Q/K/V projection + RoPE + attention + output projection, one 8-head (D=64) or 4-head (D=128) group per program; any number of images per launch; padded keys skipped |
+| `attention_block_l1` | LNC=1: one program over **all** head groups, one output (no per-group partials) |
+| `attention_block_l1_ln` | `attention_block_l1` plus residual add and LayerNorm in the epilogue |
+| `res_layernorm_fold` | residual add (sum of up to 4 branches) + affine-free LayerNorm |
+| `mlp` | GELU MLP (fc1 -> GELU -> fc2) |
+| `mlp_res_ln` | `mlp` plus residual add and the next block's LayerNorm in the epilogue |
+| `swiglu_stream` | SwiGLU FFN that streams the weights in chunks of the intermediate dimension (so large FFNs fit on chip) |
+| `swiglu_stream_ln` | `swiglu_stream` plus residual add and the next block's LayerNorm |
+Host-side packing (`pack_weights`, `pack_tables`, `lane_permutation`, `pack_block`) supports both
+RoPE pairings via `rope=`:
+* `"half"`: pairs lane `j` with `j + D/2` (rotate-half; DINOv3, EVA-02, HF Llama style).
+* `"interleaved"`: pairs lane `2j` with `2j + 1` (V-JEPA2 style).
+## Supported shapes and limits
+* **Bidirectional (encoder) self-attention with RoPE.** No causal masks, KV cache, cross-attention,
+  or learned / absolute position embeddings. A model without RoPE could pass identity tables (cos=1,
+  sin=0), but that is untested.
+* **Pre-norm blocks with standard LayerNorm** (not RMSNorm), optional LayerScale.
+* hidden size a multiple of 128; head_dim 64 or 128; any head count (padded to a multiple of 8
+  at D=64 or 4 at D=128).
+* **<= 2048 tokens per image** after padding to a multiple of 128 x LNC (e.g. 720 x 720 px at
+  patch 16). Longer inputs fall back to the stock forward.
+* FFN: GELU MLP (`nl.gelu`, the erf form) or SwiGLU (SiLU gate).
+* Inference only (no autograd). The layers refuse in training mode.
+* ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism.
+  The streaming SwiGLU kernel is slower than the compiler's FFN at that width.
+* Validated end to end on DINOv3 ViT-B/L/H+ and V-JEPA2 ViT-L only.
+## Requirements
+* trn2 (verified on trn2.3xlarge, LNC=1 and LNC=2).
+* PyTorch Native (TorchNeuron) with NKI 0.7 and `nkilib`. Verified on Beta 6 (torch 2.13.0,
+  torch-neuronx 2.13.3, NKI 0.7.0b1, neuronx-cc 2.0.404056).
+* Transformers with DINOv3 / V-JEPA2 (verified on 4.57.6); `kernels` for `get_kernel`.
+## Repository layout
+```
+build/torch-neuron/
+├── __init__.py              # public API
+├── metadata.json
+├── layers.py                # NeuronDINOv3ViTModel, NeuronVJEPA2Encoder
+├── encoder.py               # BlockWeights, pack_block, rope_tables, EncoderRunner
+└── nki_kernels/
+    ├── attention_block.py       # attention (per head group) + host-side packing
+    ├── attention_block_l1.py    # attention, all head groups (LNC=1)
+    ├── attention_block_l1_ln.py # ... + residual + LayerNorm epilogue
+    ├── attention_inner.py       # modified copy of nkilib attention_cte (vendor attention core)
+    ├── res_layernorm.py         # residual + affine-free LayerNorm
+    ├── mlp.py                   # GELU MLP
+    ├── mlp_res_ln.py            # GELU MLP + residual + next LayerNorm
+    └── swiglu_stream.py         # weight-streaming SwiGLU (+ residual/LayerNorm variant)
+```
+## Attribution and license
+Apache-2.0 (see `LICENSE`, `NOTICE`). The attention, MLP and residual/LayerNorm kernels derive from
+the V-JEPA2 encoder kernels in
+[`jburtoft/vjepa2-neuron-kernels`](https://huggingface.co/kernels/jburtoft/vjepa2-neuron-kernels),
+developed by Antonio Mena ([@menaman123](https://huggingface.co/menaman123)). This repo generalizes
+them: both RoPE pairings, any head count, head_dim 128, batching, padded-key skipping, an LNC=1
+all-group variant, and the fused residual/LayerNorm epilogues. `swiglu_stream.py` is new.
+`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.
+DINOv3 and V-JEPA2 model weights are distributed under their own licenses.

build/torch-neuron/__init__.py ADDED Viewed

	@@ -0,0 +1,37 @@

+"""NKI kernels for RoPE ViT encoders on AWS Trainium2 (trn2).
+Model layers (drop-in forwards for Hugging Face Transformers):
+  NeuronDINOv3ViTModel   -- DINOv3ViTModel (S/S+/B/L/H+)
+  NeuronVJEPA2Encoder    -- VJEPA2Encoder
+Model-agnostic building blocks:
+  BlockWeights, pack_block, rope_tables, EncoderRunner   -- pre-norm encoder loop on the kernels
+  kernels: attention_block, attention_block_l1, attention_block_l1_ln, res_layernorm_fold,
+           mlp, mlp_res_ln, swiglu_stream, swiglu_stream_ln
+"""
+from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
+from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder
+from .nki_kernels import (
+    attention_block,
+    attention_block_l1,
+    attention_block_l1_ln,
+    lane_permutation,
+    mlp,
+    mlp_res_ln,
+    pack_out_weights,
+    pack_tables,
+    pack_v_weights,
+    pack_weights,
+    res_layernorm_fold,
+    swiglu_stream,
+    swiglu_stream_ln,
+)
+__all__ = [
+    "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder",
+    "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC",
+    "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
+    "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
+    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
+]

build/torch-neuron/encoder.py ADDED Viewed

	@@ -0,0 +1,237 @@

+"""Model-agnostic ViT encoder loop on the NKI kernels.
+A pre-norm transformer block
+    x = x + ls1 * attn(LN1(x))          attn: q/k/v projections, RoPE on q/k, softmax, out proj
+    x = x + ls2 * ffn(LN2(x))           ffn:  GELU MLP  or  SwiGLU
+is described once per block as a :class:`BlockWeights` (plain fp32 tensors, any source), packed
+once by :func:`pack_block` (LayerNorm affines and LayerScale folded into the weights), and run by
+:class:`EncoderRunner` on a stacked ``[B*N, H]`` row buffer.
+Kernel choice per block (all automatic):
+* attention: LNC=1 -> one program over all head groups, with the residual add + LN2 in its
+  epilogue; LNC=2 -> per-group programs whose partials the residual/LN kernel sums.
+* FFN: GELU -> fused MLP; SwiGLU -> weight-streaming SwiGLU. Both also do the residual add and
+  the NEXT block's LN1 in their epilogue, so the stand-alone residual/LN kernel runs only before
+  block 0 (and after attention at LNC=2).
+INFERENCE ONLY (no autograd). Constraints: hidden size % 128 == 0, head_dim in {64, 128},
+tokens per image <= 2048 after padding to a multiple of 128 * LNC.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import importlib as _il
+_ab = _il.import_module(__package__ + ".nki_kernels.attention_block")
+_ab1 = _il.import_module(__package__ + ".nki_kernels.attention_block_l1")
+_ab1ln = _il.import_module(__package__ + ".nki_kernels.attention_block_l1_ln")
+from .nki_kernels.attention_block import LNC
+from .nki_kernels.mlp import nki_mlp_k4_xp
+from .nki_kernels.mlp_res_ln import nki_mlp_res_ln
+from .nki_kernels.res_layernorm import nki_res_layernorm_fold
+from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln
+__all__ = ["BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC"]
+_SBUF_BUDGET = 180 * 1024        # per-partition bytes the LNC=1 resident-weights decision allows
+_SW_MIN_ROWS = 256               # below this many rows per core, SwiGLU runs in torch (measured faster)
+_L1_MAX_N = 1280                 # LNC=1 all-group attention kernels verified up to this many tokens per
+                                 # image; beyond it (e.g. 2048-token video clips) they exceed SBUF and
+                                 # the per-group kernel + residual/LN kernel are used instead
+_RESIDENT_MAX_N = 1280           # LNC=2 attention: weights staged once per launch only up to this many
+                                 # tokens per image (2048 tokens + resident weights exceed SBUF)
+@dataclass
+class BlockWeights:
+    """One transformer block, as plain fp32 CPU tensors in ``nn.Linear`` layout ([out, in])."""
+    ln1_w: torch.Tensor
+    ln1_b: torch.Tensor
+    ln1_eps: float
+    wq: torch.Tensor
+    bq: Optional[torch.Tensor]
+    wk: torch.Tensor
+    bk: Optional[torch.Tensor]
+    wv: torch.Tensor
+    bv: Optional[torch.Tensor]
+    wo: torch.Tensor
+    bo: Optional[torch.Tensor]
+    ln2_w: torch.Tensor
+    ln2_b: torch.Tensor
+    ln2_eps: float
+    # GELU MLP: fc1 [I, H], fc2 [H, I]. SwiGLU: gate [I, H], up [I, H], down [H, I].
+    fc1_w: Optional[torch.Tensor] = None
+    fc1_b: Optional[torch.Tensor] = None
+    fc2_w: Optional[torch.Tensor] = None
+    fc2_b: Optional[torch.Tensor] = None
+    gate_w: Optional[torch.Tensor] = None
+    gate_b: Optional[torch.Tensor] = None
+    up_w: Optional[torch.Tensor] = None
+    up_b: Optional[torch.Tensor] = None
+    down_w: Optional[torch.Tensor] = None
+    down_b: Optional[torch.Tensor] = None
+    ls1: Optional[torch.Tensor] = None       # LayerScale (None = 1)
+    ls2: Optional[torch.Tensor] = None
+    scale: Optional[float] = None            # attention scale (None = head_dim ** -0.5)
+def _f(t):
+    return None if t is None else t.detach().float().cpu()
+def _z(t, n):
+    return torch.zeros(n) if t is None else t
+def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
+    """Fold and pack one block for the kernels. ``rope`` in {"half", "interleaved"}."""
+    wq, wk, wv, wo = _f(bw.wq), _f(bw.wk), _f(bw.wv), _f(bw.wo)
+    H = wq.shape[1]
+    F_ = wq.shape[0]
+    D = F_ // num_heads
+    gsz = 512 // D                                   # heads per kernel program
+    nh_p = ((num_heads + gsz - 1) // gsz) * gsz      # padded head count
+    Fp = nh_p * D
+    ngrp = nh_p // gsz
+    g1, t1, g2, t2 = _f(bw.ln1_w), _f(bw.ln1_b), _f(bw.ln2_w), _f(bw.ln2_b)
+    bq, bk, bv = _z(_f(bw.bq), F_), _z(_f(bw.bk), F_), _z(_f(bw.bv), F_)
+    ls1 = _f(bw.ls1) if bw.ls1 is not None else torch.ones(H)
+    ls2 = _f(bw.ls2) if bw.ls2 is not None else torch.ones(H)
+    bo = _z(_f(bw.bo), H) * ls1
+    wo = wo * ls1[:, None]
+    padr = lambda w: torch.cat([w, w.new_zeros(Fp - w.shape[0], *w.shape[1:])], 0)
+    p = {"D": D, "heads": num_heads, "nh_p": nh_p, "ngrp": ngrp, "H": H,
+         "eps1": float(bw.ln1_eps), "eps2": float(bw.ln2_eps),
+         "scale": float(bw.scale if bw.scale is not None else D ** -0.5)}
+    p["wqk"], p["bqk"] = _ab.pack_weights(padr(wq * g1[None]), padr(wk * g1[None]),
+                                          padr(bq + wq @ t1), padr(bk + wk @ t1), nh_p, D, rope)
+    p["wvt"], p["bv"] = _ab.pack_v_weights(padr(wv * g1[None]), padr(bv + wv @ t1))
+    p["wot"] = torch.cat([wo, wo.new_zeros(H, Fp - F_)], 1).t().contiguous().to(torch.bfloat16)
+    p["bo"] = (bo / ngrp).reshape(1, H).contiguous().to(torch.bfloat16)      # per-group partials
+    p["bo_full"] = bo.reshape(1, H).contiguous().to(torch.bfloat16)          # all-group kernel
+    if bw.gate_w is not None:                                                # SwiGLU
+        gw, uw, dw = _f(bw.gate_w), _f(bw.up_w), _f(bw.down_w)
+        I = gw.shape[0]
+        Ip = ((I + 127) // 128) * 128
+        pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0)
+        p["swiglu"] = True
+        p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16)          # gate (silu)
+        p["s_b1"] = pi(_z(_f(bw.gate_b), I) + gw @ t2).contiguous()
+        p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16)          # up
+        p["s_bg"] = pi(_z(_f(bw.up_b), I) + uw @ t2).contiguous()
+        p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16)
+        p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous()
+    else:                                                                     # GELU MLP
+        w1, w2 = _f(bw.fc1_w), _f(bw.fc2_w)
+        p["swiglu"] = False
+        p["w1t"] = (w1 * g2[None]).t().contiguous().to(torch.bfloat16)
+        p["b1"] = (_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous()
+        p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16)
+        p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous()
+    return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()}
+def rope_tables(cos: torch.Tensor, sin: torch.Tensor, n_prefix: int, n_pad: int, num_heads: int,
+                D: int, rope: str, device) -> torch.Tensor:
+    """Kernel RoPE tables for a ``[n_prefix + n_rot, ...]`` token layout padded to ``n_pad`` rows.
+    ``cos``/``sin``: ``[n_rot, D]`` the model's own tables for the rotated tokens (prefix tokens
+    such as CLS/registers and padding get cos=1, sin=0). For ``rope="interleaved"`` pass the
+    SIGNED sin table (see ``pack_tables``).
+    """
+    C = torch.ones(n_pad, D)
+    S = torch.zeros(n_pad, D)
+    n_rot = cos.shape[0]
+    C[n_prefix:n_prefix + n_rot] = cos.float().cpu()
+    S[n_prefix:n_prefix + n_rot] = sin.float().cpu()
+    gsz = 512 // D
+    nh_p = ((num_heads + gsz - 1) // gsz) * gsz
+    return _ab.pack_tables(C, S, nh_p, D, rope).to(device)
+class EncoderRunner:
+    """Runs packed blocks on a stacked ``[B*N, H]`` fp32 residual stream (rows of image b at b*N)."""
+    def __init__(self, packed: list):
+        self.L = packed
+        p0 = packed[0]
+        self.D, self.H, self.ngrp = p0["D"], p0["H"], p0["ngrp"]
+        self.plan = []                                   # (group base, grid) per LNC=2 attention call
+        g = 0
+        while g < self.ngrp:
+            gr = 2 if (LNC == 2 and self.ngrp - g >= 2) else 1
+            self.plan.append((g, gr))
+            g += gr
+    def _res(self, mod, N):
+        return int(mod.sbuf_estimate(self.H, self.ngrp * 512, N, self.D, True) <= _SBUF_BUDGET)
+    def _all_group(self, N):
+        return LNC == 1 and self.ngrp > 1 and N <= _L1_MAX_N
+    def _attn_parts(self, z1, p, tab, n_valid, B):
+        if self._all_group(z1.shape[0] // B):
+            N = z1.shape[0] // B
+            return _ab1.attention_block_l1(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"],
+                                           p["bo_full"], p["scale"], self.D, n_valid, B, self._res(_ab1, N))
+        outs = [_ab.attention_block(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo"],
+                                    p["scale"], self.D, n_valid, hb, gr, B,
+                                    int(z1.shape[0] // B <= _RESIDENT_MAX_N))
+                for hb, gr in self.plan]
+        P = outs[0] if len(outs) == 1 else torch.cat(outs, 0)
+        if P.shape[0] > 4:                                # residual/LN kernel takes <= 4 branches
+            P = P.float().sum(0, keepdim=True).to(torch.bfloat16)
+        return P
+    def _ffn(self, z2, p):
+        if p["swiglu"]:
+            if z2.shape[0] // LNC >= _SW_MIN_ROWS:
+                return nki_swiglu_stream(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
+                                         p["s_w2t"], p["s_b2"]).unsqueeze(0)
+            g = torch.nn.functional.silu((z2 @ p["s_w1t"]).float() + p["s_b1"])
+            u = (z2 @ p["s_wgt"]).float() + p["s_bg"]
+            return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0)
+        return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0)
+    def _ffn_fused(self, z2, p, x, eps_next):
+        if p["swiglu"]:
+            if z2.shape[0] // LNC < _SW_MIN_ROWS:
+                return None
+            return nki_swiglu_stream_ln(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
+                                        p["s_w2t"], p["s_b2"], x, eps_next)
+        return nki_mlp_res_ln(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], x, eps_next)
+    def __call__(self, x: torch.Tensor, tab: torch.Tensor, n_valid: int, B: int) -> torch.Tensor:
+        """x: [B*N, H] fp32 (padded rows zero). Returns the pre-final-norm residual stream [B*N, H] fp32."""
+        H = self.H
+        branch = x.new_zeros(1, x.shape[0], H, dtype=torch.bfloat16)
+        z1 = None
+        nL = len(self.L)
+        for li, p in enumerate(self.L):
+            if z1 is None:
+                x, z1 = nki_res_layernorm_fold(x, branch, p["eps1"])
+            N = x.shape[0] // B
+            if self._all_group(N):
+                x, z2 = _ab1ln.attention_block_l1_ln(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"],
+                                                     p["bo_full"], x, p["scale"], self.D, n_valid, p["eps2"], B,
+                                                     self._res(_ab1ln, N))
+            else:
+                parts = self._attn_parts(z1, p, tab, n_valid, B)
+                x, z2 = nki_res_layernorm_fold(x, parts, p["eps2"])
+            fused = self._ffn_fused(z2, p, x, self.L[li + 1]["eps1"]) if li + 1 < nL else None
+            if fused is not None:
+                x, z1 = fused
+                branch = None
+            else:
+                branch = self._ffn(z2, p)
+                z1 = None
+        if branch is not None:
+            x = x + branch[0].float()
+        return x

build/torch-neuron/layers.py ADDED Viewed

	@@ -0,0 +1,220 @@

+"""KernelConfig-mappable drop-in forwards for Hugging Face Transformers ViT encoders.
+``NeuronDINOv3ViTModel``  -> ``DINOv3ViTModel``   (all DINOv3 ViT sizes: S/S+/B/L/H+; 7B needs TP)
+``NeuronVJEPA2Encoder``   -> ``VJEPA2Encoder``    (V-JEPA2 video encoder, up to 2048 tokens / clip)
+Both run the whole transformer stack on the NKI kernels (``encoder.EncoderRunner``) and keep the
+model's own patch embedding, RoPE position computation and final LayerNorm. Weights are folded
+once, on the first device forward, from the stock module's parameters (checkpoint loading is
+untouched) and cached on the module. Anything the kernels do not support (off-device input,
+training mode, > 2048 tokens, unsupported head size) falls back to the stock forward.
+Usage (graft, works on any Transformers version)::
+    import types
+    model.forward = types.MethodType(kernels.NeuronDINOv3ViTModel.forward, model)
+    model = torch.compile(model, backend="neuron")
+``torch.compile(backend="neuron")`` is part of the intended configuration; the classes set
+``can_torch_compile = True`` so the ``kernels`` library keeps them under ``Mode.TORCH_COMPILE``.
+"""
+from __future__ import annotations
+import torch
+from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
+_PACK = "_vit_enc_neuron_packed"
+_TABS = "_vit_enc_neuron_tabs"
+_MAX_TOKENS = 2048
+def _rup(n: int, m: int) -> int:
+    return ((n + m - 1) // m) * m
+def _supported(H: int, heads: int) -> bool:
+    D = H // heads
+    return H % 128 == 0 and D in (64, 128) and H == heads * D
+def _stack(hs: torch.Tensor, N: int):
+    B, n, H = hs.shape
+    x = hs.float()
+    if N > n:
+        x = torch.cat([x, x.new_zeros(B, N - n, H)], 1)
+    return x.reshape(B * N, H).contiguous()
+# ------------------------------------------------------------------------------------------------
+# DINOv3
+# ------------------------------------------------------------------------------------------------
+def _dinov3_layers(model):
+    enc = getattr(model, "model", None)              # newer Transformers: model.model.layer
+    return enc.layer if enc is not None and hasattr(enc, "layer") else model.layer
+def _dinov3_block(lyr) -> BlockWeights:
+    a = lyr.attention
+    bw = BlockWeights(
+        ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps,
+        wq=a.q_proj.weight, bq=a.q_proj.bias, wk=a.k_proj.weight, bk=a.k_proj.bias,
+        wv=a.v_proj.weight, bv=a.v_proj.bias, wo=a.o_proj.weight, bo=a.o_proj.bias,
+        ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps,
+        ls1=lyr.layer_scale1.lambda1, ls2=lyr.layer_scale2.lambda1, scale=float(a.scaling))
+    m = lyr.mlp
+    if hasattr(m, "gate_proj"):
+        bw.gate_w, bw.gate_b = m.gate_proj.weight, m.gate_proj.bias
+        bw.up_w, bw.up_b = m.up_proj.weight, m.up_proj.bias
+        bw.down_w, bw.down_b = m.down_proj.weight, m.down_proj.bias
+    else:
+        bw.fc1_w, bw.fc1_b = m.up_proj.weight, m.up_proj.bias
+        bw.fc2_w, bw.fc2_b = m.down_proj.weight, m.down_proj.bias
+    return bw
+def _dinov3_ok(model) -> bool:
+    c = model.config
+    if not _supported(c.hidden_size, c.num_attention_heads):
+        return False
+    if getattr(c, "use_gated_mlp", False):
+        return getattr(c, "hidden_act", "silu") == "silu"
+    return getattr(c, "hidden_act", "gelu") == "gelu"
+@torch._dynamo.disable
+def _dinov3_freeze(model):
+    if model.training:
+        raise RuntimeError("NeuronDINOv3ViTModel: these kernels have no autograd rule; call model.eval() first.")
+    dev = model.norm.weight.device
+    heads = model.config.num_attention_heads
+    packed = [pack_block(_dinov3_block(l), heads, "half", dev) for l in _dinov3_layers(model)]
+    setattr(model, _PACK, EncoderRunner(packed))
+    setattr(model, _TABS, {})
+@torch._dynamo.disable
+def _dinov3_tab(model, pixel_values, n_prefix, N):
+    key = (tuple(pixel_values.shape[-2:]), N)
+    tabs = getattr(model, _TABS)
+    if key not in tabs:
+        cos, sin = model.rope_embeddings(pixel_values[:1].float())        # [n_patches, D] each
+        c = model.config
+        tabs[key] = rope_tables(cos, sin, n_prefix, N, c.num_attention_heads,
+                                c.hidden_size // c.num_attention_heads, "half", model.norm.weight.device)
+    return tabs[key]
+class NeuronDINOv3ViTModel(torch.nn.Module):
+    """Drop-in forward for ``transformers.DINOv3ViTModel`` (bf16 NKI encoder, fp32 residual stream).
+    Returns ``last_hidden_state`` (final-normed, ``[B, 1 + registers + patches, H]``) and
+    ``pooler_output`` (CLS), like the stock model. Any batch size; images up to 2048 tokens
+    (e.g. 720x720 at patch 16).
+    """
+    can_torch_compile = True
+    has_backward = False
+    def forward(self, pixel_values, bool_masked_pos=None, head_mask=None, **kwargs):
+        from transformers.modeling_outputs import BaseModelOutputWithPooling
+        pv = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
+        hs = self.embeddings(pv, bool_masked_pos=bool_masked_pos)           # [B, n, H]
+        B, n, H = hs.shape
+        N = _rup(n, 128 * LNC)
+        if (hs.device.type != "neuron" or self.training or head_mask is not None or N > _MAX_TOKENS
+                or not _dinov3_ok(self)):
+            import transformers.models.dinov3_vit.modeling_dinov3_vit as mv
+            if head_mask is not None:
+                kwargs["head_mask"] = head_mask
+            return mv.DINOv3ViTModel.forward(self, pixel_values, bool_masked_pos=bool_masked_pos, **kwargs)
+        if not hasattr(self, _PACK):
+            _dinov3_freeze(self)
+        n_prefix = 1 + self.config.num_register_tokens
+        tab = _dinov3_tab(self, pixel_values, n_prefix, N)
+        x = getattr(self, _PACK)(_stack(hs, N), tab, n, B)
+        x = x.reshape(B, N, H)[:, :n]
+        out = self.norm(x.to(self.norm.weight.dtype))
+        return BaseModelOutputWithPooling(last_hidden_state=out, pooler_output=out[:, 0, :])
+# ------------------------------------------------------------------------------------------------
+# V-JEPA2
+# ------------------------------------------------------------------------------------------------
+def _vjepa2_block(lyr) -> BlockWeights:
+    a = lyr.attention
+    return BlockWeights(
+        ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps,
+        wq=a.query.weight, bq=a.query.bias, wk=a.key.weight, bk=a.key.bias,
+        wv=a.value.weight, bv=a.value.bias, wo=a.proj.weight, bo=a.proj.bias,
+        ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps,
+        fc1_w=lyr.mlp.fc1.weight, fc1_b=lyr.mlp.fc1.bias, fc2_w=lyr.mlp.fc2.weight, fc2_b=lyr.mlp.fc2.bias,
+        scale=float(a.scaling))
+@torch._dynamo.disable
+def _vjepa2_tab(encoder, N, n):
+    """Signed-sin 3D RoPE tables from the model's own position ids (interleaved pairing)."""
+    a = encoder.layer[0].attention
+    H, D = a.all_head_size, a.attention_head_size
+    dev = a.query.weight.device
+    pos = a.get_position_ids(torch.zeros(1, n, H, device=dev), masks=None)
+    to1d = lambda t: (t if t.dim() == 1 else t.reshape(-1, n)[0]).float().cpu()
+    cos = torch.ones(n, D)
+    sin = torch.zeros(n, D)
+    s = 0
+    for sub, p in zip((a.d_dim, a.h_dim, a.w_dim), (to1d(pos[0]), to1d(pos[1]), to1d(pos[2]))):
+        half = sub // 2
+        omega = 1.0 / (10000.0 ** (torch.arange(half, dtype=torch.float32) / float(half)))
+        f = p[:, None] * omega
+        c, sn = f.cos(), f.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(D); g[0::2] = -1.0
+    return rope_tables(cos, sin * g, 0, N, H // D, D, "interleaved", dev)
+@torch._dynamo.disable
+def _vjepa2_freeze(encoder):
+    if encoder.training:
+        raise RuntimeError("NeuronVJEPA2Encoder: these kernels have no autograd rule; call model.eval() first.")
+    a = encoder.layer[0].attention
+    dev = a.query.weight.device
+    heads = a.all_head_size // a.attention_head_size
+    setattr(encoder, _PACK, EncoderRunner([pack_block(_vjepa2_block(l), heads, "interleaved", dev)
+                                           for l in encoder.layer]))
+    setattr(encoder, _TABS, {})
+class NeuronVJEPA2Encoder(torch.nn.Module):
+    """Drop-in forward for ``transformers.VJEPA2Encoder``.
+    Any batch size; clips up to 2048 tokens (e.g. 16 frames at 256 px, tubelet 2, patch 16).
+    """
+    can_torch_compile = True
+    has_backward = False
+    def forward(self, pixel_values_videos=None, head_mask=None, output_attentions=False,
+                output_hidden_states=False, **kwargs):
+        from transformers.modeling_outputs import BaseModelOutput
+        hs = self.embeddings(pixel_values_videos)                           # [B, n, H]
+        B, n, H = hs.shape
+        N = _rup(n, 128 * LNC)
+        a = self.layer[0].attention
+        if (hs.device.type != "neuron" or self.training or head_mask is not None or output_attentions
+                or output_hidden_states or N > _MAX_TOKENS
+                or not _supported(H, a.all_head_size // a.attention_head_size)):
+            import transformers.models.vjepa2.modeling_vjepa2 as mv
+            return mv.VJEPA2Encoder.forward(self, pixel_values_videos, head_mask=head_mask,
+                                            output_attentions=output_attentions,
+                                            output_hidden_states=output_hidden_states, **kwargs)
+        if not hasattr(self, _PACK):
+            _vjepa2_freeze(self)
+        tabs = getattr(self, _TABS)
+        if N not in tabs:
+            tabs[N] = _vjepa2_tab(self, N, n)
+        x = getattr(self, _PACK)(_stack(hs, N), tabs[N], n, B)
+        out = self.layernorm(x.reshape(B, N, H)[:, :n])
+        return BaseModelOutput(last_hidden_state=out)

build/torch-neuron/metadata.json ADDED Viewed

	@@ -0,0 +1 @@


1	+ { "name": "vit-encoder-neuron-kernels", "id": "vit_encoder_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,23 @@

+"""The NKI kernels and their host-side weight packing."""
+from .attention_block import (
+    attention_block,
+    lane_permutation,
+    pack_out_weights,
+    pack_tables,
+    pack_v_weights,
+    pack_weights,
+)
+from .attention_block_l1 import attention_block_l1
+from .attention_block_l1_ln import attention_block_l1_ln
+from .mlp import nki_mlp_k4_xp as mlp
+from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
+from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
+from .swiglu_stream import nki_swiglu_stream as swiglu_stream
+from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
+__all__ = [
+    "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
+    "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
+    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
+]

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

	@@ -0,0 +1,801 @@

+"""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
+the repo-root incumbent 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 incumbent (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, the V-JEPA2 kernels' verification notes), 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 the V-JEPA2 kernels' verification notes). 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: the V-JEPA2 kernels' verification notes).
+# 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
+import os as _os
+#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
+LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
+_SPLIT = 2  # heads are always split in two halves (8 heads per program)
+_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, rope: str = "half") -> torch.Tensor:
+    """Column order that puts each quad-of-heads' RoPE pair halves in their own 128-row group.
+    RoPE rotates pairs of lanes ``(a, b)`` of each head. ``rope`` picks the pairing:
+      * ``"half"``        -- (j, j + D/2), the rotate-half convention (DINOv3, EVA-02, HF Llama style).
+      * ``"interleaved"`` -- (2j, 2j + 1), adjacent lanes (V-JEPA2 / GPT-NeoX-interleaved style).
+    The projection's output columns are reordered so group ``2a`` holds the FIRST lane of every
+    pair for heads ``4a..4a+3`` (at D=64; 2 groups per 4 heads at D=128) and group ``2a+1`` the
+    SECOND lane, so RoPE is six full-width elementwise ops with no transposes. The permutation
+    cancels in the Q/K contraction, so nothing downstream un-permutes.
+    Returns ``perm`` of length ``heads*D`` s.t. ``w.t()[:, perm]`` is the reordered weight.
+    """
+    _kernel_assert(rope in ("half", "interleaved"), f"rope must be 'half' or 'interleaved', got {rope!r}")
+    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):
+            lo, hi = (j, half + j) if rope == "half" else (2 * j, 2 * j + 1)
+            perm[(2 * a) * _P_MAX + b * half + j] = hd * D + lo
+            perm[(2 * a + 1) * _P_MAX + b * half + j] = hd * D + hi
+    return perm
+def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
+                 bq: torch.Tensor | None, bk: torch.Tensor | None,
+                 heads: int, D: int, rope: str = "half"):
+    """``(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, rope).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, sin: torch.Tensor, heads: int, D: int, rope: str = "half"):
+    """``tab [4, 128, N] fp32`` RoPE tables, quad-tiled to match ``lane_permutation(..., rope)``.
+    Inputs are the model's own ``[N, D]`` cos / sin tables (shared by every head; prefix tokens that
+    take no rotation: cos=1, sin=0). With ``(a, b)`` a lane pair the kernel computes
+    ``a' = a*tab0 + b*tab2`` and ``b' = b*tab1 + a*tab3``, so:
+      * ``"half"``:        ``sin`` is the PLAIN sin table (rotate-half: ``x*cos + rotate_half(x)*sin``);
+                           tab = (cos_lo, cos_hi, -sin_lo, +sin_hi).
+      * ``"interleaved"``: ``sin`` is the SIGNED sin table (``-sin`` on even lanes, ``+sin`` on odd);
+                           tab = (cos_even, cos_odd, sinsigned_even, sinsigned_odd).
+    """
+    _kernel_assert(rope in ("half", "interleaved"), f"rope must be 'half' or 'interleaved', got {rope!r}")
+    half = D // 2
+    quad = _P_MAX // half
+    if rope == "half":
+        parts = [cos[:, :half], cos[:, half:], -sin[:, :half], sin[:, half:]]
+    else:
+        parts = [cos[:, 0::2], cos[:, 1::2], sin[:, 0::2], sin[:, 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, n_valid, hb, grid, nimg=1, resident=0):
+    """out[16, N, D] bf16 = attention(rope(h@Wq), rope(h@Wk), h@Wv), one kernel (per-core partials).
+    BATCHED (nimg>1): h_hbm is the STACKED [nimg*N, H] buffer (images back to back, N rows each).
+    The per-image body below runs once per image inside this ONE launch, reading image b's rows at
+    row offset b*N and writing partials to result[slot, b*N:(b+1)*N]. ONE launch, no host-side
+    slice/cat glue.
+    resident=0: the loop wraps the whole body, so weights/biases/rope tables are re-DMA'd per image.
+    resident=1 (N<=2048): Q/K/V/O weights, biases and rope tables are staged ONCE at the bottom of
+    SBUF before the image loop (~38 KiB/lane ViT-L + tables); each image's buffers are allocated
+    above them. Same math, bit-identical outputs."""
+    NT, H = h_hbm.shape
+    N = NT // nimg
+    Fp = wqk_hbm.shape[2]                 # padded attention features = 512 * n_groups
+    D = head_dim
+    half = D // 2
+    quad = _P_MAX // half                 # heads per 128-partition parity group = 4
+    heads = Fp // D                       # padded head count (multiple of 8)
+    local_heads = _PSUM_FMAX // D         # heads per program: 8 at D=64, 4 at D=128
+    local_quads = local_heads // quad     # 2
+    n_chunks = H // _P_MAX                # 8 contraction chunks
+    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX   # token tiles of 512 (last may be ragged)
+    n_ktiles = (N + _K_TILE - 1) // _K_TILE     # K tiles / head (last may be partially filled)
+    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 the V-JEPA2 kernels' verification notes.
+    _kernel_assert(wqk_hbm.shape == (2, H, Fp), f"wqk must be [2,{H},{Fp}]")
+    _kernel_assert(bqk_hbm.shape == (2, Fp, 1), f"bqk must be [2,{Fp},1]")
+    _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), f"tab must be [4,{_P_MAX},{N}]")
+    _kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
+    _kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
+    _kernel_assert(H % _P_MAX == 0, f"H={H} must be a multiple of 128")
+    _kernel_assert(quad * half == _P_MAX, f"D={D} must give 4 heads per group")
+    _kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512")
+    _kernel_assert(local_heads * D == _PSUM_FMAX,
+                   "inline V projection wants all local heads in one 512-wide stream")
+    _kernel_assert(N % _P_MAX == 0, f"N={N} must be a multiple of 128")
+    _kernel_assert(NT == nimg * N, f"rows {NT} must be nimg({nimg}) x N")
+    _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((grid, NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    # Each program handles one 8-head group: group = hb + program_id. grid==1 -> single program.
+    if grid == 1:
+        pid = hb
+        out_slot = 0
+    else:
+        out_slot = nl.program_id(0)       # SYMBOLIC: index arithmetic only
+        pid = hb + out_slot
+    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)
+    # ---- RESIDENT MODE (resident=1, N<=2048 only): stage the per-launch-invariant inputs ONCE,
+    # at the BOTTOM of SBUF, before the image loop. Everything per image is then allocated above
+    # `res_top`. Invariant: Q/K/V/O weights, their biases, and the rope tables (same for every image).
+    # Cost ~38 KiB/lane (ViT-L) + tables; per-image phase peaks were <=106 KiB, budget 192 KiB.
+    _res = bool(resident) and not (N > _PSUM_FMAX * 4)
+    res_top = 0
+    if _res:
+        ra = ModularAllocator(initial_address=0)
+        r_w_sb = ra.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=r_w_sb[qk][c],
+                              src=wqk_hbm[qk, c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_quads * 2 * _P_MAX])
+        r_wv_sb = ra.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=r_wv_sb[c], src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, vcol0:vcol0 + local_heads * D])
+        r_b_sb = ra.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)
+        r_b_stage = ra.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=r_b_stage, src=bqk_hbm[qk, col0 + g * _P_MAX:col0 + (g + 1) * _P_MAX, 0:1])
+                nisa.tensor_copy(dst=r_b_sb[qk][g], src=r_b_stage)
+        r_bv_bc = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
+        nisa.dma_copy(dst=r_bv_bc, src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=vcol0))
+        r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                     block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+        for s in range(n_tt):
+            for t_i in range(4):
+                _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw], src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        r_wo_sb = ra.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=r_wo_sb[oc], src=wot_hbm[col0 + oc * _P_MAX:col0 + (oc + 1) * _P_MAX, 0:H])
+        r_bo_bc = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+        nisa.dma_copy(dst=r_bo_bc, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
+        res_top = ra.get_current_address()
+    for img in range(nimg):              # unrolled at trace time (python int)
+        rofs = img * N * H                # element offset of image img's rows in h / result
+        allocator = ModularAllocator(initial_address=res_top)
+        # ---- 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, min(_PSUM_FMAX, N - s * _PSUM_FMAX)]], offset=0),
+                        src=h_hbm.ap([[H, min(_PSUM_FMAX, N - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
+                                     offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
+                    )
+        if _res:
+            w_sb, wv_sb, b_sb, bv_bc = r_w_sb, r_wv_sb, r_b_sb, r_bv_bc
+        else:
+            # 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.
+        if _res:
+            tab = r_tab
+        else:
+            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):
+                        _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                        nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
+                                      src=tab_hbm[t_i, 0:_P_MAX,
+                                                  s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        # ---- 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
+            tw = min(_PSUM_FMAX, N - t0)          # ragged last token tile (N % 512 != 0)
+            # 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=rofs + 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(tw // _V_TILE):
+                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[:, 0:tw],
+                                stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
+                                moving=x_t[c][tt][:, 0:tw],
+                                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. ========
+                        # NKI 0.7: a Scalar-engine pin here fails ISA opcode_engine_check on newer compilers
+                        # TOT compiler -> unpinned (compiler picks the engine).
+                        nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
+                                           op0=nl.add, operand0=b_sb[qk][g])
+                    # rope: the six ops of nki_rope, same order, tables sliced per tile
+                    nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
+                                       data2=tab[0][tt][:, 0:tw], op=nl.multiply)
+                    nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
+                                       data2=tab[2][tt][:, 0:tw], op=nl.multiply)
+                    nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw], data2=tmp[:, 0:tw], op=nl.add)
+                    nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
+                                       data2=tab[1][tt][:, 0:tw], op=nl.multiply)
+                    nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
+                                       data2=tab[3][tt][:, 0:tw], op=nl.multiply)
+                    nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw], data2=tmp[:, 0:tw], op=nl.add)
+                    # 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:tw],
+                                    src=roped[par2][r0:r0 + half, 0:tw])
+                        else:
+                            for par2 in range(2):
+                                nisa.dma_copy(
+                                    dst=q_priv[lh][0, par2 * half:(par2 + 1) * half,
+                                                   t0:t0 + tw],
+                                    src=roped[par2][r0:r0 + half, 0:tw])
+        # ---- PHASE 2: reclaim the dead producer buffers. NO ALIASING. -------------------
+        allocator.set_current_address(phase2_base)
+        if _res:
+            wo_sb, bo_bc = r_wo_sb, r_bo_bc
+        else:
+            # ---- 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,
+                    seqlen_k_valid=n_valid,
+                )
+            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,
+                    seqlen_k_valid=n_valid,
+                )
+            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)
+        n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
+        for tb in range(N // _P_MAX):
+            for ob in range(n_ob):
+                ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX)   # ragged last column block (H%512!=0)
+                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[:, 0:ow],
+                        stationary=o_feat[oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
+                        moving=wo_sb[oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
+                        accumulate=(oc > 0),
+                    )
+                nisa.tensor_tensor(dst=out_sb[:, 0:ow], data1=ps_out[:, 0:ow],
+                                   data2=bo_bc[:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
+                                   op=nl.add)
+                nisa.dma_copy(
+                    dst=result.ap(pattern=[[H, _P_MAX], [1, ow]],
+                                  offset=out_slot * NT * H + rofs + tb * _P_MAX * H + ob * _PSUM_FMAX),
+                    src=out_sb[:, 0:ow])
+    return result
+_ATTN = {g: wrap_nki_to_torch_hop(_attention_block_kernel[g]) for g in sorted({1, 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("vit_enc::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,
+    n_valid: int,
+    hb: int = 0,
+    grid: int = LNC,
+    nimg: int = 1,
+    resident: int = 0,
+) -> 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 the V-JEPA2 kernels' verification notes. NKI kernels cannot `raise`, so the guard lives here.
+    N = h.shape[0] // nimg
+    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"the V-JEPA2 kernels' verification notes. Use N<=2048 (up to 16-frame clips) -- the bit-for-bit "
+            f"shipping path."
+        )
+    return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident)

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

	@@ -0,0 +1,402 @@

+"""Attention for LNC=1: ONE program over ALL head groups.
+The 2-core kernel in ``attention_block.py`` owns one 8-head group per program. At LNC=1 that
+means one launch per group per layer (2 for ViT-L, 3 for ViT-H+), each re-transposing x, each
+re-loading the rope tables, each writing an [N, H] bf16 out-projection PARTIAL to HBM that the
+LN-fold kernel then sums.
+This kernel, per image (``nimg`` images stacked on the row axis, one launch):
+  * transposes x once and loads the rope tables once (shared by every group: tables are
+    quad-tiled per head, identical for all heads);
+  * for each group: projects Q/K/V + RoPE (same producer as attention_block), runs the vendor
+    attention per head, and writes O feature-major into a PERSISTENT per-group SBUF slab;
+  * after all groups: ONE out-projection contracting over ALL features (G*4 chunks accumulate
+    in one fp32 PSUM tile), adds the FULL bias once, writes ONE [N, H] result.
+No partials, no host sum; the LN-fold kernel receives a single branch.
+``resident=1``: Q/K/V/O weights, biases and rope tables are staged ONCE per launch (bottom of
+SBUF) and reused by every image. Host decides by an SBUF estimate (does not fit at N>=1024 for
+ViT-L); ``resident=0`` restages each group's weights per image.
+N<=2048 only (the shipping single-section vendor path). Numerics: identical producer and vendor
+attention; the out-projection now accumulates all groups in fp32 PSUM instead of summing bf16
+per-group partials, so results differ from attention_block by bf16 rounding of the partials.
+"""
+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
+from .attention_inner import _attention_cte
+_P_MAX = 128
+_PSUM_FMAX = 512
+_PSUM_BANK = 2048
+_K_TILE = 512
+_V_TILE = 128
+def _kernel_assert(cond: bool, msg: str):
+    assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
+def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int:
+    """Conservative bytes/lane at the kernel's peak (host-side decision helper, not exact)."""
+    G = Fp // _PSUM_FMAX
+    n_chunks = H // _P_MAX
+    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
+    n_ktiles = (N + _K_TILE - 1) // _K_TILE
+    lh = _PSUM_FMAX // D
+    n_och = (lh * D) // _P_MAX
+    w_grp = 3 * n_chunks * _PSUM_FMAX * 2 + _PSUM_FMAX * 4 + 64     # wqk + wv + bv_bc + biases
+    w_out = G * n_och * H * 2 + H * 4                                # wo + bo_bc
+    tab = 4 * n_tt * _PSUM_FMAX * 4
+    o_feat = G * n_och * N * 2
+    x_t = n_chunks * n_tt * _PSUM_FMAX * 2
+    kv = lh * (n_ktiles * _K_TILE * 2 + (N // _V_TILE) * D * 2)
+    work = 2 * _PSUM_FMAX * 2 * 2 + 2 * _PSUM_FMAX * 4 + _PSUM_FMAX * 4
+    inner = 48 * 1024
+    resid = (G * w_grp + w_out + tab) if resident else 0
+    base = resid + o_feat + x_t + (0 if resident else tab)
+    p1 = kv + (0 if resident else w_grp) + work
+    p2 = kv + inner
+    p3 = (0 if resident else w_out) + 2048
+    return base + max(p1, p2, p3)
+@nki.jit
+def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
+                               wot_hbm, bo_hbm, scale, head_dim, n_valid, nimg=1, resident=0):
+    NT, H = h_hbm.shape
+    N = NT // nimg
+    Fp = wqk_hbm.shape[2]
+    D = head_dim
+    half = D // 2
+    quad = _P_MAX // half
+    local_heads = _PSUM_FMAX // D
+    local_quads = local_heads // quad
+    G = Fp // _PSUM_FMAX                  # head groups, ALL handled by this one program
+    n_chunks = H // _P_MAX
+    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
+    n_ktiles = (N + _K_TILE - 1) // _K_TILE
+    n_vtiles = N // _V_TILE
+    n_ochunks = (local_heads * D) // _P_MAX   # feature chunks per group
+    v_per_tt = _PSUM_FMAX // _V_TILE
+    _kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
+    _kernel_assert(bqk_hbm.shape == (2, Fp, 1), "bad bqk")
+    _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
+    _kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
+    _kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
+    _kernel_assert(H % _P_MAX == 0, "H must be a multiple of 128")
+    _kernel_assert(quad * half == _P_MAX, "D must give quad*half=128")
+    _kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512")
+    _kernel_assert(N % _P_MAX == 0, "N must be a multiple of 128")
+    _kernel_assert(N <= _PSUM_FMAX * 4, "N>2048 not supported")
+    _kernel_assert(NT == nimg * N, "rows must be nimg x N")
+    result = nl.ndarray((1, NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    # ---------------- resident (per-launch invariant) staging ----------------
+    res_top = 0
+    r_w = []
+    r_wv = []
+    r_b = []
+    r_bv = []
+    if resident:
+        ra = ModularAllocator(initial_address=0)
+        for g in range(G):
+            col0 = g * _PSUM_FMAX
+            w_sb_ = ra.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_ = ra.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, col0:col0 + local_heads * D])
+            b_sb_ = ra.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_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.bfloat16)
+            for qk in range(2):
+                for gg in range(local_quads * 2):
+                    nisa.dma_copy(dst=b_stage_,
+                                  src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
+                    nisa.tensor_copy(dst=b_sb_[qk][gg], src=b_stage_)
+            bv_bc_ = ra.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=col0))
+            r_w.append(w_sb_)
+            r_wv.append(wv_sb_)
+            r_b.append(b_sb_)
+            r_bv.append(bv_bc_)
+        r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                     block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+        for s in range(n_tt):
+            for t_i in range(4):
+                _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw],
+                              src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        r_wo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
+                                    block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
+        for oc in range(G * n_ochunks):
+            nisa.dma_copy(dst=r_wo[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
+        r_bo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+        nisa.dma_copy(dst=r_bo, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
+        res_top = ra.get_current_address()
+    n_grp = N // _P_MAX
+    for img in range(nimg):
+        rofs = img * N * H
+        allocator = ModularAllocator(initial_address=res_top)
+        # O for ALL groups stays in SBUF until the single out-projection.
+        o_feat = []
+        for g in range(G):
+            of_ = allocator.alloc_sbuf_tensor(shape=(_P_MAX, N), dtype=nl.bfloat16,
+                                              block_dim=[n_ochunks], num_free_tiles=[n_ochunks],
+                                              align_to=32)
+            o_feat.append(of_)
+        # x^T: shared by every group (transposed ONCE per image).
+        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, n_tt],
+                                          align_to=32)
+        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, min(_PSUM_FMAX, N - s * _PSUM_FMAX)]], offset=0),
+                    src=h_hbm.ap([[H, min(_PSUM_FMAX, N - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
+                                 offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
+                )
+        # rope tables: shared by every group.
+        if resident:
+            tab = r_tab
+        else:
+            tab = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                              block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+            for s in range(n_tt):
+                for t_i in range(4):
+                    _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                    nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
+                                  src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        img_base = allocator.get_current_address()
+        for g in range(G):
+            allocator.set_current_address(img_base)
+            col0 = g * _PSUM_FMAX
+            k_tiles = []
+            v_tiles = []
+            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_)
+            phase2_base = allocator.get_current_address()
+            if resident:
+                w_sb = r_w[g]
+                wv_sb = r_wv[g]
+                b_sb = r_b[g]
+                bv_bc = r_bv[g]
+            else:
+                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, col0:col0 + local_heads * D])
+                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 gg in range(local_quads * 2):
+                        nisa.dma_copy(dst=b_stage,
+                                      src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
+                        nisa.tensor_copy(dst=b_sb[qk][gg], 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=col0))
+            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_priv = []
+            for _lh in range(local_heads):
+                q_priv.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
+            # ---- producer (identical to attention_block's N<=2048 path) ----
+            for tt in range(n_tt):
+                t0 = tt * _PSUM_FMAX
+                tw = min(_PSUM_FMAX, N - t0)
+                for vsub in range(tw // _V_TILE):
+                    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):
+                        nisa.tensor_copy(dst=v_tiles[lh][tt * v_per_tt + vsub],
+                                         src=vout[:, lh * D:(lh + 1) * D], engine=engine.scalar)
+                for qk in range(2):
+                    for j in range(local_quads):
+                        for par in range(2):
+                            gq = j * 2 + par
+                            cols = gq * _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[:, 0:tw],
+                                               stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
+                                               moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
+                            nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
+                                               op0=nl.add, operand0=b_sb[qk][gq])
+                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
+                                           data2=tab[0][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
+                                           data2=tab[2][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw],
+                                           data2=tmp[:, 0:tw], op=nl.add)
+                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
+                                           data2=tab[1][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
+                                           data2=tab[3][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw],
+                                           data2=tmp[:, 0:tw], op=nl.add)
+                        for b in range(quad):
+                            lh = j * quad + b
+                            r0 = b * half
+                            if qk == 1:
+                                for par2 in range(2):
+                                    nisa.tensor_copy(
+                                        dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half, 0:tw],
+                                        src=roped[par2][r0:r0 + half, 0:tw])
+                            else:
+                                for par2 in range(2):
+                                    nisa.dma_copy(
+                                        dst=q_priv[lh][0, par2 * half:(par2 + 1) * half, t0:t0 + tw],
+                                        src=roped[par2][r0:r0 + half, 0:tw])
+            # ---- attention for this group's heads -> o_feat[g] (SBUF, feature-major) ----
+            allocator.set_current_address(phase2_base)
+            o_sb = []
+            for _lh in range(local_heads):
+                oc = (_lh * D) // _P_MAX
+                rr = (_lh * D) % _P_MAX
+                grp = []
+                for _g in range(n_grp):
+                    grp.append(o_feat[g][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):
+                _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,
+                    seqlen_k_valid=n_valid,
+                )
+                allocator.set_current_address(attn_base)
+        # ---- ONE out-projection over ALL G*n_ochunks feature chunks, full bias once ----
+        allocator.set_current_address(img_base)
+        if resident:
+            wo_sb = r_wo
+            bo_bc = r_bo
+        else:
+            wo_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
+                                                block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
+            for oc in range(G * n_ochunks):
+                nisa.dma_copy(dst=wo_sb[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
+            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))
+        out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16, align_to=32)
+        n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
+        for tb in range(n_grp):
+            for ob in range(n_ob):
+                ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX)
+                ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
+                                    address=(0, 0))
+                for g in range(G):
+                    for oc in range(n_ochunks):
+                        nisa.nc_matmul(
+                            dst=ps_out[:, 0:ow],
+                            stationary=o_feat[g][oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
+                            moving=wo_sb[g * n_ochunks + oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
+                            accumulate=((g > 0) or (oc > 0)),
+                        )
+                nisa.tensor_tensor(dst=out_sb[:, 0:ow], data1=ps_out[:, 0:ow],
+                                   data2=bo_bc[:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow], op=nl.add)
+                nisa.dma_copy(
+                    dst=result.ap(pattern=[[H, _P_MAX], [1, ow]],
+                                  offset=rofs + tb * _P_MAX * H + ob * _PSUM_FMAX),
+                    src=out_sb[:, 0:ow])
+    return result
+_ATTN_L1 = wrap_nki_to_torch_hop(_attention_block_l1_kernel[1])
+@nki_op("vit_enc::attention_block_l1", mutates_args={})
+def attention_block_l1(
+    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,
+    n_valid: int,
+    nimg: int = 1,
+    resident: int = 0,
+) -> torch.Tensor:
+    """All head groups in one program. Returns ONE [1, nimg*N, H] bf16 result (full bias)."""
+    for t in (h, wqk, bqk, tab, wvt, bv, wot, bo):
+        if t is not None and t.requires_grad:
+            raise RuntimeError("attention_block_l1 is inference-only; use torch.no_grad().")
+    if h.shape[0] // nimg > 2048:
+        raise NotImplementedError("attention_block_l1 supports N<=2048 only")
+    return _ATTN_L1(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, nimg, resident)

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

	@@ -0,0 +1,451 @@

+"""Attention for LNC=1, all head groups, with the residual add + LN2 fused into the epilogue.
+Same body as ``attention_block_l1``. The epilogue differs: per 128-token block the
+out-projection for ALL H columns lands in an fp32 SBUF row tile, then
+    x_new = x + (attn @ Wo + bo)          (fp32; attention output never rounded to bf16)
+    z2    = LN_no_affine(x_new)           (bf16; same stat chain as res_layernorm.py)
+and both are written to HBM. This replaces K3 (one launch + a bf16 [N,H] HBM write/read per layer).
+--- attention_block_l1 notes ---
+Attention for LNC=1: ONE program over ALL head groups.
+The 2-core kernel in ``attention_block.py`` owns one 8-head group per program. At LNC=1 that
+means one launch per group per layer (2 for ViT-L, 3 for ViT-H+), each re-transposing x, each
+re-loading the rope tables, each writing an [N, H] bf16 out-projection PARTIAL to HBM that the
+LN-fold kernel then sums.
+This kernel, per image (``nimg`` images stacked on the row axis, one launch):
+  * transposes x once and loads the rope tables once (shared by every group: tables are
+    quad-tiled per head, identical for all heads);
+  * for each group: projects Q/K/V + RoPE (same producer as attention_block), runs the vendor
+    attention per head, and writes O feature-major into a PERSISTENT per-group SBUF slab;
+  * after all groups: ONE out-projection contracting over ALL features (G*4 chunks accumulate
+    in one fp32 PSUM tile), adds the FULL bias once, writes ONE [N, H] result.
+No partials, no host sum; the LN-fold kernel receives a single branch.
+``resident=1``: Q/K/V/O weights, biases and rope tables are staged ONCE per launch (bottom of
+SBUF) and reused by every image. Host decides by an SBUF estimate (does not fit at N>=1024 for
+ViT-L); ``resident=0`` restages each group's weights per image.
+N<=2048 only (the shipping single-section vendor path). Numerics: identical producer and vendor
+attention; the out-projection now accumulates all groups in fp32 PSUM instead of summing bf16
+per-group partials, so results differ from attention_block by bf16 rounding of the partials.
+"""
+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
+from .attention_inner import _attention_cte
+_P_MAX = 128
+_PSUM_FMAX = 512
+_PSUM_BANK = 2048
+_K_TILE = 512
+_V_TILE = 128
+def _kernel_assert(cond: bool, msg: str):
+    assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
+def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int:
+    """Conservative bytes/lane at the kernel's peak (host-side decision helper, not exact)."""
+    G = Fp // _PSUM_FMAX
+    n_chunks = H // _P_MAX
+    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
+    n_ktiles = (N + _K_TILE - 1) // _K_TILE
+    lh = _PSUM_FMAX // D
+    n_och = (lh * D) // _P_MAX
+    w_grp = 3 * n_chunks * _PSUM_FMAX * 2 + _PSUM_FMAX * 4 + 64     # wqk + wv + bv_bc + biases
+    w_out = G * n_och * H * 2 + H * 4                                # wo + bo_bc
+    tab = 4 * n_tt * _PSUM_FMAX * 4
+    o_feat = G * n_och * N * 2
+    x_t = n_chunks * n_tt * _PSUM_FMAX * 2
+    kv = lh * (n_ktiles * _K_TILE * 2 + (N // _V_TILE) * D * 2)
+    work = 2 * _PSUM_FMAX * 2 * 2 + 2 * _PSUM_FMAX * 4 + _PSUM_FMAX * 4
+    inner = 48 * 1024
+    resid = (G * w_grp + w_out + tab) if resident else 0
+    base = resid + o_feat + x_t + (0 if resident else tab)
+    p1 = kv + (0 if resident else w_grp) + work
+    p2 = kv + inner
+    p3 = (0 if resident else w_out) + 2 * (H * 4 * 2 + H * 2 + 32) + H * 4
+    return base + max(p1, p2, p3)
+@nki.jit
+def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
+                                  wot_hbm, bo_hbm, x_hbm, scale, head_dim, n_valid, eps, nimg=1, resident=0):
+    NT, H = h_hbm.shape
+    N = NT // nimg
+    Fp = wqk_hbm.shape[2]
+    D = head_dim
+    half = D // 2
+    quad = _P_MAX // half
+    local_heads = _PSUM_FMAX // D
+    local_quads = local_heads // quad
+    G = Fp // _PSUM_FMAX                  # head groups, ALL handled by this one program
+    n_chunks = H // _P_MAX
+    n_tt = (N + _PSUM_FMAX - 1) // _PSUM_FMAX
+    n_ktiles = (N + _K_TILE - 1) // _K_TILE
+    n_vtiles = N // _V_TILE
+    n_ochunks = (local_heads * D) // _P_MAX   # feature chunks per group
+    v_per_tt = _PSUM_FMAX // _V_TILE
+    _kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
+    _kernel_assert(bqk_hbm.shape == (2, Fp, 1), "bad bqk")
+    _kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
+    _kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
+    _kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
+    _kernel_assert(H % _P_MAX == 0, "H must be a multiple of 128")
+    _kernel_assert(quad * half == _P_MAX, "D must give quad*half=128")
+    _kernel_assert(Fp % _PSUM_FMAX == 0, "padded features must be a multiple of 512")
+    _kernel_assert(N % _P_MAX == 0, "N must be a multiple of 128")
+    _kernel_assert(N <= _PSUM_FMAX * 4, "N>2048 not supported")
+    _kernel_assert(NT == nimg * N, "rows must be nimg x N")
+    _kernel_assert(x_hbm.shape == (NT, H), "x must be [nimg*N, H] fp32")
+    xnew_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
+    z_hbm = nl.ndarray((NT, H), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    inv_h = 1.0 / float(H)
+    neg_inv_h = -inv_h
+    # ---------------- resident (per-launch invariant) staging ----------------
+    res_top = 0
+    r_w = []
+    r_wv = []
+    r_b = []
+    r_bv = []
+    if resident:
+        ra = ModularAllocator(initial_address=0)
+        for g in range(G):
+            col0 = g * _PSUM_FMAX
+            w_sb_ = ra.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_ = ra.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, col0:col0 + local_heads * D])
+            b_sb_ = ra.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_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.bfloat16)
+            for qk in range(2):
+                for gg in range(local_quads * 2):
+                    nisa.dma_copy(dst=b_stage_,
+                                  src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
+                    nisa.tensor_copy(dst=b_sb_[qk][gg], src=b_stage_)
+            bv_bc_ = ra.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=col0))
+            r_w.append(w_sb_)
+            r_wv.append(wv_sb_)
+            r_b.append(b_sb_)
+            r_bv.append(bv_bc_)
+        r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                     block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+        for s in range(n_tt):
+            for t_i in range(4):
+                _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                nisa.dma_copy(dst=r_tab[t_i][s][:, 0:_tw],
+                              src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        r_wo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
+                                    block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
+        for oc in range(G * n_ochunks):
+            nisa.dma_copy(dst=r_wo[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
+        r_bo = ra.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+        nisa.dma_copy(dst=r_bo, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
+        res_top = ra.get_current_address()
+    n_grp = N // _P_MAX
+    for img in range(nimg):
+        rofs = img * N * H
+        allocator = ModularAllocator(initial_address=res_top)
+        # O for ALL groups stays in SBUF until the single out-projection.
+        o_feat = []
+        for g in range(G):
+            of_ = allocator.alloc_sbuf_tensor(shape=(_P_MAX, N), dtype=nl.bfloat16,
+                                              block_dim=[n_ochunks], num_free_tiles=[n_ochunks],
+                                              align_to=32)
+            o_feat.append(of_)
+        # x^T: shared by every group (transposed ONCE per image).
+        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, n_tt],
+                                          align_to=32)
+        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, min(_PSUM_FMAX, N - s * _PSUM_FMAX)]], offset=0),
+                    src=h_hbm.ap([[H, min(_PSUM_FMAX, N - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
+                                 offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
+                )
+        # rope tables: shared by every group.
+        if resident:
+            tab = r_tab
+        else:
+            tab = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
+                                              block_dim=[4, n_tt], num_free_tiles=[4, n_tt], align_to=4)
+            for s in range(n_tt):
+                for t_i in range(4):
+                    _tw = min(_PSUM_FMAX, N - s * _PSUM_FMAX)
+                    nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
+                                  src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        img_base = allocator.get_current_address()
+        for g in range(G):
+            allocator.set_current_address(img_base)
+            col0 = g * _PSUM_FMAX
+            k_tiles = []
+            v_tiles = []
+            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_)
+            phase2_base = allocator.get_current_address()
+            if resident:
+                w_sb = r_w[g]
+                wv_sb = r_wv[g]
+                b_sb = r_b[g]
+                bv_bc = r_bv[g]
+            else:
+                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, col0:col0 + local_heads * D])
+                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 gg in range(local_quads * 2):
+                        nisa.dma_copy(dst=b_stage,
+                                      src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
+                        nisa.tensor_copy(dst=b_sb[qk][gg], 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=col0))
+            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_priv = []
+            for _lh in range(local_heads):
+                q_priv.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
+            # ---- producer (identical to attention_block's N<=2048 path) ----
+            for tt in range(n_tt):
+                t0 = tt * _PSUM_FMAX
+                tw = min(_PSUM_FMAX, N - t0)
+                for vsub in range(tw // _V_TILE):
+                    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):
+                        nisa.tensor_copy(dst=v_tiles[lh][tt * v_per_tt + vsub],
+                                         src=vout[:, lh * D:(lh + 1) * D], engine=engine.scalar)
+                for qk in range(2):
+                    for j in range(local_quads):
+                        for par in range(2):
+                            gq = j * 2 + par
+                            cols = gq * _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[:, 0:tw],
+                                               stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
+                                               moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
+                            nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
+                                               op0=nl.add, operand0=b_sb[qk][gq])
+                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
+                                           data2=tab[0][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
+                                           data2=tab[2][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=roped[0][:, 0:tw], data1=acc[:, 0:tw],
+                                           data2=tmp[:, 0:tw], op=nl.add)
+                        nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[1][:, 0:tw],
+                                           data2=tab[1][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[0][:, 0:tw],
+                                           data2=tab[3][tt][:, 0:tw], op=nl.multiply)
+                        nisa.tensor_tensor(dst=roped[1][:, 0:tw], data1=acc[:, 0:tw],
+                                           data2=tmp[:, 0:tw], op=nl.add)
+                        for b in range(quad):
+                            lh = j * quad + b
+                            r0 = b * half
+                            if qk == 1:
+                                for par2 in range(2):
+                                    nisa.tensor_copy(
+                                        dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half, 0:tw],
+                                        src=roped[par2][r0:r0 + half, 0:tw])
+                            else:
+                                for par2 in range(2):
+                                    nisa.dma_copy(
+                                        dst=q_priv[lh][0, par2 * half:(par2 + 1) * half, t0:t0 + tw],
+                                        src=roped[par2][r0:r0 + half, 0:tw])
+            # ---- attention for this group's heads -> o_feat[g] (SBUF, feature-major) ----
+            allocator.set_current_address(phase2_base)
+            o_sb = []
+            for _lh in range(local_heads):
+                oc = (_lh * D) // _P_MAX
+                rr = (_lh * D) % _P_MAX
+                grp = []
+                for _g in range(n_grp):
+                    grp.append(o_feat[g][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):
+                _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,
+                    seqlen_k_valid=n_valid,
+                )
+                allocator.set_current_address(attn_base)
+        # ---- ONE out-projection over ALL G*n_ochunks feature chunks, full bias once ----
+        allocator.set_current_address(img_base)
+        if resident:
+            wo_sb = r_wo
+            bo_bc = r_bo
+        else:
+            wo_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16,
+                                                block_dim=[G * n_ochunks], num_free_tiles=[G * n_ochunks])
+            for oc in range(G * n_ochunks):
+                nisa.dma_copy(dst=wo_sb[oc], src=wot_hbm[oc * _P_MAX:(oc + 1) * _P_MAX, 0:H])
+            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))
+        # ---- epilogue: out-proj (all groups, fp32) + bias + residual -> x_new ; LN -> z2 ----
+        _RING = 2
+        row = []
+        a_t = []
+        z_t = []
+        st = []
+        for _i in range(_RING):
+            row.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4))
+            a_t.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4))
+            z_t.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.bfloat16, align_to=32))
+            st.append(allocator.alloc_sbuf_tensor(shape=(_P_MAX, 1), dtype=nl.float32,
+                                                  block_dim=[5], num_free_tiles=[5], align_to=4))
+        sink = allocator.alloc_sbuf_tensor(shape=(_P_MAX, H), dtype=nl.float32, align_to=4)
+        n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
+        for tb in range(n_grp):
+            k = tb % _RING
+            r0 = img * N + tb * _P_MAX
+            nisa.dma_copy(dst=a_t[k], src=x_hbm[r0:r0 + _P_MAX, 0:H])
+            for ob in range(n_ob):
+                ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX)
+                ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
+                                    address=(0, ((tb * n_ob + ob) % 4) * _PSUM_BANK))
+                for g in range(G):
+                    for oc in range(n_ochunks):
+                        nisa.nc_matmul(
+                            dst=ps_out[:, 0:ow],
+                            stationary=o_feat[g][oc][:, tb * _P_MAX:(tb + 1) * _P_MAX],
+                            moving=wo_sb[g * n_ochunks + oc][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
+                            accumulate=((g > 0) or (oc > 0)),
+                        )
+                nisa.tensor_tensor(dst=row[k][:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow],
+                                   data1=ps_out[:, 0:ow],
+                                   data2=bo_bc[:, ob * _PSUM_FMAX:ob * _PSUM_FMAX + ow], op=nl.add)
+            nisa.tensor_tensor(dst=row[k], data1=row[k], data2=a_t[k], op=nl.add)
+            nisa.dma_copy(dst=xnew_hbm[r0:r0 + _P_MAX, 0:H], src=row[k])
+            # stat chain (identical ops/order to res_layernorm.py)
+            neg_shift = st[k][0]
+            acc_s = st[k][1]
+            neg_mu = st[k][2]
+            var = st[k][3]
+            rstd = st[k][4]
+            nisa.tensor_scalar(dst=neg_shift, data=row[k][0:_P_MAX, 0:1],
+                               op0=nl.multiply, operand0=-1.0)
+            nisa.activation_reduce(dst=sink, op=nl.copy, data=row[k],
+                                   reduce_op=nl.add, reduce_res=acc_s, bias=neg_shift, scale=1.0)
+            nisa.tensor_scalar(dst=neg_mu, data=acc_s, op0=nl.multiply, operand0=neg_inv_h,
+                               op1=nl.add, operand1=neg_shift)
+            nisa.activation_reduce(dst=sink, op=nl.square, data=row[k],
+                                   reduce_op=nl.add, reduce_res=acc_s, bias=neg_mu, scale=1.0)
+            nisa.tensor_scalar(dst=var, data=acc_s, op0=nl.multiply, operand0=inv_h,
+                               op1=nl.add, operand1=eps)
+            nisa.activation(dst=rstd, op=nl.rsqrt, data=var)
+            nisa.tensor_scalar(dst=z_t[k], data=row[k], op0=nl.add, operand0=neg_mu,
+                               op1=nl.multiply, operand1=rstd)
+            nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_t[k])
+    return xnew_hbm, z_hbm
+_ATTN_L1_LN = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
+@nki_op("vit_enc::attention_block_l1_ln", mutates_args={})
+def attention_block_l1_ln(
+    h: torch.Tensor,
+    wqk: torch.Tensor,
+    bqk: torch.Tensor,
+    tab: torch.Tensor,
+    wvt: torch.Tensor,
+    bv: torch.Tensor,
+    wot: torch.Tensor,
+    bo: torch.Tensor,
+    x: torch.Tensor,
+    scale: float,
+    head_dim: int,
+    n_valid: int,
+    eps: float,
+    nimg: int = 1,
+    resident: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """All-group attention + residual + affine-free LN2. Returns (x_new [R,H] fp32, z2 [R,H] bf16)."""
+    for t in (h, wqk, bqk, tab, wvt, bv, wot, bo, x):
+        if t is not None and t.requires_grad:
+            raise RuntimeError("attention_block_l1_ln is inference-only; use torch.no_grad().")
+    if h.shape[0] // nimg > 2048:
+        raise NotImplementedError("attention_block_l1_ln supports N<=2048 only")
+    return _ATTN_L1_LN(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident)

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

	@@ -0,0 +1,1220 @@

+# 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 the 4-kernel V-JEPA2 encoder reaches, with the four h8/h10 engine pins
+applied in place. Self-contained: imports only ``nki``, ``nkilib``, ``torch`` (none here) --
+NO imports from ``experiment/`` or the repo-root incumbent modules.
+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,
+    seqlen_k_valid: int = None,
+):
+    """One head of context-encoding attention (SBUF-KV path).
+    EDIT: ``seqlen_k_valid`` overrides the active key length (default len(k_cache)*512)
+    so keys at positions >= seqlen_k_valid (token padding) are skipped by every loop bound.
+    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,
+    )
+    if seqlen_k_valid is not None:
+        kernel_assert(seqlen_k_valid <= seqlen_k_active, "seqlen_k_valid exceeds K cache length")
+        seqlen_k_active = seqlen_k_valid
+    # 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,296 @@

+"""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.
+import os as _os
+#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
+LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "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("vit_enc::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)
+#: Requested export alias. Same object as the incumbent-named nki_op, so the runtime
+#: op-name string ("vit_enc::nki_mlp_k4_xp") is unchanged whichever name is used.
+mlp = nki_mlp_k4_xp

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

	@@ -0,0 +1,308 @@

+"""GELU MLP with the residual add + NEXT layer's affine-free LN1 in the epilogue.
+Same fc1/gelu as ``mlp.py``. fc2 is reordered to (row subtile rs, out block ko) so each 128-row tile's
+full [128, K] fc2 output lands in one fp32 SBUF row, then:
+    x_new = x + (g @ w2.T + b2)            fp32 (MLP output never rounded to bf16)
+    z1    = LN_no_affine(x_new; eps)        bf16 (same stat chain as res_layernorm.py)
+Returns (x_new [R,K] fp32, z1 [R,K] bf16). Replaces K4 + the next layer's K1.
+--- original mlp.py notes ---
+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.
+import os as _os
+#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
+LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "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_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
+    """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
+    Args:
+        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
+    _kernel_assert(tuple(a_hbm.shape) == (R, K), "residual a must be [R, K] fp32")
+    xnew_hbm = nl.ndarray((R, K), dtype=nl.float32, buffer=nl.shared_hbm)
+    z_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    inv_k = 1.0 / float(K)
+    neg_inv_k = -inv_k
+    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)
+    # b2 broadcast ONCE over all K columns (fp32), plus LN per-lane scratch.
+    b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
+    acc_s = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+    neg_shift = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+    neg_mu = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+    var = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+    rstd = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+    # 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 + b2 + residual -> x_new (fp32), LN -> z1 (bf16), per 128-row subtile ----
+        for rs in range(n_rs):
+            rr0 = rs * _P_MAX
+            prs = min(_P_MAX, pr - rr0)
+            r0 = r_abs + rr0
+            row = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)   # residual, accumulated in place
+            z_t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+            nisa.dma_copy(dst=row[0:prs, 0:K], src=a_hbm[r0:r0 + prs, 0:K])
+            for ko in range(n_ko):
+                ko0 = ko * _F_MOV_MAX
+                nk = min(_F_MOV_MAX, K - ko0)
+                ps = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
+                for ht in nl.affine_range(n_h):
+                    h0 = ht * _P_MAX
+                    ph = min(_P_MAX, H - h0)
+                    nisa.nc_matmul(
+                        dst=ps[0:prs, 0:nk],
+                        stationary=gt_tiles[ht][0:ph, rr0:rr0 + prs],
+                        moving=w2t_tiles[ht][0:ph, ko0:ko0 + nk],
+                    )
+                nisa.tensor_tensor(dst=row[0:prs, ko0:ko0 + nk], data1=row[0:prs, ko0:ko0 + nk],
+                                   data2=ps[0:prs, 0:nk], op=nl.add)
+            nisa.tensor_tensor(dst=row[0:prs, 0:K], data1=row[0:prs, 0:K], data2=b2_bc[0:prs, 0:K], op=nl.add)
+            nisa.dma_copy(dst=xnew_hbm[r0:r0 + prs, 0:K], src=row[0:prs, 0:K])
+            # stat chain (identical ops/order to res_layernorm.py)
+            nisa.tensor_scalar(dst=neg_shift[0:prs, 0:1], data=row[0:prs, 0:1],
+                               op0=nl.multiply, operand0=-1.0)
+            nisa.activation_reduce(dst=z_t[0:prs, 0:K], op=nl.copy, data=row[0:prs, 0:K],
+                                   reduce_op=nl.add, reduce_res=acc_s[0:prs, 0:1],
+                                   bias=neg_shift[0:prs, 0:1], scale=1.0)
+            nisa.tensor_scalar(dst=neg_mu[0:prs, 0:1], data=acc_s[0:prs, 0:1], op0=nl.multiply,
+                               operand0=neg_inv_k, op1=nl.add, operand1=neg_shift[0:prs, 0:1])
+            nisa.activation_reduce(dst=z_t[0:prs, 0:K], op=nl.square, data=row[0:prs, 0:K],
+                                   reduce_op=nl.add, reduce_res=acc_s[0:prs, 0:1],
+                                   bias=neg_mu[0:prs, 0:1], scale=1.0)
+            nisa.tensor_scalar(dst=var[0:prs, 0:1], data=acc_s[0:prs, 0:1], op0=nl.multiply,
+                               operand0=inv_k, op1=nl.add, operand1=eps)
+            nisa.activation(dst=rstd[0:prs, 0:1], op=nl.rsqrt, data=var[0:prs, 0:1])
+            nisa.tensor_scalar(dst=z_t[0:prs, 0:K], data=row[0:prs, 0:K],
+                               op0=nl.add, operand0=neg_mu[0:prs, 0:1],
+                               op1=nl.multiply, operand1=rstd[0:prs, 0:1])
+            nisa.dma_copy(dst=z_hbm[r0:r0 + prs, 0:K], src=z_t[0:prs, 0:K])
+    return xnew_hbm, z_hbm
+def _no_grad(*tensors) -> None:
+    for t in tensors:
+        if t is not None and t.requires_grad:
+            raise RuntimeError("nki_mlp_res_ln is inference-only; use torch.no_grad().")
+@nki_op("vit_enc::nki_mlp_res_ln", mutates_args={})
+def nki_mlp_res_ln(
+    x: torch.Tensor,
+    w1t: torch.Tensor,
+    b1: torch.Tensor,
+    w2t: torch.Tensor,
+    b2: torch.Tensor,
+    a: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """MLP(x) + residual a, then affine-free LN. x [R,K] bf16, a [R,K] fp32.
+    Returns (x_new [R,K] fp32, z [R,K] bf16)."""
+    _no_grad(x, w1t, b1, w2t, b2, a)
+    return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps)

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

	@@ -0,0 +1,213 @@

+"""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.
+import os as _os
+#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
+LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "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("vit_enc::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)
+#: Requested export alias. Same object as the incumbent-named nki_op, so the runtime
+#: op-name string ("vit_enc::nki_res_layernorm_fold") is unchanged whichever name is used.
+res_layernorm_fold = nki_res_layernorm_fold

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

	@@ -0,0 +1,291 @@

+"""Weight-streaming SwiGLU FFN, with an optional fused residual + next-layer LN epilogue.
+    out = (silu(x @ w1.T + b1) * (x @ wg.T + bg)) @ w2.T + b2           (SwiGLU FFN)
+WHY STREAMING. ``swiglu.py`` keeps gate+up+down weights SBUF-resident (3 x K x I bf16). At ViT-H+
+(K=1280, I=3416) that is ~205 KiB/partition, over the ~192 KiB budget, so the compiler spills.
+Here the intermediate dim I is processed in chunks of CI columns: only one chunk's gate/up columns
+and down rows are live (double-buffered by the scheduler: chunk c+1's DMA has no dependency on
+chunk c's matmuls), and the down-projection is accumulated across chunks in an fp32 SBUF
+accumulator with rows on partitions.
+LOOP ORDER. Rows (per core) are processed in super-blocks of S rows (host-chosen, <=1024). Within
+a super-block the chunk loop is OUTER, so each weight chunk is DMA'd once per super-block and
+reused by every 512-row slice in it -- weight traffic per row falls as S grows.
+FUSED EPILOGUE (``fuse=1``). The accumulator is initialised from the fp32 residual ``a``; after the
+last chunk: x_new = acc + b2 (fp32), z = LN_no_affine(x_new; eps) (bf16, same stat chain as
+res_layernorm.py). Returns (x_new, z). ``fuse=0`` returns bf16 ``acc + b2`` like ``swiglu.py``.
+Rows are sharded across LNC cores (R % (LNC*128) == 0). INFERENCE ONLY.
+"""
+from __future__ import annotations
+import torch
+from torch_neuronx import nki_op
+import nki
+import nki.isa as nisa
+import nki.language as nl
+import os as _os
+#: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
+LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
+_P_MAX = 128
+_F_MOV_MAX = 512
+_BF16 = nl.bfloat16
+def _div_ceil(n: int, d: int) -> int:
+    return (n + d - 1) // d
+def _kernel_assert(cond: bool, msg: str):
+    assert cond, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {msg}"
+def plan(K: int, I: int, rows_per_core: int, budget: int = 176 * 1024):
+    """Host-side choice of (S, CI) from a per-partition SBUF estimate (bytes)."""
+    n_k = _div_ceil(K, _P_MAX)
+    best = None
+    for CI in (1024, 512, 256, 128):
+        for S in (1024, 512, 256, 128):
+            if S > _div_ceil(rows_per_core, _P_MAX) * _P_MAX and S != 128:
+                continue
+            nh = CI // _P_MAX
+            w = 2 * (2 * n_k * CI * 2 + nh * K * 2)                  # 2 chunks in flight
+            acc = (S // _P_MAX) * K * 4                              # fp32 accumulator rows
+            xt = n_k * S * 2 + (S // _P_MAX) * K * 2                 # x^T + staging rows
+            g = nh * _F_MOV_MAX * (2 + 4)                            # g bf16 + silu fp32 per tile
+            ep = 2 * K * 2 + K * 4                                   # z tiles + b2 broadcast
+            tot = w + acc + xt + g + ep + _div_ceil(I, _P_MAX) * 8 + 4096
+            if tot <= budget:
+                n_sb = _div_ceil(rows_per_core, S)
+                key = (-n_sb, CI, -S)      # fewest weight re-streams, then big CI, then tight S
+                if best is None or key > best[0]:
+                    best = (key, S, CI, tot)
+    _kernel_assert(best is not None, f"no SBUF plan for K={K} I={I} rows={rows_per_core}")
+    return best[1], best[2], best[3]
+def _stage_xt(x_hbm, r_abs, S, n_k, K):
+    xt = []
+    for _ in range(n_k):
+        xt.append(nl.ndarray((_P_MAX, S), dtype=_BF16, buffer=nl.sbuf))
+    for rs in range(S // _P_MAX):
+        rr0 = rs * _P_MAX
+        xb = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+        nisa.dma_copy(dst=xb[0:_P_MAX, 0:K], src=x_hbm[r_abs + rr0:r_abs + rr0 + _P_MAX, 0:K])
+        for kt in range(n_k):
+            k0 = kt * _P_MAX
+            pk = min(_P_MAX, K - k0)
+            nisa.dma_transpose(dst=xt[kt][0:pk, rr0:rr0 + _P_MAX], src=xb[0:_P_MAX, k0:k0 + pk])
+    return xt
+def _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, fuse):
+    R, K = x_hbm.shape
+    K1, I = w1t_hbm.shape
+    _kernel_assert(K1 == K and tuple(wgt_hbm.shape) == (K, I), "w1t/wgt must be [K, I]")
+    _kernel_assert(tuple(w2t_hbm.shape) == (I, K), "w2t must be [I, K]")
+    _kernel_assert(R % (LNC * _P_MAX) == 0, f"R={R} must be a multiple of LNC*128")
+    rows_per = R // LNC
+    _kernel_assert(S % _P_MAX == 0, f"S={S} must be a multiple of 128")
+    _kernel_assert(CI % _P_MAX == 0, "CI must be a multiple of 128")
+    pid = nl.program_id(0)
+    r_base = pid * rows_per
+    n_k = _div_ceil(K, _P_MAX)
+    n_h = _div_ceil(I, _P_MAX)
+    n_c = _div_ceil(I, CI)
+    n_ko = _div_ceil(K, _F_MOV_MAX)
+    n_sb = _div_ceil(rows_per, S)
+    if fuse:
+        xnew_hbm = nl.ndarray((R, K), dtype=nl.float32, buffer=nl.shared_hbm)
+        z_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    else:
+        out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    # biases (tiny, resident): b1/bg per 128-row tile of I, b2 broadcast over partitions.
+    b1_t = []
+    bg_t = []
+    for ht in range(n_h):
+        h0 = ht * _P_MAX
+        ph = min(_P_MAX, I - h0)
+        t1 = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+        nisa.dma_copy(dst=t1[0:ph, 0:1], src=b1_hbm[h0:h0 + ph])
+        b1_t.append(t1)
+        tg = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+        nisa.dma_copy(dst=tg[0:ph, 0:1], src=bg_hbm[h0:h0 + ph])
+        bg_t.append(tg)
+    b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
+    inv_k = 1.0 / float(K)
+    neg_inv_k = -inv_k
+    for sb in range(n_sb):
+        r_sb = r_base + sb * S
+        Ssb = min(S, rows_per - sb * S)          # ragged last super-block (multiple of 128)
+        n_rs = Ssb // _P_MAX
+        xt = _stage_xt(x_hbm, r_sb, Ssb, n_k, K)
+        acc = []
+        for rs in range(n_rs):
+            t = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
+            if fuse:
+                nisa.dma_copy(dst=t[0:_P_MAX, 0:K],
+                              src=a_hbm[r_sb + rs * _P_MAX:r_sb + (rs + 1) * _P_MAX, 0:K])
+            acc.append(t)
+        for c in range(n_c):
+            i0 = c * CI
+            ci = min(CI, I - i0)
+            nh_c = _div_ceil(ci, _P_MAX)
+            # ---- this chunk's weights (gate/up columns i0:i0+ci, down rows i0:i0+ci) ----
+            w1c = []
+            wgc = []
+            for kt in range(n_k):
+                k0 = kt * _P_MAX
+                pk = min(_P_MAX, K - k0)
+                t1 = nl.ndarray((_P_MAX, CI), dtype=_BF16, buffer=nl.sbuf)
+                nisa.dma_copy(dst=t1[0:pk, 0:ci], src=w1t_hbm[k0:k0 + pk, i0:i0 + ci])
+                w1c.append(t1)
+                tg = nl.ndarray((_P_MAX, CI), dtype=_BF16, buffer=nl.sbuf)
+                nisa.dma_copy(dst=tg[0:pk, 0:ci], src=wgt_hbm[k0:k0 + pk, i0:i0 + ci])
+                wgc.append(tg)
+            w2c = []
+            for hl in range(nh_c):
+                h0 = i0 + hl * _P_MAX
+                ph = min(_P_MAX, I - h0)
+                t2 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+                nisa.dma_copy(dst=t2[0:ph, 0:K], src=w2t_hbm[h0:h0 + ph, 0:K])
+                w2c.append(t2)
+            for rb in range(_div_ceil(Ssb, _F_MOV_MAX)):
+                q0 = rb * _F_MOV_MAX
+                rw = min(_F_MOV_MAX, Ssb - q0)
+                # ---- gate/up for this chunk x this 512-row slice -> g [ci, rw] bf16 ----
+                g_t = []
+                for hl in range(nh_c):
+                    hg = (i0 // _P_MAX) + hl
+                    ph = min(_P_MAX, I - (i0 + hl * _P_MAX))
+                    c0 = hl * _P_MAX
+                    y1 = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
+                    for kt in nl.affine_range(n_k):
+                        pk = min(_P_MAX, K - kt * _P_MAX)
+                        nisa.nc_matmul(dst=y1[0:ph, 0:rw],
+                                       stationary=w1c[kt][0:pk, c0:c0 + ph],
+                                       moving=xt[kt][0:pk, q0:q0 + rw])
+                    yg = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
+                    for kt in nl.affine_range(n_k):
+                        pk = min(_P_MAX, K - kt * _P_MAX)
+                        nisa.nc_matmul(dst=yg[0:ph, 0:rw],
+                                       stationary=wgc[kt][0:pk, c0:c0 + ph],
+                                       moving=xt[kt][0:pk, q0:q0 + rw])
+                    sa = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
+                    nisa.activation(dst=sa[0:ph, 0:rw], op=nl.silu, data=y1[0:ph, 0:rw],
+                                    bias=b1_t[hg][0:ph, 0:1])
+                    g = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
+                    # g = (yg + bg) * silu(y1 + b1)   -- one fused op, PSUM operand
+                    nisa.scalar_tensor_tensor(dst=g[0:ph, 0:rw], data=yg[0:ph, 0:rw],
+                                              op0=nl.add, operand0=bg_t[hg][0:ph, 0:1],
+                                              op1=nl.multiply, operand1=sa[0:ph, 0:rw])
+                    g_t.append(g)
+                # ---- down-projection partial for this chunk, accumulated into acc (fp32 SBUF) ----
+                for rs_l in range(rw // _P_MAX):
+                    rs = (q0 // _P_MAX) + rs_l
+                    for ko in range(n_ko):
+                        ko0 = ko * _F_MOV_MAX
+                        nk = min(_F_MOV_MAX, K - ko0)
+                        ps = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
+                        for hl in nl.affine_range(nh_c):
+                            ph = min(_P_MAX, I - (i0 + hl * _P_MAX))
+                            nisa.nc_matmul(dst=ps[0:_P_MAX, 0:nk],
+                                           stationary=g_t[hl][0:ph, rs_l * _P_MAX:(rs_l + 1) * _P_MAX],
+                                           moving=w2c[hl][0:ph, ko0:ko0 + nk])
+                        if c == 0 and not fuse:
+                            nisa.tensor_copy(dst=acc[rs][0:_P_MAX, ko0:ko0 + nk], src=ps[0:_P_MAX, 0:nk])
+                        else:
+                            nisa.tensor_tensor(dst=acc[rs][0:_P_MAX, ko0:ko0 + nk],
+                                               data1=acc[rs][0:_P_MAX, ko0:ko0 + nk],
+                                               data2=ps[0:_P_MAX, 0:nk], op=nl.add)
+        # ---- epilogue per 128-row subtile ----
+        for rs in range(n_rs):
+            r0 = r_sb + rs * _P_MAX
+            row = acc[rs]
+            nisa.tensor_tensor(dst=row[0:_P_MAX, 0:K], data1=row[0:_P_MAX, 0:K],
+                               data2=b2_bc[0:_P_MAX, 0:K], op=nl.add)
+            zt = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
+            if fuse:
+                nisa.dma_copy(dst=xnew_hbm[r0:r0 + _P_MAX, 0:K], src=row[0:_P_MAX, 0:K])
+                neg_shift = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
+                acc_s = 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)
+                nisa.tensor_scalar(dst=neg_shift, data=row[0:_P_MAX, 0:1], op0=nl.multiply, operand0=-1.0)
+                nisa.activation_reduce(dst=zt, op=nl.copy, data=row[0:_P_MAX, 0:K],
+                                       reduce_op=nl.add, reduce_res=acc_s, bias=neg_shift, scale=1.0)
+                nisa.tensor_scalar(dst=neg_mu, data=acc_s, op0=nl.multiply, operand0=neg_inv_k,
+                                   op1=nl.add, operand1=neg_shift)
+                nisa.activation_reduce(dst=zt, op=nl.square, data=row[0:_P_MAX, 0:K],
+                                       reduce_op=nl.add, reduce_res=acc_s, bias=neg_mu, scale=1.0)
+                nisa.tensor_scalar(dst=var, data=acc_s, op0=nl.multiply, operand0=inv_k,
+                                   op1=nl.add, operand1=eps)
+                nisa.activation(dst=rstd, op=nl.rsqrt, data=var)
+                nisa.tensor_scalar(dst=zt, data=row[0:_P_MAX, 0:K], op0=nl.add, operand0=neg_mu,
+                                   op1=nl.multiply, operand1=rstd)
+                nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:K], src=zt)
+            else:
+                nisa.tensor_copy(dst=zt, src=row[0:_P_MAX, 0:K])
+                nisa.dma_copy(dst=out_hbm[r0:r0 + _P_MAX, 0:K], src=zt)
+    if fuse:
+        return xnew_hbm, z_hbm
+    return out_hbm
+@nki.jit
+def _swiglu_stream_kernel(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, S, CI):
+    return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, None, 0.0, S, CI, 0)
+@nki.jit
+def _swiglu_stream_ln_kernel(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI):
+    return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, 1)
+def _no_grad(*ts):
+    for t in ts:
+        if t is not None and t.requires_grad:
+            raise RuntimeError("swiglu_stream is inference-only; use torch.no_grad().")
+def _plan_for(x, w1t):
+    S, CI, _ = plan(x.shape[1], w1t.shape[1], x.shape[0] // LNC)
+    S = int(_os.environ.get("VIT_ENC_SW_S", S))
+    CI = int(_os.environ.get("VIT_ENC_SW_CI", CI))
+    return S, CI
+@nki_op("vit_enc::nki_swiglu_stream", mutates_args={})
+def nki_swiglu_stream(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, wgt: torch.Tensor,
+                      bg: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
+    """x [R,K] bf16; w1t/wgt [K,I] bf16; b1/bg [I] f32; w2t [I,K] bf16; b2 [K] f32 -> [R,K] bf16."""
+    _no_grad(x, w1t, b1, wgt, bg, w2t, b2)
+    S, CI = _plan_for(x, w1t)
+    return _swiglu_stream_kernel[LNC](x, w1t, b1, wgt, bg, w2t, b2, S, CI)
+@nki_op("vit_enc::nki_swiglu_stream_ln", mutates_args={})
+def nki_swiglu_stream_ln(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, wgt: torch.Tensor,
+                         bg: torch.Tensor, w2t: torch.Tensor, b2: torch.Tensor, a: torch.Tensor,
+                         eps: float) -> tuple[torch.Tensor, torch.Tensor]:
+    """SwiGLU(x) + residual a, then affine-free LN. Returns (x_new [R,K] fp32, z [R,K] bf16)."""
+    _no_grad(x, w1t, b1, wgt, bg, w2t, b2, a)
+    S, CI = _plan_for(x, w1t)
+    return _swiglu_stream_ln_kernel[LNC](x, w1t, b1, wgt, bg, w2t, b2, a, eps, S, CI)

examples/dinov3_features.py ADDED Viewed

	@@ -0,0 +1,22 @@

+"""DINOv3 ViT-L/16 features on trn2 with vit-encoder-neuron-kernels.
+    NEURON_LOGICAL_NC_CONFIG=2 python dinov3_features.py
+"""
+import types
+import torch
+from kernels import get_kernel
+from transformers import AutoModel
+k = get_kernel("jburtoft/vit-encoder-neuron-kernels", version=1, trust_remote_code=True)
+model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
+ref_model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
+model = model.to("neuron")
+model.forward = types.MethodType(k.NeuronDINOv3ViTModel.forward, model)
+model = torch.compile(model, backend="neuron")
+torch.manual_seed(0)
+pixel_values = torch.randn(4, 3, 224, 224)
+with torch.no_grad():
+    out = model(pixel_values.to("neuron")).last_hidden_state.cpu()
+    ref = ref_model(pixel_values).last_hidden_state
+print("features", tuple(out.shape), "rel_l2 vs fp32 CPU:", ((out - ref).norm() / ref.norm()).item())