This only makes sense as the sequel to the occupancy note. The short version: occupancy on this kernel was architecturally closed at one CTA per SM, the kernel was latency-bound, and with no second CTA to switch to, the mainloop's software pipeline was the only thing standing between the tensor cores and a memory stall.
Stage depth is how far ahead the mainloop prefetches. More stages means more tiles in flight and more latency hidden. Stages are bought with shared memory, one tile buffer each, and shared memory per SM is fixed. So the depth you can afford is whatever is left once the kernel's other shared-memory tenants have taken their share.
Which makes it a budgeting problem. The interesting move is not to spend the budget better. It is to make one of the tenants smaller.
The atom changes the fragment
A 2-SM MMA atom, SM100_MMA_F16BF16_2x1SM_SS paired with
SM100_TMA_2SM_LOAD_IM2COL, splits one logical MMA across a
cluster of two SMs. Each SM then holds half the A operand instead of all of
it. The per-CTA A fragment halves, and the shared memory that was holding
the other half is free for stage buffers.
At a tile of 256x128x64 and a cluster of 2x1x1:
1-SM atom Stages 6 231,424 B/CTA
2-SM atom Stages 8 231,424 B/CTA identical budget
StageCount<9> fails the capacity assert
conv launch 450.86 us -> 375.33 us (-16.8%)Two extra stages at the same shared-memory footprint. Nothing was traded. The tile is the same, occupancy is the same one CTA per SM, and the arithmetic the kernel performs is unchanged. The operand simply stopped being duplicated across a pair of SMs that were going to cooperate anyway.
Why a search would not have found it
This is the part I want to underline, because it is an argument about what tuning can and cannot do for you.
The usual way to pick stage depth is StageCountAutoCarveout,
which measures what is left after the fixed allocations and fits as many
stages as will go. It is a good default and it is exactly wrong here, because
it treats the operand fragment as fixed and carves around it. The win comes
from changing the thing it carves around.
A sweep over stage counts does not find it either.
StageCount<8> simply fails to build on the 1-SM atom, so a
search along that axis sees a wall where there is a door one axis over. You
have to already know what a 2-SM atom does to fragment size to know the wall
can be moved, and that is knowledge about the hardware rather than about the
search space.
Where this generalises
- When a resource binds, look at who is holding it, not at how it is divided. The shared-memory budget looked fully allocated. One of the allocations turned out to be avoidable.
- Cluster-level features change per-CTA arithmetic. Anything that lets two SMs cooperate on one operand is a candidate for halving something you had been treating as fixed.
- A compile failure is a boundary, not a verdict.
StageCount<9>failing tells you where the budget actually runs out, which is a measurement. Read the assert instead of backing off to the last value that built.
Auto-tuning finds the best answer inside the constraints you hand it. Knowing which constraint is negotiable is not something it can do for you.
Sources
- CUTLASS, efficient GEMM: pipelining, stage counts and where the shared-memory budget goes.
- CuTe MMA atoms: what an atom is and how a 2-SM one distributes its operands.
- Blackwell tuning guide: clusters and the SM-pair features this depends on.
- Grid 24 is not headroom: why stage depth was the only remaining lever.