| """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, calibrate_fp8, pack_block, rope_tables |
| from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder |
| from .nki_kernels import ( |
| attention_block, |
| attention_block_l1, |
| attention_block_l1_ln, |
| lane_permutation, |
| mlp, |
| mlp_res_ln, |
| pack_out_weights, |
| pack_tables, |
| pack_v_weights, |
| pack_weights, |
| res_layernorm_fold, |
| swiglu_stream, |
| swiglu_stream_ln, |
| ) |
|
|
| __all__ = [ |
| "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder", |
| "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8", |
| "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold", |
| "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln", |
| "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation", |
| ] |
|
|
