Difflet Engineering Report: Diffusion Inference on AWS Trainium

Architecture, compilation, parallel execution, model loading, resident serving, and the measurements used to validate the implementation.

Yotta Labs systems engineering. Code reviewed through 2026-07-19. Measured devices: Trainium2, Trainium3, H100, B300.

This is the full engineering report behind Difflet, Yotta Labs’ diffusion serving engine for AWS Trainium. For the overview and headline results, start with the Difflet announcement post.

Abstract

Difflet implements image and video diffusion inference with a backend-neutral model layer and a Trainium backend. The code covers ahead-of-time (AOT) compilation, content-addressed artifacts, single program, multiple data (SPMD) execution, tensor/context/classifier-free-guidance (CFG)/sequence parallelism, optional request replication, model-specific attention paths, per-rank checkpoint loading, optional step caching, and a resident HTTP worker. This report describes those mechanisms first. Benchmark and serving data are included to state what has been measured and to mark the remaining gaps.

Code and measurement scope

The local checkout and the latest mainline history contain different parts of the implementation. The distinction matters when reading the serving sections.

Measurement terms used below

• DiT (diffusion transformer) per-step: Device-synchronized denoiser compute, excluding model load. This is the primary cross-device comparison.
• warm E2E (legacy CLI): A new process with weights already in the OS page cache. It still reloads and initializes the model.
• resident request: An HTTP request accepted after worker readiness, with the compiled model already loaded.

Important: device prewarm is a worker-start optimization, not a request optimization. It was merged after the original benchmark matrix, and a full prewarm-on/off end-to-end A/B has not yet been recorded. We therefore show its measured ~6.7 s overlap opportunity without claiming an unmeasured E2E speedup.

Terms used in this report

Short definitions for readers familiar with GPU inference but not the AWS Neuron stack or diffusion-model naming.

• AOT · ahead-of-time compilation: Compile a fixed graph before serving rather than compiling each operation when a request runs.
• NEFF · Neuron Executable File Format: The compiled executable that Neuron Runtime loads onto a Neuron device. It contains the compute graph and its device execution plan.
• NRT · Neuron Runtime: The host runtime that discovers Neuron devices, loads NEFF files, manages device resources, and launches compiled graphs.
• NKI · Neuron Kernel Interface: A Python language and library for writing kernels that run directly on Trainium NeuronCores.
• NxD · NeuronX Distributed: The Neuron SDK library used here to shard models and coordinate distributed inference across NeuronCores.
• LNC · Logical NeuronCore: The core unit visible to software. On Trainium2, one logical core may group two physical NeuronCores; an LNC=2 kernel therefore uses two physical cores as one logical execution unit. The compile-time setting must match the runtime setting.
• HLO · High Level Optimizer IR: XLA’s compiler intermediate representation: the graph form inspected or transformed before target-specific lowering.
• SPMD · single program, multiple data: Every rank runs the same traced program while operating on its own tensor or sequence shard.
• DiT · diffusion transformer: The transformer denoiser that predicts the update applied to image or video latents at each diffusion step.
• TP / CP / SP: Tensor, context, and sequence parallelism. They shard weights or activations, token context, and selected sequence regions respectively.
• CFG · classifier-free guidance: Combines conditional and unconditional denoiser predictions. CFG parallelism evaluates those two branches on separate ranks.
• SDPA · scaled dot-product attention: PyTorch’s general attention operator for softmax(QKᵀ / √d)V, with optional masks.
• Attention CTE · Context Encoding: An NKI Library attention kernel for prompt/prefill-style attention. It tiles the same attention calculation for Trainium and accepts per-query valid-key bounds, avoiding a dense padding mask.
• MMDiT · multimodal diffusion transformer: A diffusion transformer that jointly processes image/video and text token streams.
• AdaLN · adaptive layer normalization: Layer normalization whose scale and shift are conditioned on the diffusion timestep and, depending on the model, other conditioning signals.
• UniPC · Unified Predictor-Corrector: A multistep diffusion sampler/scheduler. Its latent and history state remain fp32 in the corrected video path.

Module boundaries and artifact lifecycle

A single public pipeline resolves a model, derives a content-addressed compile identity, and dispatches to a model application. Model code imports a frozen backend-neutral op surface; Trainium-specific kernels, collectives, and runtime details live underneath it.

Backend operation boundary

