A Go library that provides mappings between AWS Neuron devices and their corresponding EFA (Elastic Fabric Adapter) devices. These mappings enable aligned device allocation for Neuron + EFA workloads on supported EC2 instance types.
When allocating Neuron and EFA devices to a pod, the devices must be aligned — each Neuron device must be paired with its corresponding EFA device for optimal performance. Standard PCI root-based matching does not produce optimal allocations for Trainium instances.
This library provides pre-computed mappings from Neuron device IDs to EFA Bus:Device:Function (BDF) identifiers, enabling correct aligned allocation.
go get github.com/aws-neuron/connected-device-maps-over-efa-for-neuronimport (
"github.com/aws-neuron/connected-device-maps-over-efa-for-neuron/lookup"
)Get the device group IDs for an EFA device:
deviceGroupIDs, err := lookup.GetEFADeviceGroupIDs(deviceID)Get the device group IDs for a Neuron device:
deviceGroupIDs, err := lookup.GetNeuronDeviceGroupIDs(deviceID)See lookup/config/efa-enabled-instances.json for the full list of supported EFA-enabled Neuron instance types.
New instance type mappings are maintained by the AWS Neuron team and will be added as new instance types are introduced.
The supported instance types are listed in lookup/config/efa-enabled-instances.json.
New instance family mappings are maintained by the AWS Neuron team and will be added as new instance families are introduced.
Mappings from Neuron device IDs to EFA BDFs are stored as JSON files in the lookup/neuron directory. These mappings are constructed from the Neuron collectives topology definitions and the Neuron driver device mappings.
lookup/config/instance-family-mappings.json maps instance families to their corresponding mapping files.
See CONTRIBUTING for more information.
This project is licensed under the Apache-2.0 License. See the LICENSE file.