[RFC] CUTLASS CollectiveMMA for cublasDx #233
Labels
Comments
|
Hi, we're working on optimizations that should help with addressing issue you raised. Feel free to drop me an e-mail on [email protected] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hello!
What was the motivation for using cooperative_gemm for the internals of cublasDx? I ask, because due to the points below, CollectiveMMA seems to be a much more compelling option for general-pupose, large GEMMs, which I believe to be the goal of cublasDx.
TL;DR
✅ = winner for that category, doubly is a tie
🥹 = loser for that category
c_gemm = cooperative_gemm
Comparison
All experiments run on an A100 80GB GPU.
Performance At Scale
M = 4096, N = 4096, K = 64 (small K due to cublasDX limits), precision = tf32, tf32, float, bM = 128, bN = 64, bK = 64.
Diffs
Here is my diff for the CUTLASS code, in case you want to replicate these results.
diff
// inside gemm_tn auto bM = Int<128>{}; - auto bN = Int<128>{}; + auto bN = Int<64>{}; auto bK = Int< 8>{}; auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) - auto bP = Int<3>{}; // Pipeline + auto bP = Int<1>{}; // Pipeline // Define the smem layouts (static) auto sA_atom = make_layout(make_shape ( bM, bK), make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major [[maybe_unused]] auto sB_atom = make_layout(make_shape ( bN, bK), make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major - auto sA = tile_to_shape(sA_atom, make_shape(bM, bK, bP)); - auto sB = tile_to_shape(sA_atom, make_shape(bN, bK, bP)); + auto sA = make_layout(make_shape(bM, bK, bP), LayoutRight{}); + auto sB = make_layout(make_shape(bN, bK, bP), LayoutRight{}); auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx // Define the thread layouts (static) TiledCopy copyA = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<TA>, TA>{}, - Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout<Shape<_16,_8>,Stride<_8,_1>>{}, // Thr layout 16x8 k-major Layout<Shape< _1,_1>>{}); // Val layout 1x1 TiledCopy copyB = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<TB>, TB>{}, - Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout<Shape<_16,_8>,Stride<_8,_1>>{}, // Thr layout 16x8 k-major Layout<Shape< _1,_1>>{}); // Val layout 1x1 - TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{}, - Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA + TiledMMA mmaC = TiledMMA< + MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>, + Layout<Shape<_2, _2, _1>, + Tile<_32, _32, _8> + > + >{};Below is the diff for cublasDx. I had to make changes to enable different types.
diff to support tf32, tf32, float
Results
Also here is the result summary of running multi_block.
cublasDx result summary
CUTLASS is below
CUTLASS result summary
Prologue
It is clear that the CollectiveMMA is an order of magnitude faster here (bK = 8). On the one hand, this is not necessarily an apples to apples comparison, given the invocation overhead of cublasDx and that the underlying cooperative_gemm does more work (albeit unrequested), namely a register to shared copy, so the question becomes do these alone account for the severe performance gap?
I do not think so, I hope to elaborate why in the following points. Note that CollectiveMMA has a more richer space of architecture-specfic parameters like gmem tiled copy layouts, pipelining, via register double-buffering (Volta) or shared memory (Ampere and above), and much more sophisticated configurations for Hopper.
Ease-of-use
Given the vast parameter space for configuring ColletiveMMA templates, it is definitely much less trivial to get it right compared to the simpler API of the cooperative_gemm. For example, CollectiveMMA demands that you specify the CopyAtom and GmemTiledCopy while cooperative_gemm doesn't.
However, in my opinion, the above constraint applies to the lay user not the experienced engineers developing cublasDx. This is the case, because they have already tackled this complexity through their
suggest_*APIs, which translate a high-level GEMM description to optimal, architecture-specific template parameters.Thus, neither implementation seems to be more complex for the expert CUDA developer.
Shared Memory Demand
This is where CollectiveMMA really shines.
For example, let's use these block dimensions bM=128, bN=128, bK=8 and floating-point precision and blockDim of 128 threads.
With these parameters, the CollectiveMMA with no shared memory pipelining only demands 8K =$(128\cdot 8 \cdot 4 \cdot 2)$ shared memory per block to store A and B matrices. Alternatively, the memory amount scales with the pipeline stages, so at 2 stages (minimum amount for Ampere), shared memory would go up to 16K.
On the other hand, cublasDx, in addition stores the C matrix in shared memory; thus, requires 8K + ($128\cdot 128 \cdot 4$ ) = 72 K, a 9X increase.
This may not seem like a lot, but the implications for block scheduling and occupancy are worth considering.
Continuing with this example, let's use the Volta architecture, where maximum shared memory per SM is 96K. In this architecture, only one block would be scheduled per SM, when I run a GEMM using cublasDx.
On the other hand, with the CollectiveMMA, you can fit 5 blocks per Volta SM. Note, you cannot fit 96/8 = 12 blocks because of register pressure, which I explain next.
But the point is clear, more active blocks per SM translates to better performance.
I believe this is the major cause for the poor performance of cublasDx vs collectiveMMA, as both APIs have largely identical underlying code, invoking cute::gemm underneath.
K Limitations
More importantly, this limitation constrains cublasDx currently to only matrix sizes, where k <= 196.
Any larger k requires the user to write their own k-reduction logic, defeating the purpose of using a library as cublasDx, because most likely than not, writing high-performance GEMM CUDA code is a futile exercise for non-experts.
Registers
Both collectiveMMA and cooperative_gemm require the same amount of registers, specifically, the registers used as A, B, and C operands for the same underlying gemm function, namely using CuTe's notation, tCrA, tCrB, and tCrC. For the CollectiveMMA done in the performance comparison above, register demand (size(tCrA) + size(tCrB) + size(tCrC)) is about 96 per thread (about 128 for no spill loads), hence why we can only fit 5 blocks per Volta SM.
Key
The key difference is that after the GEMM, collectiveMMA leaves matrix C in register memory, this can allow for more efficient epilogue operations as the roundtrip to and from shared memory is bypassed. On the other hand, cooperative_gemm, does only one axpby and then stores the result in shared memory forcing cublasDx to only support only one element-wise epilogue on register memory, while subsequent epilogues must be done on shared memory by the user.
For further motivation, note from here that the shared memory for the C matrix is actually used only for the epilogue axpby and not the core gemm. So, we can completely eliminate that shared memory allocation and instead operate directly on registers, or if needed reuse the shared memory for A and B, since the gemm would be complete by the time the epilogue begins.
Next Steps
To reiterate the initial point, if cublasDx aims to be a library for executing high-performance, general-pupose, large GEMMs, then it seems that CollectiveMMA is the way to go; otherwise, then the existing setup is sufficient. Also, we must keep in mind that cublasDx is still in early access mode, so there may already be ongoing efforts revisiting this design choice.
Finally, I am aware that cublasDx is not necessarily open-source, but as a CUDA geek :), I am more than happy to contribute to making this happen, if the cublasDx team decides in that regard.
The text was updated successfully, but these errors were encountered: