How Deep Learning Computes

SSM 추론과 하이브리드 모델

이 경로는 SSM을 “Attention을 대체하는 새 구조”로만 보지 않습니다. 핵심 질문은 이것입니다.

Transformer serving은 왜 KV cache를 중심으로 짜였고,
SSM/Mamba는 그 가정을 어디서 바꾸는가?

먼저 Transformer block을 두 개의 mixer로 보는 관점을 잡고, KV cache가 왜 vLLM 같은 추론 엔진의 중심 자원이 되었는지 복습합니다.

그다음 SSM의 recurrent state 관점으로 넘어갑니다. SSM은 과거 token마다 K/V를 쌓는 대신, 일정한 크기의 state를 갱신합니다. 이 차이가 prefill/decode, prefix caching, batching, scheduler, cache memory layout을 어떻게 바꾸는지 봅니다.

마지막에는 순수 SSM보다 실전에서 더 자주 보이는 hybrid 모델을 봅니다. Jamba, Zamba, Falcon-H1, Nemotron-H 같은 모델은 attention layer와 SSM/Mamba layer를 섞습니다. 이때 serving system은 KV cache와 SSM state cache를 동시에 관리해야 합니다.

이 path의 산출물은 Mamba 수식 암기가 아닙니다. 목표는 vLLM의 Mamba/Hybrid 관련 코드를 읽을 때 다음 질문에 답하는 것입니다.

왜 MambaSpec이 별도로 필요한가?
왜 mamba_cache_mode가 있는가?
왜 hybrid cache manager가 필요한가?
SSM은 KV cache 문제를 없애는가, 아니면 다른 cache 문제로 바꾸는가?
  1. Transformer Block은 두 Mixer다 — Attention은 token mixing, MLP/FFN은 feature transformation이라는 관점으로 block을 본다.
  2. KV Cache — Decode에서 과거 token의 Key/Value를 저장해 재계산을 줄이고, 그 대가로 어떤 메모리 병목이 생기는지 이해한다.
  3. SSM 직관 — State Space Model을 긴 sequence를 고정 크기 state로 요약하며 갱신하는 모델로 이해한다.
  4. Mamba와 Selective State — Mamba가 입력에 따라 state update를 선택적으로 바꾸는 이유와, 이 점이 S4류 SSM과 어떻게 다른지 이해한다.
  5. SSM State vs KV Cache — Transformer의 KV cache와 SSM/Mamba의 state cache가 추론 시스템에서 어떻게 다른 자원인지 비교한다.
  6. Mamba Prefill과 Decode — Mamba/SSM에서 prefill은 state를 만들어 가는 구간이고 decode는 state를 갱신하는 구간이라는 관점으로 본다.
  7. vLLM Mamba Cache Spec — vLLM이 Mamba/SSM state를 일반 KV cache와 다르게 표현하기 위해 MambaSpec과 mamba cache 설정을 두는 이유를 이해한다.
  8. Mamba Selective State Update Kernel — Mamba decode에서 selective state update가 왜 별도 커널 관심사가 되는지 serving 관점으로 이해한다.
  9. 왜 Hybrid SSM-Transformer인가 — 실전 모델들이 순수 SSM이나 순수 Transformer 대신 attention layer와 SSM/Mamba layer를 섞는 이유를 이해한다.
  10. Hybrid Cache Manager — Attention cache와 Mamba state cache가 한 모델 안에 공존할 때 serving engine이 풀어야 하는 문제를 정리한다.