activity
20192026
most citedCheckmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

37 citations · 43 across the 9 of their papers we have counts for

collaborators

10 papers

cs.CY2026

Open Technical Problems in Open-Weight AI Model Risk Management

Stephen Casper, Kyle O'Brien, Shayne Longpre +19

Frontier AI models with openly available weights are steadily becoming more powerful and widely adopted. However, compared to proprietary models, open-weight models pose different…

cs.LG2025

FlashFormer: Whole-Model Kernels for Efficient Low-Batch Inference

Aniruddha Nrusimha, William Brandon, Mayank Mishra +4

The size and compute characteristics of modern large language models have led to an increased interest in developing specialized kernels tailored for particular training and infere…

cs.LG2024★ 2 cited

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

William Brandon, Mayank Mishra, Aniruddha Nrusimha +2

Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to…

cs.LG2024

Mitigating the Impact of Outlier Channels for Language Model Quantization with Activation Regularization

Aniruddha Nrusimha, Mayank Mishra, Naigang Wang +3

We consider the problem of accurate quantization for language models, where both the weights and activations are uniformly quantized to 4 bits per parameter, the lowest bitwidth fo…

cs.LG2024★ 1 cited

Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding

Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha +3

To combat the memory bandwidth-bound nature of autoregressive LLM inference, previous research has proposed the speculative decoding frame-work. To perform speculative decoding, a…

cs.LG2023★ 1 cited

Striped Attention: Faster Ring Attention for Causal Transformers

William Brandon, Aniruddha Nrusimha, Kevin Qian +4

To help address the growing demand for ever-longer sequence lengths in transformer models, Liu et al. recently proposed Ring Attention, an exact attention algorithm capable of over…

cs.CL2023★ 2 cited

Towards Verifiable Text Generation with Symbolic References

Lucas Torroba Hennigen, Shannon Shen, Aniruddha Nrusimha +3

LLMs are vulnerable to hallucinations, and thus their outputs generally require laborious human verification for high-stakes applications. To this end, we propose symbolically grou…

cs.CV2021

Self-Supervised Pretraining Improves Self-Supervised Pretraining

Colorado J. Reed, Xiangyu Yue, Ani Nrusimha +9

While self-supervised pretraining has proven beneficial for many computer vision tasks, it requires expensive and lengthy computation, large amounts of data, and is sensitive to da…

eess.AS2021

Integer-only Zero-shot Quantization for Efficient Speech Recognition

Sehoon Kim, Amir Gholami, Zhewei Yao +7

End-to-end neural network models achieve improved performance on various automatic speech recognition (ASR) tasks. However, these models perform poorly on edge hardware due to larg…

cs.LG2019★ 37 cited

Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Paras Jain, Ajay Jain, Aniruddha Nrusimha +5

We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing stra…