Mamba-2 e State Space Duality (SSD): Eliminando o Gargalo de KV-Cache em Produção

O Mamba-2 unifica Modelos de Espaço de Estados e Atenção Linear em uma formulação matricial otimizada para Tensor Cores. A teoria de State Space Duality (SSD) permite treinar modelos com complexidade linear sem abrir mão da alta intensidade aritmética das GPUs modernas.

Em cargas de inferência de contexto longo (128k a 1M tokens), a alocação dinâmica de KV-Cache nos Transformers tradicionais consome dezenas de gigabytes de VRAM por usuário. Essa sobrecarga de memória reduz o batch size e encarece o custo por milhão de tokens gerados em servidores de produção.

O custo oculto do KV-Cache em contextos longos

O crescimento linear do KV-Cache limita a capacidade de atendimento simultâneo de servidores de inferência em contextos extensos. Cada novo token decodificado exige consultar chaves e valores anteriores, transformando a geração em uma operação estritamente limitada pela largura de banda de memória (HBM).

Em uma arquitetura Transformer padrão de 70 bilhões de parâmetros (como Llama 3.1 com GQA), o volume de memória exigido para armazenar o histórico de contexto é governado pela equação:

Memória_KV_por_token = 2 × n_camadas × n_heads_kv × d_head × bytes_por_elemento

Substituindo os parâmetros reais em precisão FP16 (80 camadas × 8 cabeças KV × dimensão 128 × 2 bytes):

Memória_KV_por_token = 2 × 80 × 8 × 128 × 2 bytes = 327.680 bytes ≈ 320 KB/token

Quando um único usuário processa uma janela de 128.000 tokens, o consumo de VRAM dedicado exclusivamente ao KV-Cache atinge:

320 KB × 128.000 tokens ≈ 40,96 GB de VRAM

Em um cluster servindo centenas de requisições simultâneas, a memória da GPU esgota antes mesmo de os núcleos de computação atingirem 40% de utilização.

Como o Mamba-1 alcançou complexidade linear

O Mamba-1 resolveu a dependência de memória ao substituir a atenção quadrática por um Modelo de Espaço de Estados Seletivo (Selective SSM / S6). Em vez de reter tensores de cada token passado, o modelo comprime o histórico sequencial em um estado oculto fixo h_t ∈ ℝ^N.

A formulação matemática do sistema contínuo e sua discretização via Zero-Order Hold (ZOH) é estruturada em:

h_t = exp(Δ_t · A) · h_{t-1} + (Δ_t · B_t) · x_t
y_t = C_t · h_t + D · x_t

Onde as matrizes B_t, C_t e o passo temporal Δ_t tornam-se funções dinâmicas do token atual x_t:

# Mecanismo de Seleção Dinâmica (S6) no Mamba-1
import torch
import torch.nn as nn

class SelectiveSSM(nn.Module):
    def __init__(self, d_model: int, d_state: int = 16):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.x_proj = nn.Linear(d_model, d_state + d_state + 1, bias=False)
        self.dt_proj = nn.Linear(1, d_model, bias=True)
        self.A_log = nn.Parameter(torch.log(torch.randn(d_model, d_state).abs()))
        self.D = nn.Parameter(torch.ones(d_model))

    def forward_step(self, x_t: torch.Tensor, h_prev: torch.Tensor):
        # Projeção dependente do conteúdo de entrada
        proj = self.x_proj(x_t)
        b_t = proj[:, :self.d_state]
        c_t = proj[:, self.d_state:2 * self.d_state]
        delta_raw = proj[:, 2 * self.d_state:]
        
        delta = torch.nn.functional.softplus(self.dt_proj(delta_raw))
        a_mat = -torch.exp(self.A_log)
        
        # Discretização ZOH
        a_bar = torch.exp(delta.unsqueeze(-1) * a_mat)
        b_bar = delta.unsqueeze(-1) * b_t.unsqueeze(1)
        
        # Atualização de estado O(1) sem KV-Cache
        h_t = a_bar * h_prev + b_bar * x_t.unsqueeze(-1)
        y_t = (h_t * c_t.unsqueeze(1)).sum(dim=-1) + self.D * x_t
        return y_t, h_t

Na inferência autoregressiva, o estado h_t ocupa espaço constante O(1). O modelo consome exatamente a mesma quantidade de memória gerando o token 10 ou o token 100.000.

O gargalo de hardware no scan associativo do Mamba-1

O Mamba-1 encontrou um limite de eficiência em GPUs devido à baixa intensidade aritmética do scan associativo paralelo. O algoritmo dependia de núcleos CUDA convencionais e realizava muitas leituras de memória para poucas operações de ponto flutuante.

Os aceleradores modernos (NVIDIA Hopper e Blackwell) atingem pico de desempenho através de Tensor Cores dedicados à multiplicação densa de matrizes 16 × 16 ou maiores. O scan paralelo do Mamba-1 não conseguia alimentar os Tensor Cores, ficando restrito pela largura de banda da SRAM.

Além disso, a dimensão de estado d_state precisava ser mantida pequena (d_state = 16) para evitar saturação de registradores. Isso limitava a capacidade do modelo de memorizar detalhes factuais densos em tarefas de recuperação em contexto.

A teoria do State Space Duality (SSD) no Mamba-2

