Question about L2 cache swizzling in FP8 GEMM

[luongthecong123]

Thank you for sharing your amazing solutions. Could you give more information on the logic of L2 cache swizzling used in your kernel ?
Is this similar to the "Scheduling and L2 cache" in Kernel 6 of this blog, where the author improved L2 cache hit:

https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog

Thanks,
Cong

[jpy794]

Thanks for your interest in this repo!

We've experimented with the thread block swizzling strategy mentioned in the blog you referenced. You can find the relevant code here:

RadeonFlow_Kernels/src/gemm/gemm_kernel.cpp

Lines 412 to 415 in 5983fe5

	// threadblock swizzle
	auto log_tile = 1;
	auto block_idx_x = blockIdx.x >> log_tile;
	auto block_idx_y = (blockIdx.y << log_tile) + ((blockIdx.x) & ((1 << (log_tile)) - 1));

RadeonFlow_Kernels/src/gemm/gemm_launcher.cpp

Line 69 in 5983fe5

dim3 grid(ceil_div(N, BN) << 1, ceil_div(M, BM) >> 1, SPLITK_FACTOR);

We've also tried the swizzling strategy designed for MI300's disaggregated L2 cache, as implemented in aiter

However, in our experiments, neither approach showed a significant performance difference. So in our final implementation, enabling or disabling swizzling made little practical impact.

[luongthecong123]

Thank you for your response.
I also found from this legendary blog https://siboehm.com/articles/22/CUDA-MMM that the thread block swizzling didn't benefit him very much. Quote:

It didn’t increase performance, presumably because L2 hit rate is already fairly high at 80%, so I ended up removing the swizzling code.

Also, I would like to ask 2 other questions:

1. In your solution, how many blocks run per CU, and what is their block size. I'm looking into low occupancy kernel to maximize VGPR usage. And I don't know how to explicitly force the compiler to know that I want 128 VGPRs per thread. I'm currently running 2 blocks of 1024 threads per CU (each block will output a 128x128 elements in C) with 100% occupancy, which limit me to only 64 VGPRs per thread. I'm looking into running 2 blocks of 512 threads per CU = 50% occupancy, so that each thread can have 128 VGPRs, then each block can spit out 256x128 elements in C.

2. Can you explain a bit about your second most important optimization: "Batch GDS-to-LDS memory loads." ?

Thanks.

[jpy794]

In your solution, how many blocks run per CU, and what is their block size.

In our GEMM case, we increase occupancy by increasing the number of warps per block instead of blocks per CU. So I assume you're referring to warps per block here — we use 8 warps per block (arranged as 4x2), which collectively compute a 256x128 output tile of the GEMM.

I'm looking into low occupancy kernel to maximize VGPR usage. And I don't know how to explicitly force the compiler to know that I want 128 VGPRs per thread. I'm currently running 2 blocks of 1024 threads per CU (each block will output a 128x128 elements in C) with 100% occupancy, which limit me to only 64 VGPRs per thread. I'm looking into running 2 blocks of 512 threads per CU = 50% occupancy, so that each thread can have 128 VGPRs, then each block can spit out 256x128 elements in C.

It sounds like your goal is to write a kernel where each thread uses 128 VGPRs, and launch 2 blocks of 512 threads per CU. This is easy to achieve as the compiler limits register usage to 128 VGPRs per thread by default. (If you're launching too many blocks, they simply won't be scheduled concurrently on the same CU. ) So you just need to decrease blockDim.

You can also use __launch_bounds__ to increase or decrease register limit as you want.

That said, in the gemm case, I'm not sure maximizing registers usage of a CU would help with performance.

Can you explain a bit about your second most important optimization: "Batch GDS-to-LDS memory loads." ?

This refers to vectorized global memory loads. MI300 supports up to 16-byte global load instructions (e.g., global_load_dwordx4). Using vectorized loads significantly improves global memory bandwidth utilization. You can see an example of this here:

RadeonFlow_Kernels/src/gemm/gemm_kernel.cpp

Line 458 in 5983fe5

reg_a[reg] = *(LoadVec *)&a[m + i][k + j];

[luongthecong123]

Thanks.

In our GEMM case, we increase occupancy by increasing the number of warps per block instead of blocks per CU. So I assume you're referring to warps per block here — we use 8 warps per block (arranged as 4x2), which collectively compute a 256x128 output tile of the GEMM.

So each block of 512 threads (8 warps) will calculate 256 * 128 elements, and you're using 16x16x32 MFMA. 256 * 128 / (16 * 16) = 128 16x16 tiles to store the output, each block has 8 warps, so each warp has to hold 128 / 8 = 16 frag_c. That's equivalent to 16 * 4 VGPRs = 64 VGPRs, which already used up all the available VGPRs per thread for 100% occupancy. How do you have extra VGPRs for other operations in the kernel ?

[jpy794]

You can check out AMD's doc — each CU on MI300 has a 512KB register file. So in our case, each thread can use up to 256 VGPRs. We use __launch_bounds__ to bump up the default limit.