amd-agi/geak/mla-tilelang-to-triton

When to use

The MLA (multi-head latent attention) decode path is the bottleneck on a DeepSeek-class model on MI300X (gfx942) and the live implementation is a TileLang kernel. Port it to Triton to (a) get an editable, CUDA-graph-capturable core the workflow can further tune, and (b) close the gap to the hand-tuned aiter/asm path while staying author-friendly.

Mechanism

MLA decode is memory-bound on the KV-latent read with a low-rank up-projection fused into the attention. TileLang expresses the tiling explicitly but is harder to autotune inside the kernel layer and can carry a host-sync / graph-capture hazard. A Triton port lets you: keep the latent in registers/LDS across the absorbed up-proj, pick MFMA-friendly BLOCK_M/BLOCK_N for the small decode M, and split the KV loop (split_k / flash-style online softmax) — all as Triton autotune knobs. The win is removing the per-step host sync and matching the MFMA tile to the latent head_dim, not raw FLOPs.

Procedure

1. Extract the MLA decode op into the immutable unittest task dir (shapes spanning decode M∈{1,64} and

the latent dims; bf16 and fp8_e4m3_fnuz operands; reference I/O oracle from the TileLang core).

1. Author a Triton kernel: flash-style online softmax over the KV-latent loop; absorb the up-projection

(W_UK/W_UV) into the QK / OV matmuls so the latent never expands to full head_dim in HBM.

1. Graph-capture safety is mandatory (this kernel overlays the live decode path): steady-state hot

path must be host-sync-free — no .item()/.cpu()/.synchronize(), no Python branch on a GPU scalar. Cache any weight prep (preshuffle/requant) once by weight.data_ptr(), never per call.

1. Tune: BLOCK_M small (decode), BLOCK_N/split_k over the KV loop, num_warps, kpack, fp32

accumulate. Keep the fp8 path fused (fold block-scale into the operand scale → one fp8 MFMA).

1. Gate isolated vs the oracle; then let the e2e layer overlay it (bare core, capture-safe).

Knobs & pitfalls

• Decode M is tiny — wide tiles waste lanes; start BLOCK_M=16/32.
• split_k>1 reductions with small BLOCK_M can go numerically wrong on this build — verify parity per

config, not just speed.

• Do NOT re-materialize bf16 weights for the absorbed proj across all layers (memory blow-up → forces

mem-fraction down → KV-cache starves → net e2e regression even if the kernel is faster).

Do-no-harm notes

• This is a decode-regime skill; do not apply the decode tiling to prefill MLA (different shape regime).
• If parity fails at any candidate config, drop that config — never trade correctness for speed.

Sources

• Recipe distilled from the MLA / DeepSeek case studies in `perf_knowledge/case_studies/by_model/

deepseek_mla_mi300x.md and deepseek_v3v4_attention.md`.

• STATUS: draft — needs on-box validation. Run `_contribute/validate_skill.py mla_tilelang_to_triton

--emit-plan then --record` with the kernel_workflow isolated A/B eval dir.