write-cuda-gemm-kernel

Skill: Write CUDA GEMM Kernel

Purpose

Guide the agent through designing and implementing a correct, performance-aware CUDA GEMM kernel (C = alpha * A * B + beta * C) for a specific problem configuration, including decisions about tiling strategy, memory hierarchy usage, tensor core eligibility, and when to defer to cuBLAS or CUTLASS instead.

Use this when

• You need a custom GEMM or GEMM-like operation that cuBLAS does not support (fused epilogue, custom accumulation, non-standard layouts, sparse masks)
• You are implementing a batched GEMM variant with irregular batch structure
• You need to fuse the GEMM with a downstream operation (bias add, activation, quantization) and cannot tolerate the memory round-trip of a separate kernel
• You are on a constrained embedded or inference target where you control the tiling strategy precisely
• You are building a learning or research kernel and need explicit control over every memory access

Do not use this when

• Standard sgemm/hgemm/dgemm without custom epilogue: use cuBLAS (cublasGemmEx) — it will outperform any first-attempt custom kernel on all shipping hardware
• GEMM with tensor core acceleration and standard epilogues: use CUTLASS — it exposes MMA-level tiling with a composable epilogue framework that is already highly tuned
• Batch GEMM with fixed batch sizes and standard shapes: use cublasGemmStridedBatchedEx
• FP8 GEMM on Hopper: use cublasLtMatmul with FP8 descriptors or CUTLASS 3.x FP8 kernels
• The problem is memory-bound rather than compute-bound (small K relative to M and N): tiling will not help, and a simpler kernel may be better

Inputs the agent should gather first

• M, N, K: exact or typical sizes; whether they are statically known or runtime-dynamic
• dtypes: A dtype, B dtype, accumulator dtype, output C dtype (e.g., A=fp16, B=fp16, acc=fp32, C=fp16)
• Layout of A and B: row-major or column-major; leading dimension stride if non-standard
• Transpose flags: transA, transB (determines which dimension is the inner product dimension)
• Hardware target: SM architecture (SM70/Volta, SM75/Turing, SM80/Ampere, SM89/Ada, SM90/Hopper) — determines warp MMA availability, shared memory capacity, async copy support
• Precision requirements: is fp32 accumulation required for fp16 inputs, or is fp16 accumulation acceptable
• Epilogue requirements: is there an alpha/beta scaling, bias addition, activation function, quantization step, or other per-element operation to fuse
• Batch dimension: is this a single GEMM or batched; are batch strides uniform
• Tensor core eligibility: does the problem shape and dtype satisfy alignment requirements (K divisible by 8 for fp16 wmma, 16 for MMA with specific fragment sizes)