Characterizing State Space Model and Hybrid Language Model Performance with Long Context

TL;DR

This paper benchmarks and compares the performance of Transformer, State Space Models (SSMs), and hybrid models for long-context inference on consumer GPUs, revealing SSMs' advantages at very long sequences due to their linear complexity and reduced memory footprint.

Contribution

It provides the first comprehensive benchmarking of these models on consumer GPUs, highlighting SSMs' suitability for on-device long-context AI applications.

Findings

SSMs outperform Transformers at very long sequences (~57K tokens).

Transformers are faster at short sequences (<8K tokens).

Custom SSM kernels dominate inference runtime on edge platforms.

Abstract

Emerging applications such as AR are driving demands for machine intelligence capable of processing continuous and/or long-context inputs on local devices. However, currently dominant models based on Transformer architecture suffers from the quadratic computational and memory overhead, which hinders applications required to process long contexts. This has spurred a paradigm shift towards new architectures like State Space Models (SSMs) and SSM-Transformer hybrid models, which provide near-linear scaling. The near-linear scaling enabled efficient handling of millions of tokens while delivering high performance in recent studies. Although such works present promising results, their workload characteristics in terms of computational performance and hardware resource requirements are not yet thoroughly explored, which limits our understanding of their implications to the system level optimizations. To address this gap, we present a comprehensive, compara-ive benchmarking of carefully selected Transformers, SSMs, and hybrid models specifically for long-context inference on consumer and embedded GPUs. Our analysis shows that SSMs are well-suited for on-device AI on consumer and embedded GPUs for long context inferences. While Transformers are up to 1.9x faster at short sequences (<8K tokens), SSMs demonstrate a dramatic performance inversion, becoming up to 4x faster at very long contexts (~57K tokens), thanks to their linear computational complexity and ~64% reduced memory footrprint. Our operator-level analysis reveals that custom SSM kernels like selective scan despite being hardware-aware to minimize memory IO, dominate the inference runtime on edge platforms, accounting for over 55% of latency due to their sequential, element-wise nature. SSM-Scope is open-sourced at https://github.com/sapmitra/ssm-scope