掘金2026년 11월 5일 21:41중국어 → 한국어

KIMI K3 경량 복제에서 본 KDA의 AI 인프라 과제

경량 대형 모델 복제 과정에서 KDA가 AI 인프라에 던지는 과제와 각사의 대응을 하이브리드 어텐션 실행 구조 중심으로 정리했다.

중국어 원문을 AI로 번역했습니다. 고유명사와 수치는 원문 표기를 우선하며, 중요한 판단에는 아래 출처 원문을 함께 확인하세요.

0x00 개요

0x01. 먼저 올바른 사고모델(mental model)을 세우기

- 1.1 혼합 어텐션의 전체 실행 체인

- 1.2 세 가지 상태 비교

- 1.3 KDA의 순차성: 분할은 가능하지만 순서는 반드시 유지해야 함

0x02. KDA의 장점과 구조적 한계

- 2.1 장점

- 2.2 대가: 고정 상태는 완전한 기록이 아님

0x03. Prefix Cache와 KV Cache가 더 이상 등가가 아닌 이유

- 3.1 일반 Prefix Cache가 단순한 이유

- 3.2 KDA가 이 등가 관계를 깨는 이유

0x04. KDA가 추론 시스템에 가져오는 새로운 과제

- 4.1 캐시 대상이 하나에서 여러 개로 변함

- 4.2 상태 복구는 경계 일관성을 만족해야 함

- 4.3 활성 상태는 직접 공유할 수 없음

- 4.4 PD/EPD 분리는 더 이상 KV만 옮기는 것이 아님

0x05. SGLang의 대응 방안

- 5.1 KDA‑aware Radix Tree

- 5.2 Sparse Checkpoint: 왜 단순한 등간격이 아닌가

- 5.3 Compressed Checkpoint

- 5.4 COW, Snapshot과 Donate

- 5.5 Unified Memory

- 5.6 Mooncake의 책임 경계

0x06. SGLang 소스 코드 안내

- 6.1 소스 코드 지도

- 6.2 모델 계층: KimiK3DeltaAttention은 행렬곱만 하는 것이 아님

- 6.3 통합 진입점: RadixLinearAttention.forward

- 6.4 KDAKernelDispatcher: 단계별로 다른 kernel 사용

- 6.5 SGLang의 state는 recurrent state만 있는 것이 아님

- 6.6 Checkpoint Pool의 실제 구현

- 6.7 활성 상태, Donate와 경계

- 6.8 checkpoint는 언제 생성되는가

- 6.9 Unified Memory의 소스 코드 근거

- 6.10 SGLang 복구 흐름 개요

0x07. vLLM의 대응 방안

- 7.1 State block이란 무엇인가

- 7.2 논리적 prefix와 물리적 state의 디커플링

- 7.3 Fine‑grained Partial Hit

- 7.4 Interval‑based와 Marconi‑style Retention

0x08. vLLM 소스 코드 안내

- 8.1 소스 코드 지도

- 8.2 KDA 실행과 state cache

- 8.3 metadata가 state를 찾는 방법

- 8.4 Fine‑grained Partial Hit의 소스 코드 근거

- 8.5 Retention의 소스 코드 근거

- 8.6 vLLM 복구 흐름 개요

0x09. 심층 논의

- 9.1 SWA의 checkpoint 방법을 KDA에 그대로 쓸 수 없는 이유

- 9.2 Checkpoint 밀도의 딜레마

- 9.3 인스턴스 간 재사용과 PD/EPD 분리

- 9.4 개념과 소스 코드 대응표

- 9.5 소스 코드를 재현할 때의 최소 읽기 순서

0x10. 최종 결론

0xFF 참고

KDA의 고정 상태 설계는 표준 serving 스택의 여러 컴포넌트를 무효화시켰고, 이에 따라 각 추론 엔진이 이에 적응하고 있다:

이 글의 전반부는 개념적 언어로 KDA가 왜 추론 캐싱의 정확성 조건을 바꾸었는지, 그리고 추론 프레임워크가 정확성을 훼손하지 않으면서 어떻게 VRAM과 재계산을 절약하는지를 설명하고, 후반부는 SGLang과 vLLM의 구체적인 파일과 핵심 함수에 대응하는 소스 코드 안내를 제시한다.

참고:

- 프로젝트 github.com/RossiLH/Mod… 와 시리즈 블로그는 나의 하나의 시도이자 학습 과정이다. 독자의 요구를 충족시키는지 확신할 수 없으며, 나 스스로 성찰하고 개선할 것이니 많은 분들의 소중한 의견을 바란다.

- KIMI K3, DeepSeek V4, PI0.5, PI0.7 등의 모델을 다룬다. 후속으로 새 모델을 추적할 수 있다.

- 전 CPU 복각이며, GPU는 다루지 않고, 대규모 학습(이는 또 다른 방대하고 심오한 분야다)과 하드웨어도 다루지 않는다.

- 모델의 컴포넌트는 최대한 원본 대형 모델의 축소판으로 하며, 컴포넌트에 오픈소스 버전이 없으면 근사 컴포넌트로 대체한다.

- 각 컴포넌트가 어떤 문제를 해결하는지, 이 문제가 어디가 어려운지, 왜 이렇게밖에 풀 수 없는지, 해법의 원리는 무엇인지, 그리고 컴포넌트들이 어떻게 맞물려 하나의 전체를 이루는지를 설명할 것이다. 핵심은 용어를 외우는 것이 아니라 하나의 주된 흐름을 이해하는 것이다.

- 각 주요 컴포넌트에는 바로 실행 가능한 Gradio demo가 있으며 GPU가 필요 없다. 그 밖에 정적 이미지를 생성할 수 있는 여러 시각화 스크립트(plot_*.py)가 있다.

- 모든 파일의 줄 번호는 로컬 소스 코드 스냅샷을 기준으로 한다; 이후 소스 코드가 변경되면 줄 번호는 이동할 수 있지만, 함수명과 호출 관계는 여전히 위치를 찾는 진입점으로 쓸 수 있다.

- 각 글이 최대한 독립적으로 성립할 수 있도록 일부 내용이 서로 다른 글에 모두 존재할 수 있어, 시리즈 차원에서 중복이 있을 수 있으니 양해를 바란다.

Kimi K3의 혼합 어텐션은 대략 다음과 같이 그릴 수 있다:

입력 token ├─ MLA ──> token-addressable KV/latent history └─ KDA ──> convolution state └── recurrent state └── KDA 출력 └──────────> 혼합 출력

K3 특징

MLA는 여전히 token 또는 block 단위로 접근할 수 있는 역사를 보존한다;

KDA는 각 역사 token의 완전한 주소를 저장하지 않고, 고정 크기의 상태를 점진적으로 갱신한다:

S_t = M_t S_{t- 1 } + b_t

- S_{t-1} : 직전 시점의 장기 기억;

- M_t : 현재 token이 결정하는, 기존 내용을 얼마나 보존·감쇠·변환할지;

- b_t : 현재 token이 기록하는 내용;

- S_t : 현재 token을 처리한 뒤의 장기 상태.

KDA는 또한 짧은 convolution을 자주 사용하므로, 최근 여러 token의 convolution state를 저장해야 한다.

차이

따라서 전통적인 Transformer의 KV Cache는 "token 번호별 역사 기록"으로 볼 수 있다:

token 0 -> K,V token 1 -> K,V token 2 -> K,V ...

새 Query가 도착하면, 시스템은 임의의 역사 token을 읽을 수 있다. prefix 캐시는 보통 이렇게 이해할 수 있다: token prefix가 같으면 대응하는 KV block을 재사용할 수 있다.

반면 KDA는 역사를 고정 크기의 recurrent state에 지속적으로 기록한다.

그래서 KDA는 단순히 KV Cache를 다른 텐서로 바꾸는 것이 아니라, "모델 상태"를 하나에서 의미가 서로 다른 여러 객체로 확장한 것이다.

상태 기억하는 것 새 token에 따라 갱신되는가 전형적 복구 방식 token 단위 임의 절단 가능 여부 MLA KV 역사 token의 주소 지정 가능한 표현 추가 KV block 읽기 기본적으로 가능 KDA recurrent state 전체 prefix가 압축된 장기 요약 예 checkpoint 복구 임의 절단 불가 KDA convolution state 최근 국소 윈도우 예 국소 윈도우 스냅샷 복구 합법적 경계에서 저장해야 함

Recurrent state는 장기 기억을 결정하고, convolution state는 다음 단계의 국소 convolution 출력을 결정한다. 둘 다 후속 token에 영향을 미치므로 MLA KV만 복구할 수는 없다.

KDA의 경우, token 구간 C가 상태에 미치는 전체 영향을 다음과 같이 표기한다:

F_C (S) = M_C S + b_C

연속으로 실행되는 두 구간의 합성은:

F_{C_2}(F_{C_1}(S)) = M_{C_2} M_{C_1} S + M_{C_2} b_{C_1} + b_{C_2}

중간의 M_{C_2} b_{C_1}이 중요하다: 첫 구간이 기록한 정보는 두 번째 구간의 감쇠나 소거를 다시 받는다.

이것은 두 가지를 설명한다:

- 한 구간의 token을 하나의 (M, b)로 압축할 수 있으므로 chunkwise 병렬이 가능하다;

- 임의의 두 token 구간의 순서를 바꿀 수는 없는데, 뒷 구간이 앞 구간의 기록에 작용하기 때문이다.

