Optimize Triton Block Parameters

Skill: Optimize Triton Block Parameters

Purpose

Guide the agent through the systematic process of choosing and tuning block size parameters in Triton kernels — BLOCK_M, BLOCK_N, BLOCK_K for GEMM-style kernels; BLOCK_SIZE for reduction and pointwise kernels; and the associated num_warps and num_stages values that control parallelism and pipeline depth. This is an optimization skill, not a write-kernel skill. It assumes a correct kernel exists and asks: what configuration makes it fast?

Use this when

• A Triton kernel is functionally correct and you need to choose launch configuration parameters to maximize throughput or minimize latency.
• You are using triton.autotune and need to define a principled config search space rather than guessing random power-of-2 combinations.
• The kernel performs well on one shape but degrades on another (e.g., trained config for large M/N performing poorly on small M/N).
• You are tuning for a specific hardware target (A100, H100, RTX 4090) and need architecture-aware config decisions.
• You suspect shared memory pressure, register pressure, or pipeline stalls are limiting achieved throughput and want a systematic approach to diagnose and fix them.

Do not use this when

• The kernel has a correctness bug. Fix correctness first. Tuning a broken kernel wastes time and produces misleading measurements.
• The bottleneck is not the Triton kernel itself — for example, Python overhead, data transfer, or memory allocation are dominating the profile. Identify the actual bottleneck before tuning kernel parameters.
• You are running on hardware with radically different SM architecture (e.g., porting a config optimized for A100 to a T4 or V100 without re-profiling). Configs do not transfer across GPU generations without re-measurement.
• The kernel is not on a hot path and a 2x improvement in its execution time has negligible end-to-end impact. Do not optimize prematurely.

Inputs the agent should gather first

Before tuning, confirm:

1. Kernel type — GEMM, reduction (softmax/layernorm), pointwise/elementwise, or attention. Each has different SRAM footprint models.
2. Hardware target — GPU model, SM count, max shared memory per SM (A100: 164KB, H100: 228KB), max registers per SM, L2 cache size.
3. Problem shapes — the actual M/N/K or sequence length / hidden dim values at runtime. If these vary, identify the representative shape for the primary deployment case.
4. Input dtype — fp16, bf16, fp32, or mixed. Affects register usage per element and tensor core eligibility.
5. Whether autotuning is permitted — production kernels that ship to end users sometimes need a fixed config for reproducibility. If autotuning is used, clarify whether the autotuned cache is shipped with the code.
6. Baseline measurement — what is the current achieved throughput (GB/s for memory-bound, TFLOP/s for compute-bound)? Without a baseline, you cannot know whether tuning helped.

Required reasoning process

Step 1: Classify the kernel as memory-bound or compute-bound

Compute the arithmetic intensity (FLOPs per byte of memory traffic) for your kernel at the target problem size.

• For a GEMM of shape (M, N, K) with fp16: FLOPs = 2MNK, bytes = 2(MK + KN + M*N). Intensity = FLOPs / bytes.
• For A100: fp16 tensor core peak ~312 TFLOP/s, HBM bandwidth ~2 TB/s. Ridge point: 312e12 / 2e12 = 156 FLOPs/byte. If your kernel's arithmetic intensity < 156, it is memory-bound; otherwise compute-bound.
• For reduction kernels (softmax, layernorm): these are typically memory-bound at most problem sizes. Intensity is approximately 2-3 ops per byte.

This classification determines which parameter matters most: memory-bound kernels benefit from larger tiles that improve cache reuse; compute-bound kernels benefit from tile shapes that maximize tensor core occupancy.

Step 2: Estimate shared memory (SRAM) usage per block

For GEMM-style kernels, the dominant SRAM consumer is the A and B tiles awaiting the dot product:

smem_bytes = BLOCK_M * BLOCK_K * sizeof(dtype)   # A tile
           + BLOCK_K * BLOCK_N * sizeof(dtype)   # B tile
           + (num_stages - 1) * above             # pipeline double/triple-buffering

For softmax/layernorm, SRAM is used implicitly by Triton for tiled loads. The dominant cost is register pressure from fp32 accumulators.

Check: smem_per_block * max_concurrent_blocks <= SM_smem_capacity. On A100 with 164KB, a config using 96KB per block allows at most 1 concurrent block per SM — acceptable for compute-bound kernels, but starves occupancy for memory-bound ones.

Step 3: Estimate register pressure

Triton does not expose register counts directly. Use these heuristics:

• Each fp32 element in a tile held in registers costs 1 register (4 bytes).
• For GEMM, the accumulator tile alone is BLOCK_M * BLOCK_N * 4 bytes. For 128x128: 64KB of registers. A100 has 256KB of registers per SM, shared across all resident warps.
• High register pressure reduces occupancy. Use triton-bench or Nsight Compute to get actual register counts after compilation.