Model implementations call difflet.ops for attention, linear layers, normalization, embeddings, collectives, and platform queries. They do not import NKI or Neuron runtime modules directly. Backend dispatch selects CPU reference implementations for unit and numerical tests or Trainium implementations for compiled execution.

• Graph ownership: difflet/models owns model structure and tensor layouts; difflet/backends/trainium owns kernels, NxD wrappers, process groups, tracing, and runtime configuration.
• Why the boundary exists: The same model-level operation can be compared against a CPU definition without putting accelerator imports into the public model graph.
• Extension rule: The public op surface is additive. A new accelerator implementation is registered behind the dispatcher rather than selected with device checks inside model code.
• Failure boundary: Unsupported layouts fail at the backend or model adapter seam; they do not silently change the public pipeline contract.

Compile identity and artifact contract

The compile-cache key is a 16-character prefix of a SHA-256 digest over normalized compile inputs: model and pinned revision, parallel configuration, dtype, fixed height/width/frame shape, compile-affecting application options, and relevant toolchain versions. This prevents reuse across incompatible sequence lengths, world sizes, attention modes, or compiler ABIs.

• Stored artifact: A schema-versioned manifest, component Neuron Executable File Format (NEFF) files, and—where supported—one presharded checkpoint per rank.
• Excluded from identity: The resolved local model path and Python patch version remain diagnostic metadata, so equivalent hosts do not miss the cache only because their directory or micro-version differs.
• Runtime-only options: Host text/decode component loading is excluded from the key because it does not change the compiled NEFF graph.
• Validation: A missing, unreadable, schema-mismatched, or input-mismatched manifest is a cache miss. The runtime does not load an artifact on filename alone.

Code ownership

Application composition

Applications are composed in two forms because model placement and compiler limits differ. Both forms still implement the same download, compile, load, and call lifecycle.

Opaque pipeline application. FLUX and LTX-2 are exposed as one logical callable. LTX-2 remains hybrid: the TP4 transformer runs on Neuron while text, connectors, VAE, and vocoder execute on the host.

Extracted stage application. Wan, HunyuanVideo, and Qwen-Image expose prompt encoder, denoiser, and decoder boundaries. CLI stages exchange files; resident serving passes exact typed payloads in process.

Selection rule: use an extracted stage when placement, core count, or failure recovery differs by component; use an opaque application when the upstream pipeline owns tightly coupled host logic that is not usefully split.

Intra-request parallelism and request replication

The logical device mesh is {dp, cfg, cp, tp}. Model collectives use named process groups. Mainline also provides a separate CLI router that starts independent replicas on disjoint core ranges for request-level data parallelism.

Tensor parallel: shard projections and attention heads

TP divides projection weights and their output features across ranks. Column-parallel projections produce local feature shards; row-parallel projections combine partial results. Attention heads are divided with the corresponding Q/K/V projections instead of being replicated.

• Models: FLUX, Qwen-Image, Wan 2.1/2.2, HunyuanVideo, and LTX-2 all have TP implementations. The measured Trainium matrix uses TP4.
• Rank contract: Every loaded component shares one process world size. The largest-TP component loads first because it establishes the communicator used by later components.
• Model-specific work: Wan also computes Q/K RMS from a cross-rank sum of squares so normalization remains equivalent after head sharding.
• Cost: TP reduces per-rank weight and compute volume but introduces collectives at row-parallel boundaries.

Context parallel: shard image and video tokens

CP leaves a rank with a local query-token range and supplies the K/V information needed to attend outside that range. Text tokens remain replicated where the model uses joint image/text attention.

• gather_kv: All-gathers K/V inside the CP group, then computes attention for local queries. It is the default and simplest correctness reference.
• ring: Rotates K/V shards through the CP group and merges partial results with the online-softmax max/sum correction, avoiding a full K/V materialization on each rank.
• ulysses: Mainline all-to-all changes sequence shards into head shards for attention and reverses the exchange afterward. Head count must divide tp_degree × cp_degree.
• Boundaries: Masked Ulysses is rejected for Qwen-Image and HunyuanVideo. CP consumes the same extra mesh lanes as CFG parallel and cannot be enabled with SP in the current implementation.

CFG parallel: split conditional and unconditional evaluations

True classifier-free guidance evaluates an unconditional and a conditional denoiser branch. CFG parallel places those branches on separate mesh lanes, gathers both predictions, and applies the guidance combination once.

