Difflet: Serving Diffusion Models on AWS Trainium

Six production image and video models, end-to-end, matching or beating NVIDIA H100 per step on most models, at up to 4.10x the throughput per dollar.

Diffusion inherited a CUDA-only serving stack

Most image and video diffusion pipelines were built around PyTorch, Hugging Face diffusers, and CUDA. That made NVIDIA GPUs the default deployment target, even though the workload itself is dominated by dense matrix multiplication and attention rather than CUDA-specific primitives.

Diffusion model inference can be expensive. Take HunyuanVideo as an example. With stock diffusers on an H100, one denoising step takes 1.50 seconds. A 30-step generation therefore spends roughly 45 seconds in the DiT alone, before text encoding and VAE decoding.

At the same time, Trainium2 capacity is priced at $2.235 per accelerator-hour, or about 43.1% of the $5.191 H100 rate used in our comparison. The hardware is available and well matched to the computation; what has been missing is a diffusion software path that can use it.

The result is higher dollars per inference, diffusion traffic competing with LLMs for the same GPUs, and Trainium capacity sitting idle instead of serving requests.

Making systolic silicon run diffusion fast

Trainium’s compute core is a systolic array: operands flow through a grid of processing elements and are reused at every hop, so one on-chip fetch is amortized over thousands of multiply-accumulates. Data movement is the dominant performance cost and energy cost of inference, and this attacks cost at its source. And because Trainium compiles ahead-of-time, Difflet builds the whole graph before it runs: kernel selection, memory layout, and engine schedules tuned per model, with predictable latency and a memory footprint known before deployment.

None of this required bending the models. A diffusion transformer spends almost all of its time in dense GEMMs and attention, precisely the work a systolic array exists for. What was missing was the software in between: kernels that stage operands the way the array wants them, schedules that keep the tensor and vector engines in Trainium busy at the same time, and a compiler that sees the whole pipeline rather than one op at a time.

Difflet is that layer: an AOT-compiled path from a diffusers pipeline down to the NeuronCores, with four mechanisms doing the heavy lifting.

AOT-compiled DiT pipelines. Full pipelines compiled ahead-of-time with static memory plans. Whole-graph optimization per model: predictable latency, memory footprint known before deployment. Operationally, a static plan means capacity math is settled before launch: no allocator jitter, no fragmentation surprises. The same request costs the same milliseconds tonight as it did this morning.

Custom systolic attention. A systolic array amortizes one on-chip fetch over thousands of multiply-accumulates, but only if operands arrive in the order the array wants them. Our kernel stages attention exactly that way, replacing PyTorch’s stock path: 2.06x speedup on Wan attention; HunyuanVideo per-step decreases from 3719 ms to 851 ms.

Engine decomposition. GEMM-heavy blocks run on the tensor engine while softmax, normalization, and rotary embeddings run on the vector engines, concurrently, not in sequence. Phases that queue up behind each other on a GPU overlap here, so the systolic array, the expensive part of the chip, almost never sits idle waiting on elementwise work.

Training-free step caching. Adjacent denoising steps produce nearly identical results, so Difflet reuses the previous step’s output and skips some DiT forward passes. Adapted from TeaCache, with no retraining and no weight changes: measured at 1.91x end-to-end speedup on HunyuanVideo, with output quality checked against uncached runs; Wan reaches 1.77x with fixed cadence at 0.9997 final-latent cosine. Three skip policies are available. The benchmarks in the performance section below are measured with caching disabled, so this speedup stacks on top.

Running end-to-end today on Trainium2 and Trainium3: FLUX.1-dev, Qwen-Image, Wan 2.1, Wan 2.2, HunyuanVideo, and LTX-2, covering image, video, and video + audio.

Throughput per dollar

Raw peak speed belongs to the newest, priciest, scarcest silicon, and we show it anyway. What runs a datacenter is output per dollar, and there Trainium wins on every one of the six models.

A word on method: per-step DiT latency is where diffusion spends its time, and one backend-generic harness measures it identically on four devices (NVIDIA H100, B300, Trainium2 and Trainium3) under pinned model revisions. Trainium runs Difflet’s AOT path at tp=4; H100 and B300 run the stock Hugging Face diffusers path a team would deploy today. Step caching is disabled in every benchmark in this post. All per-step and end-to-end numbers compute the full step count, so the 1.91x TeaCache speedup stacks on top of them. Every result below can be rerun from the commands at the end of this post.

Wan 2.2 note: its Trainium3 figures reuse the Wan 2.1 result because the two models share the same execution architecture and compiled graph; 442.5 ms per step and the derived 1.25x figures are reused rather than independently measured.

And yes, the newest Blackwell GPU (B300) wins raw steps on most models. It is also the most expensive and hardest-to-get silicon in the fleet.

