publications

2026

1. NeurIPS

   

   MonarchRT: Efficient Attention for Real-Time Video Generation

   Krish Agarwal, Zhuoming Chen, Cheng Luo, Yongqi Chen, Haizhong Zheng, Xun Huang, Atri Rudra, and Beidi Chen

   In Advances in Neural Information Processing Systems (NeurIPS), Sydney, Australia, Dec 2026

   Abs arXiv Bib Blog Code

   Real-time video generation with Diffusion Transformers is bottlenecked by the quadratic cost of 3D self-attention, especially in real-time regimes that are both few-step and autoregressive, where errors compound across time and each denoising step must carry substantially more information. In this setting, we find that prior sparse-attention approximations break down, despite showing strong results for bidirectional, many-step diffusion. Specifically, we observe that video attention is not reliably sparse, but instead combines pronounced periodic structure driven by spatiotemporal position with dynamic, sparse semantic correspondences and dense mixing, exceeding the representational capacity of even oracle top-k attention. Building on this insight, we propose Monarch-RT, a structured attention parameterization for video diffusion models that factorizes attention using Monarch matrices. Through appropriately aligned block structure and our extended tiled Monarch parameterization, we achieve high expressivity while preserving computational efficiency. We further overcome the overhead of parameterization through finetuning, with custom Triton kernels. We first validate the high efficacy of Monarch-RT over existing sparse baselines designed only for bidirectional models. We further observe that Monarch-RT attains up to 95% attention sparsity with no loss in quality when applied to the state-of-the-art model Self-Forcing, making Monarch-RT a pioneering work on highly-capable sparse attention parameterization for real-time video generation. Our optimized implementation outperforms FlashAttention-2, FlashAttention-3, and FlashAttention-4 kernels on Nvidia RTX 5090, H100, and B200 GPUs respectively, providing kernel speedups in the range of 1.4-11.8X. This enables us, for the first time, to achieve true real-time video generation with Self-Forcing at 16 FPS on a single RTX 5090.

   @inproceedings{agarwal2026monarchrt,
     author = {Agarwal, Krish and Chen, Zhuoming and Luo, Cheng and Chen, Yongqi and Zheng, Haizhong and Huang, Xun and Rudra, Atri and Chen, Beidi},
     title = {MonarchRT: Efficient Attention for Real-Time Video Generation},
     booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
     year = {2026},
     month = dec,
     location = {Sydney, Australia},
   }

