Views
No views yet
conv2d kernel for AWS Trainium / Inferentia,
packaged for the HuggingFace kernels library + the KernelConfig API.torch.nn.Conv2d with an implicit-GEMM NKI implementation that runs
on the NeuronCore Tensor Engine.build/torch-neuron/ — pure-Python NKI kernel (compiled by neuronx-cc at
load time). Requires the Neuron SDK (nki) to be installed in the runtime.(sH, sW)(pH, pW)R x S), 1x1, 3x3, 5x5, ...bf16 and fp32stride >= 1, dilation = 1, groups = 1, padded plane
Hp*Wp <= 32767 (single-tile). Correctness validated against
torch.nn.functional.conv2d (cosine = 1.0; fp32 max-abs ~1e-5).1from transformers import AutoModelForCausalLM, KernelConfig # or any model with nn.Conv2d
2
3kernel_config = KernelConfig({"Conv2d": "<owner>/conv2d-neuron-kernels:NeuronConv2d"})
4model = AutoModelForCausalLM.from_pretrained(
5 "<model-id>",
6 use_kernels=True,
7 kernel_config=kernel_config,
8)Conv2d (the key) is the original module class name that gets replaced.
NeuronConv2d (the value) is the KernelName; the repo also provides the
companion NeuronConv2dLayout that holds parameters and declares the
[Cout,Cin,R,S] -> [Cin,R,S,Cout] weight relayout via conversion_mapping.