In B200 gpu, if I use shared-memory for both A and B operands of a gemm operation (such as fp32 += bf16 * bf16), does tensor core load the data from smem into its own internal memory once before starting operations? Or does it continuously re-access all rows of A and all columns of B?
If its constantly accessing smem for A and B, it must be decreasing available bandwidth to these:
- TMA-load: must be badly affected, can even move the roofline to more difficult place and enforce a larger tile to be used
- Multicast: SM - to - SM connection lands at shared-memory so any broadcast within a cluster must also depend on the available shared-memory bandwidth
- TMA-store: writing epilogue is also another bandwidth consumption but only when its hidden behind gemm k-loop so maybe not that important
- Registers spilled to L1: because L1 and shared-memory are unified and share the same hardware
- Warp-shuffle/vote: another bottleneck point that can be affected during elect-leader function for selecting the thread that will launch TMA or UMMA
- Loading - storing metadata about a gemm operation on shared-memory or L1-cached global memory, even from another cuda-block running on same SM unit
- Waiting on a shared-memory barrier (maybe doesn’t depend on bandwidth too much if its not spin-wait implementation)
So, if it makes roofline model better, I want to experiment loading A-operand into tmem (tensor memory) to see if it reduces smem pressure during tensor core work and let tiles be loaded quicker, epilogues written faster, multicast scales better, etc but I couldn’t find an example that casually defines multiple A matrices in tensor memory to support N stages of a pipeline. It’s a bit complex to understand ptx to define a tmem access function for placement of A matrices in there.
Can I use just 1 warp per pipeline stage to store A from smem to tmem? I’m testing a pipeline with 7 stages of 128x128x64 tile that has only 1100 tflops performance of bf16. But cublas does 1670 tflops. 128 x 128 x 64 tile with smem-only A & B is leaving 570 teraflops on table. Maybe loading A into tmem makes tiles load faster, or tensor-core operation faster. (multicast already enabled)
If 1 warp per pipeline stage is not usable for tmem-based A, what is the suggested way of doing this? Full warpgroup based producer consumer? Or just double-buffered input? (using cutlass 4.5.2)
Which tensor core instructions do you use?
Some load from registers, some from tensor memory.
_SS suffix with 128-N option. B-tile is multicasted. 2 SM per cluster. 1 block per SM. Using all SM units. So its shared-memory for both A and B. I don’t know how to use tensor memory, and especially need multiple A slots for experiments.
7 warps independently doing this:
- load tiles
- wait
- compute
- wait
- repeat with no communication to other warps
this should be overlapping some of compute with some of tile-load with no inter-warp communication cost but maybe not optimal.
The version that I need to test:
- load tiles A,B
- wait
- write A to tensor-memory (is this a high-latency operation?)
- compute (faster?)
- wait
- repeat with no comm to other warps
B200 has nearly 2x gemm performance of H100 but smem performance per SM is nearly same. I assume at least there is a tensor-core row-cache or column-cache that keeps loaded smem col/row inside the tensor core but streams the other data from other tile (so smem vs tmem matters). Tutorials are suggesting maximizing N dimension of tile, so I think B-tile column is cached inside tensor-core and the A-tile data is streamed from smem unless _TS suffix is used. But I don’t know how to do this with Cutlass 4.5.2 for multiple A (1 per stage of pipeline).
Within this paper
the Blackwell tensor core bandwidth has been analyzed.
Especially section 3.1.1 shared memory traffic hints that the operands are read once per MMA instruction.
If operands are read only once after a fixed time (in cycles), can we use the same shared-memory tile outside of tensor core for other tasks such as getting another tile asynchronously without waiting for mma to complete? Is shared-memory frequency is coupled with tensor core frequency?
I don’t quite understand.
Each SM can do 128B/cycle of shared memory, which is shared between all 4 SM partitions.
Between several TC instructions you can do other shared memory transactions, if bandwidth allows.
You probably cannot reuse the operands read from shared memory in another way, but have to read them again.
You probably can write something else into shared memory as soon as the operands have been read.
The frequencies are coupled. There is a certain frequency of the shaders vs. the global memory. And two SMs are even more in lockstep with each other.
I meant to re-use same smem area to load next tile, right after tensor core completes reading that tile. Maybe with some binary-search + nanosleep, I can find how many cycles are required for a tile size, to safely “release” the tile for next TMA-load. But this is more of a hack than just testing A-from-TMEM. I couldn’t find such an example kernel for TMEM-A-tile.
Are you limited in shared memory size?
I’m using around 192kB - 220 kB to keep pipeline depth maximum. Even with this smem usage, 128 x128 tile size doesn’t reach 60% of peak tensor core performance and I was hoping that TMEM-based A-tile could change the roofline to get more performance out of same tile size.
tcgen05 is an asynchronous operation. So instead of counting cycles, you probably need an mbarrier, as in this example at the end of commit:
The PTX docs are explicit about TensorCore collector buffers, you can give reuse hints (e.g. .collector::a::use) for tiles that get used several times.
A sequence of MMA instructions may reuse the same A matrix with a sequence of B matrices or may reuse the same B matrix with a sequence of A matrices. In these patterns the TensorCore may be able to laod the unchanged matrix once and reuse it through the sequence without multiple reloads. The A or B matrices are loaded into a TensorCore collector buffer (i.e., special cache).
An MMA instruction has an optional collector qualifier to specify when an A or B matrix is new to the sequence and should be loaded, unchanged within the sequence and should be reused, or the last use in the sequence and should be discarded. The collector qualifier is used to give the TensorCore permission to reuse a previously loaded A or B matrix; however reuse is opportunistic in that the TensorCore may reload a matrix even when it has permission to reuse that matrix. Thus, the source memory of an A or B matrix must not be modified while the MMA instruction using those matrices has not completed - regardless of collector qualifier permissions.
For the rest:
- The SMEM pipeline shared between the four SM partitions can do either a load or a store per cycle, but not both. However, that doesn’t necessarily mean that TMA writes to SMEM, or TC reads from SMEM, have the same limitation. What prevents simultaneous loads and stores for ordinary SMEM accesses (e.g. LDS, STS, LDSM) might well be the stage mapping warp-lane addresses to wavefronts. TMA or TC operations accessing SMEM don’t have per-lane addresses.
- While B200 has twice the TC performance of H100 for the same dtypes, it also reduces SMEM traffic by doubling (
cta_group::1) or quadrupling (cta_group::2) the maximum M dimension, as well as broadcasting the halves of B for cta_group::2.
256x256xK shape executes a 128x256xK gemm on each SM, but each SM only needs to read a 128xK A-tile and a Kx128 B-tile from SMEM.
elect.sync is a fixed latency instruction that doesn’t use the SMEM path.- My mental model of the Sync Unit (SYNCS - the instruction used for most mbarrier operations) is that it has a cache for mbarriers, and all external sync signals go through the Sync Unit instead of accessing SMEM directly. I might be wrong on this, but it would mean that if you only use a few mbarriers per SM they will hardly ever touch SMEM.
IIRC the documentation stated that the TC smem acceses are routed through the async engine, which is per lane. Therefore it could have similar limitations.
collector qualifier looks good when there are 2 fragments to compute and the second fragment will not require to load the common tile of the two gemm calls again. But couldn’t find its switch in Cutlass.