{"slug": "implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl", "title": "Implementing a High-Performance Custom Diffusion Attention Kernel with FlyDSL", "summary": "AMD published a step-by-step FlyDSL workflow for implementing a custom diffusion attention kernel, targeting vLLM's TiDAR (Think in Diffusion, Talk in Autoregression) mode with paged KV caches and scratch storage for speculatively generated tokens. The guide covers AMD GPU generation differences — MI350 and MI355 (gfx950/CDNA 4) offer more LDS capacity per compute unit, transpose-load instructions, and 16x16x32 and 32x32x16 MFMA instructions (mfma_f32_16x16x32_f16 and mfma_f32_32x32x16_f16) versus MI300 and MI325 (gfx942/CDNA 3) — and recommends a Triton baseline for correctness checks. AMD said it tested the workflow by giving an LLM agent its customer requirements and that the resulting implementation performed well.", "body_md": "# Implementing a High-Performance Custom Diffusion Attention Kernel with FlyDSL[#](#implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl)\n\nReaders may be familiar with traditional Transformer models and their attention mechanisms. The traditional autoregressive transformers generate tokens iteratively. Since this feature significantly limits the inference throughput, researchers have begun exploring approaches such as diffusion models that can generate multiple tokens in each iteration.\n\nDiffusion models present unique challenges for kernel implementation, for the following reasons:\n\n1. **Flexible KV-cache layouts.** Our customers want kernels that work with vLLM in TiDAR mode (Think in Diffusion, Talk in Autoregression). This requires support for paged KV caches and scratch storage for speculatively generated tokens.\n2. **Configurability.** Customers are exploring different settings to optimize real-world workload performance, so the kernel must be highly configurable.\n3. **High performance yet flexible.** Although FlexAttention offers considerable flexibility, it can have performance limitations. We aim to deliver optimizations across the stack, including FlashAttention and split-K at the algorithm and dataflow levels, as well as register-usage and data-movement optimizations at the low-level hardware layer.\n\nThis is where FlyDSL can help. FlyDSL addresses these needs by exposing low-level hardware details that kernel developers can optimize while retaining the flexibility required for customization.\n\nAs AI coding agents are increasingly used to implement GPU kernels, this post presents an step-by-step workflow with rich references. Developers can use this post to guide an agent for implementing or optimizing kernels. We also tested this workflow by giving an LLM agent our customer requirements and guide it as listed below. The resulting implementation performed well.\n\n## Optimizing Attention With FlyDSL, Step by Step[#](#optimizing-attention-with-flydsl-step-by-step)\n\n### Learn the FlyDSL Basics[#](#learn-the-flydsl-basics)\n\nThe FlyDSL repository provides a comprehensive starting guide:\n\n### Choose a Starting Point[#](#choose-a-starting-point)\n\nAMD’s public repositories provide several attention kernel implementations. Review these examples and reuse relevant code when possible. In addition to the examples in the [FlyDSL repository](https://github.com/ROCm/FlyDSL/tree/main/examples), FlyDSL kernels are available in:\n\n### Understand Architectural Differences Across AMD GPU Generations[#](#understand-architectural-differences-across-amd-gpu-generations)\n\nThe examples target different GPU architectures. It is important to understand how hardware differences affect kernel implementations. Key considerations include:\n\n- **LDS buffer size.** MI350 and MI355 GPUs (gfx950/CDNA 4) provide more LDS capacity per compute unit (CU) than MI300 and MI325 GPUs (gfx942/CDNA 3), enabling better data prefetching and pipelining.\n- **Transpose-load instructions.** MI350 and MI355 GPUs provide specialized instructions that transpose data while loading it from LDS into vector general-purpose registers (VGPRs). MI300 and MI325 GPUs require explicit transposition.\n- **MFMA instructions.** MI350 and MI355 GPUs introduce 16x16x32 and 32x32x16 MFMA instructions (`mfma_f32_16x16x32_f16` and`mfma_f32_32x32x16_f16` ). These provide higher throughput than the previous 16x16x16 and 32x32x8 variants.\n\n### Debug Systematically[#](#debug-systematically)\n\n- **Use a Triton baseline.** End users often prototype in Triton before moving to FlyDSL for better performance. A Triton implementation therefore provides a useful correctness baseline. Compare intermediate results at steps such as the log-sum-exp (LSE) calculation and before and after register operations such as permutations and XOR reductions.\n- **Start simple and add one feature at a time.** Begin with a straightforward implementation, then add features such as paged-attention support, the split-K algorithm, and data pipelining.\n\n### Optimize Performance[#](#optimize-performance)\n\n#### Match MFMA Operands to the Desired Fragment Layout[#](#match-mfma-operands-to-the-desired-fragment-layout)\n\nFor fused attention, MFMA operand order determines how score and probability fragments are distributed across lanes. The mathematical QK GEMM is `A = Q`, `B = K^T`, and `D = QK^T`, where M is the query dimension and N is the token dimension. However, logical row-major contiguity is not the same as per-lane register contiguity.\n\nFor gfx942’s `V_MFMA_F32_16X16X16_BF16`, refer to the general output layout in section 7.1.4 of the [AMD Instinct MI300 CDNA3 ISA Reference Guide](https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-mi300-cdna3-instruction-set-architecture.pdf). Use AMD’s [Matrix Instruction Calculator](https://github.com/ROCm/amd_matrix_instruction_calculator/tree/2ef91896bcdc4d26624f952e5c905c787cd9bc9e) to inspect the mapping:\n\n```\n./matrix_calculator.py \\\n  -a gfx942 \\\n  -i v_mfma_f32_16x16x16_bf16 \\\n  --matrix-layout --D-matrix\n```\n\nConsequently, one lane fixes the N coordinate `j` and stores four consecutive M rows. As shown in the output:\n\n```\nlane 0:  v0=D[0][0]  v1=D[1][0]  v2=D[2][0]  v3=D[3][0]\nlane 16: v0=D[4][0]  v1=D[5][0]  v2=D[6][0]  v3=D[7][0]\n```\n\nOn gfx942, using K as operand A and Q as operand B produces score fragments in the layout required by the subsequent P×V MFMA. This operand-swapped QK formulation allows the probability fragment to feed P×V directly, eliminating the probability LDS transpose, its synchronization barrier, and the associated LDS traffic. This optimization is illustrated in the figure below.\n\nThe AMD FlyDSL FlashAttention kernel demonstrates the same layout. The QK loop invokes [`mfma_acc(k_pack, q_pack, accumulator)`](https://github.com/ROCm/FlyDSL/blob/b8ed73fe6d9e17e101093324b1a7af518b1a0f29/kernels/attention/flash_attn_generic.py#L415-L416), placing K before Q. Its P×V invokes [`mfma_acc(v_transposed_pack, p_pack, accumulator)`](https://github.com/ROCm/FlyDSL/blob/b8ed73fe6d9e17e101093324b1a7af518b1a0f29/kernels/attention/flash_attn_utils.py#L2809-L2810).\n\n#### Analyze Register Usage[#](#analyze-register-usage)\n\nTo dump the generated assembly, set `FLYDSL_RUNTIME_ENABLE_CACHE=0` to avoid reusing stale cache entries and set `FLYDSL_DUMP_IR=1`. You can also set `FLYDSL_DUMP_DIR=/tmp/xx` to select the output directory; the default is `/root/.flydsl/debug/`. Because FlyDSL uses just-in-time (JIT) compilation, run the kernel at least once to generate the output.\n\nRelevant fields in `<DUMP_DIR>/kernel_<name>_0/21_final_isa.s` include:\n\n- `.set kernel.num_vgpr` /`.vgpr_count` and`num_agpr` /`.agpr_count`\n- `accum_offset` and`next_free_vgpr`\n- `group_segment_fixed_size` : LDS size\n- `.vgpr_spill_count` : memory spills\n\n#### Analyze Performance[#](#analyze-performance)\n\nUse `rocprofv3` to collect runtime statistics. The following hardware counters can help identify bottlenecks:\n\n| Purpose | Counters | \n|---|---|\n| L2 cache hits and coalescing | `TCC_HIT_sum` ,`TCC_MISS_sum` ,`TCC_REQ_sum` | \n| HBM efficiency and traffic distribution | `TCC_EA0_RDREQ_sum` ,`TCC_EA0_RDREQ_32B_sum` ,`TCC_EA0_RDREQ_DRAM_sum` ,`TCP_TCC_READ_REQ_sum` | \n| MFMA utilization | `SQ_INSTS_MFMA` ,`SQ_INSTS_VALU_MFMA_MOPS_*` ,`MfmaUtil` (MFMA busy percentage) | \n| VMEM and LDS activity | `SQ_INSTS_VMEM_*` ,`SQ_INSTS_LDS` ,`SQ_LDS_BANK_CONFLICT` ,`LDSBankConflict` | \n\nSee these guides for detailed collection and analysis instructions:\n\n#### Address LDS Bank Conflicts[#](#address-lds-bank-conflicts)\n\nSee the [GEMM optimization guide’s section on LDS bank conflicts](https://github.com/ROCm/FlyDSL/blob/421935cc6f09fd9b27d5d5ae52e0960e18834bd5/.claude/skills/gemm-optimization/SKILL.md?plain=1#L181).\n\n#### Prefetch Data[#](#prefetch-data)\n\nSee the [prefetch data-load guide](https://github.com/ROCm/FlyDSL/blob/421935cc6f09fd9b27d5d5ae52e0960e18834bd5/.claude/skills/prefetch-data-load/SKILL.md).\n\nFor more data-movement optimizations, see the [LDS optimization guide](https://github.com/ROCm/FlyDSL/blob/421935cc6f09fd9b27d5d5ae52e0960e18834bd5/.claude/skills/lds-optimization/SKILL.md).\n\n## Integration Notes[#](#integration-notes)\n\n### Graph Capture[#](#graph-capture)\n\nBecause FlyDSL uses JIT compilation, run the kernel once before starting graph capture. When designing the kernel interface, expose configuration values that can change between invocations as runtime parameters. Configuration values that remain fixed, such as the KV-cache page size and split-K factor, can be kernel template parameters.\n\n### Additional Reference[#](#additional-reference)\n\nSee the [FlyDSL kernel tuning guide](https://github.com/ROCm/FlyDSL/blob/421935cc6f09fd9b27d5d5ae52e0960e18834bd5/docs/kernel_tuning_guide.md).\n\n## Summary[#](#summary)\n\nAs AI models and attention algorithms evolve rapidly, traditional kernel libraries face challenges in both performance and flexibility. More and more customers require customized kernels for their algorithms.\n\nIn this blog, you gained a high-level yet actionable understanding of how we used FlyDSL to develop a customized diffusion attention kernel. We recommend that readers use this guide as a reference when implementing their own customized attention kernels. We tested this approach by providing an LLM agent with our customer kernel specification and carefully guide the agent step-by-step, and the resulting implementation worked well.\n\nFor other FlyDSL guide, please refer to other posts on our [ROCm blog](https://rocm.blogs.amd.com/software-tools-optimization/flydsl-python-native/README.html).\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. However, 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. THIS 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. AMD, the AMD Arrow logo, AMD Instinct, AMD 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. © 2026 Advanced Micro Devices, Inc. All rights reserved", "url": "https://wpnews.pro/news/implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/flydsl-customized-attention/README.html", "published_at": "2026-09-17 00:00:00+00:00", "updated_at": "2026-09-17 15:54:34.568131+00:00", "lang": "en", "topics": ["ai-infrastructure", "ai-chips", "large-language-models", "ai-agents", "developer-tools"], "entities": ["AMD", "FlyDSL", "vLLM", "TiDAR", "Triton", "MI350", "MI355", "MI300"], "alternates": {"html": "https://wpnews.pro/news/implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl", "markdown": "https://wpnews.pro/news/implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl.md", "text": "https://wpnews.pro/news/implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl.txt", "jsonld": "https://wpnews.pro/news/implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl.jsonld"}}