Optimizing Krea 3 on AMD Instinct MI355X with AI researchers

TL;DR: On AMD Instinct MI355X GPUs, Hedra’s optimized stack, combined with model-level optimizations Krea provided, brought Krea 3’s diffusion transformer (DiT) time per edit from 22.9 seconds in the BF16 baseline to 3.3 seconds. The main changes were MXFP6 precision for the multilayer perceptrons (MLPs), Hedra’s own MXFP6 matrix multiplication (GEMM) kernel, a custom attention kernel, fused operations and a transfer path that moves GPU-to-GPU exchanges onto MI355X’s SDMA copy engines.

At Hedra, we believe that current and upcoming generations of language models are not only capable of improving themselves, but capable of iteratively improving and training other models given the appropriate infrastructure.

Our new inference engine allows us to rapidly iterate and bring optimized versions of models online in days instead of weeks. In this article, we break down the optimization work and the performance gains we measured.

We optimized Krea 3, Krea’s image editing model, on AMD Instinct MI355X. For each edit, its diffusion transformer runs many times to produce a 4-megapixel image, so a small saving in one pass repeats across the whole edit.

MI355X runs 6-bit MXFP6 matrix math at four times its BF16 rate, and its dedicated SDMA copy engines move data between GPUs without taking compute units away from the model. Turning those capabilities into faster inference required software built for the workload.

For Krea 3, agents tested our researchers’ ideas on specific parts of the model. We combined those experiments with library tuning and custom implementations across numerical precision, matrix kernels, attention and communication. Each improvement changed where time was spent and what needed attention next.

Quality set the precision limit

MI355X runs MXFP6 at the same rate as MXFP4, so choosing between them comes down to quality. MXFP6 is a microscaled format that represents values using 6-bit floating-point elements, with a shared scaling factor for each block.

We compared lower-precision implementations with PyTorch's default BF16 path. We evaluated the format changes using structural similarity (SSIM) against the full BF16 model's outputs, across 23 image edits supplied by Krea.

Even with the same input and random seed, repeated BF16 runs produce slightly different outputs, so we used a BF16 rerun as the reference: it scored 0.962. Our target was to stay within 2% of that reference, and we also reviewed the output differences visually.

FP8 scored 0.954 (99.2% of the reference), MXFP6 0.946 (98.3%) and MXFP4 0.858 (89.2%). MXFP4 missed the target. FP8 stayed closest to BF16, but AITER's MXFP6 implementation ran the MLP faster than our best FP8 version while staying within the target, so we took MXFP6 forward into kernel optimization.

We tuned before writing custom kernels

General matrix multiplication (GEMM) libraries provide multiple kernels and use dispatch tables to select among them. Those tables are tuned for common matrix shapes. A library can support a numerical format without having a dispatch entry for the particular dimensions a model uses. In that case, the workload can fall back to a generic configuration. That happened with an 8-bit activation MLP kernel: its exact shapes weren't in AITER's dispatch table, so it fell back to a generic configuration.

We integrated automated tuning into our optimization framework. We timed candidates on real data and retained the fastest that passed quality and correctness checks. Saving those choices made the results available to later processes. For a new request size, the search could start from previous optimizations for the same GEMM. This phase selected among existing kernels and configurations; writing a new implementation was the next step.

We used agents to test the tradeoffs behind our implementation choices. On third-party cloud GPU VMs, they compared tuning RCCL with moving GPU-to-GPU copies onto MI355X’s dedicated SDMA engines. They tested how finely to split work for transfer overlap: two pieces helped; four were slower. They also tested fusing small operations into matrix multiplications, which brought no gain in those experiments.

We built a faster MXFP6 kernel

The choice of implementation mattered even within MXFP6. In a profiler comparison of one GPU's share of the MLP's GEMMs, a generic MXFP6 kernel, the Composable Kernel example, took 4.51 ms, about the same as PyTorch's default BF16 path at 4.34 ms. It stored each 6-bit value in a full byte, and neighboring workgroups re-read the same data. AMD's AITER MXFP6 implementation avoids both and took 1.72 ms, 2.5x faster than BF16.

After tuning the library kernels, we wrote our own MXFP6 GEMM, tailored to MI355X's cache hierarchy. It took 1.54 ms, 2.8x faster than BF16 and about 12% faster than AITER, and its output is bit-identical to AITER's.

Figure 1. GEMM kernel time for one GPU's share of the MLP. Timings come from a profiler trace on MI355X and cover GEMM kernels only. The BF16 reference uses the same matrix shapes. Lower is better.

Fusion cut intermediate memory traffic

The optimization effort also targeted intermediate memory traffic. Separate kernels can write a temporary result to memory only for the next operation to read it back. Fusion combines operations so that some intermediate results can stay within the combined computation. This addresses work around the GEMM that selecting a faster standalone matrix kernel does not necessarily remove.

For example, the query, key, value and attention-gate projections all read the same normalized activations. Stacking them into one MXFP6 projection means those activations are quantized and read once instead of four times. Other fusions merge consecutive operations: modulation with RMSNorm, the attention output projection with its gate and residual, and the MLP’s gate and up projections with the SwiGLU activation. The down projection reads the quantized intermediate directly, which avoids an extra BF16 write and requantization.