• Active model paths: Wan and LTX-2 have true two-pass CFG paths. Without CFG parallelism they execute two batch-1 calls sequentially because the compiled artifact remains batch 1.
• Not applicable: The report’s FLUX, Qwen-Image, and HunyuanVideo profiles are guidance-distilled or otherwise expose no CLI-level second branch to split.
• World size: The CFG axis has degree 2, so world_size = tp × 2 when it is enabled.
• Interaction: CFG parallel is mutually exclusive with CP. The serial residual used by DiT caching is also disabled when CFG branches are distributed across ranks.

Sequence parallel: shard non-matmul activation regions

SP reuses the existing TP group to shard normalization, modulation, and residual regions along the sequence axis. It does not add a new mesh dimension or increase world size.

• Collective change: A row-parallel all-reduce becomes reduce-scatter so each rank retains a sequence shard; an all-gather restores the full sequence before the next column-parallel projection.
• Implemented models: FLUX, Wan, and HunyuanVideo are device-verified. Qwen-Image is excluded because its upstream forward does not expose the per-rank identifier as a live graph input; LTX-2 has no SP path.
• Compatibility: sp_enabled requires cp_degree == 1. Both modes shard sequence work, but over different process groups.
• Validation: The recorded dense-versus-SP device comparisons have cosine similarity at or above 0.999 for the enabled models.

ulysses requires cp_degree > 1 and the number of attention heads to be divisible by tp_degree × cp_degree. Qwen-Image and HunyuanVideo reject masked Ulysses attention. CP and CFG parallelism are mutually exclusive; SP also requires cp_degree == 1.

Implemented model-parallel feature surface

* Ulysses is present in origin/main after the checked-in device verification matrix. This table says the code path exists; it does not add Ulysses performance claims.

Request replicas and preset modes

The mainline DP router writes a request manifest, divides visible cores into non-overlapping ranges, starts one CLI process per replica, and gathers per-request completion or failure markers. Workers retain dp=1 model semantics; routing is outside the compiled model. An HBM-fit check runs when local weights are available.

Runtime invariant: current diffusion artifacts use one Python process with multiple visible NeuronCores. All components loaded together must share the same world size, and the largest-TP component must initialize first because it fixes the process-wide communicator.

Model-specific graph and runtime choices

The public lifecycle is shared, but the graph structure, attention layout, component placement, and compiler constraints differ by model. These differences are encoded in each model application and its CLI or serving adapter.

FLUX.1-dev

The guidance-distilled MMDiT alternates dual-stream text/image blocks and single-stream joint blocks. Difflet owns the Trainium graph, so TP, CP gather/ring/Ulysses, and SP are inserted at the model’s native tensor boundaries.

• Placement: The transformer is TP4. Replicated CLIP/VAE components still load under the same world-size-4 process contract, after the TP component initializes the communicator.
• Attention layout: Dual-stream and single-stream blocks pack text and image tokens differently; CP scatters only the image region and preserves the joint ordering required by each block type.
• Guidance: The measured profile uses guidance distillation and therefore has no second CFG branch to parallelize.
• DiT cache path: An optional fused block-0 probe keeps its previous signal on device and returns only a scalar change estimate to the host controller.

Qwen-Image

Qwen-Image wraps the upstream diffusers transformer instead of maintaining a separate full forward. Parallel behavior is therefore inserted only at seams controlled by Difflet.

• CP insertion: Scatter packed image hidden states and image RoPE, retain full replicated text tokens, gather K/V for attention, then gather the packed image output.
• Serving stages: Mainline separates text encoder, denoiser, and VAE runners. The accepted serving topology keeps all three compatible with TP4/world-size 4.
• SP boundary: SP is disabled because the monkey-patched upstream forward does not carry Difflet’s per-rank identifier as a live traced input; every rank would otherwise read rank zero.
• Mask boundary: Ring and Ulysses paths reject masks that their distributed attention implementation cannot preserve.

Wan 2.1 and Wan 2.2 A14B

Wan 2.1 runs one 14B transformer. Wan 2.2 selects a high- or low-noise 14B expert from the current timestep; only one expert executes during a denoise step.

