Implementing a High-Performance Custom Diffusion Attention Kernel with FlyDSL 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. Implementing a High-Performance Custom Diffusion Attention Kernel with FlyDSL implementing-a-high-performance-custom-diffusion-attention-kernel-with-flydsl Readers 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. Diffusion models present unique challenges for kernel implementation, for the following reasons: 1. 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. 2. Configurability. Customers are exploring different settings to optimize real-world workload performance, so the kernel must be highly configurable. 3. 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. This 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. As 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. Optimizing Attention With FlyDSL, Step by Step optimizing-attention-with-flydsl-step-by-step Learn the FlyDSL Basics learn-the-flydsl-basics The FlyDSL repository provides a comprehensive starting guide: Choose a Starting Point choose-a-starting-point AMD’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: Understand Architectural Differences Across AMD GPU Generations understand-architectural-differences-across-amd-gpu-generations The examples target different GPU architectures. It is important to understand how hardware differences affect kernel implementations. Key considerations include: - 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. - 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. - 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. Debug Systematically debug-systematically - 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. - 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. Optimize Performance optimize-performance Match MFMA Operands to the Desired Fragment Layout match-mfma-operands-to-the-desired-fragment-layout For 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. For 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: ./matrix calculator.py \ -a gfx942 \ -i v mfma f32 16x16x16 bf16 \ --matrix-layout --D-matrix Consequently, one lane fixes the N coordinate j and stores four consecutive M rows. As shown in the output: lane 0: v0=D 0 0 v1=D 1 0 v2=D 2 0 v3=D 3 0 lane 16: v0=D 4 0 v1=D 5 0 v2=D 6 0 v3=D 7 0 On 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. The 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 . Analyze Register Usage analyze-register-usage To 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. Relevant fields in