Decoding NKI Kernel Performance using AWS Neuron Profiler
This blog demonstrates how to use Neuron Profiler with NKI (Neuron Kernel Interface) to analyze custom kernel performance, showing through a practical example how the profiler reveals optimization opportunities across NeuronCore engines.
NKI (Neuron Kernel Interface ) is a powerful low-level programming framework that enables direct access to AWS NeuronCore hardware on Trainium and Inferentia instances. This bare-metal language and compiler system empowers machine learning developers to create and optimize custom operators, maximizing the potential of NeuronCore's compute and memory capabilities. Starting with NeuronCores v2 (Trainium1) and beyond, NKI allows developers to independently explore and implement novel approaches to hardware utilization, ultimately pushing the boundaries of what's possible with AWS ML accelerators.
Neuron Profiler serves as our window into the performance characteristics of NKI kernel execution on Neuron devices. By leveraging dedicated on-chip profiling hardware, it provides detailed performance metrics with near-zero overhead. The tool captures comprehensive data about instruction execution across compute engines and monitors data movement through DMA engines, while measuring critical performance indicators like compute engine utilization and memory bandwidth efficiency. This detailed visibility enables developers to precisely identify performance bottlenecks and optimization opportunities, ensuring maximum hardware efficiency for their applications.
Neuron Profiler with NKI Kernels
For this example we are going to use a kernel that implements a tiled element-wise
exp operation on a 2D tensor using NKI. You can profile a NKI kernel in many ways, for this example we will use @nki.benchmark decorator.Step 1 - Save this python file as
exp.py1
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
27
28
29
30
31
32
33
34
35
36
37
38
"""
Example kernel used to demmonstrate Neuron Profile with nki.benchmark.
"""
from neuronxcc import nki
from neuronxcc.nki.typing import tensor
import neuronxcc.nki.language as nl
import math
def tensor_exp_kernel_(in_tensor):
"""NKI kernel to compute elementwise exponential of an input tensor
Args:
in_tensor: an input tensor of ANY 2D shape (up to SBUF size)
Returns:
out_tensor: an output tensor of ANY 2D shape (up to SBUF size)
"""
out_tensor = nl.ndarray(in_tensor.shape, dtype=in_tensor.dtype,
buffer=nl.shared_hbm)
sz_p, sz_f = in_tensor.shape
i_f = nl.arange(sz_f)[None, :]
for p in nl.affine_range(math.ceil(sz_p / nl.tile_size.pmax)):
# Generate tensor indices for the input/output tensors
# pad index to pmax, for simplicity
i_p = p * nl.tile_size.pmax + nl.arange(nl.tile_size.pmax)[:, None]
# Load input data from external memory to on-chip memory
# only read up to sz_p
in_tile = nl.load(in_tensor[i_p, i_f], mask=(i_p<sz_p))
# perform the computation
out_tile = nl.exp(in_tile)
# store the results back to external memory
# only write up to sz_p
nl.store(out_tensor[i_p, i_f], value=out_tile, mask=(i_p<sz_p))
return out_tensor
if __name__ == "__main__":
tensor_exp_kernel_(tensor[[250, 512], nl.float32])Step 2 - Run this python file
1
python3 exp.pyAfter the execution is complete, you will see that the profile for the NKI kernel has been captured.
1
2
3
4
5
+---+----+---------+---------+---------+---------+---------+------+-------+-------+--------+---------+---------+-------+
B NC NC USED WEIGHTS MODE INF/S IRES/S L(1) L(50) L(99) NCL(1) NCL(50) NCL(99) %USER
1 1 1 dynamic LIBMODE 5123.60 5123.60 53 80 93 20 20 21 N/A
+---+----+---------+---------+---------+---------+---------+------+-------+-------+--------+---------+---------+-------+
Profile results are written to profile.ntffStep 3 - Visualize the profile for this NKI kernel using Neuron Profiler.
1
neuron-profile view -n file.neff -s profile.ntff 
If you look at the middle section of this profile, you will notice significant DMA (Direct Memory Access) activity. By hovering the cursor over this region, you can view the precise DMA throughput measurements at any specific timestamp. This interactive feature allows for detailed analysis of memory transfer patterns and performance characteristics.
To identify specific operations running on different engines, simply hover the mouse over any timeline. A tooltip will appear displaying the current operation's opcode, providing detailed insight into which tasks are occupying each engine. This interactive feature allows us to easily track and analyze the execution flow across our system's engines in real-time.
You can easily identify all instructions that are linked to a specific line of NKI source code by searching for the pattern
/./ in the trace as shown below. This technique is particularly useful when analyzing and optimizing NKI kernels, as it helps pinpoint potential performance bottlenecks.
Moreover, when you hover over the highlighted instruction bars, the tool not only displays the standard instruction details but also reveals the corresponding NKI source code location. This mapping provides clear visibility into which exact line of NKI code triggered the operation on the given engine, making it much easier to connect low-level execution behavior with high-level kernel code.
Complementing its graphical interface, Neuron Profiler provides a detailed tabular view of performance data. This alternative representation presents the same metrics in a structured format, making it easy to analyze specific values like
total_time , event_count etc. The table format is particularly useful for precise analysis and comparison of performance metrics, offering a quantitative complement to the visual insights provided by the graphs above.
In this example, we used Neuron Profiler to profile an NKI kernel and gain insights into its performance characteristics. Profiling with Neuron Profiler not only provides a detailed view of how the kernel behaves at runtime but also helps identify potential bottlenecks, inefficiencies, and optimization opportunities. From the profiling results, we can see that this kernel primarily utilizes three Neuron core engines: the vector, scalar, and GPSIMD engines. However, there is no overlap among them. Additionally, there is no overlap between DMA and compute—the data is fully loaded before computation begins. This indicates that tiling could be applied to better pipeline and parallelize operations across the different compute engines, improving overall efficiency.
This information serves as a valuable guide for tuning and improving kernel performance in the right direction. For a deeper understanding of best practices and optimization strategies, please refer to the NKI Performance Guide which provides a comprehensive collection of patterns and optimization strategies that serve as an excellent starting point when tuning NKI kernels for better performance.
This blog was authored by Esha Lakhotia, Scott Perry, and Sadaf Rasool.
Esha Lakhotia is a Product Manager in Annapurna Labs at AWS. She enables customers to achieve their ML performance optimization goals by building developer tools for AWS AI chips. She focuses on creating intuitive tooling experiences that help customers profile, debug, and accelerate their machine learning workloads on AWS Trainium and AWS Inferentia.
Scott Perry is a Solutions Architect on the Annapurna ML accelerator team at AWS. Based in Canada, he helps customers deploy and optimize deep learning training and inference workloads using AWS Inferentia and AWS Trainium. His interests include large language models, deep reinforcement learning, IoT, and genomics.
Sadaf Rasool is a Solutions Architect in Annapurna Labs at AWS. Sadaf collaborates with customers to design machine learning solutions that address their critical business challenges. He helps customers train and deploy machine learning models leveraging AWS Trainium or AWS Inferentia chips to accelerate their innovation journey.
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article