We moved transfers off the compute units

Faster matrix kernels made communication visible. Sequence parallelism splits the tokens across the GPUs. Most of the transformer works on each GPU’s tokens, but attention needs the entire sequence. Before attention, an all-to-all exchange gives each GPU every token for its share of the attention heads and a second exchange returns results.

Classifier-free guidance evaluates the model twice per denoising step, once conditioned and once unconditioned. We first ran the two classifier-free guidance branches in parallel, each on a separate GPU group using sequence parallelism. Running both across the full GPU group instead let each GPU communicate over more peer links. Overhead limited the bandwidth gain to about 2.6x per GPU (about 153 versus 58 GB/s). This layout made the transformer as a whole about 12% faster. Batching the two branches into one pass saved about another 5%.

Without overlap, exposed transfer time (transfers not hidden behind any compute) was 18% of each generation in our profiling run. Splitting the heads into two groups, so that one group’s data moves while the other computes, halved it to 9%. It went no further because the collective library runs its copies as kernels on the same compute units as the model, so the small kernels around attention ran much slower while a transfer was in flight.

MI355X has dedicated SDMA copy engines, so we built a transport that moves the exchange onto them, one engine per peer. The copy engines pay off with overlap: on their own, exposed transfer time was still 16%. Together with the head grouping, they cut it to 4%, and sending each group’s results back as soon as its attention finished brought it down to 2%. Useful overlap therefore depends on both the schedule and the resources used to move data.

Figure 2. Transformer time per generation in the profiling run, split into computing and waiting on transfers, which here means exposed transfer time. It was 18% with the collective library (RCCL), one head group and no overlap; 9% with RCCL and two head groups; 16% with the SDMA copy engines, one head group and no overlap; 4% with SDMA and two head groups; and 2% with each group's results sent back as soon as its attention finished. Right: how much slower small kernels next to attention ran beside RCCL and SDMA transfers than alone.

Chart data for Figure 2

Kernel time beside an in-flight transfer divided by the same kernel’s time alone. A ratio of 1.0 means no change; lower is better.

We chose accuracy over raw attention speed

Attention sits between the two all-to-all exchanges, so the final stack uses our own attention kernel, built around that layout and generated as gfx950 assembly for MI355X. It takes queries and keys (Q and K) as 8-bit integers (INT8) and values (V) in MXFP6. We chose this combination because it was more accurate than plain MXFP6 attention. Our attention kernel took 1.3 ms in isolation, compared with AITER’s 0.97 ms. Q and K are quantized before the exchange, which halves the data they send. The kernels that produce Q and K apply RMSNorm and rotary position embedding (RoPE) and write the results directly into the all-to-all send buffers. Inside the attention kernel, two groups of four wavefronts alternate between matrix instructions for one tile and softmax work for another. That inner schedule is separate from the two head groups, which overlap one group’s transfer with the other group’s compute.

The combined stack cut DiT time to 3.3 seconds

Figure 3 follows four cumulative configurations. The BF16 baseline used PyTorch's default attention. The second configuration combined AITER's low-precision MLP with FP8 projections, fusions and compilation, bringing DiT time per edit from 22.9 to 16.6 seconds. Adding AITER's low-precision attention brought it to 13.8 seconds. Our full stack reached 3.3 seconds. The final stack includes our own MXFP6 GEMM, our attention kernel, the communication changes and fused operations, combined with model-level optimizations Krea provided.

The balance of work shifted along the way. In the BF16 baseline, matrix multiplications took 35% of DiT time and attention 43%. After the AITER changes, matrix multiplications took 53% and waiting on communication 20%. Waiting on communication counts the time the compute stream sat idle or blocked on transfers, so it differs from the exposed transfer time in Figure 2.

Figure 3. Diffusion transformer time per edit and each configuration's reported profile shares. Each configuration adds to the previous; the second also adds FP8 QKV and output projections, output-projection gate and RMSNorm fusions, and compilation. Benchmark medians and profile shares are from separate runs, so multiplying the two does not give component times. Percentages are rounded and may not sum to 100%. These are DiT timings, not full-request latency. Times are medians across 23 image edits supplied by Krea, measured after warm-up on MI355X GPUs with a 1 MP reference input and a 4 MP output.

Chart data for Figure 3

Reported profile shares of DiT time for each cumulative configuration.

Benchmark medians and profile shares come from separate runs. Multiplying them does not give component times. Percentages are rounded and may not sum to 100%.

The 3.3-second result reflects the combined stack rather than any single change. Each fix moved the bottleneck: faster matrix kernels exposed communication, and moving transfers onto the copy engines let communication and compute overlap.

Build with Hedra Model Factory

The work we did with Krea shows the decisions an autonomous optimizer has to handle: whether lower precision meets the quality target, when greater accuracy justifies a slower kernel, and where the bottleneck moves after each change. Our goal is for AI researchers to improve both model implementations and their own tools and strategies, then use those improvements in the next round of experiments.

Hedra Model Factory combines reinforcement learning research, GPU infrastructure and inference optimization for multimodal models. Our researchers and agents work with your team to develop environments, run post-training and evaluate results.

Bring us the problem you need to solve.

We help you build, deploy and keep improving a model you own.