• Attention: Q/K/V projections and heads are TP-sharded. Q/K RMS uses a cross-rank sum of squares so each shard applies the global normalization value.
• Guidance: True CFG runs as sequential batch-1 calls or on two CFG lanes. TeaCache uses the combined post-CFG prediction as its residual unit and is disabled for the distributed CFG path.
• Expert switch: Wan 2.2 resets cached residual and gate state when the active expert changes; state from the high-noise network cannot be reused by the low-noise network.
• Decode: The Neuron VAE uses a one-core artifact. Longer clips can exceed the single-shot compiler instruction limit, so host causal decode remains the rollback path.

HunyuanVideo 1.0

Separate CLIP and Llama prompt paths feed a joint image/text DiT. The text key-padding mask is contiguous right padding, which permits an exact bounds representation for Attention CTE.

• Mask conversion: The model derives per-query [bound_min, bound_max) ranges outside the generic traced attention resolver and passes them into the NKI kernel.
• Serving stages: Mainline exposes CLIP, Llama, denoiser, and decoder runners with explicit host/Neuron placement bindings.
• Parallel boundary: Masked ring and Ulysses are rejected. The checked tp2cp2 CLI profile remains a compiler failure rather than a supported result.
• DiT cache signal: Block-0 modulation includes timestep and pooled text conditioning, making it a viable adaptive signal in the recorded gate.

LTX-2

The TP4 transformer jointly advances video and audio streams. Its 48 dual-stream blocks execute through a segmented runtime because one monolithic graph exceeds compiler limits.

• Placement: The transformer runs on Neuron; text encoder, connectors, video/audio decoders, and vocoder remain on the host. Serving therefore treats the hybrid pipeline as one opaque application.
• Parallel boundary: LTX-2 has TP and true CFG, but no CP or SP. Non-parallel CFG uses two sequential batch-1 calls to match the compiled shape.
• Segment execution: Blocks are loaded and run in segments to stay within compiler limits; this changes runtime scheduling but not the external pipeline contract.
• DiT cache state: The cache stores raw video and audio velocity residuals—the values consumed by the schedulers—rather than caching x0 and converting it at small sigma.

Compiler- and runtime-driven fixes

• Masked attention routing: an initial generic in-graph mask-to-bounds conversion failed an XLA broadcast during Hunyuan compilation. The final implementation keeps the generic resolver but performs Hunyuan’s contiguous-mask conversion in a trace-safe model path. Arbitrary or soft masks continue to fall back to PyTorch scaled dot-product attention (SDPA).
• Wan VAE logical-core configuration: the decoder runs as a one-core stage. Its artifact is explicitly compiled with one Logical NeuronCore (LNC=1); inheriting Trainium2’s LNC=2 default made the artifact expect a logical execution unit formed from two physical NeuronCores, so Neuron Runtime (NRT) refused to load it under the one-core runtime configuration.
• LTX-2 CFG shape: the default artifact is batch 1. Without CFG parallelism, conditional and unconditional bundles are evaluated sequentially rather than passing an invalid batch-2 input to the compiled graph.
• Text-encoder presharding: Hunyuan Llama and Qwen LLM stages use the stock NeuronX Distributed (NxD) configuration rather than Difflet’s fork. Their compile path explicitly requests shard writes, while load enables presharding only after confirming that shard files exist.
• Ulysses artifact identity: staged directory names originally omitted cp_mode, so ring and gather artifacts could collide despite different content hashes. Mainline adds the mode token while keeping the default gather name backward compatible.

Video output corrections

Stage-by-stage comparison against diffusers found issues outside the DiT kernel: MP4 export received pre-converted uint8 frames and wrapped pixel values; Wan padding rows were not zeroed before unmasked cross-attention; initial latents used a smoke-test scale rather than prepare_latents; Unified Predictor-Corrector (UniPC) loop latents were held in bf16 instead of fp32; and negative CFG embeddings were replaced with zeros. The corrected paths export float [0,1] frames, use the reference latent preparation and negative-prompt encoding, keep scheduler state in fp32, and retain CPU attention bounds. These fixes changed output fidelity rather than the per-step benchmark method.

Implemented execution and startup changes

AOT graph specialization

Difflet traces fixed model, shape, dtype, and parallel topology into NEFF files before generation. Static tensor extents let the compiler choose layouts, schedule engines, and allocate device memory without a dynamic-shape decision in the request loop.

