publications

2026

1. COLM

   

   Quokka: Accelerating Program Verification with LLMs via Invariant Synthesis

   Anjiang Wei, Tianran Sun, Tarun Suresh, and 3 more authors

   COLM 2026, Oct 2026

   Paper Code

   Program verification relies on loop invariants, yet automatically discovering strong invariants remains a long-standing challenge. We investigate whether large language models (LLMs) can accelerate program verification by generating useful loop invariants. We introduce Quokka, a framework for LLM-based invariant synthesis with soundness guarantees and state-of-the-art performance. Unlike prior work that treats LLM outputs as noisy symbolic material requiring substantial post-processing, Quokka adopts a simpler algorithm design that directly validates whether each LLM-generated invariant helps prove the target assertion. We construct a benchmark of 866 evaluation instances and 3589 training instances derived from SV-COMP, and evaluate 9 LLMs spanning multiple model families. We demonstrate that supervised fine-tuning and Best-of-N sampling yield measurable improvements, and we show that Quokka consistently outperforms prior LLM-based verifiers. Our code and data are publicly available at https://github.com/Anjiang-Wei/Quokka

   @article{wei2025quokka,
     title = {Quokka: Accelerating Program Verification with LLMs via Invariant Synthesis},
     author = {Wei, Anjiang and Sun, Tianran and Suresh, Tarun and Wu, Haoze and Wang, Ke and Aiken, Alex},
     journal = {COLM 2026},
     year = {2026},
     booktitle = {COLM},
     month = oct,
   }

2025

1. DL4C@NeurIPS

   

   Astra: A Multi-Agent System for GPU Kernel Performance Optimization

   Anjiang Wei, Tianran Sun, Yogesh Seenichamy, and 5 more authors

   DL4C@NeurIPS 2025, Dec 2025

   Paper Code

   GPU kernel optimization has long been a central challenge at the intersection of high-performance computing and machine learning. Efficient kernels are crucial for accelerating large language model (LLM) training and serving, yet attaining high performance typically requires extensive manual tuning. Compiler-based systems reduce some of this burden, but still demand substantial manual design and engineering effort. Recently, researchers have explored using LLMs for GPU kernel generation, though prior work has largely focused on translating high-level PyTorch modules into CUDA code. In this work, we introduce Astra, the first LLM-based multi-agent system for GPU kernel optimization. Unlike previous approaches, Astra starts from existing CUDA implementations extracted from SGLang, a widely deployed framework for serving LLMs, rather than treating PyTorch modules as the specification. Within Astra, specialized LLM agents collaborate through iterative code generation, testing, profiling, and planning to produce kernels that are both correct and high-performance. On kernels from SGLang, Astra achieves an average speedup of 1.32x using zero-shot prompting with OpenAI o4-mini. A detailed case study further demonstrates that LLMs can autonomously apply loop transformations, optimize memory access patterns, exploit CUDA intrinsics, and leverage fast math operations to yield substantial performance gains. Our work highlights multi-agent LLM systems as a promising new paradigm for GPU kernel optimization. Our code is publicly available at https://github.com/Anjiang-Wei/Astra.

   @article{wei2025astra,
     title = {Astra: A Multi-Agent System for GPU Kernel Performance Optimization},
     author = {Wei, Anjiang and Sun, Tianran and Seenichamy, Yogesh and Song, Hang and Ouyang, Anne and Mirhoseini, Azalia and Wang, Ke and Aiken, Alex},
     journal = {DL4C@NeurIPS 2025},
     year = {2025},
     month = dec,
   }

2. Preprint

   

   Translating Large-Scale C Repositories to Idiomatic Rust

   Saman Dehghan^*, Tianran Sun^*, Tianxiang Wu, and 2 more authors

   arXiv preprint, Nov 2025

   Paper Code

   @article{dehghan2025rust,
     title = {Translating Large-Scale C Repositories to Idiomatic Rust},
     author = {Dehghan, Saman and Sun, Tianran and Wu, Tianxiang and Li, Zihan and Jabbarvand, Reyhaneh},
     journal = {arXiv preprint},
     year = {2025},
     month = nov,
   }

3. EMNLP Findings

   

   LastingBench: Defend Benchmarks Against Knowledge Leakage

   Yixiong Fang^*, Tianran Sun^*, Yuling Shi, and 2 more authors

   EMNLP Findings, Nov 2025

   Paper Code

   @inproceedings{fang2025lastingbench,
     title = {LastingBench: Defend Benchmarks Against Knowledge Leakage},
     author = {Fang, Yixiong and Sun, Tianran and Shi, Yuling and Wang, Min and Gu, Xiaodong},
     booktitle = {EMNLP Findings},
     pages = {18304--18317},
     address = {Suzhou, China},
     year = {2025},
     month = nov,
   }

4. ACL Findings

   

   Building A Proof-Oriented Programmer That Is 64% Better Than GPT-4o Under Data Scarcity

   Dylan Zhang, Justin Wang, and Tianran Sun

   ACL Findings 2025, Jul 2025

   Paper

   Existing LMs struggle with proof-oriented programming due to data scarcity, which manifest in two key ways: (1) a lack of sufficient corpora for proof-oriented programming languages such as F*, and (2) the absence of large-scale, project-level proof-oriented implementations that can teach the model the intricate reasoning process when performing proof-oriented programming. We present the first on synthetic data augmentation for project level proof oriented programming for both generation and repair. Our method addresses data scarcity by synthesizing basic proof-oriented programming problems for proficiency in that language; incorporating diverse coding data for reasoning capability elicitation and creating new proofs and repair data within existing repositories. This approach enables language models to both synthesize and repair proofs for function- and repository-level code. We show that our fine-tuned 14B parameter model, PoPilot, can exceed the performance of the models that outperforms GPT-4o in project-level proof-oriented programming by 64% relative margin, and can improve GPT-4o’s performance by 54% by repairing its outputs over GPT-4o’s self-repair.

   @inproceedings{zhang2025popilot,
     title = {Building A Proof-Oriented Programmer That Is 64\% Better Than GPT-4o Under Data Scarcity},
     author = {Zhang, Dylan and Wang, Justin and Sun, Tianran},
     booktitle = {ACL Findings 2025},
     pages = {23101--23118},
     address = {Vienna, Austria},
     year = {2025},
     month = jul,
   }

5. KnowFM@ACL

   

   AttentionRAG: Attention-Guided Context Pruning in Retrieval-Augmented Generation

   Yixiong Fang^*, Tianran Sun^*, Yuling Shi, and 1 more author

   KnowFM@ACL 2025, Jul 2025

   Paper

   While RAG demonstrates remarkable capabilities in LLM applications, its effectiveness is hindered by the ever-increasing length of retrieved contexts, which introduces information redundancy and substantial computational overhead. Existing context pruning methods, such as LLMLingua, lack contextual awareness and offer limited flexibility in controlling compression rates, often resulting in either insufficient pruning or excessive information loss. In this paper, we propose AttentionRAG, an attention-guided context pruning method for RAG systems. The core idea of AttentionRAG lies in its attention focus mechanism, which reformulates RAG queries into a next-token prediction paradigm. This mechanism isolates the query’s semantic focus to a single token, enabling precise and efficient attention calculation between queries and retrieved contexts. Extensive experiments on LongBench and Babilong benchmarks show that AttentionRAG achieves up to 6.3x context compression while outperforming LLMLingua methods by around 10% in key metrics.

   @article{fang2025attentionrag,
     title = {AttentionRAG: Attention-Guided Context Pruning in Retrieval-Augmented Generation},
     author = {Fang, Yixiong and Sun, Tianran and Shi, Yuling and Gu, Xiaodong},
     journal = {KnowFM@ACL 2025},
     year = {2025},
     month = jul,
   }