Step 4: Enumerate valid configurations

A valid configuration must satisfy:

• BLOCK_M, BLOCK_N, BLOCK_K are powers of 2.
• For tl.dot to use tensor cores: BLOCK_K >= 16, and both matrix dimensions >= 16.
• BLOCK_SIZE for reduction kernels: power of 2, typically between 64 and 4096.
• smem per block <= SM smem capacity (usually set to 164KB on A100 via triton.compiler hints or capped by hardware limits).
• num_stages >= 1. Typical range: 2 to 5. Higher values require more SRAM for double-buffering.
• num_warps: 1, 2, 4, or 8. More warps improve latency hiding but consume more registers per block. For small tiles (BLOCK_M=32), num_warps=4. For large tiles (BLOCK_M=128), num_warps=8.

Step 5: Define the autotune config space

For a GEMM kernel, a representative search space:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'num_stages': 3, 'num_warps': 8}),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 32, 'num_stages': 3, 'num_warps': 8}),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'num_stages': 4, 'num_warps': 4}),
        triton.Config({'BLOCK_M':  64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'num_stages': 4, 'num_warps': 4}),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N':  64, 'BLOCK_K': 32, 'num_stages': 4, 'num_warps': 4}),
        triton.Config({'BLOCK_M':  64, 'BLOCK_N':  64, 'BLOCK_K': 64, 'num_stages': 4, 'num_warps': 4}),
        triton.Config({'BLOCK_M':  32, 'BLOCK_N':  32, 'BLOCK_K': 32, 'num_stages': 2, 'num_warps': 4}),
    ],
    key=['M', 'N', 'K'],
)

The key argument is critical. It specifies which problem dimensions trigger a re-autotune. Include all shape dimensions that affect tile assignment or smem usage. Omitting a key dimension causes the autotuner to reuse a cache entry from a different shape.

For reduction kernels, the config space is simpler (only BLOCK_SIZE and num_warps vary). Start with:

configs=[
    triton.Config({'BLOCK_SIZE': 1024, 'num_warps': 4}),
    triton.Config({'BLOCK_SIZE': 2048, 'num_warps': 4}),
    triton.Config({'BLOCK_SIZE': 4096, 'num_warps': 8}),
],
key=['D'],

Step 6: Benchmark and interpret

Use triton.testing.do_bench:

ms = triton.testing.do_bench(lambda: my_kernel(a, b, c), warmup=25, rep=100)
gbps = (total_bytes / 1e9) / (ms / 1e3)
tflops = (total_flops / 1e12) / (ms / 1e3)

Compare achieved GB/s against device peak HBM bandwidth (for memory-bound) or achieved TFLOP/s against device peak (for compute-bound). Efficiency = achieved / theoretical. Aiming for > 70% efficiency is a reasonable target before claiming the kernel is well-tuned.

Kernel design rules

• Never use a non-power-of-2 BLOCK_SIZE. Triton's vectorized load logic and tl.dot dispatch both require this.
• Always declare tile size parameters as tl.constexpr. Variable tile sizes prevent compile-time specialization and eliminate tensor core dispatch.
• num_stages controls the number of in-flight async loads. Setting it too high (e.g., 5 or 6) on A100 will exceed smem limits for large tile sizes, causing silent register spilling or kernel launch failure. Compute the smem budget explicitly.
• For kernels with irregular shapes (M or N not a multiple of BLOCK_M or BLOCK_N), test at boundary shapes (e.g., M=17, M=63, M=129) during tuning to confirm the config handles them correctly and with acceptable performance.
• Autotune key must include all shape dimensions that change the optimal config. If your kernel is called with both (M=128, K=4096) and (M=4096, K=128), these should produce different cache entries.
• Do not autotune on the first call in production. The autotuner measures multiple configs and takes several hundred milliseconds. Cache the result and warm it before production deployment.

Correctness requirements

• Changing block parameters must never change the output. Verify with torch.allclose at multiple shapes before and after tuning.
• When testing boundary shapes during tuning, confirm that masking logic is correct. A config with BLOCK_M=128 on a problem with M=64 must still produce correct results — masking must cover the full tile.
• Do not introduce tl.constexpr tile sizes that assume specific input alignment. If BLOCK_K=64 but K=48, the masking logic must handle the last partial block correctly at any BLOCK_K.

Performance requirements

The agent must reason through the following:

• Tile size and L2 reuse. Larger BLOCK_N increases the reuse of A-tile data within a single program instance. Larger BLOCK_M increases reuse of B-tile data. For memory-bound GEMMs, maximizing these reduces effective memory traffic.
• num_stages and latency hiding. Async loads in Triton (on A100/H100) have a latency of ~100 cycles. With num_stages=3, the compiler can issue 3 loads ahead of the computation, hiding most of this latency. Fewer stages leave the tensor cores stalling on data. More stages use more SRAM.
• num_warps and instruction-level parallelism. Each warp issues instructions independently. With 4 warps per block, if one warp stalls on a load, the other 3 can continue issuing compute instructions. 8 warps provide more ILP but consume more register file budget per SM.
• Occupancy. concurrent_blocks_per_SM = min(smem_capacity // smem_per_block, register_budget // registers_per_block, hardware_max_blocks_per_SM). For compute-bound kernels, 1-2 blocks per SM is fine (smem can be large). For memory-bound kernels, you want more concurrent blocks to hide memory latency — target >= 4 blocks per SM.
• Shape-specific tuning. A config optimal for square large GEMMs (M=N=K=4096) will be suboptimal for tall-thin GEMMs (M=1, N=4096, K=4096). Do not assume one config generalizes across shapes. Use autotune with shape-specific keys.

Output format

The agent should produce:

1. A table of candidate configurations with estimated smem usage per config and whether it fits within SM smem limits.
2. The @triton.autotune decorator with a justified config list and correct key parameters.
3. Benchmark code using triton.testing.do_bench that reports achieved GB/s or TFLOP/s alongside theoretical peak.
4. A final recommendation stating the winning config, the achieved efficiency percentage, and whether further tuning is expected to yield meaningful gains.
5. Documentation of any eliminated configs and why (e.g., smem overflow, non-power-of-2, below tensor core threshold).

Common failure modes

• Autotuning with too few configs. A list of 2-3 configs often misses the performance peak. For GEMM, the optimal config depends on the balance of M, N, K and the L2 capacity. Include at least 6-8 configs spanning different BLOCK_M/BLOCK_N/BLOCK_K ratios.
• Non-power-of-2 BLOCK_SIZE. Using BLOCK_SIZE=96 or BLOCK_SIZE=192 will cause Triton to reject the config or silently fall back to a scalar path. All tile dimensions must be powers of 2.
• num_stages too high. Setting num_stages=5 with BLOCK_M=128, BLOCK_N=256, BLOCK_K=64 on A100 requires 5 * (128*64 + 64*256) * 2 bytes = 5 * 49KB = 245KB of smem, exceeding the 164KB limit. The kernel fails at compile time or produces incorrect results if limits are not enforced.
• Tuning on wrong problem shape. Autotuning with M=N=K=4096 and deploying on inference with M=1 gives a config optimized for large square GEMMs. The small-M case typically needs much smaller BLOCK_M (32 or 64) to achieve any parallelism.
• Missing shape dimensions in autotune key. If M and N are included but K is omitted, the same config is reused for K=64 and K=4096. The optimal BLOCK_K and num_stages differ substantially between these cases.
• Interpreting autotune time as kernel time. The first call to an autotuned kernel runs all configs in sequence. This takes much longer than a single kernel call. Do not log or report this time as representative kernel latency.
• Comparing achieved TFLOP/s against peak without accounting for dtype. A100's fp16 tensor core peak is ~312 TFLOP/s. fp32 GEMM peak is ~19.5 TFLOP/s. Comparing your fp32 kernel against the fp16 peak and claiming 6% efficiency is a comparison error.
• Not testing boundary shapes after tuning. A config optimized for M=4096 may fail or produce NaN for M=17 if the masking logic has a subtle bug that is exposed only by certain M % BLOCK_M values.

Review checklist

• All BLOCK_* parameters are powers of 2.
• All BLOCK_* parameters are declared tl.constexpr in the kernel signature.
• For tl.dot usage: BLOCK_K >= 16 and both matrix-multiply dimensions >= 16.
• smem per block is computed explicitly for each config and verified against SM capacity.
• num_stages budget for smem is included in the smem estimate (multiply tile smem by num_stages).
• Autotune key includes all shape dimensions that affect tile geometry or smem usage.
• Config list has >= 6 entries for GEMM-style kernels, >= 3 for reduction/pointwise.
• Benchmark uses triton.testing.do_bench with sufficient warmup (>= 25 iterations) and repetitions (>= 100).
• Achieved efficiency is reported as achieved / theoretical, not as absolute TFLOP/s alone.
• The correct theoretical peak (matching the kernel's dtype and operation type) is used as the denominator.
• Tuning was performed on the representative production shape, not a convenient round number.
• Boundary shapes (problem dimensions not divisible by BLOCK_*) were tested for correctness after selecting the final config.
• Autotune cache warm-up is documented for production deployment.
• No claim that the final config is "optimal" without profiling at least the next-best config to establish the margin.