File size: 780 Bytes
eadc767 | 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 | """V-JEPA2 ViT-L encoder NKI kernels for AWS Trainium2 (trn2).
Public API:
NeuronVJEPA2Encoder -- KernelConfig-mappable layer ("VJEPA2Encoder" -> this)
res_layernorm_fold, attention_block, mlp -- the three kernel entry points
pack_* / build_rope_tables / lane_permutation -- host-side weight packing
"""
from .layers import NeuronVJEPA2Encoder
from .nki_kernels import (
attention_block,
build_rope_tables,
lane_permutation,
mlp,
pack_out_weights,
pack_tables,
pack_v_weights,
pack_weights,
res_layernorm_fold,
)
__all__ = [
"NeuronVJEPA2Encoder",
"res_layernorm_fold", "attention_block", "mlp",
"pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
"lane_permutation", "build_rope_tables",
]
|