Programming AWS Trainium with the Neuron Kernel Interface - What the Compiler Stops Deciding for You, Which Names Do Which Jobs, and When Not to Descend
First Published:
Last Updated:
Read NKI as an API for going faster and there is nothing in this article to gain. What NKI hands over is not speed; it provides the right to make decisions that the compiler previously made automatically. And with that right comes responsibility. Accepting the right to make these decisions also means accepting the obligation to make them.
The official documentation lists three key elements:
The Neuron Kernel Interface (NKI) offers low-level access to the Trainium instruction set, memory, and execution scheduling so performance engineers can build custom kernels and push performance beyond standard frameworks.
Instructions, memory, and execution scheduling. While writing models through a framework, you would never directly interact with these elements. They worked fine because the Neuron compiler made those decisions for you. Descending into NKI means taking those three decisions over from the compiler.
This article will begin by outlining what you should understand before taking that step, then will systematically explain what becomes your responsibility once you have taken control. Finally comes the case for not descending at all.
This article is intended for performance engineers who are running models on Trainium and are currently limited by the performance of the standard implementation. It assumes experience with Triton and CUDA on the GPU side, and familiarity with tiling, shared memory, and reduction. It assumes, on the other hand, that the Neuron stack is new. Terms such as NeuronCore, SBUF, PSUM, and NEFF are defined where they first appear.
This article is not an introductory guide, nor does it provide step-by-step instructions on how to write kernels. It covers the decision of whether to descend and what becomes your own work afterward. The official tutorials and API reference provide the necessary information on how to write kernels.
This article does not cover situations where compute kernels are not batch-invariant, nor the numerical consequences of such behavior. The phenomenon where the same input does not always produce the same output, the ways in which different partitioning strategies for reduction can affect results, and the costs associated with achieving determinism, are all discussed in the existing article Reproducible LLM Inference, Section 5. This article will not reiterate those explanations. That article looks at kernels from the side that reads them; this one looks from the side that writes them.
This article also does not address the history and generations of the chips themselves. The details of when Inferentia and Trainium were introduced and what changes they brought are documented in the existing article AWS Custom Silicon History and Timeline. Similarly, the evolution of instance families is detailed in Amazon EC2 Instance Types History and Timeline. This article will not present a timeline.
This article also does not cover the design of providing models on Neuron. The process of providing your own models on Amazon EKS using vLLM is described in the existing article Self-Managed LLM Inference on Amazon EKS, which includes details on Neuron's pre-compilation procedures, specifically the practice of distributing and reusing compiled artifacts. This article focuses on the internals of the compiler, not the operation of the resulting compiled artifacts. The design of separating prefill and decode is covered in Disaggregated Prefill and Decode for LLM Serving on AWS.
This article also does not address weight quantization. The existing article LLM Weight Quantization on AWS details how different formats impact accuracy. The design of communication between nodes is covered in Elastic Fabric Adapter and the AWS Network Fabric.
And there are two layers above this where descending this far is unnecessary in the first place. Discussions on creating models within a managed environment are found in Model Customization on Amazon Bedrock. The topic of how far to roll back from failures when running training on your own cluster is addressed in Distributed Training Resilience on AWS. The latter article explicitly states that it does not cover instruction-level optimization, and hands the topic to this one. This article addresses the lowest layer of that scope.
This article does not discuss pricing. It will not mention any figures, costs, or unit prices. No benchmarks were run for this article either. A claim that writing a kernel makes things faster only holds with the author's own measurement behind it. Nothing here was measured on real hardware, so no multipliers and no timings appear. Numbers published by AWS are quoted with the attribution made explicit. No performance comparison against GPUs appears either.
The specifications presented in this article were verified on September 1, 2026. This article differs from others in that it identifies two distinct paths for obtaining primary source information.
The authoritative source for NKI is the AWS Neuron documentation, which is hosted on readthedocs rather than on
docs.aws.amazon.com. Conversely, the product pages – specifically, the Neuron SDK page, the Trainium research page, the Trn3 UltraServers page, and the Prescriptive Guidance – are located on aws.amazon.com. This article aims to present both sources, noting any discrepancies without attempting to reconcile them.The process of comparing these two sources proved to be the most valuable aspect of this article. The product pages and the developer documentation say different things about NKI on three points. These differences concern the number of components, the targeted chips, and whether the compiler source code has been publicly released. Sections 8 through 10 address each of these points.
Any item that could not be verified is left marked as unverified. Nothing is filled in by speculation.
Table of Contents
- 1. Standard Frameworks Run Unchanged, and That Is the Default
- 2. What the Compiler Decides When You Do Nothing
- 3. Two Things to Do Before You Write Anything
- 4. Where NKI Cuts Into the Pipeline
- 5. The Three Things You Take Over
- 6. Tile-Level Semantics, and Why Triton-Like Is Not Triton
- 7. It Looks Like Python and It Is Not Python
- 8. Four Names, and the Sources Do Not Agree on How Many
- 9. What the Open Source Claim Actually Covers
- 10. Which Chips: Three Answers in One Documentation Set
- 11. What NKI Does Not Give You
- 12. When Not to Descend
- 13. Failure Modes and Anti-Patterns
- 14. Frequently Asked Questions
- 15. Summary
- 16. References
1. Standard Frameworks Run Unchanged, and That Is the Default
The first thing to establish is that not descending into NKI is the default state.AWS Neuron is a developer stack for running models on Trainium and Inferentia, encompassing a compiler, runtime, libraries for training and inference, as well as tools for monitoring, profiling, and debugging. The product page repeatedly emphasizes that standard frameworks function without modification. PyTorch code runs without needing to be rewritten, and features like
torch.compile, as well as distributed mechanisms such as FSDP, DDP, and DTensor, continue to work as expected. It is also compatible with JAX.In this state, where no modifications are required, you will not need to concern yourself with instruction selection, memory allocation, or execution scheduling. The Neuron documentation describes this pre-NKI state as follows:
For comparison, before NKI was introduced, the only way to program NeuronDevices was through defining high-level ML models in frameworks such as PyTorch and JAX.
So NKI opened a second entrance beside the one that had been the only one. The original entry point remains active and functional.
1.1 The Model Path and the Kernel Path Coexist
Just because a second option is available does not mean everything needs to be rewritten. Prescriptive Guidance describes the relationship between the standard process and a custom kernel as follows:Models can be compiled before deployment using different libraries like Optimum Neuron, NeuronX Inference or PyTorch, which all use the Neuron compiler and leverage Neuron Kernel Interface (NKI) to apply custom and optimized kernels.
The entire model is still developed within the framework, and only specific operators within it are written using NKI. This is the standard approach. The intention is not to rewrite the entire model using NKI.
So the decision to adopt NKI is not a platform migration but a decision about scope: which operators inside the model you will decide for yourself. The narrower the scope, the smaller the commitment required.
2. What the Compiler Decides When You Do Nothing
When you do not explicitly define any behavior, it is the Neuron compiler that makes the decisions. Without first knowing what it decides, there is no way to know what is being taken over.The Neuron documentation lists four stages of optimization that the compiler runs sequentially. It takes high-level models defined within the framework as input and ultimately outputs a NEFF (Neuron Executable File Format) that can be executed on the NeuronDevice.
The first stage involves hardware-independent graph-level optimization. This is performed using XLA in the compiler's front-end and includes optimizations such as constant propagation, re-materialization, and operator fusion.
The second stage focuses on loop-level optimization. It transforms the graph, optimized in the first stage, into a series of nested loops, and performs optimizations such as layout optimization, tiling, and loop fusion.
The third stage involves mapping to hardware-specific instructions. It translates the architecture-independent nested loops obtained in the second stage to architecture-specific instructions.
The fourth stage involves hardware-specific optimization. Primarily performed in the compiler's back-end, it focuses on reducing memory pressure and increasing instruction-level parallelism. The documentation names what happens in this stage. Memory allocation and instruction scheduling are carried out in this fourth stage.
2.1 Of the Four Stages, Only the Last One Is Taken Over
This is the most important fact presented in this article. The three things the opening definition listed, the instructions, the memory, and the execution schedule, are concentrated in the third and fourth of the four stages. What the first and second stages decide is the shape of the graph and the structure of the loops, not which instructions are issued in what order.So when the compiler's decisions are described as being taken over, not all of them are. What is taken over is the lowest layer. That asymmetry governs everything that follows.
3. Two Things to Do Before You Write Anything
Before you start writing, there are two things you need to do. If you skip these, your efforts will often be wasted.The Neuron documentation puts this sentence in the opening lines of the NKI top page:
Before you write a kernel, check the NKI Library — it may already have an optimized kernel for your operation.
The official recommendation is not to start writing immediately, but rather, to borrow first.
3.1 The NKI Library, and the Chance Someone Has Already Written It
The NKI Library is a collection of kernels optimized for AWS Neuron-powered devices. Alongside the kernels themselves, the documentation lists design specifications, references for supported kernels, utilities for memory management and tensor views and iteration, and release notes. The source code is publicly available asaws-neuron/nki-library.The scope of the library is quite extensive. The NKI FAQ describes what the included kernels cover:
the NKI Library provides optimized kernels for attention, MLP, RMSNorm with quantization, and collective communication that outperform compiler-generated equivalents.
These include attention, MLP, RMSNorm with quantization, and collective communication. The library already covers most of the places where a modern Transformer tends to become a performance problem.
And the library keeps growing. AWS Neuron versions 2.30.0, 2.31.0, and 2.32.0 announced the addition of 22 kernels, 14 experimental kernels, and 13 kernels for Mixture of Experts training and sparse attention, respectively. The operator you are considering writing might be included in the next release.
3.2 Two Repositories That Are Easy to Confuse
There is one point that can be easily confused: the NKI FAQ clarifies that there are two public repositories, and that they serve different purposes.The NKI Library is intended for production use. It is optimized and designed to be directly integrated into models.
The NKI Samples are for learning purposes. The FAQ describes these as "reference and tutorial kernels," explicitly stating that they are intended for learning and experimentation, and not for production use.
Confusing these two can lead to using tutorial kernels in a production environment. Both are official repositories, so it can be difficult to distinguish them based solely on the repository source.
3.3 Neuron Explorer, and Measuring Where the Slowness Is First
Another crucial aspect is measurement. Before making any decisions about kernel modifications, it is essential to first understand where the slowdown is occurring.Neuron Explorer is a suite of tools for profiling and debugging. It can be accessed from the command-line interface, the user interface, and as a VS Code extension. The Neuron documentation provides a table indicating where to begin troubleshooting based on specific symptoms. When a training step is slower than expected, a system profile shows which phase dominates. When hardware utilization is low, a device profile shows engine utilization, DMA activity, and idle time. When collective operations are slow, a multi-worker system profile compares the timing across ranks.
And, importantly, one line directly relates to the subject of this article: When the NKI kernel is not performing as expected, use the device profile with source code linking enabled to map instructions to the kernel code.
3.4 Device Level and System Level Are Different Things
Neuron Explorer captures two different types of profiles. Confusing these two can lead you to analyze the wrong area.Device-level profiles capture hardware execution data from the NeuronCore. This includes instructions from the compute engine, DMA operations, and hardware utilization. If you are looking to optimize the kernel implementation, this is where you should look.
System-level profiles capture software execution data. This includes framework operations, API calls to the Neuron runtime, CPU utilization, and memory usage. If you want to identify framework overhead or CPU-side bottlenecks, this is the area to investigate.
Performance bottlenecks that cannot be resolved by modifying the kernel often appear in the system-level profile. Examining this profile first may allow you to identify and address issues without needing to write any code.
3.5 The Profiler Generation Has Already Turned Over
One point to note regarding the product lifecycle: Neuron Explorer has replaced previous profilers.Neuron ended support for the earlier Neuron Profiler 2.0 and Neuron Profiler in the Neuron 2.29 release.
With the Neuron 2.29 release, support for Neuron Profiler and Neuron Profiler 2.0 was discontinued. The user guides have been archived, and a migration guide is available. An older article or an internal runbook written around
neuron-profile is a migration item.That migration has not finished spreading through the documentation. As of September 1, 2026, the performance optimization page for NKI still assumes the use of
neuron-profile. Within the same set of documentation, the architecture guide refers to Neuron Explorer. If you find that the screen descriptions do not match the tools you are using, it is likely due to this generational difference.4. Where NKI Cuts Into the Pipeline
Suppose there is nothing to borrow, and the decision to write has been made after measuring. Where in the compiler does the kernel that gets written go?The Neuron documentation states the relationship to the four stages from Section 2 in a single sentence:
NKI kernels bypass the first 3 steps, and are compiled into IRs (intermediate representations) that the compiler's back-end (Step 4 above) can directly consume.
NKI kernels skip the initial three stages. They are then compiled into intermediate representations, and the compiler back end, which is the fourth stage, consumes them directly.
Nor is everything inside the fourth stage left automatic.
Advanced features in NKI, such as direct allocation, also allow programmers to bypass certain compiler passes in Step 4.
Using an advanced feature such as direct allocation, specific passes inside the fourth stage can be skipped as well.