O Mamba-2 elimina o gargalo de hardware ao comprovar que Modelos de Espaço de Estados e Atenção Linear são duas faces da mesma estrutura matemática. Tri Dao e Albert Gu demonstraram que, restringindo a matriz de transição A a uma estrutura escalar diagonal, o cálculo vira uma multiplicação de matrizes 1-semiseparável.

Seja uma atenção linear causal com máscara semi-separável M ∈ ℝ^{N × N}:

Y = (M ⊙ (Q · K^T)) · V

Onde M_{i,j} representa a taxa de decaimento acumulada entre as posições j e i:

M_{i,j} = exp(∑_{k=j+1}^{i} Δ_k · a_k) para i ≥ j; 0 caso contrário

Essa relação matemática unifica duas formas complementares de computação:

1. Visão Recorrente: Manter o estado h_t de tamanho fixo para geração token a token com latência mínima e zero alocação de KV-Cache.
2. Visão Matricial Dual: Processar sequências inteiras no treinamento e prefill via multiplicação densa de matrizes (GEMM), aproveitando 100% dos Tensor Cores da GPU.

       Dualidade de Espaço de Estados (SSD)
       
  [Prefill / Treinamento]         [Geração / Decode]
    Visão Matricial GEMM            Visão Recorrente
   (Tensor Cores / WGMMA)            (Estado O(1))
            │                              │
            ▼                              ▼
┌───────────────────────┐      ┌───────────────────────┐
│ Multiplicação Bloco Q │      │ Atualização de Estado │
│   Y = (M ⊙ C B^T) X   │      │ h_t = A h_{t-1} + B x │
└───────────────────────┘      └───────────────────────┘

O algoritmo de decomposição em blocos do SSD

O algoritmo do Mamba-2 divide a sequência em blocos de tamanho fixo Q (tipicamente 64 tokens) para equilibrar computação densa e propagação de estado. O cálculo ocorre em dois níveis paralelos:

Dentro de cada bloco de 64 tokens, o processamento ocorre via multiplicação matricial pura em Tensor Cores. O kernel computa a interação interna com intensidade aritmética máxima.

Entre blocos adjacentes, o estado final de cada bloco h_b é transmitido para o próximo através de um scan recorrente de baixa dimensionalidade. Essa decomposição eleva a utilização da GPU (MFU) de 35% para até 75% em aceleradores A100 e H100.

Com a aceleração em Tensor Cores, a dimensão de estado d_state pode ser expandida de 16 para 64, 128 ou 256 sem perda de vazão, ampliando a capacidade de retenção do modelo.

Arquiteturas híbridas em produção: o caso Jamba

Modelos baseados exclusivamente em SSMs apresentam dificuldades em tarefas de recuperação associativa estrita (Associative Recall / cópia exata de tokens vistos no início do texto). Em contrapartida, Transformers puros tornam-se proibitivamente caros em VRAM.

A solução adotada em produção consiste em intercalar blocos Mamba-2 com blocos de Atenção e camadas Mixture of Experts (MoE), como implementado na família Jamba da AI21 Labs e no Nemotron da NVIDIA.

Essa combinação híbrida reduz o consumo de KV-Cache em cerca de 87,5%, viabilizando contextos de 256k tokens em uma única placa NVIDIA H100 de 80GB.

Comparativo: Softmax Attention vs Mamba-1 vs Mamba-2 vs Híbrido

A tabela a seguir compara o comportamento assintótico, requisitos de memória e eficiência de hardware entre as principais arquiteturas de modelos de linguagem:

Implementando Mamba-2 em pipelines de inferência

Integrar Mamba-2 em ambientes de produção exige runtimes compatíveis com kernels de scan blocado e decodificação contínua. Frameworks como vLLM e SGLang suportam modelos híbridos através de despachantes especializados.

No ecossistema da Produtora MaxVision, a esteira de agentes e sistemas RAG de alta densidade emprega modelos híbridos para reduzir custos de infraestrutura em tarefas de análise documental extensa, síntese de repositórios de código e logs de telemetria.

O pipeline de execução em PyTorch com o pacote oficial mamba-ssm opera de forma direta:

import torch
from mamba_ssm import Mamba2

# Configuração do bloco Mamba-2 com State Space Duality
batch_size = 4
seq_len = 8192
d_model = 2048
d_state = 128   # Dimensão de estado expandida graças ao SSD
headdim = 64

device = "cuda" if torch.cuda.is_available() else "cpu"

layer = Mamba2(
    d_model=d_model,
    d_state=d_state,
    d_conv=4,
    expand=2,
    headdim=headdim
).to(device)

# Entrada de tokens em BF16
input_tensor = torch.randn(batch_size, seq_len, d_model, dtype=torch.bfloat16, device=device)

# Prefill acelerado via Tensor Cores (GEMM intra-bloco + scan inter-bloco)
output = layer(input_tensor)
print("Saída do bloco Mamba-2 processada com sucesso. Dimensões:", output.shape)

Fontes e Referências Primárias

• arXiv:2405.21060 — Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Dao & Gu, ICML 2024).
• arXiv:2312.00752 — Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu & Dao, 2023).
• arXiv:2403.19887 — Jamba: A Hybrid Transformer-Mamba Language Model (AI21 Labs, 2024).
• arXiv:2111.00396 — Efficiently Modeling Long Sequences with Structured State Spaces (Gu, Goel & Ré, ICLR 2022).
• ICML 2020 — Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (Katharopoulos et al., 2020).
• GitHub: state-spaces/mamba — Repositório oficial com kernels CUDA e implementações de Mamba-1 e Mamba-2.