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.

Required reasoning process

1. Determine what is loaded into shared memory per block. For a GEMM-style kernel with tiles BLOCK_M x BLOCK_N and accumulation tile BLOCK_K:

   • smem per block = (BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N) * element_size
   • For an elementwise kernel with no smem reuse, smem per block = 0 (tiles reside in registers only).
   • For a softmax kernel (row-wise): smem per block = BLOCK_N * element_size (one row of the softmax).

2. Compute smem usage for candidate tile sizes. Start from the hardware smem limit and work downward. A common starting constraint:

   • smem per block <= 48 KB (safe default, no cudaFuncSetAttribute needed)
   • smem per block <= 100 KB (requires explicit dynamic smem request on sm_80+)
   • smem per block <= 228 KB (sm_90 Hopper maximum with smem partitioning)
   • Reject tile sizes that exceed the smem budget.

3. Estimate register usage per thread. A precise count requires compiling and inspecting with --ptxas-info-registers or cuobjdump. As a first estimate:

   • Each live fp32 value requires one register. Each live fp16 value typically still uses one register.
   • Tile accumulation in a GEMM with a BLOCK_M x BLOCK_N output tile and 128-bit vectorized loads requires roughly (BLOCK_M * BLOCK_N) / 32 registers per thread for the accumulator alone (assuming warp-level decomposition).
   • Total register pressure = accumulator + input tiles in registers + loop induction variables + addressing offsets. Budget ~4–8 registers for overhead per active load.
   • Registers per SM = 65536. At 32 threads per warp and W warps per block: registers per thread budget = 65536 / (W * 32 * max_blocks_per_SM).

4. Compute occupancy for candidate configurations. Occupancy = (active warps per SM) / (max warps per SM). Constraints that limit occupancy:

   • smem: max_blocks_per_SM = floor(smem_per_SM / smem_per_block). Warps = max_blocks_per_SM * warps_per_block.
   • registers: max_blocks_per_SM = floor(regs_per_SM / (regs_per_thread * threads_per_block)). Then same warp calculation.
   • hardware block limit: typically 32 blocks per SM maximum (sm_86), sometimes 16.
   • Take the minimum across all constraints. This is the achieved occupancy.
   • Use cudaOccupancyMaxActiveBlocksPerMultiprocessor (CUDA) for accurate computation after writing the kernel.

5. Check warp efficiency in the last tile. If the problem dimension is not a multiple of the tile size, the last block has fewer valid elements than the tile size. The fraction of useful work in that block = (N % TILE_SIZE) / TILE_SIZE. If this fraction is very small (e.g., < 0.1) and many blocks hit this case, the waste is significant. Pad to a tile multiple or choose a smaller tile size.

6. Check alignment and divisibility for vectorized loads. Vectorized loads (float4, half2) require the pointer to be aligned to the vector width. For a tile of width W loaded with 128-bit (float4) loads, W must be a multiple of 4 (for float4). If W is not aligned, scalar fallback loads are needed, which reduces memory throughput.

7. Choose the work partitioning shape. Options:

   • 1D grid: one block per output vector (row). Simple. Good for row-wise ops (softmax, layernorm). Poor if the number of rows is small relative to SM count.
   • 2D grid: one block per output tile (block_row, block_col). Standard for GEMM, attention. Grid dimensions are ceil(M/BLOCK_M) x ceil(N/BLOCK_N).
   • 3D grid: adds a batch dimension. Standard for batched GEMM, multi-head attention. Grid = ceil(M/BLOCK_M) x ceil(N/BLOCK_N) x batch.
   • Persistent kernel: one block per SM; each block loops over multiple tiles. Reduces launch overhead, improves load balance for irregular shapes. Required for reduction-then-broadcast patterns where a second kernel launch would be expensive.
   • Streaming decomposition: for problems too large to hold all intermediate state, decompose along the K or sequence dimension and accumulate across kernel calls. Used in multi-pass attention or tiled reduction.

8. For Triton: choose num_warps and num_stages.

   • num_warps: number of warps per program instance. For compute-bound kernels, 4–8 warps is typical. For memory-bound kernels, 1–2 warps can be sufficient.
   • num_stages: pipeline stages for the software prefetch loop. More stages hide memory latency but consume more smem for the prefetch buffers. For A100-class hardware, 3–5 stages is a common range. Verify smem budget: each additional stage adds one tile's worth of smem.
   • Use triton.testing.Benchmark to sweep configurations and confirm empirically.

Kernel design rules

• Tile sizes must be powers of 2 in each dimension when using shared memory and warp-level reductions, to avoid non-uniform warp partitioning and bank conflicts.
• The tile width (innermost dimension) must be a multiple of 32 (one warp) to avoid underutilizing warp lanes. A tile width of 16 means half the warp is idle at each step.
• For GEMM-style kernels, BLOCK_M, BLOCK_N, and BLOCK_K must individually not exceed the smem budget divided by the number of tiles resident in smem simultaneously.
• For Triton, BLOCK_SIZE must be a tl.constexpr to enable compile-time specialization. Do not pass it as a dynamic argument unless the kernel explicitly handles dynamic tile sizes.
• Do not choose tile sizes that result in a grid with fewer blocks than the number of SMs on the target GPU — this leaves SMs idle for the entire kernel duration.
• For persistent kernels, distribute work to blocks in tile index order, not in arbitrary order, so that adjacent tiles share L2 cache lines when possible.
• When using Triton autotuning, define the autotune search space explicitly. Do not leave it as a fully open search — constrain to architecturally valid combinations (e.g., smem budget respected, num_warps divides threads evenly).

