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)

## Required reasoning process
1. **Decide whether to write a custom kernel at all.** If cuBLAS or CUTLASS covers the use case, state that explicitly and stop. Only proceed if there is a concrete reason a custom kernel is necessary.

2. **Determine tensor core eligibility.** For SM70+: WMMA API requires M/N/K tiles of 16x16x16 (fp16) or 8x32x16 / 32x8x16 variants. For SM80+: PTX MMA instructions offer more shapes. Requirements: fp16 or bf16 inputs, fp32 accumulator, K dimension divisible by the fragment K size, A and B pointers 16-byte aligned. If these are not met, fall back to scalar (SIMT) accumulation and document why.

3. **Choose the tiling hierarchy.** A standard 3-level tiling strategy:
   - **Thread block tile** (BM x BN x BK): typically 128x128x8 or 64x64x16 or 128x256x32 depending on target occupancy and register budget. Larger tiles = more reuse but fewer active blocks.
   - **Warp tile** (WM x WN): subdivides the thread block tile among warps. Each warp handles a contiguous WM x WN sub-tile.
   - **Thread tile / register tile**: each thread accumulates a small sub-matrix in registers (e.g., 8x8 for SIMT, or one MMA fragment pair for tensor cores).
   Justify the chosen tile sizes given smem capacity and register pressure.

4. **Design the shared memory layout.** Load A tile as BM x BK and B tile as BK x BN into shared memory. Determine whether row-major or column-major smem layout reduces bank conflicts for the warp access pattern. Add padding columns/rows to eliminate bank conflicts if needed (typically +1 or +2 columns).

5. **Plan the load pattern.** Each thread must participate in loading the A and B tiles collaboratively. Assign contiguous global memory addresses to contiguous threads to ensure coalesced loads. Use vectorized loads (float4, uint4, int4) when alignment and size allow. On SM80+, use `cp.async` (or `cuda::memcpy_async`) to overlap global loads with computation via double buffering.

6. **Design double buffering.** Allocate two smem buffers for A and two for B. While computing on buffer[ping], issue async loads for buffer[pong]. This hides global memory latency behind arithmetic. Requires `cp.async.wait_group` or `__pipeline_commit` / `__pipeline_wait_prior` fence placement to be correct.

7. **Write the inner loop.** The K-loop iterates over BK-wide slices. Each iteration: load tile into smem, sync, compute warp-level MMA or thread-level outer product, sync before next load. If double buffering: restructure so loads and computes are interleaved without extra syncs blocking overlap.

8. **Handle boundary conditions.** For M % BM != 0, N % BN != 0, K % BK != 0: guard loads with bounds checks. Either pad tensors to multiples of tile size externally, or predicate the load/store with per-thread index checks. Predicated loads should return 0 for out-of-bounds positions so partial tiles accumulate correctly.

9. **Write the epilogue.** Apply alpha scaling to the accumulator, load the C tile, apply beta scaling, add, then store. If beta == 0, skip the C load entirely. If fusing a bias or activation, apply it here in registers before the final store. Use vectorized stores when alignment allows.

10. **Estimate register pressure.** The accumulator array dominates. For a thread tile of TM x TN with fp32 accumulators, that is TM*TN registers. Add registers for A fragments and B fragments in the inner loop. Total must stay under the per-thread register limit (255 on current hardware). If over, reduce tile size or annotate `__launch_bounds__`.