• Specialized inputs: Resolution, frame count, sequence lengths, TP/CP/CFG/SP topology, attention mode, dtype, and compile-affecting application options.
• Runtime result: A request loads an already compiled graph; changing a specialized dimension requires a compatible cached artifact.
• All models: Every Trainium model path uses AOT artifacts, including text encoders and one-core VAE stages where those components run on Neuron.
• Guard: Manifest validation and the content-derived cache key prevent a graph compiled for one topology or shape from being used by another.

Attention CTE routing and mask representation

Unmasked attention routes directly to the NKI Library Attention CTE kernel. A hard mask can use the same path only when every query’s valid keys form one contiguous interval; arbitrary sparse masks and soft attention biases remain on the PyTorch SDPA fallback.

• Hunyuan: Right-padded text keys are converted in a trace-safe model path into per-query bound_min/bound_max. This corrected the 3719 → 850.6 ms routing case.
• LTX-2: Self-attention uses CTE. Text cross-attention dropped a benign padding mask only after full-generation parity, changing 477.1 → 441.8 ms in the recorded profile.
• Safety rule: Boolean/integer hard masks and large-negative additive masks are accepted only after a contiguity check. Positive values, mild negative biases, empty ranges, or head-varying masks return to SDPA.
• Compiler rule: Generic mask analysis is not executed inside the traced forward: that version caused an XLA broadcast failure. Models pass resolved bounds explicitly.

Wan attention head sharding and global RMS

Wan originally replicated attention work after TP projections. The corrected graph retains local Q/K/V head shards, computes attention locally per head shard, and combines only at the projection boundary.

• Normalization: Each rank contributes its local squared sum; a TP reduction produces the global Q/K RMS used by every shard.
• Why it matters: Applying RMS independently per shard would change the attention scores even if the subsequent matrix operations were correctly sharded.
• Measured profile: 1144 → 554.8 ms per step, a 2.06× reduction, with cosine similarity 0.999768 against the replicated baseline.
• Scope: This change applies to Wan 2.1 and both Wan 2.2 experts because they share the same execution architecture.

Per-rank presharded checkpoints

Compilation serializes weights/tp{rank}_sharded_checkpoint.safetensors. A later process reads only the shard required by its rank instead of loading the full checkpoint and repeating the model-to-rank transformation.

• Diffusion components: Difflet’s Trainium application base writes shards by default and checks that every expected file exists before selecting the presharded path.
• Text encoders: Hunyuan Llama and Qwen LLM use stock NxD configuration. Their compile stage explicitly requests shard writes; load enables the shard path only after checking all rank files.
• Recovery: A partial, old, or cleared cache falls back to shard-on-load instead of raising a raw missing-file error.
• Effect: This reduces time-to-ready and repeated-process E2E; it does not change steady-state DiT per-step latency.

jemalloc for concurrent rank loading

The Neuron loader creates one host loading thread per rank. Those threads allocate tensors and fault pages concurrently, which caused contention in glibc allocation arenas and the process mmap lock.

• Activation: generate and run locate the jemalloc library shipped with torch_neuronx, set LD_PRELOAD, and re-exec once before importing the runtime.
• Compile exclusion: The compile command does not preload jemalloc because inherited preload state caused neuronx-cc worker failures.
• Fallback: If the library is missing, already loaded, or DIFFLET_NO_JEMALLOC=1 is set, execution continues with the system allocator.
• Measured case: The clean LTX-2 median moved from 63.4 to 58.4 s; the isolated weight-load portion was about 17% shorter.

Device prewarm

The first device allocation lazily runs nrt_init → tdrv_init, resets cores, and waits for readiness. Difflet starts that work on a daemon thread while the main thread reads and initializes host weights.

• Single-process paths: FLUX and LTX-2 call prewarm before loading their resident application.
• Staged paths: The internal stage dispatcher calls prewarm for generate stages using NEURON_RT_NUM_CORES or the stage TP×CP count; compile stages do not touch the device.
• Failure behavior: Every import, device, and busy-core error is swallowed. The main thread’s first device operation remains authoritative and pays the normal initialization cost if prewarm did not finish.
• Evidence boundary: The observed lazy initialization is about 6.7 s. It is an overlap opportunity; a full prewarm-on/off worker-start A/B is still missing.

Resident workers

A serving worker binds one immutable profile to exact artifacts, loads the compiled components and rank weights once, runs a smoke request, then remains alive for accepted requests.

