arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2610.13474

CompileMask: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking

Hongzheng Li·Jiacheng Wan·Zhaolong Xing·Zhida Jiang·Yuanhang Yang·Junhua Shu·Jiaxing Wang·Xupeng Miao·Yingxia Shao·Zhen Chen·Ke Zhang
View original source ↗Read source document ↗Export metadata ↓

Abstract

Modern attention workloads require diverse masking patterns whose boundaries depend on runtime metadata, yet efficient attention kernels natively support only a small set of patterns. Existing methods for supporting custom masks in these kernels face an expressiveness-efficiency trade-off. Predicate-based methods are expressive but incur substantial preprocessing and execution overhead, whereas interval-based methods execute efficiently but restrict expressiveness through a predefined interval count. We present CompileMask, which addresses this trade-off by connecting expressive mask specifications with efficient attention execution. We design a compiler that analyzes visibility rules during the offline stage and generates a mask-specific preprocess kernel. We further propose MaskSpec, a region-based DSL that lets users express these rules. During the online stage, the generated kernel reads batch metadata and produces a compact Execution Descriptor that lets the attention kernel efficiently check visibility by comparing element indices against precomputed endpoints. Experiments show that CompileMask achieves up to $3.8\times$ speedup over state-of-the-art methods.

Explore related subjects

cs.DC

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Hongzheng Li, Jiacheng Wan, Zhaolong Xing, Zhida Jiang, Yuanhang Yang, Junhua Shu, Jiaxing Wang, Xupeng Miao, Yingxia Shao, Zhen Chen, Ke Zhang. 2026-10-09. CompileMask: Bridging the Gap Between Expressiveness and Efficiency in Attention Masking. https://arxiv.org/abs/2610.13474

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Optimizing Frequent Checkpointing via Low-Cost Differential for Distributed Training Systems

Distributed training of large deep-learning models often leads to failures, so checkpointing is commonly employed for recovery. State-of-the-art studies focus on frequent checkpointing for fast recovery from failures. However, frequent checkpointing generates numerous checkpoints, incurring substantial costs and thus degrading training performance. Recently, differential checkpointing has been proposed to reduce costs, but it is limited to recommendation systems, so its application to general distributed training systems remains unexplored. In this paper, we find that gradients generated during distributed training can be reused to construct differential checkpoints, while the former's size is smaller than the latter's, motivating us to reuse gradients for low-cost differential checkpointing. Based on this main idea, we propose \sysname, a frequent checkpointing framework for compression-enabled training systems that reuses compressed gradients as differential checkpoints, eliminating redundant differential computation and reducing checkpoint transmission cost. Furthermore, we extend gradient reuse to scenarios without gradient compression and propose \sysnameplus, which employs layer-wise-reuse snapshotting and incremental-merging persistence to overlap checkpointing with training execution. Experiments on diverse workloads, including billion-parameter-scale models, demonstrate that \sysname and \sysnameplus significantly reduce checkpointing overhead and enable checkpointing at frequencies as high as once per iteration, reducing training time by up to 89.2\% and 81.2\%, respectively.

cs.DC↗
arXiv2026

P-PAS: Prefill-Pressure Adaptive Scheduling for Long-Context LLM Serving

Long-context LLM applications such as retrieval-augmented generation (RAG) and agentic systems often process tens of thousands of input tokens to produce short outputs, making end-to-end request latency an important serving objective. We show that the maximum number of batched tokens (MBT), which controls the token scheduling budget in vLLM, has a scheduling-pressure-dependent effect on latency. Larger token budgets can reduce latency under low scheduling pressure, while smaller budgets become preferable under higher pressure. Consequently, no single static MBT performs best across load regimes. We introduce Prefill-Pressure Adaptive Scheduling (P-PAS), a lightweight policy that dynamically adapts the global scheduling budget using concurrent prefill and decode state. P-PAS retains a large token budget under low pressure and switches to a smaller budget when multiple prefills compete with active decoding. Across models, workload configurations, and GPUs, P-PAS maintains low end-to-end latency across changing load regimes. Kernel-level profiling shows that the crossover arises from a trade-off between large-chunk prefill efficiency and prefill--decode interference. We further extend the pressure signal to heterogeneous request shapes with P-PAS+. On mixed-shape workloads using Qwen3.8-27B and Nemotron-3.5-Lightning-30B-A3B, P-PAS+ reduces mean end-to-end latency by 3.0--3.1% relative to the best tested static token budget. Across exactly matched logical requests, 91.2% of Qwen requests and 88.5% of Nemotron requests complete faster with P-PAS+. Code and artifacts are available at https://github.com/TimoSaemann/ppas-vllm .

cs.DC↗
arXiv2026

Democratizing MoE inference on commodity GPUs with CoMoE

Deploying Mixture-of-Experts (MoE) models relies heavily on Expert Parallelism, which generates intense inter-GPU communication. Consequently, state-of-the-art inference systems require high-bandwidth, P2P interconnects (e.g., NVLink) in datacenter GPUs to handle massive token routing, making deployment prohibitively expensive. Consumer GPUs offer comparable compute power at significantly lower cost, promising to democratize MoE inference for individuals and enable privacy-preserving local deployments. However, their bandwidth-limited (only weak PCIe bus bandwidth) and host-mediated interconnects (no P2P support) introduce severe communication bottlenecks. We present CoMoE, a communication-efficient MoE inference system that resolves this mismatch through novel host-centric routing. Our key insight is that the unique communication topology provides the opportunity to elevate the host to an active routing hub, which can fundamentally reduce communication volume and eliminate global synchronization-induced stalls. Specifically, for token dispatch, we introduce host-backed token multicast to write shared tokens to the host exactly once, eliminating outbound transmission redundancy. For token combine, we propose a fine-grained, token-level aggregation mechanism using host staging buffers, which replaces rigid global synchronization and mitigates straggler effects. Evaluation on RTX 5090 GPUs shows that CoMoE improves inference throughput by up to 1.46x, approaching the performance of NVLink-capable A800 GPUs at only 23.4% of the hardware cost.

cs.DC↗