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.