• Removed from warm requests: Checkpoint reads, preshard selection, communicator setup, NEFF load, stage construction, and first-use runtime initialization.
• Still per request: Validation, admission wait, prompt work, denoise/decode execution, encoding, and configured artifact publication.
• Readiness: The service becomes ready only after the profile matches the artifact set and a real model-specific smoke output passes validation.
• Recovery: A timed-out or cancelled request enters a recovery fence; an unresponsive worker is terminated and recreated before readiness returns.

DiT caching: adaptive TeaCache, online delta, and fixed cadence

The shared controller stores the difference between the two most recent full denoiser predictions. On a skipped step it advances the prediction as previous prediction + cached residual and avoids the full DiT call. Warmup, cooldown, and two completed full steps are required before a skip is allowed.

• Adaptive TeaCache: Measure change in the block-0 modulated input, rescale it with a fitted polynomial, accumulate the estimate, and run a full step when the threshold is crossed.
• Online delta: Use the previous full step’s output relative-L1 change against a baseline. It uses no block-0 probe and forces a full step after every skip so the trajectory is re-measured.
• Fixed cadence: Skip by step index inside the warmup/cooldown window. It is calibration-free and bypasses both the device probe and the CPU block-0 shadow.
• Mode selection: A real-conditioning gate chooses adaptive when block-0 signal/output Pearson is at least 0.8, online delta when output-delta lag-1 autocorrelation is at least 0.7, and fixed cadence otherwise.

Reducing host/device communication

For fused adaptive paths, the previous block-0 signal is a persistent on-device parameter updated through input/output aliasing. The probe computes the change on device and returns only one scalar to the host, rather than materializing the roughly 63 MB Hunyuan block-0 tensor on the host every step. A committed multi-step skip run avoids additional probe launches. Fixed cadence needs no signal at all; online delta derives its decision from the already-produced denoiser output and likewise avoids a model-specific block-0 transfer.

How the first-block signal is handled by model

Deployment status: these are opt-in runtime paths. TeaCache is disabled in every main cross-device benchmark table and is not assumed in the resident-request numbers; a serving profile must select it explicitly.

Presharding A/B on Trn2 — warm repeated-process CLI

* HunyuanVideo is not a clean load-only comparison because the newer run also moved VAE decode on-chip. Presharding does not alter DiT per-step latency; it changes time-to-ready and repeated-process E2E only.

Measured step-cache cases

HunyuanVideo records 1.91× end-to-end with output checked against an uncached run. Wan’s 50-step cadence-2 A/B records 1.77× denoise speedup at final-latent cosine 0.9997. These measurements are separate from the full-step benchmark matrix and are not multiplied into its results.

A backend-generic harness with pinned workload identity

The harness owns the metric schema and timing logic but imports no accelerator SDK. A Trainium adapter drives Difflet AOT execution; a CUDA adapter drives stock Hugging Face diffusers. Model revision, shape, step count, dtype, prompt, guidance, and seed are pinned in one matrix.

• Prepare and compile Download the pinned revision. Trainium builds an AOT artifact; eager CUDA reports zero compile time.
• Measure the real denoise loop Record device-synchronized inter-step deltas and drop step zero. This matches the CUDA callback_on_step_end quantity where supported.
• Record cold and warm process paths Trn2 cold runs drop the OS page cache. Warm runs are subsequent new processes with the page cache hot; they still reload the pipeline.
• Qualify resident serving separately Measure the complete localhost HTTP request after readiness, including validation, inference, encoding and configured artifact publication where stated.

Cross-device boundary: per-step DiT latency is the intended comparable metric. Cold/warm E2E includes backend-specific loading, process topology, host work, and page-cache behavior. Qwen, Wan, and Hunyuan Trn2/Trn3 per-step values retain an older isolated-forward timer, while FLUX/LTX and GPU values use real-loop deltas; those mixed-method rows carry extra uncertainty.

Trainium. TP4 AOT execution on one device. Trn2: 4 logical cores, 96 GB. Trn3: 4 cores, 144 GB.

H100. PCIe 80 GB, stock diffusers, single-GPU dense. Warm runs use a separate process because a second in-process pipeline does not fit.

B300. SXM6 275 GB, same CUDA adapter and software versions as H100. Runs cold plus one warm iteration in process.

DiT per-step and process-level results

The table reports raw measurements from the pinned matrix. B300 has the lowest per-step time on five rows; Trainium3 has the lowest HunyuanVideo value. Trn2 and H100 are approximately equal on both Wan rows, while Trn2 is lower on FLUX and HunyuanVideo.

