Why Custom Silicon?

NVIDIA GPUs are general-purpose parallel processors. They're great, but they pay a "generality tax" — transistors spent on features AI workloads never use. AWS asked: what if we built a chip that does nothing but tensor math, and does it insanely well?

The answer is Trainium (for training) and Inferentia (for inference). These are purpose-built AI accelerators with a radically different memory hierarchy than GPUs.

NKI is the programming interface that lets you write custom kernels for these chips. If CUDA is how you talk to NVIDIA hardware, NKI is how you talk to Neuron hardware.

The Memory Hierarchy (A Shipping Port)

Trainium's memory system has three levels. To understand them, imagine a shipping port:

The golden rule: data must travel HBM → SBUF before any computation, and results must travel SBUF → HBM to be visible outside the kernel. PSUM is an intermediate stop only for matmul results.

Data Flow: Watch It Move

Every NKI kernel follows this data flow. Click "Next Step" to see how a tile of data moves through the chip:

Tile Semantics: Why 128?

Everything on Trainium operates on tiles. A tile is a 2D block of data with a fixed shape: 128 rows (the partition dimension) by up to 512 columns (the "free dimension").

Why 128? Because SBUF is physically organized into 128 parallel memory banks. Each bank holds one row of your tile. This is not a software choice — it is baked into the silicon.

Two API Levels

NKI gives you two ways to talk to the hardware. Think of them as "Python" vs "Assembly":

In this course, we start with nki.isa so you understand exactly what's happening. Once you've built intuition, the high-level nki.language API will feel like a natural shortcut.

DMA: The Forklift

DMA stands for Direct Memory Access. It is a dedicated hardware engine whose only job is moving tiles between HBM and SBUF. Think of it as an autonomous forklift that carries containers between the cargo ship and the dock warehouse.

Key properties of DMA on Trainium:

• It runs asynchronously — compute can happen while DMA is moving the next tile
• It always moves full 128-row tiles (partition dimension)
• It can handle strided access patterns (non-contiguous memory)

Check Your Understanding

The Universal Kernel Pattern

Every NKI kernel — from simple addition to Flash Attention — follows the same three-step pattern:

That's it. Load, compute, store. The complexity of advanced kernels comes from what you compute and how you overlap these steps — not from a different pattern.

The Complete Kernel

Here is a tensor addition kernel that adds two matrices element-wise. Left side is code, right side is plain English:

@nki.jit
def nki_tensor_add(a_input, b_input):
  c_output = nl.ndarray(
    a_input.shape, dtype=a_input.dtype,
    buffer=nl.shared_hbm)
  M, N = a_input.shape
  TILE_M = 128
  TILE_N = 512

  for m in nl.affine_range(M // TILE_M):
    for n in nl.affine_range(N // TILE_N):
      a_tile = nl.ndarray(
        (128, 512), dtype=a_input.dtype,
        buffer=nl.sbuf)
      b_tile = nl.ndarray(
        (128, 512), dtype=b_input.dtype,
        buffer=nl.sbuf)

      nisa.dma_copy(
        dst=a_tile,
        src=a_input[m*128:(m+1)*128,
                    n*512:(n+1)*512])
      nisa.dma_copy(
        dst=b_tile,
        src=b_input[m*128:(m+1)*128,
                    n*512:(n+1)*512])

      c_tile = nl.ndarray(
        (128, 512), dtype=a_input.dtype,
        buffer=nl.sbuf)
      nisa.tensor_tensor(
        dst=c_tile, data1=a_tile,
        data2=b_tile, op=nl.add)

      nisa.dma_copy(
        dst=c_output[m*128:(m+1)*128,
                     n*512:(n+1)*512],
        src=c_tile)
  return c_output

Line-by-Line Breakdown

Watch the Components Talk

Here's how HBM, SBUF, and Compute coordinate during one iteration of our addition kernel. Click "Next" to advance the conversation:

Buffer Allocation: Where Things Live

Notice the buffer= parameter in every nl.ndarray call. This tells the compiler where to place the tensor:

If you forget to specify buffer=nl.sbuf for your working tiles, the compiler won't know where to put them and will error. Always be explicit about where data lives.

Check Your Understanding

What You've Learned

In the next module, we'll tackle matrix multiplication — where PSUM finally comes into play and we see the systolic array in action.

Why Tiling Is Not Optional

Tiling is the single most important concept in NKI programming. Without it, nothing works.

Consider a typical tensor in a transformer: 4096 rows by 4096 columns of float16 values. That is 4096 * 4096 * 2 bytes = 32 MB. SBUF holds a few kilobytes. The mismatch is roughly 1000x.

The metaphor: tiling is like reading a newspaper. The full broadsheet page won't fit on your desk, so you fold it into sections and read one section at a time, moving your "window" across the page. The news is the same — you just process it in pieces.

Partition Dimension vs. Free Dimension

Every NKI tile has exactly two axes, and they are not interchangeable:

So a tile is always shaped (128, N) where N is at most 512. The partition dimension is hardware-mandated. The free dimension is your choice (within limits).

The Basic Tiling Pattern

Here is the fundamental loop structure that appears in every NKI kernel. It breaks a large tensor into 128x512 tiles and processes them one at a time:

M, N = tensor.shape
TILE_M = 128   # partition dim
TILE_N = 512   # free dim

for m in nl.affine_range(M // TILE_M):
  for n in nl.affine_range(N // TILE_N):
    # Slice: rows [m*128..(m+1)*128]
    #        cols [n*512..(n+1)*512]
    tile = tensor[m * TILE_M:(m+1) * TILE_M,
                  n * TILE_N:(n+1) * TILE_N]

Visualizing the Tile Grid

A 1024x2048 tensor split into 128x512 tiles produces a grid of 8 rows by 4 columns — 32 tiles total. Each tile is one loop iteration:

affine_range vs. sequential_range

NKI provides two loop constructs. They look similar but give the compiler very different freedoms:

Rule of thumb: default to affine_range. Only switch to sequential_range when you have a true data dependency between iterations — like accumulating a running maximum for online softmax.

Dynamic Slice: nl.ds(start, size)

Standard Python slicing (a:b) works in NKI, but for computed offsets, dynamic slice notation is cleaner:

# These two are equivalent:

# Python slice notation
data = tensor[:, i*512 : i*512+512]

# Dynamic slice notation
data = tensor[:, nl.ds(i * 512, 512)]

nl.ds(start, size) means "starting at start, take size elements." It is especially useful in nested loops where both the start and size are expressions of loop variables.

Index Arithmetic Decoded

The expression m * TILE_M : (m+1) * TILE_M appears everywhere. Here is what it computes for each value of m:

The same index arithmetic applies to the free dimension with n * TILE_N. Together, the two loop variables select a unique 128x512 rectangle from the full tensor.

When Tiles Don't Divide Evenly

What happens if your tensor has 1000 rows instead of 1024? 1000 is not divisible by 128 (1000 / 128 = 7.8125). You have two options:

In practice, most ML tensors have dimensions that are already multiples of 128 (batch sizes, hidden dimensions, sequence lengths are typically powers of 2). When they are not, padding is the simpler approach.

Check Your Understanding

What You've Learned

The Four Element-wise Primitives

In Module 2 we used nisa.tensor_tensor for addition. That is just one of four element-wise primitives that NKI provides. Together, they form the complete set of per-element building blocks:

Binary Operations: tensor_tensor

The simplest element-wise primitive. Two tiles go in, one tile comes out. Every element at position (i, j) is combined independently:

# Element-wise multiply
nisa.tensor_tensor(
  dst=result,
  data1=a_tile,
  data2=b_tile,
  op=nl.multiply
)

Supported operations: nl.add, nl.subtract, nl.multiply, nl.minimum, nl.maximum. All inputs and output must be in SBUF with matching shapes.

Scalar Broadcast: tensor_scalar

Often you need to apply the same value across an entire row — subtract the row maximum, multiply by a scale factor, add a bias. This is where nisa.tensor_scalar shines:

# Subtract row-max for numerical stability
# row_max shape: (128, 1)
# scores shape:  (128, 512)
nisa.tensor_scalar(
  dst=normalized,
  data=scores,
  op0=nl.subtract,
  operand0=row_max
)

The key insight: row_max has shape (128, 1) but gets broadcast to (128, 512). Each row's single max value is subtracted from all 512 columns in that row. This is identical to NumPy broadcasting — just happening on custom hardware.

The Secret Weapon: Chained Operations

nisa.tensor_scalar can perform two operations in sequence with a single instruction. First op0 is applied, then op1 is applied to the intermediate result:

# Fused multiply-add: out = x * scale + bias
nisa.tensor_scalar(
  dst=out,
  data=x,
  op0=nl.multiply,
  operand0=scale,
  op1=nl.add,
  operand1=bias
)

Without chaining, you would need two instructions: one tensor_scalar for multiply, then another for add. Each instruction reads the tile from SBUF and writes it back. Chaining eliminates the intermediate read-write cycle.

Activations: The Nonlinear Functions

Activation functions are the nonlinear transformations that make neural networks powerful. NKI provides them as a single instruction:

# Compute exp(data * scale)
nisa.activation(
  dst=exp_result,
  op=nl.exp,
  data=delta,
  scale=A_val
)

Available activations: nl.exp, nl.log, nl.tanh, nl.reciprocal, nl.sigmoid. The optional scale and bias parameters let you compute op(data * scale + bias) in a single instruction.

The Operations Argue

Each element-wise operation has its own specialty. Here they are, debating who is the most important:

Match the Operation to Its Use Case

Drag each NKI instruction to its correct real-world application:

Type Conversion: float32, float16, bfloat16

NeuronCore uses mixed precision internally. The Tensor Engine always outputs float32 into PSUM. But SBUF often holds float16 or bfloat16 to save space. You need to convert between them:

The general rule: compute in float32 (for precision), store in float16/bfloat16 (for space). The Tensor Engine enforces this — PSUM is always float32.

Check Your Understanding

What You've Learned

The Hardware That Does the Heavy Lifting

Every neural network is, at its core, a chain of matrix multiplications. NeuronCore has a dedicated Tensor Engine that can multiply a 128×128 matrix by a 128×512 matrix in one shot — producing a 128×512 result.

But the Tensor Engine has rules. The left-hand side (LHS) is stationary — it stays in place inside the engine. The right-hand side (RHS) is moving — it streams through the engine. The result always lands in PSUM as float32, regardless of input types.

Four Steps to an Optimized Matmul

NKI matmul kernels evolve through four stages. Each version addresses a limitation of the previous one.

Reading the Tiled Matmul

Here is the tiled version (Version 2) — the foundation every later optimization builds upon.

for m in nl.affine_range(M // TILE_M):
  for n in nl.affine_range(N // TILE_N):
    res_psum = nl.ndarray(
      (TILE_M, TILE_N), nl.float32,
      buffer=nl.psum)
    for k in nl.affine_range(K // TILE_K):
      lhsT_tile = nl.ndarray(
        (TILE_K, TILE_M), dtype=lhsT.dtype,
        buffer=nl.sbuf)
      rhs_tile = nl.ndarray(
        (TILE_K, TILE_N), dtype=rhs.dtype,
        buffer=nl.sbuf)
      nisa.dma_copy(dst=lhsT_tile,
        src=lhsT[k*TILE_K:(k+1)*TILE_K,
                 m*TILE_M:(m+1)*TILE_M])
      nisa.dma_copy(dst=rhs_tile,
        src=rhs[k*TILE_K:(k+1)*TILE_K,
                n*TILE_N:(n+1)*TILE_N])
      nisa.nc_matmul(dst=res_psum,
        stationary=lhsT_tile,
        moving=rhs_tile)
    res_sb = nl.ndarray(
      res_psum.shape, dtype=result.dtype,
      buffer=nl.sbuf)
    nisa.tensor_copy(dst=res_sb,
      src=res_psum)
    nisa.dma_copy(
      dst=result[m*TILE_M:(m+1)*TILE_M,
                 n*TILE_N:(n+1)*TILE_N],
      src=res_sb)

Data Flow: One Tile Iteration

Watch how data moves through the hardware during a single K-tile iteration of the matmul.

The Optimization That Matters Most

In the basic tiled version (Version 2), both LHS and RHS are loaded inside the innermost loop. Version 3 restructures the loops so the LHS load sits in the M-K loop, outside of N. The Tensor Engine keeps the Tensor Engine's stationary operand loaded while new RHS tiles stream through.

Version 4 goes further by loading RHS in larger contiguous blocks. Instead of fetching one 128×512 tile per iteration, it grabs multiple adjacent tiles in a single DMA burst. This keeps the memory controller busy and hides latency behind computation.

Check Your Understanding

The Formula That Powers Every LLM

Attention computes softmax(Q @ KT) @ V. That single formula is responsible for the magic of ChatGPT, image generators, and protein folders. It is also the most expensive operation in any transformer.

The challenge: softmax requires knowing the maximum value across an entire row before you can compute any single element. This creates a data dependency that makes tiling and fusion difficult.

The NKI samples show how to evolve from a straightforward implementation to one that never materializes the full attention matrix — cutting memory from O(n²) to O(n).

Version 1: The Naive Attention

Version 1 uses high-level NKI operations. It is slow but crystal-clear about what attention actually computes.

@nki.jit
def attn_fwd_v1(q, k, v):
  q_sbuf = nl.load(q)
  k_sbuf = nl.load(k)
  v_sbuf = nl.load(v)

  qk_psum = nl.matmul(
    x=q_sbuf, y=k_sbuf,
    transpose_x=True)
  qk = nl.ndarray(qk_psum.shape,
    dtype=nl.float32, buffer=nl.sbuf)
  nisa.tensor_copy(dst=qk, src=qk_psum)

  row_max = nl.max(qk, axis=1,
    keepdims=True)
  norm_row = nl.ndarray(qk.shape,
    dtype=nl.float32, buffer=nl.sbuf)
  nisa.tensor_scalar(dst=norm_row,
    data=qk, op0=nl.subtract,
    operand0=row_max)

  exp_row = nl.exp(norm_row)
  sum_row = nl.sum(exp_row,
    axis=1, keepdims=True)
  inv_sum = nl.reciprocal(sum_row)

  scores = nl.ndarray(exp_row.shape,
    dtype=nl.float32, buffer=nl.sbuf)
  nisa.tensor_scalar(dst=scores,
    data=exp_row, op0=nl.multiply,
    operand0=inv_sum)
  # ... then scores @ V

The Optimization Journey: v1 to v8

Each version addresses a specific bottleneck. The jump from v1 to v8 delivers orders-of-magnitude improvement.

The Trick Behind Flash Attention

Standard softmax needs two passes: one to find the max, another to compute exp and sum. Online softmax does it in one pass by maintaining corrections.

Check Your Understanding

The Big Picture

The attention optimization journey is really about one thing: avoiding materialization. Version 1 creates the full N×N attention matrix in memory. Version 8 never creates it at all.

Each step is independently understandable. Together, they turn a textbook formula into production-grade hardware-optimized code running on NeuronCore.

Beyond Attention: The Recurrence Alternative

Transformers dominate sequence modeling, but their O(n²) attention cost becomes painful at long sequences. Mamba offers a fundamentally different approach: instead of letting every token look at every other token, it propagates information forward through a recurrence.

The core equation of a state space model is deceptively simple:

This runs in O(n) time — linear in sequence length. But naive sequential execution wastes hardware parallelism. The key insight: this recurrence can be computed as an associative scan.

The Associative Scan: Parallel Recurrence

An associative scan is like a prefix sum, but instead of just addition, we use the pair of operations (multiply, add). This lets us compute all positions of the recurrence in parallel.

NKI provides direct hardware support via a single instruction:

            nisa.tensor_tensor_scan(
                dst=scan_res,
                data0=deltaA,
                data1=deltaBu,
                initial=0.0,
                op0=nl.multiply,
                op1=nl.add
            )
          

Think of it like a relay race. Each runner (position) inherits speed from the previous runner (multiply by deltaA) and adds their own boost (add deltaBu). Normally you'd wait for each runner to finish before the next starts. The hardware scan unit does all handoffs simultaneously — computing the entire race result in one shot.

Anatomy of the Mamba Kernel

The full Mamba kernel on NKI has five stages. The discretization step converts continuous parameters to discrete ones. Then the scan does the heavy lifting.

Notice the loop structure: the kernel iterates over batch, channel tiles, and state dimensions. For each state dimension, it runs the full discretize-scan-project pipeline, then accumulates results. This maps naturally to NKI's tile-based execution.

Why exp(delta * A)?

The matrix A represents a continuous-time decay rate. To use it in discrete steps, we need to convert it: if A says "decay at rate -0.5 per unit time" and our step size is delta, then the per-step decay factor is exp(delta * A).

NKI handles this elegantly with a fused instruction:

            nisa.activation(
                dst=deltaA,
                op=nl.exp,
                data=delta_i,
                scale=A_i
            )
          

Mamba's "selective" innovation is that B and C depend on the input — they are not fixed matrices. This means the model can choose what to remember and what to forget at each step, giving it the expressiveness to match Transformers while keeping linear cost.

Check Your Understanding

Three Power Patterns

With the fundamentals in place, we can explore the patterns that separate competent NKI code from truly efficient kernels. Each exploits a different axis of the hardware.

Pattern 1: 2D Transpose via Element Scatter

NKI has no built-in transpose instruction for free dimensions. Instead, you scatter individual elements to their transposed positions using index math.

For a matrix conceptually shaped (sz_f1, sz_f2) stored in a tile's free dimension, transposing means moving element at position (i, j) to position (j, i). In flattened coordinates: position i*sz_f2 + j moves to j*sz_f1 + i.

            for i_f1 in nl.affine_range(sz_f1):
              for i_f2 in nl.affine_range(sz_f2):
                nisa.tensor_copy(
                  dst=out[:, nl.ds(
                    i_f2*sz_f1 + i_f1, 1)],
                  src=in_tile[:, nl.ds(
                    i_f1*sz_f2 + i_f2, 1)]
                )
          

Pattern 2: Average Pooling with Access Patterns

Pooling requires reading overlapping windows from a 2D grid. Rather than copying data into a reshaped buffer, NKI's access pattern feature creates a virtual multi-dimensional view.

Pattern 3: Software Pipelining

Software pipelining is the NKI equivalent of doing laundry efficiently. Instead of waiting for each load to wash, dry, and fold before starting the next, you overlap them: while load 2 washes, load 1 dries.

The hardware has separate units for DMA (data movement), compute (matrix math), and accumulation that can all run simultaneously. A pipelined Flash Attention kernel exploits this:

The key constraint: you must manage the pipeline prologue (filling the pipe) and epilogue (draining it). The first two iterations only load and partially compute. The last two only compute and store. The steady state in between is where you get full overlap.

Match the Operation

Drag each NKI operation to its purpose.

Pipelining Concepts

With these three patterns — scatter transpose, access pattern views, and software pipelining — you have the tools to write kernels that fully exploit NKI hardware. Transpose handles data layout. Access patterns eliminate copy overhead. Pipelining hides latency. Combined, they let you write kernels that approach the theoretical peak of the silicon.