mla_tilelang_to_triton

Ports MLA decode attention kernels from TileLang to Triton on gfx942 GPUs.

178|52|Updated Jul 30, 2025
One-click install
npx skills add https://github.com/AMD-AGI/GEAK --skill mla-tilelang-to-triton-amd-agi
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: mla_tilelang_to_triton
Source: https://github.com/AMD-AGI/GEAK/tree/main/perf_knowledge/expert_skills/skills/mla_tilelang_to_triton
Command: npx skills add https://github.com/AMD-AGI/GEAK --skill mla-tilelang-to-triton-amd-agi

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve? DeepSeek-class MLA decode attention on AMD MI300X (gfx942) is memory-bound on the KV-latent read, and the live TileLang kernel is hard to autotune and carries host-sync / CUDA-graph-capture hazards that block further optimization. ## Core Features & Use Cases - TileLang-to-Triton porting procedure: Step-by-step recipe to reimplement the MLA decode core in Triton with flash-style online softmax and absorbed up-projections (W_UK/W_UV) so the latent never expands to full head_dim in HBM. - Graph-capture safety rules: Enforces a host-sync-free hot path (no .item()/.cpu()/.synchronize()) with weight prep cached by data_ptr, so the kernel can overlay the live decode path. - Tuning and parity guidance: Concrete knob ranges (BLOCK_M=16/32, split_k, num_warps, fp8 fused scaling) plus do-no-harm pitfalls such as split_k numerical errors and bf16 weight re-materialization memory blow-up. - Use Case: While running the GEAK kernel_workflow on a DeepSeek model's decode bottleneck, apply this advisory prior to port the TileLang MLA kernel to a tunable, capture-safe Triton version targeting at least 1.15x isolated speedup with required parity. ## Quick Start Ask the agent to apply the mla_tilelang_to_triton expert skill to port the TileLang MLA decode kernel to a graph-capture-safe Triton kernel on gfx942 and validate parity against the TileLang oracle.

Frequently Asked Questions about mla_tilelang_to_triton

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I port a TileLang MLA decode kernel to Triton?

Extract the MLA decode op into an isolated task with a TileLang reference oracle, then write a Triton kernel using flash-style online softmax over the KV-latent loop and absorb the W_UK/W_UV up-projections into the QK/OV matmuls. Tune BLOCK_M, split_k, and num_warps while gating every config on parity.

Why port MLA attention from TileLang to Triton on MI300X?

TileLang expresses tiling explicitly but is harder to autotune inside the kernel layer and can carry host-sync and graph-capture hazards. A Triton port removes per-step host syncs and lets you match MFMA tiles to the latent head_dim as autotune knobs.

Does this Triton MLA kernel support CUDA graph capture?

Yes, graph-capture safety is mandatory for this kernel since it overlays the live decode path. The steady-state hot path must avoid .item(), .cpu(), and .synchronize() calls, and weight prep must be cached once by weight.data_ptr().

What are common pitfalls when tuning Triton MLA decode kernels?

split_k reductions with small BLOCK_M can produce numerically wrong results, so verify parity per config rather than only speed. Also avoid re-materializing bf16 weights for absorbed projections across all layers, which blows up memory and starves the KV cache.

Can this decode MLA recipe be used for prefill attention?

No, this is a decode-regime skill only and its tiling choices target tiny decode M values. Applying the decode tiling to prefill MLA is explicitly discouraged because prefill operates in a different shape regime.