SSM State vs KV Cache

Transformer serving의 중심 자원은 KV cache다.

Mamba/SSM serving의 중심 자원은 state cache다.

둘 다 “과거를 다시 쓰기 위해 저장한다”는 점은 같지만, 저장하는 것과 커지는 방식이 다르다.

KV cache

Transformer attention은 decode step마다 새 query가 과거 모든 key/value를 봐야 한다.

그래서 과거 token마다 K와 V를 저장한다.

KV cache size ~= layers * tokens * kv_heads * head_dim * 2

context가 길어지면 cache도 길어진다.

SSM state cache

SSM은 과거 token 전체를 직접 다시 보지 않고, recurrent state를 갱신한다.

Mamba 계열에서는 보통 두 종류의 state를 생각해야 한다.

conv state:
최근 token 주변의 convolution 계산을 위한 state

SSM / temporal state:
state-space recurrence를 위한 state

이 state들은 과거 token마다 K/V를 전부 저장하는 방식과 다르다.

serving에서 달라지는 질문

KV cache 시스템에서는 이런 질문이 중요하다.

새 K/V block을 어디에 붙일 것인가?
긴 context의 block table을 어떻게 관리할 것인가?
여러 request의 KV block을 어떻게 batch로 읽을 것인가?

SSM state 시스템에서는 질문이 바뀐다.

각 request의 current state는 어디에 있는가?
prefill 이후 decode state를 어떻게 넘길 것인가?
prefix caching을 하려면 어느 지점의 state를 저장해야 하는가?
state update kernel은 batch 안에서 어떻게 호출되는가?

즉 SSM은 cache를 없애는 것이 아니라 cache의 의미를 바꾼다.

연결

• kv-cache: Transformer serving의 기본 cache
• mamba-prefill-and-decode: SSM에서 prefill/decode가 어떻게 바뀌는지
• vllm-mamba-cache-spec: vLLM이 Mamba cache를 별도 spec으로 다루는 이유

확인

• KV cache는 token 수가 늘 때 어떻게 커지는가?
• Mamba 계열에서 conv state와 SSM state를 나눠 생각하는 이유는 무엇인가?
• “SSM은 cache 문제를 없앤다”보다 “cache의 의미를 바꾼다”가 더 정확한 이유는 무엇인가?