이는 결합법칙을 만족하므로 작은 구간을 먼저 계산한 뒤 합칠 수 있고; 일반적으로 교환법칙을 만족하지 않으므로 token 순서를 임의로 바꿀 수 없다. 이 "구간 분할은 가능하나 순서는 반드시 유지해야 한다"는 성질이 바로 KDA가 병렬화되면서도 시간 순서를 부호화할 수 있는 이유다.

Checkpoint가 저장하는 것은 특정 prefix 경계의 결과다, 예를 들어:

token boundary 4096 ├─ MLA KV prefix ├─ KDA recurrent state S_4096 └─ KDA convolution state C_4096

KDA는 긴 역사를 고정 크기의 state로 압축하며, 주요 이점은:

- decode가 전체 역사를 반복적으로 스캔할 필요가 없다;

- 역사 읽기 대역폭이 더 이상 컨텍스트 길이에 따라 선형으로 증가하지 않는다;

- 상태를 chunk 단위로 계산하고 아핀 합성으로 병합할 수 있다;

- 감쇠가 느린 망각을 표현할 수 있다;

- delta 갱신이 현재 key를 따라 방향성 덮어쓰기를 할 수 있다;

- 장거리 mixing 분기로 적합하다.

상태는 압축된 요약이며, 완전한 기록이 아니다. 하나의 고정 state는 임의로 긴 역사가 한정된 용량을 두고 경쟁하게 해야 한다:

많은 token ↓ 압축 하나의 고정 상태 행렬 ↓ 유사한 key가 충돌할 수 있음 오래된 정보가 감쇠될 수 있음 미래 query가 원래 token을 다시 찾을 수 없음

두 개의 유사한 key가 가까운 방향으로 기록되어 서로 간섭할 수 있다; 오래된 정보가 한번 감쇠되거나 덮어쓰이면, 미래 query는 Full Attention처럼 원래 token을 다시 찾을 수 없다.

따라서 KDA의 state는 "지속적으로 갱신되는 레지스터 파일"에 더 가깝고, Full/MLA KV는 "주소로 읽을 수 있는 캐시"에 더 가깝다. Kimi K3가 둘을 동시에 사용하는 것은 바로 KDA가 효율적인 mixing을 담당하고 MLA가 정밀 검색을 담당하게 하기 위해서다.

KDA의 강점은 "모든 작업에서 Full Attention을 대체"하는 것이 아니라, 고정 상태로 대부분의 장거리 mixing을 담당하고 다시 MLA 등 token‑addressable 분기가 정밀 검색을 보완하는 것이다.

일반 Transformer의 KV Cache는 대략 append식 배열이다:

prefix tokens: [A, B, C, D] KV cache: [K_A,V_A] [K_B,V_B] [K_C,V_C] [K_D,V_D]

이후 E를 생성할 때는 [K_E,V_E]만 추가할 뿐, A, B, C, D의 KV를 수정하지 않는다. 따라서:

token prefix가 같음 ≈ 대응 KV가 같음 ≈ 직접 재사용 가능

이것이 "Prefix Cache와 KV Cache가 거의 등가"라는 말의 유래다. 일반 Transformer에서 prefix의 KV는 "쓰기가 끝나면 안정"된다: 후속 token은 앞 token의 K/V를 수정하지 않는다. token prefix만 일치하면 캐시 내용도 일치한다.

KDA의 state는 각 token의 독립적 기록이 아니라, 전체 prefix를 점화식으로 누적한 결과다:

S_n = F_n(F_{n- 1 }(… F_2(F_1(S_0))))

여기서 각 F_t는 하나의 token의 상태 갱신을 나타낸다.

두 요청이 다음과 같이 공유한다고 가정하자:

요청 1: A B C D E 요청 2: A B C D X

이들은 A B C D 지점의 KDA state를 공유할 수 있지만, E 또는 X 이후에는 반드시 분리되어야 한다. 공유되는 것은 안정적인 프리픽스 스냅샷이며, 두 요청이 동시에 제자리에서 수정할 수 있는 활동 상태가 아니다.

더 중요하게는, 시스템이 토큰 프리픽스가 D까지 매칭되었다는 것만 알고 S_D를 저장하지 않았다면, D의 토큰만으로는 KDA를 복원할 수 없다. 시스템은 더 이른 checkpoint부터 재실행해야 한다:

checkpoint at B ↓ replay C, D ↓ obtain S_D ↓ continue E or X

따라서 KIMI K3에서 위치 n을 복원하려면 반드시 동시에 알아야 한다:

- MLA의 위치 n 이전 전체 KV;

- KDA의 위치 n 시점 recurrent state;

- 권적 분기의 위치 n 시점 convolution state;

- 모든 상태가 실제로 동일한 토큰 프리픽스와 동일한 모델 실행 경로에서 비롯되었는지 여부.

토큰 프리픽스만 매칭하는 것은 "텍스트가 같다"는 것만 설명할 뿐, "재귀 상태가 저장되었고 복원 가능하다"는 것을 증명하지 못한다.

흔한 의문: 왜 토큰 프리픽스가 매칭되었다고 해서 KDA 캐시가 복원될 수 있는 것은 아닌가?

Radix Tree나 prefix hash는 token 0 … token N이 같다는 것만 증명할 수 있을 뿐이며, 복원에는 추가로 S_N이 저장되었는지, C_N이 저장되었거나 재구축 가능한지, MLA KV_N이 사용 가능한지가 필요하다. 만약 state가 0, 4096, 8192에서만 저장되었는데 토큰이 6000까지 매칭된다면, 시스템은 텍스트만으로 S_6000을 직접 얻을 수 없고, 4096으로 후퇴하여 4096‑6000을 replay해야 한다.

시스템은 더 이상 KV block만 관리하지 않고, 다음을 관리해야 한다:

도전은 "캐시가 하나 더 늘었다"는 것이 아니라, 이러한 객체들의 의미, 갱신 방식, 공유 가능성이 모두 다르다는 점이다.

확장 가능한 캐시 엔트리는 최소한 다음을 포함해야 한다:

CacheEntry { model_version prefix_range tokenization_fingerprint state_kind: mla_kv | kda_recurrent | kda_conv physical_location dtype / layout writable parent_checkpoint }

이 중 가장 핵심은 prefix_range와 state_kind이다. 경계가 없으면 복원이 연속적인지 판단할 수 없고, 타입이 없으면 "주소 지정 가능한 이력"과 "순차 요약"을 구분할 수 없다.

따라서 핵심은 캐시 키를 "token prefix"에서 다음으로 업그레이드하는 것이다:

(model/version, tokenizer, prefix boundary, MLA KV, KDA recurrent state , convolution state )

실제 시스템은 또한 다음을 처리해야 한다:

- 활동 state와 공유 checkpoint의 분리;

- state 압축과 복원;

- checkpoint 제거;

- 로컬과 원격 캐시의 조정;

- 서로 다른 병렬 rank 간 객체 정렬.

합법적 복원 지점은 "가장 긴 토큰 프리픽스"처럼 단순한 것이 아니라, 모든 관련 상태가 사용 가능한 공통 경계이다. 만약 MLA KV가 4096 토큰까지 복원되었는데 KDA state는 3968 토큰까지만 복원되었다면, 이후 계산은 더 이상 연속 실행과 등가가 아니다.

혼합 모델의 출력은 여러 분기가 함께 결정한다. 가령:

MLA 복원 N = 6000 KDA 복원 N = 4096

그렇다면 두 분기가 보는 이력 길이가 다르므로, 이후 혼합 결과는 당연히 N=6000에서 연속 실행한 결과가 아니다. 따라서 복원 대상은 길이만 비슷해 보이는 여러 캐시 조각이 아니라, "동일 시점의 시스템 스냅샷"과 같아야 한다.

흔한 의문: 왜 KDA 상태는 KV처럼 임의로 절단할 수 없는가?

KV를 절단한다는 것은 이후 일부 토큰을 더 이상 읽지 않는다는 의미이며, 남은 KV는 여전히 유효한 기록이다. KDA state를 절단하는 것은 다르다. state의 각 요소는 여러 역사 토큰이 혼합된 결과일 수 있다. 큰 state에서 앞쪽 10%의 "역사"를 마음대로 삭제할 수 없는데, 그것에는 독립적인 토큰 주소가 없기 때문이다. 올바른 방법은 두 가지뿐이다: 이미 저장된 더 이른 checkpoint를 사용해 replay하거나, 목표 경계에 이미 저장된 완전한 state를 사용하는 것.

MLA KV는 보통 기록된 후 변하지 않으며 여러 요청이 참조할 수 있다. KDA state는 decode가 계속됨에 따라 갱신되므로, 여러 요청이 동일한 쓰기 가능 state를 직접 가리키면 한 요청의 갱신이 다른 요청을 오염시킨다.

따라서 "공유되는 안정적 checkpoint"와 "요청 자신의 활동 사본"을 분리해야 한다. 공유 checkpoint는 읽기 전용이고, 활동 state는 사적으로 쓰기 가능하다.

전통적인 Prefill‑Decode 분리는 주로 KV blocks를 전송한다. KDA는 recurrent state, convolution state를 전송하거나 재구축해야 하며, 멀티모달과 결합되면 Encoder가 생성하는 시각 embedding도 고려해야 한다. 따라서 인터페이스 추상화는 다음과 같이 확장된다:

KV Cache transfer

에서

model state transfer

로. 통신 계층, 스케줄링 계층, 캐시 계층 모두 상태 타입, 버전, 경계, 생명주기를 알아야 한다.

SGLang은 Radix Tree의 프리픽스 매칭 능력 위에 KDA‑aware checkpoint를 추가했다. Radix Tree는 여전히 "어떤 토큰 프리픽스가 매칭되는지"를 답하는 역할을 맡지만, 대응하는 KDA 상태를 가진 노드만이 직접 복원할 수 있는 실행 지점이 된다. 이는 캐시 노드를 prefix -> KV blocks에서 다음으로 확장하는 것과 같다:

prefix boundary -> MLA KV + KDA recurrent state (+ convolution state)

SGLang의 기존 Radix Tree는 요청 간 token 프리픽스 매칭에 강점이 있다. KDA-aware가 되면 각 후보 노드는 복구 가능한 state가 있는지도 알아야 한다:

Radix Tree root └─ system prompt └─ tools ├─ user branch A └─ user branch B

각 노드는 다음을 가질 수 있다: token prefix, MLA KV blocks, KDA checkpoint, checkpoint boundary.

매칭 흐름은 다음과 같이 바뀐다:

- 가장 긴 token prefix를 찾는다

- 매칭 경로를 따라 앞으로 나아가며 가장 가까운 합법적 KDA checkpoint를 찾는다

- 같은 경계의 MLA KV + KDA state를 가져온다

- 히트된 이후의 꼬리 부분을 다시 계산한다

따라서 가장 긴 token 히트 지점과 최종 복구 지점은 다를 수 있다. 이는 캐시 정확성을 문자열 매칭에서 모델 상태 매칭으로 끌어올린다.

흔한 의문: SGLang의 KDA-aware Checkpoint란 무엇인가?

일반 Radix Tree 노드는 "프리픽스가 어디까지 매칭되는가"에 답한다. KDA-aware 노드는 "여기에 대응하는 복구 가능한 state가 있는가"도 답한다. token 매칭과 상태 매칭이 동시에 성립해야만 노드가 유효한 복구 지점이 된다. 히트 흐름은 다음과 같다: token 경로 매칭 → 가장 가까운 합법적 checkpoint로 후퇴 → MLA/KDA 공동 복구 → 나머지 token replay.

모든 token마다 완전한 KDA state를 저장하면 공간이 빠르게 팽창한다. 그래서 SGLang은 희소 checkpoint를 채택한다. 일부 경계에서만 상태를 저장하며, 단순히 전역적으로 균등 간격으로 샘플링하는 것도 아니다.

전략은 세 가지 위치를 포함한다:

- Prefill: chunk 경계를 우선한다;

- Decode: 고정 간격으로 checkpoint를 생성한다;

- Radix 분기: 여러 요청이 공유/분기하는 위치를 우선 보존한다.

이유는 세 위치의 재사용 확률이 다르기 때문이다:

- Prefill chunk 경계는 원래 실행 경계이므로, 이를 저장하는 데 별도의 동기화 지점을 만들 필요가 없다;

- Decode는 모든 token이 미래의 프리픽스가 될 수 있지만, token마다 저장하는 것은 너무 비싸므로 샘플링한다;

- 분기점은 이미 여러 요청에 의해 재사용 가치가 입증된 곳이다.

이후 단일 경로의 checkpoint 개수를 제한하고, LRU로 오랫동안 접근되지 않은 상태를 회수해야 한다.

따라서 본질적으로 이것은 캐시 예산과 재계산 비용 사이의 최적화이며, 단순히 최대 히트율을 추구하는 것이 아니다. 모든 위치를 균등하게 저장하는 대신, 실행 구조와 접근 패턴을 활용해 checkpoint를 한정된 예산 아래의 "캐시 투자"로 다룬다.

흔한 의문: SGLang의 희소 Checkpoint는 공간/히트율 모순을 어떻게 해결하는가?

모든 token의 state를 저장하지 않고, 소량의 고가치 경계만 선택해 저장한다. 이렇게 상주 오버헤드를 줄이고, 히트 후 국소 재계산으로 다소 거친 복구 입도를 보완한다. Prefill chunk 경계는 자연히 저장에 적합하고; Decode는 간격으로 개수를 제어하며; Radix 분기점은 대개 재사용 확률이 더 높고; LRU로 차가운 checkpoint를 회수하며; 경로 예산으로 단일 긴 경로가 VRAM을 잠식하는 것을 방지한다.

모든 checkpoint가 똑같이 가치 있는 것은 아니다. SGLang은 상태를 다음으로 나눈다:

런타임 활동 풀: 현재 요청이 갱신 중이며, 원래 정밀도를 유지한다

캐시 풀: Radix Tree에서 재사용을 기다리며, temporal recurrent state를 INT8로 압축할 수 있다

국소 윈도 상태는 부피가 작아 보통 원래 정밀도를 유지한다. 압축된 checkpoint를 히트하면:

INT8 checkpoint ↓ dequantize 새로운 활동 state slot ↓ 원래 정밀도로 decode 계속

이는 KDA를 영구히 INT8에서 재귀시키는 것이 아니라, "캐시에 머무는 스냅샷"만 압축하여 연속적인 양자화/역양자화로 인한 오차 누적을 피하는 것이다.

SGLang은 실행 중인 활동 상태와 Radix Tree의 캐시 상태를 분리한다. 비활성 checkpoint는 독립 캐시 풀에 넣을 수 있고, 공간을 주로 차지하는 KDA temporal state를 INT8로 압축하며, 더 작은 국소 윈도 상태는 원래 정밀도를 유지한다.

압축된 checkpoint를 히트하면 먼저 새로운 활동 state 슬롯으로 해제한 뒤, 원래 정밀도로 token 갱신을 계속한다. 이로써 양자화 오차는 주로 캐시 상주와 복구 지점에서 발생하며, 연속 decode에서 반복 누적되지 않는다. 여전히 손실 복구이며, 품질과 용량 사이의 트레이드오프가 필요하다.

흔한 의문: Compressed Checkpoint는 무엇을 하는가?

저인기 checkpoint의 저장 비용을 낮추고, 한정된 VRAM을 재사용률 높은 상태에 남겨 둔다. 주요 temporal recurrent state는 INT8을 사용하고, 더 작은 국소 윈도 상태는 원래 정밀도를 유지한다. 히트 시 먼저 활동 슬롯으로 해제한 뒤 정상 정밀도로 실행을 복구한다. 이렇게 상주 공간을 절약하지만 양자화 복구 오차를 받아들여야 하고, 채널별로 적절한 양자화 파라미터를 저장해야 한다.

유의할 점: INT8 checkpoint는 선택적 최적화이며, 모든 SGLang 배포에서 기본 활성화되는 것은 아니다.

세 가지는 함께 "공유 캐시는 안정적이어야 하고, 실행 상태는 가변적이어야 한다"는 모순을 해결한다.

Radix Tree shared checkpoint ├─ request A: copy-on-write -> private active state └─ request B: copy-on-write -> private active state

- Snapshot: 요청의 현재 상태를 동결하여 Radix Tree에 붙이고, 미래의 재사용 가능한 스냅샷이 되게 한다;

- Copy-on-write: 공유 스냅샷을 히트하면 먼저 복사하거나 사설 복사본을 만들고, 이후 갱신이 공유 스냅샷을 오염시키지 않게 한다;

- Donate: 스냅샷 완료 후 state slot의 소유권을 직접 이전하고, 인덱스만 바꾸며 데이터 블록 전체를 다시 복사하지 않는다.

두 요청이 모두 S_4096을 히트한다고 가정하면:

공유 checkpoint S_4096 ├─ A: private S_A └─ B: private S_B

Snapshot은 안정적인 공유 버전을 만드는 역할을 하고, COW는 최초 쓰기 시 공유 버전을 보호하는 역할을 하며, Donate는 이미 완료된 쓰기 가능 state를 캐시 트리로 직접 넘겨 복사를 피하는 역할을 한다. Donate의 가치는 "커다란 state 블록 하나를 복사하는 것"을 "메모리 슬롯 하나를 넘기는 것"으로 바꾸는 데 있다. 핵심 원칙: 공유 버전은 읽기 전용, 활동 버전은 쓰기 가능.

흔한 의문: COW, Snapshot, Donate는 어떻게 상태를 격리하는가?

Snapshot은 공유된 이력을 고정하고, 활동 복사본이 계속 재귀 연산을 수행한다. COW는 최초 쓰기 시 복사하며, Donate는 이전 요청이 끝난 뒤 재사용 가능한 활동 저장 공간을 넘긴다. 이렇게 하면 공유 checkpoint는 수정되지 않고, 실행 시에도 매 단계마다 전체 복사를 할 필요가 없다.

Unified Memory는 MLA KV block과 KDA state block을 동일한 물리적 VRAM 풀에 넣지만, 각자의 논리 구조와 할당 단위는 유지한다.

VRAM 풀 MLA KV blocks -----> <----- KDA state blocks 공유 유휴 영역/동적 경계

MLA KV와 KDA state는 여전히 서로 다른 논리 타입, 크기, 수명 주기를 가지며, 통일되는 것은 저수준 용량 관리다. 이렇게 하면 VRAM을 미리 두 개의 비율이 어긋난 풀로 영구히 분할하는 일을 피할 수 있다.

