Streaming process

How the SM works inside a transformer

Don't get panic and run away, bro. Just wait, I will explain you everything.

SM is just one self-contained mini-processor on the GPU, that's it. And the H100 contains 132 of this kind.

Part 1

The flow, step by step

Tap any underlined word to see what it means.

Step 1

Everything starts with HBM

Everything actually starts with HBM, the KV cache or all things are just stored in HBM.

Step 2

Kernel launch, then the warp scheduler

Then the second process starts with the kernel launch: the block scheduler gives thread blocks to the SMs, and inside each SM the warp scheduler, every clock tick (one tiny step of the chip, about half a nanosecond), picks one ready warp to run its next instruction. It keeps doing this through all the steps below. It does not decide where to transfer or what to transfer; the kernel code decides that.

Step 3

TMA loads the Q tile, then K and V tiles

And then the TMA loads the Q tile (a small sub-matrix, one block of the full matrix) once, and then K and V tiles one after another, basically query tile and key/value tiles. I would say these are like matrices, so a producer warp gives the TMA a tile description, basically the Q and K tiles that are streamed from HBM through L2 cache, and then TMA transfers to shared memory with no threads spending time on it.

Step 4

Q·Kᵀ goes inside the tensor core

Then the third thing is happening: only the matrix multiply Q·Kᵀ goes inside the tensor core (softmax does not), and the consumer warpgroup uses WGMMA and tensor cores like Q, K, and from the shared memory it just picks to the tensor core, and attention is all you need. The result S goes into the register as f32 accumulator, basically, and we know that registers are some of the fastest things to store anything.

Step 5

Softmax on CUDA cores and SFU

Then there is a softmax calculation from the registers to CUDA cores: CUDA core finds each row max and subtracts, and the SFU computes exponentially running the max by full score matrices. When the max changes, the CUDA cores also rescale the old O and the running sum. You know how that softmax does that stuff, converting into the probability, the attention weights: how much this token should look at each earlier token (not the next-token probabilities; those come from a separate softmax at the end of the model).

Step 6

Meanwhile, the next K/V tile is already loading

Meanwhile, at the same time, the next K/V tile is already loading, so TMA already fetches the new K and V tiles that are again inside the HBM, and this is how the tensor core keeps busy. I'm kind of talking in the sense of flash attention.

Step 7

P dot V, accumulate the O register

And then again it loads shared memory and then again it does tensor core again do like P dot V basically the probability that comes from softmax into the value vector and the value vector comes from obviously you know shared memory because one tile of the KV cache was copied there (the KV cache itself is stored in HBM), and P basically comes from the registers where we have kind of stored it after softmax basically, and multiply them and accumulate the O register, and then again the step repeats and repeats and repeats and we write the result back, so the forward pass is done.

Step 8

Write the result back

So now CUDA cores divide by softmax sum and cast to bf16, again it goes to registers and then register to shared memory and shared memory to TMA and TMA to L2 and L2 to HBM,

Step 9

Then the MLP, same pattern

and then the same MLP same pattern goes again again: so TMA loads, tensor core multiply from the shared memory they put, and then after multiplication then it goes to register file and then register file goes to CUDA cores, everything SiLU / SwiGLU or whatever, residual, and then after each kernel (not only at the end of each transformer layer) it goes back and is put inside the HBM, unless the kernels are fused.

Remember: KV cache is stored in HBM, not here. Shared memory is only for a small KV cache tile, or I can say tile as small matrices.

Code snippet

# one thread block of the attention kernel, on one SM (simplified)
Q = tma_load(Q_tile)                 # HBM → L2 → TMA → shared memory, once
m, l, O = -inf, 0, 0                  # running max, running sum, O: in registers
for j in range(num_kv_tiles):
    tma_load(K[j+1], V[j+1])          # next K/V tile is already loading
    S = wgmma(Q, K[j].T)             # tensor core → register (f32)
    m_new = max(m, rowmax(S))         # CUDA cores
    P = exp(S - m_new)                # SFU
    scale = exp(m - m_new)            # rescale old O and running sum
    l = l * scale + rowsum(P)
    O = O * scale + wgmma(P, V[j])   # tensor core: P dot V
    m = m_new
O = (O / l).to(bf16)                  # CUDA cores: divide, cast to bf16
tma_store(O)                         # register → shared memory → TMA → L2 → HBM

Part 2 uses the text exactly as written. The diagrams, step titles and code snippet show the same flow visually.