Searcharxiv⌕ Search
Home/Archive/arXiv

arXiv · 2610.23047

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

Bin Ma·Min Yu·Vijay Kodamalla·Jialin Liu·Dong Li
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

cs.DCcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Bin Ma, Min Yu, Vijay Kodamalla, Jialin Liu, Dong Li. 2026-10-16. DIFFLET: A Fast Inference Engine for Diffusion Transformer Models on AWS Trainium. https://arxiv.org/abs/2610.23047

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

One Communicator, Three Vendors: Extending the MoA Predictive Theory to a Real NVIDIA+AMD+CPU MPI Ensemble

The companion MPI-extension paper validated the first communication-bearing machine-shape ensemble: four NVIDIA A100 ranks and four CPU ranks on one NCSA Delta node under one Open MPI job, reassembled by MPI_Gatherv. This paper extends it to a multi-vendor communicator. A Slurm heterogeneous job spanning Delta's A100 and AMD MI100 partitions places A100, MI100 and CPU ranks (nine in all) under one communicator for the first time. We give the theory (a w_i proportional to 1/t_i weight split for three concurrent realizations; a het-job specification as an admissible directive) and report real results, including what the run exposed: the 1/t split assumes time linear in rows, and this fails for the MI100 rank (6 rows took 597 ms against a linear prediction of 1.6 ms). A launch-geometry model, with predictions recorded before measuring, succeeds for the A100 (time flat from 128 to 3456 rows, stepping at exactly 3520 and 6976) and places the MI100 floor in the key loop; changing its launch settings cut about 600 ms to 14.6 ms and the energy per call at 6 rows by 32x, though the MI100 remains about 95x slower than an A100 at 8192 rows. Controlled reruns show that the companion paper's 8.1x A100 result and 22 ms Gatherv were artifacts of four ranks sharing one GPU under Open MPI's default transport. A key/value-split partition whose message sizes are derived from shapes alone is then run: the ring Allreduce of the 8.5 MB message took 11.3, 16.0 and 15.3 ms against a predicted 11-12 ms, the small Allreduce missed its prediction in every run, and compute times reproduce to about 1% while communication times vary up to 4x between runs. Each idea is introduced with a figure and closed with worked examples.

cs.DC↗
arXiv2026

Understanding KV Residency in Production-Scale Agentic LLM Workloads

Agentic large language model (LLM) applications turn one user request into a sequence of model and tool calls over a growing context. Retaining this context's key-value (KV) state in GPU memory lets later calls reuse it instead of recomputing it, but at fleet scale this needs substantial memory. We characterize four production Microsoft M365 workload traces (Deep Research, Enterprise Assistant, Coding Agent and Enterprise Chat) spanning more than 121 million model calls. In Deep Research, the application issues 89% of model calls and the median context grows from about 10K to 68K tokens by the 30th call. About 80% of context tokens are served from the prefix cache. With the public Llama-2-70B KV layout standing in for confidential models, one 99th-percentile Deep Research session exceeds the 80 GB HBM of an NVIDIA H100, and holding all live sessions at their peak needs an estimated 501.8 TB of HBM, although inference occupies only 5.6% of that retention time. Trace-driven replay shows that a one-hour idle cap cuts this ceiling by 3.6x while raising the share of calls that recompute context only from 4.07% to 4.53%. Under tighter budgets, eviction alone is insufficient: at one tenth of the busiest hour's demand, LRU eviction makes 63% of Deep Research calls recompute and clairvoyant furthest-next-use 31%, whereas multi-resolution retention (MRR), which demotes the LRU session's KV state to int4, lowers this to 13.6%. Placement imposes a separate limit: when the next call reaches a different deployment, cached-context reuse drops from 89% to 29%, and treating KV state as deployment-local raises the recompute floor under unlimited memory from 5.4% to 20.2%. These findings motivate retention that changes the representation of KV state, together with deployment-aware routing or migration. We plan to release AgenticTraceSim, a parametric trace generator, and a residency replay framework.

cs.DC↗
arXiv2026

MailRail: Mitigating Incast Congestion with User-Space, Node-Based Flow Control

Irregular workloads, such as those stemming from Graph Neural Networks (GNNs) and Mixture-of-Experts (MoE) models, frequently induce severe incast congestion in large-scale parallel environments. Existing fabric-level congestion control mechanisms rely on reactive, link-level rate limiting that degrades overall throughput. While receiver-driven, credit-based flow control offers a proactive alternative, as process counts scale memory slots per rank fall below usable thresholds, rendering traditional rank-based accounting untenable. To address this challenge, we introduce MailRail, an out-of-order, asynchronous messaging framework, along with a novel user-space, node-based credit flow control system. By shifting credit accounting from individual processes to shared per-node pools, our proactive, fabric-agnostic approach maintains sufficient credit availability even as process density grows, enabling processes to dynamically utilize freed credits without waiting for receiver-side redistribution. Experimental evaluations demonstrate that node-based flow control achieves up to a 1132x speedup over unconstrained execution at an adversarial incast ratio of 4096:1 while adding minimal overhead on uniform workloads.

cs.DC↗