Why split matrix A/B between CTAs for 2SM MMA instruction? Code and documentation seem contradictory?

[XieXiating]

Context: I’m trying to understand how matrices A and B are stored in shared memory for the 2SM MMA instruction, but the code and documentation seem contradictory:

Code

• In CollectiveBuilder, the SmemShape_N field suggests that matrix B is partitioned along the N dimension between two CTAs.
• The function collective::detail::sm100_cluster_shape_to_tma_atom_B() further implies that if the CTA cluster consists of just one CTA pair, then matrix B isn’t multicast.
• Taken together, these points imply that in a 2SM MMA computation, matrix B is stored only once across the two CTAs’ shared memory—each CTA holds only half of B (a sub-matrix of size (N/2, K)).
• Similarly, matrix A would also be split between the two CTAs, with each CTA storing a sub-matrix of size (M/2, K) of A.

Documentation

• However, PTX documentation does not indicate that the 2SM MMA computation (tcgen05.mma..cta_group::2) share matrix A/B between CTAs. In other words, each CTA can only use the data present in its own shared memory during the MMA operation.
• If that’s the case, I’m confused about how the full (M, N) output matrix gets produced. It seems like each CTA would only compute a partial result (for example, an (M/2, N/2) tile) using the half of A and B it has. Two such partial tiles (one from each CTA) wouldn’t cover the entire (M, N) output, so something doesn’t add up.

Question: Am I misunderstanding something about how the 2SM MMA instruction works? Any clarification on how the data is stored and used or how the full result is obtained would be greatly appreciated. Thanks!

[thakkarV]

The two CTAs are indeed cooperating to produce the full MxN tile.

However, PTX documentation does not indicate that the 2SM MMA computation (tcgen05.mma..cta_group::2) share matrix A/B between CTAs. In other words, each CTA can only use the data present in its own shared memory during the MMA operation.

Where is the documentation incorrect? Your link does not point to a specific line.

[XieXiating]

Thank you for your reply!

The documentation isn't incorrect, but it doesn't explicitly mention that when performing MMA at the CTA-pair granularity, the shared memory of both CTAs is accessed as distributed shared memory.

Initially, I assumed each CTA's shared memory was only accessible by its own MMA operation, leading me to believe that matrix B needed to be multicast to each CTA within the CTA pair for correct 2-SM MMA operation.
However, this assumption appears to be incorrect.

[SimonZh1234]

PTX documentation indicates that B matrix in the peer-CTA can be shared for tcgen05.mma:
"the B matrix has shape KxN, in Shared Memory of the current CTA and optionally in peer CTA".
See https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma