We need to write a custom CUDA kernel for batched matrix multiplication. Let's analyze the problem: - Input A: [batch_size=128, m=512, k=1024] - Input B: [batch_size=128, k=1024, n=2048] - Output C: [batch_size=128, m=512, n=2048] This is a large batched GEMM. Writing a from-scratch CUDA kernel that beats cuBLAS (which torch.bmm uses) is extremely difficult. However, we can try using CUTLASS or writing a custom kernel that leverages Tensor Cores. But wait - the reference uses torch.bmm which calls into cuBLAS. On H100, cuBLAS is heavily optimized and uses Tensor Cores. We need to beat it with a custom kernel. One approach: The standard torch.bmm might have some overhead. We could try using CUTLASS or a custom kernel that fuses operations or uses better tiling for these specific shapes. However, writing a competitive GEMM kernel from scratch is very hard. Let me think about alternative approaches: 1. We could try using the CUTLASS library through PyTorch's CUDA extensions. But we need to write the kernel code inline. 2. We could write a simple kernel that calls cublasGemmBatched or similar - but that won't be faster than torch.bmm since torch.bmm already does that. Wait, let me reconsider. For these specific dimensions (batch=128, m=512, k=1024, n=2048), the total FLOPs are: 2 * 128 * 512 * 1024 * 2048 = 274,877,906,944 ≈ 275 GFLOPs At 1979 TFLOPS for BF16 tensor cores, this would take ~0.14 ms in theory. Memory bandwidth is 3.35 TB/s. The total memory moved is: A: 128 * 512 * 1024 * 4 = 268 MB B: 128 * 1024 * 2048 * 4 = 1 GB C: 128 * 512 * 2048 * 4 = 512 MB Total: ~1.75 GB At 3.35 TB/s, memory-bound time would be ~0.52 ms. So this is somewhat compute-bound. Actually, we need to check if torch.bmm in FP32 uses tensor cores. On H100, TF32 tensor cores give 495 TFLOPS (without sparsity). That would still give theoretical time of ~0.56 ms. So FP32 with tensor cores might be around memory bound. Let me think about what we can do. On H100, writing a custom kernel that matches cuBLAS performance is hard. But maybe there's room for optimization: 1. We can use CUTLASS or custom WMMA code 2. We could try to use shared memory more efficiently 3. We could try to use block-level tiling with warp-level matrix multiply Actually, let me look at whether we can use the CUTLASS library. PyTorch's load_inline can compile CUDA code, but linking against CUTLASS might be tricky if it's not in the standard include path. However, CUTLASS headers are sometimes available in PyTorch installations. Alternatively, we could write a warp-level matrix multiply kernel using `mma.sync.aligned.m64n8k16` or similar PTX instructions, or use the WMMA CUDA runtime APIs. Wait, H100 has newer Tensor Core instructions (wgmma). We could use PTX wgmma instructions for Hopper. But this is extremely complex. Let me think about a simpler approach: can we beat torch.bmm by using a kernel that processes multiple batch elements together or uses better occupancy? For these dimensions: - m=512, n=2048, k=1024 - batch_size=128 One idea: instead of calling 128 separate GEMMs (which torch.bmm might do if not fused), we could treat this as one large GEMM. But torch.bmm already uses batched GEMM in cuBLAS, which is very efficient. Another idea: Use CUTLASS's grouped GEMM or custom batching. Let me check if we can just write a very simple tiled GEMM kernel in shared memory. For FP32, the FLOPS are 67 TFLOPS (non-tensor core). The tensor cores for TF32 are ~495 TFLOPS. If torch.bmm uses tensor cores, a simple shared-memory GEMM without tensor cores won't be competitive. But wait - torch.bmm for fp32 should use TF32 tensor cores on Ampere/Hopper if enabled (which it is by default). Actually, let me reconsider: torch.bmm is very well optimized. However, maybe for these specific shapes, there's a custom approach. Let me look at the shapes again: - batch=128 - m=512 - k=1024 - n=2048 Total M = 128 * 512 = 65536 K = 1024 N = 2048 This could be viewed as a single GEMM with M=65536, K=1024, N=2048. But the batches are independent in terms of addressing. We could do a single strided_bmm. Actually, let me check: the A matrix for each batch is contiguous? Yes, A is [128, 512, 1024], so each batch matrix is contiguous, and batch stride is 512*1024. What if we just call a custom CUTLASS kernel for batched GEMM? That should match cuBLAS performance but not beat it. Hmm, what if the reference torch.bmm actually has some overhead that we can avoid? For example, if torch.bmm does extra checking or doesn't use persistent kernels? On H100, there are new features like warp-group matrix multiply (wgmma) that allow for much higher performance. If cuBLAS is not fully exploiting H100 features for these specific shapes, we might be able to write a better kernel. But realistically, writing a kernel from scratch that beats cuBLAS is extremely difficult. Let me think if there's a trick. Wait - could we use the fact that A and B are generated with torch.rand and we can assume they're contiguous? We could write a kernel that uses `ldmatrix` and `mma` PTX instructions for tensor cores. Actually, let me reconsider the approach. For the competition, we often see that writing a CUTLASS kernel or using custom tiling can yield good results. Let me try writing a simple but efficient tiled GEMM kernel using shared memory and see how close we can get. For a tiled GEMM in FP32: - Tile size: maybe 128x128 per block - Each block uses shared memory for A and B tiles - No tensor cores, just FMAs Compute required: 275 GFLOPs FP32 non-tensor peak: 67 TFLOPS Minimum time: ~4.1 ms But with memory bandwidth, we might be slower. Actually, the H100 has 67 TFLOPS for FP32, but that's non-tensor core. Tensor cores for TF32 are ~495 TFLOPS (without sparsity). If torch.bmm uses tensor cores, it should be much faster than any non-tensor-core kernel. So we MUST use tensor cores. One promising approach: use the CUTLASS library. PyTorch includes CUTLASS headers in its CUDA installation sometimes. Let me check if we can include them. Actually, in many of these tasks, the standard PyTorch CUDA extension doesn't have CUTLASS headers. But we could bundle simple CUTLASS-like code or use the `mma` PTX instructions. For H100, the equivalent of mma.sync.aligned.m64n8k16.row.col.f32.tf32.tf32.f32 instructions can be used. But more appropriately for H100, we could use the CUDA C++ wmma API for TF32. Wait, wmma on Hopper supports TF32 with m16n16k8. Let me check: Ampere/Hopper wmma: - m16n16k16 for fp16 - m16n16k8 for tf32 For H100, we could use wmma with TF32 to write a custom kernel. But is this faster than cuBLAS? Probably not for arbitrary sizes unless we optimize the tile sizes perfectly. Let me try a different approach. What if we use Triton? We can't use Triton in load_inline, but we could write a CUDA extension that calls Triton-compiled kernels. No, that's too complex. Actually, let me reconsider whether torch.bmm is truly optimal for these shapes. On H100 with batch_size=128, m=512, k=1024, n=2048, cuBLAS might use a batch algorithm that has some overhead. What if we launch a single large GEMM using cublasGemmStridedBatchedEx? torch.bmm already does something similar. Hmm, let me think about this differently. Maybe the task expects us to write a custom kernel using shared memory and 2D tiling that's competitive. For example: - Each block computes a tile of C of size BM x BN - Each thread computes multiple output elements - Use shared memory for A and B tiles For FP32 on H100, if we use the full register file and shared memory, we might achieve decent performance. Actually wait - I just realized that torch.bmm on fp32 might NOT be using tensor cores by default. Let me think about this more carefully. In PyTorch, matmul/bmm operations use cuBLAS. cuBLAS uses Tensor Cores for fp16, bf16, and tf32. For pure fp32, cuBLAS typically uses tensor cores for TF32 when TF32 is enabled (which is the default in PyTorch >= 1.7). But the actual numerical output of TF32 matmul vs pure FP32 matmul should still match within the fp32 tolerance (since TF32 rounds to 10-bit mantissa, the difference is small but still within fp32 range). Wait, is that right? Actually TF32 preserves the full fp32 output, but with reduced precision in the multiply-add, so it should be numerically close to fp32. But will the correctness checker accept it? The tolerance is atol=rtol=1e-4 for fp32, which is quite loose, so TF32 should pass. But we might be able to get even closer to pure FP32 while still using tensor cores. Actually, there's another consideration: for these problem sizes, maybe a custom kernel with register-tiled approach can outperform cuBLAS batched GEMM because we can better schedule the work across SMs for the specific dimensions. Let me try a more practical approach. I'll write a kernel that: 1. Uses a 2D grid of thread blocks 2. Each thread block computes a tile of C 3. Uses shared memory for A and B tiles 4. Loads data coalesced from global memory 5. Each thread computes multiple outputs using registers to maximize FMA throughput For H100, we have 132 SMs (for H100 SXM). Each SM can run up to 2048 threads. With good occupancy and enough warps to hide latency, we might get close to peak. Let me try a tile size approach: - Block size: 16x16 threads = 256 threads - Each thread computes an 8x8 tile of C = 64 elements - So each block computes 128x128 of C - Shared memory for A tile: 128 x (some K tile) - Shared memory for B tile: (some K tile) x 128 If K tile = 16, shared memory = (128*16 + 16*128) * 4 bytes = 16 KB. That's very small. If K tile = 32, shared memory = 32 KB. Still reasonable. If K tile = 64, shared memory = 64 KB. Good. Actually, let me use a different approach. Since m=512, n=2048, batch=128: - Total tiles in M dimension: 128 batches * 512 / 128 = 512 blocks in M (if tile M = 128) - Total tiles in N dimension: 2048 / 128 = 16 blocks in N - Total blocks: 512 * 16 = 8192 blocks That's a good number of blocks for 132 SMs. But without tensor cores, can we beat cuBLAS which uses tensor cores? It seems unlikely. Wait, I should reconsider. Maybe torch.bmm for fp32 does NOT use tensor cores because the tolerance check might fail? No, actually PyTorch uses TF32 by default and cuBLAS definitely uses it. So torch.bmm should be very fast. For H100: - Theoretical TF32 throughput: ~495 TFLOPS - Actual achievable with cuBLAS: maybe 300-400 TFLOPS - Our problem: 275 GFLOPs - Time with cuBLAS: ~0.7-0.9 ms Can we write a custom kernel that does better? On H100, we might be able to use the new wgmma (warp group matrix multiply accumulate) instructions which operate on larger tiles. But writing that in inline CUDA is very challenging. Let me think about whether we can call CUTLASS from the extension. If PyTorch is built with CUTLASS, we might be able to include its headers. But the include path for CUTLASS in standard PyTorch isn't guaranteed. Another idea: What if we use multiple warps per block and have each warp compute a sub-tile? For example, a block of 256 threads (8 warps) computing a 128x128 tile, where each warp computes 64x32 or similar. Actually, I recall that on Hopper, there's a concept of "warp groups" and TMA (Tensor Memory Accelerator). But utilizing these in inline PTX is very complex. Let me try an even simpler approach: Can we just use the `cutlass` PyTorch extension through `torch.utils.cpp_extension.load_inline`? Probably not since we need the headers. Wait, I just had another thought. What if the real optimization is not in the GEMM itself but in fusing operations or memory layout? The reference is just `torch.bmm(A, B)`, so there's nothing to fuse. Let me look at this from a different angle. What if we use `cublasGemmStridedBatchedEx` with a custom data type or algorithm? Actually, we could try calling cuBLAS directly with a different algorithm setting, but that's unlikely to beat PyTorch's default choice since PyTorch already picks good algorithms. Hmm, but what if we can use the fact that the shapes are fixed and compile a kernel specifically tuned for 512x1024x2048? That could indeed help. Specialized kernels often outperform generalized cuBLAS for specific sizes, especially in deep learning inference where shapes are known. For known shapes, we can: 1. Unroll loops completely 2. Use optimal tile sizes 3. Avoid boundary checks 4. Use vectorized loads/stores This is actually a common optimization strategy! So I'll write a specialized batched GEMM kernel for: - batch = 128 - m = 512 - k = 1024 - n = 2048 With these exact sizes, we can write: - Tile M = 64, Tile N = 64, Tile K = 16 (or larger) - Use 2D thread blocks Actually, for H100, let me consider using a larger tile. With 64KB shared memory, we could do: - Tile M = 128, Tile N = 128 - For A_smem: 128 x 32 = 4096 elements = 16 KB - For B_smem: 32 x 128 = 4096 elements = 16 KB - Total: 32 KB Then each block computes a 128x128 tile. With m=512, n=2048: - M tiles per batch: 4 - N tiles: 16 - Total tiles: 128 * 4 * 16 = 8192 tiles 8192 blocks is good. Each block has k/32 = 32 iterations over K. But without tensor cores, this would use: - FLOPs per block: 2 * 128 * 128 * 32 = 1,048,576 - Total blocks: 8192 - Total FLOPs: 8.6 GFLOPs... wait, that's wrong. Actually: 2 * 128 * 128 * 1024 = 33.5M FLOPs per block * 8192 blocks = 275 GFLOPs total. Correct. Each block does 128*128*1024 = 16,777,216 MACs. At 67 TFLOPS FP32 peak, that's about 0.25 ms for all blocks if perfectly parallelized. But realistically, without tensor cores we'd probably get 20-40 TFLOPS. For a tensor-core kernel, each block does the same work but in wmma fragments. OK, I think I need to try to write a kernel that uses tensor cores if possible. Let me write a WMMA-based kernel. For TF32 on Ampere/Hopper, wmma supports m16n16k8. Each warp computes tiles using these fragments. Here's a design: - Each warp computes a 64x64 tile using 4x4 wmma tiles of 16x16 - 4 warps per block compute a 128x64 tile or similar - Use shared memory for prefetching Actually, let me look at this more carefully. For wmma TF32 (m16n16k8): - Each mma instruction does 16*16*8*2 = 4096 FLOPs - 4 warps * 32 threads = 128 threads - If each warp handles a (4*16) x (4*16) = 64x64 tile, that's 16 mma tiles per warp - Per K=1024, there are 1024/8 = 128 k-steps - FMAs per warp per k-step: 16 mma instructions - Total mma per warp: 128 * 16 = 2048 - FLOPs per warp: 2048 * 4096 = 8,388,608 - Total warps: (128 batches) * (512/64) * (2048/64) * (warps per tile) Wait, this is getting complicated. Let me just write a simple, well-optimized non-tensor-core kernel first and see how it performs relative to the reference. If it's close or slower, I'll try tensor cores. Actually, let me reconsider the problem statement. For fp32, torch.bmm may actually not use tensor cores if the user sets `torch.backends.cudnn.allow_tf32 = False` or similar. But by default on H100, torch.bmm does use TF32 tensor cores because `torch.backends.cuda.matmul.allow_tf32` defaults to True. Wait, I just realized that our output is fp32. TF32 tensor cores compute in tf32 precision but output fp32. The numerical tolerance is 1e-4, which is generous enough for tf32. But here's a key insight: if we write a kernel that uses FP16 or BF16 tensor cores and then converts back to FP32, we might get much higher throughput. However, the numerical tolerance for fp32 is 1e-4, and fp16 might accumulate too much error unless we do careful accumulation in fp32. Actually, mixed precision GEMM with FP16 tensor cores and FP32 accumulation is standard and gives high accuracy. On H100, BF16 tensor cores with FP32 accumulation at 1979 TFLOPS would be incredibly fast. But wait, can we do that? The reference does pure fp32. If we use BF16 tensor cores, the multiplication inputs are rounded to BF16, but with FP32 accumulation, the result should still be within 1e-4 relative tolerance for inputs in [0, 1]. Let me think... For random inputs in [0, 1]: - A and B values are up to 1 - Each output element is sum of 1024 products - Each product in BF16 might have relative error up to ~1e-3 - Sum of 1024 terms: error accumulates but with random signs, might be around sqrt(1024) * 1e-3 ≈ 0.03 - Absolute values of outputs: mean ~512, max ~1024 - Relative error: 0.03/512 ≈ 6e-5 So yes, BF16 tensor cores with FP32 accumulation should pass the 1e-4 tolerance! And with 1979 TFLOPS theoretical throughput, our kernel could be extremely fast. The challenge is writing a kernel that achieves good utilization of the tensor cores. On Hopper, the story is even better. With wgmma, we can get much closer to peak. So my strategy: write a kernel that uses `mma.sync.aligned.m64n8k16` for BF16 or `wmma` subfragment instructions. Actually, on H100, the standard wmma API might not give the best performance. The new wgmma instructions are better. But wgmma is warp-group level (4 warps = 128 threads) and operates on descriptors in shared memory. For simplicity, let me try using the standard CUTLASS 3.x wgmma kernels if possible. But we don't have CUTLASS headers... Actually, there's a middle ground. I can write a kernel that uses the `mma` PTX instructions via inline PTX. For BF16 on Hopper, the mma instruction is `mma.sync.aligned.m64n8k16.row.col.f32.bf16.bf16.f32` which computes: - A: 64x16 BF16 - B: 16x8 BF16 - C: 64x8 FP32 Each mma instruction requires 4 warps (a warp group of 128 threads). This is perfect for Hopper! With this instruction: - A is loaded via ldmatrix from shared memory - B is broadcast/shared - C accumulates in registers For our problem: - M = 512 per batch, 128 batches = 65536 total M - N = 2048 - K = 1024 Tile sizes using m64n8k16: - Each warp group computes a 64x8 tile of C per mma - In K dimension, needs 1024/16 = 64 iterations To cover N=2048 efficiently: - We want larger tiles. Each warp group could compute a 64x64 tile by issuing 8 mma instructions in the N dimension. - Then we need warp groups for M dimension and batch dimension. Let me design the grid: - Each block = 1 warp group = 128 threads - Each block computes a 64x64 tile of C - N tiles per batch = 2048 / 64 = 32 - M tiles per batch = 512 / 64 = 8 - Batch tiles = 128 - Total blocks = 128 * 8 * 32 = 32768 32768 blocks for 132 SMs = ~248 blocks per SM. That's plenty. Now, shared memory per block: - A tile: 64 rows x 16 cols BF16 = 1024 elements = 2 KB per k-iteration (but we can double buffer) - B tile: 16 rows x 64 cols BF16 = 1024 elements = 2 KB per k-iteration - Total: ~4 KB per iteration, with double buffering: ~8 KB That's very manageable! We can even use larger K tiles to amortize the mma overhead. Wait, with m64n8k16, K tile is fixed at 16. So each k-iteration processes 16 elements of K. We need 1024/16 = 64 iterations. In each iteration, we need to load a 64x16 slice of A and a 16x64 slice of B into shared memory. Then the warp group executes 8 mma instructions (for N=64) each consuming the loaded k-slice. For A: 64x16 BF16 = 1024 elements. Each thread loads 1024/128 = 8 elements. We can use vectorized loads (4 elements per thread = 2 iterations, or 8 elements = 4 iterations with 16-bit loads). Actually, 1024 elements with 128 threads = 8 elements per thread. With 4-element vectorized loads (int2/bf16x2), each thread does 2 loads. For B: 16x64 BF16 = 1024 elements. Same, 8 elements per thread. This seems very efficient! Now, the main question is: can we write this kernel in load_inline format? Yes, using inline PTX for mma instructions. Actually, CUDA 12 has standard C++ interfaces for `mma` on Hopper. We can use the PTX wrappers. Wait, there's a simpler approach. Hopper supports `wmma` style operations with larger fragments. Or we can use the `cute` or `mma` APIs from CUTLASS. But without headers, inline PTX is safer. Actually, I just remembered that CUDA's `wmma` API supports `matrix_a`, `matrix_b`, and `accumulator` fragments. But for larger tiles and Hopper, it's often better to use `mma.sync.aligned.m64n8k16` directly. Let me check: does PyTorch's CUDA extension compilation support inline PTX? Yes, you can use `asm volatile` or PTX instructions in the .cu code. Here's a plan for the kernel: ```cuda __global__ void bmm_kernel( const __nv_bfloat16* __restrict__ A, const __nv_bfloat16* __restrict__ B, float* __restrict__ C, int batch_size, int m, int k, int n) { // Each block is 128 threads (4 warps) // Each block computes a 64x64 tile of C int warp_group_id = threadIdx.x / 128; // always 0 for 1 warp group per block int lane_id = threadIdx.x % 128; int warp_id = lane_id / 32; int lane = lane_id % 32; // Determine output tile coordinates int total_m_tiles = (m + 63) / 64; int total_n_tiles = (n + 63) / 64; int block_idx = blockIdx.x; int batch_idx = block_idx / (total_m_tiles * total_n_tiles); int mn_idx = block_idx % (total_m_tiles * total_n_tiles); int m_tile = mn_idx / total_n_tiles; int n_tile = mn_idx % total_n_tiles; int m_start = m_tile * 64; int n_start = n_tile * 64; // Pointers for this batch const __nv_bfloat16* A_batch = A + batch_idx * m * k; const __nv_bfloat16* B_batch = B + batch_idx * k * n; float* C_batch = C + batch_idx * m * n; // Shared memory // A_smem: 64 x 16 bf16 (with padding) // B_smem: 16 x 64 bf16 (with padding) extern __shared__ char smem[]; __nv_bfloat16* A_smem = (__nv_bfloat16*)smem; __nv_bfloat16* B_smem = (__nv_bfloat16*)(smem + 64 * 16 * sizeof(__nv_bfloat16)); // Accumulator registers: each thread holds part of 64x64 tile // For mma.m64n8k16 with 4 warps, each warp handles part of the result. // According to PTX ISA, for mma.sync.aligned.m64n8k16.row.col.f32.bf16.bf16.f32 // with 4 warps, the result is distributed across all 128 threads. // Specifically for m64n8k16 with .f32 output: // Each thread holds some number of f32 elements. // The exact distribution depends on the mma variant. // Actually, let me look at the PTX ISA for the register layout. // For mma.sync.aligned.m64n8k16.row.col.f32.bf16.bf16.f32: // - A: .row, 64x16, each thread holds 8 bf16 (since 64*16/128 = 8) // - B: .col, 16x8, each thread holds 1 bf16 (since 16*8/128 = 1)... wait that doesn't make sense. // Let me recalculate from the PTX ISA spec. // For mma.m64n8k16: // A is MxK = 64x16. Total elements = 1024. 1024 / 128 threads = 8 elements per thread. // B is KxN = 16x8. Total elements = 128. 128 / 128 threads = 1 element per thread. // C is MxN = 64x8. Total elements = 512 f32. 512 / 128 threads = 4 f32 per thread. // So each thread in the warp group holds 8 bf16 from A, 1 bf16 from B, and accumulates 4 f32 in C. // For a N=64 tile, we need 8 such mma operations in the N dimension. // So each thread accumulates 8 * 4 = 32 f32 values. float acc[32]; // 32 floats per thread for a 64x64 tile #pragma unroll for (int i = 0; i < 32; i++) acc[i] = 0.0f; // Actually, we also need to carefully manage the B matrix across sub-tiles. // For each of the 8 sub-tiles in N (0..7), we do an mma with a different B slice. // Iterate over K in steps of 16 for (int k_tile = 0; k_tile < k; k_tile += 16) { // Load A slice: [m_start:m_start+64, k_tile:k_tile+16] into smem // 64x16 = 1024 elements // Each of 128 threads loads 8 elements // Use vectorized loads for coalescing // Compute global indices for A int a_row = m_start + (threadIdx.x / 2); // 0..63, each row assigned to 2 threads int a_col = k_tile + (threadIdx.x % 2) * 8; // each thread loads 8 consecutive elements? // Hmm, for coalesced loads along K dimension: // A has shape [m, k], row-major // We want threads in a warp to access consecutive memory addresses // So thread t should load column (t * 8) % 16 + some_offset, row = t / 2 // Actually, let me think about this more carefully. // For A_smem layout, we want it to match the mma expectation. // The mma instruction expects A in a specific layout in shared memory. // For simplicity, I'll use a straightforward layout and ldmatrix. } } ``` This is getting quite complex. Let me think about whether there's a simpler way. Actually, what if I use a kernel that calls the standard `wmma` sub-warp operations? For Ampere/Hopper: ```cpp #include <mma.h> using namespace nvcuda::wmma; // For TF32, we can use: // fragment<matrix_a, 16, 16, 8, precision::tf32, row_major> a_frag; // fragment<matrix_b, 16, 16, 8, precision::tf32, col_major> b_frag; // fragment<accumulator, 16, 16, 8, float> c_frag; mma_sync(c_frag, a_frag, b_frag, c_frag); ``` Wait, I need to check: does WMMA support TF32 on Hopper? Yes, it does with `precision::tf32` and k=8. But the throughput of wmma with 16x16x8 tiles is lower than what we could get with m64n8k16. Still, it might be easier to write and could still outperform cuBLAS if the tile sizes are perfectly chosen. Actually, let me reconsider whether we even need tensor cores. The H100 has 67 TFLOPS of pure FP32. For 275 GFLOPs of work, that would take at least 4.1 ms at peak. But tensor cores at 495 TFLOPS (TF32) would take ~0.55 ms. If cuBLAS achieves ~300 TFLOPS, that's ~0.9 ms. So to beat cuBLAS, we probably need tensor cores. Pure FP32 won't have a chance unless cuBLAS is somehow not using tensor cores (which it is). So let me write a WMMA-based kernel using TF32 precision. The key parameters: - batch = 128, m = 512, k = 1024, n = 2048 - WMMA tile: 16x16x8 in TF32 - Each warp handles one 16x16 tile - Block size: 32 warps = 1024 threads (max per block) - If each block has 32 warps, it can compute: - Option 1: 32 separate 16x16 tiles arranged as 4x8 grid = 64x128 tile per block - Option 2: cooperative loading but independent warps Let me design the kernel with 1 warp = 1 tile for simplicity: Grid: - M tiles per batch: 512 / 16 = 32 - N tiles per batch: 2048 / 16 = 128 - Batch: 128 - Total tiles: 32 * 128 * 128 = 524,288 tiles - Total warps needed: 524,288 - Total blocks (128 threads each = 4 warps): 131,072 blocks 131,072 blocks on 132 SMs = ~992 blocks per SM. That's fine. But each block only computes 4 tiles = 4 * 16 * 16 = 1024 output elements. The overhead of launching so many blocks and loading shared memory might be high. Let me use larger blocks: 8 warps (256 threads) computing an 8-tile arrangement like 2x4 = 32x64 tile. Total blocks: 524,288 / 8 = 65,536 blocks Still a lot but better. Or 16 warps (512 threads) computing 4x4 = 64x64 tiles. Total blocks: 524,288 / 16 = 32,768 blocks. That's 32,768 blocks on 132 SMs = ~248 blocks per SM. Good. Shared memory per block for 64x64 tile, K=8: - A: 64 x 8 = 512 elements TF32 = 2KB - B: 8 x 64 = 512 elements TF32 = 2KB - But wait, for wmma TF32, each k-step is 8. So: - A_smem: 64 x 8 TF32 = 2KB - B_smem: 8 x 64 TF32 = 2KB - Actually, if we use TF32, the data is originally FP32. We could load FP32 and the hardware converts to TF32 internally during mma. But that wastes memory bandwidth. - Better: we could convert A and B to TF32 when loading? No, TF32 is just FP32 with truncated mantissa. The hardware does that automatically in the tensor core. Wait, for wmma::load_matrix_sync with TF32, do we pass FP32 arrays? Let me check the docs. Actually, wmma::fragment for `precision::tf32` stores data as `float`. So yes, we load FP32 data into fragments and the hardware rounds to TF32 during mma_sync. But this means we're loading FP32 from memory, which uses more bandwidth and shared memory. An optimization: we could process the kernel in FP16 or BF16, since the inputs are FP32 but we only need 1e-4 tolerance. For BF16: - We convert A and B from FP32 to BF16 when loading to shared memory - Use BF16 tensor cores (much faster than TF32 on H100) - Accumulate in FP32 On H100, BF16 tensor core throughput is ~1979 TFLOPS vs ~495 TFLOPS for TF32. That's a 4x speedup! And with our tolerance analysis above, BF16 should pass 1e-4. So let's use BF16 wmma if available. Does wmma support BF16? Yes, on Ampere and later: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 16, __nv_bfloat16, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, 16, 16, 16, __nv_bfloat16, wmma::col_major> b_frag; wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag; ``` Note: for BF16, K=16 in the fragment (not 8 like TF32). With BF16 16x16x16 tiles: - Total tiles: (512/16) * (2048/16) * 128 = 32 * 128 * 128 = 524,288 tiles - Per block with 16 warps (512 threads), each block handles 16 tiles. - We can arrange 16 warps as a 4x4 grid of tiles, computing a 64x64 tile per block. - Per batch: (512/64) * (2048/64) = 8 * 32 = 256 blocks - Total blocks: 128 * 256 = 32,768 Good block count. Shared memory: - A_smem for 64x64 tile with K=16: 64 x 16 = 1024 BF16 elements = 2 KB - B_smem for K=16 x 64: 16 x 64 = 1024 BF16 elements = 2 KB - Total: 4 KB per iteration. Very small! We can increase K tile. Wait, if we want to load a larger K tile to reduce global memory load overhead: - Load K=64 at once - A_smem: 64 x 64 = 4096 BF16 = 8 KB - B_smem: 64 x 64 = 4096 BF16 = 8 KB - Total: 16 KB. Still very small. Actually, with wmma, we process K in steps of 16 anyway. So we can load K=32, 48, or 64 and then have the warps process it in sub-steps. For K=1024, if we load K=32 per iteration, we have 32 outer iterations. In each iteration, all warps in the block load a 64x32 A slice and 32x64 B slice to shared memory. Then each warp loops 2 times over the K=16 steps. This seems pretty efficient. Let me also think about vectorized loads. When loading from global memory to shared memory: - A is [batch, m, k] = [128, 512, 1024] - To load A[m_start:m_start+64, k_tile:k_tile+32]: - Each row has 32 BF16 = 64 bytes - 64 rows - We have 512 threads, so each thread loads (64*32)/512 = 4 elements = 8 bytes - We can use `__nv_bfloat162` loads (4 bytes each), so 2 loads per thread Actually, (64*32) = 2048 elements, divided by 512 threads = 4 elements per thread. Perfect for `__nv_bfloat162` loads. For B: - B is [batch, k, n] = [128, 1024, 2048] - To load B[k_tile:k_tile+32, n_start:n_start+64]: - Layout is row-major, so consecutive elements in N dimension are adjacent in memory - Each row has 64 BF16 = 128 bytes - 32 rows - Total elements: 2048 - Each thread loads 4 elements = 8 bytes - Also good for vectorized loads One issue: bank conflicts in shared memory. For A_smem layout [64][32], consecutive threads accessing consecutive columns should be fine (each row is 64 bytes = 32 banks for BF16, so threads in warp access consecutive elements without conflict... actually let me check). For BF16, each element is 2 bytes. A warp of 32 threads accessing consecutive 2-byte elements in a row: elements 0..31 are at offsets 0, 2, 4, ... 62 bytes. Banks are 4 bytes each on modern GPUs. So element 0 is bank 0, element 1 is bank 0 (offset 2), element 2 is bank 1 (offset 4). Wait: Bank = (address / 4) % 32 for 4-byte banks. BF16 at byte offsets: - element 0: offset 0 → bank 0 - element 1: offset 2 → bank 0 - element 2: offset 4 → bank 1 - element 3: offset 6 → bank 1 ... - element 2i: bank i - element 2i+1: bank i So consecutive threads t loading element t would access: - thread 0: element 0 offset 0, bank 0 - thread 1: element 1 offset 2, bank 0 ← CONFLICT! - thread 2: element 2 offset 4, bank 1 - thread 3: element 3 offset 6, bank 1 ← CONFLICT! This is a 2-way bank conflict for BF16. We can pad to avoid this. For A_smem, if we make the row stride 33 or 34 instead of 32: - element i at offset i * 2 (padding would need to be odd to avoid conflicts) - Actually with 4-byte banks, if we pad by 1 element (2 bytes), then: - element 0: offset 0 → bank 0 - element 1: offset 2 → bank 0 - ... - element 31: offset 62 → bank 15 - element 32: offset 64 → bank 16... wait, that's row 1 Let me think again. If we have a 2D array [64][33]: - Row i starts at offset i * 33 * 2 = i * 66 bytes - Element [i][j] is at byte offset i * 66 + j * 2 - Bank = ((i * 66 + j * 2) / 4) % 32 - For a fixed row i, bank for thread j: - bank = ((i * 66)/4 + j/2) % 32 approximately - Actually: j*2/4 = j/2 - So threads 0,1 have same bank, threads 2,3 have same bank, etc. The conflict comes from adjacent threads accessing adjacent BF16 elements that fall in the same 4-byte bank. One solution is to transpose in shared memory or have each thread load 2 consecutive elements (a `__nv_bfloat162`). Actually, if each thread loads a `__nv_bfloat162` (2 BF16 values = 4 bytes) which occupies exactly one bank, then 32 threads in a warp each load from different banks (since each load is 4 bytes and consecutive), as long as the starting addresses are aligned to 128 bytes (warp size * 4 bytes). Let me reconsider. If each thread loads a `__nv_bfloat162` from consecutive memory: - Thread 0 loads bytes [0, 3] → bank 0 - Thread 1 loads bytes [4, 7] → bank 1 - ... - Thread 31 loads bytes [124, 127] → bank 31 This has NO bank conflicts for loading, provided the starting address is 128-byte aligned. But for the WMMA load, the fragments need to be loaded in a specific pattern. The `wmma::load_matrix_sync` handles the layout internally when loading from shared memory. Actually, wait. `wmma::load_matrix_sync` can load from shared memory or global memory. The shared memory layout needs to match the expected fragment layout. The GPU hardware handles the distribution. For shared memory access during `wmma::load_matrix_sync`: - For row-major A fragments, the shared memory should contain the matrix in row-major order, but with potential padding. - Bank conflicts during wmma loads are handled by the hardware. Let me not overthink the bank conflicts and instead focus on writing the kernel using `wmma::load_matrix_sync` from shared memory. Here's my revised kernel design: 1. Convert A and B from FP32 to BF16 on the fly when loading to shared memory 2. Use WMMA BF16 16x16x16 fragments 3. Each block has 512 threads (16 warps) computing a 64x64 tile 4. Each warp computes a 16x16 tile 5. K is processed in outer loops of K_TILE=32, inner wmma loops of 16 Wait, 16 warps computing 16 tiles each of 16x16 = total 4096 elements. But 64x64 = 4096 elements. Good. For the 16 warps, we can arrange them as: - warp_row = warp_id / 4 (0..3) - warp_col = warp_id % 4 (0..3) - Each warp computes C[m_start + warp_row*16 : m_start + (warp_row+1)*16, n_start + warp_col*16 : n_start + (warp_col+1)*16] For outer K loop with K_TILE=32: - Load A_smem[64][32] and B_smem[32][64] (with padding) - Then each warp does 2 iterations of wmma (k=0..15 and k=16..31) For loading global memory: - A: [batch, m, k], want A[batch, m_start:m_start+64, k_tile:k_tile+32] - Total: 64*32 = 2048 BF16 - 512 threads, each loads 4 BF16 = 2 `__nv_bfloat162` - For thread t: - Row = t / 8 (0..63), since 512/8 = 64 rows, each row has 8 threads - Col = (t % 8) * 4 (0, 4, 8, ..., 28), each thread loads 4 consecutive BF16 - This gives coalesced access since threads in warp access consecutive memory. - B: [batch, k, n], want B[batch, k_tile:k_tile+32, n_start:n_start+64] - Total: 32*64 = 2048 BF16 - 512 threads, each loads 4 BF16 - For thread t: - Row = t / 16 (0..31), each row has 16 threads - Col = (t % 16) * 4 (0, 4, 8, ..., 60), each thread loads 4 consecutive BF16 - Also coalesced. For converting from FP32 to BF16 during load: - Read 2 FP32 values per thread (or 4?) - Convert each to BF16 using `__float2bfloat16` or hardware conversion Actually, if each thread loads 4 BF16, and we're loading from FP32 global memory, each thread needs to load 4 FP32 values = 16 bytes. That's easily done with 2 `float2` loads. Wait, `float2` loads would be vectorized loads of 8 bytes each, so 2 loads = 16 bytes = 4 floats. So: each thread loads 4 floats from A, converts to 4 BF16, stores to shared memory. Let me verify the indexing: For A global load: ```cuda int t = threadIdx.x; int rowA = t / 8; // 0..63 int colA = (t % 8) * 4; // 0, 4, 8, ..., 28 float4 a_val = *(float4*)(A_ptr + (m_start + rowA) * k + k_tile + colA); ``` But wait, k_tile + colA goes from 0 to 31 in steps of 4, so a `float4` loads 4 consecutive floats. That's 16 bytes. Threads in a warp (t=0..31): - t=0: row=0, col=0, loads A[0][0..3] - t=1: row=0, col=4, loads A[0][4..7] - ... - t=7: row=0, col=28, loads A[0][28..31] - t=8: row=1, col=0, loads A[1][0..3] - t=9: row=1, col=4, loads A[1][4..7] - ... In a warp, threads 0..7 access A[0][0..31], which is perfectly coalesced (consecutive memory for the warp). Threads 8..15 access A[1][0..31], also coalesced but different row. This should be efficient since different warps handle different rows, and within a warp it's coalesced. For B global load: ```cuda int rowB = t / 16; // 0..31 int colB = (t % 16) * 4; // 0, 4, 8, ..., 60 float4 b_val = *(float4*)(B_ptr + (k_tile + rowB) * n + n_start + colB); ``` B is [k, n] row-major. Threads in warp: - t=0: row=0, col=0, loads B[0][0..3] - t=1: row=0, col=4, loads B[0][4..7] - ... - t=15: row=0, col=60, loads B[0][60..63] - t=16: row=1, col=0, loads B[1][0..3] Again, coalesced within each row, but with 16 threads per row, that's 2 warps per row. Hmm, for rowB=0, threads 0..15 (half warp) access B[0][0..63]. That's perfectly coalesced. Actually, threads 0..31 in warp 0: - 0..15: row 0 - 16..31: row 1 Warp 0 accesses B[0][0..63] and B[1][0..63]. That's 2 rows. This is still fine. For converting and storing to shared memory: ```cuda __nv_bfloat162 a_smem[2]; // each thread stores to shared memory // Actually, we need to write back to shared memory positions. ``` Actually, the shared memory layout should just be row-major with some padding. Let's allocate: - A_smem with shape [64][32+padA] - B_smem with shape [32][64+padB] For A_smem, if we pad by 1 (33 columns), that's 66 bytes per row. For BF16, bank calculation: 2 bytes per element, 4-byte banks. Row i, element j: offset = i * 66 + j * 2 Bank = (offset / 4) % 32 For row i, threads loading consecutive `float4` (which becomes 4 BF16 = 8 bytes)? Actually, each thread writes 4 BF16 elements to shared memory. Let's say thread t writes elements at columns colA to colA+3. - colA = (t % 8) * 4 - Elements at j = colA, colA+1, colA+2, colA+3 - Offsets = i*66 + colA*2, i*66 + colA*2 + 2, i*66 + colA*2 + 4, i*66 + colA*2 + 6 - Banks: (i*66/4 + colA/2) % 32, (i*66/4 + colA/2) % 32, (i*66/4 + colA/2 + 1) % 32, (i*66/4 + colA/2 + 1) % 32 For t=0: colA=0, banks: 0, 0, 1, 1 For t=1: colA=4, banks: 2, 2, 3, 3 ... For t=7: colA=28, banks: 14, 14, 15, 15 In a warp, threads 0..7 write to banks {0,0,1,1}, {2,2,3,3}, ..., {14,14,15,15}. Since each thread writes 4 elements, but the first two are bank 0, that's a bank conflict. But because each thread is writing its own 4 elements independently (not broadcasting), there might not be a true conflict... actually, there is a 2-way conflict for each pair of banks if multiple threads try to write to the same bank simultaneously. Wait, threads 0 and 1 write to different banks (0,0,1,1 vs 2,2,3,3), so no conflict between threads. Each thread writes 4 elements, and within a single thread's writes, the first two elements go to the same bank. But since it's the same thread writing to the same bank in sequence, there's no conflict, just serialization within that thread. However, threads in a warp execute together, and if they all write to different banks, that's fine. The issue is only when different threads in the same warp access the same bank with different addresses. For t=0: banks 0, 0, 1, 1 For t=1: banks 2, 2, 3, 3 ... For t=7: banks 14, 14, 15, 15 All threads in the warp access distinct banks for their first element, distinct banks for their second element, etc. No conflicts! Wait, but for B_smem with shape [32][65] (padding by 1): - rowB = t / 16, colB = (t % 16) * 4 - Thread t writes B elements at columns colB to colB+3 - Offset = rowB * 65 * 2 + colB * 2 - For threads 0..15 in a warp: - All have rowB=0 or rowB=1 - t=0..15: rowB=0, colB=0,4,8,...,60 - Banks for first element: (0 + colB/2) % 32 = 0, 2, 4, ..., 30 - All distinct! - t=16..31: rowB=1 So again no conflicts. This shared memory layout seems good. Now for the wmma loads from shared memory: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 16, __nv_bfloat16, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, 16, 16, 16, __nv_bfloat16, wmma::col_major> b_frag; wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag; wmma::load_matrix_sync(a_frag, A_smem + warp_row * 16 * A_PAD + k_inner, A_PAD); wmma::load_matrix_sync(b_frag, B_smem + k_inner * B_PAD + warp_col * 16, B_PAD); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); ``` Wait, B is loaded as col_major for wmma. But our B_smem is stored in row-major order. We need to handle this carefully. When using `wmma::load_matrix_sync` for `matrix_b` with `col_major` layout, the function expects the data in column-major format in memory. But our B matrix in the GEMM is mathematically used as B^T if we store it row-major. Standard GEMM: C = A * B - A is [M, K] row-major - B is [K, N] row-major - For wmma, if we load A as row_major and B as col_major, then: - B col_major means the memory layout is column-major, so element (i,j) is at memory[j*K+i] - But our B row-major has element (i,j) at memory[i*N+j] - So if we read B with col_major semantics from a row-major array, it's effectively reading B^T Actually, for standard GEMM C = A * B: - If A is row-major, load A fragment as `row_major` - If B is row-major, we can load B fragment as `col_major` but treat the stride appropriately. Hmm, let me think more carefully. wmma::fragment<matrix_b, M, N, K, T, col_major> means that when we call load_matrix_sync, the data is interpreted as being stored in column-major order. If we have B stored row-major as [K][N], and we want to load a KxN tile into a B fragment of shape KxN: - For col_major load, the memory access pattern accesses elements as if they were stored with column-major indexing. - This means element at memory offset `j * K + i` corresponds to logical element (i, j). - If our actual memory has element (i, j) at offset `i * N + j`, then using col_major semantics is like transposing. Actually, for C = A * B where A is MxK and B is KxN: - In column-major BLAS, C = A * B means C_{i,j} = sum_k A_{i,k} * B_{k,j} - If A is stored row-major (offset = i*K + k) and we load it as row_major, wmma interprets it correctly. - If B is stored row-major (offset = k*N + n) and we load it as col_major, wmma interprets offset as if it were n*K + k. But our actual offset is k*N + n, which is NOT equal to n*K + k. So loading a row-major B as col_major does NOT work correctly unless K=N (not our case). The correct approach: 1. Load B as row_major if that's how it's stored. But wmma for matrix_b supports col_major and row_major. Let me check: yes, `wmma::fragment` for `matrix_b` can be either `row_major` or `col_major`. For `C = A * B`: - A fragment: MxK, row_major (since A is MxK row-major) - B fragment: KxN, row_major (since B is KxN row-major) Then mma_sync computes C = A * B correctly. Actually, I need to verify this. In the original Ampere/Hopper wmma docs: - `fragment<matrix_a, M, N, K, T, LayoutA>` for A matrix - `fragment<matrix_b, M, N, K, T, LayoutB>` for B matrix - `fragment<accumulator, M, N, K, T>` for C matrix `mma_sync` computes D = A * B + C where: - A is MxK - B is KxN - C and D are MxN So for our case with M=16, N=16, K=16: - A_frag represents 16x16 of A (sub-tile of the MxK matrix) - B_frag represents 16x16 of B (sub-tile of the KxN matrix) For row-major A (which is [M][K]), a 16x16 tile starting at (m0, k0) is loaded with `row_major`. For row-major B (which is [K][N]), a 16x16 tile starting at (k0, n0) is loaded with `row_major`. Then `mma_sync(c_frag, a_frag, b_frag, c_frag)` computes: - c_frag[m,n] += sum_k a_frag[m,k] * b_frag[k,n] But wait, B_frag loaded as row_major from B[k0..k0+15][n0..n0+15]: - b_frag[k,n] = B[k0+k][n0+n] And A_frag loaded as row_major from A[m0..m0+15][k0..k0+15]: - a_frag[m,k] = A[m0+m][k0+k] Then the product is: sum_k A[m0+m][k0+k] * B[k0+k][n0+n] But in standard matrix multiplication C[m,n] = sum_k A[m,k] * B[k,n], the indices should be: sum_k A[m0+m][k0+k] * B[k0+k][n0+n] Yes, that's correct! So we can use `row_major` for both A and B fragments. I previously thought matrix_b needed col_major, but that's only if we want to transpose B. For standard GEMM with both matrices in row-major, we use `row_major` for A and `row_major` for B. Wait, let me double-check with the WMMA documentation... The WMMA PTX instruction `mma.sync.aligned.m16n16k16.row.row` exists and does exactly what we need: A is row-major, B is row-major. So yes, we use row_major for both. Actually, there's a subtlety. The wmma API in CUDA C++ uses `col_major` for `matrix_b` in many examples because those examples are doing C = A * B^T or the matrices are in column-major format (Fortran/BLAS style). For our case with row-major tensors from PyTorch: - A: row_major - B: row_major So we use `wmma::row_major` for both `matrix_a` and `matrix_b`. Great! So the kernel becomes: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 16, __nv_bfloat16, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, 16, 16, 16, __nv_bfloat16, wmma::row_major> b_frag; wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc_frag; wmma::fill_fragment(acc_frag, 0.0f); for (int k_outer = 0; k_outer < K; k_outer += K_TILE) { // Load A and B to shared memory __syncthreads(); // Then for k_inner in 0, 16 within K_TILE for (int k_inner = 0; k_inner < K_TILE; k_inner += 16) { wmma::load_matrix_sync(a_frag, A_smem + (warp_row * 16) * A_PAD + k_inner, A_PAD); wmma::load_matrix_sync(b_frag, B_smem + (k_inner) * B_PAD + (warp_col * 16), B_PAD); wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); } } // Store result wmma::store_matrix_sync(C_ptr + warp_row * 16 * n + warp_col * 16, acc_frag, n, wmma::mem_row_major); ``` Wait, wmma::store_matrix_sync for accumulator needs the memory layout parameter too. Actually, for `wmma::store_matrix_sync`, the layout parameter doesn't exist in the same way. The accumulator is always MxN. We store it to row-major memory with: ```cpp wmma::store_matrix_sync(C_smem_or_global + offset, acc_frag, leading_dimension, wmma::mem_row_major); ``` Now, there's another issue. The wmma::load_matrix_sync expects the shared memory pointer to point to the start of the tile. For A fragment of shape 16x16 from a [64][K_TILE] shared memory: - Warp at row `warp_row` of the 4x4 grid loads A tile at rows [warp_row*16, (warp_row+1)*16) and some k columns. - For each k_inner step: columns [k_inner, k_inner+16) - The shared memory pointer should be `&A_smem[warp_row * 16][k_inner]` - The leading dimension is `A_PAD` For B fragment of shape 16x16 from a [K_TILE][64] shared memory: - Warp at col `warp_col` loads B tile at rows [k_inner, k_inner+16) and cols [warp_col*16, (warp_col+1)*16) - The shared memory pointer should be `&B_smem[k_inner][warp_col * 16]` - The leading dimension is `B_PAD` This looks correct! Let me also check: we need to handle the batch dimension. Each block handles one 64x64 tile for one batch element. With 128 batches, 8 m-tiles, 32 n-tiles = 32768 blocks total. For blockIdx.x to (batch, m_tile, n_tile): ```cpp int n_tiles = n / 64; // 32 int m_tiles = m / 64; // 8 int total_tiles_per_batch = m_tiles * n_tiles; // 256 int batch_idx = blockIdx.x / total_tiles_per_batch; int tile_idx = blockIdx.x % total_tiles_per_batch; int m_tile = tile_idx / n_tiles; int n_tile = tile_idx % n_tiles; int m_start = m_tile * 64; int n_start = n_tile * 64; ``` Now for K_TILE. We want K_TILE to be large enough to amortize the __syncthreads() cost but small enough for shared memory. For K=1024, we could do K_TILE=32, 64, or even 128. With K_TILE=64: - A_smem: 64 x 64 BF16 = 4096 * 2 = 8 KB - B_smem: 64 x 64 BF16 = 8 KB - Total: 16 KB. Great. But then we have 1024/64 = 16 outer iterations, each with 64/16 = 4 inner wmma iterations. Total 16*4 = 64 wmma calls per warp. With K_TILE=128: - A_smem: 64 x 128 = 16 KB - B_smem: 128 x 64 = 16 KB - Total: 32 KB. Still within limits. But 32 KB might limit occupancy. Let me check H100 shared memory: - Per SM: 228 KB - Max per block: 227 KB - But to achieve good occupancy, we want shared memory per block to allow multiple blocks per SM. With 512 threads per block and shared memory config: - If each block uses 32 KB shared memory, then 228/32 = ~7 blocks per SM maximum. - With 512 threads, 2048 threads per SM max, so 4 blocks per SM maximum due to thread count. - So occupancy is thread-limited, not shared-memory-limited. Let's use K_TILE=64 for a balance. But wait, we also need to think about double buffering to hide latency. With K_TILE=64, we load 64 K values at once, then do 4 wmma iterations. The loads are amortized over 4 wmma calls. Alternatively, K_TILE=32 with 2 inner iterations might also work well. Let me try K_TILE=64 first. Now, let me also think about whether we need to convert the inputs from FP32 to BF16. For correctness, we can convert on the fly using: ```cpp __nv_bfloat16 __float2bfloat16(float a); ``` Or in CUDA 12, we can use: ```cpp __nv_bfloat16 h = __float2bfloat16_rn(f); ``` Actually, the standard function is `__float2bfloat16` which might round to nearest even. Now, the numerical question: will converting from FP32 to BF16 and computing the GEMM in BF16 with FP32 accumulation give results within 1e-4 of torch.bmm in FP32? torch.bmm in FP32 by default uses TF32 on H100 (since allow_tf32 is True). TF32 has 10 bits of mantissa compared to BF16's 7 bits and FP32's 23 bits. So TF32 is more accurate than BF16. But our tolerance is 1e-4 = 0.01%. For outputs that can be as large as K (since A, B ~ [0,1]), max output is about K = 1024. Absolute tolerance is 1e-4. Relative tolerance is 1e-4. For an output of ~512 (mean case), absolute error tolerance is 512 * 1e-4 = 0.0512. So we need absolute error less than about 0.05 for typical values. With BF16 accumulation in the tensor cores but FP32 output accumulator: - Each multiply is BF16 x BF16 → the error per product is about 1-2 ULP of BF16. - BF16 ULP at value ~1 is 2^-7 ≈ 0.0078. - 1024 products summed: worst-case error could be 1024 * 0.0078 ≈ 8, but that's worst-case. - With random inputs, errors average out. The standard deviation of rounding errors is much smaller. - In practice, mixed-precision GEMM with BF16 multiply and FP32 accumulate typically achieves <0.01 relative error on random data in [0,1]. So we should be fine. Actually wait, I just realized: torch.bmm in FP32 with TF32 enabled will also have reduced precision (10-bit mantissa vs 7-bit for BF16). TF32 has精度 somewhere between BF16 and FP16. The difference between TF32 and BF16 might be noticeable but should still be within tolerance. Let me be more precise. TF32 mantissa is 10 bits (same as FP16's range, but with FP32 exponent). BF16 has 7 bits of mantissa. For values ~1: - TF32 ULP: 2^-10 ≈ 0.001 - BF16 ULP: 2^-7 ≈ 0.0078 The per-product error for TF32 is smaller than BF16 by a factor of ~8. But both should still pass the 1e-4 relative tolerance, especially since the tolerance is pretty loose (1e-3 for BF16 tasks, but here it's FP32 output with 1e-4 tolerance). For random [0,1] inputs: - Mean output = K/4 = 256 - Std of output = sqrt(K * var(product)) ≈ sqrt(1024 * 1/9) ≈ 10.7 (using uniform distribution variance ~1/12) - So most outputs are in [200, 300] range. - Relative tolerance of 1e-4 means absolute tolerance of ~0.03. - BF16 GEMM with FP32 accumulation typically has errors around 0.01-0.1 for matrices of this size. Hmm, on the edge. Let me think about whether to use TF32 or BF16. TF32 tensor cores: ~495 TFLOPS. Still much faster than pure FP32. BF16 tensor cores: ~1979 TFLOPS. Even faster. Maybe I should try BP16 first and see if it passes correctness. Actually, the tolerance is `atol=rtol=1e-4` for fp32. If outputs are around 256, rtol=1e-4 means atol requirement is 0.0256. And atol itself is 1e-4. Let me do a quick calculation for BF16 error. Each product A[i][k] * B[k][j] has error up to ~0.01 (since BF16 ULP at value ~1 is ~0.008). With 1024 terms, by the central limit theorem, the error distribution has std dev ~0.01 * sqrt(1024) / sqrt(3) ≈ 0.01 * 32 / 1.7 ≈ 0.19. Wait, that's too large. Let me recalculate. The rounding of each product to BF16: the exact product p is rounded to BF16. Error e = round(p) - p. |e| <= 0.5 ULP(BF16) ≈ 0.004. e is approximately uniform in [-0.004, 0.004]. Sum of 1024 independent uniform errors: - Mean = 0 - Variance = 1024 * (0.004)^2 / 3 ≈ 1024 * 5.3e-6 ≈ 0.0055 - Std dev ≈ 0.074 So errors should mostly be within ±0.2 or so. The tolerance is max(1e-4, 1e-4 * |expected|). For expected=256, tolerance is 0.0256. But error std is 0.074, so we'd have many elements outside tolerance! Hmm, that's concerning. But actually: 1. The inputs A and B are also in BF16 in our kernel, whereas in the reference they're FP32 (or TF32). So A and B themselves have rounding errors compared to the reference. 2. torch.bmm with TF32 also rounds A and B to TF32, so it has similar but smaller rounding errors. Wait, torch.bmm with TF32 doesn't convert A and B to TF32 explicitly; the tensor cores do it internally during each mma. But since A and B are stored in FP32, each load gets the full FP32 value, and the mma rounds it to TF32 on-the-fly. In our kernel, we convert A and B to BF16 in shared memory. This means: - A_smem[i] = round_to_bf16(A_fp32[i]) - B_smem[i] = round_to_bf16(B_fp32[i]) Then the products use these rounded values. So the total error is: 1. Error from rounding A to BF16 2. Error from rounding B to BF16 3. Error from rounding products to BF16 This might be too much for the 1e-4 tolerance. But wait, I need to think about what torch.bmm actually computes with TF32. If A and B are FP32, then torch.bmm with TF32 internally rounds each operand to TF32 in the tensor cores. The error per operand is at most 0.5 ULP(TF32) ≈ 0.0005. In our case with BF16 shared memory, the error per operand is at most 0.5 ULP(BF16) ≈ 0.004. And product error: - For TF32: product of two FP32 values rounded to TF32. The product p is exact in some higher precision, then rounded to TF32. Error ≈ 0.0005. - For BF16: product rounded to BF16. Error ≈ 0.004. So BF16 is about 8x less accurate per operation. Across 1024 operations, the accumulated error might be too large. Let me try to be more precise about the tolerance. For A, B in [0,1]: - Exact output C_exact = sum A_k * B_k - Reference output C_ref = torch.bmm(A, B) which uses TF32. C_ref ≈ C_exact + error_tf32 - Our output C_ours ≈ C_exact + error_bf16 - We need |C_ours - C_ref| < 1e-4 * |C_ref| + 1e-4 - |C_ours - C_ref| ≈ |error_bf16 - error_tf32| ≈ |error_bf16| (since error_tf32 is smaller) For sum of 1024 products with BF16 operands: - Each product: A_bf16 * B_bf16 - A_bf16 = A_fp32 + e_A, |e_A| < 0.004 - B_bf16 = B_fp32 + e_B, |e_B| < 0.004 - Product = (A + e_A)(B + e_B) = AB + A*e_B + B*e_A + e_A*e_B - Error ≈ A*e_B + B*e_A (first order), bounded by ~0.008 in magnitude - Then the product is rounded to BF16: additional error ~0.004 - Total product error bound: ~0.012 - But the errors are random and mostly independent - Standard deviation of sum: sqrt(1024) * 0.012 / sqrt(3) * factor... this is roughly 0.2 as estimated. For normal distribution with std=0.2, probability of |error| > 0.025 is about 88%. That's way too much. So BF16 might not pass the 1e-4 tolerance for fp32! I should use TF32 instead, which should have error ~8x smaller and likely pass. For TF32 with FP32 accumulation: - Per operand rounding: 0.0005 - Per product rounding: 0.0005 - Total per-product error std: ~0.0005 / sqrt(3) * sqrt(2) ≈ 0.0004 (very rough) - Sum of 1024: std ≈ 0.0004 * sqrt(1024) ≈ 0.013 For expected outputs of ~256, atol tolerance is 0.0256. With error std of 0.013, most elements should pass! The 99.7% range is ±3*std = ±0.039, which is slightly above the tolerance for large outputs. But since the reference ALSO uses TF32 (error std ~0.004), the difference should be smaller. Actually wait - if both our kernel and the reference use TF32, then the outputs should match very closely! Because we're doing the same computation (just our custom kernel vs cuBLAS). Perfect! So I should use TF32 tensor cores, same as what cuBLAS uses, but with a custom-tiled kernel. Using TF32 wmma: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag; wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_frag; ``` Note: for TF32, K=8 in the fragment dimensions. The `precision::tf32` data type stores as `float` in the fragment. With TF32 16x16x8 fragments, and our desired block tiling: - 16 warps per block = 512 threads - 4x4 warp grid = 64x64 tile - Each warp computes one 16x16 tile - K_TILE=64: outer loop has 1024/64 = 16 iterations - Inner loop: 64/8 = 8 wmma iterations Each outer iteration: 1. Load A[64][64] and B[64][64] to shared memory from FP32 global memory - No conversion needed! Just load as FP32 into shared memory. - Shared memory size: 64*64*4*2 = 32 KB for FP32 - Or we could use smaller K_TILE=32: 64*32*4*2 = 16 KB With K_TILE=32 (16 KB shared memory): - Outer loop: 1024/32 = 32 iterations - Inner loop: 32/8 = 4 wmma iterations 32 outer iterations might have significant overhead from __syncthreads. Let me try K_TILE=64 (32 KB shared memory): - Outer loop: 16 iterations - Inner loop: 8 wmma iterations 16 is a reasonable number. And since each warp does 8 wmma_sync calls per outer iteration, the mma work dominates. Actually, 32 KB shared memory is fine for H100. Let me verify occupancy: - Threads per block: 512 - Max threads per SM: 2048 - Max blocks per SM due to threads: 2048/512 = 4 - Shared memory per SM: 228 KB - Max blocks per SM due to SMEM: 228/32 = 7 - So threads limit: 4 blocks per SM. Occupancy = 4/8 (for H100, 8 blocks per SM is sometimes the theoretical max, but with 512 threads, max is 4) = 50%. That's not bad. Could we use 256 threads (8 warps) with a 64x32 tile? - Threads per block: 256 - Blocks per SM: 8 (thread-limited) - SMEM per block: let's say 16 KB for K_TILE=32 - Blocks per SM due to SMEM: 228/16 = 14 - Thread-limited: 8 blocks per SM. Occupancy: 8/8 = 100%! That might be better. But the total work per block is halved, so there might be more overhead per FLOP. However, with 2x the blocks per SM, we might hide latency better. Let's analyze throughput: - With 512-thread blocks: 32768 blocks total. 32768 / 132 SMs = ~248 blocks per SM. At 4 concurrent blocks per SM, that's 62 waves of blocks. Good parallelization. - With 256-thread blocks and 64x32 tiles: m_tiles = 512/64 = 8, n_tiles = 2048/32 = 64. Total tiles per batch = 512. Total blocks = 128 * 512 = 65536. 65536 / 132 = 496 blocks per SM. At 8 concurrent blocks per SM, that's 62 waves. Same number of waves, but higher occupancy. Actually, with 256 threads and K_TILE=32: - Each block computes 64x32 = 2048 output elements - FLOPs per block: 2 * 64 * 32 * 1024 = 4,194,304 = 4.2 MFLOPs - Total blocks: 65536 - Total FLOPs: 275 GFLOPs Shared memory: - A_smem: 64 x 32 float = 8 KB - B_smem: 32 x 32 float = 4 KB (wait, N tile is 32, so B_sm is 32x32) - Total: 12 KB Actually, B_smem for K_TILE=32 and N=32: 32 x 32 = 1024 floats = 4 KB. Total shared: 12 KB. Or K_TILE=64 with N=32: - A_smem: 64 x 64 = 16 KB - B_smem: 64 x 32 = 8 KB - Total: 24 KB With 256 threads per block and 24 KB shared: - Blocks per SM: min(2048/256, 228/24) = min(8, 9.5) = 8 - Great occupancy! Let me use: - BlockDim = 256 threads (8 warps) - Tile M = 64, Tile N = 32 - K_TILE = 64 - M warps: 64/16 = 4 warps in M - N warps: 32/16 = 2 warps in N - Total: 4 x 2 = 8 warps. Perfect. Grid: - batch: 128 - m_tiles: 512/64 = 8 - n_tiles: 2048/32 = 64 - Total blocks: 128 * 8 * 64 = 65536 Each block: - Loads A[64][64] and B[64][32] into shared memory - 8 warps, each handles a 16x16 tile - Warp grid: 4x2 - M iterations: 4 (warp rows) - N iterations: 2 (warp cols) - K iterations: 64 (outer), 8 (inner, per wmma) Let me verify total FLOPs per block: 2 * 64 * 32 * 1024 = 4,194,304 FLOPs At 495 TFLOPS TF32, each block should take ~8.5 ns in theory, but that's unrealistic. In practice, with memory overhead, maybe a few microseconds. Hmm, 65536 blocks at say 10 µs each on 132 SMs = 65536 * 10 / 132 / 1000 ms = ~5 ms. That seems slow. Wait, if each block takes 10 µs and we have 8 blocks per SM concurrently, the actual time is: 65536 blocks / (132 SMs * 8 blocks/SM) * 10 µs = 62 waves * 10 µs = 620 µs = 0.62 ms. For 275 GFLOPs in 0.62 ms, that's 275 / 0.62 = 444 TFLOPS. That's very close to theoretical peak if achievable! But can our custom kernel actually achieve that? The issue is memory bandwidth for loading A and B tiles. Memory loads per block: - A tile: 64 * 64 = 4096 floats = 16 KB - B tile: 64 * 32 = 2048 floats = 8 KB - Total per outer K iteration: 24 KB - For 16 outer iterations: 384 KB per block - Total memory traffic for all blocks: 65536 * 384 KB = 24,576 MB = 24.6 GB At 3.35 TB/s memory bandwidth, memory-bound time would be 24.6 / 3350 = 0.0073 s = 7.3 ms. Wait, that can't be right. Let me recalculate: - 65536 blocks - Each block loads 16 KB + 8 KB = 24 KB - Total reads: 65536 * 24 KB = 1,572,864 KB = 1.57 GB Hmm, that's much less. I was multiplying by 16 outer iterations but that's wrong because each block processes the full K dimension. So each block indeed loads 24 KB total (not per iteration; the tiles are loaded once per K_TILE, and K_TILE=64 means 16 iterations, each loading 24 KB). Wait, no. With K_TILE=64, each outer iteration loads A[64][64] and B[64][32]. With 16 outer iterations, each block loads: - A: 16 * 64 * 64 = 65,536 floats = 256 KB - B: 16 * 64 * 32 = 32,768 floats = 128 KB - Total: 384 KB per block Total for all blocks: 65536 * 384 KB = 25,165,824 KB ≈ 24 GB. But wait - in a batched GEMM, each A and B element is only read once by some block. Total data in A = 128 * 512 * 1024 * 4 = 268 MB. Total data in B = 128 * 1024 * 2048 * 4 = 1,073 MB. Total = 1,341 MB ≈ 1.34 GB. Why does my per-block calculation give 24 GB? Because I forgot that the tiles overlap in the K dimension for different MN tiles, but not in the M and N dimensions. Actually, each A element at position [batch][m][k] is read by exactly one block: the block that computes the MN tile containing m. Since K is iterated over, A[batch][m][k] for all k values are read by the same block. So total A reads = total A size = 268 MB. Similarly, each B element is read by exactly one block. Total B reads = 1,073 MB. Total reads = 1,341 MB. Total writes = C size = 128 * 512 * 2048 * 4 = 536 MB. Total memory traffic = ~1.88 GB. At 3.35 TB/s, this is 0.56 ms. So the kernel is compute-bound if we can achieve >200 TFLOPS. Good, so our calculation of memory traffic being a bottleneck was wrong - the 24 GB apparent total was counting each element multiple times across different blocks for different K_TILEs, but actually each element is only needed by specific blocks. OK great, so with 1.88 GB of memory traffic and 275 GFLOPs of compute, the arithmetic intensity is 275/1.88 = 146 FLOPs/byte. For TF32 at 495 TFLOPS and 3.35 TB/s bandwidth, the break-even point is 495/3.35 = 148 FLOPs/byte. So we're right at the break-even! This means both memory and compute matter. If we can achieve 80% of TF32 peak (396 TFLOPS), we'd be compute-bound and take 0.69 ms. If we only achieve 50% (247 TFLOPS), we'd be somewhat memory-bound at ~1.1 ms. cuBLAS torch.bmm likely achieves good utilization. We need to match or beat it. One advantage of our custom kernel: we can process the specific shape more efficiently than cuBLAS's general batched kernel, potentially achieving better register and shared memory utilization. Now, there's another optimization: instead of using shared memory for both A and B with K_TILE=64, we could use the new TMA (Tensor Memory Accelerator) on Hopper to load directly to shared memory. But that's quite complex. Alternatively, we can increase the tile sizes to increase arithmetic intensity: - Tile M = 128, Tile N = 64 - Threads per block: 512 (32 warps in 8x4 grid) - K_TILE = 64 - A_smem: 128 x 64 = 32 KB - B_smem: 64 x 64 = 16 KB - Total: 48 KB - Blocks per SM: min(2048/512, 228/48) = min(4, 4.75) = 4 - Total blocks: 128 * 4 * 32 = 16384 - Waves: 16384 / (132 * 4) = 31 Each block does 2 * 128 * 64 * 1024 = 16.7 MFLOPs. Memory per block: 128*64*4 + 64*64*4 = 32KB + 16KB = 48KB for K_TILE=64, over 16 iterations. Total memory per block: 16 * 48KB = 768KB. Total memory traffic: 16384 * 768KB = 12.6 GB... wait, again this counts each element multiple times. Actually, the correct total memory traffic is still the same: 1.88 GB. The per-block analysis is misleading because different blocks read disjoint MN tiles. With larger tiles: - Fewer blocks (16384) - More work per block - Better amortization of overhead - Same memory traffic - But fewer blocks might mean worse load balancing if some SMs finish early With smaller tiles: - More blocks (65536) - Less work per block - More overhead - But more parallelism and better latency hiding I think 64x32 or 64x64 are good sweet spots. Let me go with 64x64 tiles and 256 threads (8 warps): Wait, 64x64 requires 16 warps (one per 16x16 tile). So that's 512 threads. OK, 512 threads with 64x64 tile, K_TILE=64: - Blocks: 128 * 8 * 32 = 32768 - Blocks per SM: 4 (thread-limited) - Total waves: 32768 / (132*4) = 62 - Each block FLOPs: 2 * 64 * 64 * 1024 = 8.4M - Shared memory per K_TILE: 64*64*4 + 64*64*4 = 32KB - But with 16 K-iterations, total SMEM used concurrently is just 32KB. Wait, A_smem for K_TILE=64 is 64x64 floats = 16KB. B_smem is also 64x64 = 16KB. Total = 32KB. With 512 threads and 32KB SMEM: - Blocks per SM = min(2048/512, 228/32) = min(4, 7.1) = 4 - 4 blocks per SM is fine. Actually, let me double check whether `wmma::fragment` with `precision::tf32` is available on Hopper. TF32 tensor cores were introduced with Ampere (SM80) and are supported on Hopper (SM90). So yes. However, I should also consider using a simpler approach: write a kernel that uses the CUTLASS library. PyTorch sometimes includes CUTLASS headers, but it's not guaranteed. Actually, let me check if we can use `cublasGemmStridedBatchedEx` directly. Since `torch.bmm` already uses cuBLAS with BF16 or TF32, calling it directly won't give us any advantage unless torch.bmm has overhead. But wait - torch.bmm might be doing extra work like checking dimensions, preparing descriptors, etc. A direct cuBLAS call could avoid some of that overhead. However, for a batch size of 128 with large matrices, the overhead of torch.bmm is negligible compared to the compute time. Hmm, another idea: what if we use `cutlass` through `torch`'s built-in support? Some PyTorch builds include CUTLASS. But this is risky. OK, let me commit to writing a custom WMMA kernel with TF32. I'll implement it step by step. Actually, I just realized there might be an even simpler and faster approach. On H100, PyTorch might be using cuBLAS's default algorithms which aren't always optimal for batched GEMM. But cuBLAS does have batched GEMM APIs that are very efficient. Wait, let me try another approach entirely: TMA-based kernel with wgmma. But writing that from scratch without CUTLASS is extremely hard. Let me write a simpler kernel using the WMMA API and see how it performs. We can iterate and optimize. Here's my proposed kernel: ```cuda #include <cuda.h> #include <cuda_runtime.h> #include <mma.h> using namespace nvcuda; __global__ void bmm_tf32_kernel( const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int batch_size, int m, int n, int k) { // Block dimensions: 512 threads = 16 warps // Each block computes a 64x64 tile of C // Warp layout: 4x4 grid (4 in M, 4 in N) const int WARP_M = 4; const int WARP_N = 4; const int BM = 64; // Block tile M const int BN = 64; // Block tile N const int BK = 64; // Block tile K const int WM = 16; // Warp tile M const int WN = 16; // Warp tile N const int WK = 8; // Warp tile K (TF32) int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; int warp_m = warp_id / WARP_N; int warp_n = warp_id % WARP_N; int n_tiles = n / BN; int m_tiles = m / BM; int tiles_per_batch = m_tiles * n_tiles; int batch_idx = blockIdx.x / tiles_per_batch; int tile_idx = blockIdx.x % tiles_per_batch; int m_tile = tile_idx / n_tiles; int n_tile = tile_idx % n_tiles; int m_start = m_tile * BM; int n_start = n_tile * BN; const float* A_ptr = A + batch_idx * m * k; const float* B_ptr = B + batch_idx * k * n; float* C_ptr = C + batch_idx * m * n; // Shared memory layout // A_smem: [BM][BK] floats with padding for alignment // B_smem: [BK][BN] floats with padding __shared__ float A_smem[BM][BK + 4]; // pad to avoid bank conflicts __shared__ float B_smem[BK][BN + 4]; wmma::fragment<wmma::matrix_a, WM, WN, WK, wmma::precision::tf32, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, WM, WN, WK, wmma::precision::tf32, wmma::row_major> b_frag; wmma::fragment<wmma::accumulator, WM, WN, WK, float> acc_frag; wmma::fill_fragment(acc_frag, 0.0f); for (int k_outer = 0; k_outer < k; k_outer += BK) { // Load A to shared memory // A[batch][m_start + row_a][k_outer + col_a] for row_a in [0, BM), col_a in [0, BK) // Total: BM * BK = 4096 floats // 512 threads, each loads 8 floats int t = threadIdx.x; int a_row = t / (BK / 4); // 512 / 16 = 32... wait // Better: each thread loads 8 floats via float4 // 4096 floats / 512 threads = 8 floats per thread = 2 float4 int load_count = (BM * BK) / blockDim.x; // 8 floats per thread float4* A_smem_vec = (float4*)&A_smem[0][0]; const float4* A_ptr_vec = (const float4*)(A_ptr + m_start * k + k_outer); // Let me use a simpler approach: each thread loads 2 float4, covering 8 floats #pragma unroll for (int i = 0; i < 2; i++) { int idx = t + i * blockDim.x; // 0..1023 int row = idx / (BK / 4); // each row has BK/4 = 16 float4s int col4 = idx % (BK / 4); // which float4 in the row // But BM = 64, so row goes 0..63 // idx goes 0..1023, row = idx / 16 = 0..63 // col4 = idx % 16 = 0..15 if (row < BM) { float4 val = A_ptr_vec[row * (k / 4) + col4]; A_smem_vec[row * ((BK + 4) / 4) + col4] = val; } } // Load B to shared memory // B[batch][k_outer + row_b][n_start + col_b] for row_b in [0, BK), col_b in [0, BN) // Total: BK * BN = 4096 floats const float4* B_ptr_vec = (const float4*)(B_ptr + k_outer * n + n_start); float4* B_smem_vec = (float4*)&B_smem[0][0]; #pragma unroll for (int i = 0; i < 2; i++) { int idx = t + i * blockDim.x; int row = idx / (BN / 4); int col4 = idx % (BN / 4); if (row < BK) { float4 val = B_ptr_vec[row * (n / 4) + col4]; B_smem_vec[row * ((BN + 4) / 4) + col4] = val; } } __syncthreads(); // Compute using wmma #pragma unroll for (int k_inner = 0; k_inner < BK; k_inner += WK) { int a_row_offset = warp_m * WM; int b_col_offset = warp_n * WN; wmma::load_matrix_sync(a_frag, &A_smem[a_row_offset][k_inner], BK + 4); wmma::load_matrix_sync(b_frag, &B_smem[k_inner][b_col_offset], BN + 4); wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); } __syncthreads(); } // Store results int c_row_offset = m_start + warp_m * WM; int c_col_offset = n_start + warp_n * WN; wmma::store_matrix_sync(C_ptr + c_row_offset * n + c_col_offset, acc_frag, n, wmma::mem_row_major); } ``` Wait, I need to check if `wmma::precision::tf32` is the right type for TF32 fragments. Let me verify in the CUDA docs. In CUDA 11+, for TF32 WMMA: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> ``` Yes, `precision::tf32` is correct. And the K dimension is 8 for TF32 (vs 16 for FP16/BF16). Also, for `wmma::load_matrix_sync` with TF32, the pointer should be `float*` and the stride is in elements. One issue: TF32 fragments only support K=8, whereas my shared memory uses BK=64. The inner loop will iterate 64/8 = 8 times, which is fine. Another issue: when loading from shared memory with `wmma::load_matrix_sync`, the shared memory must not have bank conflicts for the access pattern used by the warp. The padding in my shared memory declarations `(BK + 4)` should help avoid conflicts. Actually, for A_smem[BM][BK+4], with BK=64, the row stride is 68 floats = 272 bytes. For a warp accessing 16 columns of floats (64 bytes), with 4-byte banks: - Bank for column j: (row_offset + j*4) / 4 % 32 = (row_offset/4 + j) % 32 - For j=0..15, consecutive columns hit consecutive banks. No conflicts! For B_smem[BK][BN+4], with BN=64, row stride is 68 floats = 272 bytes. - Bank for column j: same logic, consecutive banks for consecutive j. - But wmma for matrix_b accesses elements differently depending on layout. Actually, I realize `wmma::load_matrix_sync` might have special memory access patterns optimized by hardware, so mild bank conflicts are often not a problem. Let me think about whether there are issues with my vectorized loads to shared memory. For A global → shared: ```cpp float4* A_smem_vec = (float4*)&A_smem[0][0]; A_smem_vec[row * ((BK + 4) / 4) + col4] = val; ``` This writes directly to shared memory as float4. The row stride in float4 is (68/4) = 17 float4s. For a row of 64 floats = 16 float4s, we have padding of 1 float4 per row. This should make consecutive rows start at different bank alignments, helping avoid conflicts. But wait, the indexing: - A_smem_vec is a float4 pointer to the start of A_smem. - Element [row][col4] in float4 terms: row * 17 + col4. - For row=0, col4=0: offset 0 - For row=1, col4=0: offset 17 * 16 bytes = 272 bytes - Bank = 272/4 = 68 % 32 = 4 - So consecutive rows start at banks 0, 4, 8, 12, 16, 20, 24, 28, 0, 4, ... This seems fine. Actually, I want to use a simpler approach for the kernel to reduce bugs. Let me use `__nv_bfloat16` with K=16 and see if it passes the tolerance. If not, I'll switch to TF32. But I've already convinced myself that BF16 might fail the tolerance. Let me think once more... For random inputs A, B in [0, 1]: - torch.bmm(A, B) computes with TF32 precision - Our kernel with BF16 computes with 7-bit mantissa for operands - The expected values range from ~0 to ~1024 - Tolerance: max(1e-4, 1e-4 * expected) ≈ max(1e-4, 0.1) = 0.1 Wait, for expected=1024, 1e-4 * 1024 = 0.1024. So the tolerance is 0.1 for large values! And for expected=256, tolerance is 0.0256. For expected=512, tolerance is 0.0512. For expected=1024, tolerance is 0.1024. So the tolerance grows with the magnitude of the output! For random inputs, the maximum output is about 1024 (if all A[i]=1 and B[i]=1, but for uniform [0,1], it's about sum of 1024 products of U[0,1] values, which averages 256 and has a max around 600-700 at 3 sigma, but could theoretically reach 1024). Standard deviation of the sum is about sqrt(1024 * 1/36) = sqrt(28.4) ≈ 5.3 for products of independent U[0,1]. Wait, E[XY] = 1/4, Var(XY) = E[X^2 Y^2] - E[XY]^2 = (1/3)(1/3) - (1/4)^2 = 1/9 - 1/16 = 7/144 ≈ 0.0486. Std = sqrt(1024 * 0.0486) = sqrt(49.8) ≈ 7.1. Wait, I calculated this before. So most values are within 256 ± 21 (3 sigma). Error from BF16 should be well within tolerance. Actually, the key question is: what fraction of elements have error > 0.1? With std dev of output ~7.1, values above 300 are very rare (more than 6 sigma). Tolerance at 300 is 0.03. The BF16 error per element was estimated to have std dev ~0.07-0.2. So some elements might exceed 0.03. But for elements near 256 (the mean), tolerance is 0.0256. If error std is 0.07, then about 35% of elements would have |error| > 0.0256. Hmm, that's a lot. But let me recalculate the error more carefully. Actually, there's a crucial detail I missed: torch.bmm with TF32 also introduces errors! So we're not comparing to exact arithmetic, but to TF32 arithmetic. Both our kernel (if we use BF16) and the reference use lower precision than exact FP32. Wait no - the reference is `torch.bmm(A, B)` which might be using TF32. The correctness checker compares our output to the reference output. So if both introduce errors, but different ones, the difference could be large. For reference (TF32): - Each operand rounded to TF32: error ~ 2^-10 * value ≈ 0.001 - Product rounded to TF32: error ~ 0.001 - Accumulated in FP32 exactly For our kernel (BF16): - Each operand rounded to BF16: error ~ 2^-7 * value ≈ 0.008 - Product rounded to BF16: error ~ 0.008 - Accumulated in FP32 exactly The difference between BF16 and TF32 per product: ~0.007 in the operand, leading to product error difference of ~0.014. Sum of 1024 terms with error difference ~0.014 per term: - If signs are random, std of difference ~ 0.014 * sqrt(1024) / sqrt(3) ≈ 0.26 This could easily exceed the tolerance of 0.03 for most elements! So BF16 is likely to FAIL. We need TF32. Actually, wait. I just thought of something. What if torch.bmm does NOT use TF32? What if it's disabled? In PyTorch, `torch.backends.cuda.matmul.allow_tf32` defaults to True on Ampere/Hopper for matmul. But batched matmul might have a different code path. Actually, `torch.bmm` calls into cuBLAS where the TF32 policy in PyTorch is set globally. For FP32 input on H100, it should use TF32. To be safe, let me use `wmma::precision::tf32` which should match cuBLAS's precision. Alternatively, we could write a pure FP32 kernel using shared memory tiles without tensor cores. For these problem sizes, the theoretical time at 67 TFLOPS peak would be ~4.1 ms. But的实际 achievable might be 20-30 TFLOPS, giving ~9-14 ms. cuBLAS with TF32 should be around 1-2 ms. So pure FP32 won't be competitive. So TF32 it is. Let me refine my TF32 kernel. I'll use a more careful implementation. Actually, I realize there's another approach. Instead of wmma, I can write a kernel that calls the CUTLASS library functions. Let me check if CUTLASS headers are available. In PyTorch, CUTLASS headers are sometimes available. But in the `load_inline` context, I don't know if they'll be accessible. The CUDA include path for `load_inline` typically includes `/usr/local/cuda/include` and PyTorch's own include directories. CUTLASS might not be there. Another approach: use `cublasGemmStridedBatchedEx` with CUBLAS_COMPUTE_32F_FAST_TF32. This is exactly what torch.bmm does. But we could potentially call it with a better algorithm selection. Actually, let me check: in PyTorch, the TF32 heuristic is set globally. If we create our own cuBLAS handle, we could set the math mode to TF32 and potentially choose a better algorithm. But torch.bmm already uses the default cuBLAS heuristic which is very good. Wait, what if PyTorch doesn't use `cublasGemmStridedBatchedEx` but instead calls `cublasSgemmBatched` or loops over individual `cublasSgemm` calls? For some older versions or specific shapes, PyTorch might not use the most optimal path. But on H100 with recent PyTorch, it should definitely use the best available path. Let me try a completely different strategy. What if we write a kernel using a register-tiled approach for FP32 without tensor cores, but with very large tiles and aggressive unrolling? On H100, pure FP32 might actually be faster than expected if we can use the full SM. Let me calculate: H100 has 132 SMs, 67 TFLOPS FP32 = ~0.5 TFLOPS per SM. For 275 GFLOPs total, we need 275 / 67 = ~4.1 ms at peak. But peak is rarely achieved. 50% peak would be 8.2 ms. For cuBLAS TF32 at 300 TFLOPS effective: 0.9 ms. So pure FP32 is likely 5-10x slower. Not competitive. Let me commit to the TF32 WMMA kernel. Actually, I should also consider using multiple CTAs per SM via cooperative groups or just relying on high occupancy. One more thing: in Hopper, `wmma` might not be the fastest path. The `mma.sync.aligned.m64n8k16` PTX instruction is specifically for Hopper and is much more powerful. But `wmma` should compile down to the right instructions. Let me check if `wmma::precision::tf32` actually exists and is supported. From the CUDA C++ Programming Guide: - `precision::tf32` is supported on compute capability >= 80 (Ampere, Hopper). - Fragment dimensions: 16x16x8. Yes, it exists. Now, one practical concern: compiling with CUDA 12 on H100. The extension compilation should work. Let me also think about whether I should try a mixed approach. For example, what if I first convert A and B to TF32 in shared memory? Actually, they're already FP32 in global memory, and wmma loads them directly to the fragment where the hardware rounds to TF32. So no conversion needed. Wait, for `wmma::load_matrix_sync` with `precision::tf32`, do we pass float arrays or do we need to convert first? From the docs: For `precision::tf32`, the fragment stores data as `float`. The `load_matrix_sync` takes a pointer to the appropriate type (`float*` for TF32) and loads from memory. The hardware automatically interprets the float data as TF32 during the mma operation. So we can just load from float arrays directly. OK, let me write the kernel. But I realize my vectorized load approach using `float4` has an issue: the pitch of A in global memory is `k * sizeof(float) = 1024 * 4 = 4096 bytes`, which is divisible by 16. So `float4` loads are aligned. Similarly for B: pitch `n * sizeof(float) = 2048 * 4 = 8192 bytes`, divisible by 16. Good. For A_smem shared memory, with padding `BK + 4 = 68`, pitch is `68 * 4 = 272 bytes`, divisible by 16. So `float4` stores to shared memory are also aligned. For B_smem, same: `BN + 4 = 68`, pitch is 272 bytes, divisible by 16. Now for the `wmma::load_matrix_sync` calls: - For A: `&A_smem[warp_m * WM][k_inner]`, leading dimension `BK + 4` - For B: `&B_smem[k_inner][warp_n * WN]`, leading dimension `BN + 4` These should work since the pointer is aligned (k_inner=0,8,16,... and the row start is aligned to at least 16 bytes because `BK+4` is not a multiple of 4 in the dimension that matters... wait. The pointer `&A_smem[warp_m * WM][k_inner]` points to byte offset: `(warp_m * WM) * (BK + 4) * sizeof(float) + k_inner * sizeof(float)` For warp_m = 0..3, WM = 16: - warp_m * WM = 0, 16, 32, 48 - Offset = row * 272 + k_inner * 4 For row=0, k_inner=0: offset 0 (aligned) For row=16, k_inner=0: offset 16 * 272 = 4352 = 272 * 16 = 16 * 272. Is this divisible by 16? Yes, 272 = 16 * 17, so 4352 = 16 * 16 * 17. Aligned. For `&B_smem[k_inner][warp_n * WN]`: - Offset = k_inner * (BN + 4) * sizeof(float) + warp_n * WN * sizeof(float) - = k_inner * 272 + warp_n * 64 For k_inner=0, warp_n=0: 0 (aligned) For k_inner=8, warp_n=0: 8 * 272 = 2176 = 16 * 136. Aligned. For k_inner=0, warp_n=1: 64 = 16 * 4. Aligned. Great, alignment is fine. Now, there might be an issue with `wmma::store_matrix_sync`. The output pointer `C_ptr + c_row_offset * n + c_col_offset` has leading dimension `n`. For C, each row is `n * sizeof(float) = 2048 * 4 = 8192 bytes`, divisible by 16. And `c_col_offset = warp_n * 16`, also divisible by 16. So alignment is good. One more concern: the total number of blocks with 64x64 tiles is 32768. On H100 with 132 SMs, that's ~248 blocks per SM. At 4 concurrent blocks per SM, that's ~62 waves. This is plenty to keep all SMs busy. But the launch configuration might be tricky. I should use `dim3 grid(32768)` and `dim3 block(512)`. Actually, let me use a smaller tile to increase parallelism. What about 32x32 tiles with 128 threads (4 warps)? - Blocks: 128 * 16 * 64 = 131072 - Threads: 128 - Blocks per SM: min(2048/128, ...) = 16 - Waves: 131072 / (132 * 16) = 62 - Each block computes 32x32 = 1024 elements - But the overhead of so many blocks might be significant. I'll stick with 64x64 and 512 threads. Actually, let me try with 64x32 and 256 threads first. This gives: - Blocks: 128 * 8 * 64 = 65536 - Threads: 256 - Blocks per SM: 8 - Waves: 65536 / (132 * 8) = 62 - Each block: 64x32 = 2048 elements - A_smem: 64x64 = 16KB - B_smem: 64x32 = 8KB - Total: 24KB With 256 threads and 24KB shared, blocks per SM = min(2048/256, 228/24) = min(8, 9.5) = 8. 100% occupancy! And each warp handles 16x16: 2 warps in M, 4 warps in N? Wait: - Block tile M=64, N=32 - Warp tile M=16, N=16 - Warps in M: 64/16 = 4 - Warps in N: 32/16 = 2 - Total warps: 8 = 256 threads This works! And higher occupancy might help. Let me go with: - BM=64, BN=32, BK=64 - 8 warps = 256 threads Kernel code: ```cpp __global__ void bmm_tf32_kernel( const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int batch_size, int m, int n, int k) { const int BM = 64; const int BN = 32; const int BK = 64; const int WM = 16; const int WN = 16; const int WK = 8; int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; int warp_m = warp_id / 2; // 4 warps in M int warp_n = warp_id % 2; // 2 warps in N int n_tiles = n / BN; // 64 int m_tiles = m / BM; // 8 int tiles_per_batch = m_tiles * n_tiles; // 512 int batch_idx = blockIdx.x / tiles_per_batch; int tile_idx = blockIdx.x % tiles_per_batch; int m_tile = tile_idx / n_tiles; int n_tile = tile_idx % n_tiles; int m_start = m_tile * BM; int n_start = n_tile * BN; const float* A_ptr = A + (long long)batch_idx * m * k; const float* B_ptr = B + (long long)batch_idx * k * n; float* C_ptr = C + (long long)batch_idx * m * n; __shared__ float A_smem[BM][BK + 4]; __shared__ float B_smem[BK][BN + 4]; wmma::fragment<wmma::matrix_a, WM, WN, WK, wmma::precision::tf32, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, WM, WN, WK, wmma::precision::tf32, wmma::row_major> b_frag; wmma::fragment<wmma::accumulator, WM, WN, WK, float> acc_frag; wmma::fill_fragment(acc_frag, 0.0f); for (int k_outer = 0; k_outer < k; k_outer += BK) { // Load A to shared memory: A[m_start..m_start+BM, k_outer..k_outer+BK] // 64x64 = 4096 floats. 256 threads, each loads 16 floats = 4 float4 #pragma unroll for (int i = 0; i < 4; i++) { int t = threadIdx.x + i * blockDim.x; // 0..1023 int row = t / (BK / 4); // 0..63 (since 1024/16 = 64) int col4 = t % (BK / 4); // 0..15 if (row < BM) { const float4* A_global = (const float4*)(A_ptr + (m_start + row) * k + k_outer + col4 * 4); float4* A_shared = (float4*)(&A_smem[row][col4 * 4]); *A_shared = *A_global; } } // Load B to shared memory: B[k_outer..k_outer+BK, n_start..n_start+BN] // 64x32 = 2048 floats. 256 threads, each loads 8 floats = 2 float4 #pragma unroll for (int i = 0; i < 2; i++) { int t = threadIdx.x + i * blockDim.x; // 0..511 int row = t / (BN / 4); // 0..63 (512/8 = 64) int col4 = t % (BN / 4); // 0..7 if (row < BK) { const float4* B_global = (const float4*)(B_ptr + (k_outer + row) * n + n_start + col4 * 4); float4* B_shared = (float4*)(&B_smem[row][col4 * 4]); *B_shared = *B_global; } } __syncthreads(); #pragma unroll for (int k_inner = 0; k_inner < BK; k_inner += WK) { wmma::load_matrix_sync(a_frag, &A_smem[warp_m * WM][k_inner], BK + 4); wmma::load_matrix_sync(b_frag, &B_smem[k_inner][warp_n * WN], BN + 4); wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); } __syncthreads(); } // Store C int c_row = m_start + warp_m * WM; int c_col = n_start + warp_n * WN; wmma::store_matrix_sync(C_ptr + c_row * n + c_col, acc_frag, n, wmma::mem_row_major); } ``` Wait, I need to double check the A load indexing. Let me trace through: - t ranges from 0 to 1023 in steps of 256 (4 iterations total, covering all 1024 values) - Actually, threadIdx.x ranges 0..255 - i=0: t = 0..255 - i=1: t = 256..511 - i=2: t = 512..767 - i=3: t = 768..1023 For each t: - row = t / 16 (since BK/4 = 64/4 = 16) - col4 = t % 16 For t=0: row=0, col4=0 For t=1: row=0, col4=1 ... For t=15: row=0, col4=15 For t=16: row=1, col4=0 Each thread (with a fixed t) loads 1 float4 = 4 floats from A_global and stores to A_shared. Total floats loaded: 1024 threads * 4 floats = 4096 floats. Coverage: 64 rows * 64 cols = 4096. Good. For thread 0 in warp 0: - t=0, row=0, col4=0 - A_global address: A_ptr + (m_start + 0) * k + k_outer + 0*4 = A_ptr + m_start * k + k_outer - Loads A[m_start][k_outer..k_outer+3] Thread 1: - t=1, row=0, col4=1 - A_global: A_ptr + m_start * k + k_outer + 4 - Loads A[m_start][k_outer+4..k_outer+7] Threads 0..15 in the first warp all load from row m_start, consecutive columns. Perfect coalescing! Similarly for B: - t ranges 0..511 (2 iterations of 256) - row = t / 8 (BN/4 = 32/4 = 8) - col4 = t % 8 For t=0: row=0, col4=0 For t=7: row=0, col4=7 For t=8: row=1, col4=0 Total floats: 512 * 4 = 2048. Coverage: 64 rows * 32 cols = 2048. Thread 0 loads B[k_outer][n_start..n_start+3] Thread 1 loads B[k_outer][n_start+4..n_start+7] ... Thread 7 loads B[k_outer][n_start+28..n_start+31] Thread 8 loads B[k_outer+1][n_start..n_start+3] Again, threads 0..7 in a warp load from row k_outer, consecutive columns. Coalesced! This looks good. Now let me think about whether there are any bank conflicts when writing to shared memory. For A_smem writes: - A_smem[row][col4*4] to A_smem[row][col4*4+3] are 4 consecutive floats = 16 bytes. - Threads in a warp: each writes a `float4` = 16 bytes. - Thread 0 writes to address 0 (bank 0) - Thread 1 writes to address 16 (bank 4) - Thread 2 writes to address 32 (bank 8) - ... - Thread 7 writes to address 112 (bank 28) - Thread 8 writes to address 128 (bank 0) since 128/4 = 32 % 32 = 0 Wait, thread 8 writes to row=0, col4=8, which is address 8*4*4 = 128 bytes = bank 0. Thread 0 also writes to bank 0 (address 0). But thread 0 and thread 8 are in the same warp! They both try to write to bank 0 but different addresses (0 vs 128). On modern GPUs, there are no bank conflicts for `float4` (128-bit) accesses as long as the accesses are aligned and contiguous? Actually, in shared memory, a 128-bit access from each thread is split into 4 x 32-bit accesses, each going to a different bank. Or is it treated as one transaction? Actually, for `float4` stores to shared memory, each thread issues 4 x 32-bit stores. The first store from thread 0 goes to bank 0, thread 1 to bank 4, etc. The second store from thread 0 goes to bank 1, thread 1 to bank 5, etc. With 32 threads, some might conflict. But actually, shared memory bank conflicts happen at the 32-bit level for compute capability >= 70. Each 32-bit word maps to a bank. For `float4`, thread t accesses banks (t*4 + 0, t*4 + 1, t*4 + 2, t*4 + 3) mod 32. For threads 0..31: - Thread 0: banks 0,1,2,3 - Thread 1: banks 4,5,6,7 - ... - Thread 7: banks 28,29,30,31 - Thread 8: banks 0,1,2,3 → CONFLICT with thread 0! Yes, there would be bank conflicts for the shared memory writes. But this is only if threads 0 and 8 execute their stores in the same cycle. Actually, with 128-bit accesses, I believe the hardware can handle them without bank conflicts if they're properly aligned, because the access is atomic at the 128-bit level. Wait, no. Shared memory bank conflicts are at the 32-bit level. For a `float4` store, there are 4 separate 32-bit stores, and each can have conflicts. However, since threads 0-7 write to banks 0-31 (each thread writes 4 consecutive banks), and threads 8-15 also write to banks 0-31, there is a 2-way conflict for each bank (threads t and t+8 both write to banks 4t through 4t+3, but shifted). Actually, the exact mapping depends on the starting address. Let me be more careful. A_smem[row][col] where col = col4 * 4. For a fixed row, the address of A_smem[row][col] in bytes: `row * (BK + 4) * 4 + col * 4 = row * 272 + col * 4` Bank = (address / 4) % 32 = (row * 68 + col) % 32 For row=0, col = col4 * 4: Bank = (0 + col4 * 4) % 32 = col4 * 4 % 32 For thread t=0..15 (first 16 threads in warp): - col4 = t % 16 = 0..15 - Bank = (t % 16) * 4 % 32 t=0: col4=0, bank=0 t=1: col4=1, bank=4 t=2: col4=2, bank=8 ... t=7: col4=7, bank=28 t=8: col4=8, bank=0 (32 % 32 = 0) ← CONFLICT with t=0 t=9: col4=9, bank=4 (36 % 32 = 4) ← CONFLICT with t=1 ... t=15: col4=15, bank=60%32=28 ← CONFLICT with t=7 So threads 0-7 and 8-15 have bank conflicts when writing float4! This is a 2-way conflict. To fix this, we need padding that ensures different col4 values don't map to the same bank. We want: `(row * 68 + col4 * 4) % 32` to be unique for col4=0..15, or at least not have duplicates in the same 16-thread group. Currently: (col4 * 4) % 32 for col4=0..15 gives: 0,4,8,12,16,20,24,28,0,4,8,12,16,20,24,28. Duplicates at col4=0,8; 1,9; etc. If we change the row pad such that the effective columns per row is not 68 but something else. For per-row offset 272 bytes = 68 words. The issue is that 4 * 16 = 64, and 64 % 32 = 0 (since 64 is a multiple of 32). If instead the row stride in words is coprime with 32, then consecutive col4 values would hit different banks. But col4 * 4 means we step by 4 in word index. We need (pad + 16*4) % 32 != 0, i.e., (pad + 64) % 32 != 0, which means pad % 32 != 0. Our current pad is 4 words, so total words per row is 68. 68 % 32 = 4. Let's check: Bank = (row * 68 + col4 * 4) % 32 = (row * 4 + col4 * 4) % 32 = 4 * (row + col4) % 32 For row=0: bank = 4 * col4 % 32 This still gives: 0,4,8,12,16,20,24,28,0,4,8,12,16,20,24,28. Same duplicates! The issue is that 4 * col4 mod 32 has period 8, so col4 and col4+8 give the same bank. To avoid this, we need the constant multiplier to be odd (coprime with 32). If the word step was 1, then col4=0..15 gives 0..15, no duplicates. But our columns are in groups of 4. Actually, if each thread writes a `float4`, the 4 elements map to banks base, base+1, base+2, base+3. For thread t in a warp, the base bank is: base(t) = (row * stride + t * 4) % 32 We want base(t) and base(t+8) to differ. base(t+8) = (row * stride + (t+8) * 4) % 32 = base(t) + 32 % 32 = base(t). So with 32 threads in a warp and float4 accesses, threads t and t+8 will ALWAYS have the same base bank, causing conflicts if they write to the same row at the same time. Wait, but threads t=0..15 might not all be in the same warp processing the same row. In our load code, threads 0..15 all have row=0 (since t/16 = 0 for t=0..15). So warp 0 has threads 0..31, and threads 0..15 all write to row=0, while threads 16..31 write to row=1. So threads 0..15 in warp 0 write to row=0, and within these 16 threads, t and t+8 write to the same base bank. Actually, the issue is: - t=0: row=0, col4=0, banks 0,1,2,3 - t=8: row=0, col4=8, banks 32%32=0,1,2,3 So both threads 0 and 8 write to banks 0,1,2,3. This is a 2-way bank conflict. Similarly: - t=1: row=0, col4=1, banks 4,5,6,7 - t=9: row=0, col4=9, banks 4,5,6,7 This is indeed a 2-way conflict for each group of 4 banks. To fix this, we need either: 1. Use a different access pattern so that no two threads in a warp access the same bank for any element. 2. Accept the 2-way conflict (it doubles the shared memory store time, but stores might not be the bottleneck). Is a 2-way conflict terrible? It means each warp takes 2x as long for the shared memory store. But our kernel is compute-bound, so a 2x slowdown in shared memory stores might be acceptable. Alternatively, we can change how we distribute the work. If instead of having 16 threads per row, we have fewer, or if we transpose the shared memory layout. Let me try a different approach: have each thread load 1 float (not float4), so that consecutive threads access consecutive floats. Then: - Thread t loads A_smem[row][col] where col = t % 64 (for some row distribution) - Bank = (row * 68 + col) % 32 = (row * 4 + col) % 32 - For 32 threads in a warp with consecutive cols: banks are (base, base+1, ..., base+31) % 32 = all 32 distinct banks. No conflict! But loading 1 float per thread means 4x more instructions, potentially hurting performance. However, global memory loads would still be coalesced (each thread loads 4 bytes from consecutive addresses). Let me try this: each thread loads a single float. Then in the store to shared memory, there are no bank conflicts. With BM=64, BK=64, total elements = 4096. 256 threads, each loads 4096/256 = 16 floats. That's 16 loads per thread. With float loads instead of float4 loads, we have more instructions but no bank conflicts. Alternatively, we could use `__syncthreads()` and let the hardware handle it; a 2-way conflict might not be too bad. Actually, for A_smem, I realize there's another issue: threads 16..31 have row=1. So warp 0 has: - threads 0..15: row=0 - threads 16..31: row=1 For row=1, banks = (1 * 68 + col4 * 4) % 32 = (4 + col4 * 4) % 32. Thread 16 (col4=0): banks 4,5,6,7 Thread 24 (col4=8): banks 36%32=4,5,6,7 So conflicts also happen within threads 16..31 (row=1). But threads 0..15 and 16..31 don't conflict with each other because they write to different rows and thus different banks (except at boundaries). Thread 0 writes to banks 0,1,2,3; thread 16 writes to banks 4,5,6,7. No conflict between these groups. So the conflict is only 2-way within each group of 16 threads. I think this is acceptable. Let me keep the float4 loads for efficiency and accept the mild bank conflicts. For B_smem: BK=64, BN=32. Total elements = 2048. 256 threads * 2 loads * 4 floats = 2048. t = 0..511 row = t / 8 (BN/4 = 8) col4 = t % 8 warp 0 threads 0..31: - t=0..7: row=0 - t=8..15: row=1 - t=16..23: row=2 - t=24..31: row=3 For row=0, col4=0..7: Bank = (0 * 68 + col4 * 4) % 32 = col4 * 4 % 32 = 0,4,8,12,16,20,24,28. All distinct! For row=1, col4=0..7: Bank = (1 * 68 + col4 * 4) % 32 = (4 + col4*4) % 32 = 4,8,12,16,20,24,28,0. All distinct! For row=2: 8,12,16,20,24,28,0,4. All distinct! For row=3: 12,16,20,24,28,0,4,8. All distinct! Wait, BN + 4 = 36 words per row. 36 * 4 = 144 bytes per row. Bank = (row * 36 + col4 * 4) % 32 = (4*row + 4*col4) % 32 = 4*(row + col4) % 32. For row=0, col4=0..7: 0,4,8,12,16,20,24,28. Distinct! For row=1, col4=0..7: 4,8,12,16,20,24,28,0. Distinct! For row=2, col4=0..7: 8,12,16,20,24,28,0,4. Distinct! For row=3, col4=0..7: 12,16,20,24,28,0,4,8. Distinct! And threads in a warp: - t=0: row=0, col4=0, banks 0,1,2,3 - t=1: row=0, col4=1, banks 4,5,6,7 - ... - t=7: row=0, col4=7, banks 28,29,30,31 - t=8: row=1, col4=0, banks 4,5,6,7 Wait, t=1 writes to banks 4,5,6,7 and t=8 writes to banks 4,5,6,7! CONFLICT between threads 1 and 8! Oh no! So for B_smem, there are conflicts too. Bank for B_smem[row][col] = (row * 36 + col) % 32. Thread t=1: row=0, col=4..7 → banks 4,5,6,7. Thread t=8: row=1, col=0..3 → banks 36%32=4, 37%32=5, 38%32=6, 39%32=7. Yes, threads 1 and 8 conflict on all 4 banks! Hmm, this is worse than A_smem. We have a 2-way conflict for B_smem as well. Let me think about how to fix B_smem. We need the starting bank for each thread's float4 to be distinct within the warp. With row_stride = 36 words = 144 bytes: base_bank(t) = (row(t) * 36 + col4(t) * 4) % 32 = 4 * (row(t) * 9 + col4(t)) % 32 Wait, 36 % 32 = 4. So row(t) * 36 ≡ row(t) * 4 (mod 32). base_bank = (4 * row + 4 * col4) % 32 = 4 * (row + col4) % 32. row = t / 8, col4 = t % 8. base_bank = 4 * (t/8 + t%8) % 32 = 4 * f(t) % 32. For t=0..31: t=0: 4*(0+0)=0 t=1: 4*(0+1)=4 t=2: 4*(0+2)=8 t=3: 4*(0+3)=12 t=4: 4*(0+4)=16 t=5: 4*(0+5)=20 t=6: 4*(0+6)=24 t=7: 4*(0+7)=28 t=8: 4*(1+0)=4 ← duplicate with t=1 t=9: 4*(1+1)=8 ← duplicate with t=2 ... Yes, duplicates at t and t+7 (when row increases by 1 and col4 wraps). Wait, t+7: row increases by ~1, col4 decreases by ~1. So row+col4 might stay similar. To avoid this, we need the row stride in words to be odd (or at least such that 4*(row*stride_w + col4) gives unique values mod 32 for all t in 0..31). Actually, since we're multiplying by 4, the base bank is always a multiple of 4. There are only 8 distinct base banks: 0,4,8,12,16,20,24,28. With 32 threads, some will definitely share base banks. But the question is whether threads with the same base bank execute their stores in the same cycle. Since each thread writes a float4 (4 consecutive banks), if two threads have the same base bank, they write to exactly the same 4 banks. This is a true bank conflict. Can we fix this by having each thread write only 1 float instead of 4? Then: - Thread t writes to bank = (row * 36 + col) % 32 - col = (t % 8) * 4 + offset, where offset = 0 for the first of 4 stores. For t=0: row=0, col=0, bank=0 For t=1: row=0, col=1, bank=1 ... For t=31: row=3, col=7, bank=(108+7)%32=115%32=19 All banks 0..31 might be covered without conflict if we write 1 float at a time. But we need 4 writes per thread. Alternatively, what if we change the B_smem layout to be column-major or transposed? If B_smem is stored as [BN][BK] (transposed), then each thread loading from B_global writes to a different pattern in shared memory. But wmma::load_matrix_sync for matrix_b expects row-major or col-major based on the fragment type. Since we use row_major for B, B_smem should match the logical KxN layout of B. Wait, another idea: what if we just accept the bank conflicts? Modern GPUs have mechanisms to handle some conflicts efficiently, and our kernel is compute-bound anyway. A 2-way conflict might add only a few cycles to the shared memory operations, which are hidden by the compute. I think it's fine to proceed with the float4 loads and accept the mild bank conflicts. The critical path is the wmma operations, not the shared memory stores. Actually, one more optimization: we can pipeline by using double-buffered shared memory. While the current K_TILE is being computed, we load the next K_TILE. This hides the global memory latency behind the computation. For double buffering with K_TILE=64: - A_smem: 2 * 64 * 64 * 4 = 32 KB - B_smem: 2 * 64 * 32 * 4 = 16 KB - Total: 48 KB - Blocks per SM: min(2048/256, 228/48) = min(8, 4.75) = 4 Occupancy drops to 50%. With single-buffer: - Total: 24 KB - Blocks per SM: min(8, 9.5) = 8 - Occupancy: 100% Since our working set for each block is small and we have many blocks, single-buffer with high occupancy is probably better. The memory latency can be hidden by having other warps/blocks execute while one waits. So I'll keep single-buffer. Now let me think about whether there are other optimization opportunities. One thing: using `__launch_bounds__(256)` to help the compiler optimize register usage. Another: using `#pragma unroll` on the k_inner loop. Also, I should make sure to use the correct absolute indexing with long long for the pointers, since the batch matrices might be large. In our case, batch_idx * m * k can be up to 128 * 512 * 1024 = 67M, which fits in int. But better to use long long just in case. Wait, one more thing: I should check if 32768 blocks or 65536 blocks is enough for the H100. 65536 is the max gridDim.x on some GPUs, but modern GPUs support much larger, up to 2^31-1. The H100 definitely supports large grids. 32768 blocks is fine. OK, I think the kernel is ready. Let me also add a host wrapper function that launches the kernel and handles the grid/block dimensions. The Python file will look like: ```python import torch from torch.utils.cpp_extension import load_inline cuda_source = """ #include <cuda.h> #include <cuda_runtime.h> #include <mma.h> using namespace nvcuda; __global__ void ... // kernel torch::Tensor bmm_forward(torch::Tensor A, torch::Tensor B) { int batch_size = A.size(0); int m = A.size(1); int k = A.size(2); int n = B.size(2); auto C = torch::empty({batch_size, m, n}, A.options()); int n_tiles = n / 32; int m_tiles = m / 64; int tiles_per_batch = m_tiles * n_tiles; int total_blocks = batch_size * tiles_per_batch; dim3 grid(total_blocks); dim3 block(256); bmm_tf32_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>( A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), batch_size, m, n, k); return C; } """ cpp_sources = "..." # Usually empty or just includes for pybind functions = ["bmm_forward"] module = load_inline( name="bmm_cuda", cpp_sources=cpp_sources, cuda_sources=cuda_source, functions=functions, extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_90"], ) ``` Wait, I need to be careful about the `cpp_sources` and `cuda_sources` parameters. Since we need to include `<mma.h>` which is a CUDA header, the kernel code should be in `cuda_sources`. Any host-side Torch binding code that calls the kernel can also be in `cuda_sources` (as .cu files can contain host code). Or we can put the kernel in `cuda_sources` and the wrapper in `cpp_sources`. Actually, for `load_inline`, `cpp_sources` and `cuda_sources` are separate. The kernel should go in `cuda_sources`, and the host wrapper should probably also go in `cuda_sources` since it needs to call the CUDA kernel. PyTorch headers can be included in `cuda_sources` too. Let me structure it properly. One more thing: I should make sure we're compiling for sm_90 (Hopper). The extra CUDA flags should include `-gencode=arch=compute_90,code=sm_90`. Actually, if the CUDA compiler defaults to a lower arch, wmma might still work but not use Hopper-specific optimizations. We definitely want sm_90. Let me also add `"-maxrregcount=128"` or similar to control register usage, but the compiler should manage this. Actually, I want to check if `wmma::precision::tf32` is supported with the default nvcc that PyTorch uses. If PyTorch is using CUDA 11.8 or later, it should support Hopper and TF32. For H100, we need CUDA 11.8+ or CUDA 12.x. The kernel should compile fine. Let me also think about edge cases and correctness: - The inputs A and B must be contiguous in the relevant dimensions. - The shapes must be divisible by our tile sizes: batch=128 (irrelevant), m=512 (divisible by 64), k=1024 (divisible by 64), n=2048 (divisible by 32). - Our kernel assumes these divisibility conditions. Since the problem gives fixed shapes, this is fine. Wait, I need to handle the case where the inputs might not be exactly these shapes. The problem says the shapes are fixed, but it's good practice to either assert or handle generically. For simplicity, I'll assert or just assume the shapes. Actually, the ModelNew forward should work on any inputs of the right shape, but for performance and simplicity, I'll write a kernel that works for the specific sizes. One more optimization: we can precompute some indices and use them. But the compiler should optimize this. Let me also consider whether to use an even larger tile size. With BM=128 and BN=64: - A_smem: 128x64 = 32KB - B_smem: 64x64 = 16KB - Total: 48KB - Threads: 512 (32 warps: 8x4 grid for 128x64 tile with 16x16 warps) - Blocks: 128 * 4 * 32 = 16384 - Blocks per SM: min(2048/512, 228/48) = min(4, 4.75) = 4 Occupancy is lower (50% vs 100%), but each block does more work. Fewer blocks might mean less overhead. Let me try both. Actually, let me keep 64x32 for now and see the performance. There's one more consideration. TF32 wmma fragments use the type `precision::tf32`, but I need to make sure the `wmma::fill_fragment` and `wmma::store_matrix_sync` work correctly with float accumulators. Yes: ```cpp wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_frag; wmma::fill_fragment(acc_frag, 0.0f); wmma::store_matrix_sync(C_ptr + offset, acc_frag, n, wmma::mem_row_major); ``` This should be correct. One issue: when we store via `wmma::store_matrix_sync` to global memory, the hardware handles the store. But the output C array is in FP32, which is what we want. Now let me think about whether we need any additional synchronization between blocks. No, each block computes a disjoint tile of C. Another subtlety: if A or B are not contiguous in memory, our pointer arithmetic might be wrong. The problem says they're PyTorch tensors, which are typically contiguous when created with `torch.rand`. And `torch.bmm` requires specific strides anyway. Actually, `torch.bmm` works with non-contiguous inputs too. For correctness, we should probably check if A and B are contiguous, and if not, make copies. But since `get_inputs()` returns contiguous tensors, and the task says produce numerically equivalent outputs, we should handle it. I'll add a check in the PyTorch wrapper: ```python if not A.is_contiguous(): A = A.contiguous() if not B.is_contiguous(): B = B = B.contiguous() ``` OK, I think the kernel is ready to be written. Let me draft it and compile. Actually, one more thing: I realize `wmma::precision::tf32` might trigger some compiler warnings or might need a specific CUDA version. As a backup, if TF32 doesn't compile or work, I could fall back to pure FP32 tiles. But let me try TF32 first. Let me also check if there are any issues with the `float4` loads. In CUDA, `float4` requires 16-byte alignment. For A: - A_ptr + (m_start + row) * k + k_outer + col4 * 4 - (m_start + row) * k is divisible by 4 (k=1024) - k_outer is divisible by 4 (multiples of 64) - col4 * 4 is divisible by 4 - So each address is divisible by 16 for `float4` alignment? - Yes: all terms are multiples of 4, and the starting offset is already aligned to 16 bytes because A is allocated with 256-byte alignment by PyTorch. For B: - B_ptr + (k_outer + row) * n + n_start + col4 * 4 - (k_outer + row) * n: n=2048, divisible by 4 - n_start: multiples of 32, divisible by 4 - col4 * 4: divisible by 4 - Total: divisible by 16. Alignment is good. One more subtlety: the shared memory arrays `A_smem` and `B_smem` are declared as 2D arrays with padding. For `float4` accesses to shared memory, the starting address must also be aligned to 16 bytes. For A_smem: - A_smem[row][col4*4]: the address is `&A_smem[0][0] + row * (BK+4) + col4*4` - Row 0, col4=0: base address. Aligned because the array starts at shared memory base which is aligned. - Row 1, col4=0: offset = 1 * 68 * sizeof(float) = 272 bytes = 16 * 17. Divisible by 16. ✓ - Row 2: 544 = 16 * 34. Divisible by 16. ✓ - In general: row * 68 * 4 = row * 272. 272 = 16 * 17. So every row is aligned to 16 bytes. ✓ For B_smem: - B_smem[row][col4*4]: offset = row * (BN+4) * sizeof(float) + col4*4 - BN+4 = 36. Row offset = row * 36 * 4 = row * 144. - 144 = 16 * 9. So every row is aligned to 16 bytes. ✓ Great, shared memory alignment is also good. Now, please let me think about potential numerical issues one more time. Both our kernel and torch.bmm use TF32 precision (assuming PyTorch default). The reference implementation computes `torch.bmm(A, B)` which on H100 with default settings uses TF32. Our kernel also uses TF32 via WMMA. So the numerical results should be very close, likely well within 1e-4 tolerance. However, there's a subtlety: torch.bmm through cuBLAS might use a different algorithm or ordering of operations than our WMMA kernel. The associativity of floating-point addition means different accumulation orders can produce different results. TF32 has even more pronounced differences. For example, cuBLAS might use a divide-and-conquer approach or split K, while our kernel accumulates sequentially in K. This could cause differences beyond 1e-4. One way to reduce this: ensure we're doing the exact same sequence of operations, or use a kernel that matches cuBLAS's approach as closely as possible. Actually, for a simple blocked GEMM, both cuBLAS and our kernel should accumulate in similar order: each output element accumulates K values in sequence. The difference might come from: - How the K dimension is tiled (e.g., cuBLAS might use larger K tiles with different accumulation order) - Whether cuBLAS uses split-K (parallelizing the K reduction) For our specific problem sizes (K=1024), split-K might or might not be used by cuBLAS. It's more common for very large K. I think the sequential accumulation within each thread/warp should match closely enough. But it's something to be aware of. Let me also consider using `float` accumulator initialized to 0, and the wmma operations. The WMMA accumulator is a set of 16x16 floats, and multiple mma_sync calls accumulate into it in order. Actually, one cool feature: we can use block-level synchronization inside the kernel to enable split-K, where multiple blocks compute partial sums for the same output tile and then reduce. But this adds complexity. For K=1024, split-K is probably unnecessary and could hurt more than help due to atomic operations or extra synchronization. OK, I'm going to write the kernel and test it. Wait, I just realized I need to think about one more thing. The `wmma::load_matrix_sync` for row-major matrix_a expects the shared memory to contain the elements in row-major order, and the leading dimension parameter tells it the pitch between rows. But WMMA fragments use a specific internal layout that maps threads to elements in a complex way. The key question is: when I call `wmma::load_matrix_sync(a_frag, ptr, ld)`, does it correctly load a 16x8 tile (for TF32) from the shared memory starting at `ptr` with row stride `ld`? Yes, according to the CUDA docs, `load_matrix_sync` handles the remapping from the dense row-major layout in memory to the fragment's internal distributed layout across the warp. For matrix_a with row_major layout: - The memory pointed to by `ptr` should contain a MxK matrix in row-major order. - `ld` is the leading dimension (number of columns, i.e., the row stride). - For our case: ptr = `&A_smem[warp_m * 16][k_inner]`, ld = `BK + 4 = 68`. Wait, is `ld` the number of elements or bytes? The docs say "leading dimension", and for row-major it's the number of columns (in elements). Yes, it's in elements, not bytes. For A_smem, each row has BK+4 = 68 elements. So ld = 68. For B_smem, each row has BN+4 = 36 elements. So ld = 36. Wait, B_smem is [BK][BN+4], and we're loading a KxN tile from it (or actually, for matrix_b with row_major, it's K rows by N columns). So ld should be BN+4 = 36. Actually, for the B fragment in matrix_b, the logical dimensions are KxN, and it's stored in row-major format. So the leading dimension is the number of columns = N = BN (for our tile), but since we have padding, it's BN+4. So: ```cpp wmma::load_matrix_sync(a_frag, &A_smem[warp_m * WM][k_inner], BK + 4); wmma::load_matrix_sync(b_frag, &B_smem[k_inner][warp_n * WN], BN + 4); ``` Hmm wait, `&B_smem[k_inner][warp_n * WN]` - is this valid C++ syntax for a 2D array? Yes. Actually, in CUDA kernel code, `B_smem` is declared as `__shared__ float B_smem[BK][BN + 4];`. So `&B_smem[k_inner][warp_n * WN]` gives the address of element [k_inner][warp_n * WN]. And the leading dimension is BN+4. This is correct. Similarly for A: `&A_smem[warp_m * WM][k_inner]` with leading dimension BK+4. Actually wait, I want to double check the fragment dimensions. For TF32: - matrix_a: 16 x 8 (M x K) - matrix_b: 8 x 16 (K x N) ? Or 16 x 16? Let me verify in the CUDA docs. For `mma.sync.aligned.m16n16k8`: - A is MxK = 16x8 - B is KxN = 8x16 - C is MxN = 16x16 So: - `wmma::fragment<matrix_a, 16, 16, 8, ...>` represents A as 16x8 - `wmma::fragment<matrix_b, 16, 16, 8, ...>` represents B as 8x16 For the B fragment, the layout applies to the 8x16 matrix. If we use row_major, the memory should be an 8x16 matrix in row-major order. In our shared memory B_smem[k_inner][warp_n * WN], we're accessing starting at row k_inner, column warp_n * WN. The memory region from there is a BK x BN submatrix (64x32), but we're only loading an 8x16 tile from it. With row_major layout and ld = BN+4 = 36: - wmma reads B[k_inner + i][warp_n * WN + j] for i in [0,8), j in [0,16). - Address = &B_smem[k_inner][warp_n * WN] + i * 36 + j. - This is correct row-major access. For A fragment, row_major, ld = BK+4 = 68: - wmma reads A[warp_m * WM + i][k_inner + j] for i in [0,16), j in [0,8). - Address = &A_smem[warp_m * WM][k_inner] + i * 68 + j. - This is correct. Great. Now for the store: ```cpp wmma::store_matrix_sync(C_ptr + c_row * n + c_col, acc_frag, n, wmma::mem_row_major); ``` This stores a 16x16 tile starting at C_ptr + c_row * n + c_col, with leading dimension n. And we use `wmma::mem_row_major` because C is row-major. This seems correct. One thing: C_ptr points to the start of the batch's C matrix. c_row = m_start + warp_m * WM, c_col = n_start + warp_n * WN. So the absolute address is correct. Now I should also handle the case where c_row * n + c_col might overflow int. In our case: - c_row <= 512, n = 2048, c_col <= 2048 - c_row * n <= 512 * 2048 = 1,048,576 - Plus c_col: max ~1,050,624 - Batch offset: max ~128 * 512 * 2048 = 134,217,728 This fits in a 32-bit signed int (max 2,147,483,647). But I'll use long long for safety. Actually, for the kernel launch, the total number of blocks is batch_size * m_tiles * n_tiles = 128 * 8 * 64 = 65536. This fits in a 32-bit int. Alright, let me write the complete code and compile it. Wait, one more optimization thought. Can I increase block tile size to BM=128, BN=128, BK=16 with a different warp arrangement? For BM=128, BN=128: - Warps: 8x8 = 64 warps = 2048 threads. That's the max per block. - A_smem for BK=16: 128x16 = 2KB - B_smem for BK=16: 16x128 = 2KB - Total: 4KB. Very small! - Inner loop: 1 iteration (since BK=16 = WK*2 for TF32). Wait, K=1024, so outer loop has 64 iterations. - Each outer iteration: load A[128][16] and B[16][128], then 2 wmma iterations per warp. With 2048 threads per block: - A load: 128*16 = 2048 floats. Each thread loads 1 float. That's slow. - Or each thread loads 1 float2 = 2 floats, so 1024 active threads. But we have 2048 threads... Actually, with 2048 threads, each could load 1 float, covering 2048 elements. A[128][16] has exactly 2048 elements! Perfect coalescing if arranged properly. But having 2048 threads per block means only 1 block per SM (thread-limited). Occupancy is very low. We'd need many blocks (128 * 4 * 16 = 8192 blocks), giving 8192/132 = 62 blocks per SM, but only 1 concurrent. So 62 waves. The kernel might still be fast due to high utilization per block, but low occupancy means less latency hiding. With 512 threads per block and BM=128, BN=64: - Warps: 8x4 = 32 warps = 1024 threads. Hmm, that's 1024 threads. - Or arrange as 8x2 = 16 warps = 512 threads, computing 128x32. - Then blocks = 128 * 4 * 64 = 32768. With 512 threads per block (16 warps), blocks per SM = 4 (thread limited). Total waves = 32768 / (132 * 4) = 62. A_smem: 128x32 = 16KB, B_smem: 32x32 = 4KB, total = 20KB per K_TILE. With K_TILE=32, shared = 20KB. Blocks per SM = min(4, 228/20=11) = 4. This seems reasonable. But honestly, it's getting complicated with all these trade-offs. Let me just write one kernel with 64x32 tiles and 256 threads, compile it, and see if it passes correctness. If it does but is slow, I'll tune. Actually, wait. I want to try a slightly different approach: use a smaller BK to reduce shared memory and allow more computation reuse. With BK=16, the shared memory is tiny, and we can load everything very quickly. But then we have many more outer iterations. With BK=16 and K=1024, outer loop = 64 iterations. Each iteration loads A[64][16] and B[16][32]. The loads are small but frequent. I think BK=32 or 64 is better. OK, I'll go with the 64x32 tile, 256 threads, BK=64 kernel. Let me write the code now. One more thought: I should verify that `wmma::precision::tf32` is available in the nvcc version used by PyTorch. If the PyTorch installation is compiled with CUDA 11.0-11.7 for some reason, TF32 wmma might not be available. But H100 definitely needs CUDA 11.8+, so it should be fine. Another thought: instead of using `load_inline`, I could put the host wrapper in `cpp_sources` and the device code in `cuda_sources`. But `load_inline` actually treats them differently - `cuda_sources` are compiled with nvcc as `.cu` files, and `cpp_sources` are compiled as `.cpp` files. The wrapper function that creates the torch::Tensor and calls the kernel needs to be in a `.cu` file since it calls a `__global__` function. So I'll put everything in `cuda_sources`. Let me draft the kernel:

Two errors: 1. `bmm_forward` needs to be declared in a header visible to the C++ side for pybind11 wrapping. I need to provide a declaration in the cpp_sources or use a different approach. Actually with `load_inline`, when cpp_sources is empty but functions is specified, it's supposed to work if the function is in the cuda source. But there seems to be a scoping issue. Actually, I need to put the function declaration in cpp_sources or use `extern "C"` or something. But actually, the standard way is to put C++ API declarations in `cpp_sources` and the CUDA implementation in `cuda_sources`. But that requires splitting... Actually for `load_inline` with `functions=...`, if we only have `cuda_sources`, the function should be accessible. But the main.cpp generated by pybind11 can't see it because it's in a .cu file that may not be linked properly, or because it's not in a header. Wait, another approach: I can put BOTH the kernel and the wrapper in cpp_sources if I use `.cu` extension, but that won't compile as CUDA... Actually, looking at the PyTorch docs, for `load_inline`: - `cpp_sources`: C++ source code containing the host-side code and the pybind11 declarations/wrappers - `cuda_sources`: CUDA source code containing the kernels But the wrapper function needs to be declared in the cpp_sources if we use `functions=[...]`. Actually no - `load_inline` should handle this. But there seems to be a compilation order issue. Let me try putting the wrapper in cpp_sources and the kernel declaration also in cpp_sources, with the kernel implementation in cuda_sources. Also need to fix `at::cuda::getCurrentCUDAStream()` - maybe use `c10::cuda::getCurrentCUDAStream()`. Let me restructure: put function declarations in cpp_sources, kernel implementation in cuda_sources. Wait, actually for load_inline, if we specify `functions=['bmm_forward']`, pybind11 tries to find `bmm_forward` in the scope. But if the only place it's defined is in a .cu file which gets compiled to a different object, and if there's no prototype in the .cpp file, the C++ compiler won't know about it. Let me provide a declaration in `cpp_sources`: ```cpp #include <torch/extension.h> torch::Tensor bmm_forward(torch::Tensor A, torch::Tensor B); ``` And everything else in `cuda_sources`. Or even simpler, put the whole thing in cpp_sources but compile with nvcc by having the kernel as a string and using `-x cu` flag... No, load_inline handles file extensions. Let me try: - cpp_sources: just the forward declaration of bmm_forward - cuda_sources: everything else Also let me use `c10::cuda::getCurrentCUDAStream()` instead of `at::cuda::getCurrentCUDAStream()`.

Let me try using cudaStreamLegacy or just pass 0 for the default stream. Alternatively I could use `at::cuda::getCurrentCUDAStream` with a proper include. But for simplicity, let me just pass `0` which means the default stream. But first, let me think of a simpler approach. Actually, I think the task is asking for fp32. And since this is H100 with 989 TFLOPS in TF32 Tensor Core, using TF32 would be a good optimization. But if they require exact fp32 match, TF32 might not give the required precision (atol=1e-4 might be okay though since TF32 still has 1e-7 resolution typically). Wait, let me reconsider - the simplest and most reliable way is to use cuBLAS. But we need a custom CUDA kernel. Let me try a very simple approach first: just use a direct kernel with shared memory and no wmma. Use vectorized loads/stores and a simple tiling approach. Given the sizes are large (batch=128, m=512, k=1024, n=2048), we need good performance. Let me try a simpler tiled GEMM kernel without wmma, and see if it can match PyTorch's performance. Actually, a cuBLAS-based wrapper via raw CUDA might be the way to go. But since the task asks for "custom CUDA kernel", maybe they expect actual CUDA code. Actually, let me check: can we use CUDA's `cublasGemmBatched`? The task says "Write a custom CUDA kernel that replaces a PyTorch reference". Using cuBLAS is technically a CUDA library call, not a custom kernel. Let me write a proper custom kernel. For H100 with Hopper architecture, let me write a simpler but effective kernel using: 1. Shared memory tiling 2. wmma or just plain float accumulation with good memory coalescing 3. Vectorized loads (float4) But first, let me try compiling a basic tiling kernel to see if we can even get a compile. Actually, looking at the problem: FP32 tolerance atol=1e-4. TF32 can typically achieve this. But let me use fp32 accumulation with shared memory. Let me revise and use a cleaner approach. I need proper includes for the CUDA stream. On newer PyTorch, `c10::cuda::getCurrentCUDAStream()` should work but let me check - maybe it needs `#include <ATen/cuda/CUDAContext.h>`. I'll add the proper include and use `at::cuda::getCurrentCUDAStream()`. But actually, wmma with tf32 can be tricky. Let me think about it again. For the sizes: - batch=128 - m=512 - k=1024 - n=2048 Total FLOPs = 128 * 512 * 1024 * 2048 * 2 = 274877906944 ≈ 275 GFLOPs At 67 TFLOPS fp32 or 989 TFLOPS tf32. Even if we get 30% of peak, that's ~300 TFLOPS, so this should take < 1ms. Wait, let me consider using CUTLASS or Triton. Triton might be easier but requires additional setup. Actually, for this problem, I think the most reliable approach is to use a simple but well-optimized custom kernel. Given the tolerance is 1e-4, I could use TF32 tensor cores. But to be safe, let me first try with plain FP32. Actually, the simplest kernel that could be fast: use cuBLAS via a raw handle. But let's think about whether "custom CUDA kernel" allows this. The description says "custom CUDA kernel that replaces the PyTorch reference". I think a raw kernel launch instead of PyTorch's bmm is what they want. Let me try a simple approach: Hopper has warp group kernel support, but that's complex. Let me use a classic tiled GEMM approach with: - Block size: 128x128 tiles - Each block uses shared memory for A and B tiles - Each thread computes a 4x4 or 8x8 tile of output Actually, let me simplify even further. Maybe I can just use the raw `cublasGemmBatched` through direct FFI. The point is to be faster than PyTorch, and while PyTorch likely also uses cuBLAS, we might reduce overhead. But wait - "custom CUDA kernel" probably means writing actual CUDA C++ kernel code. Let me just do a solid tiled GEMM kernel in pure CUDA. For Hopper (sm_90), let me think about the optimal configuration. A classic config: - Block = 256 threads - Tile = 128x128 (each thread computes 8x8 = 64 elements, actually let's do smaller) - Tile = 128x64 with each thread computing 4 elements (2x2) Actually, for fp32, the memory bandwidth is 3.35 TB/s and math is 67 TFLOPS. The arithmetic intensity is: - For each element: read k elements from A and k elements from B, write 1 element to C - Total bytes = batch * m * k * 4 + batch * k * n * 4 + batch * m * n * 4 ≈ 128 * (512*1024 + 1024*2048 + 512*2048) * 4 ≈ 128 * (524288 + 2097152 + 1048576) * 4 ≈ 128 * 3670016 * 4 ≈ 1.9 GB - Total FLOPs = 275 GFLOPs - AI = 275 / 1.9 ≈ 145 FLOPs/byte - At 67 TFLOPS compute and 3.35 TB/s bandwidth, the roofline says compute-bound (since we need 67/3.35 = 20 FLOPs/byte to saturate compute). So this is compute-bound. We need tensor cores or very efficient SIMT execution. For fp32 SIMT, at best we might reach 20-30% of peak = 13-20 TFLOPS. That would be 275/13 ≈ 21ms. At 3.35 TB/s, memory-bounded would be about 0.6 ms but we're compute bound. With TF32 tensor cores, we might reach 30-50% of 989 = 300-500 TFLOPS, giving 0.5-1ms. Wait but PyTorch's bmm likely also uses cuBLAS with tensor cores (if allowed). PyTorch might already be quite fast. But with custom kernel, we might be able to match or beat it. Actually wait, let me reconsider the data sizes: - A: 128 × 512 × 1024 = 67,108,864 floats = 256 MB - B: 128 × 1024 × 2048 = 268,435,456 floats = 1 GB - C: 128 × 512 × 2048 = 134,217,728 floats = 512 MB Total memory: ~1.75 GB. At 3.35 TB/s theoretical, that's 0.52 ms for memory bandwidth. But the actual FLOPs are: 128 × 512 × 1024 × 2048 × 2 = 274,877,906,944 = 274.9 GFLOPs With TF32 tensor cores at say 300 TFLOPS (30% of peak), that's 0.92 ms. So the kernel should be around 1ms if done well. Let me write a kernel using wmma with TF32, since that's the fastest approach. Config: - Block size: 256 threads = 8 warps - Each block computes 128x64 tile of C - Each warp computes 64x32 (using 2 warps along m, 2 along n? or just let each warp handle a sub-tile) - Actually, with wmma m16n16k8, each warp can do one mma per k-step. - 8 warps can cover multiple sub-tiles. Wait, let me reconsider. With wmma::mma_sync for tf32: - M=16, N=16, K=8 - Each mma does 16*16*8*2 = 4096 FLOPs (2 because multiply-add) For a 128x64 tile with K=1024: - Total mma per block = (128/16) * (64/16) * (1024/8) = 8 * 4 * 128 = 4096 mmas - With 8 warps, each warp does 512 mmas - At ~4 cycles per mma (Tensor Core throughput), that's 2048 cycles per warp - At 1.98 GHz, that's about 1 microsecond per warp - Total kernel time: ~1-2 microseconds? That seems too optimistic because we also need memory loads. Actually, accounting for shared memory loads: each mma needs to load A (16x8) and B (8x16) from shared memory. Shared memory bandwidth is high but still takes some cycles. Actually let me be more realistic. For a 128x64 tile with K=1024: - We load A: 128 x BK per k-step, B: BK x 64 per k-step - For BK=64: A_smem = 128*64 = 8K floats = 32KB. B_smem = 64*64 = 4K floats = 16KB. Total = 48KB. That's well within 227KB shared memory limit. - 1024/64 = 16 k-steps. At each k-step: - Load A: 128*64 floats. With float4, that's 2048 loads. With 256 threads, each thread does ~8 loads. That's fast. - Load B: 64*64 floats. With float4, that's 1024 loads. With 256 threads, each thread does ~4 loads. - Total loads per k-step per thread: ~12 float4 loads = 48 floats. - Then 8 warps each do (128/16)*(64/16)/8 = 4 mmas per k-step. - Each warp does 4 mmas. At ~4 cycles each, ~16 cycles. - 16 k-steps * 16 cycles = 256 cycles per warp. - At 1.98 GHz, that's ~130 ns. Plus loads... memory bandwidth is much higher but let's say ~500 ns. - Total time per block: ~1 us. - Total blocks: batch * (m/128) * (n/64) = 128 * 4 * 32 = 16384 blocks - At 132 SMs (H100), each SM handles ~124 blocks. - Total time: ~124 * 1 us = ~124 us. Hmm that's pretty fast. Actually at 128 SMs, 16384 blocks / 128 = 128 blocks per SM. If each takes ~1 us, total ~128 us. PyTorch's cuBLAS might take ~100-200 us. So we need a good kernel. Let me write a cleaner TF32 version. The key issue was the CUDA stream function. Actually, let me try using just the default stream by passing 0 or using `at::cuda::getCurrentCUDAStream()` with `#include <ATen/cuda/CUDAContext.h>`. Hmm, but the compile error said `c10::cuda` has no member. Let me check - maybe in this PyTorch version it's in `at::cuda`. And we need to include the proper header. Let me try. Actually, maybe the simplest way is to just not pass a stream at all (which uses the default stream 0): ``` bmm_tf32_kernel<<<grid, block>>>(...) ``` That should work fine! The default stream is 0 (legacy default stream). Let me try this. Also, let me simplify the kernel a bit. Actually, I realize I should also consider whether TF32 gives enough accuracy. With TF32, we have 10 bits of mantissa vs 23 in FP32. The accumulated error for matrix multiplication of this size might be around 1e-3 to 1e-4. The tolerance is 1e-4. This might be borderline but likely okay. Let me also consider using pure FP32 accumulation but with TF32 input if possible, or just pure FP32. Actually, wmma with `wmma::precision::tf32` does computation in TF32 precision. If that's not accurate enough, we need FP32 which is much slower in SIMT. Let me try the TF32 approach first and check correctness. If it fails, I'll switch to pure FP32. Actually, I just realized: the problem says "produce numerically equivalent outputs at fp32 precision" but tolerance is 1e-4. So TF32 should be fine. Let me proceed. Let me also fix the indexing issues. In TF32 wmma: - wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> - wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> - wmma::fragment<wmma::accumulator, 16, 16, 8, float> Wait, the shape for tf32 is different from fp16. For tf32, the only supported sizes are 16x16x8. Let me rewrite the kernel more carefully. Block configuration: - Each block: 256 threads = 8 warps - Tile: 128x64 - Warps arranged as 8 warps: in a 4x2 grid? or 2x4? - Let's say 4 warps in M, 2 in N. So each warp handles 32x32. - But wmma works with 16x16 tiles. So each warp does 2x2 = 4 mmas per k-step. - Total per block: 8 warps * 4 mmas = 32 mmas per k-step. - Total sub-tiles: (128/16) * (64/16) = 8 * 4 = 32. Good. Wait, but wmma for tf32 requires K=8. So k_inner loops with step 8. Shared memory: - A_smem: 128 x 64. With padding to avoid bank conflicts. - B_smem: 64 x 64. With padding. Actually, for better performance, I should transpose B in shared memory or use column_major for B. But for simplicity, let me use row_major for both with proper padding. For A_smem[128][68], B_smem[64][68]. That uses: - 128 * 68 * 4 bytes = 34,816 bytes ≈ 34KB - 64 * 68 * 4 bytes = 17,408 bytes ≈ 17KB - Total ≈ 51KB. Fine. But actually, for wmma::load_matrix_sync, the leading dimension must be a multiple of 16 for WMMA. Wait, that's for the memory layout. For row_major A, the stride must be 16-byte aligned. float4 is 16 bytes, so stride of multiples of 4 is fine. Actually, the wmma docs say for row_major matrix, the memory must be contiguous in columns with no specific alignment requirement for the leading dimension except that it must fit the matrix. Hmm, maybe better to just use simple pitch. Let me use [BM][BK+4] and [BK][BN+4] which ensures each row starts at a 16-byte boundary and avoids bank conflicts. But for wmma, let me re-read: for `wmma::load_matrix_sync`, the first parameter is the fragment, the second is a pointer to memory, and the third is the "leading dimension". For row_major, this is the number of elements between consecutive rows (i.e., the width/stride). So if A_smem is declared as `float A_smem[BM][BK+4]`, then the leading dimension is `BK+4`. Let me write this kernel properly. Wait, actually I want to reconsider whether I should use TF32 at all. For a tolerance of 1e-4, TF32 should be fine, but let me test if it's correct first. Let me compile a simpler kernel without wmma: a straightforward tiled SIMT kernel first. This will be easier to get right, and if it's fast enough, great. If not, I'll switch to wmma. Actually, for H100, a well-tuned SIMT kernel for fp32 might reach 10-20 TFLOPS, which would give 275/20 = 14ms. That's quite slow. PyTorch's bmm with cuBLAS likely uses TF32 tensor cores and gets ~100-500 TFLOPS, giving < 3ms. So a SIMT kernel would be much slower. So I definitely need tensor cores. Let me write a proper Hopper/H100 kernel using wmma with tf32. One more thing: I should check if the PyTorch version supports `TORCH_CUDNN_V8_API_ENABLED` or similar, and if TF32 is enabled by default. But our goal is to beat PyTorch, not just match it. Actually, wait. Let me reconsider the wmma approach. The issue might be that PyTorch's bmm already uses cuBLAS with TF32 tensor cores, so it might already be very fast. Can a custom kernel actually beat cuBLAS? Maybe yes, if we: 1. Reduce Python/overhead 2. Use optimal tile sizes for our specific problem 3. Fuse operations or optimize memory layout Or maybe no - cuBLAS is extremely optimized. In that case, a direct cuBLAS call might be our best bet. But we're asked to write a "custom CUDA kernel". Hmm, let me think about this differently. Maybe I can use CUTLASS or write a kernel that's competitive. But actually, maybe the simplest way is to do a direct cublas call via raw pointers. PyTorch's bmm has some overhead (checking inputs, converting to proper format, etc.). A direct cublas call might be faster. But the requirement says "Write a custom CUDA kernel". Let me try wmma first. Actually, I just realized something important. For wmma with tf32, we have: ``` wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> ``` This loads 16x8 = 128 elements from A. But tf32 wmma actually interprets these as tf32 values. However, there's a subtlety: when loading from memory, wmma expects the data in a specific format. For tf32, the memory should contain 32-bit floats, but only the upper 19 bits (sign, exponent, and upper 10 mantissa bits) are used. This should work fine with our float arrays. The hardware automatically converts the float32 to tf32 format during the load. Let me rewrite the kernel more carefully and compile. I'll simplify a bit. ```cpp #include <cuda_runtime.h> #include <mma.h> #include <torch/extension.h> using namespace nvcuda; #define BM 128 #define BN 64 #define BK 64 __global__ void bmm_tf32_kernel( const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int batch_size, int m, int n, int k) { // Warp layout: 8 warps total = 4x2 grid const int warp_id = threadIdx.x / 32; const int warp_m = warp_id / 2; // 0..3 const int warp_n = warp_id % 2; // 0..1 // Grid mapping const int n_tiles = n / BN; const int m_tiles = m / BM; const int tiles_per_batch = m_tiles * n_tiles; const int batch_idx = blockIdx.x / tiles_per_batch; const int tile_idx = blockIdx.x % tiles_per_batch; const int m_tile = tile_idx / n_tiles; const int n_tile = tile_idx % n_tiles; const int m_start = m_tile * BM; // 0, 128, 256, 384 const int n_start = n_tile * BN; // 0, 64, 128, ... // Base pointers for this batch const float* A_batch = A + (long long)batch_idx * m * k; const float* B_batch = B + (long long)batch_idx * k * n; float* C_batch = C + (long long)batch_idx * m * n; // Shared memory __shared__ float A_smem[BM][BK + 4]; __shared__ float B_smem[BK][BN + 4]; // WMMA fragments wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag; wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag; wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_frag[2][2]; // Each warp handles 2x2 sub-tiles of 16x16 // Initialize accumulators #pragma unroll for (int i = 0; i < 2; i++) { #pragma unroll for (int j = 0; j < 2; j++) { wmma::fill_fragment(acc_frag[i][j], 0.0f); } } // Main loop over K for (int k_outer = 0; k_outer < k; k_outer += BK) { // Load A to shared memory using all threads // Each thread loads multiple float4 values // Total A elements: BM * BK = 128 * 64 = 8192 floats = 2048 float4 // 256 threads, so each thread loads 8 float4 #pragma unroll for (int i = 0; i < 8; i++) { int t = threadIdx.x + i * 256; int row = t / 16; // 0..127 (BK/4 = 16) int col4 = t % 16; // 0..15 if (row < BM) { const float4* src = (const float4*)(A_batch + (long long)(m_start + row) * k + k_outer + col4 * 4); float4* dst = (float4*)(&A_smem[row][col4 * 4]); *dst = *src; } } // Load B to shared memory // Total B elements: BK * BN = 64 * 64 = 4096 floats = 1024 float4 // 256 threads, each thread loads 4 float4 #pragma unroll for (int i = 0; i < 4; i++) { int t = threadIdx.x + i * 256; int row = t / 16; // 0..63 (BN/4 = 16) int col4 = t % 16; // 0..15 if (row < BK) { const float4* src = (const float4*)(B_batch + (long long)(k_outer + row) * n + n_start + col4 * 4); float4* dst = (float4*)(&B_smem[row][col4 * 4]); *dst = *src; } } __syncthreads(); // Compute #pragma unroll for (int k_inner = 0; k_inner < BK; k_inner += 8) { // Each warp computes a 32x32 tile of C (2x2 of 16x16 WMMA tiles) #pragma unroll for (int i = 0; i < 2; i++) { int a_row = warp_m * 32 + i * 16; wmma::load_matrix_sync(a_frag, &A_smem[a_row][k_inner], BK + 4); #pragma unroll for (int j = 0; j < 2; j++) { int b_col = warp_n * 32 + j * 16; wmma::load_matrix_sync(b_frag, &B_smem[k_inner][b_col], BN + 4); wmma::mma_sync(acc_frag[i][j], a_frag, b_frag, acc_frag[i][j]); } } } __syncthreads(); } // Store results #pragma unroll for (int i = 0; i < 2; i++) { #pragma unroll for (int j = 0; j < 2; j++) { int c_row = m_start + warp_m * 32 + i * 16; int c_col = n_start + warp_n * 32 + j * 16; wmma::store_matrix_sync(C_batch + (long long)c_row * n + c_col, acc_frag[i][j], n, wmma::mem_row_major); } } } ``` Wait, this has an issue. `wmma::load_matrix_sync` for matrix_a is loading 16x8 elements. But `a_frag` is `matrix_a` with shape 16x16x8. Wait, the declaration is `wmma::fragment<wmma::matrix_a, M, N, K, ...>` where M=16, N=16, K=8 means it's an MxK matrix, i.e., 16x8. So `a_frag` is 16 rows × 8 columns. `b_frag` is 8 rows × 16 columns. `acc_frag` is 16 rows × 16 columns. So each warp handles (32/16) × (32/16) = 2 × 2 = 4 sub-tiles of 16x16. Wait, but each sub-tile needs a different a_frag and b_frag. In my code above, I'm loading a_frag once for each i (row), then b_frag for each j (col), and reusing a_frag. That should be correct. Actually wait, `a_frag` is declared outside the loops and reused. But `wmma::mma_sync` takes `a_frag`, `b_frag`, `acc_frag`. Can I reuse a_frag for different b_frags? Yes, I think so - the mma_sync just reads from the fragments. Actually, wait. In the code above, I load `a_frag` and inside the j loop I load `b_frag` and call `mma_sync`. But I'm overwriting `b_frag` each time. That's fine since `mma_sync` reads the values. But I'm also reusing `a_frag` without reloading. That's correct since the same A tile is used with different B tiles. However, wait - I need to be careful about the liveness of `a_frag`. Let me check if `load_matrix_sync` populates the fragment and `mma_sync` reads it. Yes, that should work. Actually, wait. I notice that a_frag is loaded for i, then inside j loop I load b_frag and call mma_sync. But each j iteration overwrites b_frag and calls mma_sync with the same a_frag. That is correct. One concern: is `k_inner` loop well-structured? The A_smem is accessed as `&A_smem[a_row][k_inner]`. Since A_smem is declared as [BM][BK+4], and `a_row` ranges from warp_m*32 + i*16, this should point to a 16x8 submatrix starting at (a_row, k_inner). Similarly for B_smem. Let me double-check the shapes: - A_smem has dimensions [128][68]. - For a_frag (16×8), we pass pointer to `&A_smem[a_row][k_inner]` and stride `BK+4 = 68`. - The memory layout is row-major, so consecutive elements in a row are contiguous. - wmma::load_matrix_sync should interpret this as a 16×8 matrix with leading dimension 68. This seems correct. One issue: `wmma::store_matrix_sync` for accumulator writes to C. We need to ensure no race conditions. Each warp writes its own 32×32 tile, and different warps handle different parts of the block. No two warps should write to the same location. With warp layout 4×2: - warp_m ranges 0..3, warp_n ranges 0..1 - Each warp handles warp_m*32 to warp_m*32+31 in M, warp_n*32 to warp_n*32+31 in N. - Total M coverage: 4*32 = 128 = BM. Total N coverage: 2*32 = 64 = BN. - No overlap. Good. Now let me make sure the block and grid sizes are correct: - Block: 256 threads (8 warps) - Grid: batch_size * (m/BM) * (n/BN) = 128 * 4 * 32 = 16384 blocks Wait, m=512, BM=128, so m/BM = 4. n=2048, BN=64, so n/BN = 32. tiles_per_batch = 4 * 32 = 128. Total blocks = 128 * 128 = 16384. This is on the high side but H100 has a lot of SMs and should handle it. Actually, I wonder if I should increase block size or tile size to reduce the number of blocks. With 16384 blocks and 132 SMs, that's ~124 blocks per SM. But each block has very little work, so the dispatch overhead might be significant. Alternatively, I could use larger tiles. For example, BM=128, BN=128. Then tiles_per_batch = 4 * 16 = 64. Total blocks = 128 * 64 = 8192. Better. With BN=128 and 8 warps: - warp_n could be 0..1 if we keep warp layout 4x2. But then each warp handles 32x64... wait, that doesn't work because 2*64 = 128. But 32x64 can't be divided into 16x16 tiles. 32x64 = 2*4 = 8 tiles of 16x16. With 8 warps, each warp does 1 tile? No, each warp should do multiple tiles for better utilization. Actually, let's keep BN=64. Or maybe use more warps per block. With 256 threads = 8 warps, we have 8 warps. If tile is 128x128: - Need (128/16)*(128/16) = 8*8 = 64 tiles of 16x16 - 8 warps, each handles 8 tiles. - Each warp could handle 2x4 = 8 tiles, i.e., 32x64 per warp. So warp layout could be 4x2 for a 128x128 block. warp_m: 0..3 (covers M: 128), warp_n: 0..1 (covers N: 64... wait, only 2*32=64 in N, but we need 128). Let me reconsider: if block is 128x128: - warp_m could go 0..3 (4 values), each covers 32 in M → 4*32 = 128. Good. - warp_n could go 0..3 (4 values), each covers 32 in N → 4*32 = 128. - Total warps needed: 4*4 = 16 warps = 512 threads. So block size would be 512 threads. Is that okay? Yes, up to 1024 threads per block. But H100 supports up to 32 blocks per SM, and with 512 threads per block, you can have at most 2 blocks per SM (since 1536 threads per SM limit on some architectures, but H100 might have more). Actually, H100 has 2048 threads per SM. So 4 blocks of 512 threads would fit. Total blocks: 128 * 4 * 16 = 8192. With 132 SMs and 4 blocks each = 528 active blocks. That leaves many blocks queued. Actually, let me stick with the simpler 128x64 tile and 256 threads for now. We can adjust later. Wait, I also need to be careful about the float4 loads. The A and B tensors are guaranteed to be contiguous (we call .contiguous()), but they might not be 16-byte aligned at arbitrary offsets. When we do: ``` const float4* src = (const float4*)(A_batch + (long long)(m_start + row) * k + k_outer + col4 * 4); ``` This requires that `(m_start + row) * k + k_outer + col4 * 4` is a multiple of 4 (which it is since col4 * 4 is a multiple of 4). But it also needs to be 16-byte aligned for the float4 load. Since k_outer can be any multiple of BK=64, and k=1024, the offset is always a multiple of 4, but not necessarily of 16. Actually, for float4 alignment, the pointer needs to be 16-byte aligned. Since we're offsetting by col4 * 4 where col4 is an integer, the base is `(m_start + row) * k + k_outer`. If `k` is a multiple of 4 (which 1024 is), then `(m_start + row) * k` is a multiple of 4. And `k_outer` is a multiple of 64, which is a multiple of 4. So all offsets are multiples of 4. But for 16-byte alignment, we need offsets to be multiples of 16, i.e., `col4 * 4` needs the base to be 16-byte aligned. The base offset for each row is `(m_start + row) * k + k_outer`. Since k=1024 and k_outer is a multiple of 64, both are multiples of 16. So the base is 16-byte aligned, and `col4 * 4` gives offsets 0, 4, 8, 12, ... which are all multiples of 4 but not all multiples of 16. Wait, `float4` requires 16-byte alignment. So if `col4 = 0`, offset is base + 0 = aligned. If `col4 = 1`, offset is base + 4 = not 16-byte aligned! This would cause unaligned loads which might be slow or crash. Oh! I see the issue. I should load `float4` only when the address is 16-byte aligned. For `col4 = 0, 1, 2, 3`, the offsets are 0, 4, 8, 12 from the row start. Only col4=0 is 16-byte aligned. So my float4 loading is potentially wrong for col4 > 0! Wait, let me reconsider. `(const float4*)(A_batch + offset)` where offset could be base + 4. This would indeed be an unaligned float4 access. To fix this, I should load only at positions where the offset is a multiple of 16 (in bytes, i.e., multiple of 4 in floats). So col4 should only take values 0, 4, 8, 12, etc. Let me fix the loading: For A: - Total columns in groups of 4: BK / 4 = 16 groups. - But we want to load float4, so we need `col4` to range from 0 to BK/4 - 1 = 15. - Wait, but the address needs to be 16-byte aligned. Since the row start is aligned, `col4 * 4` in floats = `col4 * 16` in bytes. So as long as we iterate `col4` from 0 to 15, the address is `base + col4 * 16` bytes, which is always 16-byte aligned! Wait, I made a confusion. `col4` in my code is the index of the float4 element, not the float index. Let me re-read: ```cpp int col4 = t % 16; // 0..15 ... const float4* src = (const float4*)(A_batch + (long long)(m_start + row) * k + k_outer + col4 * 4); ``` Here, `col4` ranges 0..15, and we add `col4 * 4` to the pointer. Since each float is 4 bytes, `col4 * 4` floats = `col4 * 16` bytes. So the address is `base + col4 * 16` bytes, which is indeed 16-byte aligned for all values of col4. Phew! It's actually correct. The pointer arithmetic `A_batch + offset` adds `offset * sizeof(float)` bytes. So `col4 * 4` gives `col4 * 16` bytes. Good. Similarly for B. So the float4 loads are properly aligned. Let me continue. Now, for the shared memory stores, I'm doing: ```cpp float4* dst = (float4*)(&A_smem[row][col4 * 4]); ``` Here `&A_smem[row][col4 * 4]` gives a pointer to a float, and we cast it to float4*. Since A_smem[row] starts at an address that's aligned based on the row and the array layout. `A_smem` is declared as `float A_smem[BM][BK+4]`. Each row has `BK+4 = 68` floats. Since 68 * 4 = 272 bytes, which is not a multiple of 16! So row 1 starts at byte 272, which is not 16-byte aligned. This means storing float4 to `&A_smem[1][0]` would be unaligned. But wait, I need to check if CUDA allows unaligned shared memory accesses. Actually, for shared memory on modern GPUs, unaligned accesses are generally okay for 16-byte loads/stores. But bank conflicts might occur. Actually, for shared memory, the alignment requirement for float4 is that the address should be 16-byte aligned for optimal performance. But if each row is 272 bytes, then: - Row 0 starts at byte 0 (16-byte aligned) - Row 1 starts at byte 272 (not 16-byte aligned, 272 % 16 = 0, wait, 272 / 16 = 17 exactly. So 272 is 16-byte aligned!) - Row 2 starts at byte 544 (544 / 16 = 34, aligned) - In general, 68 floats * 4 bytes = 272 bytes per row, and 272 = 16 * 17. So every row starts at a 16-byte boundary! Great! So shared memory is also 16-byte aligned per row. This is because I chose `BK + 4 = 68`, and 68 * 4 = 272 = 16 * 17. So all float4 accesses are properly aligned. Excellent. Now let me also check B_smem. `float B_smem[BK][BN + 4]`. BN = 64, BN + 4 = 68. Same situation, 68 * 4 = 272 = 16 * 17. So each row is 16-byte aligned. Good. OK, the kernel should be compileable and functional. Let me also fix the includes and try to compile. Actually, let me reconsider if we can even get better performance by tweaking parameters. The current config: - BM=128, BN=64, BK=64 - Block = 256 threads - Each block loads 128*64 + 64*64 = 12288 floats from global memory per k-step - K dimension has 1024/64 = 16 steps - Total loads per block: 12288 * 16 = 196608 floats = 786432 bytes - Per thread: 786432 / 256 = 3072 bytes = 768 floats = 192 float4 loads - For the compute: each thread is part of a warp doing mma operations Actually, this seems like a lot of loads per thread. Let me see if we can load more vectorized or in fewer instructions. Currently, each thread loads: - A: 8 float4 per k-step (32 floats) - B: 4 float4 per k-step (16 floats) - Total per k-step: 12 float4 loads - Over 16 k-steps: 192 float4 loads per thread That's fine. Each float4 load is a single instruction. Now, the compute per warp: - 4 mma_sync per k-step (2x2 tiles of 16x16) - 16 k-steps: 64 mma_sync per warp - At ~4 cycles each (best case tensor core throughput), ~256 cycles - Plus overhead for fragment loads/stores, syncs, etc. Maybe 500-1000 cycles. - At 1.98 GHz, ~0.25-0.5 microseconds per block. - Actually wait, each warp does 2*2=4 mma per k_inner step. There are BK/WK = 64/8 = 8 k_inner steps per k_outer. And 16 k_outer steps. So 16*8 = 128 k_inner steps. Each with 4 mmas. So 512 mmas per warp. At 4 cycles each, 2048 cycles. - Plus loads and syncs. Maybe 3000-5000 cycles total. - At 1.98 GHz, that's ~2 microseconds per block. So 16384 blocks * 2 us = total Warp-time of 32768 us. With 132 SMs * 8 warps each = 1056 warps, total time is approximately 32768 / 1056 = 31 us. Hmm, that seems optimistic. There are overheads I'm not accounting for, but this gives a rough idea. Anyway, let me compile this and test. If it's not fast enough, I'll tune further. But actually, let me also consider: what if PyTorch's bmm already takes ~50-100 us on this problem? Then our kernel needs to be faster. Actually, I just realized there's a potential issue with my wmma code. I need to double-check the dimensions. For wmma with tf32, the supported shapes are: - wmma::fragment<wmma::matrix_a, 16, 16, 8, ...> → loads a 16×8 tile - wmma::fragment<wmma::matrix_b, 16, 16, 8, ...> → loads an 8×16 tile - wmma::fragment<wmma::accumulator, 16, 16, 8, float> → accumulates 16×16 Wait, so `a_frag` is 16×8, `b_frag` is 8×16, and `acc_frag` is 16×16. In my code, I'm calling: ```cpp wmma::load_matrix_sync(a_frag, &A_smem[a_row][k_inner], BK + 4); ``` Here, `a_frag` is matrix_a with shape 16×8. The pointer points to `&A_smem[a_row][k_inner]`. The memory should represent a 16×8 matrix in row-major format with leading dimension `BK+4`. For row major, element (i, j) is at `ptr[i * lda + j]`. In my shared memory, `A_smem[row][col]` is at `&A_smem[0][0] + row * (BK+4) + col`. So the leading dimension is indeed `BK+4`. And `a_row` is the starting row, `k_inner` is the starting column. So `&A_smem[a_row][k_inner]` points to the start of a 16×8 submatrix. This seems correct. Similarly for B: ```cpp wmma::load_matrix_sync(b_frag, &B_smem[k_inner][b_col], BN + 4); ``` Wait, `b_frag` is matrix_b with shape 8×16. But `B_smem` is declared as `float B_smem[BK][BN+4]`. `&B_smem[k_inner][b_col]` points to row `k_inner`, column `b_col`. The leading dimension is `BN+4`. Is `B_smem` in row-major format? Yes, declared as `B_smem[BK][BN+4]`. And matrix_b is 8×16. So row-major 8×16 with leading dimension BN+4. This seems correct. For the accumulator store: ```cpp wmma::store_matrix_sync(C_batch + (long long)c_row * n + c_col, acc_frag[i][j], n, wmma::mem_row_major); ``` `acc_frag` is 16×16. Stored to `C_batch + c_row * n + c_col` with leading dimension `n` in row-major. This is correct. Now, I also need to include `<cuda_fp16.h>` or similar for the precision type. Actually, `wmma::precision::tf32` is available in newer CUDA versions. Let me check - H100 is sm_90 and uses CUDA 12.x, so tf32 should be supported in wmma. Wait, is tf32 supported in wmma? I know it's supported in cublas and cutlass. For wmma, I think it's available starting from Ampere (sm_80) with `wmma::precision::tf32`. Let me check: yes, `wmma::precision::tf32` is available in CUDA 11.0+ for sm_80+. But wait, the compilation error I got earlier was about `at::cuda::getCurrentCUDAStream`, not about wmma. So let me try again with a fixed stream issue. Actually, I just realized - instead of using the stream function, let me just pass `0` for the default stream: ```cpp bmm_tf32_kernel<<<grid, block>>>(...) ``` This should work and avoids the include issue entirely. OK let me try a cleaner version with simplified stream handling. Actually, let me also double-check: do I need `extern "C"` for the forward function? Not necessarily since we're using C++ and pybind11 will generate bindings. Let me write the complete code now. One more thing: I should verify that m, n, k are multiples of the tile sizes. Given: - m = 512, BM = 128 → 512/128 = 4 ✓ - n = 2048, BN = 64 → 2048/64 = 32 ✓ - k = 1024, BK = 64 → 1024/64 = 16 ✓ All are exact multiples, so no boundary conditions needed. Let me write and compile the code.

The shared memory exceeded the limit. Let me reduce the tile sizes to fit within 48KB shared memory: - BM=128, BK=64: 128 * 64 * 4B = 32,768 bytes - BK=64, BN=64: 64 * 64 * 4B = 16,384 bytes Total: 49,152 bytes, which exceeds 48KB (but I said max 48KB earlier, actually H100 has 228KB per SM but per block it's limited to 48KB or 96KB with special config) Let me use smaller tiles: BM=128, BN=64, BK=32: - A_smem: 128 * 32 * 4 = 16,384 bytes - B_smem: 32 * 64 * 4 = 8,192 bytes Total: 24,576 bytes - well within limit Or BM=64, BN=64, BK=64: - A_smem: 64 * 64 * 4 = 16,384 bytes - B_smem: 64 * 64 * 4 = 16,384 bytes Total: 32,768 bytes Let me also try using BK=32 with BM=128, BN=128 to get better compute density, but that might be too many registers. Actually let me try a simpler approach first - using cublas directly through the extension API. But actually I can just try a more conservative tile size. Let me go with BM=64, BN=64, BK=64 and 8 warps (256 threads), each warp handles 16x16x8 MMA, so each block computes 64x64 output: - 2x2 warp layout (4 warps) with each warp handling 32x32, but then 64x64 / 32x32 = 4 accumulators... Actually for 8 warps in a 4x2 layout with BM=64, BN=64: - Each warp handles 16x16 output - 64x64 = 16 warps worth But we only have 8 warps. Let me do BM=64, BN=32 with 4 warps (128 threads) in a 2x2 layout, each with 32x16, but then MMA is 16x16... - Warp m: 0..1 (2), Warp n: 0..1 (2), 4 warps total - Each warp does 2 A loads (32x8 -> two 16x8) and 2 B loads at a time... Let me keep it simple. BM=128, BN=64, BK=32 with 8 warps: - Warp layout: 4x2 (warp_m 0..3, warp_n 0..1) - Each warp handles 32x32 output (two 16x16 MMA in each dimension = 4 MMAs per warp) Actually, each warp does: - For i in 0,1 (A dimension): load 16x8 from A at row warp_m*32+i*16 - For j in 0,1 (B dimension): load 16x8 from B at col warp_n*32+j*16 - MMA them together Wait, MMA is 16x16x8. So for 32x32 output per warp: - Need 2x2 = 4 MMA operations per warp - Accumulators: 4 - A loads per k_inner: 2 (32 / 16 = 2 MMAs along A's row dim, 8 along K dim) - B loads per k_inner: 2 (32 / 16 = 2 MMAs along B's col dim, 8 along K dim) Total MMA per k loop per warp: 2 * 2 = 4. With BK=32 and WK=8, that's 4 iterations of the k_inner loop. So 16 MMAs per warp per outer K loop. With 8 warps, 128 MMAs per block. Number of blocks: batch_size=128, m=512, n=2048, BM=128, BN=64 -> 1*16*32 = 512 blocks. Total MMAs: 512 * 128 = 65,536. Each MMA is: 16 * 16 * 8 = 2,048 FMAs? No wait, TF32 MMA on Hopper does... well the standard MMA is a warp-level operation that does some FMAs. With tf32 precision, MMA accumulates to float32. Actually, let me just test out correctness first. I'll reduce shared memory. Let me also just try a simpler manual unrolled kernel. Actually, I realize wmma::precision::tf32 may not actually exist on Hopper in the same way, or the fragment loads may need the inputs to be in specific precisions. Let me try a simpler traditional shared memory approach without wmma first, or try nvcuda::wmma properly. Actually wmma::precision::tf32 should be available on Hopper (sm_90). Let me use BM=64, BN=32, BK=16 to be very safe on shared memory, then expand if correct. Or just use BM=64, BN=32, BK=32: - A_smem: 64 * 32 * 4 = 8,192 - B_smem: 32 * 32 * 4 = 4,096 Total: 12,288 With 4 warps (128 threads): - Warp layout 2x2: warp_m 0,1; warp_n 0,1 - Each warp output: 32x16 = 2 MMA along A's row * 1 MMA along B's col = 2 MMA per k_inner - k_inner iterations: BK/WK = 32/8 = 4 - Total MMA per warp: 2 * 4 = 8 That should work. Let me also simplify and avoid the array of fragments.

The issue is that my cpp_sources don't provide the actual implementation. The `cpp_sources` argument is compiled as C++ and should provide the declarations that can be found. But the problem is that the cuda function `bmm_forward` is compiled by nvcc and the C++ linker needs to see it. Let me think... When using `load_inline`, I need to provide declarations in `cpp_sources` for the functions I expose in `functions`. But the issue is my declaration had the right signature but the definition was only in cuda_sources. Wait, looking at the error again, it's a link error - undefined symbol. That means my function name was mangled differently between C++ compilation and CUDA compilation. Actually, I suspect the issue is that I'm not returning anything from the kernel because it's incomplete. But the real function that's undefined is `bmm_forward`, which was only in cuda_source. Wait, I think the issue with load_inline is that the C++ and CUDA are compiled separately and need to link together. The cpp_sources should have the declaration but the function body should also be accessible. Let me look more carefully... Looking at other people's usage and how load_inline works: - `cpp_sources`: raw C++ source. These are compiled as C++ and can include headers. This can contain the host-side wrapper functions that call CUDA kernels, or just declarations. But function bodies defined here should be accessible. - `cuda_sources`: CUDA source. These are compiled by nvcc. Any function defined here will be in the CUDA compilation unit. For proper linking, I need to make sure the functions in `functions=[...]` are visible between translation units. In load_inline, the code is assembled into what's essentially a single extension. The C++ and CUDA get compiled and linked together. So having `bmm_forward` defined in `cuda_sources` should be fine, and having its declaration in `cpp_sources` should also be fine. BUT WAIT - my previous compilation actually failed because the CUDA code was incomplete (it doesn't include a complete `bmm_forward` implementation - it ends mid-stream with a comment). The compile failed earlier with the undefined symbol error, but actually maybe the compilation succeeded for the cpp part but the CUDA part failed? Let me look again: The error was `undefined symbol: _Z11bmm_forwardN2at6TensorES0_`. This is a linker error. And yes, my cuda_source was incomplete - it didn't have a return statement or closing brace for the kernel. The nvcc would have compiled this file to an object, but `bmm_forward` in it wasn't complete so... no, a function body without a closing brace should be a compile error too. Actually I think `nvcc` might have compiled it somehow? No, that should fail. But the error only shows the undefined symbol... Let me just write a complete, working kernel. The simplest approach for correctness is to implement a standard tile-based FP32 SGEMM kernel, not using WMMA right away to avoid complexity. Given the matrix sizes: - batch=128, M=512, K=1024, N=2048 Actually, I should just call cuBLAS from the extension, which would be near-optimal. But that's probably not the spirit of the problem. Or maybe I can use a well-tuned cutlass kernel via PyTorch's interface. Wait, a simpler approach: `torch.bmm` might internally call cuBLAS or use PyTorch's implementation. But for maximum speed, I should use a custom kernel that handles the specific dimensions well. Actually, let me try a straightforward thread-tiled MM kernel first to validate the extension loading mechanism, then optimize. Or better yet - on H100, I could try using the new WGMMA or just a very well-tuned block tiling kernel with a large amount of thread-level parallelism. Let me try a simpler approach: just use cuBLAS from the CUDA extension. That would definitely be fast: ```cpp #include <cublas_v2.h> torch::Tensor bmm_forward(torch::Tensor A, torch::Tensor B) { int batch = A.size(0), m = A.size(1), k = A.size(2), n = B.size(2); auto C = torch::empty({batch, m, n}, A.options()); cublasHandle_t handle; cublasCreate(&handle); float alpha = 1.0f, beta = 0.0f; cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, B.data_ptr<float>(), n, k*n, A.data_ptr<float>(), k, m*k, &beta, C.data_ptr<float>(), n, m*n, batch); cublasDestroy(handle); return C; } ``` But PyTorch's bmm already uses cuBLAS internally! So calling cuBLAS directly won't be faster (in fact it'll be slower because we create/destroy the handle). The key is that we need to write a better kernel. Actually, wait - PyTorch might not use the most optimal cuBLAS settings or tile sizes. But it's very likely already optimal for fp32 on H100. The only way to beat it might be with specific tuning. Actually, I realize I should check what the reference runtime is. The problem says `torch.bmm` is the reference. On H100 with fp32, cuBLAS is highly optimized. But maybe I can use TF32 tensor cores or some specific optimization. Wait, but the output must be numerically equivalent at fp32 precision. TF32 might not pass the tolerance check if it's too strict? The tolerance for fp32 is atol=rtol=1e-4, and TF32 has about 1e-5 precision, so it should be fine. Actually, no - TF32 only has 10 bits of mantissa vs 23 for FP32, so it might accumulate error. Let me think... Actually the tolerance is 1e-4, which is relatively loose. Tensor core operations with TF32 should be within this tolerance. For H100 (sm_90), I could use wgmma (warp group-level matrix multiply accumulate) with 64x8x16 tile sizes... But that's complex. Let me try a simpler but well-optimized kernel: a thread block that computes a larger tile using shared memory and vectorized loads, targeting the specific dimensions. Given: - A: (128, 512, 1024) - B: (128, 1024, 2048) Total flops: 128 * 512 * 1024 * 2048 * 2 = 2.75e11 FLOPs At 67 TFLOPS FP32, this takes ~4ms At 989 TFLOPS TF32 tensor core, this takes ~0.28ms PyTorch bmm on H100 probably uses cuBLAS which can hit close to peak, maybe 2-3ms. To beat it, I need tensor cores. Let me try WMMA but more carefully. I'll write a complete kernel. Tiles: - Block: 128 threads = 4 warps - Each block computes BM=128 x BN=64 tile - Each warp computes 32x32 tile (8 WMMAs of 16x16 per warp) - Actually, let's use WGMMA or just a clean WMMA kernel. Actually wait, Hopper H100 has new instructions. But the simple approach: use WMMA with a proper accumulator layout. Let me design: - Block: 256 threads = 8 warps - Each block computes BM=128 x BN=64 - Warp layout: 4x2 (4 warps along M, 2 along N) - Each warp computes 32x32 tile - Actually 128/4 = 32 M per warp, 64/2 = 32 N per warp - Each 32x32 tile = 4 WMMAs of 16x16 BK=32 to keep shared memory small. Shared memory: - A_smem: 128 * 32 * 4 = 16KB - B_smem: 32 * 64 * 4 = 8KB Total: 24KB, well within limits. Let me write a complete 16x16 WMMA kernel with 2x2 acc_frag per warp. Actually, I realize the previous WMMA approach with acc_frag[2][2] is correct. Let me just make sure to define the kernel properly. But there's another thing: WMMA with tf32 precision on sm_90 should be available. Let me use proper tile dimensions. WMMA on sm_80+ supports 16x16x8 for tf32. Let me write this properly now. The previous compilation with BM=128 BN=64 BK=64 failed because shared memory was too large (0xcc00 = 52KB). With BK=32: - A_smem: 128 * 36 * 4 = ~18KB - B_smem: 32 * 68 * 4 = ~8.7KB Total: ~27KB, fits. With BM=128, BN=64, BK=32, total blocks: - M tiles = 512/128 = 4 - N tiles = 2048/64 = 32 - Total tiles per batch = 128 - Total blocks = 128 * 128 = 16384 - Threads per block = 256 - Total threads = 4,194,304 This should be enough parallelism for H100. Let me write the complete kernel now. I need to be careful with: 1. Complete function bodies 2. Proper return 3. Correct indexing 4. Aligned shared memory access for WMMA Also, for the function to be properly linked, I should make sure the declaration and definition match. Since the definition is in cuda_sources with C++ linkage (it's compiled by nvcc with --std=c++17), name mangling should be compatible... but actually nvcc compiles CUDA as C++ with some mangling rules. As long as there's no `extern "C"`, it should be fine. The issue before was that my code was incomplete. Let me write a fully complete, correct kernel. If I still get undefined symbol, I might try putting the definition in cpp_sources and the kernel in cuda_sources, or just put everything in cpp_sources with `__global__` etc. (nvcc handles this). Actually, looking at load_inline: - `cpp_sources` is compiled as C++ - `cuda_sources` is compiled as CUDA - They're linked together So having the declaration in cpp_sources and definition in cuda_sources is standard. The previous failure was due to the CUDA code being syntactically incomplete (no closing brace for `bmm_forward` function). Let me write a complete kernel.

The kernel compiled successfully. Now let me run the correctness check. The problem sizes are: - batch_size = 128 - m = 512 - k = 1024 - n = 2048 With my tiling BM=128, BN=64, BK=32: - n_tiles = 2048/64 = 32 - m_tiles = 512/128 = 4 - tiles_per_batch = 32 * 4 = 128 - total_blocks = 128 * 128 = 16384 blocks Each block processes one 128x64 tile of C, using 8 warps (256 threads) with each warp computing a 32x32 tile. Each warp uses 4 MMAs (2 along M * 2 along N). The dimensions are all divisible by our tile sizes so this should work. Let me run correctness check.

The correctness check failed. This is likely due to TF32 being imprecise, or there might be an issue with my memory access patterns. Since the tolerance for fp32 is 1e-4, TF32 has ~1e-3 relative precision so that could be the issue. Let me check if the problem requires fp32 precision and if TF32 is causing the mismatch. The issue says "numerically equivalent outputs at fp32 precision" but fp32 has tolerance 1e-4, whereas TF32 can have higher error. Actually wait, let me reconsider. Maybe there's a more subtle bug. Let me trace through what the kernel does more carefully. Actually, the issue could also be that I'm not handling edge cases or there might be indexing issues. But since the matrices are all nice multiples of the tile size, that shouldn't be an issue. Another thought: maybe `wmma::precision::tf32` requires the inputs to be in a specific TF32 format. When we load float values, we should call `wmma::__float_to_tf32` or similar? Actually, looking at the WMMA docs, the precision::tf32 overload should handle this automatically by rounding. Wait, actually the issue might be that I'm using `float` directly with WMMA and `tf32` precision. The values should be automatically converted when loaded with `load_matrix_sync`. But the tolerance for fp32 is `1e-4`, and TF32 has around 10 bits of mantissa precision (vs 23 in fp32), which gives about 3 decimal digits of precision vs 7 for fp32. For random values in [0,1], the absolute error from TF32 could easily be around 0.001. Let me check: - TF32 mantissa: 10 bits ≈ 3 decimal digits - Expected relative error: roughly 2^(-10) ≈ 0.001 - For values around 0.1-1.0, absolute error ≈ 0.0001 to 0.001 Wait, the errors we're seeing are ~0.17, which is way too large. That's got to be a bug in the kernel, not just TF32 precision. Let me think about this more carefully... Actually, the average difference is 0.168, which is huge. This means the output is completely wrong, not just TF32 precision issues. Looking at my load patterns: - For B: I'm loading from `B_smem[k_inner][b_col]`. Wait, WMMA for matrix B in row_major expects the leading dimension to be the N dimension. B is loaded from shared memory as [BK][BN+4] and I'm loading with `BK + 4` as the stride. Actually, I think the issue is in how I'm loading B. Let me reconsider. For wmma::fragment of type matrix_b, with row_major layout: - The matrix should be BK x BN with leading dimension BN. - When I call `wmma::load_matrix_sync(b_frag, &B_smem[k_inner][b_col], BN + 4)`, the 3rd argument is the leading dimension. - This looks correct. Wait, but `b_smem` is declared as `__shared__ float B_smem[BK][BN + 4]`. So in memory, the row dimension is the first dimension (BK), and the column dimension is the second dimension (BN+4). Since CUDA uses row-major storage, the element `B_smem[i][j]` is at address `&B_smem[0][0] + i * (BN + 4) + j`. When calling `load_matrix_sync` with leading dimension `BN + 4`, the elements are correctly accessed as `base + row * (BN + 4) + col`. This looks right. Hmm, but maybe the issue is that `B_smem[k_inner]` is a pointer to a float array of size BN+4, which starts at `&B_smem[0][0] + k_inner * (BN + 4)`. Then `&B_smem[k_inner][b_col]` is `&B_smem[0][0] + k_inner * (BN + 4) + b_col`. This should be correct. Actually, wait. Let me re-examine the B loading code. I have: ```cpp const float4* src = (const float4*)(B_batch + (long long)b_row * n + b_col); float4* dst = (float4*)(&B_smem[row][col4 * 4]); *dst = *src; ``` Here `b_row = k_outer + row` and `b_col = n_start + col4 * 4`. So `B_smem[row]` contains the data for k-dimension row `row` at position `col4 * 4` in the n dimension. This should be correct since B_smem is [BK][BN] where the first dimension is k and the second is n. Wait, but WMMA expects a row-major matrix. For the B matrix in a matrix multiplication, the data in shared memory should be interpreted as having rows along the K dimension and columns along the N dimension. Let me check if I might be computing matrix transposes incorrectly... Actually, for the WMMA operation: - A_frag represents a 16x8 tile from matrix A (M x K) - B_frag represents an 8x16 tile from matrix B (K x N) - acc_frag is a 16x16 tile from matrix C (M x N) This is all correct. OH WAIT! I think I found the issue. Look at the loop: ```cpp for (int k_inner = 0; k_inner < BK; k_inner += WK) { for (int i = 0; i < 2; i++) { wmma::load_matrix_sync(a_frag[i], &A_smem[a_row][k_inner], BK + 4); for (int j = 0; j < 2; j++) { wmma::load_matrix_sync(b_frag[j], &B_smem[k_inner][b_col], BN + 4); wmma::mma_sync(acc_frag[i][j], a_frag[i], b_frag[j], acc_frag[i][j]); } } } ``` I need to reconsider whether `a_frag[i]` is loaded correctly. For the 16x8 WMMA fragment of A: - It reads a 16x8 submatrix from pointer `&A_smem[a_row][k_inner]` - The leading dimension is `BK + 4` - So element (r, c) comes from `&A_smem[a_row][0] + r * (BK + 4) + (k_inner + c)` - Which is `&A_smem[a_row + r][k_inner + c]` This means it's reading rows a_row to a_row+15 and columns k_inner to k_inner+7. That's correct. For B fragment: - It reads a 8x16 submatrix from `&B_smem[k_inner][b_col]` - Leading dimension is `BN + 4` - Element (r, c) comes from `&B_smem[k_inner][0] + r * (BN + 4) + (b_col + c)` - Which is `&B_smem[k_inner + r][b_col + c]` So it reads rows k_inner to k_inner+7 and columns b_col to b_col+15. Correct for B. For MMA, C[i][j] += A[i] * B[j]. Correct. Hmm, I'm not seeing an obvious bug. Let me try a simpler kernel first - instead of trying to optimize with WMMA + complex shared memory tiling, let me try using a simpler approach with cublas directly, or try a simpler shared-memory GEMM kernel. Actually, let me try to just call cuBLAS through torch, or try a simpler kernel to verify understanding first. Wait, looking at the problem again. The simplest approach that will definitely work: just use `torch.bmm(A, B)` but add some optimization? No, the reference already does that. We need to beat it. Let me think about what's faster than PyTorch's bmm. PyTorch uses cuBLAS under the hood, and cuBLAS is highly optimized. To beat it, we'd typically need to: 1. Fuse operations 2. Use specialized kernels for specific sizes But our problem is literally just BMM with no operations to fuse. Hmm, that's hard. Actually, there could be overhead in PyTorch's dispatch for bmm. If we call cublas directly from a custom CUDA kernel, that might reduce overhead. Or if the problem sizes are unusual, a custom kernel might match. But the sizes are: batch=128, m=512, k=1024, n=2048. These are quite large. For large GEMMs, cuBLAS is very well optimized, especially on H100. Wait, actually for large GEMMs, the compute is the bottleneck. Let me compute the arithmetic intensity: - Flops per batch: 2 * m * n * k = 2 * 512 * 2048 * 1024 ≈ 2.1 GFLOP - Total flops: 128 * 2.1 GFLOP = 274 GFLOP - Memory reads: approximately (batch_size * (m*k + k*n + m*n)) * 4 bytes ≈ 128 * (512*1024 + 1024*2048 + 512*2048) * 4 ≈ 128 * 3.1M * 4 ≈ 1.6GB This is a compute-bound problem (TF32 tensor core is 989 TFLOPS). The arithmetic intensity is about 274 / 1.6 ≈ 170 FLOP/byte. With H100's 3.35 TB/s bandwidth, we'd need about 568 TFLOPS to be compute bound. We're at 274 GFLOP total, so this is actually not huge - we're looking at individual batch items taking ~0.3ms if we hit full TF32 rate. Hmm, actually let me reconsider. We might not be able to beat cuBLAS with a custom WMMA kernel. But perhaps we can beat it with a simpler approach that avoids some overhead. Actually, wait - maybe the trick is to use the TensorFloat-32 capability which cuBLAS uses by default in PyTorch. In torch.matmul/bmm, TF32 is used by default on Ampere+ for speed, but with full precision accumulation (float32 for the output). Let me try a different approach: use CUTLASS or just call cuBLAS batched GEMM directly. Actually, the simplest thing might be to just use torch.matmul internally but with better stream handling or something... no, that won't help. Let me try a different kernel approach - use standard shared memory GEMM without tensor cores, or try to match/beat cuBLAS by using a very efficient layout. Actually, another idea: the overhead of PyTorch's bmm dispatch might be significant for this batch size. If I call cuBLAS directly, I might save some overhead. But more importantly, maybe I can transpose the data for better layout? Let me try using cublasGemmBatched or cublasGemmStridedBatched directly from the CUDA extension. That's essentially what PyTorch does, but we can control it directly. Actually, wait - let me look at whether there's a simpler win. Maybe I can use `torch.mm` on a reshaped tensor, or use `torch.matmul` with different memory layouts that are faster. In PyTorch, `bmm` on contiguous tensors in NHWC or other layouts can sometimes be slower than expected. If we can ensure optimal layout... Actually, let me check: in PyTorch 2.0+, bmm uses cuBLAS under the hood. For the given sizes, it should be very fast. But on H100, TF32 gives ~989 TFLOPS theoretical, and we need practical performance close to that. Let me try a completely different approach: use CUTLASS or just a very simple kernel to see if there are PyTorch overheads we can remove. Actually, you know what, let me just try calling cublas directly. That should match or slightly beat PyTorch's bmm due to reduced overhead. But first, let me debug my WMMA kernel. The correctness issue suggests something fundamental is wrong. Actually, looking more carefully at WMMA documentation: - The `wmma::precision::tf32` fragment type requires data to be float. - But I'm doing 8x float4 loads per 256 threads, which means each thread does 32 float loads. That should be fine. Wait, I think I see a potential issue with the B tensor loading. Let me re-examine: ```cpp for (int i = 0; i < 2; i++) { int t = threadIdx.x + i * 256; int row = t / 16; int col4 = t % 16; if (row < BK) { int b_row = k_outer + row; int b_col = n_start + col4 * 4; const float4* src = (const float4*)(B_batch + (long long)b_row * n + b_col); float4* dst = (float4*)(&B_smem[row][col4 * 4]); *dst = *src; } } ``` For BN=64, each row of B has 64 floats. We use float4 loads, so each thread loads 4 contiguous floats. With 256 threads and 2 iterations per thread, we have 512 float4 loads = 2048 floats loaded. But BK=32 and BN=64, so we need 32 * 64 = 2048 floats total. That's exactly right. For A: BM=128, BK=32. Each thread loads 4 floats with float4. 256 threads * 4 iters * 4 floats = 4096 floats. 128 * 32 = 4096 total. Correct. Let me trace through for warp 0 (warp_m=0, warp_n=0): - a_row = 0, 16 - b_col = 0, 16 For i=0, j=0: - a_frag[0] loads A_smem[0:16][k_inner:k_inner+8] with stride BK+4=36 - b_frag[0] loads B_smem[k_inner:k_inner+8][0:16] with stride BN+4=68 - MMA produces C[0:16][0:16] For i=0, j=1: - a_frag[0] same - b_frag[1] loads B_smem[k_inner:k_inner+8][16:32] with stride 68 - MMA produces C[0:16][16:32] For i=1: - a_frag[1] loads A_smem[16:32][k_inner:k_inner+8] - etc. This all looks correct! So why is the output incorrect? OH! I think I see it! When I compute the block dimensions: - BM = 128 means m_tile covers 128 rows - BN = 64 means n_tile covers 64 columns - Each block has 8 warps - Each warp computes 32x32 = 1024 elements - Total per block: 8 * 1024 = 8192 elements - Total elements in block tile: 128 * 64 = 8192. Correct. But wait, the total number of blocks: - batch_size = 128 - m = 512, BM = 128, so m_tiles = 4 - n = 2048, BN = 64, so n_tiles = 32 - tiles_per_batch = 4 * 32 = 128 - total blocks = 128 * 128 = 16384 Each block gets its own output area, no overlap. So there shouldn't be race conditions either. Let me try running a minimal test kernel. Actually, I have a better idea: let me write a simpler kernel that uses direct global memory accesses to verify the basic algorithm first, before adding shared memory and WMMA. Actually, here's another thought: the issue might be that `wmma::precision::tf32` doesn't exist in the CUDA version, or it does something weird. Let me check: on H100 (sm_90), the Tensor Cores support TF32. But the precision enum might need to be exact. Actually, I just realized: on Hopper, we should use `wmma::precision::tf32` which is available. But there's also `wmma::__float_to_tf32` or similar casting. Maybe I need to explicitly cast the values? Looking at the CUDA documentation: - When using `wmma::precision::tf32`, the values loaded with `load_matrix_sync` should be in float format, and the hardware will convert them to TF32. Hmm, maybe the issue is different. Let me look at the pointer arithmetic more carefully. For `A_batch`, `A_batch + (long long)a_row * k + a_col`: - `A_batch` is a `const float*` - `a_row * k + a_col` computes the offset in floats - So `A_batch + offset` points to the correct float - Casting to `const float4*` for vectorized load This looks correct. For B_smem: ```cpp float4* dst = (float4*)(&B_smem[row][col4 * 4]); ``` `&B_smem[row][col4 * 4]` is `float*`, cast to `float4*`. Since the memory is contiguous within the row (due to [BN + 4] being a multiple of 4 for BN=64), this should be aligned. `B_smem` starts at an aligned shared memory address. Similarly for A_smem. Wait, maybe the issue is with the float4 loads. `float4` requires 16-byte alignment. `&A_smem[row][col4 * 4]` points to `col4 * 4 * sizeof(float) = col4 * 16` bytes into the row, which is 16-byte aligned. Good. `&B_smem[row][col4 * 4]` similarly: `col4 * 16` bytes into row. Also aligned. OK, let me try a completely different approach. Instead of trying to debug the complex WMMA kernel, let me write a simpler kernel using standard CUDA programming and see if that works and is competitive. Actually, I just realized something. The tolerance for fp32 is `1e-4`. The errors reported were ~0.17, which is huge. If TF32 were the cause, we'd expect ~1e-3 errors, not 0.17. So there's definitely a bug. Let me look once more at the block indexing: ```cpp int batch_idx = blockIdx.x / tiles_per_batch; int tile_idx = blockIdx.x % tiles_per_batch; int m_tile = tile_idx / n_tiles; int n_tile = tile_idx % n_tiles; int m_start = m_tile * BM; int n_start = n_tile * BN; ``` Wait, `tile_idx / n_tiles` gives m_tile, but shouldn't it be `tile_idx / m_tiles` or something? Actually no, with `tiles_per_batch = m_tiles * n_tiles`, and tile_idx going from 0 to tiles_per_batch-1: - m_tile = tile_idx / n_tiles ranges from 0 to m_tiles-1 - n_tile = tile_idx % n_tiles ranges from 0 to n_tiles-1 This gives each combination of (m_tile, n_tile) exactly once. Correct! Let me reconsider... maybe the issue is that B_smem is declared as `__shared__ float B_smem[BK][BN + 4];` but BK=32 and we're loading B_smem with row up to BK-1 = 31. OK. Oh wait! I think I might have found it. Look at the pointer arithmetic in the loads again: For B: ```cpp const float4* src = (const float4*)(B_batch + (long long)b_row * n + b_col); ``` Wait, `(long long)b_row * n + b_col` - the n here is the full n dimension of the output, which is 2048. So for each row of B (which has n elements), the offset is `b_row * n + b_col`. But B is shape (batch_size, k, n), so each row of B has n elements. `b_row` is the k-index (0 to 1023), and `b_col` is the n-index (0 to 2047). So `b_row * n + b_col` is the correct linear index for element B[batch, b_row, b_col]. This looks right. For A: ```cpp const float4* src = (const float4*)(A_batch + (long long)a_row * k + a_col); ``` A is shape (batch_size, m, k). `a_row` is m-index (0 to 511), `a_col` is k-index (0 to 1023). `a_row * k + a_col` gives A[batch, a_row, a_col]. Correct. For C: ```cpp wmma::store_matrix_sync(C_batch + (long long)c_row * n + c_col, ...) ``` C is shape (batch_size, m, n). `c_row` is m-index, `c_col` is n-index. `c_row * n + c_col` gives C[batch, c_row, c_col]. Correct. I'm going in circles. Let me just try a much simpler kernel to test. Actually, let me examine if there could be bank conflicts or alignment issues with WMMA. Wait - here's a thought. On Hopper (sm_90), the WMMA API might behave differently than on Ampere. Or maybe `wmma::memory_row_major` doesn't work the way I think. Actually, let me check one more thing: the template arguments to wmma::fragment. For matrix_a with row_major layout: ```cpp wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> ``` This represents a 16x8 tile (since A is MxK). When I load it from shared memory using `lead_dim = BK + 4 = 36`, the elements are: - (i, j) -> base + i * 36 + j For matrix_b with row_major layout: ```cpp wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> ``` Wait! Is matrix_b supposed to be loaded in row_major or col_major? In GEMM (C = A * B), A is MxK and B is KxN. For the tensor core operation, matrix_b should be KxN with row_major layout. But in WMMA, does matrix_b use a different interpretation? Actually, I think there might be an issue with how the tensor core interprets matrix_b. In the tensor core, the multiplication is essentially: - A is MxK, loaded in some format - B is KxN, but for efficiency it's often loaded as NxK and transposed, or directly as KxN Let me look up the WMMA documentation more carefully. In standard cuBLAS/cublasLt, matrix B can be in a transposed or non-transposed format. For WMMA: - `wmma::matrix_a` with `row_major`: the rows are along the M dimension, columns along K - `wmma::matrix_b` with `row_major`: the rows are along the K dimension, columns along N This should be correct for C = A * B. But wait, in some WMMA implementations, matrix_b needs to be loaded in a specific memory layout. For instance, on older architectures, matrix B was expected to be in a specific swizzled format. Hmm, actually I wonder if the issue is that WMMA on sm_90 (Hopper) uses a different fragment format or if I should use the newer `mma.sp.sync.aligned.*` PTX instructions instead of the WMMA API. Let me try a different approach: instead of using WMMA, let me write a standard CUDA kernel that just does the multiplication in a straightforward way with good shared memory and register blocking, without tensor cores. Actually, maybe I should also try a simpler approach first: just call cublasGemmStridedBatched directly. Or better yet, let me first check if there's a simpler reason for the failure. I'll write a kernel that literally just does `C = A * B` using direct element-wise computation to verify everything else is correct. Actually, you know what? Let me write the kernel using direct cublasGemmStridedBatched call and see if that helps. That's guaranteed to be correct (same as PyTorch) and might be slightly faster due to lower overhead. Actually, for cuBLAS direct calls, the speed should be essentially identical to PyTorch. The question is whether I can beat it. For large batch sizes and regular shapes, PyTorch's bMM should already be very efficient. But maybe there's an opportunity with memory layout or fusing. Let me think differently: maybe I can use CUTLASS or write a custom kernel that achieves better occupancy or better pipelining than what cuBLAS automatically configures for these sizes. Actually, looking at the problem more carefully - the sizes are batch=128, m=512, k=1024, n=2048. This is 128 independent matmuls of 512x1024 by 1024x2048. Each individual matmul is reasonably large (1M output elements, ~1 GFLOP). cuBLAS might not be perfectly tuning for 128 parallel instances of this size. A custom kernel that process multiple batches per SM might have better occupancy. But I need to get correctness first. Let me try the simplest possible custom kernel just to verify I understand pointer arithmetic: ```python __global__ void simple_bmm(float* C, float* A, float* B, int m, int k, int n) { int b = blockIdx.z; // batch int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; if (row >= m || col >= n) return; float sum = 0; for (int kk = 0; kk < k; kk++) { sum += A[b*m*k + row*k + kk] * B[b*k*n + kk*n + col]; } C[b*m*n + row*n + col] = sum; } ``` This would be very slow but correct. Let me use it to verify the basic setup, then optimize from there. Actually, looking at the error of ~0.17 on average, could it be that the kernel is not even computing correctly at all? Like maybe the memory loads are incorrect. Let me look at this one more time. Oh! The issue might be with shared memory indexing and the warp layout. Wait, I have 8 warps per block. Each warp has 32 threads. Block size is 256 = 8 * 32. Inside the kernel, I compute: ```cpp int warp_id = threadIdx.x / 32; int warp_m = warp_id / 2; // 0..3 int warp_n = warp_id % 2; // 0..1 ``` So warp 0: (0,0), warp 1: (0,1), warp 2: (1,0), warp 3: (1,1), warp 4: (2,0), warp 5: (2,1), warp 6: (3,0), warp 7: (3,1). With 4 warp_m positions and 2 warp_n positions, each warp handles a 32x32 tile. - warp_m=0 handles rows 0-31 of the block's M dimension - warp_m=1 handles rows 32-63 - warp_m=2 handles rows 64-95 - warp_m=3 handles rows 96-127 - warp_n=0 handles cols 0-31 of the block's N dimension - warp_n=1 handles cols 32-63 This covers the full 128x64 block tile. Correct. Inside each warp, the fragment splitting is: - i=0: rows 0-15 of warp's 32, i.e., rows 0-15 (or 32-47, etc.) - i=1: rows 16-31 of warp's 32 - j=0: cols 0-15 of warp's 32, i.e., cols 0-15 (or 32-47, etc.) - j=1: cols 16-31 of warp's 32 So each warp computes four 16x16 output tiles. Correct. Hmm, let me think about potential issues with the `__syncthreads()` placement or potential divergence. All threads execute the same path with no conditionals inside the k_outer loop except the initial loads, and the loads are guarded by `if (row < ...)` which means some threads might skip loads. But since these threads still reach `__syncthreads()`, that's fine. Wait, actually, all 256 threads participate in loads, and with 4*256 float4 loads for A and 2*256 for B, that's more loads than needed but with the `if` guards it should be correct. Each warp loads exactly some of the data. Actually, I want to double check: are there 256 threads per block? Yes, `dim3 block(256)`. For A: 4 iterations, each with 256 threads. That's 1024 thread-loads. Each load is a float4 (4 floats). So 4096 floats loaded total. BM*BK = 128*32 = 4096. With the guard `if (row < BM)`, the first 4 iters load rows 0..1023, but BM=128 and each row gets 8 float4 loads (32 floats). Let me check: `row = t / 8` where `t = threadIdx.x + i * 256`. So for i=0, t ranges from 0 to 255, row from 0 to 31, col4 from 0 to 7. For i=1, t from 256 to 511, row from 32 to 63. For i=2, row from 64 to 95. For i=3, row from 96 to 127. So rows 0-127 each get loaded for col4=0..7 (i.e., cols 0-31). Correct! For B: 2 iterations, 256 threads each. 512 thread-loads of float4 = 2048 floats. BK*BN = 32*64 = 2048. For i=0, row from 0 to 15, col4 from 0 to 15. For i=1, row from 16 to 31. So rows 0-31 each get loaded for col4=0..15 (cols 0-63). Correct! Now for the WMMA part... wait, could there be an issue with WMMA requiring all threads in the warp to participate? In a warp, all 32 threads should participate in load_matrix_sync and mma_sync. My code doesn't have any thread divergence in the k_inner loop. So all threads in each warp should execute the same instructions. Wait actually, all 256 threads execute the k_inner loop. Does WMMA require only threads within a warp, or does it matter that threads from multiple warps are executing in lock-step? The warp scheduler handles this automatically - different warps are independent. Let me think about this differently. Maybe the issue is that `wmma::load_matrix_sync` uses all threads in the warp to cooperatively load, and if there's shared memory bank conflict, it could corrupt the data. But bank conflicts usually just slow things down, not corrupt data. Actually, maybe the issue is with how I handle the accumulator across k_outer iterations. The `acc_frag` accumulates over k_outer, which is correct since C += A * B for each k chunk. Wait! I just noticed something. In my kernel, I'm accumulating in `acc_frag` across k_outer iterations. But I'm NOT zeroing acc_frag at the start of each k_outer iteration. The initial `wmma::fill_fragment(acc_frag, 0.0f)` is outside the k_outer loop, so the accumulation happens correctly. But wait, this means the accumulation is correct. So why would the values be wrong? One more idea: maybe there's a CUDA synchronization issue between the grid/blocks. But since each block handles a different part of the output and there's no overlap, there shouldn't be any race conditions. And each block's warps are synchronized with `__syncthreads()`. OK, I'm going to try a completely different tactic. Let me write a kernel using explicit shared memory tiled GEMM without WMMA, and use direct computation with float atomicAdd or direct writes. This will let me verify the basic indexing is correct. Actually, atomicAdd isn't needed since each thread/block writes to non-overlapping output regions. Here's my plan: write a simple block-tiled GEMM kernel: 1. Each block computes a BM x BN tile of C 2. Load a BM x BK tile of A and BK x BN tile of B into shared memory 3. Each thread computes a TM x TN tile of the output using registers 4. Write results directly to global memory This is the standard CUDA GEMM approach. Let me start with a simple version. ```cpp __global__ void bmm_kernel(float* C, float* A, float* B, int batch_size, int m, int n, int k) { // Block tile: BM x BN // Each thread computes one element const int BM = 32; const int BN = 32; const int BK = 32; int batch_idx = blockIdx.x / ((m/BM) * (n/BN)); int tile_idx = blockIdx.x % ((m/BM) * (n/BN)); int m_tile = tile_idx / (n/BN); int n_tile = tile_idx % (n/BN); int m_start = m_tile * BM; int n_start = n_tile * BN; int row = m_start + threadIdx.y; int col = n_start + threadIdx.x; // ... } ``` Actually, a simpler test: let me just check if the problem is TF32 precision by temporarily switching to `wmma::precision::tf32` and expecting a tolerance issue, or if the problem is actually a bug. Actually I just realized: when using WMMA with `wmma::precision::tf32`, the hardware automatically converts float inputs to tf32. But the accumulation is still in float32. So the max error should be on the order of ~0.001 for random [0,1] inputs, not 0.17. The error of 0.17 clearly indicates a bug, not precision loss. Let me reconsider if there could be a memory aliasing issue. The shared memory arrays are: - A_smem[128][36]: 128 * 36 * 4 = 18432 bytes - B_smem[32][68]: 32 * 68 * 4 = 8704 bytes Total: 27136 bytes. Well under the 48KB limit. No issue. Wait! 27136 bytes is under 48KB per block. With 256 threads per block and 1024 threads per SM, we could have up to 4 blocks per SM. Total shared memory per SM is 228KB = 233472 bytes. So we could fit up to 8 blocks per SM. Actually on H100/Hopper, shared memory per SM is up to 228KB. With 27136 bytes per block, that's about 8 blocks per SM if only shared memory were the constraint. But the block count is limited by thread count too. With 256 threads per block and 2048 threads per SM, we can have up to 8 blocks per SM. So shared memory is not the limiting factor. Hmm, let me think again about what could go wrong with such large errors... Actually, I think I should check if maybe the issue is related to alignment of the float4 loads from global memory. PyTorch tensors are typically aligned to at least 256 bytes, so float4 alignment should be fine. But maybe for some intermediate addresses there could be issues? Wait, since m=512, n=2048, k=1024, and we're loading at offsets like `a_row * k + a_col`, where `a_col = k_outer + col4 * 4`. When `k_outer = 0`, `a_col` starts at 0, 4, 8, etc. So float4 loads start at addresses that are multiples of 16 bytes. Good. When `k_outer > 0`, `a_col = k_outer + col4 * 4`. For k_outer = 32, we load at offsets 32, 36, 40, etc. These are also multiples of 16 if k_outer is a multiple of 4. BK=32, so all k_outer values are multiples of 32, hence multiples of 16. Good. For B: `b_col = n_start + col4 * 4`. n_start = n_tile * BN where BN=64, so n_start is a multiple of 64. `col4 * 4` gives values 0, 4, 8, ..., 60. All multiples of 4, so `b_col` values are multiples of 4. For float4, we need the address `B_batch + b_row * n + b_col` to be a multiple of 16 bytes. `b_row * n` where n=2048, so `b_row * 2048 * 4` bytes = `b_row * 8192` bytes. This is always a multiple of 16 since 8192 is divisible by 16. `b_col * 4` bytes is also a multiple of 16 when `b_col` is a multiple of 4. So all float4 loads are properly aligned. OK, alignment is fine. At this point, I'm going to try a different approach. Let me not use WMMA at all, and instead write a straightforward CUDA kernel. That way I can verify the basic algorithm works, then optimize if needed. Actually, another idea: maybe the large error is because the kernel isn't running at all and the output tensor contains uninitialized garbage? Let me check: I'm doing `auto C = torch::empty(...)` which allocates memory but doesn't initialize it. If the kernel fails silently, C will contain garbage. Wait, but PyTorch should throw if the CUDA kernel fails... unless it's running asynchronously and the error hasn't been caught yet. Actually, the correctness check might catch the async error and treat it as a mismatch rather than a runtime error. Let me verify by adding `AT_CUDA_CHECK(cudaGetLastError());` after the kernel launch. Or better yet, let me add `cudaDeviceSynchronize()` after the kernel to ensure it completes before returning. Actually, the kernel launch syntax `<<<grid, block>>>` returns immediately (asynchronous). If there's a launch failure (like invalid configuration), it might return an error code that's not checked. But more importantly, if there's a runtime error during execution (like out-of-bounds access), it would be detected when the stream is synchronized. Let me add proper error checking: ```cpp bmm_tf32_kernel<<<grid, block>>>(...); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { printf("Kernel error: %s\n", cudaGetErrorString(err)); } ``` And also synchronize: ```cpp cudaDeviceSynchronize(); ``` Actually, in PyTorch extensions, there's usually an implicit synchronization at some point. But let me be explicit. Hmm, but if the kernel was failing to launch or execute, PyTorch would typically throw a CUDA error during the correctness check. The fact that we got a numerical mismatch suggests the kernel did execute, just with wrong results. Let me try one more debugging approach. Instead of trying to fix the complex WMMA kernel, let me write a simple kernel and then try to optimize it. Here's a simple kernel that should definitely be correct: ```cpp __global__ void simple_bmm_kernel(const float* A, const float* B, float* C, int batch_size, int m, int n, int k) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total_elements = batch_size * m * n; if (idx >= total_elements) return; int batch_idx = idx / (m * n); int rem = idx % (m * n); int row = rem / n; int col = rem % n; float sum = 0.0f; for (int kk = 0; kk < k; kk++) { sum += A[batch_idx * m * k + row * k + kk] * B[batch_idx * k * n + kk * n + col]; } C[idx] = sum; } ``` This would be very slow but correct. Let me verify correctness with this, then optimize. Or better yet, let me write a simple tiled kernel that's reasonably efficient but much simpler than the WMMA one. Actually, let me take a step back and think about the overall strategy. The reference is `torch.bmm(A, B)`. On H100 with these sizes, PyTorch will call cuBLAS with TensorFloat-32 precision. cuBLAS is heavily optimized and should achieve near-peak TF32 performance. Can I actually beat cuBLAS with a custom kernel for this specific case? It's extremely difficult for pure GEMM. But maybe there are cases where I can: 1. Avoid PyTorch dispatch overhead (minor for large GEMMs) 2. Better exploit the batch structure 3. Use different precision modes 4. Better memory layout For these specific sizes (128 x 512 x 1024 x 2048), each individual GEMM is fairly large. The batch size of 128 means there are 128 independent GEMMs. Could I interleave them better to improve occupancy? Actually, another thought: maybe the reference uses FP32 accumulation but the cuBLAS default uses TF32. In PyTorch, TF32 is on by default for matmul on Ampere+ (since 1.7). This means the reference computation itself is in TF32 precision. Oh wait! If the reference is in TF32 and my kernel is also in TF32, then I should be able to match the reference exactly. But if my kernel has a bug, I'll get different results. Actually, I just realized: the tolerance is 1e-4 for FP32, which is quite strict. If the reference is computed with TF32 precision (which has ~1e-3 relative precision), then my custom kernel should produce the same approximate results. But my kernel's ~0.17 difference is way too big. OK let me write a super simple correct kernel first. Actually, looking at the numbers again: average difference of 0.168. For random inputs in [0,1], if the kernel weren't running at all, we'd see differences around 0.5 on average (since uninitialized memory would have random values, not quite but close). If the kernel produces all zeros, the difference would be about 0.5 (mean of absolute difference between random [0,1] and 0). An average difference of 0.168 is suspicious — it's too large for TF32 precision error but too small for completely wrong output. Wait, maybe the kernel is partially working? Or maybe the shared memory layout somehow causes different elements to be multiplied? Let me reconsider the possibility that the kernel isn't running at all. If the kernel fails to launch (misconfigured grid/block), PyTorch would throw. But what if some blocks fail while others succeed? Actually, one issue: I'm not checking whether m, n, k are divisible by BM, BN, BK. In this case: - m = 512, BM = 128. 512/128 = 4. OK - n = 2048, BN = 64. 2048/64 = 32. OK - k = 1024, BK = 32. 1024/32 = 32. OK So the sizes are exact multiples. Good. Let me try a different approach. I'll write a kernel that uses direct computation from global memory (no shared memory, no WMMA), just to verify everything works. Then I'll optimize from there. ```cpp __global__ void bmm_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int batch_size, int m, int n, int k) { int total_threads = gridDim.x * blockDim.x; int idx = blockIdx.x * blockDim.x + threadIdx.x; int total_elements = batch_size * m * n; for (int linear_idx = idx; linear_idx < total_elements; linear_idx += total_threads) { int batch_idx = linear_idx / (m * n); int rem = linear_idx % (m * n); int row = rem / n; int col = rem % n; float sum = 0.0f; #pragma unroll 4 for (int kk = 0; kk < k; kk++) { sum += A[(long long)batch_idx * m * k + row * k + kk] * B[(long long)batch_idx * k * n + kk * n + col]; } C[(long long)batch_idx * m * n + row * n + col] = sum; } } ``` This is straightforward and should be correct. It'll be very slow due to lack of coalescing and no shared memory, but it'll verify our understanding of the memory layout. Actually, for better coalescing, threads in a block should process adjacent columns: ```cpp // Block processes rows [block_start, block_start + block_rows) and all columns in chunks ``` Let me write a slightly better version: ```cpp __global__ void bmm_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int batch_size, int m, int n, int k) { int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; int batch_idx = blockIdx.z; if (row >= m || col >= n) return; float sum = 0.0f; #pragma unroll 4 for (int kk = 0; kk < k; kk++) { sum += A[(long long)batch_idx * m * k + row * k + kk] * B[(long long)batch_idx * k * n + kk * n + col]; } C[(long long)batch_idx * m * n + row * n + col] = sum; } ``` With block size e.g., (32, 8, 1) and grid (n/32, m/8, batch_size). This should be correct but slow. Let me use it to verify the basic algorithm.

Let me try to optimize the kernel further before submitting. The current naive kernel might not be fast enough. Let me try a shared memory tiled approach for better performance. Given the large matrix sizes (m=512, k=1024, n=2048) and batch_size=128, this is a large compute-bound problem. Let me try a more optimized approach with: 1. Tiling with shared memory 2. More threads per block for better occupancy 3. Better memory access patterns Actually, let me think about what might outcompete torch.bmm. PyTorch uses cuBLAS which is highly optimized. For our custom kernel to beat cuBLAS, we need something quite sophisticated. However, torch.bmm has some overhead from Python/PyTorch dispatch. If we can minimize that overhead and have a reasonably efficient kernel, we might still get a speedup on smaller problems. Given the sizes: - batch_size = 128 - m = 512 - k = 1024 - n = 2048 Total FLOPS = 128 * 2 * 512 * 1024 * 2048 = 274.9 GFLOPS At 67 TFLOPS FP32, this is ~4ms compute-bound. At 3.35 TB/s, reading A (256MB), B (1GB), C (512MB) = ~1.75GB, would be ~0.5ms memory-bound if only reading once. But we read A and B multiple times, so it's compute-bound. For cuBLAS, this size should be very efficient. Let me try to optimize the kernel as much as possible with shared memory tiling. Let me try a block size of 128x128 with shared memory: - BM = 128, BN = 128, BK = 8 - Each block computes a 128x128 tile of C - 256 threads per block - Each thread computes 64 output elements (8x8 or similar) Actually let me write a more sophisticated kernel with good shared memory usage.