intro to kernel writing on blackhole

For my own clarity and as an improvement to the previous blackhole architecture blog—which went over the chip from a more abstract, tt-metal pov—this blog will cover lowering an RMSNorm kernel onto a Tenstorrent blackhole chip, first on the SFPU and then moving part of the computation to the FPU. I’ve learned a lot more about the chip since the last article, and there are a lot of constraints, rules, and oddities that we need to be aware of when deciding how to lower kernels. Here is a tldraw diagram of the entire chip that you can look at while reading this.

By lowering kernels I mean splitting an operation into independent operations, deciding which parts can be executed simultaneously, and whether a part should be run using SFPU or FPU operations. A few reasons to choose each one:

• the SFPU can run precise float32 arithmetic.
• the FPU uses tf32 / bfloat16 inputs for the operations discussed here, with float32 accumulation
• FPU re-use is limited; getting accumulations back into srcA or srcB costs an extra instruction.
• the FPU can process more elements per instruction, but the number of fidelity phases affects how much faster it is

The approach I’ve come up with starts with picking the most precision-maximizing lowering, then moving as many operations to the FPU as possible before the relative-error becomes too high for accurate inference. Usually I try to match what a GPU would produce by default. We’ll start with all the math on the SFPU, then move the weight multiplication to the FPU. Take RMSNorm:

class RMSNorm:
  def __init__(self, dim:int=2048, eps=1e-5):
    self.eps = eps
    # learned parameters, initialized to ones here
    self.weight = Tensor.ones(dim)

  def __call__(self, x:Tensor) -> Tensor:
    # for simplicity: x.shape=(2048,)
    dtype = x.dtype # store original tensor dtype
    xf = x.float()
    x = (
      xf * (
        xf.square().mean(axis=-1, keepdim=True) + self.eps
      ).rsqrt()
    ).cast(dtype)
    return x * self.weight

which explicitly makes most of the rmsnorm compute in float32. Omitting xf=x.float() can change the accuracy of the reduction. This is the operation we’ll lower to a blackhole kernel.

Formally written:

Our inputs and output are bf16 with float32 intermediates in dst and the SFPU. The unpacker/packer handle conversions; we skip the explicit bf16 cast before the weight multiply, so results won’t be bit-accurate to the tinygrad code.

You can split RMSNorm into a few operations (the weighted input and inverse RMS can be computed independently, which will matter when we overlap FPU and SFPU operations):

note: I’m leaving out a lot of setup and scheduling here: SFPU config writes, SFPNOPs to wait for results, semaphores between threads, source-bank switching, and unpacker/packer configuration (including output rounding). The real kernel also uses SFPLOADMACRO to issue configured load/compute/store sequences with fewer instructions from trisc1. These are needed to run this efficiently on hardware, but for now we’re just following the data and the math. You can see the details in rmsnorm_hybrid.py.

getting data in and out of the SFPU

There are two sfpu load/store instructions, SFPLOAD and SFPSTORE. Both load and store only to dst, the big 1024x16 shared sfpu/fpu register. The way it loads and stores is exceptionally weird, take a look:

There is a set of 8 bits in the instruction that allows you to replace an even index load/store with an odd index (and vice versa), but it applies to the loads in every row. So instead of loading indices 0,2,4,6,8,10,12,14 from each row, you can load 0,3,4,7,8,10,13,14. If you set all 8 bits to 1, your even load becomes an odd load.

Keep in mind that the SFPU has lane-predication support; before an instruction, you can disable or enable lanes (0-31) within each LReg. That way, you can only load the first row of dst instead of all 4, for example.

getting data into dst

Our inputs start in DRAM the 8 GDDR6 banks. For RMSNorm, that’s two tensors of 2048 bf16s bfloat16: the weights γ and the activation x. They take two hops to reach dst.

1. DRAM -> L1: brisc issues the NoC reads that pull both tensors into the core’s L1.
2. L1 -> dst: trisc0 drives the unpacker by convention. this is not enforced by hardware; any trisc has full access to the entire coprocessor.

Once x and γ are in dst, we can start issuing SFPU instructions on trisc1.

where things go in dst

As shown in the tldraw diagram, dst is normally 1024x16, but we always keep float32 elements in it, which halves the size to 512x16. You can think of that as 64 allocations of 128 elements (8x16) each. ELWMUL accumulates into one contiguous 8x16 block, so it’s the natural unit for planning layouts. I’ll write ranges as dst[rows, columns], with the end of each range excluded.

For this kernel, we can fit x in dst[0:128, 0:16] (allocations 0-15) and γ in dst[128:256, 0:16] (16-31) leaving the bottom half free for intermediates. When the inputs and intermediates no longer fit in dst, we have to start streaming data in and out and keeping partial sums between chunks. The 2048-element example fits, so we don’t need to handle that here.

sfpu registers

