arXiv subjects

Kaijian Wang

Publications and source records attributed to Kaijian Wang.

6 recordsLinked to original sources

arXiv2026

AutoParallel: Composable, PyTorch-Native Automatic Parallelism for Production Training

Training large models requires multiple interacting forms of parallelism. As models and clusters scale, finding and maintaining efficient parallelization plans becomes essential for training performance. Diverse model architectures, workloads, and hardware topologies demand flexible strategies, yet production PyTorch systems typically rely on hand-written recipes that are costly to adapt and tend to reuse familiar patterns, which may be suboptimal for a given workload. We present AutoParallel, a PyTorch-native planner for distributed model training. AutoParallel searches parallelization choices jointly across the training computation, allowing established techniques to be combined and unconventional strategies to emerge. It expresses its decisions using native PyTorch abstractions and produces ordinary PyTorch programs, avoiding a separate execution runtime while composing directly with the PyTorch compilation and training stack. AutoParallel also composes automatic decisions with model-specific distributed code, allowing dynamic or specialized regions to remain under user control. Plans are inspectable and replayable artifacts. Across dense, hybrid, and mixture-of-experts models on meshes of up to four dimensions and 128 NVIDIA H100 GPUs, AutoParallel matches or exceeds expert-written TorchTitan recipes on most configurations, with throughput gains of up to 18.3%, and discovers plans that uniform expert recipes do not express.

cs.DC↗

arXiv2026

AoiZora: Compiler-Agnostic, Placement-Aware Auto-Parallelization for Video Diffusion Transformers

Video diffusion transformers (DiTs) process tens of thousands of spatio-temporal tokens per denoising step, so low-latency inference and training commonly spread them across several accelerators. Auto-parallelization systems price a sharding on its logical graph, but its cost depends on the program the compiler generates and, on direct networks such as TPU sub-slices, on which chip runs each logical rank; both are visible only after compilation, while the chip assignment must be fixed before it. We present AoiZora, a compiler-agnostic auto-parallelization framework built on two IRs, one for the program a partitioner shards and one for the program a compiler emits; it chooses the logical mesh, the layout of every tensor and operation, and the device assignment of an unsharded model by pricing its compiled programs, without executing any candidate. At the core of AoiZora is evidence, a compiler-neutral record of a compiled program that leaves its placement open, so that one compilation can be priced under many placements by a calibrated model of the accelerator network and the placement can determine which sharding is selected. On TPU v5e, AoiZora reduces the denoising latency of Wan 2.1 and HunyuanVideo by 1.23x and 1.32x on average over expert configurations and Alpa, finding each plan within 10 minutes; it also speeds up training on TPU v5p by 1.11x on 4 chips and 1.21x on 8 chips on average and parallelizes PyTorch programs on GPUs with the same core.

cs.DC↗

arXiv2026

Characterization-Guided GPU Fault Resilience in NVIDIA MPS

NVIDIA Multi-Process Service (MPS) enables fine-grained GPU sharing by allowing multiple processes to execute concurrently on the same GPU, making it an important mechanism for improving GPU utilization. However, MPS has weak fault resilience: a fault in one process can terminate all co-running processes, limiting its adoption in resilience-critical settings such as multi-tenant GPU clusters. In this work, we design fault-resilient MPS to solve this problem. Our design is guided by insights from a systematic characterization of GPU faults and a deep analysis of their end-to-end processing pipeline. Based on these insights, we design two complementary mechanisms. First, we design a fault isolation mechanism for the dominant memory-related faults that can be fully isolated while preserving process-level fail-stop semantics by software intervention in the open GPU driver kernel module. For other faults whose process is within proprietary software, we design a fast-recovery substrate that combines virtual-memory-based GPU-resident state sharing with pre-initialized standbys. Our evaluation across GPUs and workloads demonstrates effective fault isolation and fast recovery with minimal overhead: in an end-to-end case study, isolation incurs no visible outage, while recovery restores pre-fault throughput in 355\,ms.

cs.DC↗

arXiv2026

ALTO: Adaptive LoRA Tuning and Orchestration for Heterogeneous LoRA Training Workloads

Low-Rank Adaptation (LoRA) is now the dominant method for parameter-efficient fine-tuning of large language models, but achieving a high-quality adapter often requires systematic hyperparameter tuning because LoRA performance is highly sensitive to configuration choices. In practice, this leads to many concurrent LoRA jobs, often spanning heterogeneous tasks in multi-tenant environments. Existing systems largely handle these jobs independently, which both wastes computation on weak candidates and leaves GPUs underutilized. We present ALTO (Adaptive LoRA Tuning and Orchestration), a co-designed training system that accelerates LoRA hyperparameter tuning while enabling efficient cluster sharing across heterogeneous tasks. The central insight behind ALTO is that when multiple tuning jobs run concurrently over a shared frozen backbone, they expose optimization opportunities that single-job designs cannot exploit. Building on this, ALTO monitors loss trajectories to terminate unpromising configurations early, uses fused grouped GEMM together with a new rank-local adapter parallelism to co-locate surviving adapters and reclaim freed GPU capacity, and combines intra-task and inter-task scheduling to improve multi-task placement by leveraging the predictable duration of LoRA jobs. Extensive evaluation shows that ALTO achieves up to $13.8\times$ speedup over state-of-the-art without sacrificing adapter quality.

cs.LG↗

arXiv2026

GENSERVE: Efficient Co-Serving of Heterogeneous Diffusion Model Workloads

Diffusion models have emerged as the prevailing approach for text-to-image (T2I) and text-to-video (T2V) generation, yet production platforms must increasingly serve both modalities on shared GPU clusters while meeting stringent latency SLOs. Co-serving such heterogeneous workloads is challenging: T2I and T2V requests exhibit vastly different compute demands, parallelism characteristics, and latency requirements, leading to significant SLO violations in existing serving systems. We present GENSERVE, a co-serving system that leverages the inherent predictability of the diffusion process to optimize serving efficiency. A central insight is that diffusion inference proceeds in discrete, predictable steps and is naturally preemptible at step boundaries, opening a new design space for heterogeneity-aware resource management. GENSERVE introduces step-level resource adaptation through three coordinated mechanisms: intelligent video preemption, elastic sequence parallelism with dynamic batching, and an SLO-aware scheduler that jointly optimizes resource allocation across all concurrent requests. Experimental results show that GENSERVE improves the SLO attainment rate by up to 44% over the strongest baseline across diverse configurations.

cs.DC↗

arXiv2025

SuperGen: An Efficient Ultra-high-resolution Video Generation System with Sketching and Tiling

Diffusion models have recently achieved remarkable success in generative tasks (e.g., image and video generation), and the demand for high-quality content (e.g., 2K/4K videos) is rapidly increasing across various domains. However, generating ultra-high-resolution videos on existing standard-resolution (e.g., 720p) platforms remains challenging due to the excessive re-training requirements and prohibitively high computational and memory costs. To this end, we introduce SUPERGEN, an efficient tile-based framework for ultra-high-resolution video generation. SUPERGEN features a novel training-free algorithmic innovation with tiling to successfully support a wide range of resolutions without additional training efforts while significantly reducing both memory footprint and computational complexity. Moreover, SUPERGEN incorporates a tile-tailored, adaptive, region-aware caching strategy that accelerates video generation by exploiting redundancy across denoising steps and spatial regions. SUPERGEN also integrates cache-guided, communication-minimized tile parallelism for enhanced throughput and minimized latency. Evaluations show that SUPERGEN maximizes performance gains while achieving high output quality across various benchmarks.

cs.LG↗