Correctness requirements

• The tile decomposition must cover all output elements exactly once. Verify that ceil(N/TILE_SIZE) * TILE_SIZE >= N (no output position is skipped) and that no two tiles write to the same output position.
• Partial tiles at the boundary must be masked correctly. Refer to the handle-boundary-conditions skill for the masking strategy.
• For reductions split across tiles or kernel calls, the accumulation across tiles must produce the same result as a single-pass reduction. Verify with small test cases where the split is at a tile boundary.
• Thread assignments within a block must not alias: two threads must not read from or write to the same shared memory location simultaneously without synchronization.

Performance requirements

• Target at least 50% theoretical occupancy as a starting point. Below 25% occupancy, latency hiding is severely limited. Above 75% occupancy, returns on further increases are often marginal — prioritize smem and arithmetic throughput.
• For memory-bound kernels, the dominant metric is memory throughput, not occupancy. A low-occupancy kernel that achieves high memory bandwidth is preferable to a high-occupancy kernel with poor coalescing.
• For compute-bound kernels (GEMM), maximize the ratio of arithmetic to memory traffic. Large tiles improve reuse but increase register and smem pressure — find the crossover point empirically.
• Document the roofline position of the kernel: compute the arithmetic intensity (FLOPs / bytes) and compare to the hardware's compute-to-bandwidth ratio. This determines whether the bottleneck is memory or compute, and thus what tile size changes will help.
• State the expected occupancy for the chosen tile size configuration and the hardware-imposed constraint responsible for limiting it.

Output format

The agent should produce:

1. Tile size analysis table: for each candidate tile size configuration, a table showing: smem per block, estimated registers per thread, occupancy estimate, warp efficiency on the last tile, alignment compatibility.
2. Selected configuration with justification: the chosen tile size(s) and launch configuration, with explicit reasoning for why it was selected over alternatives.
3. Work partitioning description: the grid decomposition (1D/2D/3D/persistent), grid dimensions as a function of problem size, and block/warp dimensions.
4. Kernel implementation skeleton: the launch configuration and thread-to-data mapping code, even if the full kernel is not yet written.
5. Autotuning plan (if applicable): the parameter space to sweep and the metric to optimize (throughput, latency).

Common failure modes

• Tile width not a multiple of 32 (warp size): a tile width of 16, 24, or 48 leaves some warp lanes idle during each step, reducing warp efficiency. Always round tile widths up to a multiple of 32.
• smem exceeds hardware limit without explicit request: using more than 48 KB of shared memory per block on CUDA without calling cudaFuncSetAttribute causes the kernel launch to fail silently (0 blocks launched) or error.
• Grid smaller than SM count: launching fewer blocks than there are SMs on the GPU means some SMs are idle for the entire kernel. This is a hard floor on performance.
• Register pressure from large tile sizes: choosing BLOCK_M=128 x BLOCK_N=128 in a Triton GEMM requires a 128x128 register accumulator tile, which exceeds the register file on most hardware. The compiler will spill to local memory. Check the compiled kernel's register count.
• num_stages too large in Triton: each pipeline stage adds one tile's smem. With BLOCK_K=64 and fp16 (2B), each additional stage adds 64*2B = 128B per warp. For large BLOCK_M/BLOCK_N tiles, multiple stages can push over the smem budget.
• Persistent kernel with uneven work distribution: if the number of tiles is not evenly divisible by the number of blocks (SMs), some blocks will process one extra tile. This is acceptable but must be handled correctly — the last few blocks must not process past the tile count.
• Choosing tile size to avoid boundary handling rather than for performance: using BLOCK_SIZE=prime-number-close-to-N to make the problem fit in one tile is not a strategy; it produces different suboptimal behavior for different input shapes.
• Ignoring L2 reuse in tile ordering: for a 2D GEMM grid, iterating over (row, col) in row-major order reuses the A matrix tile across columns. Iterating in column-major order reuses the B matrix tile. Choosing the wrong iteration order reduces L2 hit rate.

Review checklist

• Shared memory per block is computed for the chosen tile size and verified to fit within the hardware limit (with or without explicit dynamic smem request as appropriate).
• Register pressure is estimated (or measured after compilation) for the chosen configuration.
• Occupancy is computed for the chosen configuration and the limiting resource is identified.
• The grid covers all output elements exactly once.
• The tile width is a multiple of 32 (warp size) in the innermost dimension.
• Partial boundary tiles are handled correctly (refer to handle-boundary-conditions skill).
• The grid size is at least equal to the number of SMs on the target GPU.
• For Triton: num_warps and num_stages are chosen with smem budget verified.
• The arithmetic intensity of the kernel is computed and the roofline position is stated.
• The chosen tile size is justified by occupancy, smem reuse, or register analysis — not just by convention.
• An autotuning or benchmarking plan is included if the tile size is not verified empirically.