| """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, attention_block_l1_ln_post | |
| 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 .mlp_res_ln import nki_mlp_res_ln_post as mlp_res_ln_post | |
| 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", "attention_block_l1_ln_post", "res_layernorm_fold", | |
| "mlp", "mlp_res_ln", "mlp_res_ln_post", "pack_b1", "swiglu_stream", "swiglu_stream_ln", | |
| "pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation", | |
| ] | |