Poster Presentations
Posters will be displayed in the Community Expo during the entire event. Presentations will take place during the Flare Party, Tuesday, October 21 6:00-7:30 PM PDT.
Deploying agentic applications using frontier LLMs incurs escalating hardware and token costs. While open sub-50B models (e.g., Gemma, Qwen) approach cloud parity on tasks, developers lack hardware-aware frameworks to evaluate quantized small model viability when swapped into production pipelines.
We present "Are We There Yet?", an open-source evaluation suite for assessing small model viability across edge environments. Using an OpenAI-compatible interception layer, AWTY evaluates target models against cloud baselines utilizing PyTorch backends (vLLM, ExecuTorch, Safetensors) as well as llama.cpp, and MLX.
Instead of evaluating model benchmarks alone, AWTY factors in telemetry: TTFT, TPOT, VRAM peak, and memory bandwidth across quantization schemes (FP16 to INT4). For agentic workflows, AWTY considers multi-turn trajectories to assess JSON-schema compliance, tool-calling reliability, and reasoning token efficiency (toks/task) to compute a “Readiness Score”.
Poster attendees will review system architecture, inspect cloud vs. local cost matrices on consumer hardware (RTX, Apple Silicon), and walk through real-world case studies. Open-source release scheduled for Sept 2026.
The PyTorch Certified Associate (PTCA) certification was announced by the PyTorch Foundation in 2Q, 2026. Whereas the PyTorch Ambassador Program, which brings together rising community leaders from around the world to support them in activating, connecting, and growing interest in PyTorch at scale, was launched in 2025 by the PyTorch Foundation. They are both new to the PyTorch Community, and this session aims at raising awareness of these initiatives from the PyTorch Foundation by providing a deep-dive poster in the PTCA and PyTorch Ambassador paths.
During the poster presentation, as one of the first PyTorch Ambassadors and the lead content creator and instructor for the PTCA training, I will share insights on how to get started on these paths to help the PyTorch community and grow as an individual in the PyTorch project. Whether you are interested in the certification or someone looking to expand your role in the PyTorch community, this poster session offers actionable insights to help you take the next step to becoming a PyTorch Certified Associate or PyTorch Ambassador.
Helion, a Python-embedded DSL that compiles to Triton and other GPU kernel backends, depends on autotuning for high-performance kernels—but it dominates the edit-run-measure loop. Full random search takes 10–20 minutes per compile; the fast default finishes in seconds yet yields kernels 14–31% slower. The current cache worsens this by tying results to exact source code, so renaming a variable or editing a comment invalidates configs. Each kernel's best config is also saved only on the machine that tuned it, so teammates can't reuse it.
We present FROMBESTAVAILABLE, a warm-start strategy now default for Helion's quick autotuning. It seeds the search with prior configs matched by input characteristics and device, using a match key decoupled from source identity. Integrated with RemoteCacheBackend, it merges and deduplicates local and remote entries by flat configuration, so a winning config from one machine becomes a warm-start candidate team-wide. In one case, autotuning dropped from 1238 to 92 seconds while matching full-search quality.
You'll leave with a reference repository and the knowledge to implement shared caching, saving your team hundreds of hours of compile time.
Encoder–decoder Transformer models such as Whisper combine dense projections, scaled dot-product attention, and nonlinear activations that present significant performance challenges for CPU inference. This poster presents an optimized inference path for Whisper on Arm Neoverse processors through cross-stack software optimizations in the PyTorch ecosystem. The optimized inference path combines runtime, compiler, and kernel optimizations, including enabling Whisper in vLLM's CPU backend, introducing low-precision INT8 and INT4 inference, optimizing scaled dot-product attention in PyTorch, enhancing BF16 matrix multiplication across OpenBLAS, oneDNN, and KleidiAI, and implementing an Arm-optimized BF16 GELU lookup-table kernel. Together, these optimizations accelerate key components of the inference pipeline. On AWS Graviton4 processors, the optimized INT4 implementation delivers up to 67% higher throughput than the BF16 baseline while maintaining 99.83% of the baseline accuracy.
Credited Co-Author(s): Fadi Arafeh
Efficient on-device inference increasingly depends on hardware-specific kernels, but maintaining optimized implementations across heterogeneous devices is unwieldy. At Hugging Face, we're developing an open-source framework that turns the Web into a decentralized environment for cross-platform agentic kernel optimization.
Native Safetensors support and a PyTorch-facing API simplify integration with existing ML workflows. Model architectures align with Hugging Face Transformers and Transformers.js, while first-class support for heavily quantized models and operators enables efficient low-bit inference.
We evaluate the framework across a wide variety of model architectures and sizes. For Gemma 4 E2B, agent-driven optimization increased decode throughput from 84 to 255 tokens per second on an Apple M4 Max. On the same hardware, LFM2.5 230M exceeded 1,400 tokens per second, while Bonsai 27B (a 1-bit language model) ran at up to 90 tokens per second.
Most multi-agent systems layer message brokers, state stores, and schedulers on top of PyTorch, adding infrastructure, serialization overhead, and operational complexity beyond what tightly coupled agent workloads require. This poster introduces a native PyTorch design pattern that treats AI agents as first-class distributed processes. Agents coordinate through tensor-native communication, using shared-memory tensors where co-located, Process Groups for membership, and RPC for point-to-point interactions. Embeddings, confidence scores, and intermediate activations remain GPU-resident, avoiding unnecessary serialization and CPU transfers. Coordination pattern and reference implementation are built from existing distributed APIs. We demonstrate the approach with a code-repair system where a Planner proposes patches, an Executor validates them, and a Verifier accepts, iterates, or escalates results. The pattern targets tightly coupled, GPU-intensive reasoning loops rather than long-running workflow automation. Attendees will leave with a practical design pattern for building multi-agent systems using native PyTorch as a distributed runtime instead of middleware-based orchestration.
On-device deployment of GenAI/ML models demands quantization for latency, memory efficiency, and hardware (CPU/GPU/NPU) compatibility. However, balancing quality-performance trade-offs in PyTorch remains complex.
AI Edge Quantizer (AEQ, presented at PyTorch Conf '24/25) is a post-training tool that takes PyTorch models to production across Android, iOS, Web, and IoT, ensuring quantized models execute via the LiteRT runtime, enabling top Android apps to optimize PyTorch models for deployment.
This poster explores recent AEQ updates:
* CLI, simplifying quantization across shell, CI/CD, and AI agents interaction
* AI agents that automate layer-by-layer error profiling and selective quantization (node denylists) without manual tuning
* Better support for on-device LLM formats (.litertlm) combined with optimizations to reduce memory overhead
* Advanced LLM Quantization Schemes: Blockwise weight scaling (int4/int8), Hadamard rotations to smooth activation outliers, enabling NPU-native per-channel weight and activation quantization
* NPU LLM Calibration: New on-the-fly fast and memory efficient calibration method plus direct calibration on .litertlm without custom unpacking scripts
Credited Co-Author(s): Renjie Wu, Pedro Gonnet, Rocky Rhodes, Mai Huynh
Small and medium-sized businesses increasingly need advanced forecasting and pricing capabilities but often lack access to enterprise-scale AI platforms. This poster presents a practical framework that combines predictive machine learning, conversational AI, and cloud-native architectures to support pricing, inventory, and demand planning.
The presentation explores how predictive pricing models, demand forecasting, and reinforcement learning improve operational decision-making, while conversational AI enables natural language access to complex analytics. It also covers the engineering architecture behind these systems, including microservices, real-time data pipelines, containerized deployment, and API-first integrations for scalable production environments.
Attendees will gain practical insights into model serving, data integration, validation, and deployment strategies that help deliver reliable AI applications for supply chain operations.
We present a design for integrating CUDA Graph support into PyTorch's Ahead-of-Time Inductor (AOTI), bringing significant inference latency improvements to production without requiring a Python runtime.
CUDA Graphs eliminate per-kernel CPU launch overhead by recording GPU operations and replaying them as a single unit. While torch.compile supports CUDA Graphs via Inductor's reduce-overhead mode, this is managed by a Python runtime—making it unavailable in AOTI's C++ deployment path, the standard for production inference at scale.
Our approach makes all CUDA Graph partitioning and eligibility decisions at compile time, embedding them as metadata in the generated artifact. The C++ runtime performs graph capture and replay without any Python dependency.
Early testing demonstrates a 25% latency improvement on Meta's Ads recommendation model, validating the approach for production workloads with many small GPU kernels.
This work combines AOTI's Python-free deployment with CUDA Graph's kernel launch amortization—previously unavailable in PyTorch's ecosystem—enabling near-optimal GPU utilization for production inference.
Credited Co-Author(s): Bin Bao
The main objective of this project is to develop a cost-effective vein viewer that identifies veins and provides enhanced visualization to assist medical personnel during venipuncture. A NoIR camera coupled with high illumination in lights of 850/650 nm is used to detect veins on the human hand, interfaced with a Raspberry Pi and integrated with vein viewer software developed using Python, OpenCV, and PyTorch. Novel image processing filters are combined with a lightweight U-Net deep learning model, trained in PyTorch, to segment and enhance the veins at the pixel level. To run efficiently on the Raspberry Pi, the model is optimized and quantized to INT8 for real-time on-device inference without a GPU, keeping patient images local for privacy. Preliminary results indicate the system detects veins accurately and would assist medical personnel for improved venipuncture, solving the issues of incorrect venipunctures and reducing unnecessary pain and chances of infections. The developed system provides faster vein detection, accurate venipuncture, and proper drug administration in a cost-effective manner.
Credited Co-Author(s): Dr. Harsh Kapadia
Training LLMs is bottlenecked by activation memory. PyTorch offers 2 ways to trade resources for memory: activation checkpointing/recompute and CPU offloading. But deciding what to keep, recompute, or offload is done manually via per-model heuristics, and the result is rarely optimal. We present an automatic joint FX-graph pass for TorchTitan's compiled training path (graph_trainer) that, given a peak-memory budget, decides per activation tensor whether to keep, recompute, or offload, trading recompute cost against offload bandwidth to fill memory as close to the budget as possible while maximizing MFU. We formulate this as an integer linear program with hard constraints on peak memory at every schedule point and on offload/prefetch overlap, via a 2-level decomposition: an outer ILP allocates a per-block memory budget, and per-block ILPs solve the per-tensor decisions. The pass is fully automatic, preserves bit-wise-identical loss, and honors user annotations. Across TorchTitan models we show memory-vs-throughput Pareto frontiers where the pass matches or beats hand-tuned policies at equal budgets, and we show gains on DeepSeek-V3 by maximizing offload bandwidth to cut recompute.
"Credited Co-Author(s): Shangdi Yu, Michael Lazos, Sanket Purandare, Sherlock Huang
"
Packed variable-length training saves padding computation but creates unequal batch costs, causing synchronous PyTorch DDP ranks to wait for the slowest rank.
We present Batch Cost Lookahead (BCL), a drop-in batch sampler that uses a bounded local lookahead window and greedy list scheduling to balance fixed-count batches without cross-rank coordination or changes to the model and runtime.
We derive a straggler model based on batch-sum concentration and an indivisible-record reducibility criterion. Experiments on production-shaped and public workloads show that local BCL approaches globally coordinated balancing, reduces end-to-end training time, and remains effective as distributed scale increases. A production Feed Sequence Model recipe improves wall-clock time while preserving evaluation parity. BCL demonstrates that local workload regularization can deliver much of the benefit of global scheduling while preserving PyTorch’s existing distributed-training interface.
Credited Co-Author(s): Chen Zhu, Antonio Alonso, Jaideep Ray, Yitong Zhou
Large-scale PyTorch training often fails in the most frustrating way: a tiny bug corrupts a few values, nothing obviously breaks, and only much later the run starts to drift. Loss curves tell you something changed. They rarely tell you where it first went wrong.
In practice, engineers often debug by comparing a suspect run against a known-good reference: the same checkpoint, the same seed, or a version before a stack change. But loss curves and gradient norms are still too coarse to pinpoint.
This talk introduces training alignment: a way to systematically compare two supposedly equivalent training runs and recover the longest prefix where they still compute the same tensors. The first mismatch becomes a precise debugging pivot. In OpGuard, we realize this with bitwise alignment at semantic operator boundaries, turning debugging from guesswork into finding the first wrong tensor.
I’ll show how this works in real production PyTorch systems with fused kernels, distributed collectives, async streams, graph capture, etc. With lightweight tensor fingerprints and schedule-tolerant trace alignment, we can localize semantic drift far more precisely than loss or gradient monitoring.
In the LLM era, community and industry mainly focus on GPU-based model optimizations. As we transition into the agent era, multi-agent architectures introduce new challenges and opportunities. A multi-agent approach inherently involves multiple models. The key challenge lies in how to jointly optimize these models from a workflow perspective to achieve optimal performance and cost efficiency.
To address this, we propose a system-level optimization solution for heterogeneous agent systems that fully utilizes both the CPU and GPU resources. We take a pipelined agent system on vLLM as an example, to show how CPUs with Advanced Matrix Extensions (AMX) capability can join the solution to improve overall agent solution efficiency by placing small LLM models on CPU, and large ones on GPU. We will show how we design the solution, how to do resource partitioning to avoid noisy neighbors and how advanced CPU technology like priority core turbo (PCT) help this solution. We also show the final solution can meet strict time-to-first-token constraints while leveraging otherwise idle CPU cores to improve TCO and support up to 1.4x more concurrent users when CPUs and GPUs run concurrently.
Triton compiles for GPUs, assuming SIMT parallelism over a flat device memory hierarchy. Dataflow AI architectures share tile-based abstractions but differ: many spatial cores, a small per-core scratchpad, and a core-to-core communication fabric instead of shared device memory.
We generate Triton kernels for dataflow hardware, keeping the PyTorch frontend and operator dispatch unchanged. At its center is KTIR (Kernel Tile IR), a new IR derived from Triton IR that abstracts the dataflow architecture. Triton cannot express core-to-core transfer, so we extend it with inter-tile communication APIs that lower to KTIR, and implement optimizations in torch.inductor.
This poster shows optimizations such as:
(1) Tiling that splits each operation into loops over tiles sized to fit the per-core scratchpad, keeping working sets on-core instead of spilling to device memory.
(2) Fusing operations with different iteration spaces, or work divisions, by inserting inter-tile communication that moves data directly core-to-core — avoiding device-memory round trips.
Visitors can see the generated kernels and codegen.
Credited Co-Author(s): Mori Ohara
What does it take to fit an accelerator with different execution constraints into vLLM without forking the serving stack? This poster presents our experience building the IBM Spyre out-of-tree plugin and examines where vLLM's extension points suffice and where device behavior requires changes beyond the execution layer.
We trace the serving path from hardware-agnostic model definitions in IBM's Foundation Model Stack through torch.compile with a Spyre backend that captures the FX graph and lowers it to the accelerator's native representation. We show how we adapted performance-critical vLLM features, including continuous batching with paged attention, chunked prefill, and prefix caching, to Spyre's execution model. We illustrate how device constraints propagate into scheduling through a plugin feature that pauses and resumes decoding sequences while retaining their on-device KV-cache and preserving serving semantics.
This poster will present architecture diagrams and scheduler timelines showing control and data flow across the plugin. They reveal where hardware assumptions surface across compilation, execution, and scheduling, and which vLLM interfaces custom accelerators need.
Reinforcement learning performance is often discussed as a GPU learner problem, but end-to-end training also depends on CPU-side simulation, rollout collection, evaluation, and coordination. This poster profiles a PyTorch multi-agent RL workflow using BenchMARL, which uses TorchRL as its backend, with MAPPO and the VMAS navigation environment. We scale from 3 to 30 to 50 agents across 64/191 parallel environments and compare CPU sampling/CPU learning, CPU sampling/GPU learning, and GPU sampling/GPU learning on an Arm-based Grace Hopper system. At small scale, learner updates dominate; at 30 and higher agents, rollout becomes a major share of iteration time, and CPU sampling with GPU learning provides the best end-to-end runtime. We extend the methodology to controlled Arm and x86 cloud CPU comparisons. The key takeaway is that PyTorch RL optimization must consider the full training pipeline, including CPU-intensive simulation, rollout collection, evaluation, and coordination, rather than optimizing learner compute alone.
Credited Co-Author(s): Na Li, Masoud Koleini
CUDA graphs deliver up to 12x decode and 2x prefill speedup in LLM inference by eliminating kernel launch overhead (on H100). But real models contain operations that cannot be captured – dynamic control flow, collective communication, CPU-GPU sync. Piecewise capture takes this further: by splitting graphs around non-capturable code, we achieve up to 1.13x prefill speedup over monolithic CUDA graph capture on models like Llama 3 and MiniMax-M2.7.
We present piecewise-cuda-graphs, a PyTorch-native library that lets users break out of CUDA graph capture for non-capturable code. Annotate any function with @no_graph and the library dynamically segments execution into graph regions interleaved with eager breaks.
Validated in SGLang with full performance parity (1.00x geometric mean) vs. SGLang's native breakable CUDA graphs on 8xH100. Integration with vLLM is underway.
Available at github.com/meta-pytorch/piecewise-cuda-graphs
At thousand-GPU scale, failures aren't edge cases — they're weekly events. Yet most PyTorch training scripts have zero resilience beyond manual checkpoint-and-restart. In this BoF, we discuss resilience patterns battle-tested on multi-thousand GPU runs: detecting straggler ranks via torch.distributed monitoring, graceful NCCL timeout handling, elastic training with torchrun's fault-tolerance primitives, and distributed checkpointing strategies (torch.distributed.checkpoint) that minimize lost compute.
We share our chaos injection methodology — deliberately failing ranks, introducing asymmetric network latency, simulating storage slowdowns — and discuss which PyTorch distributed configurations survive versus silently hang. Bring your own war stories — we want to crowdsource the community's resilience knowledge into patterns everyone can use.
Modern PyTorch models often include operators that do not map cleanly to portable lower-level dialects used by edge backends. This session addresses that challenge through an ExecuTorch Arm backend case study: enabling torch.nn.functional.gridsample by preserving aten.gridsampler2d through export and partitioning, then lowering it to tosa.CUSTOM for the Arm VGF backend. Rather than decomposing gridsample into a large and fragile standard TOSA subgraph, the implementation treats it as a backend-owned custom primitive with explicit metadata for the target operation, shader interface, descriptor bindings, workgroup configuration, and execution attributes. The key takeaway is a practical pattern for backend authors: how to decide when to decompose an unsupported PyTorch operator and when to preserve it as a custom backend primitive, while keeping the lowering path debuggable, testable, and extensible through focused rewrite tests, payload validation, partitioning checks, and integration coverage.
AWS Trainium is a purpose-built accelerator for large-scale model inference, but its previous vLLM integration relied on a patched fork and separate modeling framework – every upstream release meant manual re-integration, and users waited for new features. We present the vLLM Neuron plugin, a fully out-of-tree backend that brings Trainium into vLLM through its platform interface with no forked code. The plugin uses a torch.compile Dynamo backend that compiles standard PyTorch modules per rank, replacing whole-model tracing and enabling heterogeneous programs across ranks. NKI custom kernels compose inside compiled graphs, each with a CPU fallback path for development without hardware. Because execution stays fully PyTorch and vLLM compliant, features like automatic prefix caching, EAGLE3 speculative decoding, multimodal pipelines, and disaggregated inference work unmodified on Trainium. We also contributed Trainium support in NIXL's KV-transfer protocol and a push-mode KV connector to vLLM that benefits all hardware backends.The result shows vLLM's plugin model scales to complex hardware backends while preserving upstream compatibility.
Credited Co-Author(s): Finn Thompson, Aaron Dou
Deploying small language models on highly resource-constrained edge devices requires balancing memory, compute efficiency, and generation quality. We present an end-to-end ExecuTorch deployment of HuggingFaceTB/SmolLM2-135M for Arm Ethos-U85 NPUs, targeting embedded platforms such as Alif silicon.
Using post-training quantization in Executorch, we selectively quantize compute-intensive linear layers to 8-bit weights and 16-bit activations without retraining (W8A16). This backend-aware approach reduces the model from 623 MB to 228 MB while limiting WikiText perplexity degradation at context length 64 from 42.5 to 45.0–45.6. By contrast, whole-model W8A8 quantization produces unusable text, while linear-only W8A8 approximately doubles perplexity.
The poster presents the complete workflow from export and calibration through Ethos-U lowering, execution, perplexity evaluation, and qualitative generation analysis. The results show that selective W8A16 quantization offers a practical path for running SLMs within the memory and precision constraints of embedded Arm systems.
RL post-training has become a critical stage in modern LLM development, but deploying an end-to-end pipeline requires much more than running individual kernels efficiently. Systems must coordinate distributed training, rollout generation, and continuous weight synchronization across multiple software stacks while maintaining correctness and performance.
This poster presents our work enabling the open-source vime RL post-training framework on AMD Instinct GPUs using ROCm. vime pairs Megatron with vLLM for high throughput rollout generation, forming a tightly coupled loop where model weights are continuously synchronized between the two systems.
We focus on the engineering challenges of this bring-up, including correctness issues that emerge only during integrated training and rollout execution. We show how a lightweight train-to-rollout log-probability divergence signal, combined with controlled A/B experiments, turns hard-to-reproduce distributed failures into deterministic, diagnosable cases, and how holding that divergence low is a prerequisite for stable GRPO updates. Our experience shows how systematic debugging can strengthen open-source PyTorch infrastructure on AMD GPUs.
Credited Co-Author(s): Liz Li, Andy Luo
Modern foundation models have made AI agents increasingly capable, yet most deployments remain cloud-centric. Running these models directly on mobile devices offers advantages in privacy, latency, and reliability, but efficient on-device deployment remains challenging due to resource constraints and hardware diversity.
ExecuTorch, now part of the PyTorch Foundation, provides a deployment framework that enables developers to bring language and vision-language models to mobile and edge platforms. Through optimization, quantization, and EXIR-based workflows, models can be adapted for production environments while maintaining flexibility across target hardware.
This session presents practical approaches for integrating foundation models into mobile applications and building on-device AI agents with ExecuTorch. We will cover model optimization, deployment workflows, runtime integration, and performance trade-offs through real-world case studies. you will gain actionable guidance for deploying private, responsive, and production-ready AI experiences on Exynos SoCs.
Credited Co-Author(s): Hoon Choi, Jingya Zhang, Xiongzhan Linghu, Jintaek Noh, Bruce Kim
Generative recommenders reframe industrial ranking as sequential transduction over long user histories, with HSTU as the core attention primitive. Unlike softmax, HSTU uses a pointwise SiLU activation (no normalization) over a composite jagged mask, dominating inference cost.
We present an optimized HSTU kernel on AMD Composable Kernel (CK) for CDNA GPUs (MI300/MI350), whose low-level control captures matrix-core utilization Triton leaves unrealized. SiLU's lack of row reduction eliminates online-softmax state, enabling a register-resident KV loop. Key optimizations: register-resident Q/S/O with fused two-GEMM chaining (matched WarpGemm layouts, no LDS round-trip); a four-deep XOR-swizzled K/V LDS ring with transposed V loads and prefetch; sparsity/occupancy-aware scheduling (mask-range skipping, XCD balancing, L2 reuse); split-KV for long-context cross attention; and arch-specific MFMA/precision tuning for gfx942/gfx950. It supports self/cross-attention, jagged/grouped batching, fp16/bf16, and an optional flash-style softmax path. Shipped as a PT2/AOTInductor-compatible PyTorch operator, it drops into production models unchanged, delivering 1.2x-3x speedups over Triton.
Credited Co-Author(s): Hao Wu, Zhuoran Zhao
Mixture-of-Experts (MoE) inference is dominated by the all-to-all "dispatch" and "combine" that route tokens to experts. Inside a scale-up domain (NVLink or Infinity Fabric) it stays on fast intra-domain links, but wide expert parallelism spans more GPUs than one domain holds, pushing dispatch/combine onto the scale-out fabric. There, GPU-initiated libraries (DeepEP, MORI) assume a connection-oriented transport whose per-peer queue-pair state and doorbell/completion overhead grow with peer count.
We present a connectionless inference communication interface for vLLM targeting this scale-out leg: GPU threads issue dispatch/combine over a connectionless RDMA fabric, with no per-peer connection state and a constant number of doorbells per batch. It exposes a DeepEP-compatible API with a fused single-kernel, integrating into vLLM's MoE path with minimal change. We cover the design (GPU-initiated triggered operations, counting-event completion), integration on the Ethernet-based HPE Slingshot NIC, and early latency for wide-EP configs. We don't replace intra-domain links, but make the scale-out path efficient and open: a connectionless transport mapping onto Ultra Ethernet.
Credited Co-Author(s): Nathan Wichmann, Keith Underwood
CUDA Graphs are an important tool for reducing CPU overhead from kernel launches. However, they impose strict requirements on the model, including static shapes, stable data storage addresses, and additional memory overhead. These constraints make it challenging to deploy CUDA Graphs on large-scale recommender systems in production.
In this session, we will share our best practices for applying CUDA Graphs to Ads models at Meta, along with the techniques we developed along the way.
cuxray is a hardware-free static analyzer for CUDA kernel binaries. It works by reading decisions the compiler froze into the cubin, cross-referencing NVIDIA's exact architecture tables, and reports occupancy, register spills, and bank conflicts as facts with no GPU and no profiler. Because it only needs the binary, it can run on kernels you did not build, including kernels emitted by torch.compile emits and the cubins shipped inside PyTorch and vLLM wheels, across architectures. This enables kernel optimization on rented GPUs, where classic profiling isn't always available. cuxray goes beyond flagging problems and also can synthesize a verified conflict-free shared-memory swizzle and prove it correct, statically. This session walks the find, fix, verify loop on real kernels, from a sweep of real torch.compile and vLLM kernels to a batch-1 W4A8 decode GEMV that cuxray guided to 1.5 to 1.89x over Marlin, and whose operating-point boundary it predicted statically and hardware confirmed.
New LLM and VLM architectures appear on Hugging Face faster than RL systems can add model-specific wrappers, sharding rules, checkpoint converters, and rollout integrations. This creates a Day-0 gap: a model may be available for inference, yet remain weeks away from scalable RL post-training.
This poster presents a PyTorch-native architecture for narrowing that gap. We will demonstrate how DTensor-powered SPMD parallelism, reusable distributed training configurations, and standard Hugging Face APIs can be combined to bring newly released models into scalable RL workflows. Beyond the mechanics, we share practical engineering insights on common failure modes, such as numerical instability in long-running RL, and discuss diagnostic strategies for when standard training metrics fall short. We also outline proven techniques for scaling to larger models and longer sequence lengths. Finally, we analyze key design trade-offs and provide a concrete blueprint for implementing robust RL workflows while delegating the complexities of parallelism and environment orchestration to reusable building blocks.
Next-gen mobile GPUs ship increasingly powerful matrix-multiply hardware via the Vulkan Cooperative Matrix extension — a vendor-neutral Khronos standard. ExecuTorch's Vulkan delegate now offers day-1 support for this capability.
Targeting quantized linear operators (the dominant bottleneck in on-device LLM inference), our new Cooperative Matrix dispatch achieves a 3.90x kernel-level speedup on int4 GEMM and a 2.45–2.96x end-to-end latency improvement on Llama 3.2 1B/3B vs. conventional compute shaders (AMD Radeon 780M), with additional power/performance improvements on mobile GPU hardware.
This poster covers: (1) how ExecuTorch's Vulkan dispatch was restructured to accommodate cooperative matrix workgroups in quantized GEMM, (2) performance and power characterization across operators and full LLM inference, and (3) portability — built on VKKHRcooperative_matrix, the same code path accelerates any conformant GPU (Exynos, Adreno, Mali) without requiring vendor-specific branches.
Credited Co-Author(s): Pavan Kumar Lanka, Sicheng Stephen Jia
With the growth of XPU kernels, scenarios, and e2e performance demanding in PyTorch, a debuggable and profileable software stack for AI becomes critical for Intel GPUs. A top-down performance analysis is proposed to identify inefficiencies from model-level down to low-level runtime, kernel, and hardware behavior.
Workload-driven benchmarks are built covering memory management, graph execution, and inter-process communication to expose runtime overhead. Stability-oriented benchmarks are also constructed to capture runtime performance outliers.
To identify kernel inefficiency, roofline modeling is applied to do auto projection across diverse models. Combined with hardware performance counters, kernel-level inefficiencies can be root-caused.
Tools like oneAPI debugger and unitrace are demonstrated to resolve both performance and correctness issues. This work provides practical guidance to make PyTorch XPU workloads debuggable and profilable for long-term system efficiency improvement.
Credited Co-Author(s): Jinghui, Gu, Keping, Yan
Debugging numeric precision is a key hurdle in deploying PyTorch LLMs on edge devices. Diverse hardware, quantization, and float non-associativity yield nondeterminism across devices and runs, on top of that, quality drop often occurs after multiple conversation turns. Existing tools fail to capture numeric errors in long context cases.
We extend AI Edge Model Explorer with two new tools, Trace Diff and X-Lens. Trace Diff provides layer by layer comparison of intermediate tensors. By visualizing model logit distribution shift over multiple decode steps and conversation turns with cosine similarity and KL divergence, we assist human-in-the-loop debugging, reducing triaging time from weeks to under 2 hours. In two real-world examples, we demonstrate the tool's ability to isolate precision drift in RoPE operations and diagnose hardware-specific inconsistencies in denormal calculations. In addition, we developed X-Lens, building on top of Jacobian’s Lens and Logit Lens, which provides hidden space logit interpretation. In a real-world example, the tool provides an explanation of how the number of local attention and global attention layers impacts model quality layer by layer.
Credited Co-Author(s): Google AI Edge team
Computational Fluids Dynamics solvers often take hours/days per design, making broad design-space exploration expensive. AI surrogates deliver faster flow predictions, but fair and transparent benchmarking remains difficult: evaluations are often tied to custom datasets, metrics, scripts, and baselines. We present PhysicsNeMo-CFD, an open benchmarking harness for PyTorch-based CFD surrogates that standardizes dataset registration, model inference, model accuracy metrics, and reporting. The framework separates these components through reusable wrappers and registries, enabling consistent comparisons across architectures and datasets.
PhysicsNeMo-CFD provides baselines with pretrained models trained on large-scale CFD datasets like DrivAerML to evaluate and benchmark against custom PyTorch models/datasets. The evaluation suite goes beyond L2 errors to include engineering metrics like drag/lift, partial differential equation residuals, design-trend plots, UQ evaluations, and visual comparisons. This poster presents the API and workflow, showing how PyTorch practitioners and CFD engineers can plug in new models or datasets in PhysicsNeMo-CFD and generate reproducible benchmarks.
Credited Co-Author(s): Rishi Ranade, Ram Cherukuri, Sanjay Choudhry
Cold-starting an LLM inference worker can take minutes: loading weights, warming up kernels, and capturing CUDA graphs, while GPUs sit idle and SLAs slip. Autoscaling and failure recovery both pay this cost, as each naively requires a cold start. Dynamo Snapshot avoids this by checkpointing a worker and restoring it, using CRIU for host process state and cuda-checkpoint for GPU state.
The key enabler for fast restoration is GPU Memory Service (GMS), which decouples model weights from the worker's lifetime, so they outlive the process that loaded them. When a replica scales up on a fresh GPU, GMS loads weights in parallel with process-state restore over a fast path like RDMA/NVLink, then the worker imports them. For software fault recovery, the weights are still resident in GPU memory, so the worker re-imports them directly and skips the reload. (Re)starting gpt-oss-120b takes under 5 seconds.
In this talk, we'll cover the checkpoint/restore mechanics, GMS design, how to make workers and their PyTorch state snapshot-safe, and our approach in extending this to multi-GPU, as well as training workloads.
An early version of this is described in this blog post: https://tinyurl.com/nvsnp
Credited Co-Author(s): Vikram Sharma Mailthody, Dan Feigin
We present a power efficient working deployment of a masked diffusion language model on an AMD NPU. LLaDA-8B-Instruct generates text by iterative denoising instead of autoregressive decoding, so every diffusion step needs a full bidirectional forward pass. Existing NPU toolchains don't really account for this. Our pipeline quantizes LLaDA to INT4 (225 MatMulNBits operators) and partitions the graph with AMD's ryzenaionnxutils using a custom split-QKV strategy. The fix that finally unblocked us was CastAvx: replacing 450 AVX-based cast nodes with standard ONNX Cast ops, which was what kept segfaulting on XRT device memory.
At runtime we skip onnxruntime-genai's autoregressive Generator API and drop down to the og.Engine/og.Request interface to do single-step prefill on the XDNA2 NPU.
On a Ryzen AI MAX+ 395 (Strix Halo) APU, the NPU vendor path runs 504 ms per forward at seq=128 at 36 W, versus 843 ms and 110 W for CPU INT4: 1.67x faster, 8.9x better on energy. We also benchmark the iGPU (350 ms FP16) and hand-written IRON/MLIR-AIE kernels. Code: github.com/raviguptaamd/llada-npu-inference
Credited Co-Author(s): Shiksha Patel, Rishi Madduri
Hugging Face Diffusers has become the backbone of open generative media — the default way the community runs popular diffusion models for text-to-image and text-to-video inference. What it still lacks is production-grade recipes for fine-tuning these models efficiently at scale.
This talk shows how to do that natively in PyTorch. With DTensor, FSDP2, torch.compile, and tensor, context, and pipeline parallelism, distributed training becomes a configuration choice, not a code rewrite — one recipe scaling from a single 40GB GPU to multi-node clusters, with checkpoints and LoRA adapters that load straight back into Diffusers.
We'll cover what makes training efficient — flow matching, latent caching on pre-encoded VAE outputs and multiresolution bucketing— plus LoRA (rank/alpha, target modules) alongside full fine-tuning, with Wan 2.2, FLUX.2, HunyuanVideo, and Qwen-Image as concrete examples in NeMo Automodel, an open-source library built with Pytorch.
Attendees will leave knowing how to fine-tune billion-parameter diffusion models efficiently and scale them with PyTorch-native building blocks — without leaving the Diffusers ecosystem.
KV-cache transfer is becoming a core communication primitive and a major bottleneck for inference as context windows and disaggregated serving grow. Unlike collectives, this exchange needs elastic communication: endpoints may appear dynamically, exchange opaque metadata, move tensors one-sidedly between different memory tiers, and fail independently. PyTorch’s established communication interfaces are primarily rank/team-oriented: they suit collectives and structured parallelism but do not directly model endpoint discovery, asymmetric transfer, or request-local failure.
This poster presents a proposed host-side PyTorch interface for elastic endpoint transfers. It exposes explicit memory registration, opaque endpoint metadata, descriptor-based one-sided READ/WRITE, notifications, and request-local completion or failure, while leaving scheduling, placement, and cache policy to inference frameworks. We show how the interface maps directly to NVIDIA Inference Xfer Library (NIXL), whose endpoint model already supports KV-cache transfer in vLLM and SGLang, and how the same abstraction can extend to hierarchical KV cache, offload/restore, and request migration.
AutoParallel parallelizes PyTorch models by tracing a computation graph, generating operator-level sharding candidates over a device mesh, and solving an ILP that selects a globally consistent placement under memory and communication constraints. The selected placement is then materialized into an executable graph with explicit redistribution and collective operations.
AutoParallel uses an analytical cost-model to make placement decisions. To more closely match the runtime behavior, we replace the analytical cost-model in AutoParallel with a lightweight profiling pass that measures collective and operator costs directly on target hardware. These measured costs are injected into the existing ILP objective without modifying the optimization formulation.
We evaluate this approach on various Intel and NVIDIA GPUs in both scale-up and scale-out settings. Measured costs alter the selected placement to varying degrees and can improve end-to-end throughput by up to 20% on a dense LLM model.
Credited Co-Author(s): Syed Shahbaaz Ahmed
Hybrid models that go beyond conventional full attention — Nemotron 3 Ultra, Qwen3.6, and others — are rapidly becoming mainstream. At the same time, reasoning and agentic workloads are making Automatic Prefix Caching (APC) increasingly important: long shared prefixes create substantial opportunities for cache reuse, and APC can dramatically reduce Time-To-First-Token (TTFT). In this talk, we present our contributions along vLLM's journey to making APC practical for hybrid models. Supporting APC for hybrid models is not straightforward. Checkpointing their state can drive high memory usage or make the cache too sparse to hit reliably. Worse, early hybrid APC support introduced decode-path overhead, causing performance regressions when cache hit rates were low. Our work improved cache utilization while keeping memory usage under control, and with a specialized Triton kernel we removed ~20% decode overhead by streamlining CPU–GPU interactions. The resulting APC implementation for hybrid models delivers >10x TTFT improvement on long-prefix cache hits, introduces no overhead in low-hit-rate scenarios, and remains compatible with advanced serving features such as Multi-Token Prediction.
Credited Co-Author(s): Francesco Fusco
Enterprises run LLM agents over confidential corpora inside trusted execution environments (TEEs), composing hardware encryption, oblivious retrieval, and output differential privacy to protect content, access, and output. We show the stack leaves a channel open. Because a reason-and-act agent chooses its control flow from the private data it reads, the shape of its execution trace (step sequence, per-step timing, and egress size) is a function of the secret and stays visible to the host operating the TEE. We formalize this as (ε,δ)-trace-privacy, a DP notion over the variable-length trace, and mount a trace-reconstruction attack on a six-agent prior-authorization crew, using only PyTorch Profiler and CUDA event timings. From metadata alone, the host recovers a sensitive attribute at group-aware ROC-AUC 0.68, and 0.96 under a stronger model. Our defense, TraceShield, is a PyTorch runtime layer: static shapes and padded KV blocks under torch.compile, CUDA-graph-captured decode, and hooks that pad timing and egress to public ceilings, giving a proven (0,0) release that drives the attack to 0.50 at no measured utility loss, and binds each run to a hardware-attested receipt.
Tensor parallelism enables LLMs to run across devices, but requires intermediate activations to be communicated during inference. As models grow, communication overhead can become a bottleneck. A common approach reduces communication by keeping sensitive features in BF16 and quantizing the remaining features to Int4, but quantization introduces reconstruction error.
This poster presents Entropy-Reinvested Residual Correction (ERRC), a communication compression approach that asks whether entropy savings from activation compression can be reinvested to recover quantization accuracy. ERRC applies entropy coding to the Int4 activation stream and uses saved bits to transmit a compact residual correction stream for features selected based on reconstruction errors measured during calibration.
The method is evaluated on Gemma-2-2B using a two-device tensor-parallel simulation across MLP and attention projection layers with WikiText-2, LAMBADA, and CNN/DailyMail samples. Compared with the selected-BF16 + Int4 baseline, ERRC reduces communication by up to 12.4% while improving reconstruction quality, achieving up to 10.8% lower MSE while remaining within the original communication budget.
torch.compile automatically optimizes your code by tracing into Python code, extracting a computational graph, and optimizing it. However, graph tracing doesn't always succeed out of the box. Bugs, coverage gaps, and intentional design choices can all get in the user's way.
In this talk, we present three escape hatches for torch.compile: nonstricttrace, leaffunction, and custom operators. Each offers a different tradeoff between control, performance, and usability, giving users options to work around tracing limitations.
You'll learn what each escape hatch does, how they differ from one another, and when to reach for which so you can unblock yourself and get the most out of torch.compile.
Large-scale LLM inference systems, such as vLLM rely on distributed communication to scale models across nodes. However, specific communication paths make it difficult to support new collective backends without duplicating logic.
This work integrates TorchComms into vLLM as a communication abstraction for collective libraries, such as NCCL and XCCL, enabling distributed operations to be routed through a TorchComms -based communicator while reducing device- and backend-specific code. We use XCCL as the experimental vehicle for evaluating on XPU, focusing on communicator initialization, collective dispatch, stream synchronization, and compatibility with vLLM existing distributed inference workflow.
Preliminary results using llama3 70B on a single XPU node achieve performance comparable to the existing c10d implementation, reaching 99% of baseline end-to-end serving throughput while improving portability and reducing integration overhead for new accelerator backends.
Credited Co-Author(s): Julia Piotrowska
Triton kernels expose many tunable parameters—tile and block sizes, warp counts, pipeline stages, and backend metadata—whose interactions are hard to predict. Standard autotuning only searches a manually defined candidate set, while exhaustive search needs excessive benchmark runs.
We present Triton EvoTuner, which integrates genetic-algorithm search, an XGBoost cost model, and real benchmark feedback into Triton's autotuning workflow. Each Config is encoded as a chromosome of meta-parameters and execution metadata, enabling selection, crossover, mutation, and ranking directly over Triton configurations. The search space is inferred from user configs, Latin Hypercube Sampling seeds warmup, measured latencies train the cost model guiding exploration, and final selection rests on real measurements.
On the Triton CPU backend, EvoTuner delivers consistent geomean speedups over the user candidate set: 1.195x (blocked matmul), 1.691x (batched matmul), 1.688x (fused attention). GPU gains are smaller (1.024x), where defaults are near-optimal. EvoTuner shows cost-model-guided evolutionary search can move beyond hand-written candidate sets under a practical measurement budget.
Credited Co-Author(s): Wei Shen Huang, National Tsing Hua University, Taiwan
The web browser is the largest deployment surface in computing, reaching billions of users, yet PyTorch has no native path for GPU-accelerated inference in the browser. This work introduces a new WebGPU delegate, bringing GPU-accelerated inference to Chrome, Firefox, and Safari with no model conversion, installation, or server dependency. Models export via torch.export and execute via WGSL compute shaders.
On Apple M4 Pro in Chrome, ExecuTorch achieves 1964/148 tok/s prefill/decode on Llama-3.2-1B, compared to llama.cpp’s 617.7/93.3, WebLLM’s 1710/87.6, ONNX Runtime’s 310.9/47.3, outperforming all three across 128-8192 context lengths for quantized Llama-3.2-1B, Qwen2.5-0.5B, and Qwen3-0.6B. Beyond language models, the backend generalizes across image segmentation, object detection, super-resolution, vision-language understanding, voice, audio, and on-device 4-bit fine-tuning via LoRA, DiReFT, and full-parameter updates.
This work opens a new deployment surface to PyTorch’s open-source ecosystem addressing community demand for web-native ML inference. We showcase ExecuTorch WebGPU as a path for researchers and developers seeking performant browser-native inference from PyTorch.
Credited Co-Author(s): Digant Desai, Stephen Jia, Siddartha Pothapragada
Modern data-center GPUs increasingly feature non-uniform memory access behaviors to their last level caches and DRAM.
The latest CUDA releases (at the time of PyTorch Conference North America 2026) feature extended support for "CUDA locality domains" to exploit localized memory accesses.
This poster explains what locality domains are, why they matter for NVIDIA Blackwell and upcoming GPUs, and how both compute and memory localization are integrated into PyTorch. We walk through the full stack: compute localization via green contexts, memory localization via new custom allocators, and how to tie them together with CUDA graphs and/or torch.compile so higher-level operations — attention, linear layers, activations, normalizations — become locality-domain aware.
Finally, we cover the drawbacks and pitfalls, in particular how locality domains interact with communication kernels.
Artificial Intelligence demands efficient, portable data movement as models scale across heterogeneous GPU systems. This work introduces GPU-Aware and Kernel-Initiated SHMEM (KI-SHMEM), extensions to the standard OpenSHMEM API that enable direct communication between GPUs and allows GPU kernels to initiate communication without CPU intervention. Our research extends the OpenSHMEM specification with device-agnostic GPU support and implement it in HPE Cray OpenSHMEMX. We add multiple symmetric heaps, including heaps resident in GPU memory, so processes can allocate and exchange data directly on accelerators. It uses HPE Slingshot NIC capabilities while abstracting differences between AMD and NVIDIA hardware, and it selects communication protocols based on message size and buffer location. While motivated by broad HPC and AI needs, we also target emerging inference workloads such as vLLM serving with Mixture-of-Experts models, where expert-parallel dispatch and combine require fast GPU-to-GPU movement. KI-SHMEM shows how a standards-based OpenSHMEM extension can provide a portable substrate for these workloads without replacing the core programming model.
Credited Co-Author(s): Naveen Ravi, Keith Underwood
When a large distributed PyTorch job dies, torchrun's elastic agent restarts it with no idea whether it should. Restarts on an irrecoverable fault (bad assertion, exhausted quota, throttled dataset) burn GPU-hours looping into the same wall, while failing fast on a transient NCCL blip makes the user lose hours of progress.
In this poster, we present an automated triage system. Every failure is classified from the signals a PyTorch job emits (eg, torch.distributed/NCCL errors, CUDA OOM and illegal-memory-access, process signals, user-code exceptions), and these are correlated with hardware health from gpud, DCGM, and Kubernetes node conditions. Each category maps to an action like fail-fast, restart, or replace-node-&-restart.
If the system is unable to detect any error pattern in the log tails, it calls into an LLM, which reads every rank's log tail as one correlated collection and returns a category plus a suggested regex, promoted into permanent coverage. The LLM diagnoses a novel failure once; the pattern handles it for free forever. Doomed jobs burnt 4–5 hourly restarts before the users noticed; fail-fast now stops it in under 3 min, leading to better resource utilization.
Teams increasingly paste production material into LLM interfaces — credentials, customer records, unreleased work. Once that text crosses the network boundary, three things happen that cannot be undone: it leaves the network, it is logged under someone else's retention policy, and on consumer tiers it may be used for training.
Fenceline is a self-hosted gateway that sits between an application and whichever model serves it. It is built on one ordering guarantee: deterministic redaction runs before any model inspection, so no raw credential reaches any model — not the target model, and not the firewall's own guard agents. Most guardrail tools ship raw text to their own classifier first, which is itself a disclosure.
The system has three layers: a pure-Python detector layer that makes no network call; two structured-output guard agents that receive only already-redacted text; and a deterministic orchestrator that computes the final verdict, because a component that can be argued with should not be the last word in a security system. Real values are restored in the response so answers stay usable.
Production ML systems with billions of parameters are opaque during training: researchers observe only end-of-pipeline metrics while internal state determining stability, quality, and optimization opportunities remains invisible. Existing tools cannot survive modern PyTorch transformations or impose prohibitive overhead, limiting them to debugging, not continuous observability. We present Firetap, an open-source PyTorch library (https://github.com/meta-pytorch/firetap) providing deep, longitudinal visibility into model internals with negligible overhead by instrumenting parameters, gradients, activations, and optimizer state. Contributions: (1) automatic broad-coverage observability via standard PyTorch hooks; (2) transformation-resilient instrumentation via custom operator registration surviving torch.compile and FX graph transforms, with deduplication cutting compilation overhead by over 80%; and (3) a single-pass Triton kernel computing six statistics with cache-line-spaced atomics, enabling always-on monitoring at under 5% throughput cost. Deployed across a production fleet for two years, Firetap reduced mean time-to-root-cause for training instability from weeks to under 24h.
Credited Co-Author(s): Chris McGillicuddy, Than Hedman, Neel Vadoothker
Run a GGUF model on a Mac and most of the work happens in a few hundred lines of Metal you've never read. This talk opens up one of them.
Q5_K is among the most-used quantization formats in the GGUF ecosystem, and until recently ExecuTorch's MLX backend couldn't execute it. I'll walk through the upstream contribution that added it.
We'll cover what's actually inside a Q5K super-block, an affine d/dmin block, 6-bit packed sub-block scales, and a separate array holding each weight's 5th bit, and why that layout forces real kernel design rather than a copy of Q4K. Then the two kernels every decoder needs: a mat-vec for decode, a tiled simdgroup mat-mat for prefill, and dynamic-M dispatch between them. With benchmarks: Q4K vs Q5K vs Q6_K tokens/sec and memory on Apple silicon.
You'll leave knowing how block quantization really works on-device, and how to get your own kernel merged into a PyTorch runtime.
PyTorch has become the standard platform for modern AI development, but quantum computing workflows remain fragmented across simulators, hardware backends, and programming models.
FlagQuantum bridges this gap by providing a PyTorch-native runtime for distributed differentiable quantum computing. Developers can build quantum applications using familiar PyTorch workflows while transparently scaling across distributed statevector, MPS, and tensor-network simulators.
Powered by FlagOS, FlagQuantum supports portable execution across heterogeneous accelerator architectures and provides a unified deployment path from classical simulation to real quantum processing units (QPUs). The same PyTorch program can therefore move from local simulation to multi-accelerator execution and eventually to quantum hardware without changing its programming model.
This poster introduces the architecture behind FlagQuantum, including its distributed runtime, automatic differentiation, heterogeneous execution, and QPU deployment pipeline, and discusses how these components bring quantum computing closer to the broader PyTorch ecosystem.
Open-weight PyTorch models are "released" but rarely run out-of-the-box on non-NVIDIA accelerators — the friction is in the model × chip × software-stack integration, not the model itself. FlagRelease is a self-verifying agent platform that automates it: migrating, tuning, evaluating, and releasing models across 10+ heterogeneous accelerators. It has shipped nearly 100 models, including the DeepSeek, Qwen, GLM, and Kimi series and multimodal and embodied models like Qwen3-VL and RoboBrain.
An autonomous LLM agent makes high-level decisions, while long-running, error-prone work runs through deterministic tools. The hardest problem was reliability — the most dangerous failure was never a crash, but an agent that believed it had succeeded. We enforce two invariants it cannot bypass: a "pass" is written only by a separate tool reading the real result files, so the agent cannot certify its own success; and the pass criteria live where the agent cannot edit them.
FlagRelease builds on the open-source FlagOS stack, using vllm-plugin-FL for multi-chip vLLM serving and FlagGems, a Triton-based operator library. This poster shares the architecture and lessons from shipping at scale.
Long-context LLM inference is bottlenecked by a KV cache growing linearly with sequence length, yet attention is used unevenly. Prior KV-eviction methods (StreamingLLM, H2O, SnapKV) pair a bespoke importance heuristic with a bespoke kernel, so their scoring policies are hard to compare. FlexBudget is a training-free working-set attention built on PyTorch core flexattention: one compiled maskmod admits, per query, the top-B KV blocks under a fixed budget plus anchors, while the importance signal is an interchangeable scoremod. With budget and kernel fixed, the signal becomes a controlled variable, comparing recency, attention-mass, and query-key affinity on one path. On a single data-center GPU: bit-exact to dense at full budget; ~2.3x faster block-sparse at 25% budget; Qwen2.5-7B at 8K-32K shows 71-74% lower resident KV working set, time-to-first-token ~0.26-0.29x, and 3.2-3.7x decode throughput. Accuracy holds near dense to ~50% budget, where importance beats a fixed sink+window baseline, and arithmetic reasoning shows no degradation at 25%; plus an honest negative on mid-context recall. Uses only torch.nn.attention.flexattention with torch.compile and no custom GPU kernels.
Design space exploration for future distributed Machine Learning systems suffers from a lack of readily available workload representation / workload graph that is fed into cost models such as simulators which enable flexible exploration across the stack. We present Flint, a framework that bridges this gap by leveraging PyTorch's FX Graphs. The key contribution of Flint is to extract FX Graphs directly from source code, without having to provision expensive GPU resources. PyTorch's Dynamo frontend of the compiler does the heavy weight lifting of understanding and preserving the behavior of the original model code. Flint can collect the FX Graph of arbitrary cluster size because it interfaces with the compiler before hardware execution.
In this poster, we discuss: 1) The motivation of obtaining FX Graph from model code for design space exploration, 2) techniques to 'trick' Dynamo into thinking it is running on a real world cluster, and 3) validation and evaluation results. We validate the FX Graph against profiled tracees on a real-world execution and show the flexibility of Flint through a design space exploration case study.
Credited Co-Author(s): Meghan Cowan, NVIDIA; Zheng Du, Georgia Institute of Technology; Changhai Man, Georgia Institute of Technology; Srinivas Sridharan, NVIDIA; Tushar Krishna, Georgia Instotute of Technology
Reinforcement learning is now a key post-training stage for LLMs, but rollout dominates iteration time — ~80% in synchronous settings. FP8 is a natural lever, yet RL is harder than static inference: policy weights change every step, forcing per-step re-quantization and weight sync, and low-precision rollouts silently turn on-policy RL off-policy, degrading or collapsing training.
This poster presents FP8-RL, a practical low-precision stack for LLM RL built on the PyTorch ecosystem (vLLM, SGLang, FSDP, Megatron-LM, verl, NVIDIA NeMo-RL), covering four strategies: (1) blockwise W8A8 linear layers for rollout; (2) FP8 KV cache with per-step scale recalibration; (3) end-to-end FP8 across rollout and training; (4) end-to-end MXFP8 on Blackwell with hardware microscaling — all stabilized by token-level truncated importance sampling.
On Qwen3-8B (dense) and Qwen3-30B-A3B (MoE) with 20K-token responses, FP8-RL delivers up to 44% faster rollout at BF16-level AIME2024 accuracy. We also share what fails and why — router precision, E4M3 vs. hybrid recipes, scale formats, MoE routing replay (R3). Full technical report: alphaxiv.org/abs/2601.18150
Credited Co-Author(s): Shuang Yu, Junjie Lai
Goal: identify the cross-layer bottlenecks limiting RLHF training throughput on multi-node GPU clusters, then use that data to train an RL agent for autonomous cluster scheduling.
RLHF sits at the intersection of training, inference, and serving, making profiling uniquely hard. We instrument asynchronous RLHF pipelines built on PyTorch TorchTitan (FSDP2) for distributed training, VeRL/TRL for orchestration, and OpenEnv for containerized rollouts, on Dell PowerEdge GPU clusters. Using PyTorch Profiler, DCGM, and NCCL traces, we measure six bottlenecks: rollout throughput, reward-model serving latency, KV-cache pressure, GPU utilization, all-reduce cost, and trajectory staleness effects on PPO/GRPO convergence.
The resulting dataset yields guidance on reward-serving topology, async window sizing, and communication tuning at scale — and becomes training data for an RL agent that treats cluster state (utilization, queue depth, memory pressure, fabric congestion) as its environment and scheduling decisions (topology, batch size, worker scaling) as its action space, advancing autonomous, PyTorch-native cluster operations for RL training.
Client attrition in wealth management is a competing-risks problem: a client may leave independently or follow a financial advisor who departs for a competing firm. We present a PyTorch-based transformer model that predicts six-month attrition from longitudinal advisor–client email communications.
Email text is encoded with Qwen3-Embedding-8B, and the resulting representations are organized into time-ordered client sequences. A transformer encoder implemented in PyTorch learns changes in communication content, frequency, engagement, and relationship dynamics. Its output feeds a neural competing-risks survival layer that estimates cause-specific risks for client-initiated and advisor-led attrition while accounting for censoring.
The key contribution is a modular PyTorch architecture combining pretrained language embeddings, temporal attention, and differentiable competing-risks survival modeling. The poster will present the architecture, sequence construction, survival loss, and evaluation using cause-specific concordance. This work demonstrates how PyTorch can extend language modeling beyond classification to practical time-to-event prediction.
AI agents increasingly access files, APIs, and enterprise systems through the Model Context Protocol (MCP), making secure tool execution as important as secure model deployment. This poster presents a Zero-Trust architecture for MCP-based AI agents running on ARM edge devices, extending trust from model loading to runtime actions.
The framework combines SafeTensors for trusted model verification and ExecuTorch for on-device inference, enabling secure AI agents that operate entirely on-device. Every MCP tool request is evaluated using capability-based authorization and least-privilege policies before execution.
The architecture establishes three layers of trust: model integrity via SafeTensors verification, policy integrity via signed capability manifests, and runtime integrity via policy-enforced MCP execution. These controls help mitigate prompt injection, unauthorized tool use, privilege escalation, and data exfiltration.
The poster demonstrates how SafeTensors, ExecuTorch, and MCP can be integrated to deliver auditable, policy-driven edge AI agents, providing a practical foundation for secure deployment of autonomous AI systems on resource-constrained ARM devices.
Training end-to-end networks for autonomous driving requires decoding synchronized video from multiple cameras quickly enough that data loading does not stall the GPUs. In our PyTorch pipeline, the main costs came from repeatedly opening and demultiplexing MP4 files, decoding frames the model did not use, making many small calls across the Python/C++ boundary, and waiting for slow DataLoader workers.
We tested several ways to reduce this overhead: demultiplexing cameras in parallel, caching compressed packets after the first epoch, batching decode operations in a C++/PyTorch extension, skipping unneeded frames earlier, and batching color conversion. The poster shows where time was spent, which changes improved DataLoader and end-to-end training throughput, and the tradeoffs involved in caching intermediate data or moving more of the pipeline into native code.
Although this work comes from autonomous-driving training, the same bottlenecks appear in many PyTorch workloads that train directly from compressed video.
Credited Co-Author(s): Mihai Alden
Linear layers dominate many ML workloads, but static quantization can be brittle when activation ranges shift or representative calibration data is unavailable.
This poster presents dynamic W8A8 Linear/AddMM support in the ExecuTorch Arm backend using PyTorch PT2E and TOSA. We keep weights statically quantized to int8, compute activation quant parameters from runtime inputs, lower the core operation to integer matrix multiplication, and reconstruct float outputs after int32 accumulation.
The contribution combines backend enablement, graph decomposition, pattern detection, support for QDQ and folded quantized forms, metadata preservation, and runtime-aware bias handling. We validate on three models: a DLRM-style recommendation MLP, a Transformer/ViT preflight flow with a rank-compatible 2D Linear fallback, and a image MLP. At 16x input scale, static W8A8 image-MLP accuracy drops from 75.59% to 51.07%, while dynamic W8A8 preserves 75.59% accuracy and 0.9999 cosine similarity versus FP32. DLRM shows about 18x lower MAE than static W8A8 at 16x scale.
Attendees will learn why dynamic range quantization improves robustness and how backend authors can implement it in practice.
Sharded data-parallel training repeatedly issues parameter all-gather before layer execution and gradient reduce-scatter during backward. A framework can schedule these collectives close together, but scheduled overlap is not enough: shared process groups, communicator ordering, allocator paths, and backend state can still serialize them.
Using Megatron/NeMo-style training paths built on PyTorch Distributed as a case study, this poster follows the path from one shared data-parallel communicator to independent AG and RS communicators. The core systems insight is that parameter all-gather and gradient reduce-scatter can use complementary communication patterns, so the stack can be designed to let them progress concurrently instead of through one serialized path.
The poster is organized as a visual validation story: why serialization happens, what had to change in process-group and backend contracts, and what evidence proves real overlap. The validation ladder includes topology checks, registered-buffer requirements, PyTorch microbenchmarks, lower-level collective baselines, backend path checks, and profiler timelines.
Credited Co-Author(s): Sheng Fu
Mixture-of-Experts (MoE) models can improve training efficiency by routing each token to only a few experts, but that same routing makes CUDA Graphs difficult to use. In this session, we will explain why MoE training often introduces CPU-GPU synchronizations, dynamic expert shapes, and variable activation sizes, and how these issues add CPU overhead or force conservative padding. We will then discuss practical techniques for making MoE training more CUDA-graphable in PyTorch. This includes HybridEP for sync-free dispatch, kernel design choices that avoid host-side decisions, and paged stashing to reduce activation memory by storing actual routed activations in a shared paged buffer instead of saving worst-case buffers for every layer. We will also cover the trade-offs between padding, recomputation, activation saving, stashing, and custom kernels, since no single approach fits every model or training setup. As an example, we'll look at TorchTitan's implementation of these concepts to understand their memory and performance impact on real MoE training.
Co-Author: Nan Zheng, NVIDIA
Grouped GEMM is a foundational primitive for Mixture-of-Experts (MoE) models, where independent matrix multiplications across expert subgroups must be batched without individual kernel launch overhead. This work presents enablement and optimization of grouped GEMM on AMD ROCm in PyTorch, spanning API design, backend integration, and architecture-specific tuning.
The API torch.nn.functional.groupedmm(A, B, offs=None, outdtype=None) supports variable-length expert groups via offs, a 1D int32 tensor of cumulative group end indices. The canonical MoE pattern concatenates tokens sorted by expert into a flat 2D tensor A, multiplied against a stacked per-expert weight tensor B, with offs marking each expert's token boundary.
The backend uses AMD Composable Kernel (CK), enabled at runtime via TORCHROCMUSECKGROUPED_GEMM without recompilation, with hipBLASLt as fallback. Buffers are managed through the HIP Caching Allocator, eliminating per-call overhead critical in layer-by-layer MoE inference. A specialized equal-K kernel avoids per-group stride recomputation for uniform token routing, improving occupancy. Enablement spans MI250, MI300, and MI350.
In disaggregated LLM serving, KV-cache transfers between prefill and decode workers can become a major TTFT bottleneck. Optimizing SGLang’s NIXL path improves end-to-end TTFT by up to 80% in asymmetric configurations and outperforms the best existing staging-buffer approach by up to 12%, measured with DeepSeek-R1-Distill-Qwen-32B, ISL 1024, OSL 1024, and 2 prefill TP -> 4 decode TP.
The key observation is that asymmetric TP fragments the KV cache into many small regions, increasing processing overhead across the stack.
This poster presents two optimizations:
- A new scatter/gather API in UCX that batches KV-cache regions and reduces per-region overhead, reducing TTFT by up to ~50%.
- A hardware-accelerated strided path that identifies regular fragmented patterns, represents them compactly, and transfers them directly without staging-buffer copies, further reducing TTFT to up to ~80% overall.
The same techniques can benefit other disaggregated LLM systems, including vLLM and TensorRT-LLM, as well as reinforcement-learning workloads that move model state across different parallelism configurations.
PyTorch’s CPU inference stack uses torch.compile with TorchInductor to combine compiler code generation and optimized kernel backends across CPU architectures which is still the default for vLLM CPU backend. We focus on BF16, INT8, and INT4 inference, covering shape-aware GEMM selection, max-autotune, CPU-specific oneDNN integration, and future directions for oneDNN in PyTorch.
oneDNN provides standardized primitive APIs that allow optimized kernels to integrate consistently across CPU backends. We examine two challenges for dynamic models: weight prepacking and JIT kernel selection as tensor shapes change. We discuss current solutions and remaining gaps, including cases where improved kernel selection delivers up to 3x performance gains on selected Llama and GPT-OSS workloads.
Finally, we show how vLLM can reuse PyTorch’s ISA-optimized Vectorized class. Reusing vectorized paths for operations such as softmax brings existing PyTorch optimizations into vLLM, reduces duplicated code, and improves maintainability across the CPU inference ecosystem. Examples for Integrating kernels via oneDNN, tuning dynamic workloads, and extending shared optimization paths across PyTorch and vLLM.
Attention works best when focused. Model attention wastes compute when every sequence element is treated as equally important.
Computationally heavy models are increasingly using long sequence attention. A good example is DLRMv3, the latest generation recommendation model. Compared to lighter architectures such as DLRMv2, DLRMv3 adds HSTU/transformer-style attention over long user histories. This improves modeling power, but significantly increases compute, memory movement, and latency.
This talk presents importance aware attention module for PyTorch inference on Arm CPUs. Instead of changing numeric precision or rewriting the model, we score sequence items against the current request and reduce the sequence only when the relevance signal is strong. When the signal is week, we keep the full sequence whereas strong signal preserves a subset in the original order.
DLRMv3 is the first proof point. Using this methodology, throughput improved by ~3x QPS (queries per second). From a business perspective, that means roughly 3x more requests for the same CPU budget while passing accuracy gates. The idea could extend to long-context LLMs and other sequence workloads.
The increasing deployment of PyTorch workloads on ARM-based cloud infrastructure has amplified the importance of efficient memory operations. Memory initialization and data movement are fundamental to tensor allocation and lifecycle management, directly impacting runtime efficiency. We present SVE-accelerated implementations of memcpy and memzero in mimalloc, the allocator adopted by PyTorch, enabling improved performance on ARMv9 processors.
This work introduces Vector Length Agnostic (VLA) memory primitives using Arm Scalable Vector Extensions (SVE), runtime hardware-aware dispatch based on available vector length, and compiler-safe integration that preserves baseline AArch64 compatibility while enabling accelerated execution on supported systems.
Benchmarking on AWS Graviton3 shows up to 6.8% higher memory throughput, 12.3% lower instruction count, and more than 80% reduction in branch mispredictions for large memory operations. These optimizations improve a foundational component of PyTorch execution and contribute to more efficient AI workloads on ARM infrastructure.
This poster will cover recent efforts to develop AI agent skills focused on debugging torch.compile issues. The compiler is a complex system which can produce difficult-to-debug errors and performance problems. Although there is extensive documentation, it is often difficult to apply to specific problems. AI agents can help, but do not always give good advice out of the box. Skills help to provide a more reliable experience (especially for smaller models) while also reducing token usage. This poster will cover the skills being developed for torch.compile, as well as basic concepts of the compiler.
Paper: accepted at IEEE AIRC, ENAS, NFM, COMPSAC (2026) to be published in IEEE scopus.
Paper Abstract : Formally verifying deep learning compilers is essential to diagnose undefined behaviors in complex deep learning frameworks like PyTorch, requiring deep understanding of graph compilation, optimized code generation,tautology and SMT-based theorem proving. The author proposes Inductor-TV, a novel multi-layer formal verification framework for PyTorch Inductor's FX backend, using Z3-based theorems against existing Graph optimizations and extending formal verification to LLVM/NVVM IR for both CPU (C++/OpenMP) and GPU (Tri-ton/PTX) pathways through Alive. "FormalBench" is introduced as a novel formally verified benchmark of TorchBench spanning language, vision, and speech models. Inductor-TV formally verified 44.8% of PyTorch GitHub Inductor failures (79.5% of Inductor-specific cases) and 43% of LLVM IR failures from 900 collected issues over 365 days. This resulted in over 880 NVVM IR files from 1700+ Triton IR compiled files on Blackwell GPUs were formally verified through Inductor-TV. Inductor-TV constitute the first and only formal verification framework for PyTorch compiler.
TraceLens is an open-source tool for analyzing AI workloads (github.com/AMD-AGI/TraceLens). However, graph execution, fused kernels, and dynamic batching in modern inference serving frameworks obscure workload visibility and complicate trace-based performance analysis.
We present a major extension to TraceLens for inference workloads. We enhance trace-collection in inference serving frameworks (upstreamed in vLLM), enriching traces with the context and annotations that power TraceLens. Trace-splitting isolates steady-state and phase-specific regions, such as decode/prefill/mix, enabling targeted analysis of relevant execution window. Graph-mode analysis correlates graph-capture and replay traces to recover call-stacks and tensor shape information needed to analyze opaque graph-executed kernels. TraceLens adds analytical roofline models for key inference operators, backed by enriched trace annotations supplying the per-request context these models need for variable-length batched attention.
Together, these capabilities turn opaque inference traces into a structured performance report—reducing diagnosis effort and enabling systematic optimization across frameworks and hardware.
Credited Co-Author(s): Deval Shah, Abdul Basit Mohammad, Tharun Adithya S, Adeem Jassani, Steve Reinhardt
PyTorch's native profiler backend Kineto ships built-in support for some backends like CUDA and XPU compiled directly into its source tree, but out-of-tree (OOT) accelerators have historically lacked a path to kernel level, device correlated tracing. AWS Neuron's integration demonstrates a reusable pattern that any OOT backend using PyTorch's PrivateUse1 mechanism can follow: implementing Kineto's IActivityProfiler and IActivityProfilerSession interfaces externally, then registering with Kineto as a child profiler without modifications to PyTorch or Kineto core. The backend's own runtime profiling and tracing APIs supply system runtime events (op queueing, host-device handoff) and device execution events, tagged with correlation IDs from the framework dispatch layer. All events are post-processed and converted into a common format and backend's runtime and device events are aligned with PyTorch CPU events. As a result we get Chrome trace showing framework, runtime, and device activity along with dependencies, letting users pinpoint dispatch latency, communication/compute overlap, and host-bound stalls end-to-end providing PrivateUse1 backend reference for complete profiler support.
Credited Co-Author(s): Alvin Yin, Lucas Hendren, Joydeep Sinha
TorchTitan large-model training is communication-heavy due to its distributed nature: MoE expert parallelism runs an all-to-all every step to dispatch and combine tokens, and across nodes, that traffic competes with compute. RCCL's device API – GPU-Initiated Networking (GIN) – lets a GPU kernel post its own RDMA transfers instead of waiting on a CPU host. This fits MoE well: driving transfers from the kernel avoids CPU-proxy serialization across many peers and lets dispatch/combine overlap and pipeline with expert compute, hiding the communication portion behind the step instead of paying for it in series. We integrate different versions of GIN into torchtitan's DeepSeek-V3 expert-parallel all-to-all on AMD MI300X with minimal code changes: a runtime patch swaps the collective for a custom op with no fork, plus the config that lets the GPU-initiated and normal network paths coexist. It trains to baseline loss matching the traditional all-to-all. Turning that into the full end-to-end overlap win is open work the poster lays out.
Credited Co-Author(s): Liz Li, Nusrat Islam, Maria Garzaran, Atul Kulkarni, Kapil Shyam Pawar, Andy Luo, Shashidhar Gandham
GPU kernel optimization challenges LLMs beyond standard coding tasks, as it requires an understanding of hardware architecture, parallel computing optimization strategies, and profiling outputs. However, most existing approaches leveraging LLMs for kernel generation apply standard prompting and feedback loops, considering hardware only through profiling feedback. We introduce KernelFoundry, an evolutionary framework that efficiently explores the space of GPU kernels through (1) MAP-Elites quality-diversity search with kernel-specific behavioral dimensions to sustain exploration; (2) meta-prompt evolution that co-evolves prompts with kernels to uncover task-specific optimization strategies, and (3) a template-based parameter optimization approach to tune kernels to inputs and hardware. We evaluate this framework on KernelBench, robust-kbench and custom tasks, generating SYCL kernels as a cross-platform GPU programming paradigm, and CUDA kernels for comparison to prior work. Our approach consistently outperforms the baseline methods and achieves an average speedup of 2.3 on KernelBench for SYCL.
Credited Co-Author(s): Nina Wiedemann, Quentin Leboutet, Michael Paulitsch, Diana Wofk, Benjamin Ummenhofer
PyTorch's torch.compile cache is fundamental to inference performance, enabling compiled artifacts to be reused across repeated executions. As PyTorch expands to support diverse custom AI accelerators, existing cache validity assumptions based on tensor metadata alone become insufficient. Accelerators with hardware-specific memory layouts can require different compiled artifacts even when tensors appear identical to PyTorch, leading to silent incorrect cache reuse and unreliable production deployment. Backend developers are therefore forced to either disable compilation caching or perform extensive validation to ensure correctness. We present a layout-aware extension to PyTorch's inference compilation cache that incorporates accelerator-specific memory layout information into both runtime (Dynamo) and persistent (FxGraph) cache validation. Implemented entirely out-of-tree and without modifying PyTorch core, our approach preserves existing CPU/GPU behavior while enabling safe, reliable cache reuse for current and future third-party inference accelerators.
Credited Co-Author(s): Hema Prasanna KC, Pradipta Ghosh, Albin Joy, Anto John
Traditional benchmarks measure NCCL collectives and GEMMs in isolation, missing how they actually behave in distributed training and inference: running concurrently on the same GPU, contending for compute and bandwidth.
Maestro is an open-source framework that benchmarks operations as they truly overlap. Users describe workloads in YAML as ordered patterns of communication and compute blocks mapped onto independent process-group axes, each on its own CUDA stream, so realistic scenarios like all_gather overlapped with a Megatron MoE layer, or 2D/3D parallel strategies, can be measured directly.
Built on PyTorch and CUPTI, Maestro reports per-operation and per-pattern latency (avg/min/max/P99), bandwidth, a shared-bandwidth metric for overlapping collectives, and overlap percentage. Its extensible block registry, presets, and multi-node execution via torchrun/srun let practitioners reproduce production communication patterns and pinpoint contention bottlenecks invisible to micro-benchmarks.
The growing diversity of IRs, programming models, and hardware accelerators has dramatically expanded the AI kernel optimization space. While AI-driven systems (e.g., KernelEvolve) can automate kernel generation, the challenge shifts from generating candidates to efficiently exploring and pruning design choices as spaces grow. Hardware profiling and simulation remain essential for validation but are too costly for iterative design-space exploration. Analytical cost models are therefore critical for rapidly evaluating alternatives, guiding optimization, and enabling hardware-software co-design before hardware measurements are available.
We present MAKCI, a Python-native instrumentation framework for analytical performance modeling across abstraction levels. MAKCI instruments Python-native IRs through fine-grained compute and memory primitives, then estimates kernel cost via a concrete-execution tracer — without hardware profiling. This enables developers and AI agents to rapidly explore and prune optimization choices. As a case study, we apply MAKCI to Triton kernels targeting the IBM Spyre AI accelerator and demonstrate accurate cost estimation without hardware access.
Credited Co-Author(s): Ritik Raj
CUDA Graphs cut CPU overhead in PyTorch, but they also change how memory management works, and a workload that runs correctly can reserve far more memory than expected. PyTorch's caching allocator was built around eager execution: it relies on event polling to decide when memory is safe to reuse, and that mechanism does not work under multi-stream graph capture. This session looks at how capture changes the lifetime of both device and pinned host memory, and walks through two new allocator changes. On the device side, a capture-aware reuse path (`graphcapturerecordstreamreuse`) makes cross-stream blocks usable again. On the host side, a graph-private pool for pinned buffers (`HostPrivatePool`) keeps them alive instead, because the CPU consumer sits outside the graph and reuse cannot be proven safe. A benchmark shows large reductions in reserved memory at little cost to capture or replay. Attendees will learn why memory behaves differently under CUDA Graphs, when reuse is safe versus when correctness requires holding it, and how to reason about the trade-off in their own workloads.
Most AI safety is a scalar at the output: one good/bad check before an action ships. We keep the structure instead — who is affected, what is owed, who consented, who bears imposed risk — and contract to a decision only at the end, auditably.
Two open-source pieces realize it. erisml-compiler (alpha, on PyPI) turns natural-language moral material into a typed MoralGraph and a rank-1…6 moral tensor, read through four ethical lenses at once: consequentialist, deontic (Kantian gates via a Z3 solver), virtue, and care. When the lenses disagree, it refuses to silently aggregate and defers to a human. At runtime, ErisML’s three-layer Safety Gateway drops the same evaluation inside an agent’s plan→act loop, sealing every decision in a SHA-256 hash-chained DecisionProof and failing safe, never open.
The geometry is the payoff: ethics as a D4 symmetry problem over Hohfeld’s normative positions, with a Bond Index that tests whether a judgment survives the agent↔patient swap. The PyTorch hook is literal — forward hooks on transformer layers compare what a model says against what it internally exhibits. Running today on a Qwen/Gemma stack and two governed AI NPCs.
Have you ever hit a graph break in `torch.compile` and wondered why? Graph breaks are one of the main reasons torch.compile fails to accelerate more PyTorch programs. We are eliminating an entire class of them by bringing CPython's type slot machinery to Dynamo.
In CPython, type slots are the internal mechanism that defines an object's behavior. For example, a + b dispatches through the nb_add slot defined by the operands types. While CPython gets this behavior automatically, Dynamo historically maintained independent implementations of these operations, leading to correctness gaps and graph breaks.
We are systematically replacing Dynamo's ad hoc object model with one that mirrors CPython's slot-based dispatch, letting Dynamo faithfully reproduce CPython semantics while simplifying its implementation. The results are already visible: after introducing slot support, the pass rate on our integration of CPython's own test suite went up from 37% to 50%.
This poster walks through the design and implementation of how CPython dispatches through type slots, how Dynamo now mirrors that dispatch, and examples of graph breaks this work eliminates.
Running huge Mixture-of-Experts models like DeepSeek-V3 in production means spreading their experts across many GPUs, and often across separate pods. Once you do that, two things start to hurt: the constant expert dispatch/combine traffic, and moving the KV cache between prefill and decode. Keeping everything correct once a single data-parallel job spans multiple pods is harder still.
This poster shows how we tackled that on AMD Instinct GPUs by bringing the MoRI communication stack into vLLM. MoRI-EP handles the expert traffic over RDMA so Wide Expert Parallelism (Wide-EP) can reach across pods, and MoRI-IO moves the KV cache for disaggregated prefill/decode. Both sit behind one transfer layer, so the same setup runs whether you drive it with an llm-d sidecar or vLLM's own router.
We will walk through the correctness and stability bugs we hit spanning pods and how we fixed them, and share vLLM benchmark numbers showing better throughput and latency on long, high-concurrency workloads, with the pieces reusable by the wider PyTorch and vLLM community.
Credited Co-Author(s): Rishi Madduri, Ravi Gupta, Chaitanya Sri Krishna Lolla
Disaggregated prefill-decode LLM serving demands high-throughput, low-latency KV-cache transfer between nodes. Current solutions (Mooncake, NIXL) primarily target NVIDIA hardware, and adapting them to other accelerator platforms requires significant engineering effort. We present FlagCX's P2P engine as a vLLM KV-transfer connector delivering one-sided RDMA with native multi-chip support from the ground up.
FlagCX pushes KV blocks directly into remote memory via one-sided RDMA without remote CPU involvement. We analyze design trade-offs against Mooncake around memory registration, connection lifecycle, and completion signaling.
We validate on PPU and MetaX GPUs with GLM5.2, achieving competitive TTFT compared to Mooncake Transfer Engine.
The full stack is open-source under FlagOS.
Agentic models are becoming essential for modern AI applications, enabling systems to reason, plan, use tools, and complete complex tasks across diverse domains. In this work, we present how Multi-teacher On-Policy Distillation (MOPD) is used to train leading agentic models with more than 10 teacher models representing a broad range of domain expertise, culminating in Nemotron Ultra. We share key learnings from the training process, including strategies for teacher selection, on-policy data generation, distillation, evaluation, and alignment for agentic behavior. Finally, we discuss practical tips for practitioners building their own agentic systems and outline future directions for improving reliability, generalization, tool use, and real-world task performance.
Credited Co-Author(s): Jiaqi Zeng
We present an extension to PyTorch’s UCX based communication backend that introduces a new level of fault tolerance for large scale distributed training and inference. Leveraging UCX’s multi transport capabilities, when a link or NIC fails, ongoing collective and point to point operations are transparently migrated to alternate available routes, preserving in flight tensor transfers and avoiding job restarts. Our design monitors fabric health, maintains multiple candidate paths between ranks, and coordinates failover within PyTorch’s ProcessGroup layer so that user code and training state remain unchanged. When the failed link becomes available again, traffic is shifted back to the preferred high performance path using a controlled recovery protocol that preserves ordering and consistency. Microbenchmarks with ucx_perftest on multi NIC clusters and injected link and switch failures show that our UCX extensions maintain high bandwidth and low latency during failover and recovery, with modest overhead compared to fault free runs, providing improved robustness to network faults.
Credited Co-Author(s): Zihao Zhao, William Gallagher, Leonid Genkin, Yossi Itigin, Gal Shalom
Activation memory is one of the largest and most controllable contributors to training memory and PyTorch exposes a growing set of APIs to manage it. Today these include core primitives like torch.utils.checkpoint for recomputation and torch.autograd.graph.saveoncpu for offloading, as well as lower-level extension points such as torch.autograd.graph.savedtensorshooks that let more sophisticated variants be implemented out of tree. However, there's also a real opportunity to make deciding what to save, recompute, or offload simpler to express and reason about. This session looks at where activation memory APIs in PyTorch are heading: making these tradeoffs explicit, composable, and portable across eager and compile, informed by real training workloads.
Credited Co-Author(s): Edward Yang
A new class of memory hardware is emerging between HBM and host DRAM: CXL-attached memory pools, NVLink-reachable DDR, and other near-accelerator fabrics that are larger than HBM and faster than host DRAM. These "secondary fast memory" systems are emerging across the industry as spill targets for hot KV blocks that no longer fit in HBM. However, this hardware is not yet broadly available to the inference community. This talk shows how we extend vLLM's SimpleCPUOffloadConnector to emulate a three-level hierarchy (HBM ↔ secondary fast memory ↔ slow host DRAM) using two pinned-host CPU pools, so the community can iterate on placement, metadata, and worker plumbing today and swap in real hardware later with no scheduler changes.
We walk through: a CPU-Tier abstraction that keeps the scheduler symmetric across tiers; a partitioned placement model that admits requests by priority; per-tier copy backends, events, and prefix-cache lookup; and backward compatibility with the existing single-pool config. You leave with the design pattern, connector hooks, and a working emulator to evaluate placement policies for KV-cache workloads — ready for secondary-memory hardware as it ships.
Credited Co-Author(s): Chen Wang
Distributed RL can be fast and stable, but still wrong. When rollout and training engines assign different log-probabilities to the same tokens, they are no longer executing one policy; adding more GPUs only scales the mismatch.
We present a correctness-first workflow for DeepSeek-V4 Flash, a 284B MoE using FP8 SGLang for rollout and BF16 Megatron-LM on PyTorch for training. We use identical-token rescoring to enforce consistency at three points: model behavior, quantized weight synchronization, and inference runtime state. This reveals mismatches in hash-routed experts and Megatron's mHC post-mix, makes FP4/E8M0 transfers datatype-aware, and rebuilds SGLang's FP8 state after each update.
In the open-source Miles framework, these changes reduce the mean absolute train-rollout log-probability gap from 0.25 to 0.03 (8.3x) and enable 100+ stable online steps on 32 AMD Instinct MI355X GPUs. During the run, held-out AIME-2024 pass@1 rises from 0.39 to 0.49 and pass@8 from 0.53 to 0.67. The poster will show the same-token test, the fixes behind each mismatch, and the checks we use before scaling cross-engine RL.
Credited Co-Author(s): Liz Li, Yuankai Chen
Streaming automatic speech recognition (ASR) systems deployed on-device must satisfy strict latency constraints while maintaining recognition quality under chunked, low-latency inference. In this work, we present a practical methodology and case study for evaluating and optimizing streaming ASR on Arm-based edge devices using PyTorch and ExecuTorch. Using the Nemotron-0.6B streaming ASR model as an anchor, we profile end-to-end execution and identify key inference bottlenecks.
We then explore model and system optimizations leveraging Arm-specific capabilities and analyzing their impact on both real-time factor (RTF) and streaming WER. In particular, we examine how quantization, mixed precision, chunk sizing, and kernel selection influence the tradeoff between latency and recognition quality in streaming deployment.
The outcome of this work is a practical playbook for developers deploying streaming ASR with ExecuTorch: how to profile streaming workloads, reason about precision and kernel choices, and evaluate systems using both latency and streaming-aware accuracy metrics. Our goal is to help bridge the gap between model-centric ASR evaluation and real-world on-device deployment.
Credited Co-Author(s): Kshitij Sisodia
Communication remains the principal bottleneck for large-scale distributed AI workloads. While accelerator communication libraries are the standard choice for some bandwidth-bound large-message collectives such as AllReduce and AllGather, there remain many communication operations and application regimes where other libraries such as MPI can supplant or be used in tandem. However, to harness this improvement in real workloads, one must modify PyTorch internals to both preserve compute-communication overlap in various parallelism schemes, and to maintain synchronization semantics when mixing communication backends. In this talk, we will cover the low-level communication requirements of each stage of the model production pipeline (pretraining, RL, and inference), how to mix communication libraries within PyTorch to always use the best communication backend for a given communication operation, and methods for benchmarking torch-based communication operations.
Credited Co-Author(s): Quentin Anthony, Dhabaleswar K. (DK) Panda
Perch 2.0 (Google) classifies ~15,000 species and produces audio embeddings for conservation. It ships as TensorFlow — a problem on the NVIDIA GB10 (Grace Blackwell, sm_121, CUDA 13), where accelerated TF is unavailable but PyTorch runs natively. Rather than wrap the TF model, I reimplemented Perch 2.0's embedding model: a log mel-spectrogram frontend and EfficientNet-B3 embedder as an idiomatic torch.nn.Module, weights from the TF graph. It reproduces TF embeddings at cosine ~1.0 on the GB10. This meant reverse-engineering — the frontend uses log scaling (not PCEN), the stem uses VALID padding. With torch.compile it runs ~2.5× faster than ONNX (~635 clips/s), fully on-device. I replaced perch-hoplite's TF model loading with the native PyTorch port, so the embed→search→classify loop runs 100% TensorFlow-free. MARS hydrophone audio is quiet, so per-window amplitude normalization was essential for parity. On ~1.56M MARS embeddings (2018, 2020, 2026), a linear classifier distinguishes Bigg's orca, Pacific white-sided dolphin, humpback whale, and ship noise (ROC-AUC 0.959); orca validated April 2018 and held-out May 2018 (181 detections, no cross-class confusion). Code to follow.
Credited Co-Author(s): Duane R. Edgington
CUDAGraph has been widely adopted to mitigate CPU launch overhead, especially on Blackwell GPUs. However, CUDAGraph usually comes with high memory overhead, which hinders its adoption in large-scale memory-tight workloads such as GenAI training. The lack of existing tools make it even harder to understand and mitigate the memory overhead from CUDAGraph. One existing approach is to dump the PyTorch memory snapshot and rely on scripts or AI to analyze, which is usually unintuitive and time-consuming.
We built a Private Memory Visualizer to present CUDAGraph memory consumption. Unlike existing CUDA Memory Visualizer that only shows active memory, Private Memory Visualizer demonstrates activities in private mempool and its high watermark, which is the key to understanding CUDAGraph memory. It also visualizes the mempool fragmentation due to different mempools or CUDA streams. We have integrated it into PyTorch's CUDA Memory Visualizer (https://docs.pytorch.org/memory_viz) for easy access.
Credited Co-Author(s): Boyuan Feng
We present an end-to-end FP8 training path for multimodal transformers in PyTorch. Starting with Transformer Engine benchmarking and reproducible H100 packaging, we integrate FP8 into the trainer, enable selective layer conversion, and build a fused Qwen3-VL decoder covering QKV projections, QK normalization, attention, RMSNorm, and SwiGLU feed-forward layers. To support downstream model deployment, we introduce checkpoint-aware conversion that enables reliable ONNX export. We share practical challenges encountered while integrating Transformer Engine into our training pipeline, including increased GPU memory usage and limited support for torch.compile(). To address these constraints, we introduce a torchao backend that provides torch.compile()- and FSDP2-compatible Float8Linear training with tensorwise or axiswise scaling, while retaining Transformer Engine as an alternative. The resulting production workflow improves overall training performance by 10%.
Embedding-based retrieval at scale is bottlenecked by GPU memory: high-dimensional FP16 embeddings dominate HBM, capping documents per shard and driving up serving cost. Quantizing to FP8 halves the footprint and nearly doubles corpus density per shard, while FP8 tensor-core matmul (torch.scaledmm) adds ~36% throughput. The catch: scoring directly in FP8 drops low-order mantissa bits, so scores diverge from FP16 and recall falls to ~89–93% — unacceptable for production. We solve this with a two-stage hybrid that decouples storage precision from scoring precision. Stage 1 runs a fast, wide FP8 pass over the full corpus on the GPU and over-retrieves 2× candidates (Top-2K); Stage 2 re-ranks only those in full-precision FP16 from CPU, down to Top-1000. Results: 99.6–99.8% recall (vs ~89% single-pass), −50% memory, +36% first-stage throughput, and ~89% rank correlation preserved. The approach roughly doubles document density per shard within the same ~87GB GPU memory budget, using ~150GB system RAM per shard for the FP16 re-rank copy halving the shards needed to serve the same corpus. Certified for throughput and correctness parity in staging by online A/B tests against the baseline.
Credited Co-Author(s): Ronak Kaoshik, Dhritiman Das
Sharing a live TorchInductor cache across distributed workers is unsafe. In a preemptible 16-rank multimodal training job, workers wrote identical hashed Inductor and Triton artifacts, causing incomplete source reads and failed atomic renames.
We replaced the shared cache with a single-writer snapshot protocol. Nodes compile locally. After the first successful step, rank 0 serializes the PyTorch Mega-Cache and atomically publishes an immutable, versioned snapshot. Restarted ranks validate and import it before their first compiled forward pass. Missing, corrupt, incompatible, or unavailable snapshots fall back to cold compilation.
The integration also preserves state-dict keys for eager resume and export, and separates cold data-bound epochs from warm compute-bound epochs. Regional compilation produced a 1.26x transformer-trunk microbenchmark speedup. In warm end-to-end training, median step time fell from 0.3947 to 0.3647 seconds, a 7.6% reduction. A 16-rank cold run published a 39 MB snapshot that a requeued job successfully reused. We provide a protocol and measurement checklist for restart-safe compiled training.
Credited Co-Author(s): Boram Yoon
Out-of-tree backends must keep pace with PyTorch's bi-monthly releases, and the ecosystem responds with cross-repo CI relays and device-agnostic test instantiation. These fix when and where tests run. We address a different gap: what they run on. Op-level tests built from synthetic inputs pass while real models fail. We present an open-source, hardware-agnostic framework that derives operator tests from real models.
Takeaways:
– Real models emit dtype combinations, broadcast patterns, and non-contiguous strides that random generators rarely produce – the cases break a fresh backend
– The gap is widening: accelerator-optimized dtypes surface silent dtype promotions, MoE routing produces unique shapes, and agent-generated kernels outpace hand-written coverage
– A TorchDynamo tracer emits per-op YAML with observed shapes, dtypes, strides, and offsets
– A pytest runner on PyTorch's test ecosystem with the YAML and validates each op against a reference
– The YAML decouples the two, so a corpus captures once run on any target
– Coverage and defect detection across 116 operators from 8 models, plus a taxonomy of bring-up traps
We also chat lessons learned for bringing up accelerators.
Credited Co-Author(s): Tuan Hoang Trong, Umamaheswari Devi, Anubhav Jana, Ashok Pon Kumar Sree Prakash
Setuptools is deprecating the setup.py interface that much of the PyTorch ecosystem still builds on — a change that will eventually reach every project, not just PyTorch. We've moved PyTorch's core build to scikit-build-core, a standards-based PEP 517 backend that drives our existing CMake build, and it has now landed in main. This poster is a fast tour of what that means for you.
We'll cover what changed and why. If you build PyTorch from source, we'll show the developer-experience wins you get along the way, like a source tree that stays clean, editable installs that can rebuild on import so you never run stale code, and a new spin CLI that puts build, test, lint, and docs behind one command.
If you build downstream of PyTorch, here's the good news: nothing breaks today, our compatibility helpers and your Setuptools builds keep working. But the same deprecation that pushed us will reach you. We'll point to what PyTorch is putting in place — like a first-class find_package(Torch) — to make your own migration safe, and make the case for doing it on your schedule, before Setuptools forces the issue.
Delimiters do not solve prompt injection. The boundary between instruction and data is semantic, not structural, so nothing obliges a model to honour it. Most published work examines frontier API models, leaving self-hosted deployments underexamined despite their broader tool surface.
This poster red-teams a locally served agent running 1.5B-3B instruction-tuned models under vLLM, with filesystem, shell and network access. Compromise is measured rather than adjudicated: canary tokens planted in sensitive paths give unambiguous ground truth, sidestepping LLM-as-judge reproducibility issues.
Six attack families are evaluated, from direct override and encoded payloads to poisoned tool output and delayed multi-turn triggers. These run against a cumulative defence stack: fencing, spotlighting, guard-model screening, and sandboxed execution under a capability allowlist.
Attack success is reported with two figures defensive work rarely publishes: false refusal on legitimate tool-dependent tasks, and per-layer token and latency cost. The conclusion is architectural. Prompting lowers the odds but never reaches zero, and only constraining what the agent can execute bounds the damage.
In this poster, we introduce Regional Inductor, a PyTorch feature that enables selective Inductor lowering within torch.compile.
Today, torch.compile maps each captured graph to a single Inductor compilation. This one-to-one model works well for most users but limits the fine-grained control needed by power users, who may want to preserve an entire ATen graph, for example, to enable distributed optimizations, while applying Inductor only to selected regions. Selective compilation is useful when Inductor accelerates particular operations or fusions while preserving bitwise equivalence, even if such equivalence cannot be guaranteed across the full graph. It can also reduce compilation time when full-graph compilation is slow or bypass failures caused by non-critical operators.
To address these needs, Regional Inductor provides annotation APIs for selecting regions to compile while leaving the remainder of the graph as traced ATen operations.
Credited Co-Author(s): Animesh Jain
AOTInductor(AOTI) powers ~70% of Recommendation System (RecSys) inference model compute at Meta, but adoption today is all-or-nothing—an entire model must be exportable/compilable before any speedup materializes—demanding weeks to months of expert co-design. In practice, a few dense regions(submodules) dominate inference latency while surrounding regions (e.g. sparse embedding lookups, custom operations) are non-compilable or offer negligible benefits, leaving many high-value models unoptimized.
Regional-AOTI is a PyTorch-native partial compilation framework that solves this: 1) user-friendly annotation: engineers mark critical regions and invoke compile_regions(…) with representative inputs; 2) graceful fallback: compiled regions are replaced with AOTI-optimized artifacts, others fall back to eager execution seamlessly; 3) actionable diagnostics: a structured report shows per-region status and failure causes, allowing coverage expansion iteratively.
Regional-AOTI was enabled on 7+ key inference models across hardwares, yielding up to 126% QPS speedup with 300KW savings. Regional-AOTI significantly enhanced AOTI adoption and opened the door for AI-agent-driven onboarding.
Credited Co-Author(s): Tzu-Hsin Yang, Yidi Wu
Large language models (LLMs) are increasingly adopted in scientific domains, where post-training techniques such a direct preference optimization (DPO) and group relative policy optimization (GRPO) enable researchers to align pretrained models to domain-specific tasks and data. Running these workflows at scale requires tight integration between post-training frameworks, hardware backends, and distributed training infrastructure.
In this work, we evaluate popular PyTorch-based post-training frameworks — TRL, Torchtitan, Monarch, and Torchstore — on the Aurora supercomputer at the Argonne Leadership Computing Facility (ALCF), which is equipped with Intel Data Center GPU Max Series (Ponte Vecchio) accelerators. We apply DPO and GRPO to models targeting scientific use cases relevant to ongoing research at Argonne, and report training throughput and convergence behavior across model sizes and post-training methods. Our results show that Aurora is an effective platform for large-scale LLM post-training, while the PyTorch XPU ecosystem supports widely adopted post-training algorithms and frameworks with competitive performance and scalability.
Credited Co-Author(s): Guoqiong Song, Filippo Simini, Sam Foreman, Nathan Nichols, Varuni Sastry, Samuel Wheeler, Khalid Hossain, Huihuo Zheng, Murali Emani, Marieme Ngom, Ethan Wong, Venkat Vishwanath
Large-scale distributed workloads—including AI training, inference serving, reinforcement learning (RL), and HPC—are increasingly exposed to GPU, process, and network failures. A single unhealthy rank can stall collectives or force costly communicator reinitialization, reducing availability, goodput.
This poster presents resilience primitives for NCCL-backed PyTorch ProcessGroups. The shrink_group() API exposes
NCCL shrink functionality, allowing failed ranks to be excluded and a new ProcessGroup to be formed from the remaining healthy ranks without requiring them to participate.
We also discuss NCCLSHRINKABORT, which uses revoke-based semantics to quiesce a communicator before resource release, reducing secondary-failure risks compared with conventional abort recovery.
The poster covers API design, communicator lifecycle, rank remapping, failure-policy ownership, implementation lessons, and comparisons with split-based and abort-and-reinitialize approaches. As future work, NCCL Grow could add repaired or spare ranks to restore capacity or support rank replacement. Together, shrink and revoke provide a foundation for more resilient distributed workloads.
RL (reinforcement learning) requires bitwise identical numerics between trainer and generator to achieve stable training and deterministic results.
Attention is usually optimized differently for training vs inference because of drastic differences in the shape of the input data. However, this difference in implementation causes non-deterministic floating point reductions and numerical mismatches which leads to worse RL results.
This poster will explore how to unify the attention implementations between trainer and generator to achieve bitwise identity and batch invariance across runs.
We will focus on Flex Attention and Varlen Attention (w/ FA2 and FA3 backends) implementations and how mask mod computations, block sizes, kernel tile sizes, KV caching, and paging can all affect floating point reduction results.
Guaranteeing bitwise identity sometimes means compromising on kernel optimizations that pertain to only training or inference, so we'll also cover how to balance these requirements to minimize the performance hit.
PyTorch pipeline is capable of chaining models for secure and privacy-preserving LLM inference — PII/PHI redaction, safety classification, then the primary LLM — but running each PyTorch/vLLM stage on a physical GPU is the safe default that wastes silicon, since most stages are latency-bound and idle the compute they're pinned to.
SR-IOV carve one GPU into isolated VFs, each with dedicated VRAM but sharing the underlying compute; concurrent VFs get a fair share of GPU time under the driver's scheduler. We use this to co-locate a full three-stage pipeline on a single Intel Arc Pro B70: gliner-PII (PyTorch) redacts PII/PHI on one VF, Llama Guard (PyTorch) screens the sanitized prompt on a second VF, and the primary LLM (vLLM) serves on a third — no data ever leaves the workstation. The result is an on-prem, HIPAA-/GLBA-aligned inference stack on one physical GPU, with utilization higher than the strict one-GPU-per-stage baseline.
Attendees leave with a reusable pattern for enforcing policy stages in front of an LLM, performance of pipeline on SR-IOV vs. multiple GPUs, and a decision guide for when SR-IOV GPU sharing is the right primitive — versus MIG, MPS, or separate cards.
Credited Co-Author(s): Kushall Mittal, Chun Tao
PyTorch is widely used for model training and inference, but retrieval pipelines are often implemented using search engines & filtering systems, separate from inference systems. However retrieval systems are increasingly employing deep neural models to align with downstream objectives. This has created the need for a new type of engine that can run state-of-the-art inference at the retrieval layer.
In this poster, we present a torch-native retrieval engine that uses PyTorch as the primary runtime for retrieval, filtering, and ranking. The system combines GPU-resident torch tensor index, tensor-based retrieval operations, custom CUDA kernels for attribute filtering, and TorchScript model execution within a unified serving architecture orchestrated through a Rust and tch-rs backend.
We discuss the architectural decisions, performance optimizations, and operational challenges encountered while scaling the system to production workloads involving tens to hundreds of millions of documents per GPU. We present lessons learned around GPU memory management, custom kernel integration as Torch Library, end-to-end execution efficiency and failure modes at scale.
Credited Co-Author(s): Vishal Shah
GPU kernel autotuning requires trading quality against wall time. We show that retrieval-augmented generation over historical tuning trajectories provides a practical mechanism for controlling this trade-off. Comparing four strategies on 33 Triton kernels shapes on H100 with 165 runs provides us identical reliability confirming LLM-based proposals introduce no correctness penalty. LLM-only autotuner yields the largest gain in latency (−42.1% time [0.527–0.635], p<10⁻⁸) while the Hybrid autotuner achieves the lowest aggregate latency (−17.1% [0.748–0.910], p=0.012) when compared to the baseline LFBO autotuner. The RAG-LLM autotuner reduces latency by −11.3% [0.764–1.049], p=0.066, and wins the most individual kernels (11.5/33).
RAG-LLM requires 58.0% less readiness time (p<10⁻⁸) than LFBO. Hybrid's extended LFBO exploration costs a substantial time cost (~875K tokens, 9.2 h vs. 3.8 h elapsed). Historical evidence thus enables adaptive autotuning policies. Developers can select their operating point along the quality–readiness frontier with no reliability penalty.
Credited Co-Author(s): Karthick Panner Selvam, Angela Yi
Would you run a random shell script from the internet just because it has thousands of GitHub stars? Probably not. So why do we treat AI model checkpoints any differently?" Somewhere along the way, downloading a multi-gigabyte model and immediately loading it into production became an accepted workflow—and that's a software supply chain story waiting to happen.
As the PyTorch ecosystem increasingly relies on pretrained checkpoints distributed through public model hubs and internal artifact registries, establishing trust in model provenance has become just as important as achieving high inference throughput. While Safetensors eliminates the risks associated with arbitrary code execution during model loading, it does not answer equally important questions: Who produced this checkpoint? Has it been modified? Can its origin be independently verified?
DeepSpeed now provides end-to-end support for the Muon optimizer across all ZeRO stages (ZeRO-1/2/3). Muon has emerged as a promising optimizer and has already been adopted by several frontier AI labs for training their foundation models, including Kimi-K2-Thinking.
Muon is designed for the hidden 2D weight matrices that dominate modern neural networks. Rather than maintaining both first- and second-order moments as in Adam, Muon computes a momentum update and applies Newton–Schulz iterations to orthogonalize the momentum matrix before updating the weights. This design significantly reduces optimizer-state memory requirements while preserving strong optimization performance.
In this talk, we present the design, implementation, and optimization of Muon within DeepSpeed. We will discuss how we design the overall training systems with DeepSpeed, as well as integrate/optimize/scaling Muon Optimizer performance for LLM pre-training. We will also share extensive benchmarking and large-scale LLM pre-training results, highlighting improvements in memory efficiency, scalability, and training throughput.
LLM inference at batch size one is fundamentally memory-bound, with latency dominated by hundreds of small CUDA kernel launches. Megakernels overcome this by fusing the entire transformer forward pass into a single execution pipeline, overlapping memory movement, computation, and synchronization to eliminate GPU pipeline bubbles and maximize hardware utilization. In this session, we'll showcase interpreter-driven megakernel execution, large-scale kernel fusion, shared-memory management, and fine-grained synchronization, achieving significant latency improvements over conventional inference. Built on PyTorch, this work leverages custom CUDA operators and the PyTorch execution stack to deliver ultra-low-latency LLM inference while maintaining a familiar development workflow.
Credited Co-Author(s): Bhavya Nirav Shah
As a research frontier in Physical AI, World Model-based closed-loop evaluation is critical for uncovering compounding errors missed by open-loop testing. This introduces a massive systems challenge: autoregressively orchestrating tightly coupled PyTorch models (Policies, Rewarders, World Models) alongside simulators. While Kubernetes handles basic orchestration, it lacks the fine-grained coordination and fast tensor routing needed to keep PyTorch GPU engines fully utilized.
We present a distributed infrastructure bridging PyTorch and Ray. Wrapping PyTorch pipelines in asynchronous Ray Actors decouples executions of various models. We leverage dynamic batching and zero-copy tensor sharing to maximize efficiency, eliminating CPU-GPU bottlenecks. This architecture easily manages heterogeneous workloads, ensuring high GPU utilization even when policies and simulators operate at different frequencies.
Guaranteeing reproducibility across multi-model rollouts, our system scales across tens to hundreds of GPUs, achieving >10× higher throughput. We establish PyTorch and Ray as the definitive stack for Physical AI.
Credited Co-Author(s): Shenyuan Gao, Zi-ang Cao, Kaiyuan Zheng, Fangqi Zhu
How can local communities scale PyTorch knowledge and grow a sustainable ecosystem?
PyTorch Korea has spent the past eight years building multiple pathways for developers, researchers, and students to learn, apply, and contribute to PyTorch. Rather than relying on a single channel, we expanded knowledge sharing through Korean translations of the official PyTorch tutorials, open-source mentoring, university lectures, international meetups, and community-driven technical events.
This approach grew the community to 660K+ users in the past year (2× YoY) and 2.4M forum pageviews, while also lowering the barrier to open source with 26 first-time contributors and 51 merged PRs to tutorials-kr (160+ total contributors). We further expanded the ecosystem through the vLLM Korea Meetup (350+ registrations, 75+ companies) and are organizing PyTorch Day Korea (expected 300 attendees), connecting researchers, engineers, and practitioners across the PyTorch stack.
This poster shares a practical, reproducible framework for community-driven PyTorch knowledge dissemination, highlighting lessons regional communities can adapt to strengthen local ecosystems and encourage long-term participation.
PyTorch’s test suite is one of the framework’s strongest quality assets, but years of accelerator-specific assumptions make backend enablement harder than it should be. Codebase analysis shows 2,000+ hardcoded device strings, 3,600+ backend-related skips, and 3,000+ tolerance overrides across 180+ files. As PyTorch expands across CUDA, XPU, MPS, PrivateUse1, NPUs, TPUs, and ASICs, backend teams often triage test coupling before they can validate real correctness gaps.
This poster presents an ongoing effort to evolve PyTorch testing from backend-specific exception handling into reusable validation infrastructure for a heterogeneous accelerator ecosystem. We organize tests into accelerator-unrelated, accelerator-agnostic, and accelerator-specific categories, then show the engineering patterns that make this practical: device-type instantiation, OpInfo-driven operator coverage, capability-based execution, declarative filtering, tolerance cleanup, and migration paths that remove backend coupling.
Through before-and-after examples, the poster shows how reusable validation can reduce duplicated effort, clarify ownership, and make new hardware bring-up faster and more predictable.
Linear probes on a model's residual stream are among the most practical safety tools in production: labs use them to flag jailbreaks, misuse, and hallucination cheaply, because the activations they read are already computed in the forward pass. But they're harder to use when serving open weights models. The probe is trivial to train; getting the activation out of a fast inference engine is the hard part, because the optimizations that make serving fast (continuous batching, paged attention, CUDA graphs, tensor parallelism, quantization) each hide the internals a probe needs.
This work reports a look at that cost. Using PyTorch with a small open model, I built an end-to-end probe and measure what it costs alongside generation. Emitting hidden states and scoring on-device are both free, but reading each token's activation off-device in the decode loop costs ~24% throughput. I'll walk through the measurement, show where the residual stream lives in vLLM's and SGLang's forward path, and outline what per-layer, per-request activation access would take. You'll leave knowing what activation access costs today and a path to making safety probing nearly free.
We streamline production inference serving of video diffusion models on AWS Trainium with a Neuron plugin for vLLM Omni. Built on vLLM Omni's standardized plugin APIs to orchestrate diffusion pipeline and optimize sharding strategies and caching layers for models such as WAN, we extend the implementation for efficient serving on Trainium through fused NKI kernels that place communication, attention, and quantization inside a single compiled graph.
The key innovation is where communication lives. On Trainium, ahead-of-time-compiled, statically-scheduled (SPMD) execution places collectives inside the compiled graph, so communication and quantization fuse into single NKI kernels with deterministic, in-trace compute overlap.
A set of WAN-specific optimizations follows from this. Context-parallel self-attention becomes a single compiled kernel that fuses the ring collective, per-shard flash attention, and the online-softmax reduction. A fused Adaptive LayerNorm (adaLN) kernel enables Megatron-style tensor and sequence-parallel sharding of DiT: placing the post-adaLN activation on the sequence-parallel all-gather, so the collective carries half the payload and feeds a low-precision GEMM.
Credited Co-Author(s): Aneesh Shetty, Siyuan Tang, Yide Zou
SGLang-Plugin-FL is an open-source out-of-tree backend plugin that enables SGLang to run on diverse AI accelerators through the FlagOS unified multi-chip software stack. It provides a reusable integration layer connecting SGLang with accelerator-specific implementations, reducing the need for extensive hardware-specific modifications to the framework.
This project addresses a key challenge in the rapidly diversifying AI accelerator ecosystem: supporting multiple hardware backends without fragmenting the inference framework or duplicating hardware-specific code. SGLang-Plugin-FL integrates with SGLang's model execution and operator dispatch paths, maps framework operations to FlagOS and accelerator runtimes, and supports distributed inference for large-scale language models.
This session presents the architecture and implementation of SGLang-Plugin-FL, its integration with the FlagOS unified backend, and the engineering challenges of supporting custom operators, communication, memory management, and distributed execution in an out-of-tree backend. It also discusses performance results and deployment experience from multi-chip LLM inference workloads.
CUDA Graphs eliminate CPU overhead and are critical for maximum performance in deep learning today. However, they have historically required execution to be static: fixed launch topology, fixed tensor shapes, and no data-dependent decisions that would force a CPU roundtrip.
This talk presents a new contract for CUDA Graphs in PyTorch. Many workloads are dynamic in a structured way: so long as the control-flow region has a stable static output shape, CUDA Graphs can use a feature called conditional nodes to support dynamically computed branch choices and loop guards, as well as a bounded number of dynamic input shapes. The control flow ops torch.cond(), torch.while_loop(), and torch.switch() encode this contract precisely. If a workload can be expressed using these operators, PyTorch will lower the control flow to IF, WHILE, and SWITCH conditional nodes during graph capture.
We will demonstrate how this unlocks CUDA graphs for important workloads that were previously incompatible, such as skipping optimizer steps with non-finite gradients, MoE routing with fixed-capacity dispatch buffers, embedding-bag reductions for recsys, and numerical solvers that run until convergence.
Your training slows, and backward time piles up on all but one rank. The culprit turns out to be one rank's dataloader — its stall surfaces as backward time on every other rank at the next sync. Profilers see through that, but are too heavy to leave on.
StageFrontier is a tiny always-on signal that closes that gap. Each rank reports a short vector of coarse stage times — no kernel tracing, no cuda.synchronize(). At each stage boundary it takes the furthest-along rank's cumulative time; frontier increments sum exactly to observed step time and point to where the delay first became group-visible, not where sync made it appear.
Since March 2026 StageFrontier has run by default across NVIDIA's AV training fleet. A dashboard tracks every job on multiple clusters and shows how per-step time, bottlenecks, stragglers, and outliers shift over time, surfacing regressions; decisions about what to fix and where to focus start from StageFrontier's numbers. The poster covers the frontier idea and its guardrails, the PyTorch integration, and lessons from fleet-scale operation — and why "run a profiler" is now a cheaper question: which window, rank, and stage are worth tracing first?
Credited Co-Author(s): Wei Chen, NVIDIA; Ville Kallioniemi, NVIDIA
Multimodal serving combines stages with different resource profiles. Vision encoding is bursty and image-count dependent, while language prefill and decode are latency-sensitive and memory/KV-cache intensive. We present an SLO-driven study of heterogeneous E-PD serving on a Xeon-hosted Intel Arc Pro B50/B70 workstation using PyTorch/vLLM, placing replicated vision encoders on B50 GPUs and reserving B70 GPUs for combined prefill+decode. We compare colocated VLM serving against E-PD across text-only, multi-image, bursty mixed, and image-storm workloads using current Qwen multimodal models, and report SLO-goodput, P99 TTFT, P99 TPOT, stage queueing, transfer overhead, GPU utilization, and normalized resource efficiency.
Credited Co-Author(s): Kim, Min Sung
Helion is a PyTorch kernel DSL with an integrated autotuner, but its compilation backends primarily target GPUs and other accelerators. We present an emitter that lowers Helion's torch.fx-based DeviceIR to structured MLIR Linalg, making Helion a portable frontend for MLIR-based compilers.
The emitted Linalg IR can be consumed by IREE or an upstream MLIR pipeline. We extend Helion's search space with CPU-specific compilation parameters for both paths. For the IREE path, we expose internal compilation configurations to Helion's autotuner. For the upstream path, we express cache tiling, register tiling, and RISC-V Vector Extension (RVV) LMUL through the transform dialect. Candidates are selected using on-device measurements.
We evaluate both paths on a Banana Pi BPI-F3 with a SpacemiT X60 RISC-V CPU using matmul and reduction kernels. Across matmul size from 128 to 1024, both Helion-tuned paths outperform PyTorch eager, with IREE achieving a 2.5-17× speedup. Helion tuning also improves the upstream path by 1.8-2.5× over its untuned schedule.
Most real-world machine learning runs on tables, not text or images. Health records, transaction logs, sensor readings. Yet tabular ML mostly missed the foundation-model wave that reshaped NLP and vision. That's started to change with tabular foundation models (TFMs), but there's a practical mess in the way: each one installs differently, wants its data in its own shape, and reports results its own way, so even trying two or three side by side is more work than it should be. TabTune is an open-source library that smooths that over built using pytorch. One scikit-learn-style interface to run, fine-tune, and evaluate a bunch of current TFMs, with the per-model preprocessing handled for you. It also covers the parts people usually leave for later: calibration and fairness checks, ensembling, distilling a big model down to a smaller one, and benchmarking you can actually reproduce. We will cover what TFMs are, how to get from a raw table to a trained model in a few lines, and what we've found about when they help and when they don't. No prior TFM experience needed.
Credited Co-Author(s): Pratinav Seth, Mohamed Bouadi, Utsav Avaiya, Vinay Kumar Sankarapu
TCCL is a native PyTorch communication backend that enables distributed training and inference on Apple Silicon. It integrates directly with torch.distributed, operates on MPS tensors, and uses RDMA over Thunderbolt 5 for low-latency, high-bandwidth communication between Mac nodes. TCCL provides the collective operations required for data, tensor, and pipeline parallelism. On real workloads, it achieves approximately 3.8× average speedup for distributed LoRA fine-tuning, up to 3.0× speedup for tensor-parallel diffusion-model inference, and up to 2.3× faster LLM decoding across four nodes.
As LLM serving moves to disaggregated prefill/decode and GPU pools, KV cache must survive changing parallelism and topology.
This paper presents Tensor KV Cache, a tensor-native abstraction built on TorchStore with a prototype vLLM integration. Instead of opaque blocks, KV is stored as PyTorch tensors/DTensors with metadata: layer, token range, dtype, and sharding layout. This decouples cache identity from deployment layout, letting consumers retrieve slices even when producer and consumer parallelism differ.
The prototype includes a TorchStore-backed vLLM KVConnector that maps paged KV blocks into logical tensor views, plus TorchStore components that subscribe to vLLM KV block events and maintain a token-sequence prefix tree. This tree indexes reusable prefixes and guides eviction/tiering across host memory and persistent storage.
We will cover the architecture, prefix-tree metadata, TP-mismatched prefill/decode reuse, and tradeoffs in moving from block-oriented to tensor-native KV offload. Attendees will learn how PyTorch-native tensor storage enables elastic, heterogeneous, tiered KV cache systems and future techniques such as KV cache blending over logical KV tensors.
Credited Co-Author(s): Ugur Kaynar, Sagar Kewalramani
ExecuTorch now spans 17 backends and delegates across phones, edge devices, microcontrollers, and custom accelerators. As this ecosystem grows, backend validation becomes a scaling challenge: hundreds of test files, thousands of test methods, multiple comparison paths, scattered skip logic, and uneven adoption of shared infrastructure.
This poster introduces the 3C Validation Model: Capability, Correctness, and Consistency. Capability defines what a backend can execute and drives test selection through a backend capability registry. Correctness verifies numerical behavior through shared comparison infrastructure and target-aware tolerance overrides. Consistency tracks backend health and regressions over time in CI.
We share ongoing work from ExecuTorch backend testing improvements, including the centralized tolerance registry, backend capability registry, shared comparison logic, and CI health tracking roadmap. These efforts move ExecuTorch from fragmented backend-specific validation toward reusable infrastructure. The goal is shared tooling by default while preserving backend autonomy over numerical bounds, skip rules, and hardware-specific behavior.
torch.compile and FSDP2 each work well on their own. Combining them is where things get difficult – recent PyTorch changes mean that compiling through FSDP2 hooks without graph breaks is no longer supported, silently invalidating the approach most teams were using.
This poster is a practical guide to what works and what doesn't. We cover three problem areas. First, the strategy choice: fullgraph=False is simpler but prevents cross-collective optimization, while pre-FSDP compilation offers better throughput at the cost of fragility to post-sharding hooks. Second, compiled autograd constraints – the entire backward must be compilable, ruling out double backwards and several
DTensor backward paths. We document which patterns survive. Third, the SPMD gap: torch.compile does not assume SPMD by default and can make rank-divergent fusion decisions across ranks, causing NCCL timeouts that look like network failures but are actually compiler bugs.
We present a working recipe – regional compilation, pre-FSDP wrapping, explicit SPMD flags – and a decision framework for when the throughput gain justifies the added complexity over eager FSDP2.
We enable a new class of practical methods for optimized CPU inference for LLMs with YNNPack, the modular successor to industry-standard XNNPack. With a graph-first architecture, automatic loop fusion and tiling, YNNPack dramatically reduces memory pressure, improves memory locality and scales hardware utilization without monolithic, hand-crafted kernels. While a highly optimized, declarative Scaled Dot-Product Attention (SDPA) for the edge is the showcased example, this graph-level approach opens up a broader class of performance wins.
Modular vs Handcrafted: YNNPack uses a lightweight graph compiler and runtime (Slinky) to express and optimize inference using only basic operations. Automatic loop fusion improves memory locality without custom kernels.
Prefill – Declarative Flash Attention: We express memory-efficient SDPA natively as a computational graph to achieve the memory-bandwidth savings of Flash Attention entirely with pre-existing micro-kernels.
Decode – Transpose Reordering: Simple algebraic transformations and built-in buffer aliasing optimizations avoid expensive memory operations and significantly reduce memory traffic, achieving up to 8.3x performance boost.
Credited Co-Author(s): Marie White, Misha Gutman, Alexander Shaposhnikov
PyTorch PR #175746 adds ATen RVV (RISC-V Vector Extension) support but not TorchInductor code generation for LLM Linear workloads.
I extended TorchInductor with RVV ISA selection, BF16 M=1 GEMV and M>=2 GEMM, bounded template routing, and explicit packed weights. Full CPU graph capture was costly on the Banana Pi, so SGLang invokes generated C++ through regional torch.compile boundaries while retaining the packed layout. I also added M=2-15 serving support and a pinned riscv64 setup.
On a Banana Pi BPI-F3, three artifact-warm Llama-3.2-1B runs used real weights with batch/input/output=1/64/8. Against the same row-major TorchInductor path, explicit packing increased median prefill throughput by 4.77x and decode throughput by 1.63x. It raised median peak RSS by 2.30 GiB. The path also ran Qwen3-0.6B and DeepSeek-R1-Distill-Qwen-1.5B , and supported Linear buckets generated RVV C++ artifacts with no extern_kernels.mm fallback.
The next step is to reduce packed-weight memory . This issue #18072 primarily tracks my separate work about SGLang RVV attention backend and also details this TorchInductor RVV: https://github.com/sgl-project/sglang/issues/18072
Credited Co-Author(s): Jenq-Kuen Lee
LLM training in PyTorch typically relies on eager mode or selective `torch.compile`, with complex interactions between hooks, autograd, and compilation that are difficult to customize and compose, making advanced optimizations challenging.
GraphTrainer uses a minimal dispatcher-based tracer to capture the full training step—forward, loss, and backward—as an FX graph, creating a natural substrate for both manual optimization and agentic performance hill-climbing.
With compute/communication fully represented in the graph, GraphTrainer enables flexible per-tensor activation checkpointing and offloading, fine-grained comp/comm overlap, easier CUDA Graphs enablement, bitwise-preserving regional Inductor compilation, rank-agnostic precompilation, graph-based pipeline parallelism, support for heterogeneous GPUs and accelerators, and a path toward higher MFU through systematic graph optimization.
Supporting dense and MoE models with TorchTitan-native parallelism strategies, GraphTrainer demonstrates toolkit-style use of PyTorch compiler components, making distributed training explicit, inspectable, and optimizable with control over performance and numerics.
Intel XPUs offer a compelling and versatile platform for AI training across diverse scales and use cases. This poster provides a showcase of Intel XPU support in TorchTitan, PyTorch's reference architecture for large-scale LLM training.
We cover the following areas :
Ease of Use – getting started , zero code change, etc.
Supported Features – multi-dimensional parallelism, reduced precision , compiler parallelization.
XPU Silicon SKU support – from desktop to datacenter GPU.
Scaling – Demonstrated MoE LLM scaling from 2 ranks on a single node to 10,000+ ranks on Aurora supercomputer.
Training Examples – Real-world pre-training and post-training results.
This work delivers the engineering foundation that makes ambitious scaling of AI workloads possible, from the model showcased in "Scaling PyTorch on Intel XPU: Field Notes from AI for Science on Aurora" to very large MoE models at supercomputer scale.
Credited Co-Author(s): The Intel TorcTitan team – Xue Hu, Tanima Dey, Jiao Wang, Jerome Mitchell, Shahbaaz Syed Ahmed, Frost Mitchell, , Guoqiong Song, Kourdis, Panagiotis
This poster illustrates a profile-driven optimization methodology for large-scale MoE training and quantifies the cumulative impact of parallelism design, topology-aware communication, kernel selection, reduced-precision computation, and asynchronous execution. The work also contributed improvements to the broader PyTorch and TorchTitan ecosystem, including HybridEP integration, SM100 attention enablement, MXFP8 support, and BF16 optimizer-state support. These results demonstrate that PyTorch-native training infrastructure can efficiently scale frontier MoE models on Blackwell systems while remaining modular and reproducible.
Credited Co-Author(s): Syed Ahmed, George Kurian, Vivek Goel
TPUs offer compelling price-performance for large-scale training and inference, but using them has often meant leaving the tooling you know behind. This poster presents ongoing work to integrate the new torch TPU backend across the Hugging Face ecosystem, so you can target TPUs with the same code you already run on GPUs.
We cover four fronts: first-class TPU support in transformers; enhanced generation for the most popular LLMs, tuned for TPU execution; integration with the kernels library to bring optimized ops to the backend; and diffusers support to extend TPU acceleration beyond text to image and multimodal workloads.
The result is a path to TPU acceleration that stays idiomatic to PyTorch and Hugging Face. Come see the current state of the integration, what's working today, benchmarks, and where it's headed.
Credited Co-Author(s): Alvaro Moran, Jingya Huang, Sayak Paul
GPU capacity is scarce, yet teams often hold a mix of instance types (H100 + H200, B200 + B300) that sit in isolated pools because a job targets one type. We show you can train a single model across mixed-generation GPUs using only standard PyTorch FSDP/HSDP, with no custom framework and no code changes. The core obstacle is asymmetric inter-node networking: GPU generations expose different NIC "rail" counts, and NCCL rejects mismatched topologies. Forcing the higher rail count restores correctness. Because H100/H200 (and B200/B300) share compute and differ only in memory, the slow-GPU bottleneck that plagues mixed-speed clusters disappears. This fix improves FSDP throughput 10–33% over native rail counts. Across Llama 3 8B/70B, Qwen3-32B, and gpt-oss-20b MoE on both generations, mixed clusters reach 76–94% of homogeneous per-GPU throughput with identical convergence on C4. For Llama 3.1 405B, which OOMs on any single Blackwell node, mixing B200 + B300 unlocks 16-way FSDP, the only configuration where the model fits. We share configs, NCCL settings, strategy selection, and micro-batch tuning so you can use idle mixed capacity today.
PyTorch documentation advises against using CUDA Unified Virtual Memory (UVM) for deep learning because page faults, evictions, and repeated CPU–GPU transfers can make it significantly slower than explicit memory placement.
We show that these limitations are not inherent to UVM, but rather to how it is managed. With the right runtime policies, UVM can efficiently support LLM training on up to 50% fewer GPUs while maintaining comparable MFU.
To demonstrate this, we built TRANSIT (TRANsparent Scale-In for multi-node Training), a runtime that transparently integrates with existing PyTorch training stacks. TRANSIT minimizes CPU–GPU data movement through dynamic UVM policies, including prefetching and zero-copy, requiring no software changes.
Evaluated on a multi-node NVIDIA H100 GPU cluster over a RoCE network, TRANSIT outperforms TorchTitan Offloading, ZeRO-Offload, and ZeRO-Infinity by up to 68%, 59%, and 42%, respectively. By reducing each job's GPU footprint, TRANSIT also lowers inter-node network traffic by up to 49% and reduces cluster queueing delays by up to 95%.
Credited Co-Author(s): Hyungyo Kim, University of Illinois, Urbana-Champaign; Apoorve Mohan, IBM Research; Nicholas Satchanov, University of Illinois, Urbana-Champaign; Hrishi Shah, University of Illinois, Urbana-Champaign; Gaohan Ye, University of Illinois, Urbana-Champaign; Jiaqi Lou, University of Illinois, Urbana-Champaign; Robert Walkup, IBM Research; Shweta Salaria, IBM Research; I-Hsin Chung, IBM Research; Hubertus Franke, IBM Research; Seetharami Seelam, IBM Research; Nam Sung Kim, University of Illinois, Urbana-Champaign
Triton kernel launches incur Python-bound host overhead, hurting small-kernel latency. CUDA Graphs reduce this overhead but demand deterministic shapes, limiting dynamic workloads where eager mode is essential.
We aim for near-C++ host performance from existing Python Triton wrappers with minimal code changes, without sacrificing dynamic flexibility.
Trident JIT-compiles the host path. It captures PyTorch functions with Triton kernels via Torch Dynamo, exporting FX graphs and guards. The graphs are lowered through Torch-MLIR to MLIR and then compiled to LLVM native code.
The compiled host logic is invoked through TVM FFI wrappers auto-generated by atengen, dramatically cutting launch latency in dynamic scenarios. A guard-based specialization with an LLVM dispatcher handles shape and dtype changes by recompiling only when necessary.
Trident compiles to binaries, not traces, preserving Python ease while delivering native host performance for dynamic batch sizes and shapes. It has been validated on matrix addition, multiplication, softmax, and attention.
LLM inference performance is often constrained by critical operators and compiler-generated inefficiencies, particularly in Mixture-of-Experts workloads. Achieving consistently high performance across diverse AI accelerators requires coordinated optimization at both the operator and compiler levels.
This talk presents three Triton-based optimization techniques. First, we optimize critical MoE operators through workload-aware kernel design, improved data movement, efficient expert dispatch, and fused computation. Second, we introduce compiler-level layout optimizations that eliminate redundant layout conversions and improve memory-access efficiency across more than 50 operators. Third, we present instruction-scheduling optimizations that restructure dependency chains, increase instruction-level parallelism, and better overlap memory operations with computation.
Together, these techniques enable optimized Triton operators to outperform equivalent implementations written in native accelerator programming languages. They also deliver over 10% end-to-end LLM inference performance improvement across multiple workloads and AI accelerator platforms.
Recommendation system preprocessing ("preproc") modules run on CPU and account for significant serving cost at scale. These modules relied on TorchScript, whose hierarchical IR fundamentally limits optimization opportunities.
We present our experience upgrading production preproc to the PT2 stack via torch.export. The flat, functional PT2 IR enables static operator dispatch, graph-level optimizations, and ahead-of-time transformations—eliminating expensive JIT profiling and caching. On production models, PT2 delivers 1.2–2.5× latency improvements, 13–33× reductions in model-loading memory (59 MB vs. 2 GB), and 12–36× faster load times with no accuracy loss.
However, preproc modules present torch.export challenges. Pipelines are built around variable-length sparse features—jagged tensors, KeyedJaggedTensors, and varying batch sizes—making data-dependent expressions (DDEs) pervasive. We present a taxonomy of DDE patterns and share fix strategies: (1) guard idioms for TorchScript/PT2 coexistence; (2) constraints for the symbolic shape system; (3) vectorized rewrites replacing Python iteration with batched ops; and (4) API bypasses for framework code using PT2-incompatible operations.
LLMs scale, optimizing KV cache and storage is critical to boosting inference performance. This session explores recent storage optimizations in inference frameworks and how they affect memory allocation and registration in PyTorch. A key optimization is the use of HugePages (HPs) to streamline the scatter-gather process when writing to storage. During scatter-gather, the kernel gathers listed memory pages, translates virtual addresses to physical DRAM addresses, and sends them to the storage controller. HPs reduces the number of pages and scatter-gather entries, which improves efficiency and complies with storage systems that require shorter scatter-gather lists.
The shift to HP-based memory allocation, which enables the deployment of NVIDIA DOCA Memos (NVIDIA CMX) and NIXL storage-bound plugins, spurred activity in benchmarking and evaluating different approaches to memory allocation and their effects on storage performance.
The session will provide the PyTorch community with insights into 1) advanced memory allocation, including how to allocate HPs and use that externally allocated memory with PyTorch, and 2) the resulting gains in inference engine storage efficiency.
In LLM inference, achieving peak performance on modern hardware demands both hand-tuning custom kernels and aggressive operator fusion. Compiler approaches must integrate custom kernels and custom fusion passes, which struggle when pattern-matching against opaque kernel calls, fragmented low-level torch ops, and in-place memory semantics.
In this talk, we present vLLM IR, a functional intermediate representation (IR) and operator library that acts as a higher-level “dialect” in the torch FX graph representation, fully interoperable with torch ops and other custom kernels. Each IR op is declared with a canonical torch-based semantic definition, also serving as the default implementation. Crucially, consistent eager and compiled semantics allow for better debuggability. During compilation, this higher-level representation allows custom compiler passes to operate on the unified IR nodes instead of the various implementations. It also allows easy auto-tuning over the different implementations and out-of-tree compiler integration (e.g. megakernel compilers or LLM-based kernel generators).
Differential privacy guarantees are only useful if teams can reason about their cost. This poster presents an empirical accuracy-vs-epsilon curve for LoRA fine-tuning under DP-SGD, using Opacus for per-sample gradient clipping, noise calibration, and privacy accounting. We fine-tune the same base model across five noise multipliers, plot task accuracy against the resulting epsilon at fixed delta, and compare against a non-private baseline. We also cover the practical friction points teams hit applying DP-SGD to LoRA specifically — BatchNorm incompatibility, clipping-norm sensitivity, and the compute cost of running a full sweep rather than a single setting.
GPU cluster observability for PyTorch training and inference still relies on signals built for web services. Latency, traffic, errors, and saturation don't map onto GPU workloads, and no accepted replacement exists — every large operator invents an incomplete one internally.
We're building that replacement: five candidate golden signals for shared GPU clusters (queue-of-work reliability, GPU-compute effectiveness, HBM-pressure headroom, straggler amplification, output-quality calibration), each paired with an SLO template so operators know when to alert, warn, or act.
The foundation is VGAC, our calibration-gated observability platform, validated on Amazon EKS, AWS Slurm HPC, and public Alibaba and Borg traces. VGAC ties each intervention — from annotation to admission gating — to a calibration prerequisite, so a model's authority to act shrinks as trust degrades. But calibration alone doesn't answer the operational question: which signals should an SRE actually watch? That's the open problem we're presenting.
We want PyTorch practitioners running production GPU clusters to pressure-test these signals.
Repo: github.com/espirado/Golden-Signals-AI-GPU
FSDP gradient communication becomes a bottleneck while training at scale. At each step, the reduce-scatter collective needs to send full gradients over a shared bandwidth, which can cause the GPU to stall. Compressing gradients on the GPU can help, but it steals GPU cycles. We present DPU-FSDP, a system that offloads the entire reduce-scatter collective using a PyTorch FSDP communication hook, performs gradient reduction, quantizes gradients to int8 and 1-bit with error feedback, and exchanges inter-node updates over RDMA on a BlueField-2 DPU. We evaluate it on GPT-2 (124M-2.6B) on a 4-node cluster and present two findings. First, on-DPU int8 and 1-bit quantization with error feedback do not affect model quality and match the fp32 NCCL loss over 5000 steps.
Second, at 100 Gb/s, offload is 1.25-2.15x slower than NCCL despite sending 4x fewer bytes. The offload is memory-bound, and the DPU kernels saturate its ~10.5 GB/s of DRAM bandwidth. It pays off only below 17 Gb/s of inter-node bandwidth; we derive a general rule: a store-and-forward offload can process gradients at no more than Bmem/k (k DRAM crossings per byte), so it pays off only when network bandwidth is below Bmem/k.
Disaggregated inference, splitting LLM serving into prefill and decode stages, is supported in vLLM and llm-d. Existing guidance on when to use it is qualitative: "prefill-heavy workloads," "large models," "sufficient bandwidth." How prefill-heavy? How large? What bandwidth is sufficient?
llm-d-diagnostics (github.com/torchedhat/llm-d-diagnostics) is a zero-dependency Python toolkit that makes this concrete. It runs experiments on Kubernetes LLM deployments: 1) latency decomposition, 2) throughput scaling, 3) contention isolation, and 4) fault recovery, producing per-request data. Based on model architecture and GPU type, an advisory module predicts whether disaggregation provides value.
The toolkit quantifies what the community's qualitative guidance points toward. Disaggregation is not a fixed architecture choice but a workload-dependent operating point. Depending on concurrency, prompt length, and model architecture, the same deployment benefits from monolithic serving at low load and disaggregated serving under contention. For some models, the advantage reverses at high load. Attendees will leave with a measurement protocol they can run against their own clusters.
Prefix-cache reuse is what makes long-context and multi-turn LLM serving economical: keeping KV blocks alive across requests turns prefills into HBM reads. Once GPU pressure forces eviction, the cache moves to CPU or storage — and if the scheduler cannot see which tier holds which block on which engine, offloaded blocks become invisible to routing, requests land on cold engines, and tail TTFT erupts.
llm-d, a CNCF open-source cloud-native distributed inference framework, closes the loop: emit KVEvents as blocks to move between tiers, and keep a near-real-time map of which engine holds each block on which tier — GPU, CPU, or shared storage. The router's scorer ranks engines by longest prefix match weighted by tier, so a CPU hit beats a cold engine, and a storage backend with GPU-direct DMA lets new engines read cache on first request.
This session walks through the design end-to-end: the OffloadingConnector, the KVEvents protocol, the tier-weighted scorer, and a benchmark showing low TTFT after tier-aware routing. We close honestly: CPU offloading is the simpler default; storage offloading is the choice when prefix state needs to be shared across pods or outlive a single pod.
Agentic AI systems often look reliable in demos and offline evaluations. They complete benchmark tasks, produce strong scores, and seem ready to ship. Production is different. Inputs are messier, context is incomplete, tools fail, and latency becomes part of the experience.
Many failures begin before the final answer. An agent may choose the wrong tool, misuse retrieved information, or make a decision based on stale context. In multi-step workflows, one small mistake can affect everything that follows, while a single performance score hides where the problem began.
This session shows how to evaluate production agents built on PyTorch models and inference stacks. It covers trace-level analysis, PyTorch instrumentation, model serving, batching, latency profiling, and tradeoffs between accuracy, throughput, and cost. We will discuss how to connect model-level signals with agent behavior so teams can identify the first point of failure and understand how errors spread.
Attendees will learn how to turn production traces into failure categories, build targeted regression tests, and create an evaluation loop that improves as the PyTorch model and the surrounding agent system evolve.
Serverless LLM inference at the edge is defined by cold starts: every session may reload weights from disk under tight memory and power budgets. The usual suspects are hardware — NVMe bandwidth, PCIe lanes, VRAM. Profiling Qwen-2.5-1.5B on an RTX 5060 (Blackwell) with Nsight Systems, we show the real bottleneck is software: a single-threaded Python unpickling loop that starves the GPU and caps NVMe-to-VRAM throughput at 17% of rated bandwidth. We call this the Deserialization Wall. Format migration alone closes most of the gap: PyTorch Pickle → SafeTensors yields 2.0× faster Time-To-First-Token and 33.7% less load energy; GGUF (Q4KM) reaches 3.9× faster and 70% less. Framed as Joules-per-token (0.99 → 0.30), cold start dominates cost for short edge sessions. We argue serialization format is a first-class latency and sustainability decision, and that zero-copy loading belongs in every edge-AI framework by default.
Model developers targeting a new accelerator pay a porting tax: HuggingFace transformers and diffusers assume CUDA, so every vendor maintains a forked model zoo that lags upstream and breaks each release. Pytorch-Plugin-FL(torch_fl) removes this tax at its root.
Because torchfl unifies heterogeneous chips behind a single PrivateUse1 device (flagos) at the operator dispatch layer, model code never sees the hardware. A model written for CUDA runs unmodified — same frompretrained, same .to(device), same generation loop. Zero model-code changes; no vendor fork of transformers or diffusers.
Portability is solved beneath the model: FlagGems Triton kernels are the unified primary compute path, with per-tier fallback via libtorch.so. The model layer issues standard PyTorch ops and inherits torch_fl routing. A per-chip, per-model integration matrix collapses to one axis — get the chip onto FlagOS, and the entire HuggingFace model surface comes with it.
We demonstrate unmodified LLMs, diffusion, encoder and audio models across four production backends, validated against CUDA/CPU reference.
Leveraging LLM Agents to automatically optimize models running on TPU.
Project showcased in https://github.com/vlasenkoalexey/tpu_performance_autoresearch_wiki and being adopted in Google internally to simplify model optimization
This work presents a practical approach to building resilient, adaptive embedding models for multimodal LLM inference routing, with efficient deployment in the vLLM Semantic Router using native PyTorch on multiple AMD’s advanced MI355X GPU nodes.
Our approach uses learned semantic embeddings to route each request to the right model, cache path, or inference backend, supporting text, image, and audio inputs beyond keyword-based routing.
The key innovation is a 2D Matryoshka embedding architecture that adapts along two axes: encoder depth and embedding dimension. This enables fast, low-dimensional routing for simple queries and deeper, higher-quality representations for complex ones, without retraining.
We further extend routing from text-only to multimodal semantic routing through a 2DMSE embedder, aligning long-context text, image, and audio encoders into a shared embedding space.
The system is trained with native PyTorch on AMD GPUs using distributed training, bf16 mixed precision, ROCm-backed Flash Attention, staged checkpointing,and fault-aware recovery, showing a practical path for resilient multimodal model training and serving on AMD hardware.
Credited Co-Author(s): Huamin Chen