DIFFLET: A Fast Inference Engine for Diffusion Transformer Models on AWS Trainium

Bin Ma; Min Yu; Vijay Kodamalla; Jonny Liu; Da Li; Dong Li

cs.DCOct 16, 2026

DIFFLET: A Fast Inference Engine for Diffusion Transformer Models on AWS Trainium

Authors: Bin Ma, Min Yu, Vijay Kodamalla, Jonny Liu, Da Li, Dong Li

Organizations: Yotta Labs, USA · University of California, Merced, Merced, California, USA.

Abstract

On AWS Trainium2, every component of a diffusion model runs as a program compiled ahead of time for fixed shapes, while pipelines differ across models, prompt lengths across requests, and caching decisions across denoising steps. We present Difflet, a serving engine in which five model pipelines share one compilation backend, joint attention is split across cores by attention head, and a host controller decides at each denoising step whether to run the denoiser or reuse its earlier output. A small probe program computes the caching signal on the device and returns one number per step, and a resident model process keeps every program loaded across requests. Across five image and video model families, fused attention accelerates denoiser steps by 1.2-4.3x over the same framework using PyTorch/XLA attention, and at 32.8k visual tokens it is the only attention that compiles. On four of the five models, resident serving is 1.5-3.4x faster than standalone runs that reload weights from a warm file cache for each generation.

Figures & tables

Figure 1: Model abstraction and Trainium adaptation in Difflet. (a) Pipelines specify component dataflow and configurations. (b) Component, operator, and parallelism interfaces connect models to (c) Trainium compilation and execution. Solid arrows denote dataflow; dashed arrows denote configuration and adaptation. NeuronCore icons are schematic.

Figure 2: Joint attention with visual tokens partitioned and text replicated across two ranks. (a) Ulysses gives each head group the full joint Q/K sequence with equal lengths; bounds

[0,6)

exclude text padding. (b) Ring uses the same local queries in separate text and visual calls with unequal Q/K lengths, then merges unnormalized outputs and softmax statistics. Starred: the text bounds

[0,2)

that Ring needs for padded prompts. Both paths restore visual sequence shards and replicated text.

Figure 3: Calibrated adaptive caching with a device probe in Difflet. The probe retains features on Trainium and returns a scalar score to the host CPU. Teal denotes full denoiser execution; gold denotes cached prediction. Solid arrows show data flow; dashed arrows show execution choices.

Figure 4: Serving in Difflet. (a) Startup selects artifacts for the serving profile, loads components, and checks readiness. Dashed arrows connect loading to the model process and Trainium. (b) The API process validates and queues requests. One resident model process orchestrates encoding, denoising, and decoding and controls the device cores; component placement depends on the model configuration. (c) After a timeout or failure, dispatch waits for prior execution to stop and the model process to become ready, with a restart when needed.

model	params	H $\times$ W $\times$ F	tokens	steps	guidance
FLUX.1-dev	12B	$1024^{2}$	4,096	50	3.5 dist.
Qwen-Image	20B	$1024^{2}$ b	4,096	50	4.0 CFG
Wan 2.1	14B	$480{\times}832{\times}81$	32,760	50	5.0 CFG
HunyuanVideo	13B	$320{\times}512{\times}61$ a	10,240	50	6.0 emb.
LTX-2	19B	$480{\times}704{\times}49$ a	2,310	40	1.0 c

Table 1: Evaluated models. Tokens are the visual tokens the denoiser attends over per step; guidance is distilled (dist.), classifier-free (CFG), or embedded (emb.).

system	AOT	NKI	preshard	resident	cache	parallel
PyTorch/XLA	Y	$\times$	Y	$\times$ d	Y	TP, CP e
NxDI	Y	Y	$\times$	$\times$	$\times$	TP
Difflet	Y	Y	Y	Y	Y	5 axes f

Table 2: Systems compared and the mechanisms each provides: ahead-of-time compilation, NKI attention, presharded weights, resident serving, step caching, and parallelism. NxDI covers FLUX.1-dev only. Y: yes;

\times

: no.

	compile	latency (s)			step (ms)
model	(min)	cold	warm	resident	NKI	SDPA
FLUX.1-dev	16.7	318	47.6	14.2	270	682
NxDI	15.7	337	66.3	–	271	–
Qwen-Image	24.1	537	98.9	42.4	835	1,582
Wan 2.1	75.1	1,311	974	1,245	14,876	– g
HunyuanVideo	85.8	625	141	59.9	845	3,620

