AMD introduced GEAK, an open-source agent-driven framework for GPU kernel development and optimization, designed to automate the process of transforming PyTorch implementations into high-performance Triton kernels. The framework focuses on optimizing the DeepSeekV4 MLA kernel for AMD GPUs, demonstrating its ability to streamline the migration, profiling, and validation of kernel implementations. According to AMD, GEAK automates the kernel optimization loop, covering implementation generation, profiling analysis, issue localization, optimization iteration, and fix validation. The work highlights how the framework can enhance both kernel-level and end-to-end (E2E) performance across different workload shapes.
The results were tested on a representative MI355 Deepseek V4 workload with ISL/OSL=8k/1k, TP=8, and Concurrency=32. The optimized MLA kernel improved E2E throughput by 2.10x and reduced TTFT by 3.71x. The project followed a structured optimization flow, starting with establishing a PyTorch baseline, then using GEAK to drive the PyTorch-to-Triton migration for both prefill and decode paths. GEAK applied roofline-guided optimization across multiple iterations, integrating the optimized kernel into SGLang and validating its performance under real workloads.
The process focused on three main requirements: robust functional correctness, continuous performance improvement on the MI355 platform, and stable E2E gains after SGLang integration. The baseline implementation was based on FlashMLA/tests/ref.py, with test cases covering prefill and decode configurations. The configurations included various settings for h_q, d_qk, s_q, s_kv, topk, and batch sizes. The optimization process involved generating and iterating on Triton kernel implementations, running profiling and regression validation, and verifying fixes across both prefill and decode paths.
This shifted the workflow from manual trial-and-error to a systematic, agent-driven optimization process. The results are based on AMD testing using the described configurations, with performance potentially varying based on hardware, software versions, workload characteristics, and model configuration. Source: amd
