Optimization strategy (cuTile → Triton)

Purpose: One-page strategy distilled from translations/advanced-patterns.md and optimizing-reference.md. Agents must apply this when converting attention / FMHA / Gemma-style kernels (e.g. gemma_attention) or any kernel where Triton is expected to match an optimized in-repo baseline.

Full detail: Use this file for what to do and in what order; open the two linked docs for proofs, code samples, and edge cases.

When to read this (mandatory triggers)

Trigger	Action
Converting attention, FMHA, sliding window, soft cap, or GQA (e.g. Gemma)	Read §4 Gemma / FMHA checklist below before writing the Triton inner loop.
Host exposes `transpose` / `transpose_v` or MLA-style layout modes	Read §1 + advanced-patterns §1; use two kernels, not one + `tl.trans` in the KV loop.
Kernel uses `@triton.autotune`	Read §2 + advanced-patterns §2; grid `lambda META: (...)`; never freeze `BLOCK_*` from Python in a way that ignores autotune.
After TMA is in place but Triton still >20% slower than cuTile / baseline	Walk §3 + optimizing-reference sections 1–7 and §9 as applicable.
10–50× regression	translations/workflow.md — CRITICAL PERFORMANCE PATTERNS first (raw `tl.load` vs TMA).

§1 Layout flags → structure (from advanced-patterns)

Dual layout (transpose / transpose_v): Implement separate @triton.jit kernels when math and TMA layouts differ per mode (MLA: qk [H,N] vs [N,H], different V descriptor). Reusing the transpose path with extra tl.trans per KV block → 3–15× on the other mode.
Autotune + grid: grid = lambda META: (triton.cdiv(..., META["BLOCK_M"]), ...) — tile sizes come from META, not hard-coded forward() kwargs that bypass tuning.
Host TensorDescriptor vs in-kernel tl.make_tensor_descriptor: If descriptors must track autotuned block sizes, use a pre_hook to set block_shape per config (see Gemma FMHA host side), or follow mla_decoding.py in-kernel style — do not freeze wrong tile sizes on the host.
Multi-tensor small ops: Batched launch (program_id(1) selects tensor) — optimizing-reference §8, advanced-patterns §5.

§2 Post-TMA micro-optimizations (from optimizing-reference)

Apply in roughly this order after all 2D+ tile paths use TMA (or justified exceptions):

Priority	Pattern	Impact (typical)	Pointer
1	Autotune breadth + backend/GPU split — `get_available_triton_backend()`, `torch.cuda.get_device_capability()`, `num_stages`, `num_warps`, `occupancy`, `warp_specialize` where the stack supports it	10–20%+	optimizing-reference §4
2	EVEN_K / EVEN_* heuristics — skip masks in inner loop when divisible	5–15%	§1
3	Transpose — pointer/layout encoding; avoid `tl.trans` inside the K/V loop when avoidable	5–15%	§2
4	2D grid for pointer BMM — `(num_mn_blocks, batch)` + grouped M for L2	0–10%	§3
5	Epilogue subtile (large TMA stores)	5–15%	§5
6	`tl.assume` alignment on strides/pointers	5–15%	§6
7	Persistent vs non-persistent + occupancy tuning	10–30%	§7, §9.3

Blackwell / complex iterative kernels: TMA descriptors, tl.range(..., loop_unroll_factor=1) for heavy loops, TMEM-friendly tl.dot blocks (often ≥32), slab allocator — optimizing-reference §9, advanced-patterns §6.

§3 Fast vs slow patterns (skill gotchas, condensed)

Slow	Fast
Raw `tl.load(ptr+offs, mask=…)` for block-shaped 2D+ tiles	`tl.make_tensor_descriptor` / host `TensorDescriptor` + TMA load/store
`broadcast_to` + `tl.dot` for batched matmul	One batch (or head) per program; 2D `tl.dot` on real tiles
`qk = tl.zeros(...); tl.dot(q, k, qk)` when a fused dot exists	`qk = tl.dot(q, k)` (avoid redundant zero + 3-arg dot if the compiler path is worse)
Generic autotune for all GPUs/backends	Separate config lists for different backends and architectures; sm_90 vs sm_120 vs sm_80
Wrong TMA logical shape (e.g. K/V head dim = H when tensor is H_kv)	Descriptor `shape` matches tensor rank sizes for GQA: `H // QUERY_GROUP_SIZE` (or `num_head_kv`) on K/V
Always `libdevice.tanh` in hot path	Use fast tanh approximation where numerics allow (see `gemma_attention.py`)
Forcing `.contiguous()` on every forward	Only when required for TMA/strides; avoid extra copies
Autotune key missing `WINDOW_SIZE`, `SOFT_CAP`, `dtype`	Include keys that change optimal tile or specialization

§4 Gemma FMHA / `gemma_attention` conversion checklist (mandatory)

When converting or reworking Gemma-style attention (soft cap, sliding window, causal, GQA, BNSD), do not stop at “correct TMA” — apply these checklist items to match optimized Triton patterns:

GQA TMA metadata: Q / Out descriptors use [B, H, S, D]; K and V descriptors use [B, H // QUERY_GROUP_SIZE, S_kv, D] (same strides as the physical K/V tensor). Using H for the KV head dimension in the descriptor shape is a common bug and can hurt TMA behavior.
Autotune: get_configs() branches on get_available_triton_backend() and torch.cuda.get_device_capability() (e.g. (12,0)/(12,1) vs (9,0) vs (8,0)). Include occupancy, warp_specialize, num_stages, num_warps as in the reference — a single “SM ≥ 10 → 256×128” grid misses tuned Blackwell behavior.
@triton.autotune key: Include S_qo, S_kv, BLOCK_D, STAGE, QUERY_GROUP_SIZE, WINDOW_SIZE, SOFT_CAP, dtype (or equivalent) so different attention modes do not share one stale config.
QK matmul: Prefer qk = tl.dot(q, k) over explicit tl.zeros + tl.dot(q, k, qk) unless profiling shows otherwise.
Soft cap: Preserve scale order (tanh on logits in original scale, then align with INV_LOG_2 / exp2 softmax). Use fast tanh when supported by the backend (guard with a feature-detection flag).
Inner loop: offs_m / offs_n as tl.constexpr where the reference does; tl.max(qk, 1) / tl.sum(p, 1) axis consistent with layout.
Host: triton.set_allocator for TMA metadata; pre_hook updates descriptor block_shape for each autotune config. Avoid unconditional .contiguous() on Q/K/V unless needed.
Prune configs: For causal STAGE == 3, enforce BLOCK_M % BLOCK_N == 0 if the algorithm requires it (see reference prune_invalid_configs).

After implementation, run the same pytest + perf gates as in SKILL.md mandatory completion checklist.

Cross-references

Document	Role
translations/advanced-patterns.md	Dual kernels, autotune+META, descriptor vs autotune, diagnosis table
optimizing-reference.md	EVEN_K, transpose, grid, autotune, epilogue, alignment, persistent, batched launch, Blackwell §9
translations/workflow.md	Phases, TMA gate, catastrophic regression section
SKILL.md	Master checklist; must link optimization-strategy for attention conversions

/tilegym-converting-cutile-to-triton

referencesoptimization-strategy.md