Steps per dollar at AWS Capacity-Block prices (July 2026): Trainium2 $2.235/accel-hr (tp=4, one chip = 4 NeuronCores) vs H100 $5.191/accel-hr, a 2.32x price edge before any speed. Every model clears parity; HunyuanVideo hits 4.10x.

Trainium runs Difflet AOT at tp=4; H100 and B300 run stock Hugging Face diffusers, single-GPU dense. B300, the newest Blackwell, is fastest per raw step on most models. It is also the priciest and scarcest. On HunyuanVideo, Trainium3 (650.0 ms) beats even B300 (874.5 ms).

(all values in ms per step)

Trainium-over-H100 speedup = H100 per-step time divided by Trainium per-step time. Values above 1.0x favor Trainium. Trainium2 matches or beats H100 on 4 of 6 models; Trainium3 reaches 2.31x on HunyuanVideo.

Same engine, next silicon: Trainium3 is 1.11-1.38x faster per step across the suite.

The current best measured serving path starts from an already-loaded worker: no per-request weight load. Flux and Qwen report the stable mean of HTTP runs 2-3; Wan uses the accepted Neuron VAE profile; LTX-2 reports repeated resident API completion. Fixed output profiles differ by modality, so these are deployed request latency rather than a cross-model efficiency ranking. Device prewarm is outside request timing; it affects one-time worker startup only.

Under the hood: the engineering behind the curves

None of these numbers came free. Each one is a specific piece of Difflet engineering: a kernel rewired, an attention head resharded, the load path reworked, measured before and after.

Two startup costs are paid once when the worker starts and amortized across requests. First, weight data loading: Difflet writes per-rank weight shards during AOT compilation, then loads ranks concurrently with a per-thread allocator. That removes request-time transformer sharding and host allocator contention. Second, device prewarm: a best-effort background device touch initializes the Neuron runtime and assigned cores while host-side weights load, so roughly 6.7 seconds of device bring-up overlaps with host loading. The main path remains the fallback, so startup improves without becoming a correctness dependency. A full end-to-end prewarm on/off A/B has not yet been recorded.

From proven engine to production backend

Per-step compute is solved, and the resident serving path is now measured end to end. The worker pays model loading and device bring-up once, then serves requests from weights already in HBM.

This separates two performance surfaces cleanly: presharding, allocator tuning, and device prewarm shorten time-to-ready; residency removes that startup path from steady-state requests altogether.

The headline: a measured, same-profile 8.79x Qwen serving speedup. Stable 20-step resident HTTP requests take 9.294 s versus 81.720 s for warm repeated-process CLI. Flux measures 8.53x on the same controlled comparison. TeaCache is disabled in both.

Three things make the serving layer production-grade. Batching and multi-tenant scheduling keep the fleet reliable under bursty production traffic: many models, many tenants, one pool of NeuronCores. Automatic fleet routing dispatches each request by its per-model cost/latency profile, with automatic GPU fallback for unsupported shapes, failures, or timeouts; the same layer extends to TPU as it comes online. And adoption is a routing change: Difflet keeps the diffusers-compatible interface and the industry-standard request shape. No migration, no retraining, no model surgery. Point traffic at it.

Live on Yotta Labs’ heterogeneous fleet

Difflet is built by Yotta Labs and runs on its fleet: NVIDIA GPUs, AMD GPUs, AWS Trainium, with TPU coming. Difflet routes image and video generation to under-used Trainium capacity. The payoff is double: every request moved off a GPU frees scarce GPU capacity for CUDA-only work, and idle Trainium earns inference revenue instead of nothing.

An engine proves itself under production traffic, and Difflet is rolling out across the fleet in deliberate stages:

1. Shadow traffic. Mirror production requests to Trainium and compare cost, latency, and quality, with zero user risk.
2. Fixed-percentage cutover. Shift a controlled slice of live traffic, with automatic GPU fallback backstopping every request.
3. Three models, then six. FLUX.1-dev, HunyuanVideo, and Wan first, then the full six-model suite.

Reported business metrics: dollars per inference, joules per inference, fleet utilization, and throughput per dollar.

Every number in this post is reproducible

One harness measures identical per-step DiT latency across Trainium2, Trainium3, H100, and B300 under pinned model revisions. Adding a backend is one adapter file. All protocols and reproduction commands are public.

# AOT compile + generate on Trainium
python -m benchmark.bench --model hunyuan_video

# per-step DiT latency (warm)
python -m benchmark.step_latency --model hunyuan_video

# H100/B300 reference
python -m benchmark.bench --backend cuda --model hunyuan_video

All figures measured with the public cross-backend harness under pinned model revisions. Trainium2 $2.235/accel-hr, H100 $5.191/accel-hr (AWS Capacity Blocks, July 2026).

For the full engineering detail behind these results, including the parallelism design, model-specific graph work, and complete measurement methodology, read the Difflet engineering report.