NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

Efficient Scaling of LLM Training with Flexible Context Parallelism

Yifan Niu, Han Xiao, Dongyi Liu, Wei Zhou, Jia Li

Expand abstract

Scaling long-context capabilities is crucial for Large Language Models (LLMs). However, real-world data contain a large number of sequences with heterogeneous lengths. Existing training libraries for LLMs rely on static parallelism strategies, which suffer from severe load imbalance, redundant communication, and suboptimal hardware utilization under data heterogeneity. In this work, we propose Flexible Context Parallelism (FCP), an efficient parallelism strategy that adaptively reconfigures communication groups and context parallelism degrees during LLM training. We generalize more flexible non-power-of-two parallelism degrees and develop a polynomial-time algorithm to generate near-optimal parallelism strategies with only millisecond-level overhead per training batch. FCP is able to maintain high hardware efficiency even under extreme data heterogeneity. Experimental results demonstrate that FCP significantly outperforms Megatron-LM and DeepSpeed in both LLM and MLLM training, achieving up to 1.47 $\times$ speedup in average throughput while maintaining near-linear scaling efficiency across large-scale clusters. For extremely unbalanced batches, FCP even achieves 2.24 $\times$ speedup.

Poster
Efficient Pre-Training with Token Superposition

Bowen Peng, Théo Gigant, Jeffrey Quesnelle

Expand abstract

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.

Poster
Efficient Memory Crystallization for Graph Learning under Non-Stationary Distribution Shifts

Yue Hou, Ruomei Liu, Yingke Su, Wu Junran, Ke Xu

Expand abstract

Deep graph learning models deployed in real-world systems often need to cope with non-stationary environments, where the underlying graph distribution drifts continually over time. Prevailing solutions rely on training auxiliary generative modules to synthesize memory graphs for cross-domain adaptation, which incurs substantial computational overhead and scales poorly under prolonged distribution shifts. We argue that a more economical path exists: rather than generating memory, one can crystallize it. To this end, we propose Efficient Memory Crystallization (EMC), a training-free test-time framework that distills each incoming graph domain into a compact, semantically faithful memory through a closed-form solution to a memory-oriented distribution-matching objective, thereby eliminating redundant domain information under continual covariate shifts. To preserve both generalizability and adaptability as the model traverses a long sequence of target domains, EMC further models inter-domain dependencies through state-evolving memories and admits a theoretically grounded, tighter generalization error bound than direct adaptation. Extensive experiments demonstrate the superior performance of EMC over state-of-the-art baselines on graphs under non-stationary distribution shifts, while reducing average runtime by 87.4% and GPU memory consumption by 92.4% relative to the recent competitor, making continual graph adaptation practical at scale.

Poster
Efficient Learning of Truncated Boolean Product Distributions: Influence to the Rescue

Rohan Chauhan, Ioannis Panageas

Expand abstract

Learning the natural parameters $z \in \mathbb{R}^n$ of discrete distributions $\mu_z$ from independent samples constrained to a subset $S \subseteq \\{0,1\\}^n$ is a foundational challenge in high-dimensional statistics. Existing methods for efficiently estimating truncated Boolean product distributions, notably the work of [Fotakis et al' COLT'20, Algorithmica '22], require either strong local connectivity assumptions on $S$ -- a property denoted *fatness* -- or stringent anti-concentration assumptions and necessitate the total mass of the truncation set to be a constant with respect to $n$. Moreover, the results in [Fotakis et al' COLT'20, Algorithmica '22] suffer from sample complexities that scale as $\Omega(2^n)$ if the mass of $S$ is exponentially small in $n$. In this work, we circumvent these limitations by analyzing the geometry of $S$ under the measure $\mu_z$. We refine the existing parameter estimation guarantees under the fatness assumption, improving the prior sample complexity to $\mathcal{O}( \log n / \epsilon^2)$ for $\ell_\infty$-recovery, matching the untruncated minimax rate. We further generalize fatness using the notion of influence utilized in the analysis of Boolean functions and provide sufficient conditions for efficient inference. Notably, unlike previous work, our method does not require sampling at arbitrary parameterizations of the model and instead relies on gradient descent. Lastly, we establish a theoretical lower bound demonstrating the sample complexity exhibits an intrinsic exponential dependence on the width of the model and the minimum distance between elements in the set.

Poster
Efficient Label Distribution Inference Attack and Defense on Classifier Weights in Federated Learning

Daeyoung Choi, Gyuejeong Lee

Expand abstract

Label distribution leakage in federated learning (FL) poses a critical privacy threat, exposing population-level patterns beyond individual data breaches. Building on known correlations between class sample frequencies and classifier weight magnitudes, we demonstrate how adversaries can exploit these relationships to infer label distributions from shared model parameters in FL. To characterize this threat, we propose WLIA (Weight Norm-based Label Distribution Inference Attack), which infers client label distributions by analyzing variations in classifier weight norms across communication rounds. WLIA operates on the classifier layer without requiring auxiliary data, synthetic samples, or meta-classifiers, enabling broader applicability across heterogeneous model architectures. To counter this threat, we introduce RNR (Randomized Weight Norm Regularization), a defense mechanism that disrupts frequency-weight correlations by randomly regularizing weight norms during local training. RNR adds only a penalty term to the cross-entropy loss, making it efficient and easy to integrate. Comprehensive experiments across six datasets demonstrate that WLIA achieves superior inference accuracy compared to existing attacks, while RNR effectively mitigates leakage with minimal impact on utility.

Poster
Efficient LLM Adaptation with Forward-Only Passes

Baichuan Huang, Ananth Balashankar, Amir Aminifar

Expand abstract

Large language model (LLM) inference serving is undergoing rapid growth and large-scale deployment, motivating us to rethink how the inference process itself can be leveraged to enable efficient task-specific LLM adaptation. In this paper, we propose a forward-only approach for efficient LLM adaptation with forward-only passes in LLM inference serving. We exploit the angle concentration of activations induced by each singular value decomposition (SVD) component to measure its contribution—_dispersed_ or _concentrated_—to the angle concentration of the hidden states. Based on this, our forward-only search then efficiently identifies the weight matrix with the highest dispersion that merits rank reduction. We then selectively remove the higher-order components and retain the lower-order components in SVD. Empirical results across diverse datasets demonstrate the competitive accuracy of our forward-only approach, while theoretical analysis shows lower peak memory and greater speedup than the gradient-based approach, and greater speedup than the exhaustive search approach. The extended experiments further present the robustness and generalization of our forward-only approach to various LLMs with up to 57B parameters. The code is included in the supplementary material.

Poster
Efficient Image Synthesis with Sphere Latent Encoder

Tung Do, Thuan H Nguyen, Hao Li

Expand abstract

Few-step image generation has seen rapid progress, with consistency and meanflow-based methods significantly reducing the number of sampling steps. Despite their low inference cost, these approaches often suffer from training instability and limited scalability. Sphere Encoder is a recent alternative that produces high-quality images in only a few steps; however, it requires repeated transitions between the pixel space and latent space during inference while jointly optimizing reconstruction and generation within a single architecture. This design leads to computational inefficiency and objective conflict between reconstruction and generation. To address these limitations, we decouple the framework into a fixed pretrained image encoder and a separate latent denoising model trained entirely in a spherical latent space. Our approach eliminates repeated pixel-space operations during training and inference, improving efficiency and allowing reconstruction and generation to specialize independently. On Animal-Faces, Oxford-Flowers and ImageNet-1K datasets, our method significantly outperforms Sphere Encoder in both generation quality and inference speed, while achieving competitive results against strong few-step and multi-step baselines.

Poster
Efficient Fine-Tuning for Structured Sparsity Under Group Repartitioning

Diyang Li

Expand abstract

Structured sparsity extends classical sparse learning by encouraging sparsity at the level of predefined groups rather than individual parameters. This paradigm is widely adopted in modern machine learning and is typically enforced by sparsity-inducing norms. In real pipelines, the optimal group partition is often unknown a priori and does not remain static after initial training. Parameter groupings are often modified as domain knowledge or upstream representation evolves. Standard approaches require retraining the model from scratch under the new target partition, which discards the optimization effort already invested and becomes computationally prohibitive at scale. In this work, we focus on the group-wise $\ell_{p,1}$-norm, a flexible regularization family that encompasses classical group Lasso and other sparse variants. We develop a principled algorithmic framework that leverages the pre-trained optimum to fine-tune model weights toward a new desired group structure without full retraining. Technically, we decompose arbitrary group-wise transitions into a sequence of tractable primitives and derive the closed-form learning dynamics for each sub-action. Extensive experiments on real-world datasets demonstrate that our method achieves significant acceleration while provably recovering the correct optimal solution.

Poster
Efficient Decoder Scaling Strategy for Constructive Neural Routing Solvers

Qing Luo, Fu Luo, Ke Li, Zhenkun Wang

Expand abstract

Construction-based neural routing solvers, typically composed of an encoder and a decoder, have emerged as a promising approach for solving vehicle routing problems. While recent studies suggest that shifting parameters from the encoder to the decoder enhances performance, most works restrict the decoder size to 1–3M parameters, leaving the effects of scaling largely unexplored. To address this gap, we conduct a systematic study comparing two distinct strategies: scaling depth versus scaling width. We synthesize these strategies to construct a suite of 12 model configurations, spanning a parameter range from 1M to $\sim$150M, and extensively evaluate their scaling behaviors across three critical dimensions: parameter efficiency, data efficiency, and compute efficiency. Our empirical results reveal that parameter count is insufficient to accurately predict the model performance, highlighting the critical and distinct roles of model depth (layer count) and width (embedding dimension). Crucially, we demonstrate that scaling depth yields superior performance gains to scaling width. Based on these findings, we provide and experimentally validate a set of design principles for the efficient allocation of parameters and compute resources to enhance the model performance.

Poster
Efficient Computation and Best-Response Dynamics in Anonymous Two-Action Games with Linear Utilities

Michail Fasoulakis, Evangelos Markakis, Ioannis Panageas, Christodoulos Santorinaios, Jingming Yan

Expand abstract

The computational complexity of anonymous games has been a central theme in algorithmic game theory, with foundational contributions by Daskalakis and Papadimitriou [2015], Goldberg and Turchetta [2017], and Cheng et al. [2017] establishing that the landscape of equilibrium computation admits a PTAS. In this paper, we investigate the class of $n$-player anonymous games with linear utilities and two actions. While finding an equilibrium in anonymous games is PPAD-hard for arbitrary utility functions—even with a constant number of actions [Chen et al., 2015]—we show that this linear structure allows for significant algorithmic improvements. Exploiting the fact that payoffs depend solely on the first moment of the players’ strategy distribution, we provide the first algorithm for computing exact Nash equilibria that runs in time polynomial in the number of players $n$. We achieve this by establishing a novel combinatorial characterization of the Nash equilibria. Finally, we establish a novel connection between Nash equilibria in these games and non-convex non-concave min-max optimization. Despite the technical challenges posed by this structural property, we design a sequential best-response dynamic that provably converges to an $\epsilon$-Nash equilibrium in $\mathcal{O}(\frac{1}{\epsilon})$ steps.

Poster
Efficient Collaborative LLM Fine-Tuning over Heterogeneous Mobile Devices via Many Backbones to One Side-Network Tuning

Xingke Yang, Liang Li, Sicong Li, Liwei Guan, Hao Wang, Xiaoqi Qin, Jiang Liu, Xin Fu, Miao Pan

Expand abstract

Collaboratively fine-tuning (FT) large language models (LLMs) over heterogeneous mobile devices fosters immense potential applications of personalized intelligence. However, such a vision faces critical system challenges. Existing federated LLM fine-tuning approaches remain limited by prohibitive on-device resource overhead, straggler-prone synchronous aggregation, and the assumption of a unified pretrained backbone across heterogeneous devices. In this paper, we propose EC-MobiLLM, a novel design for efficient collaborative LLM fine-tuning across heterogeneous mobile devices and heterogeneous backbones. EC-MobiLLM adopts a server-assisted side-tuning paradigm to minimize on-device overhead and pioneers non-blocking asynchronous collaboration to accelerate training. Moreover, EC-MobiLLM introduces adaptive feature alignment to support collaboration across heterogeneous backbones, aligning with practical needs in real-world deployments. Extensive experimental results demonstrate that EC-MobiLLM can maintain robust fine-tuning performance while achieving extremely low on-device memory, with at least 95.2\% reduction in computation overhead, 93.2\% reduction in communication costs and $5.1\times$ faster convergence compared to existing methods, validating its efficacy for practical LLM adaptation over heterogeneous mobile devices.

Poster
Efficient Brain-to-Speech Decoding with Fixed-Delay Spiking Neural Networks

Aleksandra Wisniewska, Seo-Hyun Lee, Seong-Whan Lee

Expand abstract

Brain-to-speech decoding aims to restore communication by translating neural activity into speech-related representations. From a modeling perspective, this task can be formulated as a weakly supervised neural sequence decoding problem, where phoneme sequences must be recovered without frame-level annotations. Existing speech decoders commonly rely on recurrent sequence models to integrate temporal context, but their dense sequential computation can limit their suitability for resource-constrained portable or implantable neuroprosthetic systems. In this study, we propose a residual spiking neural network (SNN) for phoneme-level decoding of intracortical speech signals under constrained model capacity and weak sequence-level supervision. Our approach incorporates biologically inspired fixed connection-specific delays to integrate temporal context, providing each postsynaptic unit with access to recent presynaptic spike activity without recurrent state or dense temporal convolutions. Because each directed connection is assigned a single fixed delay, the temporal field is expanded without adding trainable delay parameters. Across 500k, 2M, and 5M core-parameter budgets, fixed delays consistently improve phoneme error rate (PER), with gains saturating once the delay range covers the task-relevant temporal horizon. Boundary-aligned CTC analysis further suggests that delayed temporal context reduces local uncertainty around phoneme transitions, indicating that fixed delays support weakly supervised phoneme alignment rather than simply increasing model capacity. Compared with recurrent, convolutional, and feedforward baselines under matched core-parameter budgets, the proposed SNN achieves the best PER at 500k and 2M parameters and the second-best PER at 5M parameters, while reducing estimated operation-level energy by approximately 2-4$\times$ relative to the strongest non-spiking baselines. These results suggest that fixed-delay SNNs provide an efficient temporal modeling mechanism for brain-to-speech decoding, potentially enabling lightweight decoding for on-device speech neuroprostheses.

Poster
Efficient Benchmarking Is Just Feature Selection and Multiple Regression

Sam Bowyer, Acyr Locatelli, Kris Cao

Expand abstract

Efficient benchmarking techniques aim to lower the computational cost of evaluating LLMs by predicting full benchmark scores using only a subset of a benchmark's questions. By reframing this problem as an instance of multiple regression with feature selection, we find that existing efficient benchmarking methods can be greatly improved by simply using kernel ridge regression at the prediction stage. Additionally, using an information-theoretic feature-selection algorithm called minimum redundancy maximum relevance (mRMR), we can further improve upon these methods by selecting question subsets that will be maximally useful for prediction. Except in very data-poor settings, these approaches consistently achieve smaller prediction errors (in both MAE and RMSE), and greater ranking correlation between predicted and true scores (in both Spearman $\rho$ and Kendall $\tau$) across a range of benchmarks using both binary and continuous metrics. Furthermore, mRMR subsampling is much faster than competitor methods (which often involve fitting probabilistic models or running clustering algorithms), and is more likely to select the same questions under different random seeds or training data splits.

Poster
Efficient Algorithms for Distributed Saddle Problems

Ruichen Luo, Anton Rodomanov, Sebastian Stich

Expand abstract

The distributed setting for Saddle Problems (SPs) has recently emerged as a framework for modern applications in machine learning and multiagent systems. Despite its relevance, the theoretical foundations of this setting have not yet been thoroughly established. In this paper, we advance this research direction by formalizing the distributed setup for SPs and providing rigorous definitions of communication and oracle costs. Further, we prove lower bounds for any distributed gradient-span algorithm, which reveals the gap from existing methods and this theoretical limit. To this end, we provide a Decoupled Method built upon a novel multi-stage reduction that reduces the SP into a sequence of decoupled minimization tasks of residual norms. Our algorithm matches the communication lower bound, thus setting the communication complexity within the gradient-span algorithms. Moreover, it yields the first strict improvement over the long-standing oracle cost of the Extragradient method for general SPs. Finally, we study the extension of distributed SP into Variational Inequality Problem (VIP), which generalizes two-player zero-sum games to multiplayer general-sum games. We show that our Decoupled Method achieves a new state-of-the-art communication complexity for this broader class.

Poster
Efficient Algorithms For Fully Dynamic Bipartite Matching In Metric Spaces

Pankaj Agarwal, Oliver Chubet, Sharath Raghvendra, Arian Zamani

Expand abstract