시스템은 실시간 부하에 따라 양쪽 점유량을 동적으로 조정할 수 있어, 특정 유형의 캐시가 고정 용량을 독점하는 것을 피할 수 있다. 통일되는 것은 물리적 용량 관리이며, 두 가지 상태를 같은 종류의 객체로 섞는 것이 아니다.

장점은 용량이 부하에 따라 치우칠 수 있다는 것이고, 대가는 저수준 할당자가 서로 다른 크기의 객체, 빈 공간 재사용, 단편화 문제를 처리해야 한다는 것이다.

흔한 의문: Unified Memory란 무엇인가?

그것은 MLA KV와 KDA state를 동일한 종류의 데이터로 만드는 것이 아니라, 둘이 같은 물리적 VRAM 예산에서 할당되도록 하는 것이다. KV manager와 State manager는 논리적으로 여전히 둘로 나뉘어 있지만, 물리적 용량 풀을 공유한다. 워크로드가 MLA에 치우치면 KV를 더 남길 수 있고, KDA에 치우치면 state를 더 남길 수 있다.

책임 경계는 이렇게 그릴 수 있다.

SGLang ├─ 무엇이 합법적인 복구 지점인지 판단 ├─ KDA checkpoint 생성/복구 └─ 로컬 요청 상태 관리 ↓ Mooncake / HiCache ├─ 합법적으로 확인된 캐시 객체 저장 ├─ 인스턴스 간 전송 └─ 더 넓은 범위의 공유 캐시 관리

Mooncake은 KDA의 재귀 수학을 새로 이해할 필요가 없다. 경계, 타입, 버전 정보가 붙은 객체를 안정적으로 저장하고 운반하기만 하면 된다.

이 장은 개념적 요약이 아니라 소스 코드 안내서다. 모든 파일 행 번호는 로컬 소스 스냅샷에 기반한다. 소스가 이후 바뀌면 행 번호는 이동할 수 있지만, 함수명과 호출 관계는 여전히 위치를 찾는 진입점으로 쓸 수 있다.

소스 코드에서 KDA 지원은 별도의 KDACache 클래스 하나가 아니라 네 계층에 분산되어 있다.

모델 계층: KDA의 q /k/v/g/beta 생성 Kernel/backend 계층: prefill, decode, conv, recurrent update 상태 풀 계층: recurrent state, convolution state, checkpoint, 활동 복사본 스케줄링/캐시 계층: prefix hit, state index, checkpoint 경계, 인스턴스 간 전송

가장 오해하기 쉬운 점은 SGLang의 "checkpoint pool"이 일반적인 KV pool이 아니라는 것이다. 그것이 저장하는 것은 복구 가능한 KDA 실행 상태다.

파일 핵심 함수 / 클래스 역할 sglang/srt/models/kimi_k3.py KimiK3DeltaAttention.forward_qkvbg, forward_qkvbg_fused, forward K3 모델 측에서 KDA 입력을 생성하고 attention backend에 연결 sglang/srt/layers/radix_linear_attention.py RadixLinearAttention.forward, unified_linear_attention_with_output 선형 attention 결과를 forward batch와 cache location에 연결 sglang/srt/layers/attention/linear/kda_backend.py KDAKernelDispatcher.decode, extend, target_verify decode/prefill/spec‑verify에 따라 KDA kernel 선택 sglang/srt/mem_cache/mamba_checkpoint_pool.py Int8CheckpointStore, MambaCheckpointPool.store_from_active, load_to_active 캐시 상태, INT8 압축, 활동 풀로 복구 sglang/srt/mem_cache/memory_pool.py HybridReqToTokenPool.__init__, donate_mamba_ping_pong_slot, clear 활동 상태 할당, checkpoint pool 연결, Donate와 해제 sglang/srt/managers/schedule_batch.py _mamba_radix_cache_v2_req_prepare_for_extend, mamba_lazy_prealloc_at_boundary checkpoint 경계 선택, ping‑pong state 준비 sglang/srt/mem_cache/unified_memory_pool.py UnifiedKVPool, MambaSubPoolSpec MLA/MHA와 Mamba state가 물리적 byte buffer 공유

파일: sglang-main/python/sglang/srt/models/kimi_k3.py

핵심 위치:

- forward_qkvbg: 1679행부터;

- forward_qkvbg_fused: 1773행부터;

- forward: 1821행부터.

핵심 로직:

hidden_states ↓ mixed_qkv, gates, beta, f/a 등 KDA 입력을 투영 ↓ KDA/MLA attention ↓ AttnRes, output gate, o_proj

파일 최상단 1‑7행은 Attention Residual, Latent MoE, MLA output gate, full‑rank KDA gate를 포함한 K3‑specific features를 직접 선언하고 있다. 이는 Infra가 K3를 단순히 "하나의 일반 GDN 레이어"로 처리할 수 없음을 의미한다.

파일: sglang-main/python/sglang/srt/layers/radix_linear_attention.py:38-121

함수: RadixLinearAttention.forward(...)

간략화한 로직:

만약 extend/prefill이고 piecewise CUDA graph가 활성화된 경우: unified_linear_attention_with_output 경로로 처리, 그렇지 않으면: get_attn_backend().forward(...)에 위임

소스코드 86‑121행이 이 분기를 보여준다. 142‑156행의 unified_linear_attention_with_output은 forward_batch.out_cache_loc을 일시적으로 좁힌 뒤 backend를 호출하고, 출력을 다시 기록한다.

엔지니어링상의 의미:

- KDA 결과는 요청에 대응하는 cache location에 반드시 기록되어야 한다;

- cache location은 batch 메타데이터에 속하며, kernel 내부의 임시 변수가 아니다;

- CUDA graph/piecewise 경로에서도 동일한 상태 기록 시맨틱을 유지해야 한다.

파일: sglang-main/python/sglang/srt/layers/attention/linear/kda_backend.py:39-202

클래스: KDAKernelDispatcher

주요 함수:

- __init__: 42‑202행, decode/prefill/verify backend 선택;

- packed_decode: 204‑235행;

- decode: 237‑264행;

- target_verify: 266‑307행;

- extend: 309‑332행.

간략화한 호출 그래프:

decode: Triton / Helion / CuTe DSL / FlashInfer 경로 prefill: Triton / FlashKDA / CuTe DSL / NVIDIA KDA / PTX 경로 speculative verify: Triton / FlashInfer / nv_cutedsl 경로

소스코드 64‑190행은 이들 backend와 fallback 조건을 명확히 나열한다. 237‑264행의 decode는 ssm_states, cache_indices를 전달받고; 309‑332행의 extend도 동일한 상태 인덱스를 전달받는다.

이는 다음을 증명한다: KDA state는 kernel의 명시적 입력이며, "모델이 Python 객체 안에 암묵적으로 저장하는" 것이 아니다.

kda_backend.py 6‑18행의 임포트:

- causal_conv1d_fn

- causal_conv1d_update

- TritonKDAKernel

- build_verify_intermediate_state_indices

따라서 한 번의 KDA 실행은 동시에 다음을 수반한다:

컨볼루션 윈도 상태 recurrent SSM/KDA 상태 speculative verify 중간 상태 인덱스

흔한 의문: recurrent state와 convolution state란 무엇인가? SGLang에 대응되는 것이 있는가?

Recurrent state는 KDA의 장기 재귀 행렬로, 압축된 긴 이력을 담고 있다; convolution state는 단기 컨볼루션에 필요한 최근 윈도우로, 국소적 연속성을 담고 있다. 둘은 서로 다른 저장 전략을 취할 수 있다: temporal state는 압축할 수 있고, 국소 윈도우는 보통 원래 정밀도를 유지한다. SGLang의 KDA‑aware checkpoint는 주로 recurrent state를 중심으로 전개되지만, 정당한 복원은 그것만 복원해서는 안 된다; 대응하는 국소 윈도우 상태도 복원하거나 명확한 규칙에 따라 재구성해야 한다.

6.6.1 설계 경계: 공유 checkpoint는 읽기 전용, 활동 state는 사적이며 쓰기 가능

파일: sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:15-45

모듈 설명은 다음을 명확히 구분한다:

cached states: radix - owned idle compressed active MambaPool: running requests full precision kernel - facing

이것이 바로 Copy‑on‑write의 근본이다: 공유 checkpoint는 읽기 전용, 활동 state는 사적이며 쓰기 가능하다.

6.6.2 Int8CheckpointStore

파일: sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:60-184

- __init__: 79‑113행, qdata와 scale 할당;

- quantize: 117‑134행;

- dequantize: 136‑138행;

- store: 142‑149행;

- load: 151‑153행;

- copy_to_pool: 155‑164행;

- store_from_pool: 166‑174행.

데이터 레이아웃:

qdata: [ layers, slots, heads, d_v, d_k ] int8 scale: [ layers, slots, heads, 1, d_k ]

第 127행은 d_v를 따라 max-abs를 수행하고, 第 128-133행은 대칭 양자화를 완료한다. 第 164행은 checkpoint를 역양자화한 뒤 곧바로 active pool에 기록하는데, 이것이 바로 cache hit 시의 COW 복원이다.

중요한 엔지니어링 결론:

checkpoint 적중: int8 → 한 번 dequantize → active state에서 decode 계속: active state는 원래 정밀도로 업데이트 유지

모든 token이 INT8 state 상에서 반복 양자화되는 것은 아니다.

6.6.3 MambaCheckpointPool

파일: sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:186-280

- __init__: 第 187-225행, temporal int8 store, conv buffers, slot allocator를 생성;

- store_from_active: 第 245-251행;

- load_to_active: 第 253-259행;

- estimate_mem_usage_bytes: 第 262행부터.

store_from_active의 의미:

active_mamba_pool.temporal → checkpoint temporal active_mamba_pool.conv → checkpoint conv

load_to_active의 의미:

checkpoint temporal → active temporal checkpoint conv → active conv

SGLang은 실제로 convolution state에 대응하는 저장소를 갖추고 있으며, recurrent state만 캐시하는 것이 아니다.

6.6.4 선택적 초기화

파일: sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:300-369

핵심 로직:

Mamba 설정이 없으면: 생성하지 않음 enable_int8_mamba_checkpoint가 켜져 있지 않으면: 생성하지 않음 그 외의 경우: 공간을 추정하고, 사용 가능한 GPU 메모리를 확인하고, 독립적인 checkpoint pool을 생성

이는 INT8 checkpoint가 선택적 최적화임을 나타내며, 모든 SGLang 배포에서 기본으로 활성화된다고 써서는 안 된다.

6.7.1 HybridReqToTokenPool.init

파일: sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1271-1289

self.mamba_allocator = MambaSlotAllocator(...) self.mamba_map = {...} self.mamba_ckpt_pool = maybe_init_int8_mamba_checkpoint_pool(...)

활성 state와 checkpoint state는 서로 다른 풀을 사용한다: 활성 요청은 mamba_allocator / active pool을 거치고, Radix cached checkpoint는 mamba_ckpt_pool을 거친다.

6.7.2 donate_mamba_ping_pong_slot

파일: sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1488-1509

함수: donate_mamba_ping_pong_slot(req, new_slot)

보존할 ping-pong slot을 찾음 ↓ 이전 slot을 캐시 가능한 state로 반환 ↓ 요청을 new_slot으로 전환하여 계속 실행

第 1491-1495행의 docstring에는 이미 "donate tracked-state slot to radix cache"라고 명확히 적혀 있다. 이것은 데이터를 복사하는 것이 아니라 ownership/slot을 radix cache에 넘겨 사용하게 하는 것이다.

6.7.3 clear

파일: sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1565-1577

第 1569-1574행은 mamba_ckpt_pool.clear()를 특별히 처리한다. 그 이유는 Radix Tree를 비운 뒤에는 checkpoint slot도 반드시 해제해야 하며, 그렇지 않으면 "트리는 없어졌는데 slot은 여전히 점유된" 누수가 발생하기 때문이다.

6.8.1 _mamba_radix_cache_v2_req_prepare_for_extend

파일: sglang-main/python/sglang/srt/managers/schedule_batch.py:2672-2735

함수: _mamba_radix_cache_v2_req_prepare_for_extend(req)

- mamba_cache_chunk_size()를 읽음;

- mamba_checkpoint_grid(self.tree_cache.page_size)에 따라 checkpoint 그리드를 얻음;

- extend_range.length >= checkpoint_grid인지 판단;

- 추적할 seqlen을 계산;

- state를 ping-pong track buffer에 넣음.

2677‑2680번째 줄은 checkpoint 그리드가 tree page에서 온다고 설명하고, 2699‑2730번째 줄은 정렬 경계를 계산한다. 이는 KDA checkpoint가 임의의 token에 자연스럽게 존재하는 것이 아닌 이유를 설명한다. 즉 스케줄러가 합법적인 추적 경계를 능동적으로 선택한다.

6.8.2 mamba_lazy_prealloc_at_boundary

파일: sglang-main/python/sglang/srt/managers/schedule_batch.py:3002-3031

함수: mamba_lazy_prealloc_at_boundary(mamba_track_interval)

3015‑3017번째 줄은 요청이 tracking interval에 속하는지 검사하고, 3026‑3031번째 줄은 필요에 따라 두 번째 ping‑pong slot을 요청한다.

이것이 바로 lazy ping‑pong이다.

평상시: 하나의 활성 slot → checkpoint 경계에서: 두 번째 slot을 임시로 요청하고, 새 상태는 다른 쪽에 기록되며, 이전 상태는 안전하게 donate될 수 있다

파일: sglang-main/python/sglang/srt/mem_cache/unified_memory_pool.py

- MambaSubPoolSpec: 181‑210번째 줄;

- UnifiedKVPool: 217번째 줄부터;

- 생성자: 222‑280번째 줄.

MambaSubPoolSpec의 185‑205번째 줄은 conv_state_shapes, temporal_state_shape, conv_dtype, temporal_dtype를 명확히 기술한다.

UnifiedKVPool의 217‑220번째 줄은 두 개의 sub‑pool이 하나의 uint8 byte buffer를 공유한다고 설명하고, 232‑243번째 줄은 현재 구현이 위로 증가하는 하나와 아래로 증가하는 하나의 두 풀로 제한된다고 밝힌다.

간략화한 레이아웃:

동일한 byte buffer MHA/MLA pool ---> 빈 영역 <--- Mamba pool

주의: 이는 물리적 용량의 통일이지, 논리적 객체의 통일이 아니다.

요청 prefix 적중 ↓ Radix 선택/준비 KDA tracking boundary ↓ MambaCheckpointPool.load_to_active(...) ├─ temporal int8 dequantize └─ conv native-dtype copy ↓ 요청이 전용 active slot 사용 ↓ 경계 도달 시 donate/snapshot ↓ checkpoint pool + radix cache

핵심 증거 파일: mamba_checkpoint_pool.py:245-259, memory_pool.py:1488-1509, schedule_batch.py:2672-2735.

vLLM은 서로 다른 상태를 동일한 스케줄링 프레임워크에 포함하지만, 억지로 동일한 저장 방식을 쓰지는 않는다.

동일한 Scheduler ├─ MLA/full-attention: Paged KV blocks └─ KDA: ├─ recurrent-state blocks/checkpoints └─ convolution-state blocks/checkpoints

따라서 "통일"은 주로 요청 스케줄링과 생명주기 조정을 뜻하며, 저수준에서는 여전히 KV와 KDA state의 의미론적 차이를 인정한다.

State block은 물리적 메모리 관리 단위이며, 하나의 token과 같지도 않고 완전한 요청과 같지도 않다.

State block B ┌─────────────────────────────┐ │ 한 구간의 연속 실행으로 생성된 KDA state │ │ recurrent state │ │ convolution state │ │ 경계/버전/레이아웃 등 메타데이터 │ └─────────────────────────────┘

그것은 할당, 회수, 참조 카운팅, 교체 및 원격 전송을 위해 동작한다. 하나의 큰 block 안에 여러 개의 논리적 Prefix Entry를 등록할 수 있다.

vLLM 맥락에서:

- recurrent state: KDA가 어느 위치까지 스캔한 뒤 얻는 장기 상태 행렬로, 장문 맥락 요약을 담는다;

- convolution state: 단기 convolution에 필요한 최근 여러 token의 국소 캐시로, 국소적 연속성을 담는다.

전자는 장기 기억을 결정하고, 후자는 다음 단계의 국소 convolution 출력을 결정한다. 둘 다 후속 token에 영향을 주므로 MLA KV만 복원해서는 안 된다.

자주 나오는 질문: vLLM의 state block이란 무엇인가?

그것은 state의 물리적 컨테이너로, KV block의 "메모리 페이지" 역할과 비슷하지만 그 안에는 token별 K/V가 아니라 KDA 상태 스냅샷이 담긴다. 여기에는 보통 state tensor, 커버/연관된 prefix 범위, 스냅샷 경계, dtype/layout, 참조 카운트, 쓰기 가능 여부가 함께 붙는다. 하나의 state block은 여러 개의 논리적 복원 지점을 가질 수 있으므로, 물리적 블록 크기가 Prefix Matching 세분성과 같을 필요는 없다.

만약 Prefix Matching의 경계가 State block의 경계와 같아야 한다면, 매칭 세분성은 물리적 분할에 의해 제한된다. 그래서 vLLM은 두 가지를 분리한다.

- 논리 계층: 임의의 완료된 token prefix를 기록한다;

- 물리 계층: 더 큰 state block으로 상태를 저장한다.

둘은 오프셋, 경계 메타데이터, state 버전을 통해 연결된다.

물리적 State block이 4096 token 단위로 관리되지만, 실제 요청이 4480 token까지 공유된다고 가정하자.

물리적 저장: State block B: [ 0 ————————————————— 8192 ] 논리적 prefix: Prefix Entry at 4096 Prefix Entry at 4480

강제로 block-aligned를 적용하면 4096만 적중할 수 있지만, 논리적 매칭이 block 내부에 떨어질 수 있다면 4480을 더 긴 합법적 프리픽스로 삼을 수 있다. 여기서 state를 "반 블록"으로 잘라내는 것이 아니라, 큰 컨테이너 안에 더 세밀한 유효 경계를 저장하거나 인덱싱하는 것이다.

흔한 의문: 논리적 프리픽스 매칭과 물리적 상태 위치의 분리(decoupling)란 무슨 뜻인가?

