Add disaggregated inference (P/D) on Neuron sample#14
Closed
yahavb wants to merge 94 commits into
Closed
Conversation
This sample demonstrates deploying Rolling Forcing (real-time video generation) on Amazon EKS using AWS Trainium2 instances with: - Dynamic Resource Allocation (DRA) and Neuron DRA driver - ResourceClaimTemplates for NeuronCore slicing (small/medium sizes) - Capacity Reservation nodegroup for trn2.48xlarge - Tensor-parallel inference (TP=4) backend with FastAPI - Gradio frontend for interactive text-to-video generation Includes: - cluster/: EKS cluster and nodegroup definitions - dra/: ResourceClaimTemplates for Neuron device allocation - app/: Application code, deployment manifests, and configs - README.md: End-to-end setup instructions
…ne call - inference_neuron_tp.py: added run_benchmark() and --benchmark flag Rank 0 calls streaming pipeline directly, measures per-block FPS, prints results, exits - rolling-forcing-job.yaml: simplified to single container running torchrun --benchmark No HTTP server, no sidecar, no endpoints — just load model, run inference, report FPS
…Neuron - Added kernels/vae_conv2d.py and kernels/vae_attention.py - AttentionBlock in vae.py dispatches to NKI when on Neuron device - run_vae_decode.py uses eager mode (no torch.compile) with NKI kernels - Same VAE architecture as Wan2.2, just different channel dims (512 vs 1024)
…o not in torch_neuron_eager env Conv2d + SDPA work natively in Neuron eager mode. NKI VAE kernels only work in wan2-ti2v-5b repo which uses neuronxcc.nki directly.
Was s-lnc2-trn2 (1 ND, 4 NCs) with nproc=4 and TP_DEGREE=4. Now m-trn2 (2 NDs, 8 NCs) with nproc=8 and TP_DEGREE=8. Matches wan2-ti2v-5b-job.yaml resource allocation.
1.3B model has 12 heads. 12/8=1.5 (not divisible) so TP=8 is impossible. 12/4=3 heads/rank → TP=4 is correct. Keep m-trn2 (2 NDs) for T5+VAE headroom.
- Add models/vae_tp.py: Megatron-style col→row TP for CausalConv3d (ColumnParallelCausalConv3d, RowParallelCausalConv3d with all-reduce) - Update inference_neuron_tp.py: VAE_TP_DEGREE env var, multi-rank VAE loading, shard decoder before moving to Neuron, worker VAE decode for TP all-reduce - Add VAE_TP_DEGREE=2 to rolling-forcing-job.yaml - Backward compatible: VAE_TP_DEGREE=1 (default) uses original single-rank path
…lBlock) Wan 2.1 Decoder3d.upsamples is a flat nn.Sequential of ResidualBlock + Resample, not nested Up_ResidualBlock objects like Wan 2.2. Iterate flat list directly.
Benchmark results: TP=1: steady-state 4.50 fps (TTFF 1194s due to full-size compilation) TP=2: steady-state 3.91 fps (TTFF 412s due to half-size compilation) TP=2 is 13% WORSE at steady-state because all-reduce after every RowParallelCausalConv3d costs more than halving 128/256/512 channels. Keep TP=1 for production. TP=2 only useful for faster compilation.
wrap_nki() on the container SDK expects raw functions, not already @nki.jit decorated ones. Error was 'must be called with a dataclass type or instance'. Now import @nki.jit functions directly — they're already callable as NKI kernels.
@nki.jit kernels import fine but crash at runtime with 'No module named torch_neuronx.pyhlo' on containers without pyhlo. Add pyhlo availability check at import time — if missing, disable NKI and fall back to PyTorch SDPA.
Port NKI kernel integration from vae2_2.py to Wan 2.1 vae.py: - pyhlo check at import time (same pattern as layers.py) - Direct @nki.jit import (no wrap_nki) - _nki_conv2d_forward helper for 1x1 and 3x3 conv2d - AttentionBlock.forward() uses NKI path when available - Falls back to PyTorch SDPA when not on Neuron or NKI unavailable ALL 6 NKI kernels now loaded: DiT: cross_attention, rope, self_attention (layers.py) VAE: conv2d_k1, conv2d_k3, self_attention (vae.py)
DiT kernels (layers.py): cross_attention, rope, self_attention VAE kernels (vae.py): conv2d_k1, conv2d_k3, self_attention No more pyhlo gate. @nki.jit kernels import as decorated functions and compile on first call via neuronxcc.
Direct @nki.jit calls fail with 'No module named torch_neuronx.pyhlo' because neuronxcc.nki._torch_xla requires pyhlo at compile time. wrap_nki() from torch_neuronx.nki_hop uses a different code path (HOP) that works on both concourse images. DiT: cross_attention, rope, self_attention — via wrap_nki() VAE: conv2d_k1, conv2d_k3, self_attention — via wrap_nki()
Both rolling-forcing and stream-diffusion now output identical format: - Compilation time (first block) - T5 encode time - DiT/block time (avg excluding first) - VAE decode time (avg per block) - OVERALL FPS, STEADY-STATE FPS, VAE decode FPS - Real-time ratio (vs 16fps) - Steady real-time ratio
… format Now both pipelines report identical metrics: - DiT/block time (5 steps) — measured as time between generator yields - VAE/batch time (N frames) — time to decode yielded latent block - Batch E2E time (N×DiT+VAE) — total per-block - STREAMING FPS = frames_per_block / batch_e2e - Stream real-time ratio Previously dit_ms was always 0 (baked into yield). Now correctly measured as the gap between yields = time generator spent computing DiT inside.
Collects decoded frames during streaming benchmark loop and writes them to mp4 at the end. No extra inference run needed.
… conflict) - Replace imageio.get_writer(mp4) with PIL Image.save(PNG) per frame - Frames saved to /var/mdl/rolling_forcing/frames/frame_XXXX.png - Job YAML stitches PNGs into mp4 via ffmpeg after torchrun exits - Enables quality inspection of individual frames on PVC
Key findings: - 351/437 NEFFs are small (10-100KB) individually compiled sub-modules - NeuronCore active only 0.417s out of 1535s total - 425,789 kernel launches creating massive scheduling overhead - Compute MFU is 34% when active - the problem is Python overhead - Next: aggressive whole-block fusion to reduce to ~67 NEFFs
…ngle NEFFs Profile showed 99.97% overhead (0.03% MFU): - 437 NEFFs, 425K kernel launches, NeuronCore active only 0.4s/1535s - 351 small NEFFs (10-100KB) from individually compiled sub-modules Fix: torch.compile(block, backend='neuron') for each of 30 blocks. Fuses: norm + modulation + self_attn(NKI) + cross_attn(NKI) + FFN + residuals into ONE NEFF per block per input shape. Expected: 437 NEFFs → ~67 NEFFs, 7-8x fewer kernel launches, ~7x speedup
…ks as single NEFFs" This reverts commit 6c98877.
… tensors) torch.compile(block, backend='neuron') fails at execution time: NKI kernels receive non-contiguous tensors from .permute() inside compiled graph. .contiguous() appears to be optimized away by the Neuron backend. Documents: profiling data, attempted fix, error, potential solutions. See docs/ISSUE_WHOLE_BLOCK_FUSION.md
…-block fusion fix)
…mpatibility The new torch_neuronx wrap_nki() requires @nki.jit to return a CompileKernel dataclass (from nki.framework.compiled). The old neuronxcc.nki.jit returned a plain callable which fails with 'must be called with a dataclass type or instance' error. Change: import neuronxcc.nki → import nki (standalone package)
…pilation, NKI nl.multiply/nl.add API, /tmp local disk output)
…ranch - Copy wan/ module package from aws-reactor-rolling-forcing (provides tokenizers, T5, VAE modules needed by inference_neuron_tp.py) - Update deploy/rolling-forcing-job.yaml to clone from aws-neuron/aws-neuron-eks-samples rolling-forcing branch instead of reactor-team/aws-reactor-rolling-forcing
The code incorrectly assumed 2 NeuronDevices with rank//2 mapping. With logical-neuroncore-config=2 on a single ND, all 4 ranks are on ND0. Device mapping is handled by the Neuron runtime based on DRA allocation, not by the application code. Removed all ND references and just log rank.
Both main forward and cache-update now pass [1, 15, C, H, W] tensors to the model (same shape). Cache-update fills preceding slots with real context from noisy_cache instead of zeros. This should eliminate the e shape mismatch (3 vs 15) recompilation without breaking quality. Using default cache_size_limit=8 (no compiler crash risk).
Per SDK team guidance: Python ints that change per iteration cause dynamo guards and recompilation. Promote current_start_frame to tensor outside the compiled block and pass it in. The Python int current_start stays for cache slicing (already causes graph breaks). This removes the 'current_start == 2574' guard from the compiled graph — one fewer recompilation trigger per block.
Kernel now handles non-multiple-of-128 seq_len internally (computes num_tiles = ceil(seq_len/P)). Eliminates F.pad(x) from PyTorch per RoPE call. cos_sin still padded (kernel loads full tiles from it). Removes 1 pad op per RoPE call × 4 calls/block × 30 blocks = 120 fewer PyTorch ops in the compiled graph.
Custom GELU with float32 upcast + manual tanh math was causing torch.compile to break the FFN graph (same issue SDE team found). F.gelu is a single op the compiler can fuse into the surrounding Linear layers without a graph break.
The kernel cannot load past tensor bounds (NKI validates statically). Revert to padding x in PyTorch before passing to kernel. F.gelu change (3e0837f) is still in place.
Move all changing Python int operations (current_start, cache_start, num_valid_frames, kv_cache dict reads) to after the tensor compute in Phase 1. This way dynamo traces Phase 1 without encountering any values that trigger guards — only static class attributes and tensors. Phase 1 now contains: QKV projection + RoPE (pure tensor ops). Phase 2+ (cache management) uses the Python ints but those already cause graph breaks from dict access and conditionals. Expected: dynamo compiles Phase 1 once (no recompilation), reducing unique guard combinations from 112 to ~10.
…kens) The 'current tokens' copy was using dynamic valid_tokens as slice size, causing different output shapes per iteration (1170 vs 1560 per Andrew's trace). Use static self.block_length instead — always copies the full block. Padding beyond valid_tokens is masked out by the attention mask. Also deferred all dynamic int usage past Phase 1 tensor compute.
Port from aws-reactor-rolling-forcing branch neuron-science-team-0528. Replace alpha APIs with compatibility stubs: - nkilib (kernel_assert, PSUM_BANK_SIZE, div_ceil, ModularAllocator, TensorView) → nkilib_compat.py - nki_op (mutating kernels) → nki_op_compat.py (wraps with wrap_nki) - torch_neuronx import wrap_nki → torch_neuronx.nki_hop import wrap_nki Remaining hard ports (rope.py reshape_dim/expand_dim/broadcast): plan to use our existing working RoPE kernel instead. All files parse. Not yet integrated or tested.
…r SDK - Replace rope.py with our proven kernel + PyTorch grid-building - Fix dit_attention.py to pad inputs for NKI tile boundary - Set cache_size_limit to default (8) instead of 128 - All files parse, ready for integration test To run: cd app/science_team && torchrun --nproc_per_node=4 generate_latents.py --tp_degree 4
Runs e2e_pipeline.py with TP=4, SP=1 on 4 NeuronCores. Clones neuron-science-port branch.
wrap_nki expects nki_kernel.func attribute which only exists on @nki.jit decorated functions.
_compile is no-op during model construction so state_dict keys match checkpoint (no _orig_mod prefix). Enable after loading. First run will be eager (no compilation) to verify correctness.
Replace whole-block torch.compile with per-sub-module compilation: - Q, K, V, O linears compiled individually (fullgraph=True) - RMSNorm (norm_q, norm_k) compiled individually - FFN compiled as a unit - NKI kernels (attention, RoPE, cache) run in eager between ops This avoids dynamo encountering Python ints inside the compiled scope — no recompilation guards, no cache_size_limit issues. Each compiled unit has static shapes and no control flow. Job manifest updated to use our existing app code (not science_team/).
…ks quality TPRMSNorm contains dist.all_reduce for global RMS computation. Compiling it with fullgraph=True likely breaks the collective op, producing incorrect normalization → blurry output. Leave norms in eager, compile only Linears + FFN.
RowParallelLinear (O projections, FFN fc2) contains all_reduce_sum. FFN as a whole also contains RowParallel fc2. Compiling these with fullgraph=True breaks TP collective ops → blurry output. Only compile Q/K/V (ColumnParallel, no communication) + embeddings.
fullgraph=True might cause numerical issues on this SDK version. Test with default graph mode (allows graph breaks) to isolate whether fullgraph is the quality problem.
Prefill/decode disaggregated serving with vLLM on Trn2/Trn3: - DRA-based Neuron + EFA device allocation (co-located devicegroup8 claim) - NIXL/LIBFABRIC KV transfer over EFA - 1P1D skeleton with dynamic xPyD kubectl scale - routing via vLLM production-stack and AIBrix (pd algorithm) - tutorial README + cluster/nodegroup/claim/deploy manifests
Fixes two open Dependabot alerts in the Llama 3.1 finetune sample: - GHSA-mw35-8rx3-xf9r (high): RCE via Parquet Arrow extension-type deserialization (ray >=2.49.0,<2.55.0) - GHSA-q5fh-2hc8-f6rq (moderate): unauthenticated dashboard DELETE DoS (ray <2.54.0) 2.55.0 covers both.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| yield f"data: {json.dumps({'error': str(e)})}\n\n" | ||
|
|
||
| import json | ||
| return StreamingResponse(generate_frames(), media_type="text/event-stream") |
Contributor
Author
|
Superseded by #15, which is cut cleanly from master (this branch was based on a feature branch and pulled in unrelated files). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a new EKS sample —
disagg_serve_vllm/— demonstrating disaggregated inference (prefill/decode separation) with vLLM on AWS Neuron (Trn2/Trn3).Contents
README.md) — announces DI on Neuron, walks through EKS setup, ends with request examples against both routers.xl-lnc2-trn2-efa-rct.yamlpairs Neuron + EFA devices in onedevicegroup8-constrained ResourceClaimTemplate so the NIXL/LIBFABRIC path is valid; references the Neuron and EFA DRA driver docs.trn2-48xl-efa-ng.yaml(efaEnabled,privateNetworking, capacity reservation) +cluster.yaml.prefill-deploy.yaml(kv_producer) /decode-deploy.yaml(kv_consumer), withkubectl scaleexamples for dynamic xPyD (prefill- or decode-dominant).production-stack-router-deploy.yaml(vLLM production-stack) andaibrix-router-deploy.yaml(AIBrixpdalgorithm, shipped in v0.7.0).Root
README.mdinference table updated with the new sample (links to vLLM production-stack and AIBrix).Notes