cs.DCFeb 24, 2026

Scaling State-Space Models on Multiple GPUs with Tensor Parallelism

Authors: Anurag Dutt, Nimit Shah, Hazem Masarani, Anshul Gandhi

Organizations: Stony Brook University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Selective state space models (SSMs) have rapidly become a compelling backbone for large language models, especially for long-context workloads. Yet in deployment, their inference performance is often bounded by the memory capacity, bandwidth, and latency limits of a single GPU, making multi-GPU execution increasingly necessary. Although tensor parallelism (TP) is widely used to scale Transformer inference, applying it to selective SSM blocks is non-trivial because the SSM mixer couples large projections with a sequence-wise recurrent state update and local mixing whose efficiency depends on preserving locality and avoiding synchronization in the critical path. This paper presents a communication-efficient TP design for selective SSM inference that addresses three practical engineering challenges: enabling TTFT improvements via an SSM state cache across prefill and decode, partitioning the mixer's packed parameter tensor so that recurrent updates remain local while minimizing communication, and reducing TP aggregation overhead with quantized AllReduce. We evaluate on three representative SSM-based LLMs spanning pure-SSM and hybrid architectures - Mamba, Falcon-Mamba, and Zamba - on NVIDIA A6000 and A100 clusters. Our experiments show substantial throughput gains from tensor-parallel SSM inference, improving batch-request throughput by ~1.6-2.1x on 2 GPUs and ~2.6-4.0x on 4 GPUs for Mamba, with the largest benefits at long context lengths, and achieving a further ~10-18% throughput improvement from quantized all-reduce by lowering synchronization bandwidth overhead.

Figures & tables

Fig. 1 : Mamba-style SSM mixer block. The input projection produces a packed tensor that is split into an SSM path and a gating path. The SSM path applies a channel-separable 1D convolution, projects activations into token-dependent SSM fields ( Δ\Delta , BB , CC ) together with per-channel parameters ( AA , DD ), performs the state update over the sequence, gates the result, and applies an output projection back to the residual stream.
Fig. 2 : Illustration of our tensor-parallel inference implementation for Mamba, consisting of our four key design components: (1) SSM cache, (2) channel splitter and depthwise convolution, (3) packed parameter handling, and (4) AllReduce quantization.
Fig. 3 : Maximum input sequence length possible at fixed (256) batch size under our TP design and under DP and 1x.
Fig. 4 : Throughput gains afforded by our tensor-parallel inference design (for 2-GPU, 4-GPU) for Mamba, Mamba-2, Falcon-Mamba, and Zamba, compared to no parallelism (“1x”, first row) and compared to data parallelism (“DP”, second row) on the A6000 cluster.
Fig. 5 : Throughput gains afforded by our tensor-parallel inference design (for 2-GPU, 4-GPU) for Mamba, Mamba-2, Falcon-Mamba, and Zamba, compared to no parallelism (“1x”, first row) and compared to data parallelism (“DP”, second row) on the A100 cluster.
Fig. 6 : TTFT gain (vs. DP) afforded by our TP on 4 GPUs.
Fig. 7 : TPOT gain (vs. DP) afforded by our TP on 4 GPUs.
Fig. 8 : GPU and memory utilization during TTFT and TPOT.
ModelMetricValue (%)
MambaTop-1 (argmax)98.81
MambaTop-5 overlap (Unordered)99.03
MambaTop-5 (Ordered)89.01
TABLE I: Accuracy impact of quantization for Mamba.
Fig. 9 : Throughput gains afforded by quantizing AllReduce in our TP inference design for Mamba as a function of output sequence length.
Fig. 10 : Ablation results for Mamba (256 input/output length).
More in this topic

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

Explore similar work

CardsList
Oct 9, 2026

ReplaySSM: Cache SSM Inputs, Not State

State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models. This eager state management causes large state I/O with low arithmetic intensity, leaves speculative decoding no cheap rollback, and feeds state quantization error into subsequent recurrent updates. We introduce ReplaySSM, which unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs. This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic. Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states. The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates. We implement ReplaySSM in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, it accelerates standard decoding by 1.201.20--1.47×1.47\times over a tuned vLLM baseline across four hybrid models, and delivers about 1.8×1.8\times the throughput of vLLM's speculative decoding on the two 120B-class models. Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.

Ze-Wei Liou, Harris Nover, Amey Naik +4

Princeton University · NVIDIA

Inference AccelerationLLM Inference AccelerationState Space Models
AskPDF
Jun 15, 2026

Long-Context Modeling via GSS-Transformer Hybrid Architecture with Learnable Mixing

Modeling long-range dependencies remains a central challenge in natural language processing. Transformer architectures achieve strong performance via self-attention but scale quadratically (O(N2)O(N^2)) with sequence length, while State Space Models (SSMs) scale linearly (O(N)O(N)) but suffer from a selective recall bottleneck, struggling to retrieve precise information from compressed states. This creates a fundamental tradeoff between efficiency and perplexity. To tackle these challenges, we propose the Parallel Hybrid Architecture (PHA), which runs Gated State Spaces (GSS), Grouped Query Attention (GQA), and Feed-Forward Networks (FFNs) as independent parallel branches fused by a learnable mixing mechanism. Instead of forcing SSMs to approximate attention or serializing the two paradigms, PHA allows each branch to specialize: GSS captures global context, while attention performs selective retrieval, with FFN providing complementary processing. On WikiText-103, PHA achieves 16.51 PPL at 125M parameters, outperforming Hedgehog (16.70) and H3-125M (23.70). Scaling to 180M parameters yields 16.42 PPL, which gives comparable results with the pure attention baseline while delivering 24% higher throughput and up to 40% lower memory usage at long contexts. On OpenWebText, our 125M model achieves 19.72 PPL, outperforming standard Transformers (20.60) and GSS hybrid baselines (19.80). These results demonstrate that separating sequence modeling paradigms into parallel specialists enables Transformer-level perplexity with substantially improved efficiency for long-context language modeling.

Kuzey Torlak, Hüseyin Arda Arslan, Anıl Dervişoğlu +2

Kadıköy Anadolu High School · Politecnico di Torino · Istanbul Technical University +2

Long-Context Language ModelingHybrid Transformer-State Space Model ArchitecturesState Space Models
AskPDF
Apr 29, 2026

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

We present tensor and sequence parallelism (TSP), a parallel execution strategy that folds tensor parallelism and sequence parallelism onto a single device axis. In conventional multi-dimensional parallelism layouts, tensor parallelism (TP) shards model weights while sequence parallelism (SP) shards tokens, reducing per-device parameter or activation memory, respectively. Traditionally, each scheme is assigned its own mesh dimension. TSP instead assigns each rank both a weight shard and a sequence shard, reducing both parameter and activation memory along the same device axis. We implement this design with two runtime schedules. For attention, ranks iterate over broadcast parameter shards and reconstruct context through a sequence-wise key/value exchange. For gated MLPs, weight shards circulate in a ring while partial outputs accumulate locally. By sharding both weights and activations across the same devices, TSP trades additional communication volume for reduced memory overhead. We provide a theoretical communication and memory analysis, describe our implementation of TSP attention and gated MLP blocks, and benchmark TSP against TP, SP, and TP+SP. These results position TSP as a hardware-aware alternative for long-context and memory-constrained model training, and as a viable axis of parallelism in concert with existing parallelism schemes such as pipeline and expert parallelism for dense and mixture-of-expert models.

Vasu Shyam, Anna Golubeva, Quentin Anthony

Zyphra

Transformer InferenceTransformerTensor ParallelismMemory-Efficient Inference
AskPDF