"지도 위의 주소"와 "창고 안의 상자"를 분리하는 것이다. 논리 계층은 "token 4480에 적중했다"고 말할 수 있지만, 물리 계층은 0‑8192 token을 동일한 state block B에 담아 둘 수 있다. B 안에 state after 4480이 등록되어 있기만 하면, 논리 계층은 4480을 복구 지점으로 삼을 수 있다. 그렇지 않다면 물리 block의 존재만으로는 4480에 복구 가능한 state가 있다는 것을 자동으로 증명하지 못한다.

하나의 State block이 128 token의 상태를 담을 수 있다고 가정하자. 그런데 요청 프리픽스는 96 token까지만 적중했다. 순진한 방식으로는 128 token짜리 block이 아직 완전하지 않으므로 96 token에서 복구할 수 없고, 더 이른 위치부터 다시 계산해야 한다.

Fine‑grained Partial Hit의 방식은 다음과 같다. 여전히 이 128‑token 물리 block을 사용하되, 그 안에 더 세밀한 여러 Prefix Entry를 기록한다.

State block B: entry at 32 -> state after token 32 entry at 64 -> state after token 64 entry at 96 -> state after token 96 entry at 128 -> state after token 128

따라서 96 token은 아직 채워지지 않은 물리 State block 안에 있음에도 합법적 복구 지점이 될 수 있다. 여기서 "partial hit"는 상태 행렬의 절반만 복구한다는 뜻이 아니라, 해당 block 안에 사전에 기록된 어떤 유효 상태 스냅샷을 복구한다는 뜻이다.

전체 과정:

요청 P: 0 … 4480을 실행하여 state S_4480 생성 물리적으로는: S_4480은 큰 State block B 안에 위치 새 요청: token 0 … 4480 적중 ↓ Prefix Entry(4480) 발견 ↓ S_4480 추출/복구 ↓ 새 요청 전용 활동 슬롯에 복사 ↓ 계속 생성

"Partial"은 논리적 적중 위치가 물리 block의 일부만 차지할 수 있다는 뜻이지, 대응하는 스냅샷이 없는 중간 위치로부터 시스템이 state를 억지로 추론하도록 허용한다는 뜻이 아니다.

이는 하나의 큰 State block 안에 여러 세밀한 Prefix Entry를 등록할 수 있게 한다. 새 요청이 4480에 적중할 때 block 시작점 4096으로 후퇴할 필요도 없고, 또 하나의 완전한 물리 block을 복사할 필요도 없다. 해당 state를 직접 찾아내고, 그것을 전용 활동 state로 복사한다.

핵심 경계: partial hit는 "논리 경계가 물리 컨테이너보다 세밀한 것"을 허용하지만, "존재하지 않는 스냅샷으로부터 state를 추론하는 것"은 허용하지 않는다.

흔한 의문: Fine‑grained Partial Hit는 어떻게 작동하는가?

큰 state block 안에 더 세밀하고 연쇄적인 Prefix Entry를 추가로 기록한다. 예를 들어 하나의 block이 4096 token을 커버하는데, 실제 공유 프리픽스는 4480 token까지 이른다. 논리적 매칭은 block 내부의 4480 경계에 떨어질 수 있고, 해당 위치에서 이미 생성된 KDA state와 연결될 수 있으며, Prefix Matching을 4096의 물리 경계로 거칠게 만들 필요가 없다. 이는 미완성 state를 "반으로 자르는" 것이 아니라, 여러 논리적 복구 지점을 더 큰 물리 컨테이너에 넣는 것이다. 즉 물리적으로는 합쳐서 저장하고, 논리적으로는 세밀하게 적중한다. 복구 시에는 여전히 해당 state를 요청 전용 공간에 복사한 뒤 재귀를 계속해야 한다.

제한된 캐시 예산에서는 모든 위치의 checkpoint를 저장할 수 없다.

- Interval‑based Retention: 고정 또는 구조화된 간격으로 보존한다. 예를 들어 일정 token마다 저장한다. 안정적이고 예측 가능하며 구현이 단순하다.

- Marconi‑style Selective Retention: 런타임 접근 빈도, 프리픽스 인기도, 재사용 가치에 따라 보존 위치를 선택한다. 더 유연하지만 온라인 통계와 퇴출 전략이 필요하다.

전자는 계획적으로 도로 표지판을 세우는 것과 비슷하고, 후자는 교통량에 따라 인기 출구를 설정하는 것과 비슷하다.

Interval-based: 0 ———— K ———— 2K ———— 3K ———— 고정적, 예측 가능, 구조화된 프리픽스에 적합. Marconi-style: 처음 프리픽스를 봤을 때: 접근만 기록. 두 번째 적중 시: 그때 state 저장. 핫스팟은 계속 보존하고 콜드스팟은 퇴출.

Interval‑based는 추가로 Prompt 끝 상태를 보존할 수 있는데, 멀티턴 대화에서 흔히 이전 라운드의 prompt를 통째로 재사용하기 때문이다. Marconi‑style은 한 번의 긴 prompt가 곧바로 대량의 state 캐시를 차지하는 것을 피한다.

흔한 의문: Interval‑based Retention과 Marconi‑style Retention의 차이는 무엇인가?

Interval‑based는 미리 정해진 이정표다. K token마다 checkpoint 하나, Prompt 끝은 추가로 보존한다. 안정적이고 예측하기 쉬우며, 멀티턴 대화처럼 경계가 명확한 시나리오에 적합하다. Marconi‑style은 먼저 관찰한 뒤 투자하는 방식이다. 처음 나타나면 접근만 기록하고 큰 state는 저장하지 않으며, 다시 적중하면 핫스팟임이 증명되므로 그제야 checkpoint를 생성한다. 한 번뿐인 prompt가 캐시를 낭비하는 것을 막을 수 있지만, 온라인 통계와 더 복잡한 퇴출 로직이 필요하다. 둘은 함께 쓸 수 있다. 구조화된 경계에는 interval을, 동적 핫스팟에는 selective retention을 쓴다.

파일 / 주요 함수 / 클래스 / 역할 vllm/models/kimi_k3/nvidia/kda.py KimiK3DeltaAttention.forward, _forward, _store_cache_checkpoints_kernel KDA conv/recurrent 실행, checkpoint 기록 vllm/models/kimi_k3/nvidia/kda_metadata.py _get_aligned_state_indices_kernel, KimiK3KDAMetadataBuilder.build block table을 state index로 변환, checkpoint metadata 구성 vllm/v1/core/single_type_kv_cache_manager.py MambaManager.find_longest_cache_hit Mamba/KDA의 prefix lookup과 세밀한 hit vllm/v1/core/kv_cache_coordinator.py KVCacheCoordinator.__init__, _cache_hit_alignment_tokens 다중 attention group 정렬, partial hit 활성화 vllm/v1/core/kv_cache_utils.py resolve_block_hashes 호출 지점 세밀한 hash를 유지하여 partial hit 지원

8.2.1 KimiK3DeltaAttention.get_kv_cache_spec

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda.py:593-599

함수: get_kv_cache_spec(vllm_config)

spec = super().get_kv_cache_spec(...) return replace( spec, num_prefill_checkpoint_blocks = int( self.kda_prefill_backend == "flashkda" ), )

KDA의 cache spec은 일반 attention KV spec만이 아니라 prefill checkpoint block 수도 포함한다.

8.2.2 forward와 _forward

파일:

- vllm-main/vllm/models/kimi_k3/nvidia/kda.py:601-638

- vllm-main/vllm/models/kimi_k3/nvidia/kda.py:640-710

KimiK3DeltaAttention.forward 601-638번 줄:

hidden_states를 projection하여 mixed_qkv, gates, beta를 분리 _forward 호출 o_proj

KimiK3DeltaAttention._forward 640-710번 줄:

- forward context에서 KDA metadata 획득;

- state_indices, checkpoint 읽기;

- self.kv_cache에서 conv_state와 recurrent_state 획득;

- 순수 decode 조건이 충족될 때 fused KDA decode 호출.

681-685번 줄에서 직접 언패킹:

conv_state, recurrent_state, * recoverssm_records = self.kv_cache

694-710번 줄에서 conv_state, recurrent_state, state_indices를 fused kernel에 전달.

vLLM 소스 코드는 KDA cache가 최소한 convolution state와 recurrent state 두 개의 독립적인 tensor 의미를 가진다는 것을 직접 증명한다.

8.2.3 prefill의 convolution과 recurrent state

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda.py:820-860

- 831-846번 줄에서 _prefill_conv 정의;

- 848-850번 줄에서 각각 q/k/v에 대해 causal convolution 수행;

- 858-860번 줄에서 gather_initial_states로 초기 recurrent state 획득.

간략화:

prefill input ├─ q -> causal conv -> q ├─ k -> causal conv -> k ├─ v -> causal conv -> v └─ gather initial recurrent state

8.2.4 _store_cache_checkpoints_kernel

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda.py:236-297

함수: _store_cache_checkpoints_kernel(...)

- 262-266번 줄에서 state index와 checkpoint offset 읽기;

- 267-285번 줄에서 해당 convolution 이력을 conv_state에 기록;

- 287-295번 줄에서 recurrent checkpoint를 recurrent_state에 기록.

vLLM의 checkpoint 기록은 추상적인 설명이 아니라 kernel 수준에서 동시에 처리한다:

conv checkpoint -> convolution state recurrent checkpoint -> recurrent state

8.3.1 _get_aligned_state_indices_kernel

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:55-101

함수: _get_aligned_state_indices_kernel(...)

