choose-tile-size-and-work-partitioning

Skill: Choose Tile Size and Work Partitioning

Purpose

Guide the agent through selecting tile sizes and work partitioning strategies for a CUDA or Triton kernel, based on shared memory budget, register pressure, occupancy targets, problem shape, and access pattern.

Use this when

• Designing a new tiled kernel and the tile size and launch configuration have not yet been chosen.
• An existing kernel has suboptimal performance and the cause may be a poor tile size, low occupancy, or inefficient work partitioning.
• The problem shape has an irregular or non-power-of-2 size that makes default tile choices potentially wasteful.
• Writing a Triton kernel where BLOCK_M, BLOCK_N, BLOCK_K, num_warps, and num_stages must be chosen.
• Profiling shows low SM utilization, high idle warp cycles, or memory throughput below roofline estimates.

Do not use this when

• Using a library (cuBLAS, cuDNN, CUTLASS with auto-tuning) that handles tile size selection internally. Trust the library's tuner unless profiling shows a clear gap.
• The kernel is purely streaming (one pass, no reuse) and tiling provides no shared memory reuse benefit. In that case, work partitioning reduces to choosing a block size that achieves good occupancy and coalesced access, which is a simpler problem.
• The problem is so small (total work fits in one or two blocks) that tile size selection is irrelevant compared to launch overhead.

Inputs the agent should gather first

• Problem dimensions (e.g., M, N, K for GEMM; sequence length and head dimension for attention; reduction length for softmax).
• Element dtype: determines element size in bytes (fp16 = 2B, fp32 = 4B, bf16 = 2B, int8 = 1B).
• Target GPU architecture: shared memory capacity per SM (48 KB–228 KB depending on architecture and configuration), L2 size, number of SMs, warp size (always 32 for CUDA).
• Maximum shared memory per block (for sm_86: up to 100 KB with cudaFuncSetAttribute(f, cudaFuncAttributeMaxDynamicSharedMemorySize, ...)).
• Number of registers available per SM (65536 for most modern architectures), and target occupancy.
• Whether the problem shape is static or dynamic at kernel launch time.
• Whether the input access pattern is coalesced in the tile dimension being chosen.
• For Triton: whether autotuning will be used to sweep tile sizes, or if a fixed tile must be chosen.