Table 3: End-to-end latency at TP=4. Cold and warm are standalone runs; resident is the request latency of the model process; step is the step latency with NKI and with SDPA attention.

Figure 5: Where the request-latency speedup comes from, adding Difflet’s mechanisms one at a time to a standalone PyTorch/XLA run. Each bar is a model’s starting latency; colored segments are the time each mechanism removes, labeled with that step’s speedup where it fits, and gray is the final latency.

\times

: Wan’s PyTorch/XLA program does not compile, so its bar starts from NKI attention. Caching is fixed cadence 2; for FLUX, Wan, and LTX-2 it is the resident caching-off request minus the denoising time cadence 2 saves. LTX-2’s first two stages are medians of two alternated runs (106–113 s).

Figure 6: Step speedup of NKI attention over PyTorch/XLA’s SDPA, each model at its evaluated shape, ordered by visual tokens per step.

\times

: Wan 2.1’s SDPA program does not compile at 81 frames.

		CP
Model	TP	Gather	Ring	Ulysses	SP	CFG	DP
FLUX.1-dev	Y	Y	$\dagger^{1}$	Y	Y	–	Y
Qwen-Image	Y	Y	Y	Y	Y	$\times$	Y
Wan 2.1/2.2	Y	Y	$\dagger^{1}$	$\dagger^{4}$	Y	$\dagger^{4}$	$\dagger^{4}$
HunyuanVideo	Y	$\dagger^{2}$	$\times^{3}$	Y	Y	–	Y
LTX-2	Y	$\times$	$\times$	$\times$	$\times$	Y	Y

Table 4: Parallel modes each model runs. Y: supported;

\dagger

: constrained;

\times

: unsupported; –: not applicable.

Figure 7: Parallel configurations for FLUX.1-dev (4.1k tokens) relative to TP4: resident request latency and throughput. DP2

\times

TP2 serves two concurrent requests; its latency is the wall time for both and its throughput counts both.

	FLUX (4.1k tok.)		Wan 81f (32.8k tok.)
config	req. (s)	step (ms)	req. (s)	step (ms)
TP4	13.97	272.5	974	14,876
TP4+SP	14.28	278.7	933	13,882
TP2CP2, gather	14.29	278.8	–	–
TP2CP2, Ulysses	13.51	263.3	HBM q	HBM q
TP2CP2, ring	13.75	268.2	–	–

Table 5: Request latency and step latency across parallel configurations. FLUX is resident; Wan is standalone. DP rows serve two concurrent requests, so their latency is the wall time for both. FLUX TP4 is a separate run from Table 3 (14.2 s, 270 ms).

Figure 8: Step caching at official step counts. Left: denoising-time speedup over caching off (caching-off time under each model); dotted: the ceiling set by cadence 2’s 20 of 50 skipped calls (15 of 40 on LTX-2); fewer steps runs cadence 2’s 30 calls (25 on LTX-2) without caching. Right: PSNR against the model’s caching-off output, mean over 16 prompts (five on Wan), whisker to the worst prompt. Online-delta’s speedup is for the timing prompt. Calibrated adaptive runs standalone, against the standalone caching-off time, and its Wan and LTX-2 points come from an earlier denoiser program, as does LTX-2’s quality; LTX-2’s online-delta mean omits one prompt that skipped no step. Wan’s device throttles during its 12-minute caching-off loop, which lifts its speedups past the ceiling; its quality is scored without calibrated adaptive.

model	device probe	ms/step	fit $R^{2}$
FLUX.1-dev	fused in the denoiser	2.3	0.60
Qwen-Image	fused in the denoiser	7.3	0.68
Wan 2.1	separate program	– o	0.13
HunyuanVideo	separate program	17.8	0.71
LTX-2	separate program	2.6	0.82

Table 6: Signal of calibrated adaptive caching: the probe’s cost per scheduler step and how well the fitted curve predicts the output change (

R^{2}

over 8 recording prompts).

model	(min)	cold	warm	recompiles
	first start	restart (s)
FLUX.1-dev	18.5	340	91	1 of 1
Qwen-Image	34.8	640	129	1 of 3
Wan 2.1	66.9 v	1,102 u	786 u	– p
HunyuanVideo	116.8	624	194	– p
LTX-2	33.8	811	155	– p

