File size: 1,360 Bytes
c855122 3b1bffa c855122 3b1bffa c855122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """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",
]
|