There are 8 general LRegs, LReg0 through LReg7, each holding 32 lanes of 32 bits, and a few registers for constants and configuration. I’ll use these names throughout to avoid confusing LReg1 with the core’s L1 SRAM.

• LReg9 is read-only and contains zero in every lane.
• LReg10 is read-only and contains float32 1.0 in every lane.
• LReg11 through LReg14 can hold constants we want to reuse. To write one, we use SFPCONFIG, which takes 8 lanes from LReg0 and repeats them across all 32 lanes of the chosen register. We won’t use these yet.

inverse root mean squared (RMS) on sfpu

square and average

We can save some instructions and compute time here by squaring and summing at the same time. Some SFPU instructions take separate source and output registers, and others modify a register in-place. The SFPMAD multiply-add instruction takes three source registers and an output register:

TTSFPMAD(LReg0, LReg1, LReg2, LReg3, instr_mod1=0)

This computes LReg0 * LReg1 + LReg2, storing the result in LReg3.

Initialize all 32 lanes of LReg0 to FP32 zero, then run a loop:

1. Load 32 elements from dst into LReg1.
2. Issue TTSFPMAD(LReg1, LReg1, LReg0, LReg0). This is x*x + prev_sum = new_sum

The weird load/store addressing of the SFPU doesn’t matter here, since we just need to sum all the elements. We just have to make sure the loads go through the entire 2048-element allocation of x in dst. This takes $\frac{2048}{32}=64$ loads. Each load pulls 8 elements from each of 4 dst rows, so one window — 4 rows of 16, or 64 elements — takes two loads: one for the even columns and one for the odd. 2048 elements is 32 windows.

After this sequence, you end up with 32 partial sums in LReg0. The LRegs are summed elementwise; imagine stacking two LRegs on top of each other and adding down, that is what ends up in LReg0. Now we have to sum all 32 lanes in LReg0.

Up until now, all the SFPU instructions we’ve run have interpreted an LReg as a flat list of 32 floats. To do a sum of elements within a register, we need to look at the register like a grid of 4 rows and 8 columns.

First, we reduce every row of 8 into its own sum. Since the adds work between LRegs, we have to copy the current sum to a 2nd LReg, rotate it by 4, and add it back. Then repeat with rotations of 2 and 1, copying the updated sum each time. SFPSHFT2 with modifier 3 only rotates by one lane, so the larger rotations are just repeated single-lane rotates.

After you have the row sums in LReg0, copy that to 3 other LRegs. Then you run SFPTRANSP (no arguments). The 4 0th rows (LReg0’s 0th row, LReg1’s 0th row, etc) move to LReg0. Same for the 1st, 2nd, and third rows. Since all four copies are identical, this leaves each row’s sum filling a register of its own, so adding the registers together gives the total in every lane.

note: SFPTRANSP transposes both groups, LReg0-LReg3 and LReg4-LReg7, in one instruction. We could have used either group for the reduction, but anything in the other group gets shuffled too.

Now, you can just add across all 4 LRegs using SFPADD. SFPADD is SFPMAD with the first operand fixed to LReg10, which contains 1.0 in every lane. This makes our multiply-add just an ADD. SFPMUL does the same thing with the addend fixed to LReg9, which contains zero. In this case our instruction would look like:

TTSFPADD(LReg10, LReg0, LReg1, LReg0) # LReg0 <- LReg0 + LReg1

Repeated 3 times to add LReg1, LReg2, and LReg3 to LReg0.

At this point, our sum is repeated in every lane inside LReg0. Next is just SFPMUL by $\frac{1}{2048}$, then another SFPADD to add ε 1e-5.

To multiply by a constant, like in our SFPMUL example, we must load the immediate constant value into an LReg. This is done using SFPLOADI:

SFPLOADI(destination, mode, immediate_16_bits)

An arbitrary float32 constant requires two SFPLOADI instructions, because there is only enough room in the 32-bit instruction encoding to supply 16 bits at a time. Mode 10 takes bits & 0xffff to load the lower bits, and mode 8 takes bits >> 16 to load the upper bits. Some constants, including bf16 values converted to float32, fit in a single instruction.

The remaining computations are the rsqrt and the two elementwise multiplications.

rsqrt (fast inverse square root)

Currently, LReg0 in the SFPU holds

and we need to calculate $1/\sqrt{m}$. We approximate it using a bit shift and a subtraction, then apply two corrections. The constant and algorithm are explained in this paper; have chatgpt explain it to you if you’re curious, this is not really the point of the article.

The SFPU instructions for this are straightforward. First, copy m from LReg0 into LReg1, keeping LReg0 for the corrections:

SFPMOV(LReg0, LReg1)        # copy m
SFPSHFT(LReg1, -1)          # shift the entire 32-bit pattern right by one
SFPLOADI(LReg2, 0x5f1110a0) # the magic constant
SFPIADD(LReg2, LReg1)       # integer subtract: LReg1 = LReg2 - LReg1

Shifts take a signed amount, so a negative one shifts right. These are trimmed down: the real instructions carry extra modifier arguments, and the magic constant needs two SFPLOADIs to get all 32 bits in. LReg1 now holds our initial guess, y.

The first correction is a cheap polynomial correction:

t = -m * y * y
p = 2.2825186 + t * (2.2533049 + t)
y = y * p

Which we can evaluate using SFPMUL, SFPADD, and SFPMAD. Then we apply the newton correction:

error = 1 - m * y * y
y = y + (0.5 * y) * error

These are pseudocode for the arithmetic; we still need to load the constants and allocate scratch LRegs. We can keep the reused constants in LReg11-LReg14 using SFPCONFIG. Once both corrections are done, store the final rsqrt approximation, r, in LReg0. It is repeated in every lane.

elementwise multiplications

The last computation remaining is simply

SFPLOAD  → LReg1 = next 32 elements of x from dst[0:128, 0:16]
SFPLOAD  → LReg2 = corresponding 32 weights γ from dst[128:256, 0:16]
SFPMUL   → LReg1 = LReg1 × LReg0 # LReg0 contains inverse rms
SFPMUL   → LReg1 = LReg1 × LReg2
SFPSTORE → write LReg1 back to dst[256:384, 0:16]

Repeated until we cover all 2048 elements. This version computes (x * r) * γ; the hybrid version will compute (x * γ) * r, using the independent weighted-input branch from earlier. Reordering the multiplications can change float32 rounding.

For lane i, that computes:

packing the final output

Now we can run the packer using trisc2, copy from dst[256:384, 0:16] to L1, and then use ncrisc to copy that back to DRAM.

moving the weight multiplication to the FPU

x * γ doesn’t depend on the sum or the rsqrt, so we can move it to the FPU. This is what rmsnorm_hybrid.py does. The main change is where we put the inputs.

We now unpack 1024 elements of x into srcA and the matching 1024 weights into srcB srcA and srcB hold 1024 elements each, so our 2048-element input takes two passes.. We still need x in dst for the SFPU square sum, so we use MOVA2D to copy srcA into a scratch region. Our new float32 dst layout is:

Each of those regions is 64 rows of 16, or 8 of the 8x16 allocations. After copying x into scratch, we point ELWMUL back at the start of srcA and srcB and issue:

TTELWMUL(
  0,  # bank flips: bit 0 = srcA, bit 1 = srcB; neither here
  0,  # reserved, must be zero
  0,  # srcB broadcast: bit 0 = column 0, bit 1 = one row; neither here
  0,  # address mode: advance both source rows by 8
  64, # destination starting row
)

This reads rows 0–7 from srcA and srcB and accumulates their elementwise products into dst rows 64–71, i.e. dst[64:72, 0:16]. We zero the accumulators beforehand because ELWMUL adds to dst. The ELWMUL ISA docs describe the operands and fidelity phases.

Each ELWMUL advances 8 rows further into srcA and srcB, so eight of them cover all 1024 elements. Then we repeat across four fidelity phases.

why four fidelity phases?

A bf16 value carries 8 significant bits, but the FPU only feeds 5 of srcA’s and 7 of srcB’s into the multiplier at a time. One pass gives you a rough product. Each later phase multiplies a different pair of bit groups and adds that contribution to dst, filling in the pieces the rough pass left out. Four phases cover all of them, which is what HiFi4 buys us: more of the input precision, at the cost of more instructions. see the fidelity phase tables.

That is 8 allocations × 4 phases = 32 ELWMUL instructions per half. We only zero dst before the first phase; the later phases add to it.

The SFPU reads x from scratch and accumulates its square sum as before. The partial sums stay in the LRegs across both halves; only the copy of x in scratch gets overwritten. Once both units are done with the first half, we reuse scratch for the second half of x and put its products at dst row 128 (dst[128:192, 0:16]).

After the reduction and rsqrt, the last loop is just SFPLOAD from dst[64:192, 0:16], multiply by r, and SFPSTORE back to the same locations. The packer reads both output regions into L1, ready to copy back to DRAM. We no longer need to load γ into an LReg or do that multiplication on the SFPU.

benchmarks

The full-SFPU kernel ran in approximately 1,471 cycles (1.09 µs), and the SFPU/FPU hybrid kernel took 1,306 cycles, about 11% faster!

With more aggressive scheduling and more careful register use, my fastest version of this kernel so far is 1,083 cycles, but getting there involves using several of the hardware features I deliberately skipped in this article.

update 9/17/2026: i’ve gotten it down to 996 cycles.

If you’re curious, the main change is using the SFPU’s macro system to run a set of repeated SFPU instructions instead of issuing them one-by-one. Issuing a tensix coprocessor instruction (SFPLOAD, ELWMUL, etc) takes around 4 cycles, and in most cases the compute will run faster than you can issue instructions, leading to delays.