We consider the problem of maintaining a minimum-cost bipartite matching between two point sets $A$ and $B$ in a metric space $(\mathbb{X},\mathsf{d})$ under insertions and deletions of input points. The Wasserstein-$p$ ($W_p$) cost of a matching $M$ is defined as $(\sum_{(a,b)\in M}\mathsf{d}(a,b)^p)^{1/p}$ and an $\alpha$-approximate matching is one whose total cost is at most $\alpha$ times the cost of a minimum-cost matching. We obtain two main results. If $A$ and $B$ are points in $\mathbb{R}^d$ and the distance between two points is measured under the $\ell_p$-metric, an $O(d\epsilon^{-3/2})$-approximate $W_1$-matching between $A$ and $B$ can be maintained with an amortized $O\big(n^\epsilon d\epsilon^{-3/2}\varphi(n,1/\sqrt{\epsilon})\big)$ update time, for any parameter $\epsilon\in(0,1]$, and where $\varphi(n,1/\sqrt{\epsilon})$ denotes the update time of a $\epsilon^{-1/2}$-approximate nearest neighbor data structure. For the $\ell_2$ norm this results in an update time of $O\big((d+ \epsilon^{-3/2}n^\epsilon)\log n\big)$. The current state-of-the-art for high-dimensional settings only focus on the $\ell_2$-metric and maintain an $O(\log n)$-approximate matching. Our result not only extends to any $\ell_p$ norm, but it also improves the approximation factor for $d=o(\log n)$ under the $\ell_2$-metric. Next, we consider the problem of maintaining a $W_p$-matching when $A$ and $B$ are point sets in an arbitrary finite metric space. We show that for any $k,p\geq 1$ and constant $\epsilon>0$, a $2k(1+\epsilon)$-approximate matching can be maintained with an amortized update time of $O(kn^{1+1/k}\epsilon^{-1}\log \Delta)$, where $\Delta$ is the spread of $A\cup B$. Prior work on finite metric spaces gave an insertion-only data structure for $W_1$-matching. In contrast, we develop a fully dynamic approach that supports both insertions and deletions and work for all $W_p$-matchings. Together, our results extend dynamic Wasserstein matching beyond fixed-dimensional $W_1$ settings: they handle high-dimensional geometric instances, extend prior insertion-only $1$-Wasserstein matching guarantees to the fully dynamic setting, and support $p$-Wasserstein costs for every integer $p\ge 1$ in arbitrary metric spaces.

Poster
Dynamics-Aware Sparse Attention for Efficient Autoregressive Video Diffusion

Yuanyu He, Zhuokun Chen, Yefei He, Zhiwei Tang, Jiasheng Tang, Yinghao Yu, Jianfei Cai, Bohan Zhuang

Expand abstract

Autoregressive Diffusion Transformers have emerged as a powerful paradigm for long video generation. However, they are often constrained by prohibitive computational overhead and substantial memory footprints. This inefficiency arises from dense computation mechanisms that allocate uniform processing resources across all spatiotemporal regions, failing to leverage the inherent temporal redundancy of video where motion is typically confined to sparse areas. In this paper, we propose **Dynamics-Aware Sparse Attention (DASA)**, a unified framework that shifts from dense computation to sparse, motion-driven generation to eliminate both storage and computational redundancies. Inspired by video codec principles, our approach explicitly decouples video content into dynamic and static regions. To mitigate memory bottlenecks, we selectively compress the historical KV cache by retaining only dynamic features within the generated chunk. Furthermore, to reduce computational overhead, we employ a region-level computation allocation strategy based on localized dynamics. Additionally, a dynamics-guided distillation method is introduced to enhance performance via a lightweight distillation training process. Experiments on VBench demonstrate that our approach accelerates autoregressive video generation by 2.16$\times$, while preserving high visual fidelity.

Poster
Dynamic Resolution Routing for Efficient Egocentric Grounding

Huixin Sun, Wangbo Zhao, Fanyue Wei, Qiuxia Lin, Pengzhan Sun, Angela Yao

Expand abstract

Egocentric visual grounding requires high-resolution inputs to localize small objects. However, scaling Multimodal Large Language Models to this domain is constrained by the excessive cost of visual token processing. We identify that current efficient strategies based on token reduction are unreliable to select object-centric spatial evidence. To overcome this, we propose SmartRes, a framework that performs efficiency optimization in the pixel space via dynamic resolution routing. SmartRes first encodes a low-resolution view for global context and uses a lightweight router to activate high-resolution patches in object-centric regions and constructs an order‑preserving visual sequence. To further enable robust routing under severe foreground‑background imbalance, we introduce a margin-regularized routing objective that increases foreground-background logit separation and improves foreground recall. Experiments on Ego4D and EgoIntention show that SmartRes reduces visual tokens by up to 67\% while retaining 86.4\% of full-resolution performance, and achieves up to $1.66\times$ faster inference than state-of-the-art token reduction methods with higher accuracy. Furthermore, strong performance on small object grounding indicates the effectiveness of SmartRes towards egocentric applications. Code will be publicly available.

Poster
Dynamic Convolutions Improve Transformers

Oliver Sieberling, Bharat Runwal, Rameswar Panda, Yoon Kim

Expand abstract

Transformers have become the dominant architecture for large language models, largely due to the scalability and flexibility of attention, feed-forward layers, residual connections, and normalization. This paper introduces dynamic short convolutions as an additional neural network primitive for improving Transformer architectures. Unlike static short convolutions, dynamic convolutions use input-dependent filters, preserving the locality bias of convolution while increasing expressivity. Motivating experiments show that applying dynamic short convolutions to key, query, and value representations improves performance on challenging associative recall tasks compared with static convolutional variants. Across language-modeling experiments ranging from 150M to 2B parameters, dynamic convolutions consistently outperform standard Transformers and Transformers augmented with static short convolutions. Scaling-law analysis indicates a 1.33$\times$ compute advantage over parameter-matched Transformers, while a custom Triton kernel enables efficient execution with minimal end-to-end slowdown. These results suggest that dynamic short convolutions are a scalable, hardware-efficient, and expressive primitive for advancing Transformer-based language models.

Poster
DyRA: Dynamic Residual Approximation for Efficient Matrix Multiplication in DNNs

Daewon Chae, Hyunwon Chung, Changwoo Lee, Hun-Seok Kim

Expand abstract

Large-scale foundation models achieve strong performance across diverse tasks, but their size makes inference costly, largely due to dense matrix multiplications. Prior work reduces this cost by replacing dense weight matrices with efficient structured forms such as low-rank factorizations. However, these methods approximate weights rather than the output activations that determine inference accuracy. Consequently, small weight-space errors can be amplified by input activations, producing large output errors. In this work, we propose DyRA, an input-adaptive method that improves structured matrix multiplication approximation by correcting residual output errors during inference. We show that matrix multiplication can be approximated more effectively by directly optimizing low-rank factors of the output. DyRA builds on this insight by dynamically approximating and correcting the output error introduced by structured weight approximations. This combines efficient structured computation with input-dependent correction, yielding a more faithful approximation of full matrix multiplication under the same computational budget. Across vision, speech, and language models, DyRA consistently improves the accuracy–efficiency trade-off over structured weight approximations alone. Notably, DyRA achieves a $1.5\times$ end-to-end GPU speedup for DINOv3 while reducing accuracy degradation by more than $3\times$ relative to weight-only baselines.

Poster
DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

Jiading Gai, Shuai Zhang, Xiang song, Yuyang (Bernie) Wang, George Karypis

Expand abstract

Modern RL post-training methods such as GRPO and DAPO train on $N$ response sequences of $R$ tokens sampled from a shared prompt of $P$ tokens, but standard FlashAttention replicates all $P$ prompt tokens $N$ times across both forward and backward passes --- duplicating compute and memory on identical hidden states. In large-rollout, long-context RL training ($N{\geq}16$, $P{\geq}8\text{K}$), this redundancy dominates the policy update cost. We observe that in decoder-only models, causal masking makes prompt representations invariant across sequences at every layer, so all per-token operations (norms, projections, MLP) and attention can process the prompt once --- a property not yet exploited at the kernel level for training. We propose \textbf{DualKV}, the first FlashAttention kernel variant that eliminates shared-prompt replication during RL training, via (1) fused CUDA forward and backward kernels that iterate over two disjoint KV regions --- shared context and per-sequence response --- in a single kernel launch, and (2)~a data-pipeline redesign in veRL that repacks $N(P{+}R)$ tokens into $P{+}NR$ tokens per micro-batch, extending the token reduction from attention to the entire model by a factor $\rho = N(P{+}R)/(P{+}NR)$. DualKV is mathematically equivalent to standard attention and introduces no approximation. On Qwen3-8B GRPO training with 8$\times$H100 GPUs ($N{=}32$, 8K-context), DualKV achieves $1.63$--$2.10\times$ policy-update speedup, enables $2\times$ larger micro-batches, and raises MFU from $36\%$ to $76\%$. Similar gains hold for DAPO ($2.47\times$ speedup, $77\%$ MFU). At 30B MoE scale on 16$\times$H100, DualKV achieves $3.82\times$ policy-update and $3.38\times$ end-to-end step speedup over FlashAttention (which requires 4-way Ulysses sequence parallelism to avoid OOM).

Poster