2. arXiv

   

   FlashRT: Agent Harness for Guiding Agents to Deploy Real-Time Multimodal Applications

   Krish Agarwal, Zhuoming Chen, Yanyuan Qin, Zhenyu Gu, Atri Rudra, and Beidi Chen

   Under review , 2026

   Abs arXiv Bib Blog Code

   Real-time multimodal applications, including voice agents and interactive video generation, compose heterogeneous models into pipelines whose efficient deployment requires application-specific decisions about placement, streaming, and intra-model parallelism. Existing serving systems and auto-parallelism compilers commit to limited transformations and fixed workload assumptions, so achieving high performance on a new application requires hand-crafting an efficient implementation. We present FlashRT, an agent harness that guides coding agents to lift simple developer-written reference implementations into optimized multi-GPU deployments that flexibly weigh target metrics like latency and throughput. Using a new chain-of-program paradigm, FlashRT directs a generic coding agent through a multi-pass transformation process where an agent transforms the reference into an intermediate representation (IR) to capture data dependencies and persistent-state scopes, validates this IR via a sequential interpreter, and performs static analyses to identify candidate transformations. Then, the agent iteratively implements, verifies, and benchmarks each candidate under a measurement-gated optimization loop to produce effective deployments that span different hardware budgets. Across various applications, including video world models and multimodal LLMs, FlashRT converts reference implementations into highly efficient deployments, delivering up to ∼70x latency reduction and 2.8x throughput improvement on NVIDIA B200 GPUs. On AMD MI355X GPUs, FlashRT matches the peak latency reduction while increasing peak throughput improvement to 3.6x, demonstrating that agent-driven optimization can be more scalable on platforms with less mature expert optimization. In fact, for Qwen3-Omni text-to-audio inference, FlashRT reduces response latency by 65% compared to the expert vLLM-Omni implementation on AMD MI355X.

   @misc{agarwal2026flashrt,
     author = {Agarwal, Krish and Chen, Zhuoming and Qin, Yanyuan and Gu, Zhenyu and Rudra, Atri and Chen, Beidi},
     title = {FlashRT: Agent Harness for Guiding Agents to Deploy Real-Time Multimodal Applications},
     year = {2026},
     eprint = {2607.18171},
     archiveprefix = {arXiv},
     primaryclass = {cs.LG},
     url = {https://arxiv.org/abs/2607.18171},
   }

2025

1. arXiv

   

   Efficient MoE Serving in the Memory-Bound Regime: Balance Activated Experts, Not Tokens

   Yanpeng Yu^*, Haiyue Ma^*, Krish Agarwal, Nicolai Oswald, Qijing Huang, Hugo Linsenmaier, Chunhui Mei, Ritchie Zhao, and 5 more authors

   Preprint , 2025

   Abs arXiv Bib

   Expert Parallelism (EP) permits Mixture of Experts (MoE) models to scale beyond a single GPU. To address load imbalance across GPUs in EP, existing approaches aim to balance the number of tokens each GPU processes. Surprisingly, we find that this objective degrades performance rather than improving it when processing is memory-bound - a common occurrence in MoE serving, especially in the decode phase. Our analysis reveals that balancing the number of tokens processed per GPU increases the number of activated experts, exacerbating memory pressure in the memory-bound regime. We propose Minimum Expert Token ROuting, a novel token-routing algorithm for high-performance expert-parallel MoE serving in the memory-bound regime that balances the number of activated experts per GPU rather than token counts. METRO achieves near-optimal routing quality with minimal computational overhead by jointly optimizing algorithmic efficiency and leveraging the GPU’s parallel processing power. To guarantee routing quality, METRO also employs a novel allGather scheme to gather global top-k knowledge, which has minimal overhead compared to conventional allToAll. Our evaluation of METRO against EPLB on both real systems (vLLM over 8 A100 GPUs) and a proprietary simulator (8-16 B200 GPUs) shows that METRO reduces decode latency by 11 - 22%, and total token throughput by 3 - 21% for Qwen3 and DeepSeek-V3 serving, where prefill and decode phases are co-deployed. In addition, by trading latency headroom for throughput, METRO improves decode throughput by up to 4.11x over EPLB at a fixed decode SLO.

   @misc{yu2025metro,
     author = {Yu, Yanpeng and Ma, Haiyue and Agarwal, Krish and Oswald, Nicolai and Huang, Qijing and Linsenmaier, Hugo and Mei, Chunhui and Zhao, Ritchie and Borkar, Ritika and Darvish Rouhani, Bita and Nellans, David and Krashinsky, Ronny and Khandelwal, Anurag},
     title = {Efficient MoE Serving in the Memory-Bound Regime: Balance Activated Experts, Not Tokens},
     year = {2025},
     eprint = {2512.09277},
     archiveprefix = {arXiv},
     primaryclass = {cs.DC},
     url = {https://arxiv.org/abs/2512.09277},
   }

2. IROS

   

   L3M+P: Lifelong Planning with Large Language Models

   Krish Agarwal^*, Yuqian Jiang^*, Jiaheng Hu, Bo Liu, and Peter Stone

   In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), Hangzhou, China, Oct 2025

   Abs arXiv Bib Code

   By combining classical planning methods with large language models (LLMs), recent research such as LLM+P has enabled agents to plan for general tasks given in natural language. However, scaling these methods to general-purpose service robots remains challenging: (1) classical planning algorithms generally require a detailed and consistent specification of the environment, which is not always readily available; and (2) existing frameworks mainly focus on isolated planning tasks, whereas robots are often meant to serve in long-term continuous deployments, and therefore must maintain a dynamic memory of the environment which can be updated with multi-modal inputs and extracted as planning knowledge for future tasks. To address these two issues, this paper introduces L3M+P (Lifelong LLM+P), a framework that uses an external knowledge graph as a representation of the world state. The graph can be updated from multiple sources of information, including sensory input and natural language interactions with humans. L3M+P enforces rules for the expected format of the absolute world state graph to maintain consistency between graph updates. At planning time, given a natural language description of a task, L3M+P retrieves context from the knowledge graph and generates a problem definition for classical planners. Evaluated on household robot simulators and on a real-world service robot, L3M+P achieves significant improvement over baseline methods both on accurately registering natural language state changes and on correctly generating plans, thanks to the knowledge graph retrieval and verification.

   @inproceedings{agarwal2025l3mp,
     author = {Agarwal, Krish and Jiang, Yuqian and Hu, Jiaheng and Liu, Bo and Stone, Peter},
     title = {L3M+P: Lifelong Planning with Large Language Models},
     booktitle = {IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
     year = {2025},
     month = oct,
     location = {Hangzhou, China},
   }

3. Thesis

   

   Tensor Core Accelerated Kernels for Hadamard and Fourier Masked Attention

   Krish Agarwal

   Undergraduate Honors Thesis, The University of Texas at Austin , May 2025

   Abs Bib PDF Code

   Softmax attention incurs quadratic time complexity in sequence length, making it the primary bottleneck in transformer training and inference for long sequences. Structured Masked Attention (SMA) offers a sub-quadratic alternative by removing softmax normalization and applying a structured matrix as the attention mask. Dropping softmax enables reformulating elementwise masking as a matrix multiplication, allowing the use of fast transform algorithms to compute the operation in sub-quadratic time when the mask is structured. Implementing SMA for a given structured mask is straightforward if a standalone fast transform implementation is available. However, standard implementations of these transforms often achieve low hardware utilization on GPUs because they do not leverage high-throughput matrix-multiply units like Tensor Cores. Furthermore, using a separate kernel to perform the transform prevents operator fusion, missing critical opportunities to reuse data in on-chip SRAM. In this work, we address these challenges for SMA when using Walsh-Hadamard and discrete Fourier transform matrices, both of which admit O(n log n) algorithms for computing matrix-vector products. Specifically, we design fused, Tensor Core accelerated kernels to efficiently perform inference and backpropagation for Hadamard and Fourier masked attention while retaining sub-quadratic time complexity.

   @techreport{agarwal2025thesis,
     author = {Agarwal, Krish},
     title = {Tensor Core Accelerated Kernels for Hadamard and Fourier Masked Attention},
     institution = {The University of Texas at Austin},
     type = {Undergraduate Honors Thesis},
     number = {HR-25-10},
     year = {2025},
     month = may,
   }

2024

1. arXiv

   

   HadaCore: Tensor Core Accelerated Hadamard Transform Kernel

   Krish Agarwal^*, Rishi Astra^*, Adnan Hoque, Mudhakar Srivatsa, Raghu Ganti, Less Wright, and Sijia Chen

   Preprint (featured on the PyTorch Blog) , 2024

   Abs arXiv Bib Blog Code

   We present HadaCore, a modified Fast Walsh-Hadamard Transform (FWHT) algorithm optimized for the Tensor Cores present in modern GPU hardware. HadaCore follows the recursive structure of the original FWHT algorithm, achieving the same asymptotic runtime complexity but leveraging a hardware-aware work decomposition that benefits from Tensor Core acceleration. This reduces bottlenecks from compute and data exchange. On Nvidia A100 and H100 GPUs, HadaCore achieves speedups of 1.1-1.4x and 1.0-1.3x, with a peak gain of 3.5x and 3.6x respectively, when compared to the existing state-of-the-art implementation of the original algorithm. We also show that when using FP16 or BF16, our implementation is numerically accurate, enabling comparable accuracy on MMLU benchmarks when used in an end-to-end Llama3 inference run with quantized (FP8) attention.

   @misc{agarwal2024hadacore,
     author = {Agarwal, Krish and Astra, Rishi and Hoque, Adnan and Srivatsa, Mudhakar and Ganti, Raghu and Wright, Less and Chen, Sijia},
     title = {HadaCore: Tensor Core Accelerated Hadamard Transform Kernel},
     year = {2024},
     eprint = {2412.08832},
     archiveprefix = {arXiv},
     primaryclass = {cs.DC},
     url = {https://arxiv.org/abs/2412.08832},
   }