## Kernel design rules
- Always accumulate in fp32 when inputs are fp16 or bf16 unless the application explicitly accepts fp16 accumulation and you have verified it does not degrade accuracy for the problem.
- Shared memory layout must be chosen to avoid 4-way or higher bank conflicts for the inner loop access pattern. Verify the warp access pattern against smem bank indices before finalizing.
- Thread block dimensions must be a multiple of 32 (warp size). Common choices: (128,1), (256,1), or (16,16). For tensor core kernels, warp-level tile assignment must align with MMA fragment shapes.
- Do not use `__syncthreads()` inside the inner compute loop if it is not needed. Each unnecessary sync stalls all warps in the block.
- Use `__ldg()` or `__restrict__` on read-only pointers to hint L1 cache for non-tiled accesses.
- The epilogue must handle alpha == 0 (output is all zeros) and beta == 0 (C input is unused) as fast paths to avoid unnecessary loads.
- For WMMA kernels: use `wmma::load_matrix_sync`, `wmma::mma_sync`, `wmma::store_matrix_sync` with correct layout specifiers (`WMMA_ROW_MAJOR` vs `WMMA_COL_MAJOR`). Fragment layouts must match the smem layout used.
- For PTX MMA kernels (SM80+): use `mma.sync.aligned.m16n8k16` (or appropriate variant) only when the surrounding tile structure ensures correct data placement. Errors in fragment lane mapping are silent and produce wrong results.

## Correctness requirements
- **Indexing**: every global A[row][col] access must correctly account for transA, leading dimension (lda), and the thread's position in the tile hierarchy. Verify the formula: `A[blockRow*BM + warpRow*WM + threadRow][k_base + k_offset]` maps to the correct flat index `(blockRow*BM + warpRow*WM + threadRow) * lda + (k_base + k_offset)`.
- **Boundary guards**: threads that map outside [0, M) x [0, K) for A or [0, K) x [0, N) for B must not issue out-of-bounds loads. Predicate these with explicit if-checks or masked vector loads.
- **Shared memory synchronization**: every smem tile write (by all threads in the block) must be followed by `__syncthreads()` before any thread reads from that tile. Every read must complete before the next write overwrites the buffer.
- **Accumulator initialization**: all accumulator registers must be initialized to 0 before the K-loop. For WMMA fragments: call `wmma::fill_fragment(acc_frag, 0.0f)`.
- **Alpha/beta application**: apply alpha to the accumulator result, not to partial sums during the K-loop. Apply beta to the loaded C value, not to the accumulator.
- **Dtype conversion**: when loading fp16 into shared memory and accumulating in fp32, the conversion to fp32 must happen before the fused multiply-add, not after.
- **Vectorized load alignment**: float4 loads require 16-byte aligned addresses. Verify that the base pointer plus per-thread offset is always a multiple of 16 bytes before using `float4` loads. Misaligned vector loads are undefined behavior.
- **Output write bounds**: the store to C must also be predicated for threads outside the valid M x N range.

## Performance requirements
- State the theoretical peak FLOP/s for the target SM and dtype (e.g., SM80 A100: ~312 TFLOPS fp16 with tensor cores, ~19.5 TFLOPS fp32 SIMT).
- Estimate arithmetic intensity for the chosen tile size: for a BM x BN x BK tile, AI = 2*BM*BN*BK / (2*(BM*BK + BK*BN) * sizeof(dtype)) in FLOPs/byte. A well-tiled GEMM should be compute-bound for large M, N, K.
- Target minimum 70-80% of cuBLAS throughput as a bar for "good enough to ship." If unable to reach 50%, defer to cuBLAS.
- Occupancy: target at least 50% theoretical occupancy (2+ active blocks per SM, or 4+ active warps per SM beyond the minimum). Compute `(smem_per_block + register_per_block)` consumption and check against SM limits.
- Double buffering should be used for all SM80+ kernels targeting high arithmetic intensity. On SM70 and SM75, double buffering is possible without `cp.async` using register prefetch or manual smem ping-pong.
- Vectorized loads (float4 for fp32, int4 for int8, float2 for fp16 pairs) should be used wherever alignment guarantees hold. This reduces instruction count and improves memory throughput.
- The inner loop (MMA or outer product) must be structured to maximize instruction-level parallelism. Avoid data dependencies between consecutive MMA calls if possible.