4.1 The Work of the Skipped Stages Does Not Vanish, It Moves
What is easy to get wrong here is that the three skipped stages did not become unnecessary.The first stage of graph optimization, the second stage of loop optimization, and the third stage of instruction mapping were previously handled automatically through the framework. Within the NKI kernel, the developer is now making those equivalent decisions. Determining the shape of the tiles corresponds to loop optimization, and deciding which instructions to call corresponds to instruction mapping.
It is simply that the automation has been removed; the total amount of decisions that need to be made has not decreased. That is what the phrase right and obligation amounts to in practice.
4.2 The NKI Compiler Is a Different Compiler
There is another structural fact to consider: compiling the NKI kernel is performed by a different compiler than the one used to compile the entire model.The Neuron documentation explicitly explains why this separation exists. The logic is that a dedicated NKI Compiler is necessary, separate from the existing Neuron Graph Compiler, to ensure that the optimizations chosen by the author are preserved during compilation. The NKI Compiler parses the kernel's code into an abstract syntax tree, performs semantic analysis, and applies a limited number of transformations to optimize resource allocation and instruction scheduling, generating an optimized intermediate representation of NKI. This intermediate representation is then integrated into the larger model.
And the location of this integration has recently changed.
This insertion of NKI kernels into the graph is done very late in the compilation process. This is different from prior versions of NKI that integrated NKI kernels earlier in the compile process.
The insertion of kernels into the graph now occurs much later in the compilation process. Previous versions integrated them at an earlier stage. The documentation explains that delaying this integration allows for providing NKI-specific behavior and results in more predictable outcomes.
This also means that knowledge based on the behavior of older versions needs to be updated.
5. The Three Things You Take Over
Let us examine each of the three points mentioned in the initial definition, in detail.5.1 Memory: Where a Tensor Lives Is Written in the Declaration
Trainium's memory architecture consists of three layers: HBM, which serves as the device's working memory; SBUF, an on-chip scratchpad; and PSUM, a dedicated buffer for accumulating partial sums in matrix operations. SBUF is short for State Buffer, and PSUM is Partial Sum Buffer. SBUF functions as a cache managed by software, providing a location to hold active tensors and serving as the destination for DMA transfers.This is the key difference between Trainium and GPUs. While typical CPUs and GPUs utilize caches managed by hardware, in Trainium, the software – specifically NKI and the Neuron Graph Compiler – explicitly manages memory allocation and data movement across the entire memory hierarchy. The documentation states:
NKI exposes all NISA primitives needed to manage the memory hierarchy.
This explicit management means that the code specifies which layer a tensor is placed in. An example of allocation, as shown in the Neuron documentation, is as follows:
import nki.language as nl
# Allocate tensor of FP32 elements in SBUF (on-chip scratchpad memory)
# using ndarray call similar to numpy
# like numpy, nl supports ndarray(), zeros() and ones() functions
x_on_chip = nl.ndarray((128, 32, 512), dtype=nl.float32, buffer=nl.sbuf)
# Allocate tensor of FP16 elements in HBM (high-bandwidth memory, off-chip)
y_in_hbm = nl.ndarray(shape, dtype=nl.float16, buffer=nl.shared_hbm)
The
buffer argument directly indicates the location. Instead of hoping the data fits in a cache, the code declares where it should reside.The capacity of each layer is limited and relatively small. In the NeuronCore-v2, SBUF has a capacity of 24MiB, while PSUM has 2MiB. The 24MiB figure may seem large compared to the shared memory found in GPUs, but it is extremely small when compared to HBM. The decision of what to place in this hierarchy directly impacts the performance of the kernel.
5.2 Instructions: Which One Runs Is Written as a Name
NeuronCore utilizes four distinct calculation engines: the Tensor Engine handles matrix multiplication and convolution; the Vector Engine manages vector operations and reductions with multiple inputs; the Scalar Engine performs element-wise operations, including non-linear functions; and the GpSimd Engine is responsible for general-purpose, programmable processing. In addition to these, there are dedicated CC-Cores for collective communication.Within the NKI's lower-level API, the name of the call indicates which engine and which instruction are being used. This is demonstrated in the examples provided in the Neuron documentation.
import nki.isa as nisa
# Matrix multiplication on Tensor Engine using nc_matmul
# nc stands for NeuronCore, and matmul is the instruction name
# stationary: [128, 128], moving: [128, 512], output: [128, 512]
# The input arguments must meet NISA requirements as defined
# in the Trainium architecture, such as data types, layout, tile sizes
# and buffer memory types (SBUF or PSUM)
# dst is explicitly defined as the first instruction parameter
nisa.nc_matmul(output, stationary, moving)
# Element-wise operations between two tensors
# in this specific example, x and y must have the same partition dimension size
# and the same number of elements per partition.
# Notice the destination (dst) is explicit defined in the instruction parameters
# and op=nl.add defines the actual element-wise operation needed
nisa.tensor_tensor(dst=output, data1=x, data2=y, op=nl.add)
Note how many conditions the comments impose. The requirements include data types, layout, tile size, and the memory type of the buffer. It is also important to note that the destination is explicitly specified as an argument, a characteristic not typically exposed through the framework. Instead of receiving a return value, you provide the location where the data should be written.
5.3 Execution Schedule, and the Synchronization That Is Automatic
The third item is the execution schedule. The part that does not have to be taken over is stated explicitly here, so the line is worth drawing precisely.The four compute engines execute four independent instruction streams asynchronously in parallel. Atomic semaphores in hardware handle the synchronization between engines that satisfies data dependencies. Furthermore, the Neuron documentation clearly states that the programmer does not need to manually write this synchronization. The Neuron compiler automatically inserts the necessary synchronization at compile time, based on data dependencies identified within the NKI kernel.
So what is taken over regarding the execution schedule is the order the instructions are placed in, not the synchronization between engines. If you mistakenly assume that you need to manually handle all synchronization, you will unnecessarily introduce complexity.
Now, to what extent is the order in which you write the code preserved? A qualifying clause trails this one, so it has to be read to the end. The documentation states:
NKI uses a sequential programming model where operations run in the order they're written. However, the compiler may change the order of operations that don't depend on each other to make the code faster. This approach gives predictable execution while letting the hardware's multiple compute engines work in parallel behind the scenes.
The default behavior is to follow the order in which you wrote the code. However, the compiler may reorder operations that are not dependent on each other to improve speed. Operations with dependencies will remain unchanged. The documentation explains that this design provides predictable execution while simultaneously allowing the hardware's multiple compute engines to operate in parallel.
Therefore, the understanding that you control the order of instructions in NKI is only partially correct. The order is preserved for operations with dependencies, while the compiler handles operations that are not dependent. If you encounter issues, the documentation provides a how-to guide that allows you to control the execution order using dependency edges and blocks that prevent reordering. If you want to prevent reordering in a specific area, you must explicitly disable it.
6. Tile-Level Semantics, and Why Triton-Like Is Not Triton
Readers coming from a GPU background often initially focus on the syntax, and that familiarity can easily lead them into a trap.6.1 The Term the Official Documentation Uses Is Triton-Like
The Trainium research page describes the NKI syntax as follows:It is a Python-based programming environment which adopts commonly used Triton-like syntax and tile-level semantics.
Triton-like. That is, like Triton, and not Triton.
The Neuron documentation uses a slightly different phrasing, adding another point of comparison:
which has syntax and tile-level semantics that are similar to Triton and NumPy
It states that it is similar to both Triton and NumPy. The inclusion of NumPy is a significant factor. The way tensors are allocated and indexed uses NumPy terminology, while selecting instructions happens at a layer below Triton. This two-layered structure is evident here.
Summarize NKI as the Trainium version of Triton and that two-layer structure disappears. Any estimate of a porting effort built on that summary will be wrong.
6.2 The Two APIs Sit on Different Layers
The APIs provided by NKI come in two sets.nki.language is the upper layer. It simplifies memory allocation, tensor indexing, and control of logical NeuronCore groups. The documentation states that data scientists and machine learning engineers familiar with NumPy and Triton should find it readily accessible.nki.isa is the lower layer. It provides direct access to the Neuron Instruction Set Architecture, or NISA. This allows for a direct mapping of operations to hardware instructions, providing complete control over instruction selection, scheduling, and allocation.These two APIs are designed to be used together, not exclusively.
nki.language simplifies indexing and memory operations, while nki.isa provides the hardware-specific details necessary to maximize efficiency. The two code examples presented in Section 5 correspond to these two layers, respectively.In addition to these two APIs, the API reference also includes the core
nki and nki.collectives. Collective communication exists as an independent namespace.6.3 The Hardware Decides the Shape of a Tile
The specific requirements of tile-based semantics are reflected in the memory structure. According to the architecture guide for the second-generation NeuronDevice, both SBUF and PSUM are treated as two-dimensional memory structures, each with 128 partitions. The Tensor Engine of the same generation is built as a systolic array composed of processing elements arranged in a 128-row, 128-column configuration, taking input from SBUF and writing output to PSUM.The number 128 is not a free variable for whoever writes the kernel. Tiles get divided to fit that shape. The reason the beginning of the tensors allocated in the code example in Section 5 was 128, and the shape of the arguments for
nc_matmul also started with 128, is due to the same constraint.Which generation a figure belongs to always has to be checked, though. The values mentioned above are taken from the architecture guides for Trainium/Inferentia2, specifically the pages dealing with the second-generation NeuronDevice. The key takeaway is not the specific value of 128 itself, but rather the underlying structure where the hardware dictates the shape of the tiles.
This constraint does not emerge automatically from experience with Triton. If there is a tendency to optimize tile sizes on the GPU for performance, that freedom is significantly limited here. Detailed information for each generation can be found in the architecture guides for Trainium/Inferentia2, Trainium2, and Trainium3. Neuron recommends that new developers begin by reading the Trainium/Inferentia2 guide. This article does not provide a generation-specific table of values, as these numbers are subject to change and vary between generations.
7. It Looks Like Python and It Is Not Python
Every code example so far looks like Python. However, the Neuron documentation explicitly states that it is not Python.While the NKI meta-programming language looks and feels like Python, it is not actually Python code.
When the Python interpreter encounters a top-level function decorated with
@nki.jit, the NKI Compiler is invoked to handle the compilation of that function.# this is a Python function that calls 'kernel', which is a NKI kernel
def a_function(x,y,z):
kernel(x, y, z)
# this is a NKI kernel that will be compiled by the NKI Compiler and
# integrated back into the overall model by the Neuron Graph compiler
@nki.jit
def kernel(x,y,z):
# this is kernel code
The function above is executed as Python, while the function below is compiled. Both appear to be Python functions.
7.1 Some Python Features Are Not Available
As a consequence of this design, using Python features that NKI does not support will result in errors. The documentation explains this behavior and the reasons behind it.Neuron has intentionally constrained the NKI meta-programming language to be as minimal as possible while serving the needs of building high performance kernels for today's popular models and will continue to grow and evolve the language over time.
This phrasing indicates a deliberate effort to keep the language minimal. It also states that the language will continue to evolve in the future.
This point has implications for estimation. If plans are based on the idea of porting existing Python code to kernels, the plan turns into a rewrite the moment it meets an unsupported feature. It is more accurate to understand that these features are not simply unimplemented; rather, they are intentionally not supported by the design.
7.2 A Loop Has Two Meanings
Loops are a particularly effective example to illustrate constraints. NKI has two different behaviors for loops, and how the loop is written decides which one applies.Python loops using
range() are fully expanded at specialization time. The compiler creates a copy of the loop body for each iteration, and the loop variable is resolved as a compile-time constant within each copy. Because expansion occurs before compilation, metaprogramming can be performed within the loop body. This allows for branching based on the loop variable, or calculating the shape of tiles, for example.On the other hand,
nl.dynamic_range creates loops that run on the device at runtime. The compiler does not expand these loops. This is used when the number of iterations is not known at specialization time, or when expanding the loop would result in too many instructions, preventing compilation. The documentation states that this is often the case when iterating over tiles of large activation tensors.By simply choosing how to write the loop, you can determine whether it will be expanded at compile time or executed at runtime. This distinction does not exist in standard Python loops.
7.3 Some Ways of Writing Have Been Deprecated
This area is currently under development, and the documentation listsnl.static_range, nl.sequential_range, and nl.affine_range as deprecated. These functions behave as aliases for range(), meaning they will be expanded during specialization and continue to function, but the documentation recommends using range() directly in new code.If you encounter these three functions in existing kernel examples or internal code, they are candidates for refactoring. They work, so nothing breaks, but there is no guarantee the deprecated state will last.
8. Four Names, and the Sources Do Not Agree on How Many
Here, the following sections address instances where the product page and developer documentation present conflicting information. The first is how many pieces there are.There are four similar names: Neuron Kernel Interface, NKI Compiler, Neuron Kernel Library, and Neuron Explorer. AWS's own documentation provides two different answers regarding the relationship between these four.
8.1 The Product Page Puts the Four Side by Side
The product page for the Neuron SDK presents three offerings designed for performance engineers. These are accessed through NKI, provide insights via Neuron Explorer, and offer an optimized kernel library called NKILib. NKI, NKILib, and Neuron Explorer are presented as three items at the same level, and the NKI Compiler joins them there.The entire page is written within this framework, even the description of NKI itself.
NKI provides APIs for memory allocation, and execution scheduling and direct access to the Trainium ISA enabling control over instruction-level programming.
Within this sentence, NKI refers solely to the APIs. The compiler and library are mentioned outside of this specific sentence.
8.2 The Developer Documentation Puts Three of Them Inside NKI
Meanwhile, the About page for Neuron documentation demonstrates a completely different nested structure.Neuron Kernel Interface (NKI) is an open source tool for developing kernels for Trainium hardware. It has three main parts:
The text states that NKI consists of three parts. The first is the NKI Programming Interface, which provides two APIs:
nki.language and nki.isa. The second is the NKI Compiler, built on MLIR, which translates NKI kernel code into optimized hardware instructions while preserving the execution order and memory allocation specified by the developer. The third is the NKI Library, which provides optimized kernels that can be used directly or studied for learning.Within this framework, both the NKI Compiler and the NKI Library sit inside NKI. Neuron Explorer, however, is not among the three; it is a separate tool.
8.3 The Question Is Not Which Is Correct

