Ao Sun
Author directoryPapers on this page may belong to the following people: Ao Sun, Ao Sun
2026
NOSA: Native and Offloadable Sparse Attention
Yuxiang Huang | Pengjie Wang | Jicheng Han | Weilin Zhao | Zhou su | Ao Sun | Hongya Lyu | Hengyu Zhao | Yudong Wang | Chaojun Xiao | Xu Han | Zhiyuan Liu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Yuxiang Huang | Pengjie Wang | Jicheng Han | Weilin Zhao | Zhou su | Ao Sun | Hongya Lyu | Hengyu Zhao | Yudong Wang | Chaojun Xiao | Xu Han | Zhiyuan Liu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Decoding throughput gains from larger inference batches are often limited by GPU memory, mainly due to the large size of KV cache. Training-free KV cache offloading mitigates this issue by storing redundant context on the CPU and fetching sparse KV subsets, but often harms long-generation quality due to training-inference mismatch of sparse patterns. Trainable sparse attention can improve quality, yet its unconstrained KV accesses are inefficient for offloading and may incur excessive CPU–GPU transfers. To this end, we propose NOSA, a trainable sparse attention mechanism designed for KV cache offloading. NOSA explicitly bounds CPU–GPU KV transfers, enabling low communication overhead and high decoding throughput. We further develop NOSI, an offloading inference system that realizes these efficiency gains. Experiments on 1,3,8B LLMs show that NOSA outperforms KV cache offloading baselines on general, long-input, and long-generation tasks, while improving decoding throughput by up to 5.04×, 1.92×, and 1.83× over FullAttn, InfLLMv2, and ShadowKV, respectively.
2025
Seq1F1B: Efficient Sequence-Level Pipeline Parallelism for Large Language Model Training
Ao Sun | Weilin Zhao | Xu Han | Cheng Yang | Xinrong Zhang | Zhiyuan Liu | Chuan Shi | Maosong Sun
Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)
Ao Sun | Weilin Zhao | Xu Han | Cheng Yang | Xinrong Zhang | Zhiyuan Liu | Chuan Shi | Maosong Sun
Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)
Training large language models (LLMs) heavily relies on distributed training strategies, among which pipeline parallelism (PP) plays a crucial role. As training sequences extend to 32k or even 128k tokens, current PP methods face severe bottlenecks, including substantial pipeline bubbles and high memory footprint, greatly hindering training throughput and model scalability. This paper introduces a sequence-level one-forward-one-backward (1F1B) PP method, named Seq1F1B, tailored for training LLMs on long sequences with high training throughput and memory efficiency. Unlike typical PP methods, which adopt batch-level pipeline schedule, Seq1F1B schedules the pipeline of training LLMs at the sequence level. It uses a computational strategy to partition sequences appropriately, significantly reducing pipeline bubbles and memory footprint. Compared to competitive PP baselines such as Megatron 1F1B PP, Seq1F1B achieves 1.14X training throughput with half memory footprint.Notably, Seq1F1B trains an LLM with 30B parameters on sequences up to 64k tokens using 64X NVIDIA A100 GPUs without using recomputation strategies, a feat unachievable with existing methods.We have released our code on GitHub to facilitate further research and development in LLM training on long sequences: https://github.com/thunlp/Seq1F1B.
ViDove: A Translation Agent System with Multimodal Context and Memory-Augmented Reasoning
Yichen Lu | Wei Dai | Jiaen Liu | Ching Wing Kwok | Zongheng Wu | Xudong Xiao | Ao Sun | Sheng Fu | Jianyuan Zhan | Yian Wang | Takatomo Saito | Sicheng Lai
Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: System Demonstrations
Yichen Lu | Wei Dai | Jiaen Liu | Ching Wing Kwok | Zongheng Wu | Xudong Xiao | Ao Sun | Sheng Fu | Jianyuan Zhan | Yian Wang | Takatomo Saito | Sicheng Lai
Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: System Demonstrations
LLM-based translation agents have achieved highly human-like translation results and are capable of handling longer and more complex contexts with greater efficiency. However, they are typically limited to text-only inputs. In this paper, we introduce ViDove, a translation agent system designed for multimodal input. Inspired by the workflow of human translators, ViDove leverages visual and contextual background information to enhance the translation process. Additionally, we integrate a multimodal memory system and long-short term memory modules enriched with domain-specific knowledge, enabling the agent to perform more accurately and adaptively in real-world scenarios. As a result, ViDove achieves significantly higher translation quality in both subtitle generation and general translation tasks, with a 28% improvement in BLEU scores and a 15% improvement in SubER compared to previous state-of-the-art baselines. Moreover, we introduce DoveBench, a new benchmark for long-form automatic video subtitling and translation, featuring 17 hours of high-quality, human-annotated data. Our demo is available here: https://vidove.willbe03.com/
APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs
Yuxiang Huang | Mingye Li | Xu Han | Chaojun Xiao | Weilin Zhao | Ao Sun | Hao Zhou | Jie Zhou | Zhiyuan Liu | Maosong Sun
Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Yuxiang Huang | Mingye Li | Xu Han | Chaojun Xiao | Weilin Zhao | Ao Sun | Hao Zhou | Jie Zhou | Zhiyuan Liu | Maosong Sun
Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
While long-context inference is crucial for advancing large language model (LLM) applications, its prefill speed remains a significant bottleneck. Current approaches, including sequence parallelism strategies and compute reduction through approximate attention mechanisms, still fall short of delivering optimal inference efficiency. This hinders scaling the inputs to longer sequences and processing long-context queries in a timely manner. To address this, we introduce APB, an efficient long-context inference framework that leverages multi-host approximate attention to enhance prefill speed by reducing compute and enhancing parallelism simultaneously. APB introduces a communication mechanism for essential key-value pairs within a sequence parallelism framework, enabling a faster inference speed while maintaining task performance. We implement APB by incorporating a tailored FlashAttn kernel alongside optimized distribution strategies, supporting diverse models and parallelism configurations. APB achieves speedups of up to 9.2×, 4.2×, and 1.6× compared with FlashAttn, RingAttn, and StarAttn, respectively, without any observable task performance degradation.
FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling
Weilin Zhao | Tengyu Pan | Xu Han | Yudi Zhang | Ao Sun | Yuxiang Huang | Kaihuo Zhang | Weilun Zhao | Yuxuan Li | Jie Zhou | Hao Zhou | Jianyong Wang | Zhiyuan Liu | Maosong Sun
Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Weilin Zhao | Tengyu Pan | Xu Han | Yudi Zhang | Ao Sun | Yuxiang Huang | Kaihuo Zhang | Weilun Zhao | Yuxuan Li | Jie Zhou | Hao Zhou | Jianyong Wang | Zhiyuan Liu | Maosong Sun
Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Speculative sampling has emerged as an important technique for accelerating the auto-regressive generation process of large language models (LLMs) by utilizing a draft-then-verify mechanism to produce multiple tokens per forward pass. While state-of-the-art speculative sampling methods use only a single layer and a language modeling (LM) head as the draft model to achieve impressive layer compression, their efficiency gains are substantially reduced for large-vocabulary LLMs, such as Llama-3-8B with a vocabulary of 128k tokens. To address this, we present FR-Spec, a frequency-ranked speculative sampling framework that optimizes draft candidate selection through vocabulary space compression. By constraining the draft search to a frequency-prioritized token subset, our method reduces LM Head computation overhead by 75% while ensuring the equivalence of the final output distribution. Experiments across multiple datasets demonstrate an average of 1.12× speedup over the state-of-the-art speculative sampling method EAGLE-2. Code is availableat https://github.com/thunlp/FR-Spec.
Search
Fix author
Co-authors
- Xu Han 4
- Zhiyuan Liu 4
- Weilin Zhao 4
- Yuxiang Huang 3
- Maosong Sun (孙茂松) 3
- Chaojun Xiao 2
- Hao Zhou 2
- Jie Zhou 2
- Wei Dai 1
- Sheng Fu 1
- Jicheng Han 1
- Ching Wing Kwok 1
- Sicheng Lai 1
- Mingye Li 1
- Yuxuan Li 1
- Jiaen Liu 1
- Yichen Lu 1
- Hongya Lyu 1
- Tengyu Pan 1
- Takatomo Saito 1
- Chuan Shi 1
- Zhou Su 1
- Jianyong Wang 1
- Pengjie Wang 1
- Yian Wang 1
- Yudong Wang 1
- Zongheng Wu 1
- Xudong Xiao 1
- Cheng Yang 1
- Jianyuan Zhan 1
- Kaihuo Zhang 1
- Xinrong Zhang 1
- Yudi Zhang 1
- Hengyu Zhao 1
- Weilun Zhao 1