85행은 seq_lens와 CACHE_BLOCK_SIZE에 따라 첫 번째 state slot을 계산하고, 89‑100행은 block table에서 state indices를 읽어온다.

그것이 하는 일은:

요청 길이 ↓ 해당 KDA state block의 시작 위치 ↓ state index tensor

8.3.2 _mamba_get_block_table_tensor

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:104-140

함수: _mamba_get_block_table_tensor(...)

mamba_cache_mode가 "all" 또는 "none"이면 110행에서 block table을 그대로 반환하고, 그렇지 않으면 114‑140행에서 state indices를 생성한 뒤 Triton kernel을 통해 KDA state slot에 정렬한다.

이것이 바로 "논리적 프리픽스"와 "물리적 state block" 사이의 적응 계층이다.

8.3.3 KimiK3DAMetadataBuilder.build

파일: vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:560-717

- 570‑575행: prefill이고, prefill checkpoint block이 활성화되었으며, cache mode가 align일 때만 checkpoint metadata를 준비한다;

- 588‑600행: sequence length, block size에 따라 checkpoint split과 checkpoint column을 계산한다;

- 602‑624행: offsets와 state indices를 KDACheckpointMetadata로 구성한다;

- 716행: checkpoint metadata를 최종 KimiK3DAMetadata에 넣는다.

이 소스는 중요한 경계를 보여준다. checkpoint는 모든 실행 모드에서 자동으로 생성되는 것이 아니며, prefill backend, cache mode, 정렬 조건에 의해 함께 제어된다.

8.4.1 MambaManager

파일: vllm-main/vllm/v1/core/single_type_kv_cache_manager.py:1268-1345

클래스: MambaManager

1268‑1295행:

- supports_fine_grained_hash_lookup = True

- mamba_cache_mode를 저장

- align 모드에서 checkpoint block 및 partial‑tail 관련 상태를 유지 관리한다.

8.4.2 find_longest_cache_hit

같은 파일 1297행부터: MambaManager.find_longest_cache_hit(...)

1315‑1321행은 resolve_block_hashes를 호출하며 다음을 전달한다:

supports_fine_grained_hash_lookup = True alignment_tokens =alignment_tokens

1327‑1345행은 alignment_tokens < block_size일 때, 하나의 큰 block을 여러 개의 더 작은 hash unit으로 쪼개어 히트를 찾는다.

이것이 바로 "물리적 block은 크고 논리적 prefix는 작다"는 것의 소스 기반이다.

8.4.3 Coordinator가 partial hit를 활성화하는 방법

파일: vllm-main/vllm/v1/core/kv_cache_coordinator.py:622-662

- 622‑637행은 MambaSpec이 존재하고 mamba_cache_mode == "align"인지 확인한다;

- 638‑651행은 다른 cache manager가 호환되는지 확인한다;

- 655‑662행은 partial hit 시 _cache_hit_alignment_tokens가 hash_block_size를 사용하도록 하고, 그렇지 않으면 scheduler_block_size를 사용하도록 한다.

다시 말해:

일반 히트에서는 scheduler block에만 떨어질 수 있고 KDA partial hit에서는 더 세밀한 hash block에 떨어질 수 있다

- vllm-main/vllm/v1/core/kv_cache_coordinator.py:32-62

- vllm-main/vllm/v1/core/kv_cache_coordinator.py:152-158

함수: _validate_prefix_cache_retention_interval(...)

40‑45행은 retention이 SlidingWindow와 Mamba/linear‑attention group에만 적용된다는 것을 설명하고, 56‑60행은 interval이 음수가 아니며 scheduler block의 배수일 것을 요구한다.

Coordinator 152‑158행은 retention_interval을 저장한다.

由此可知：

retention_interval = None: 희소화를 하지 않고 기본 밀도를 유지함 retention_interval = 0 : 적용 대상 group에 특수한 최소 유지 전략을 채택함 양수: scheduler block에 정렬된 주기로 유지함

소스 코드로 interval‑based retention을 확인할 수 있다. Marconi‑style selective retention을 계속 고증하려면, 구체적인 브랜치에서 그 독립적인 전략 구현을 찾아야 하며, coordinator의 interval 파라미터만 가지고 둘이 같은 함수라고 말해서는 안 된다.

prefix hash lookup ↓ MambaManager .find_longest_cache_hit (...) ↓ fine-grained alignment / state block lookup ↓ KimiK3DAMetadataBuilder .build (...) ↓ state_indices + checkpoint offsets ↓ KimiK3DeltaAttention ._forward (...) ├─ conv_state └─ recurrent_state ↓ kernel decode/prefill/update

핵심 증거 파일: single_type_kv_cache_manager.py:1297-1345 , kv_cache_coordinator.py:622-662 , kda_metadata.py:570-624 , kda.py:640-710 .

SWA (Sliding Window Attention)는 최근 윈도 내의 token만 필요로 한다. 윈도가 앞으로 이동하면 오래된 token은 자연히 빠져나간다. 만약 128 token마다의 윈도 경계를 저장한다면, 복구 시 기껏해야 하나의 윈도 사이에 있는 짧은 구간만 다시 계산하면 된다.

윈도 [t-W, …, t]

윈도가 앞으로 이동한 후, 오래된 token은 명확히 폐기된다. 윈도 경계 하나를 저장하는 것만으로도 인근의 국소 상태를 설명하기에 충분하다.

KDA의 recurrent state는 전체 이력을 압축한 결과다. 위치 n의 state는 최근 윈도만으로 결정되는 것이 아니라, 처음부터 n까지의 모든 전이와 쓰기에 의해 공동으로 결정된다:

S_n = M_n … M_1 S_0 + Σ_{ i = 1 }^{n} M_n … M_{ i + 1 } b_i

checkpoint가 하나 빠지면, 임의의 인근 token으로부터 완전한 state를 직접 도출할 수 없다. 더 이른 checkpoint로부터 모든 중간 업데이트를 다시 순차적으로 수행해야 한다.

게다가 KDA state는 대개 SWA의 윈도 경계보다 훨씬 크다. 고정 간격이 촘촘할수록 메모리 비용이 높아지고, 간격이 성길수록 히트 후 재계산이 길어진다. 이것이 KDA가 SWA보다 Prefix Cache를 만들기 더 어려운 근본적인 이유다.

흔한 의문: 왜 KDA는 SWA의 고정 간격 checkpoint를 쓸 수 없는가?

SWA의 상태는 유한한 하나의 윈도만을 나타내고, KDA state는 전체 프리픽스의 점화식 요약을 나타낸다. SWA는 최근 윈도 경계에서 계속 이어가며 기껏해야 짧은 구간의 윈도 계산만 보충하면 되지만, KDA는 checkpoint에서 히트 지점까지의 모든 token이 state를 바꿀 수 있으므로 반드시 순서대로 replay해야 한다. 게다가 KDA state 자체의 부피가 커서, 고정적으로 고빈도로 저장하는 비용은 SWA의 가벼운 경계 정보보다 훨씬 크다.

checkpoint 간격을 K라고 하자:

- 저장 비용은 대략 1/K에 따라 증가한다;

- 히트 후의 평균 재계산 길이는 대략 K에 따라 증가한다.

이것은 정확한 성능 공식은 아니지만 추세를 설명할 수 있다:

checkpoint 밀집 -> state 점유 높음, 복구 정밀, 재계산 적음 checkpoint 희소 -> state 점유 낮음, 복구 거침, 재계산 많음

가령 10k token마다 한 번만 저장한다고 하자:

checkpoint: 0 ——————— 10k ——————— 20k 공유 프리픽스: 0 ——————————————————— 15k

두 요청은 실제로 15k까지 공유하지만, 가장 가까운 합법적 복구 지점은 10k뿐이다. 시스템은 "히트한 15k token 텍스트"로부터 직접 복구할 수 없고, 다음처럼 해야만 한다:

S_10k 복구 replay 10k … 15k 그런 다음 새 요청을 계속함

원래는 5k 꼬리만 계산하면 되었는데, 이제 5k를 재계산해야 한다; 히트 경계가 더 불리하다면 낭비는 더 커진다.

따라서 진정한 목표는 "token prefix hit rate"를 최대화하는 것이而不是 아니라 다음을 최대화하는 것이다:

캐시 이익 − state 상주 비용 − 복구 재계산 비용 − 관리 복잡도

이것은 단순한 히트율 문제가 아니라, state 상주 비용 + replay 비용 + 요청 대기 시간 의 공동 최적화다.

흔한 의문: checkpoint 밀도의 딜레마를 어떻게 이해해야 하는가?

밀집 checkpoint는 메모리를 복구 속도와 교환하고, 희소 checkpoint는 메모리를 재계산 시간과 교환한다. 최적 전략은 프리픽스 재사용 확률, state 크기, 요청 길이, 메모리 예산에 따라 달라지며, 부하로부터 분리된 고정 최적 간격은 존재하지 않는다. 시스템 prompt, 다중 턴 대화, Agent 도구 정의는 보통 높은 우선순위로 저장할 가치가 있지만, 일회성의 무작위 긴 prompt는 저장할 가치가 반드시 있는 것은 아니다.

흔한 의문: checkpoint 간격을 키우면 어떤 문제가 있는가?

