| 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. | |