/tilegym-converting-cutile-to-triton

Advanced Triton patterns (cuTile → Triton)

Strategy hub: For an ordered summary of this file plus optimizing-reference.md (and a mandatory Gemma FMHA checklist), read references/optimization-strategy.md first when converting attention or matching gemma_attention-class perf.

Patterns that are easy to miss in conversion but cause large correctness or performance gaps. Complements translations/workflow.md (phases, TMA, autotune) and SKILL.md (checklist).

1. Dual layout flag (transpose / transpose_v): two kernels, not one + transposes {#dual-layout-flag}

Applies to: MLA-style decoding, attention variants, any op where the host exposes a boolean like transpose=True/False to match cuTile or framework layout.

Failure mode (real case: MLA decoding)

Why: The fast Triton baseline uses different math and tensor layouts per mode:

• transpose=False (cache-friendly V along batch×seq): Keep qk as [BLOCK_H, BLOCK_N], softmax state l_prev shape [BLOCK_H], tl.max(qk, 1), tl.sum(p, 1), and tl.dot(p, v, acc) with p already [H, N] and v [N, D] — one streamlined path.

• transpose=True (V read with seq leading): qk is [BLOCK_N, BLOCK_H] (e.g. tl.dot(k, q) with q as [D, H]), l_prev [BLOCK_N, BLOCK_H], V needs a separate TMA descriptor (shape=[S_kv, B, D], strides=[stride_n, stride_b, 1], block_shape=[BLOCK_N, 1, D]) so the value load is not folded away or mis-coalesced — then tl.dot(v, p, acc) with transposed v/p (see naive_absorb_mla_transpose in the same file).

Anti-pattern: Implementing transpose=False by reusing the transpose kernel’s structure (e.g. forcing l_prev to [BLOCK_N, BLOCK_H], transposing qk and p every KV block). That is correctable but much slower than the dedicated non-transpose kernel.

Agent checklist:

1. If cuTile or the PyTorch wrapper has transpose (or equivalent), grep the Triton tree for one @triton.jit handling both — verify each branch uses the same layout strategy as the in-repo reference, not “transpose path + extra tl.trans”.
2. Add two @triton.jit kernels when the reference does (or split with tl.constexpr flags only if the compiler fully specializes — prefer two kernels for clarity and perf).
3. Do not pass fixed BLOCK_H / BLOCK_N from Python into autograd.Function.forward when the kernel is @triton.autotune — see §2.

2. Autotune + grid: BLOCK_* must come from META, not from the host

Failure mode: Host passes BLOCK_H=16, BLOCK_N=128 into apply(), and launch uses grid = (cdiv(heads, BLOCK_H), B). Autotune configs that use 64×128 or 128×128 never affect grid → under-occupancy and wrong tuning vs a baseline that uses:

grid = lambda META: (triton.cdiv(num_head, META["BLOCK_H"]), B, 1)
# Launch kernel with that grid; pass tensor bases and META-sized BLOCK_* inside the kernel — not via apply().

Rule: For autotuned kernels, forward should only pass dynamic shapes / strides / pointers; tile sizes are tl.constexpr filled by autotune from META.

Cross-check: Run with TRITON_PRINT_AUTOTUNING=1 and confirm the chosen config matches the problem key (BLOCK_D, S_kv, EVEN_N, etc.).

3. When to use host TensorDescriptor vs tl.make_tensor_descriptor in-kernel

TileGym’s mla_decoding.py uses tl.make_tensor_descriptor inside the JIT function with raw tensor bases + strides (plus triton.set_allocator for TMA metadata).

An alternate style (some LLM ports) builds triton.tools.tensor_descriptor.TensorDescriptor on the host when sm >= 90 and passes descriptors into the kernel. That can be correct but:

• If you freeze tile sizes on the host descriptor to match fixed BLOCK_H/BLOCK_N, you fight autotune — either descriptors must be rebuilt per config (heavy) or you should use in-kernel tl.make_tensor_descriptor like mla_decoding.py.

Preference for new conversions: Align with mla_decoding.py: allocator + in-kernel descriptors + autotune + dual kernels for layout flags.

4. Quick diagnosis table

5. Batched kernel launch for multi-tensor ops

Applies to: cat, stack, split, multi-tensor copy/scatter operations.

When converting a cuTile kernel that processes multiple tensors (e.g., concatenation), do not naively launch one Triton kernel per tensor. Instead, batch up to 4 tensors per launch using a 2D grid:

# Grid: (blocks_over_elements, num_tensors_in_batch)
grid = (triton.cdiv(max_elements, BLOCK), min(4, num_remaining_tensors))

# Kernel uses program_id(1) to select which tensor to process
pid_y = tl.program_id(1)
if pid_y == 0:
    in_ptr, size = in_ptr_a, size_a
elif pid_y == 1:
    in_ptr, size = in_ptr_b, size_b
# ...

Impact: 2–4× speedup for small-to-medium tensors (LLM KV-cache concatenation). Kernel launch overhead (~5–10µs per launch) dominates for small tensors; batching amortizes this cost.

Full pattern and code: See §8 Batched Kernel Launch in references/optimizing-reference.md.

6. Blackwell Optimization Patterns

For complex kernels targeting Blackwell (sm_100+), additional optimization patterns can yield 2–10× speedups. These are documented in references/optimizing-reference.md §9 and include:

Key insight: On Blackwell, register pressure from loop unrolling is often the #1 performance killer. Use tl.range(..., loop_unroll_factor=1) for loops with >16 iterations and complex bodies.

7. Reference implementations

• MLA decoding with transpose flag: Implement two separate @triton.jit kernels (transpose=False path and transpose=True path). The transpose=False path: qk as [BLOCK_H, BLOCK_N], l_prev as [BLOCK_H], tl.dot(p, v, acc). The transpose=True path: separate V TMA descriptor with layout matching the transposed access pattern. See §1 above and examples/05_attention/ for a worked example.

• Cat (batched): Original TileGym cat implementation (adapted from FlagGems) demonstrates batched multi-tensor kernel launch. See cat_fwd_tensor_list_contiguous_any_dim_up_to_four_inputs_kernel pattern or §8 of references/optimizing-reference.md.