Showing 1–10 of 10 results for author: Ouyang, A

CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe

Authors: Tara Saba, Zhiyang Chen, Jikai Jason Li, Anne Ouyang, Xujie Si, Fan Long

Abstract: High-performance GPU kernels are critical to modern machine learning systems, yet developing them remains a manual, expert-driven process. Recent work has explored using LLMs to automate kernel generation, but generated kernels still fall short of carefully tuned references on standardized benchmarks. We present CuTeGen, an agentic GPU kernel synthesis framework that treats kernel development as a… ▽ More High-performance GPU kernels are critical to modern machine learning systems, yet developing them remains a manual, expert-driven process. Recent work has explored using LLMs to automate kernel generation, but generated kernels still fall short of carefully tuned references on standardized benchmarks. We present CuTeGen, an agentic GPU kernel synthesis framework that treats kernel development as a structured generate-test-refine workflow over the CuTe abstraction layer. Two design choices distinguish CuTeGen from prior work: targeting CuTe rather than raw CUDA, which exposes performance-critical structures such as tiling and data movement while remaining stable enough for iterative refinement, and a delayed profiling schedule that withholds low-level performance feedback until the kernel's high-level structure has stabilized. On the 209 tasks of KernelBench Level-1 and Level-2, CuTeGen achieves an average speedup of 1.71$\times$ over PyTorch and outperforms the prior agentic baseline CudaForge (0.89$\times$) at comparable per-task generation cost. Code available at https://github.com/taratt/cutegen.git △ Less

Submitted 3 June, 2026; v1 submitted 1 April, 2026; originally announced April 2026.

CausalT5k: Diagnosing Refusal and Failure Modes in Trustworthy Causal Reasoning Across Causal Rungs

Authors: Longling Geng, Andy Ouyang, Theodore Wu, Daphne Barretto, Matthew John Hayes, Rachael Cooper, Yuqiao Zeng, Sameer Vijay, Gia Ancone, Ankit Rai, Matthew Wolfman, Patrick Flanagan, Edward Y. Chang

Abstract: Large language models increasingly produce fluent causal explanations, yet they often fail in ways aggregate accuracy cannot diagnose: confusing association with intervention, abandoning correct judgments under pressure, over-refusing valid claims, or answering when evidence is underdetermined. We introduce CTK, a diagnostic benchmark of 5,147 cases and growing, across 10 domains and all three lev… ▽ More Large language models increasingly produce fluent causal explanations, yet they often fail in ways aggregate accuracy cannot diagnose: confusing association with intervention, abandoning correct judgments under pressure, over-refusing valid claims, or answering when evidence is underdetermined. We introduce CTK, a diagnostic benchmark of 5,147 cases and growing, across 10 domains and all three levels of Pearl's Ladder of Causation. Unlike benchmarks that only score correctness, CTK reveals why a model failed by annotating causal rung, trap type, pressure sensitivity, refusal quality, and Utility-Safety tradeoffs. Its Sheep/Wolf taxonomy separates valid causal designs from inferential traps; paired neutral/pressure variants measure sycophantic drift through Bad Flip Rate; and Wise Refusal fields test whether a model identifies the missing information needed before endorsing a claim. CTK exposes failure modes hidden by aggregate accuracy: the Skepticism Trap, Rung Collapse under scaling, pressure-induced drift, Detection-Correction gaps, and counterfactual error modes. Rather than prescribing a correction method, it provides the diagnostic substrate for studying causal-reasoning failure profiles. △ Less

Submitted 16 June, 2026; v1 submitted 9 February, 2026; originally announced February 2026.

Comments: 12 pages, 17 tables, 4 figures

ACM Class: I.2.7

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

Authors: Anjiang Wei, Tianran Sun, Yogesh Seenichamy, Hang Song, Anne Ouyang, Azalia Mirhoseini, Ke Wang, Alex Aiken

Abstract: 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 e… ▽ More 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. △ Less

Submitted 2 December, 2025; v1 submitted 9 September, 2025; originally announced September 2025.

Morse-based Modular Homology for Evolving Simplicial Complexes

Authors: Anqiao Ouyang

Abstract: The computation of homology groups for evolving simplicial complexes often requires repeated reconstruction of boundary operators, resulting in prohibitive costs for large-scale or frequently updated data. This work introduces MMHM, a Morse-based Modular Homology Maintenance framework that preserves homological invariants under local complex modifications. An initial discrete Morse reduction produ… ▽ More The computation of homology groups for evolving simplicial complexes often requires repeated reconstruction of boundary operators, resulting in prohibitive costs for large-scale or frequently updated data. This work introduces MMHM, a Morse-based Modular Homology Maintenance framework that preserves homological invariants under local complex modifications. An initial discrete Morse reduction produces a critical cell complex chain-homotopy equivalent to the input; subsequent edits trigger localized updates to the affected part of the reduced boundary operators over a chosen coefficient ring. By restricting recomputation to the affected critical cells and applying localized matrix reductions, the approach achieves significant amortized performance gains while guaranteeing homology preservation. A periodic recompression policy together with topology-aware gating and a column-oriented sparse boundary representation with a pivot-ownership map confines elimination to the affected columns and can bypass linear algebra when invariants are decidable combinatorially. The framework offers a drop-in upgrade for topology pipelines, turning costly rebuilds into fast, exact updates that track homology through local edits. Reframing dynamic homology as a locality-bounded maintenance task provides an exact alternative to global recomputation for evolving meshes and complexes. △ Less

Submitted 20 August, 2025; originally announced August 2025.

Fourier-Invertible Neural Encoder (FINE) for Homogeneous Flows

Authors: Anqiao Ouyang, Hongyi Ke, Qi Wang

Abstract: We present the Fourier-Invertible Neural Encoder (FINE), a compact and interpretable architecture for dimension reduction in translation-equivariant datasets. FINE integrates reversible filters and monotonic activation functions with a Fourier truncation bottleneck, achieving information-preserving compression that respects translational symmetry. This design offers a new perspective on symmetry-a… ▽ More We present the Fourier-Invertible Neural Encoder (FINE), a compact and interpretable architecture for dimension reduction in translation-equivariant datasets. FINE integrates reversible filters and monotonic activation functions with a Fourier truncation bottleneck, achieving information-preserving compression that respects translational symmetry. This design offers a new perspective on symmetry-aware learning, linking spectral truncation to group-equivariant representations. The proposed FINE architecture is tested on one-dimensional nonlinear wave interaction, one-dimensional Kuramoto-Sivashinsky turbulence dataset, and a two-dimensional turbulence dataset. FINE achieves an overall 4.9-9.1 times lower reconstruction error than convolutional autoencoders while using only 13-21% of their parameters. The results highlight FINE's effectiveness in representing complex physical systems with minimal dimension in the latent space. The proposed framework provides a principled framework for interpretable, low-parameter, and symmetry-preserving dimensional reduction, bridging the gap between Fourier representations and modern neural architectures for scientific and physics-informed learning. △ Less

Submitted 29 November, 2025; v1 submitted 21 May, 2025; originally announced May 2025.

KernelBench: Can LLMs Write Efficient GPU Kernels?

Authors: Anne Ouyang, Simon Guo, Simran Arora, Alex L. Zhang, William Hu, Christopher Ré, Azalia Mirhoseini

Abstract: Efficient GPU kernels are crucial for building performant machine learning architectures, but writing them is a time-consuming challenge that requires significant expertise; therefore, we explore using language models (LMs) to automate kernel generation. We introduce KernelBench, an open-source framework for evaluating LMs' ability to write fast and correct kernels on a suite of 250 carefully sele… ▽ More Efficient GPU kernels are crucial for building performant machine learning architectures, but writing them is a time-consuming challenge that requires significant expertise; therefore, we explore using language models (LMs) to automate kernel generation. We introduce KernelBench, an open-source framework for evaluating LMs' ability to write fast and correct kernels on a suite of 250 carefully selected PyTorch ML workloads. KernelBench represents a real-world engineering environment and making progress on the introduced benchmark directly translates to faster practical kernels. We introduce a new evaluation metric fast_p, which measures the percentage of generated kernels that are functionally correct and offer a speedup greater than an adjustable threshold p over baseline. Our experiments across various state-of-the-art models and test-time methods show that frontier reasoning models perform the best out of the box but still fall short overall, matching the PyTorch baseline in less than 20% of the cases. While we show that results can improve by leveraging execution and profiling feedback during iterative refinement, KernelBench remains a challenging benchmark, with its difficulty increasing as we raise speedup threshold p. △ Less

Submitted 14 February, 2025; originally announced February 2025.

Let's Think Frame by Frame with VIP: A Video Infilling and Prediction Dataset for Evaluating Video Chain-of-Thought

Authors: Vaishnavi Himakunthala, Andy Ouyang, Daniel Rose, Ryan He, Alex Mei, Yujie Lu, Chinmay Sonar, Michael Saxon, William Yang Wang

Abstract: Despite exciting recent results showing vision-language systems' capacity to reason about images using natural language, their capacity for video reasoning remains under-explored. We motivate framing video reasoning as the sequential understanding of a small number of keyframes, thereby leveraging the power and robustness of vision-language while alleviating the computational complexities of proce… ▽ More Despite exciting recent results showing vision-language systems' capacity to reason about images using natural language, their capacity for video reasoning remains under-explored. We motivate framing video reasoning as the sequential understanding of a small number of keyframes, thereby leveraging the power and robustness of vision-language while alleviating the computational complexities of processing videos. To evaluate this novel application, we introduce VIP, an inference-time challenge dataset designed to explore models' reasoning capabilities through video chain-of-thought. Inspired by visually descriptive scene plays, we propose two formats for keyframe description: unstructured dense captions and structured scene descriptions that identify the focus, action, mood, objects, and setting (FAMOuS) of the keyframe. To evaluate video reasoning, we propose two tasks: Video Infilling and Video Prediction, which test abilities to generate multiple intermediate keyframes and predict future keyframes, respectively. We benchmark GPT-4, GPT-3, and VICUNA on VIP, demonstrate the performance gap in these complex video reasoning tasks, and encourage future work to prioritize language models for efficient and generalized video reasoning. △ Less

Submitted 9 November, 2023; v1 submitted 23 May, 2023; originally announced May 2023.

Comments: Accepted to the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP 2023)

Visual Chain of Thought: Bridging Logical Gaps with Multimodal Infillings

Authors: Daniel Rose, Vaishnavi Himakunthala, Andy Ouyang, Ryan He, Alex Mei, Yujie Lu, Michael Saxon, Chinmay Sonar, Diba Mirza, William Yang Wang

Abstract: Recent advances in large language models elicit reasoning in a chain-of-thought that allows models to decompose problems in a human-like fashion. Though this paradigm improves multi-step reasoning ability in language models, it is limited by being unimodal and applied mainly to question-answering tasks. We claim that incorporating visual augmentation into reasoning is essential, especially for com… ▽ More Recent advances in large language models elicit reasoning in a chain-of-thought that allows models to decompose problems in a human-like fashion. Though this paradigm improves multi-step reasoning ability in language models, it is limited by being unimodal and applied mainly to question-answering tasks. We claim that incorporating visual augmentation into reasoning is essential, especially for complex, imaginative tasks. Consequently, we introduce VCoT, a novel method that leverages chain-of-thought prompting with vision-language grounding to recursively bridge the logical gaps within sequential data. Our method uses visual guidance to generate synthetic multimodal infillings that add consistent and novel information to reduce the logical gaps for downstream tasks that can benefit from temporal reasoning, as well as provide interpretability into models' multi-step reasoning. We apply VCoT to the Visual Storytelling and WikiHow summarization datasets and demonstrate through human evaluation that VCoT offers novel and consistent synthetic data augmentation beating chain-of-thought baselines, which can be used to enhance downstream performance. △ Less

Submitted 22 January, 2024; v1 submitted 3 May, 2023; originally announced May 2023.

Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

Authors: Tim Kaler, Nickolas Stathas, Anne Ouyang, Alexandros-Stavros Iliopoulos, Tao B. Schardl, Charles E. Leiserson, Jie Chen

Abstract: Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of pe… ▽ More Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment, under which we identify major performance bottlenecks hitherto under-explored by developers: mini-batch preparation and transfer. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler, a shared-memory parallelization strategy, and the pipelining of batch transfer with GPU computation. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised. Such an observation unifies training and inference, simplifying model implementation. We report comprehensive experimental results with several benchmark data sets and GNN architectures, including a demonstration that, for the ogbn-papers100M data set, our system SALIENT achieves a speedup of 3x over a standard PyTorch-Geometric implementation with a single GPU and a further 8x parallel speedup with 16 GPUs. Therein, training a 3-layer GraphSAGE model with sampling fanout (15, 10, 5) takes 2.0 seconds per epoch and inference with fanout (20, 20, 20) takes 2.4 seconds, attaining test accuracy 64.58%. △ Less

Submitted 16 March, 2022; v1 submitted 15 October, 2021; originally announced October 2021.

Comments: MLSys 2022. Code is available at https://github.com/MITIBMxGraph/SALIENT

Impact of COVID-19 on Public Transit Accessibility and Ridership

Authors: Michael Wilbur, Afiya Ayman, Anna Ouyang, Vincent Poon, Riyan Kabir, Abhiram Vadali, Philip Pugliese, Daniel Freudberg, Aron Laszka, Abhishek Dubey

Abstract: Public transit is central to cultivating equitable communities. Meanwhile, the novel coronavirus disease COVID-19 and associated social restrictions has radically transformed ridership behavior in urban areas. Perhaps the most concerning aspect of the COVID-19 pandemic is that low-income and historically marginalized groups are not only the most susceptible to economic shifts but are also most rel… ▽ More Public transit is central to cultivating equitable communities. Meanwhile, the novel coronavirus disease COVID-19 and associated social restrictions has radically transformed ridership behavior in urban areas. Perhaps the most concerning aspect of the COVID-19 pandemic is that low-income and historically marginalized groups are not only the most susceptible to economic shifts but are also most reliant on public transportation. As revenue decreases, transit agencies are tasked with providing adequate public transportation services in an increasingly hostile economic environment. Transit agencies therefore have two primary concerns. First, how has COVID-19 impacted ridership and what is the new post-COVID normal? Second, how has ridership varied spatio-temporally and between socio-economic groups? In this work we provide a data-driven analysis of COVID-19's affect on public transit operations and identify temporal variation in ridership change. We then combine spatial distributions of ridership decline with local economic data to identify variation between socio-economic groups. We find that in Nashville and Chattanooga, TN, fixed-line bus ridership dropped by 66.9% and 65.1% from 2019 baselines before stabilizing at 48.4% and 42.8% declines respectively. The largest declines were during morning and evening commute time. Additionally, there was a significant difference in ridership decline between the highest-income areas and lowest-income areas (77% vs 58%) in Nashville. △ Less

Submitted 5 August, 2020; originally announced August 2020.