DiT latency per denoising step — milliseconds, lower is better

* Wan 2.2’s displayed Trn3 value reuses Wan 2.1 because the execution architecture and compiled graph are shared; it was not independently measured.

Derived ratios from the per-step table

Cost model: Trn2 $2.235/accelerator-hour and H100 $5.191/accelerator-hour from AWS Capacity Blocks pricing, accessed July 2026. It excludes host, storage, networking, fleet utilization, and volume discounts.

Legacy warm E2E — seconds, new process with hot OS page cache

Note: This warm table is retained because it explains the load-path work. It should not be used as the resident-request table: every row constructs or reloads a model process.

Load and device-initialization sequence

Presharding removes repeat work, jemalloc makes the remaining host load scale across rank threads, and prewarm overlaps device bring-up with that host work. Residency then amortizes the result across every request served by the worker.

The measured load-path evidence is: presharding reduces warm repeated-process E2E by 12–40% across the recorded rows; jemalloc reduced the clean LTX-2 median from 63.4 to 58.4 s and made the weight-load portion about 17% shorter; the lazy Neuron runtime initialization observed by the prewarm change is about 6.7 s. Only the first two have full E2E A/B data.

Prewarm is best-effort. A daemon thread imports the Neuron backend and touches each assigned logical core. Exceptions are ignored. If the thread fails or finishes late, the first main-thread device operation performs initialization. A full worker-start prewarm-on/off A/B is still pending.

Profile, artifact, worker, and request lifecycle

The mainline serving layer resolves one immutable model profile at process startup. It prepares compatible artifacts, loads a worker, executes a real smoke request, and opens readiness only after the output passes model-specific validation.

Serving profile and artifacts

ServingProfile contains checkpoint identity, fixed shape, dtype, parallel topology, placement choices, and TeaCache settings. Artifact identities include these compile-affecting fields. ImmutableArtifactManager publishes versioned generation directories with manifests instead of modifying a live artifact in place. A worker binds a resolved profile to an exact artifact set; shape or world-size mismatches fail before readiness.

• Profile role: The profile is resolved once at startup and becomes the request compatibility contract; requests cannot change shape, world size, placement, or cache policy.
• Publication: Artifacts are written into versioned generation directories and exposed only after their manifest is complete, avoiding in-place mutation of a live worker dependency.
• Binding: The worker records the exact artifact generation it loaded. A later artifact publication does not silently change the running process.
• Readiness gate: Profile/artifact validation, worker load, and a real smoke request must all succeed before readiness opens.

Typed stage execution

StagePipelineEngine creates runners in the declared pipeline order and rejects missing, reordered, or extra runners. ValidatedStageRunner checks exact input and output payload types at every boundary. Stages execute sequentially and report timings through a shared request context. Shutdown runs in reverse stage order and aggregates cleanup failures.

• Construction check: The declared profile stage list and constructed runner list must match exactly before the worker accepts traffic.
• Payload check: Every runner validates its concrete input and output type, preventing a text, latent, image, or video payload from crossing the wrong stage boundary.
• Timing: A shared request context records stage timing without changing the payload contract.
• Cleanup: Shutdown walks runners in reverse construction order and reports all cleanup failures instead of stopping at the first one.

Mainline resident stage definitions

Admission, timeout, and worker recovery

The API parent maintains bounded admission using max_running_requests + max_queued_requests. A run lock serializes access to the resident worker. Queue wait and execution share an absolute request deadline: a full queue returns 429, queue timeout returns 429, and deadline expiration returns 504. Cancellation or timeout starts a recovery fence. The engine first requests clean cancellation; if that does not reach a terminal state, it terminates and recreates the worker before accepting new work. Health, readiness, recovery state, pending/running/queued counts, worker PID, last heartbeat, stage, and request ID are exposed through runtime status.

• Capacity: Running plus queued requests is bounded before work reaches the resident worker. A full queue is rejected with HTTP 429.
• Deadline: Queue wait and execution consume one absolute deadline, so a request cannot receive a fresh timeout budget after leaving the queue.
• Recovery fence: Cancellation or timeout prevents new work from entering until the current worker reaches a clean terminal state or is replaced.
• Observability: Status includes readiness/recovery, queue counts, worker PID, heartbeat, current stage, and request identifier.

Image/video API and output storage

