Yuanxin Wei

dblp:331/1353 · DBLP profile ↗
← Back
8ranked-venue papers
3as first author
8since 2021 · last 2026
0009-0000-6551-9274ORCID · corroborated

Domains — the database's venue-derived domains; a paper can count in several

Systems, architecture and hardware · 5 · 2 first-author · 5 since 2021Artificial intelligence and machine learning · 2 · 2 since 2021Software engineering, systems software and programming languages · 1 · 1 first-author · 1 since 2021Databases, data management, data science and information retrieval · 1 · 1 first-author · 1 since 2021Graphics, computer vision, multimedia, augmented reality and games · 1 · 1 first-author · 1 since 2021
YearPublicationVenuePosition
2026 MixCache: Mixture-of-Cache for Video Diffusion Transformer Acceleration
abstract
Efficient video generation models are increasingly vital for multimedia synthetic content generation. Leveraging the Transformer architecture and the diffusion process, video DiT models have emerged as a dominant approach for high-quality video generation. However, their multi-step iterative denoising process incurs high computational cost and inference latency, which limits their practical deployment in large-scale and interactive multimedia applications. Caching, a widely adopted optimization method in DiT models, leverages the redundancy in the diffusion process to skip computations in different granularities (e.g., step, cfg, block). Nevertheless, existing caching methods are limited to single-granularity strategies, struggling to balance generation quality and inference speed in a flexible manner. In this work, we propose MixCache, a training-free caching-based framework for efficient video DiT inference. MixCache first distinguishes the interference and boundary between different caching strategies, and then introduces a context-aware cache triggering strategy to determine when caching should be enabled, along with an adaptive hybrid cache decision strategy for dynamically selecting the optimal caching granularity. Extensive experiments on diverse models demonstrate that MixCache can significantly accelerate video generation (e.g., 1.94× speedup on Wan 14B, 1.97× speedup on HunyuanVideo) while delivering both superior generation quality and inference efficiency compared to baseline methods.
Yuanxin Wei, Lansong Diao, Bujiao Chen, Shenggan Cheng, Zhengping Qian, Wenyuan Yu, Nong Xiao 0001, Wei Lin 0016, Jiangsu Du
ICMR1
2025 Skrull: Towards Efficient Long Context Fine-tuning through Dynamic Data Scheduling
abstract
Long-context supervised fine-tuning (Long-SFT) plays a vital role in enhancing the performance of large language models (LLMs) on long-context tasks. To smoothly adapt LLMs to long-context scenarios, this process typically entails training on mixed datasets containing both long and short sequences. However, this heterogeneous sequence length distribution poses significant challenges for existing training systems, as they fail to simultaneously achieve high training efficiency for both long and short sequences, resulting in sub-optimal end-to-end system performance in Long-SFT. In this paper, we present a novel perspective on data scheduling to address the challenges posed by the heterogeneous data distributions in Long-SFT. We propose Skrull, a dynamic data scheduler specifically designed for efficient long-SFT. Through dynamic data scheduling, Skrull balances the computation requirements of long and short sequences, improving overall training efficiency. Furthermore, we formulate the scheduling process as a joint optimization problem and thoroughly analyze the trade-offs involved. Based on those analysis, Skrull employs a lightweight scheduling algorithm to achieve near-zero cost online scheduling in Long-SFT. Finally, we implement Skrull upon DeepSpeed, a state-of-the-art distributed training system for LLMs. Experimental results demonstrate that Skrull outperforms DeepSpeed by 3.76x on average (up to 7.54x) in real-world long-SFT scenarios.
Hongtao Xu, Wenting Shen, Yuanxin Wei, Ang Wang, Guo Runfan, Tianxing Wang 0006, Yong Li 0045, Mingzhen Li 0001, Weile Jia
NeurIPS3
2025 Co-Designing Transformer Architectures for Distributed Inference With Low Communication
abstract
Transformer models have shown significant success in a wide range of tasks. However, the massive resources required for its inference prevent deployment on a single device with relatively constrainted resources, thus leaving a high threshold of integrating their advancements. Observing scenarios such as smart home applications on edge devices and cloud deployment on commodity hardware, it is promising to distribute Transformer inference across multiple devices. Unfortunately, due to the tightly-coupled feature of Transformer model, existing model parallelism approaches necessitate frequent communication to resolve data dependencies, making them unacceptable for distributed inference, especially under relatively weak interconnection. In this paper, we propose DeTransformer, a communication-efficient distributed Transformer inference system. The key idea of DeTransformer involves the co-design of Transformer architecture to reduce the communication during distributed inference. In detail, DeTransformer is based on a novel block parallelism approach, which restructures the original Transformer layer with a single block to the decoupled layer with multiple sub-blocks. Thus, it can exploit model parallelism between sub-blocks. Next, DeTransformer contains an adaptive execution approach that strikes a trade-off among communication capability, computing power and memory budget over multiple devices. It incorporates a two-phase planning for execution, namely static planning and runtime planning. The static planning runs offline, containing a profiling procedure and a weight placement strategy before execution. The runtime planning dynamically determines the optimal parallel computing strategy from an expertly crafted search space based on real-time requests. Notably, this execution approach can adapt to heterogeneous devices by distributing workload based on devices’ computing capabilities. We conduct experiments for both auto-regressive and auto-encoder tasks of Transformer models. Experimental results show that DeTransformer can reduce distributed inference latency by up to 2.81× compared to the SOTA approach on 4 devices, while effectively maintaining task accuracy and a consistent model size.
Jiangsu Du, Yuanxin Wei, Shengyuan Ye, Jiazhi Jiang, Xu Chen 0004, Dan Huang 0001, Yutong Lu
IEEE Trans. Parallel Distributed Syst.2
2024 Communication-Efficient Model Parallelism for Distributed In-Situ Transformer Inference
abstract
Transformer models have shown significant success in a wide range of tasks. Meanwhile, massive resources required by its inference prevent scenarios with resource-constrained devices from in-situ deployment, leaving a high threshold of integrating its advances. Observing that these scenarios, e.g. smart home of edge computing, are usually comprise a rich set of trusted devices with untapped resources, it is promising to distribute Transformer inference onto multiple devices. However, due to the tightly-coupled feature of Transformer model, existing model parallelism approaches necessitate frequent communication to resolve data dependencies, making them unacceptable for distributed inference, especially under weak interconnect of edge scenarios. In this paper, we propose DeTransformer, a communication-efficient distributed in-situ Transformer inference system for edge scenarios. DeTransformer is based on a novel block parallelism approach, with the key idea of restructuring the original Trans-former layer with a single block to the decoupled layer with multi-ple sub-blocks and exploit model parallelism between sub-blocks. Next, DeTransformer contains an adaptive placement approach to automatically select the optimal placement strategy by striking a trade-off among communication capability, computing power and memory budget. Experimental results show that DeTransformer can reduce distributed inference latency by up to 2.81 x compared to the SOTA approach on 4 devices, while effectively maintaining task accuracy and a consistent model size.
Yuanxin Wei, Shengyuan Ye, Jiazhi Jiang, Xu Chen 0004, Dan Huang 0001, Jiangsu Du, Yutong Lu
DATE1
2024 Understanding the Inference Performance of Spatial Temporal Diffusion Transformer
Yuanxin Wei, Jiangsu Du, Dan Huang 0001, Nong Xiao 0001
NPC (1)2
2024 APTMoE: Affinity-Aware Pipeline Tuning for MoE Models on Bandwidth-Constrained GPU Nodes
abstract
Recently, the sparsely-gated Mixture-Of-Experts (MoE) architecture has garnered significant attention. To benefit a wider audience, fine-tuning MoE models on more affordable clusters, which are typically a limited number of bandwidthconstrained GPU nodes, holds promise. However, it is non-trivial to apply existing cost-effective fine-tuning approaches to MoE models, due to the increased ratio of data to computation. In this paper, we introduce APTMoE, which employs affinityaware pipeline parallelism for fine-tuning MoE models on bandwidth-constrained GPU nodes. We propose an affinity-aware offloading technique that enhances pipeline parallelism for both computational efficiency and model size, and it benefits from a hierarchical loading strategy and a demand-priority scheduling strategy. To improve the computation efficiency and reduce the data movement volume, the hierarchical loading strategy designs three loading phases and efficiently allocates computation across GPUs and CPUs during these phases, leveraging different levels of expert popularity and computation affinity. With the aim of alleviating the mutual interference among the three loading phases and maximizing the bandwidth utilization, the demand-priority scheduling strategy proactively and dynamically coordinates the loading execution order. Experiments demonstrate that APTMoE outperforms existing methods in most cases. Particularly, APTMoE successfully fine-tunes a 61.2B MoE model on 4 Nvidia A800 GPUs(40GB) and achieves up to $33 \%$ throughput improvement compared to the SOTA method.
Yuanxin Wei, Jiangsu Du, Jiazhi Jiang, Xianwei Zhang 0001, Dan Huang 0001, Nong Xiao 0001, Yutong Lu
SC1
2024 Scalable and Parallel Optimization of the Number Theoretic Transform Based on FPGA
abstract
In lattice-based postquantum cryptography (PQC), polynomial multiplication is complex and time-consuming, which affects the overall computational efficiency. In addition, the parameters of different lattice-based algorithms require different number theoretic transform (NTT) structures, which limits the versatility of hardware design. To this end, this article proposes scalable and parallel optimization of the NTT based on a field-programmable gate array (FPGA). By analyzing the algorithm flow of the NTT, inverse NTT (INTT), and pointwise multiplication (PWM), an FPGA loosely coupled structure is designed, which can be used to place butterfly units of multiple pipelines in parallel and supports various modulo operations of a polynomial. In addition, to improve computing efficiency and scalability, key algorithm modules such as multipliers and modular reduction are deeply optimized. Moreover, the storage optimization of multiple RAM channels is carried out, and the alternate access control of data and the multiplexing of RAM resources reduce resource consumption and improve data access efficiency. For the SHA-3 algorithm, the scalable Keccak algorithm is implemented in a serial–parallel hybrid manner and supports multiple hash modes. Finally, taking the Dilithium algorithm as an example, through the parallelization of SHA-3 and NTT, the calculation cycle of key generation, signature, and verification is shortened. The experimental results and analysis show that the scheme in this article shortens the NTT calculation period while ensuring a high frequency, and the calculation time is significantly better than that of other schemes. Furthermore, it can support the optimized parallelization of multiple moduli and give full play to the computing advantages of an FPGA.
Bin Li 0023, Yunfei Yan, Yuanxin Wei, Heru Han
IEEE Trans. Very Large Scale Integr. Syst.3
2022 MCS: An In-battle Commentary System for MOBA Games
abstract
This paper introduces a generative system for in-battle real-time commentary in mobile MOBA games. Event commentary is important for battles in MOBA games, which is applicable to a wide range of scenarios like live streaming, e-sports commentary and combat information analysis. The system takes real-time match statistics and events as input, and an effective transform method is designed to convert match statistics and utterances into consistent encoding space. This paper presents the general framework and implementation details of the proposed system, and provides experimental results on large-scale real-world match data.
Xiaofeng Qi, Zhongping Liang, Jigang Liu, Yuanxin Wei, Lanxiao Huang
COLING6