arXiv:2602.21144cs.DCcs.LG2026-02

通过优化张量并行,大幅提升长序列状态空间模型的多卡推理速度。

Scaling State-Space Models on Multiple GPUs with Tensor Parallelism

  • 设计高效通信的张量并行方案,保持状态更新局部性
  • 2卡提升1.6-2.1倍吞吐,4卡达2.6-4.0倍,长序列效果更显著
  • 适合部署大模型长文本任务的工程师和研究者

选择性状态空间模型(SSMs)正快速成为大语言模型的核心架构,尤其适用于长上下文任务。然而在部署中,其推理性能常受限于单个GPU的内存容量、带宽与延迟,促使多GPU执行变得必要。尽管张量并行(TP)广泛用于加速Transformer推理,但将其应用于选择性SSM模块颇具挑战,因其混合作用层耦合了大型投影与序列级递归状态更新及依赖局部性的本地混合操作,效率依赖于维持局部性并避免关键路径中的同步。本文提出一种面向选择性SSM推理的高效通信张量并行设计,解决三大工程难题:通过跨预填充与解码的状态缓存实现TTFT优化;将混合作用层的参数张量分块,使递归更新保持局部性同时最小化通信开销;采用量化AllReduce降低TP聚合开销。我们在NVIDIA A6000与A100集群上对三种代表性基于SSM的LLM——Mamba、Falcon-Mamba与Zamba——进行了评估。实验表明,张量并行的SSM推理带来显著吞吐提升:在2张GPU上,Mamba的批处理请求吞吐提升约1.6-2.1倍;在4张GPU上提升2.6-4.0倍,且长上下文场景下优势更明显;量化AllReduce进一步带来10-18%的吞吐增益,有效降低同步带宽开销。

原文摘要 · Abstract (English)

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.

状态空间模型张量并行多卡推理长序列

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。