Table 7: Startup of one resident model at TP=4, until the server reports ready. Recompiles counts the programs rebuilt when switching from caching off to calibrated adaptive.

	p50 / p99 (s)		SLO met
FLUX, $\rho$	TP4	DP2 $\times$ TP2	TP4	DP2 $\times$ TP2
0.78	39.4 / 65.3	39.1 / 70.0	40%	30%
0.99	33.0 / 53.8	46.7 / 69.3	40%	25%
1.18	65.9 / 105.3	74.7 / 110.7	20%	20%
HunyuanVideo, $\rho$	TP4		TP4
0.81	136.1 / 244.2	n/a	45%	n/a

Table 8: Request latency under open-loop Poisson load through difflet serve over HTTP: for FLUX, one TP4 server against two TP2 servers on separate core pairs, each request sent to the server with fewer requests in flight; for HunyuanVideo, one TP4 server (its TP2 denoiser does not fit in HBM). Both configurations receive the same arrival sequence; 20 requests per load level. Load

\rho

is the realized arrival rate times the isolated TP4 request latency (FLUX 14.2 s; HunyuanVideo 58.5 s in this run, 59.9 s in Table 3 ). SLO: completion within 2.1

\times

that latency (30 s and 123 s).

Diffusion Model Inference Acceleration

Diffusion model inference acceleration reduces the time and compute needed to generate or reconstruct outputs while preserving quality.

See all 506 papers

SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers
Jun 28, 2026Xuanhua Yin, Yuxuan Jia, Chuanzhi Xu +1Diffusion Transformer Diffusion Model Inference Acceleration
School of Computer Science, The University of Sydney, Australia
Denoising Diffusion Generative Models Secretly Calculate Attentions
Sep 1, 2026Farzan Haddadi, Leila Monfared, Ebrahim Rezaii +3Attention Control in Diffusion Models Diffusion-Based Image Generation
School of Electrical Engineering, Iran University of Science & Technology, Tehran, IRAN · Independent researcher
Beyond Fixed Formulas: Data-Driven Linear Predictor for Efficient Diffusion Models
Apr 29, 2026Zhirong Shen, Rui Huang, Jiacheng Liu +6Diffusion Transformer Diffusion Model Caching
Shanghai Jiao Tong University · University of Electronic Science and Technology of China · Shandong University +2

Jun 28, 2026J/K move · Enter open · S save

SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers

Xuanhua Yin, Yuxuan Jia, Chuanzhi Xu, Weidong Cai

School of Computer Science, The University of Sydney, Australia

Diffusion Transformer Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

High-resolution Diffusion Transformer (DiT) inference contains substantial spatial redundancy, but many spatially adaptive implementations encode regional computation as attention masks, which can inadvertently move scaled dot-product attention (SDPA) away from FlashAttention fast paths. We identify this avoidable systems bottleneck as Mask-Induced Dispatch Tax (MIDT) and show that it grows with latent sequence length. We introduce SAFE-DiT, a training-free Semantics-Aware Fast-path Execution framework that separates exact mask elision from approximation-based spatial scheduling. SAFE-DiT removes only provenance-certified image self-attention masks that induce a row-wise constant shift in attention logits, preserves semantics-bearing masks such as text-padding masks, and realizes spatial adaptation through prompt-conditioned token partitioning, selective state updates with global context, and periodic context refresh. We call this acceleration-only configuration SAFE-Core and report sensitivity-weighted classifier-free guidance separately as SAFE-DiT+SW. On the evaluated PyTorch SDPA stack, redundant masks make long-sequence attention

4.1\times

5.8\times

slower than the mask-free path. On Lumina-Next, SAFE-DiT achieves

2.69\times

end-to-end acceleration at

1024^2

resolution and

5.09\times

2560^2

, reduces peak memory at

2560^2

from 94.1 to 27.9 GB, and enables

3072^2

generation when dense inference runs out of memory. Paired metrics, component ablations, and a blinded human study support visual non-inferiority of SAFE-Core to the dense fast-path baseline, while SAFE-DiT+SW provides a separate prompt-alignment operating point without reintroducing spatial self-attention masks. Code is available at https://github.com/xuanhuayin/SAFE-DiT.

Abstract

Diffusion Model Inference Acceleration

Explore similar work

SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers

Denoising Diffusion Generative Models Secretly Calculate Attentions

Beyond Fixed Formulas: Data-Driven Linear Predictor for Efficient Diffusion Models