The issue is that when dealing with these four elements in design documents or internal agreements, it is unclear which framework they are described within. For example, when it is stated that NKI is being used, it is unclear whether that refers only to the API, or to the entire suite including the compiler and libraries. This ambiguity can significantly affect the scope of estimates.
So the practical handling is this. Call the four by four separate names, and write down once what each of them does. The tool for writing is the NKI Programming Interface, the tool that compiles is the NKI Compiler, the place to borrow from is the NKI Library, and the tool that measures is Neuron Explorer. The functions of these four elements remain the same, regardless of which framework they are described within. Only the terminology changes.
8.4 The Library Has Four Spellings
Another mechanical, but problematic, issue exists. The same library is referred to with four different spellings, depending on the documentation.The product page lists it as "Neuron Kernel Library," and uses "NKILib" as an abbreviation. The About page in the Neuron documentation lists it as "NKI Library," using "NKI-Lib" as an abbreviation. The top page and dedicated pages within the documentation consistently use "NKI Library." On GitHub, the repository is named
aws-neuron/nki-library.This actually causes problems when searching. If you search using only one spelling, you will not find pages that use a different spelling. This article uses NKI Library throughout, matching the spelling in the Neuron documentation.
9. What the Open Source Claim Actually Covers
The second divergence concerns the state of the open source release. The product page states that it is complete, while the developer documentation indicates that it is still to come.9.1 The Product Page Writes in the Present Tense
The wording on the Neuron SDK product page is this:NKI Compiler is open sourced, built on MLIR, and provides developers visibility into the complete compiler pipeline.
Another section on the page lists specific components in more detail:
While we move more parts of our stack to opensource, NKI Compiler, Neuron Kernel Driver, Neuron Kernel Library, NxD Inference, Neuron Explorer, and PyTorch, JAX and vLLM integrations are available today fully open-sourced.
The text specifies that these components are "available today," meaning they are currently open-sourced.
9.2 The Developer Documentation Writes in the Future Tense
Meanwhile, the dedicated page for the NKI Compiler addresses the same topic in the following way:Neuron is planning to release the source code for the NKI Compiler to increase awareness and transparency, to enable easier development of tools, and to invite participation and collaboration as we evolve the NKI language.
This phrase, "is planning to release," signifies a plan for public release. The subsequent text is also entirely in the future tense, describing how developers will be able to obtain the source code, modify it, build it, and utilize the self-built compiler within their development workflows.
The description of where to obtain the code states:
To do this, developers will be able to download our sources from our public git repository in the aws-neuron GitHub organization (aws-neuron/).
Only the organization name is provided; the repository name is not included. The text within the parentheses ends with
aws-neuron/.9.3 What Is Settled and What Is Not
Putting the two side by side, here is what can be said.What is confirmed is that the NKI Compiler is built on MLIR. This is consistent across the product page, the About page, and the dedicated NKI Compiler page. The About page also states that using MLIR allows connection to the LLVM ecosystem and the compiler research community, and that it provides visibility into the compilation process, from Python code to final hardware instructions.
What could not be confirmed is whether the source code is published today. Looking at the release history, the announcement for AWS Neuron 2.27.0 described the NKI Compiler, built on MLIR, as a private beta. NKI itself and Neuron Explorer then moved from Beta to Stable in 2.29.0, but no announcement stating a change of state for the compiler source release could be found.
So this article handles it as follows. That it is built on MLIR is stated. That the source release is complete is not. And neither the repository name nor the license is given. Given that the primary sources only provide the organization's name, adding the repository name would be speculative.
If open source availability is a premise of the design, that single point is worth confirming firsthand. Any plans based on the assumption that the compiler can be built and replaced independently should only be developed after confirming this point.
9.4 Some Parts Are Verifiably Open
In some instances, it is possible to confirm that certain components are publicly available. The source code for the NKI Library is available asaws-neuron/nki-library, and the kernel samples for tutorials are available as aws-neuron/nki-samples, with repository names clearly indicated. Regarding Neuron Explorer, the product page states that the user interface is open source on GitHub.Which parts of the Neuron stack are open is not uniform. It is necessary to verify the availability of each component individually.
10. Which Chips: Three Answers in One Documentation Set
The third divergence concerns which chips this can be used on. Of the three divergences, this one carries the most consequence, because it decides adoption directly.10.1 Three Answers Coexist
There are three different answers presented.The product page states that both Trainium and Inferentia are supported. The Trainium research page states:
Neuron Kernel Interface (NKI) is a new programming interface for AWS AI chips, Trainium and Inferentia.
The top page of the Neuron documentation also aligns with this. It describes the programming target as AWS Trainium and Inferentia NeuronDevices.
The FAQ inside that same documentation set draws the range more narrowly.
NKI supports Trainium2 and Trainium3 chips, available in the following instance types: Trn2 and Trn3.
It only mentions Trainium2 and Trainium3, excluding Inferentia. Furthermore, each page of the NKI documentation includes a banner indicating the target. On every one of those pages the banner reads
This document is relevant for: Trn2, Trn3. This matches the description in the FAQ.A third answer is found on the About page. The page begins by describing NKI as a tool for Trainium, but when listing the supported NeuronDevice generations, it currently states:
NKI currently supports the following NeuronDevice generations:
And the first line reads:
Trainium/Inferentia2, available on AWS trn1, trn1n and inf2 instances
It lists trn1, trn1n, and inf2. These are instance types that the FAQ excludes.
10.2 The Split Runs Inside a Single Page
Within a single page, inconsistencies can also arise. The About page, for example, presents two different messages. It mentions Trainium initially, but then includes Inferentia2 in the generation list.Therefore, reading a single page and deciding adoption from it can produce the wrong decision. If you are planning to develop a kernel for Inferentia2, relying solely on the About page might lead you forward, while only reading the FAQ could halt your progress.
10.3 By Contrast, the Neuron Explorer Side Is Explicit
Within the same documentation set, Neuron Explorer clearly specifies the supported hardware. It lists trn1, trn2, and trn3 for Trainium and inf2 for Inferentia 2, and then adds this line:Neuron Explorer is not supported on Inf1.
It explicitly lists the unsupported option. Furthermore, the banner at the top of the page also states
Inf2, Trn1, Trn2, Trn3, which aligns with the main text.Within the same documentation set, one side is explicit and the other is split. It seems natural to interpret the NKI documentation as being in a transitional state, but that is this article's inference and not something AWS has stated.
10.4 How to Handle This in Practice
This article does not definitively state which option is correct. With the primary sources split, declaring one would be speculation.It can still be turned into something usable for a decision.
For a plan on Trn2 or Trn3, every source says it is supported. There are no discrepancies here.
For a plan on trn1, trn1n, or inf2, the sources disagree. If you intend to use these, it is advisable to confirm with AWS or verify the information in the release notes for the specific version of the Neuron SDK. This article will leave this range unverified.
The regional coverage of the AWS Neuron SDK itself is a separate matter. The announcement for AWS Neuron 2.32.0 states that it is available across all AWS Regions where Amazon EC2 Trn1, Trn2, Trn3, Inf2, and Inf1 instances are supported. That the SDK is installable and which chips NKI targets are two different claims. Confusing these two can lead to the incorrect assumption that the availability of the SDK equates to the supported range of NKI.
11. What NKI Does Not Give You
Here are three things NKI does not provide. All of these are explicitly stated in the official documentation, but are not mentioned on the product page.11.1 The Instruction Set Is Not Covered in Full
The NKI FAQ addresses the question of whethernki.isa covers the entire instruction set, stating:Neuron will iteratively add support for the Neuron instruction set through adding more nki.isa (Instruction Set Architecture) APIs in upcoming Neuron releases.
This response indicates that support for the instruction set will be gradually expanded with the addition of more APIs. It does not state that all instructions are currently available.
This differs from what one might naturally assume based on the concept of direct access to instructions. Currently, only instructions supported by the published
nki.isa APIs are accessible through NKI. Before designing systems that rely on specific instructions, it is necessary to verify their existence in the API reference.11.2 There Is No Special Compatibility Guarantee for the APIs
The same FAQ addresses backward compatibility, stating:The NKI APIs follow the Neuron Software Maintenance policy for Neuron APIs.
This response indicates that the NKI APIs adhere to the maintenance policy applied to all Neuron APIs. There is no mention of any additional guarantees specific to NKI.
In practice, the API versions have been updated. NKI transitioned from Beta to Stable with AWS Neuron 2.29.0, with a version of 0.3.0 at that time. Subsequent updates include version 0.4.0 with 2.30.0, version 0.5.0 with 2.31.0, and version 0.6.0 with 2.32.0. As of the verification date of September 1, 2026, the current version is 0.6.0.
With each version update, new features are added, and as seen in Section 7, some methods have been deprecated. The kernel is not a finished product; it requires ongoing updates to align with the SDK. That maintenance load did not exist on the path through the framework.
11.3 The Documentation Does Not Always Keep Up With the Tools
Third, this is not about guarantees, but about practical challenges. As mentioned in Section 3, the page detailing NKI performance optimization still referencesneuron-profile, a tool that is no longer supported. The screen descriptions also refer to that tool.The transition is still reflected in the URLs. The NKI how-to guide includes a page describing how to profile the NKI kernel using Neuron Explorer. The title points at Neuron Explorer, and the URL is still
use-neuron-profile.html. The content is new, but the address is outdated.This is not limited to Neuron; the impact is greater at lower levels. Even if the documentation for higher-level components is outdated, there are usually alternative procedures that work. At the instruction and memory layers, a missing tool stops the work there.
Deciding to descend includes accepting that the documentation will be caught mid-transition.
12. When Not to Descend
With all of the above in hand, here is when it is better not to descend.12.1 The Official Documentation Gives Only Two Reasons to Descend
NKI's official documentation lists only two primary reasons for using it.First, it is for performance optimization. The rationale is that while the Neuron Compiler offers general optimizations, there are cases where it does not fully exploit the hardware for a specific workload. In such cases, NKI allows users to write custom kernels to maximize compute and memory throughput.
Second, it is for implementing new operators and architectures. The documentation states:
NKI enables you to implement operators that are not yet supported by the Neuron Compiler, letting you self-serve new deep learning architectures and custom operations without waiting for compiler support.
The reason is to serve yourself rather than wait for the compiler to catch up.
If neither of these reasons applies, then there is no official justification for using NKI.
12.2 What a Decision Not to Descend Looks Like
Sections 3 through 11 rearrange into a decision.If the NKI Library already holds a matching kernel, do not descend. Attention, MLP, RMSNorm with quantization, and collective communication are already in there, and the library continues to expand. Using what is already in the library is less total work than maintaining something written in-house.
If nothing has been measured yet, do not descend. When examining system profiles in Neuron Explorer, you may find that the dominant factors are framework overhead or CPU wait times. In such cases, even if you significantly speed up the kernel, the overall performance will not improve noticeably.
If the target is neither Trn2 nor Trn3, verify first. As mentioned in Section 10, the sources disagree about trn1, trn1n, and inf2. You need to confirm the information before making any assumptions.
If the goal is only to make the whole model faster, do not descend. NKI is beneficial when you have a specific operator where you know you can optimize it better than the compiler. Until you can clearly articulate the specific reasons why your implementation would outperform generic optimizations, there is no compelling reason to proceed.
If only one person will be responsible for maintaining it, proceed with caution. Kernels need to be kept in sync with the SDK versions, and deprecated coding practices may arise. The ability to have someone else read and understand your code when the original author leaves is not a performance issue; it is an operational concern.
12.3 What Justifies Descending
Descending is justified when the following conditions hold together.Measurement shows that one particular kernel dominates. A device profile shows which engine that kernel runs on and how much time it takes there.
That kernel is not available in the NKI Library, or, even if it is, it does not match the desired shape or accuracy requirements.
There are hypotheses regarding the reason for the performance issues. The documentation treats the ideal endpoint of optimization as landing on one of two outcomes, either strictly compute-bound or memory-bound. While it is unknown which of the two applies, a rewrite is guesswork.
The target hardware has been finalized. Since the architecture guidelines vary depending on the generation, confirming the target hardware also solidifies the assumptions regarding the tile structure.
12.4 The Experiment Is Cheap, and Some Things Cannot Be Tried
One point before leaving the decision open. Writing a kernel does not itself require a Trainium instance.NKI has a CPU simulator.
nki.simulate runs NKI kernels on the CPU using Python and NumPy, and does not require Trainium hardware. It executes all operations defined in nki.isa and nki.language, producing results that closely mimic the hardware's behavior. Kernels decorated with the same @nki.jit will run both in the simulator and on the actual hardware. You can use breakpoint(), PDB, or IDE debuggers to step through the code and inspect tensor values. Checking whether the kernel can be written at all is therefore a local exercise.However, this API is experimental. The documentation includes a warning at the beginning, stating that it may change in future releases.
There are four things that cannot be tested in the simulator. This is the core topic of this section.
First is performance. The simulator does not model instruction latency, engine parallelism, or hardware scheduling. The documentation explicitly states that it is not suitable for measuring performance. If performance is the reason for descending, the simulator cannot test that reason.
Second is a subset of Python. Because the simulator directly executes kernel code as Python, it lacks a compilation stage and accepts any valid Python. As Section 7 showed, what the NKI Compiler accepts is a constrained subset of Python. A kernel using a construct the NKI Compiler does not support therefore runs on the simulator and fails when compiled for hardware.
Third is memory capacity. The simulator allocates each tensor independently and does not verify overlapping regions or the capacity limits of SBUF and PSUM. Kernels that have memory conflicts will pass in the simulator, but may fail or produce incorrect results on the actual hardware.
Fourth is collective communication. The
nki.collectives API is not implemented in the simulator.In summary, it is inexpensive to test correctness, but not speed. And, passing the simulator does not guarantee that the kernel will run on the actual hardware. It is a tool to help you make decisions, not a replacement for making them.
13. Failure Modes and Anti-Patterns
Reading NKI as an API for going faster. This is the biggest misconception. What NKI hands over is the right to decide, not speed. Take the right and the obligation to decide arrives with it. Estimate the work on the assumption that writing a kernel makes things faster, and the estimate will always be wrong.Not consulting the NKI Library before writing. This skips the step the official documentation places in the opening lines of its top page. Attention, MLP, RMSNorm with quantization, and collective communication are already in there. And that included information grows with every release.
Writing before measuring. By examining the system profile in Neuron Explorer, you may discover that the source of the slowness lies outside the kernel. If speeding up the kernel does not result in overall improvement, that effort will be entirely wasted.
Assuming NKI is exclusively for Trainium, or conversely, assuming it can be used with Inferentia. As detailed in Section 10, the primary sources split three ways. Neither direction can be asserted. If you are working with trn1, trn1n, or inf2, verify your assumptions with the relevant SDK version.
Summarizing NKI as a Trainium-specific version of Triton. The official documentation uses the term "Triton-like," and explicitly states that it resembles both Triton and NumPy. This simplification overlooks the two-layered structure of
nki.language and nki.isa, and your estimates for porting will inevitably be inaccurate.Assuming it can be written in Python. The documentation clearly states "it is not actually Python code," and the meta-programming language is intentionally kept to a minimum. Plans to incorporate existing Python code will likely require a complete rewrite once you encounter unsupported features.
Planning based on the assumption that the NKI Compiler's source code is publicly available. The product page states "available today," while the dedicated page says "is planning to release." The repository name is also not mentioned in the primary sources. If you are planning a self-build, confirm that point first.
Using
nl.static_range, nl.sequential_range, and nl.affine_range in new code. All of these are deprecated and function as aliases for range(). They work and will not break, but the documentation recommends using range() directly in new code.Taking a NKI Samples kernel into production. The FAQ states that Samples is for learning and experimentation, while the NKI Library is for production use. Both are official repositories, so where the code came from does not tell them apart.
Treating a pass on the CPU simulator as evidence that the kernel runs on hardware. The simulator accepts any valid Python, so a kernel using a construct the NKI Compiler does not support runs there anyway. It does not check SBUF and PSUM capacity either. Passing does not mean the kernel will compile.
Measuring speed on the CPU simulator. The documentation states that the simulator models neither instruction latency nor engine parallelism nor hardware scheduling, and is not suitable for performance benchmarking. If performance is the reason for descending, the simulator cannot test that reason.
Writing the synchronization between engines by hand. The Neuron compiler inserts the required synchronization automatically, based on the data dependencies it identifies. What you take over is the order of the instructions, not the synchronization itself.
Continuing to use a runbook built around
neuron-profile. Neuron Profiler and Neuron Profiler 2.0 reached end of support in Neuron 2.29. Because the NKI performance optimization page still points at neuron-profile, a reader following the documentation ends up looking for the older tool.Reading the Regions where the Neuron SDK is offered as the range of chips NKI supports. That the SDK is available in the Regions where Trn1, Trn2, Trn3, Inf2, and Inf1 instances are available, and which chips NKI targets, are separate claims.
Carrying a tile figure from the generation it was verified against into a different generation. The architecture guides are split into Trainium/Inferentia2, Trainium2, and Trainium3 because the contents differ.
14. Frequently Asked Questions
Will using NKI definitely make things faster? No. The official documentation states that it allows you to write custom kernels when the compiler's general-purpose optimizations are not sufficient for your specific workload. No specific speedup is mentioned. This article has not measured anything either, so it gives no multiplier.If I can write Triton code for the GPU, can I just port it over to NKI? No. While the syntax is similar to Triton, the documentation notes that it resembles both Triton and NumPy. But the API is divided into two layers:
nki.language and nki.isa, and the hardware's 128 partitions fix the shape of a tile. The familiar syntax does not necessarily guarantee ease of porting.Can NKI be used on Inferentia? This cannot be stated either way, because the primary sources disagree. The product page and the top page of the Neuron documentation mention both Trainium and Inferentia, and the About page includes "inf2" in the list of supported generations. The FAQ, on the other hand, lists only Trainium2 and Trainium3, and the banners on each page only show Trn2 and Trn3. On Trn2 and Trn3, every source agrees.
Is NKI's code written in Python? No. The documentation explicitly states that while it may look and feel like Python, it is not actually Python code. The NKI Compiler handles any function decorated with
@nki.jit. Using Python features that are not supported will result in errors.Is the NKI Compiler source code published? This could not be confirmed. The product page currently states that it is fully open source, and the dedicated NKI Compiler page mentions that its release is planned for the future. All documentation consistently states that it is built on top of MLIR.
What should I do before writing a kernel? There are two things to do. First, check if there is a corresponding kernel available in the NKI Library. Second, use Neuron Explorer to identify performance bottlenecks. The official documentation says, in the opening lines of the NKI top page, to check the library before writing.
Does NKI allow you to call every instruction available on the hardware? No. The FAQ states that support for the instruction set will grow step by step through more
nki.isa APIs in coming releases. Currently, only instructions supported by the published API are available.Do you need to write the synchronization between engines yourself? No. The four compute engines run independent instruction streams asynchronously in parallel, but the synchronization required to satisfy data dependencies is automatically inserted by the Neuron compiler. What the author controls is the order of the instructions, and a separate mechanism is provided for that.
Can you test the kernel without renting Trainium instances? You can test its correctness.
nki.simulate runs the NKI kernel on a CPU using Python and NumPy, and does not require any hardware. However, the documentation clearly states that this API is experimental and not suitable for performance measurement. The simulator accepts any valid Python and does not check SBUF and PSUM capacity. Therefore, there is no guarantee that a kernel that passes the simulation will compile successfully on a real device, nor is there a guarantee that it will function correctly.What is the difference between the NKI Library and the NKI Samples? The NKI Library consists of production-ready, optimized kernels, while the NKI Samples are reference and tutorial kernels that demonstrate NKI programming patterns. The FAQ explicitly states that the Samples are intended for learning and experimentation, not for production use.
Is a runbook written for the older profiler still usable? No. Support for the Neuron Profiler and Neuron Profiler 2.0 ended in Neuron 2.29, and Neuron Explorer is the recommended tool. A migration guide is available. However, as of September 1, 2026, the page on NKI performance optimization still refers to those older tools.
15. Summary
NKI provides not speed, but control. The official documentation lists instruction sets, memory, and execution scheduling as targets for low-level access. As long as you are writing models within a framework, you typically will not interact with these directly. You were able to avoid interacting with them because the Neuron compiler was making those decisions for you.What you take over is the last of the four stages. The Neuron compiler runs through four stages in sequence: graph-level optimization, loop-level optimization, instruction mapping, and hardware-specific optimization. A NKI kernel skips the first three stages and is compiled into an intermediate representation that the fourth stage consumes directly. Automation is removed, but the total amount of decisions required does not decrease.
There are two things to do before you begin. First, verify whether a kernel exists within the NKI Library. Second, use Neuron Explorer to measure performance bottlenecks. The official documentation says, in the opening lines of its top page, to check the library before writing anything.
The familiar syntax does not necessarily mean easy portability. The documentation uses the term "Triton-like," and the documentation itself states that it resembles both Triton and NumPy. The API is structured in two layers:
nki.language and nki.isa, and the shape of the tiles follows the hardware's 128-partition format. And what you are writing is not Python.The four names are grouped differently depending on the documentation. The product page lists NKI, the NKI Compiler, the NKI Library, and Neuron Explorer as separate entities, while the developer documentation places the first three within the scope of NKI, and lists Neuron Explorer separately. It is not about which is correct; the problem is that when the design documents refer to "NKI," the scope is not always clear. It is more reliable to refer to each by its distinct name and clearly describe its role.
The status of the open source release could not be confirmed. The product page states that it is complete, while a dedicated page indicates that it is planned. The fact that it is based on MLIR is consistent across all sources. The repository name does not appear in the primary sources, so this article does not supply one.
The supported chips are split three ways across the sources. Every source agrees on Trn2 and Trn3. The sources disagree about trn1, trn1n, and inf2, and this article leaves that unverified.
There are tools to simplify decision-making, but they do not replace it.
nki.simulate allows you to run the kernel on the CPU without hardware. However, it is not suitable for performance measurement, it accepts Python that the NKI Compiler will not, and it does not check memory capacity. You can easily check for correctness, but not for speed.There are cases where it is better not to descend. They are the case where the NKI Library already holds a match, the case where nothing has been measured yet, the case where the goal is only to make the whole model faster without a specific target, and the case where only one person will maintain the result. The two reasons the official documentation gives for descending are performance optimization and operators the compiler does not yet support. If neither applies, no reason to descend appears in the official documentation.
16. References
- AWS Neuron - AWS
- Build on Trainium - AWS
- Amazon EC2 Trn3 UltraServers - AWS
- Neuron Kernel Interface (NKI) - AWS Neuron Documentation
- About the Neuron Kernel Interface (NKI) - AWS Neuron Documentation
- NKI FAQ - AWS Neuron Documentation
- About the NKI Compiler - AWS Neuron Documentation
- NKI Library Documentation - AWS Neuron Documentation
- NKI CPU Simulator - AWS Neuron Documentation
- NKI Guides - AWS Neuron Documentation
- NKI Deep Dives - AWS Neuron Documentation
- NKI Performance Optimizations - AWS Neuron Documentation
- Trainium/Inferentia2 Architecture Guide for NKI - AWS Neuron Documentation
- Neuron Explorer - AWS Neuron Documentation
- Model optimization techniques - AWS Prescriptive Guidance
- AWS Neuron introduces Neuron Kernel Interface (NKI), NxD Training, and JAX support for training - AWS What is New
- Announcing AWS Neuron SDK 2.27.0 - AWS What is New
- AWS Neuron SDK 2.29.0 now available with NKI and Neuron Explorer out of Beta, a new CPU Simulator, and an expanded NKI Library - AWS What is New
- AWS Neuron 2.30.0 now available with NKI 0.4.0 and 22 new NKI Library kernels - AWS What is New
- AWS Neuron 2.31.0 now available with NKI 0.5.0 and UltraServer Operator - AWS What is New
- AWS Neuron 2.32 introduces expanded NKI programming, MXFP8 training kernels, and variable-size collectives for Trn2 and Trn3 - AWS What is New
- Reproducible LLM Inference - hidekazu-konishi.com
- AWS Custom Silicon History and Timeline - hidekazu-konishi.com
- Amazon EC2 Instance Types History and Timeline - hidekazu-konishi.com
- Self-Managed LLM Inference on Amazon EKS - hidekazu-konishi.com
- Disaggregated Prefill and Decode for LLM Serving on AWS - hidekazu-konishi.com
- LLM Weight Quantization on AWS - hidekazu-konishi.com
- Elastic Fabric Adapter and the AWS Network Fabric - hidekazu-konishi.com
- Model Customization on Amazon Bedrock - hidekazu-konishi.com
- Distributed Training Resilience on AWS - hidekazu-konishi.com
References:
Tech Blog with curated related content
Written by Hidekazu Konishi