File size: 945 Bytes
c855122 ce0e45f c855122 cabaa8f c855122 cabaa8f ce0e45f 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 | """The NKI kernels and their host-side weight packing."""
from .attention_block import (
attention_block,
lane_permutation,
pack_out_weights,
pack_qk_bias,
pack_tables,
pack_v_weights,
pack_weights,
)
from .attention_block_l1 import attention_block_l1
from .attention_block_l1_ln import attention_block_l1_ln
from .mlp import nki_mlp_k4_xp as mlp
from .mlp import pack_b1
from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
from .swiglu_stream import nki_swiglu_stream as swiglu_stream
from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
__all__ = [
"attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
"mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
"pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
]
|