간격이 커지면, 프리픽스가 token 수준에서는 많이 히트하더라도 state 수준에서는 더 이른 위치까지만 복구할 수 있다. 그리하여 히트 이익이 replay에 잠식된다. 예를 들어 905k token을 공유하지만 최근 checkpoint가 900k에 있다면 여전히 5k를 재계산해야 하고, 800k에만 저장했다면 105k를 재계산해야 한다. token 히트율은 매우 높아 보이지만 실제 TTFT가 반드시 개선되지는 않는다.

크로스 인스턴스 재사용

전통적인 시스템은 가장 긴 매칭 token prefix를 찾은 다음 KV block을 이전하면 되었다. KDA는 동시에 다음을 만족할 것을 요구한다:

- MLA KV가 해당 경계에서 사용 가능할 것;

- KDA recurrent state가 해당 경계에서 사용 가능할 것;

- convolution state가 필요할 때 사용 가능하거나 안전하게 재구성될 수 있을 것;

- 모든 상태가 동일한 모델 버전, 동일한 tokenization, 동일한 프리픽스에서 비롯될 것.

최장 token prefix ↓ 공통의 합법적 경계 탐색 ↓ 연합 마이그레이션: MLA KV KDA recurrent state KDA convolution state ↓ 목표 인스턴스로 복원

서로 다른 출처의 캐시는 단순히 이어붙일 수 없다. 한 출처의 MLA KV에 다른 출처의 KDA state를 더하는 경우, 길이가 같더라도 계산 일관성이 보장되지 않는다. 인스턴스 A의 MLA KV와 인스턴스 B의 KDA state를 길이에 따라 이어붙일 수 없는 이유는, 이들이 서로 다른 tokenization, 모델 버전, 병렬 레이아웃 또는 prefix 경계에서 비롯되었을 수 있기 때문이다.

따라서 인스턴스 간 캐시 조회는 '최장 prefix 탐색'에서 '최장 연합 복원 가능 prefix' 탐색으로 바뀐다.

흔한 의문: 인스턴스 간 재사용에서 왜 최장 token prefix만 찾으면 안 되는가?

왜냐하면 해당 prefix에 대응하는 KDA state checkpoint도 찾아야 하고, 그것이 MLA KV와 같은 위치, 같은 버전, 같은 레이아웃에서 유효한지 확인해야 하기 때문이다. 서로 다른 인스턴스의 캐시는 길이만으로 단순히 이어붙일 수 없다. 인스턴스 간에는 '텍스트가 어디까지 같은가'뿐만 아니라, 어떤 경계에 KDA checkpoint가 있는가? MLA KV와 KDA state가 같은 경계에 있는가? 원격 객체가 같은 모델/병렬 구성에서 온 것인가? 복원된 state를 로컬 활성 복제본으로 안전하게 전환할 수 있는가?도 물어야 한다.

PD 분리

전통적인 PD 분리 전송:

Prefill → KV blocks → Decode

Prefill 측이 생성하는 것은 KV block 한 묶음이 아니라, 경계와 버전 정보를 지닌 모델 상태 한 묶음이다. Decode 측은 다음을 알아야 한다.

- 상태 유형;

- 대응하는 token 범위;

- tensor layout과 dtype;

- checkpoint가 쓰기 가능한지 여부;

- 이후 COW가 필요한지 여부;

- 상태가 현재 병렬 분할에 적용되는지 여부.

Prefill ├─ MLA latent history ├─ KDA recurrent state ├─ KDA convolution state └─ boundary/version/layout metadata ↓ Decode

이는 통신 프로토콜이 '캐시 페이지 운반'에서 '복원 가능한 모델 상태 운반'으로 업그레이드되어야 함을 요구한다. 모델마다 별도의 통신 프로토콜을 작성하면 인프라가 모델 구조에 의해 크게 오염된다. 통일된 모델 상태 프로토콜은 '객체 유형'과 '복원 메타데이터'를 일급 시민으로 만들 수 있다.

흔한 의문: PD 분리에서 왜 KV Cache 전송에서 모델 상태 전송으로 진화해야 하는가?

왜냐하면 Decode가 필요로 하는 입력이 더 이상 KV만이 아니라 MLA KV + recurrent state + convolution state이기 때문이다. 각 상태 유형은 서로 다른 레이아웃, 정밀도, 소유권 및 복원 규칙을 가질 수 있다. Prefill에서 Decode로의 입력은 이제 세 가지와 메타데이터를 포함한다. 이들을 각각 KV 프로토콜에 밀어 넣으면 통신 계층이 모델 세부 사항에 묶여 버리므로, 더 자연스러운 추상화는 '복원 가능한 모델 상태'를 통일적으로 전송하는 것이다.

멀티모달 EPD

멀티모달 추론은 Encoder가 출력한 시각 embedding도 데이터 플레인에 포함시켜야 한다.

Encoder embeddings ↓ Prefill이 MLA/KDA states 생성 ↓ Decode가 활성 복제본 사용

시각 embedding은 여러 요청 사이에서 재사용될 수도 있고, 현재 요청만을 위해 서비스될 수도 있다. 이는 KDA state가 아니며, KV block으로 간단히 취급할 수도 없다.

데이터 플레인은 시각 embedding, MLA token history, KDA recurrent snapshot, KDA convolution snapshot을 지원해야 하며, 각각의 수명 주기와 소유권을 유지·관리해야 한다.

통일된 데이터 플레인은 모든 내용을 KV로 위장하는 대신, 서로 다른 수명 주기, 서로 다른 크기, 서로 다른 공유 가능성을 지닌 객체를 동시에 처리해야 한다.

흔한 의문: 멀티모달 EPD의 추가적인 과제는 무엇인가?

Prefill이 생성한 어텐션 상태 외에도 Encoder의 시각 embedding이 흘러야 한다. 이들은 크기, 수명 주기, 재사용 방식이 서로 다르므로, 시스템은 유형과 소유권 정보를 유지하면서 통일된 데이터 플레인을 갖춰야 한다. EPD는 PD보다 Encoder 출력이 더해지며, 통일된 데이터 플레인은 유형, 소유권, 버전 및 복원 경계를 반드시 유지해야 하고, 단일 token prefix hash에만 의존할 수 없다.

개념 SGLang 소스 vLLM 소스 recurrent state MambaCheckpointPool.temporal self.kv_cache 내의 recurrent_state convolution state MambaCheckpointPool.conv self.kv_cache 내의 conv_state 활성 복제본 MambaPool /allocator slot request 대응 state slot checkpoint MambaCheckpointPool slot state block + metadata/checkpoint slot COW load_to_active 요청 자신의 state slot으로 복원 Donate donate_mamba_ping_pong_slot state block/slot 수명 주기 관리 세밀한 단위 적중 Radix/KDA tracking boundary supports_fine_grained_hash_lookup 희소 checkpoint schedule boundary + tracking interval retention interval / state retention 물리적 통합 메모리 UnifiedMemoryPool Hybrid cache manager / common block pool kernel 상태 인덱스 backend cache indices state_indices metadata

교육용 KDA Infra를 직접 구현한다면 다음 순서대로 읽을 것을 권한다.

- SGLang mamba_checkpoint_pool.py:117-174 : 먼저 양자화, 로딩, 쓰기 되돌림을 이해한다;

- SGLang memory_pool.py:1271-1289,1488-1509 : 활성 slot과 Donate를 이해;

- vLLM kda.py:640-710 : conv/recurrent state가 어떻게 kernel로 진입하는지 이해;

- vLLM kda_metadata.py:570-624 : checkpoint offsets/state indices를 이해;

- vLLM single_type_kv_cache_manager.py:1297-1345 : 세밀한(fine-grained) hit를 이해;

- SGLang unified_memory_pool.py:181-280 : 물리적 용량 공유를 이해.

소스 코드를 보면, KDA‑aware Infra의 최소한의 올바른 추상화는 prefix → KV blocks가 아니라 다음과 같다:

prefix boundary → token/hash identity → MLA KV blocks → recurrent state slot → convolution state slot → writable/readonly ownership → dtype/layout/version metadata

KDA가 AI Infra에 미치는 핵심 영향은 다음과 같이 요약할 수 있다:

Prefix Cache는 더 이상 "같은 token을 찾아 KV를 옮기는" 것이 아니다. 그것은 "공통 prefix 경계를 찾아, 의미상 일관되고 생명주기가 서로 다른 일련의 모델 상태를 복원하는" 것이 된다.

SGLang은 Radix Tree 위에 KDA‑aware checkpoint, 희소/압축 저장 및 공유 상태 격리를 추가하는 데 더 중점을 두며, 구현 초점은 "Radix + Mamba checkpoint pool + ping‑pong/COW + UnifiedKVPool"이다;

vLLM은 state block, 논리적 경계와 물리적 저장의 분리, 그리고 세밀한 partial hit 및 상태 보존 전략에 더 중점을 두며, 구현 초점은 "Hybrid cache manager + MambaManager fine‑grained hash alignment + KDA metadata/kernel pipeline"이다.

두 가지는 방향이 다르지만 공통된 결론은 일치한다: KDA를 겨냥한 추론 시스템은 반드시 KV‑centric cache에서 state‑centric cache로 진화해야 한다. 오직 같은 prefix 경계에서 완전한 모델 상태가 모두 위치 파악 가능하고, 복원 가능하며, 안전하게 쓸 수 있을 때에만 Prefix Cache가 진정으로 성립한다.

TransFormer-커버

본문은 markdown.com.cn으로 조판되었습니다

罗西的思考

440

글

430k

읽음

439

팬

출처 원문