vLLM Mamba Cache Spec

vLLM에는 attention KV cache만 있는 것이 아니다.

Mamba/SSM 계열 모델을 지원하려면 Mamba state를 위한 cache 표현도 필요하다.

왜 별도 spec이 필요한가

일반 attention cache는 key/value tensor의 block을 관리한다.

Mamba cache는 conv state와 SSM state처럼 다른 모양의 state를 관리한다.

AttentionSpec:
K/V cache block을 설명한다.

MambaSpec:
Mamba state cache의 shape, dtype, page size를 설명한다.

저장하는 텐서의 의미가 다르므로, 같은 KV cache abstraction만으로 처리하기 어렵다.

mamba_cache_mode

vLLM에는 Mamba cache mode가 있다.

none:
prefix caching이 꺼져 있을 때의 기본 상태

all:
block boundary마다 Mamba state를 cache한다.

align:
scheduler step의 마지막 token이 block boundary에 맞을 때 state를 cache한다.

이 설정은 SSM에서도 prefix reuse를 하고 싶기 때문에 필요하다.

Transformer prefix caching은 prefix의 KV block을 재사용한다. Mamba prefix caching은 특정 prefix 지점의 state를 재사용해야 한다.

중요한 관점

Mamba가 “KV cache를 안 쓴다”고 해서 vLLM의 cache manager 문제가 사라지지는 않는다.

오히려 serving engine은 두 종류의 세계를 모두 이해해야 한다.

attention layer:
past K/V blocks

mamba layer:
conv state + SSM state

hybrid 모델에서는 이 둘이 한 모델 안에 같이 존재한다.

연결

• ssm-state-vs-kv-cache: cache 자원의 의미 차이
• hybrid-cache-manager: attention cache와 Mamba cache가 공존할 때 생기는 문제

확인

• vLLM이 MambaSpec을 따로 두는 이유는 무엇인가?
• mamba_cache_mode는 prefix caching과 어떤 관계가 있는가?
• Mamba cache가 KV cache보다 단순하다고만 말하기 어려운 이유는 무엇인가?