{"slug": "geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel", "title": "GEAK Agent-Driven Optimization of the DeepSeekV4 MLA Kernel", "summary": "AMD's GEAK agent-driven framework automated the optimization of the DeepSeekV4 MLA kernel, achieving a 2.10x improvement in end-to-end throughput and a 3.71x reduction in time-to-first-token on MI355 GPUs. The framework handled PyTorch-to-Triton migration, profiling, and iterative optimization, integrating the kernel into the SGLang inference framework.", "body_md": "# GEAK Agent-Driven Optimization of the DeepSeekV4 MLA Kernel[#](#geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel)\n\nOptimizing LLM inference kernels requires more than a single kernel rewrite. Developers need to migrate reference implementations, analyze profiling results, validate correctness, and iterate quickly across different workload shapes. In this blog, we use DeepSeekV4 MLA as a case study to show how GEAK automates this workflow, from PyTorch-to-Triton migration to kernel-level optimization and SGLang end-to-end(E2E) validation on AMD GPUs.\n\nThis work demonstrates that GEAK can automatically drive the kernel optimization loop across implementation generation, profiling analysis, issue localization, optimization iteration, and fix validation. 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.\n\n## Background[#](#background)\n\n### GEAK Overview[#](#geak-overview)\n\nGEAK is an open-source, AMD-developed agent-driven framework for GPU kernel development and optimization, available at [AMD-AGI/GEAK](https://github.com/AMD-AGI/GEAK/tree/GEAK_v3.2). It automates the kernel optimization workflow by taking a kernel repository and an optimization prompt as input, then iteratively identifying optimization opportunities, modifying the kernel implementation, and evaluating the resulting patches.\n\nGEAK combines an agent loop with MCP and local tools. The figure below shows how GEAK takes a kernel repository and an optimization prompt as input, then iteratively optimizes kernels and evaluates patches.\n\n### DeepSeekV4 MLA Optimization Objective[#](#deepseekv4-mla-optimization-objective)\n\nThe PyTorch reference implementation of DeepSeekV4 MLA attention served as the functional and performance baseline for this project. Rather than redesigning the algorithm, our goal was to transform the existing implementation into a high-performance Triton kernel and integrate it reliably into the SGLang inference framework.\n\nThe work followed a straightforward optimization flow:\n\nEstablish the PyTorch baseline;\n\nLet GEAK drive the PyTorch-to-Triton migration for both the prefill and decode paths;\n\nLet GEAK apply roofline-guided optimization across multiple iterations;\n\nIntegrate the optimized kernel into SGLang and validate both kernel-level and E2E performance under real workloads.\n\nThroughout this process, we focused on three requirements: robust functional correctness, continuous performance improvement on the MI355 platform, and stable E2E gains after SGLang integration.\n\n## Baseline Implementation and Configuration[#](#baseline-implementation-and-configuration)\n\n### Codebase[#](#codebase)\n\nThe original reference implementation is based on [ FlashMLA/tests/ref.py](https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py).\nThe test cases are based on:\n\n### Configurations[#](#configurations)\n\nThe baseline configurations are summarized below. For the full configuration table, see Appendix.\n\nFor prefill, the benchmark covers eight settings across\n\n`h_q=64/128`\n\n,`d_qk=512`\n\n,`s_q=4096`\n\n,`s_kv=8K/32K/48K/64K`\n\n, and`topk=512/1024`\n\n.For decode, the benchmark covers 24 dual-scope configurations and 2 peak-performance cases with only the main scope enabled.\n\nFor 24 dual-scope configs, the main scope uses\n\n`s_kv=16K`\n\nwith`topk=128`\n\n,`block_size=256`\n\n, while the extra scope covers`extra_s_k=16K`\n\nwith`extra_topk=512/1024`\n\n.For the 2 peak-performance cases, the main scope uses\n\n`s_kv=32K`\n\nwith`topk=16K`\n\n,`block_size=64`\n\n.The configurations also vary\n\n`h_q=64/128`\n\n,`extra_block_size=2/64`\n\n, and`batch sizes`\n\nvalues from 2 to 256.\n\n## MLA Kernel Optimization Process[#](#mla-kernel-optimization-process)\n\nResults are based on AMD testing using the configurations described herein. Performance may vary based on hardware, software versions, workload characteristics, and model configuration.\n\nGEAK automated the main kernel optimization loop for this project. It generated and iterated on Triton kernel implementations, ran profiling and regression validation, localized performance issues, and verified fixes across both prefill and decode paths. This shifted the workflow from manual trial-and-error to a systematic, agent-driven optimization process.\n\n### Prefill Kernel[#](#prefill-kernel)\n\nThe prefill optimization focused on block-level tiling, sparse KV access patterns, and compiler optimization.\n\nIn the discussion below, `Config 1`\n\nand `Config 2`\n\nrefer to the two prefill configuration groups listed in the appendix. Each group covers four `s_kv`\n\nvalues, resulting in eight prefill configurations in total.\n\nThe initial Triton kernel was roughly on par with the baseline for `CONFIG 1`\n\n, but only reached 0.67x on `CONFIG 2`\n\n. Profiling showed that the main bottlenecks came from sparse KV gather, register pressure, and control-flow overhead.\n\nThe first optimization round focused on fusing the sparse attention data path into the Triton kernel. GEAK integrated vectorized KV gather, `topk_length`\n\nmasking, and attention sink logits handling into the kernel, reducing intermediate tensor materialization and redundant global memory traffic. It then explored tile size, `num_warps`\n\n, and `num_stages`\n\nconfigurations to balance memory throughput, parallelism, and register usage. With these changes, the speedup improved to 6.85x on `Config 1`\n\nand 4.75x on `Config 2`\n\n.\n\nIn the second round, GEAK continued the optimization by upgrading to Triton 3.6.0, expanding the autotune search space, and introducing a hybrid bf16/fp32 accumulation strategy. With these changes, the final speedup reached 9.13x on `Config 1`\n\nand 6.92x on `Config 2`\n\n.\n\nFor correctness validation, the 304 test cases from [FlashMLA/tests/test_flash_mla_sparse_prefill.py](https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_sparse_prefill.py) were reused, and all tests passed the numerical accuracy checks.\n\n### Decode Kernel[#](#decode-kernel)\n\nThe decode path has a more complex execution pattern than prefill. It runs token by token, typically under low-batch and high-concurrency workloads, and needs to retrieve sparse KV blocks from both the main and extra scopes.\n\nBecause of this, decode performance is not dominated by the attention computation alone. A large part of the cost comes from the surrounding data path, including sparse gather, dequantization, tensor layout conversion, kernel launch overhead, and coordination between the two KV scopes.\n\nGEAK optimized the decode kernel through four major iterations. The initial Triton version underperformed the baseline, with a geomean speedup of only 0.66x across the 26 decode test configurations. After the full optimization process, the geomean speedup improved to 4.94x.\n\nThe first iteration focused on kernel fusion. Fragmented Torch-side operations, including gather and dequantization, were fused into the Triton kernel. This reduced intermediate tensor materialization, enabled vectorized memory access, reused precomputed address offsets, and removed redundant operations such as unnecessary reshapes.\n\nThe second iteration focused on simplifying the dual-scope path for the main and extra KV scopes. Instead of handling the two scopes with separate or partially duplicated logic, the tensor layout and memory access pattern were reorganized so both scopes could share a more unified kernel path. This reduced redundant branches, repeated scheduling, and unnecessary data movement between the two scopes.\n\nThe third iteration focused on workload-specific kernel configuration. Different vectorization and blocking choices, such as `BLOCK_N`\n\n, were evaluated to better match the target decode shapes. Launch and scheduling overhead were also reduced so the kernel could scale more effectively across batch sizes and retrieval sizes. At this stage, the geomean speedup reached over 3.7x.\n\nThe final iteration targeted the remaining critical-path overhead and low-occupancy cases. Base pointers were computed ahead of time to reduce repeated address calculations, redundant code paths were simplified, unnecessary branches were removed, and extra data movement and recomputation were eliminated. For shapes with insufficient parallelism, the split-K strategy was expanded to expose more parallel work and improve occupancy. Together with additional autotune configuration optimization, this brought the final geomean speedup across the 26 decode configurations to a stable 4.94x.\n\nFor correctness validation, the 4,748 test cases from [FlashMLA/tests/test_flash_mla_sparse_decoding.py](https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_sparse_decoding.py) were reused, and all tests passed the numerical accuracy checks.\n\n## Roofline-Guided Optimization[#](#roofline-guided-optimization)\n\nDuring optimization, `rocprof-compute`\n\nwas used for roofline analysis, and the results were correlated with key low-level metrics, including arithmetic intensity, bandwidth utilization, occupancy, register spills, and bank conflicts.\n\nThis analysis helped us avoid several misleading optimization directions:\n\nReducing memory traffic did not always improve performance. Some changes reduced global memory accesses but increased VGPR usage, which in turn caused register spills or lower occupancy and led to performance regressions.\n\nFor peak-performance cases, further operator fusion was often less effective than exposing more parallelism and reducing register pressure.\n\nGrounding each optimization step in profiler evidence reduced trial-and-error cost and shifted the workflow from experience-driven tuning to data-driven optimization.\n\n## E2E Validation[#](#e2e-validation)\n\nResults are based on AMD testing using the configurations described herein. Performance may vary based on hardware, software versions, workload characteristics, and model configuration.\n\nThe SGLang baseline for DeepSeekV4 MLA used a TileLang implementation. The TileLang baseline was replaced with the optimized Triton sparse MLA kernel, and E2E performance was compared under the same evaluation setup. The integration PR is available here: [sglang#26208](https://github.com/sgl-project/sglang/pull/26208).\n\nAcross multiple concurrency levels, we observed consistent improvements in both performance and accuracy validation:\n\nUnder the 8k/1k, TP=8, Concurrency=2/4/8/16/32 setup, TPUT, TTFT, and ITL all improved.\n\nE2E throughput improved by approximately 16% to 110% across different concurrency levels, with larger gains at higher concurrency.\n\nAccuracy validation passed. For example, GSM8K changed from 0.936 to 0.942, which is within the expected range and showed no correctness regression.\n\nA representative result is shown below for ISL/OSL=8k/1k, TP=8, and Concurrency=32:\n\nMetric |\nBaseline |\nOptimized |\nImprovement |\n|---|---|---|---|\nTPUT |\n3002.29 tok/s |\n6293.05 tok/s |\n2.10x |\nTTFT |\n34620.41 ms |\n9321.49 ms |\n3.71x |\nITL |\n32.16 ms |\n28.63 ms |\n1.12x |\n\n## Summary[#](#summary)\n\nThis DeepSeekV4 MLA work shows that optimizing a complex inference kernel requires a full engineering loop, from implementation generation and profiling to regression validation, issue localization, and fix verification. In this project, GEAK automated this optimization loop, driving the PyTorch-to-Triton migration, iterative kernel optimization, profiling analysis, and validation across both prefill and decode paths.\n\nBy automating these steps, GEAK significantly improved optimization productivity and accelerated performance convergence. This case study demonstrates that GEAK has moved beyond agent-assisted development and can now drive automated kernel optimization for complex LLM inference workloads.\n\n## Appendix[#](#appendix)\n\n### Prefill (8 configs)[#](#prefill-8-configs)\n\nConfig |\ns_q |\nd_qk |\nh_q |\ntopk |\ns_kv |\n|---|---|---|---|---|---|\nConfig 1 |\n4096 |\n512 |\n64 |\n512 |\n8192, 32768, 49152, 65536 |\nConfig 2 |\n4096 |\n512 |\n128 |\n1024 |\n8192, 32768, 49152, 65536 |\n\n### Decode (26 configs)[#](#decode-26-configs)\n\nConfig |\nh_q |\ns_kv (main) |\ntopk (main) |\nextra_s_k |\nextra_topk |\nTotal Retrieval |\nblock_size |\nextra_block_size |\nbatch_size |\n|---|---|---|---|---|---|---|---|---|---|\nConfig 1 |\n64 |\n16K |\n128 |\n16K |\n512 |\n640 |\n256 |\n64 |\n2/64/74/128/148/256 |\nConfig 2 |\n128 |\n16K |\n128 |\n16K |\n1024 |\n1152 |\n256 |\n64 |\n2/64/74/128/148/256 |\nConfig 3 |\n64 |\n16K |\n128 |\n16K |\n1024 |\n1152 |\n256 |\n2 |\n2/64/74/128/148/256 |\nConfig 4 |\n128 |\n16K |\n128 |\n16K |\n1024 |\n1152 |\n256 |\n2 |\n2/64/74/128/148/256 |\nPeak perf Config 2 |\n64 |\n32K |\n16K |\nN/A |\nN/A |\n16K |\n64 |\nN/A |\n148 |\nPeak perf Config 1 |\n128 |\n32K |\n16K |\nN/A |\nN/A |\n16K |\n64 |\nN/A |\n148 |\n\n## Disclaimers[#](#disclaimers)\n\nThe information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information.\n\nHowever, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes.\n\nTHIS INFORMATION IS PROVIDED “AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.\n\nAMD, the AMD Arrow logo, AMD Instinct, ROCm, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies.\n\n© 2026 Advanced Micro Devices, Inc. All rights reserved.", "url": "https://wpnews.pro/news/geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/geak-mla-optimization/README.html", "published_at": "2026-07-13 00:00:00+00:00", "updated_at": "2026-07-13 15:53:12.625930+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-tools", "ai-infrastructure", "ai-chips"], "entities": ["AMD", "DeepSeek", "GEAK", "SGLang", "MI355", "Triton", "FlashMLA"], "alternates": {"html": "https://wpnews.pro/news/geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel", "markdown": "https://wpnews.pro/news/geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel.md", "text": "https://wpnews.pro/news/geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel.txt", "jsonld": "https://wpnews.pro/news/geak-agent-driven-optimization-of-the-deepseekv4-mla-kernel.jsonld"}}