Image requests use the OpenAI-style Chat Completions response. Without object storage, PNG bytes are returned as a Base64 data URL; with S3-compatible storage, the server writes a private object and returns a time-limited presigned URL. Video generation has synchronous and asynchronous routes, an in-memory job repository, file-backed media, list/retrieve/content/delete operations, retention cleanup, and a lease preventing two services from owning the same artifact root. Queued and terminal jobs can be deleted; a running job returns a conflict instead of interrupting the worker.

Optional API-key middleware checks Bearer authentication before parsing a /v1 request or reserving admission capacity. /health and /ready remain unauthenticated. This is one shared key, not tenant authorization. Request validation requires model, shape, and fixed-profile fields to match the worker before generation.

• Image result: Return an inline Base64 PNG when storage is absent, or upload a private object and return a time-limited URL when S3-compatible storage is configured.
• Video result: Synchronous and asynchronous routes share file-backed media, job retrieval, content download, deletion, and retention cleanup.
• Ownership: A filesystem lease prevents two video services from managing the same artifact root.
• Authentication: One optional Bearer key protects /v1 before parsing and admission. It is a shared service key, not tenant-level authorization.

Measured request latency

The following measurements begin after readiness. The controlled image comparison uses identical 20-step profiles and complete HTTP wall time.

Measured Trn2 resident request latency

Profiles differ in resolution, frame count, VAE placement, and output type; these values describe deployed request latency, not cross-model efficiency. Prewarm happens before readiness and is excluded.

Controlled resident HTTP vs repeated-process CLI — 20-step images

Both serving measurements include request validation, inference, PNG encoding, R2 upload, and URL construction. The first Qwen request was excluded from the stable mean because the parent process lazily loaded a request-validator tokenizer, adding about 5.29 s. TeaCache was disabled. FLUX outputs were byte-identical between CLI and serving; Qwen differed by about 0.499/255 mean absolute RGB-channel value because the two paths used different VAE topologies.

Video serving and placement

Resident video qualification on Trn2

Placement is part of immutable worker identity. A request cannot move a VAE between host and Neuron or trigger compilation. Mainline promotes the accepted Wan and Hunyuan Neuron-VAE profiles while keeping --host-vae as a startup-time rollback. LTX-2 has no Neuron VAE binding.

Numerical and output checks

Kernel and distributed-path parity

Optimized kernels are compared with the previous dense or replicated computation before their latency is used. Wan head-sharded attention records cosine 0.999768 against replicated attention; LTX-2 unmasked cross-attention records full-generation cosine 0.999934 against masked SDPA. Ring-attention and SP changes have separate device comparisons because collective ordering can preserve shape while changing values.

Load-path parity

Presharding is treated as a serialization change, not a numerical optimization. FLUX outputs are bit-identical across presharded versus shard-on-load and cold versus warm combinations. Missing rank shards intentionally select the original shard-on-load path so an incomplete cache changes startup work rather than model values.

DiT-cache quality

Cache speed is reported together with the final trajectory check: Wan cadence-2 records 1.77× denoise speedup at final-latent cosine 0.9997; HunyuanVideo records 1.91× end-to-end against an uncached quality reference. These opt-in measurements are not folded into the uncached device table.

Decoded output validation

The benchmark checks output shape, finite values, and expected range. Serving goes further by decoding emitted PNG or MP4 data, checking repeated-request behavior, and validating the configured storage response. Stage-by-stage diffusers comparison is used when a valid DiT tensor still produces visibly wrong media.

Quality and latency are reported as separate gates. Step caching is optional and excluded from the main performance tables; attention and sharding corrections remain enabled because their numerical checks meet the implementation’s parity gates.

Commands and recorded environment

# Trainium: compile + generate using the pinned model matrix
python -m benchmark.bench --model hunyuan_video
# Real denoise-loop per-step timing
python -m benchmark.step_realloop --model flux_1_dev
# CUDA reference on H100 or B300
DIFFLET_BENCH_DEVICE=h100 python -m benchmark.bench --backend cuda --model flux_1_dev --iters 0
# Warm repeated-process CUDA run
python -m benchmark.warm_e2e --backend cuda --model flux_1_dev

The checked-in JSON reports preserve raw samples, summary statistics, configuration identity, output validation, toolchain versions, and notes. The source revisions used for the report are stated in the opening scope section.

Toolchain and benchmark identity