## Output format
The final response must include:
1. **Decision rationale**: whether a custom kernel is appropriate, or whether cuBLAS/CUTLASS should be used instead — with explicit reasoning.
2. **Tile configuration table**: BM, BN, BK, WM, WN, TM, TN (or MMA fragment shape), thread block dimensions, grid dimensions, smem size per block.
3. **Memory layout diagram** (ASCII or description): how A and B tiles are laid out in smem, including any padding.
4. **Kernel code**: complete, compilable CUDA kernel with all helper functions (device functions for warp reduction, fragment load, etc.). No pseudocode.
5. **Host launch code**: grid/block dims, smem size, kernel call with all parameters.
6. **Correctness notes**: explicit list of boundary conditions handled and how.
7. **Performance notes**: expected arithmetic intensity, estimated occupancy, known bottlenecks.
8. **Known limitations**: shapes or hardware where this kernel will underperform or is incorrect.

## Common failure modes
- **Shared memory bank conflicts in the inner loop**: accessing smem with a stride equal to 32 (or a multiple) causes 32-way conflicts. Fix: add +1 or +4 padding to the smem tile column dimension.
- **Wrong tile boundary handling**: tiles at the M or N edge that are smaller than BM/BN require predicated loads. Forgetting this causes out-of-bounds reads and incorrect results on non-round shapes.
- **Misaligned `float4` loads**: base pointer + per-thread byte offset is not 16-byte aligned. This causes a misaligned memory access error at runtime. Fix: check alignment statically or fall back to scalar loads for the last partial vector.
- **Register spilling**: tile sizes are too large for the register file. The compiler spills to local memory, causing 20-100x slowdown on accesses. Fix: reduce TM x TN, add `__launch_bounds__`, or check ptxas output for `spill stores/loads`.
- **Incorrect alpha/beta**: beta is applied to the accumulator instead of C, or alpha is not applied at all, or the accumulator is stored before alpha scaling.
- **Missing `__syncthreads` after smem load**: threads in the block race: some threads start reading A/B tiles from smem while other threads are still writing them. Always sync after the collaborative load phase.
- **Wrong transposition handling**: when transA=true, the A tile loads from the transposed dimension. A common bug is using the wrong stride (lda vs K) or swapping row/col indices without updating the smem layout.
- **WMMA layout mismatch**: loading a row-major A fragment with `WMMA_COL_MAJOR` specifier or vice versa produces silently wrong results. Carefully match smem layout to the fragment load specifier.
- **Double buffer fence errors**: `cp.async.wait_group(N)` leaves N async groups in flight. Using N=0 blocks until all are done; using N=1 allows one group to still be in flight. Wrong N causes reading data before it has arrived from global memory.

## Review checklist
- [ ] Have I verified that cuBLAS or CUTLASS cannot cover this use case? If yes, is the reason documented?
- [ ] Are M, N, K boundary conditions handled for non-multiples of the tile size?
- [ ] Are all smem accesses free of bank conflicts, or has padding been added and verified?
- [ ] Is every smem write followed by `__syncthreads()` before the corresponding read?
- [ ] Are float4/float2 loads used only on verified-aligned pointers?
- [ ] Is the accumulator initialized to zero before the K-loop?
- [ ] Is alpha applied to the accumulator and beta applied to the C load, not the other way around?
- [ ] Does the epilogue skip the C load when beta == 0?
- [ ] Are output stores predicated for threads outside the valid [0,M) x [0,N) range?
- [ ] For WMMA kernels: do fragment layout specifiers match the actual smem layout?
- [ ] For double buffering: are `cp.async` commit and wait fences placed correctly?
- [ ] Has ptxas output been checked for register spills?
- [ ] Has the kernel been validated against a reference (cuBLAS or CPU GEMM) on at least: square shapes, rectangular shapes, non-power-of-two shapes, M=1 (vector-matrix), K=1?
- [ ] Is the launch configuration correct (grid covers all M x N outputs, block dims are warp multiples)?