# PackMamba: Efficient Processing of Variable-Length Sequences in Mamba training

arXiv: https://arxiv.org/abs/2408.03865 · HTML: https://mlsanity.com/paper/2408.03865/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Core SSM architectures for language, one of 10 directions in State-space sequence models (S4, Mamba).
- Builds on 12 of its 17 references that are in the map.
- First posted on arXiv 7 Aug 2024.

## Placed in

- State-space sequence models (S4, Mamba) / [Core SSM architectures for language](https://mlsanity.com/ssm/core-ssm-architectures-for-language/)

## Builds on (most cited first, 12 of 12)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Free-Form Image Inpainting with Gated Convolution](https://mlsanity.com/paper/1806.03589/) (2018)
- [InternLM2 Technical Report](https://mlsanity.com/paper/2403.17297/) (2024)
- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Recurrent Models of Visual Attention](https://mlsanity.com/paper/1406.6247/) (2014)
- [TurboTransformers: An Efficient GPU Serving System For Transformer Models](https://mlsanity.com/paper/2010.05680/) (2020)
- [ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs](https://mlsanity.com/paper/2210.03052/) (2022)
- [Boosting Distributed Training Performance of the Unpadded BERT Model](https://mlsanity.com/paper/2208.08124/) (2022)

## Built on by (most cited first, 1 of 1)

- [RubiConv -- Efficient Boundary-Respecting Convolutions](https://mlsanity.com/paper/2605.08451/) (2026)

## Abstract (arXiv)

With the evolution of large language models, traditional Transformer models become computationally demanding for lengthy sequences due to the quadratic growth in computation with respect to the sequence length. Mamba, emerging as a groundbreaking architecture in the field of generative AI, demonstrates remarkable proficiency in handling elongated sequences with reduced computational and memory complexity. Nevertheless, the existing training framework of Mamba presents inefficiency with variable-length sequence inputs. Either single-sequence training results in low GPU utilization, or batched processing of variable-length sequences to a maximum length incurs considerable memory and computational overhead. To address this problem, we analyze the performance of bottleneck operators in Mamba under diverse tensor shapes and proposed PackMamba, a high-throughput Mamba that efficiently handles variable-length sequences. Diving deep into state-space models (SSMs), we modify the parallel operators to avoid passing information between individual sequences while maintaining high performance. Experimental results on an NVIDIA A100 GPU demonstrate throughput exceeding the baseline single-sequence processing scheme: 3.06x speedup on the 1.4B model and 2.62x on the 2.8B model.