[sglang] 하이브리드 모델의 Mamba 상태를 GreedyDual-Size로 축출하기: SGLang의 gds eviction policy 분석

PR 링크: sgl-project/sglang#49430 상태: Merged | 변경: +904 / -70

들어가며

Qwen3-Next나 GLM-5.3-Flash 같은 hybrid 모델은 attention layer 사이에 KDA / GDN / Mamba 같은 linear-attention(recurrent) layer를 섞어 씁니다. 이런 모델에서 radix cache의 prefix를 재사용하려면 KV만 남아 있어서는 안 되고, 그 지점의 recurrent state checkpoint도 함께 남아 있어야 합니다. recurrent state는 토큰 단위로 쪼갤 수 없는 누적 상태이기 때문입니다.

문제는 두 pool의 크기 차이입니다. PR 설명에 따르면 GLM-5.3-Flash NVFP4를 2x B200(TP2)에 올렸을 때 state slot은 419개, KV는 2,969,600 토큰이고, 실행 중인 요청 하나가 state slot을 5개씩 씁니다. 기존 Mamba LRU는 state를 하나씩 버리면서 KV는 그대로 둡니다. 그 결과 agent 트래픽에서 tool call을 마치고 돌아온 세션은 KV는 멀쩡히 남아 있는데 최신 state가 사라져 있는 경우가 많습니다. 이때는 더 오래된 state로, 최악에는 token 0으로 돌아가 한 번도 evict되지 않은 토큰들을 다시 prefill하게 됩니다.

이 PR은 FlashState 논문을 따라 recurrent state에 GreedyDual-Size(GDS) 축출 정책을 도입한 --radix-eviction-policy gds를 추가합니다.

GDS의 핵심 아이디어

각 state의 priority는 다음과 같이 정의됩니다.

priority = (마지막 사용 시점의 watermark) + (가장 가까운 얕은 state, 없으면 token 0에서 이 state까지 재구성하는 prefill 비용)

• priority가 가장 낮은 state를 먼저 evict합니다.
• watermark는 evict된 state의 priority까지 올라갑니다. 그래서 아무도 resume하지 않는 state도 결국 aging으로 밀려납니다.
• 다른 state 바로 아래에 있는 state는 재구성 비용이 작으므로 먼저 나갑니다.
• 긴 세션의 유일한 state는 비싸서 오래 남습니다.

LRU는 "얼마나 최근에 썼는가" 하나만 봅니다. GDS는 "이걸 잃으면 얼마나 다시 계산해야 하는가"를 recency와 함께 봅니다.

코드 분석

1. choices.py — 정책 등록

# Before
RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority", "tlru", "joint_lru"]

# After
RADIX_EVICTION_POLICY_CHOICES = [
    "lru",
    "lfu",
    "slru",
    "priority",
    "tlru",
    "joint_lru",
    "gds",
]

2. 비용 함수와 설정 검증 (evict_policy.py, kv_cache_hook.py)

GDSStrategy는 token a에서 b까지 재구성하는 비용을 다음 식으로 계산합니다.

cost(a, b) = d*(b - a) + e*(b^2 - a^2)/2e5

선형항 d는 토큰당 prefill 비용입니다. 이차항 e는 context가 길어질수록 attention 비용이 커지는 부분을 반영합니다. 기본값 d=1.862e-5, e=3.33e-6는 FlashState의 GLM-5.3-Flash fit이고, 문서에 적힌 대로 순서에는 d와 e의 비율만 영향을 줍니다. 그래서 튜닝할 knob이 사실상 하나뿐이라는 점이 장점입니다. 경쟁 PR(#46512, Marconi)처럼 recency weight를 따로 맞출 필요가 없습니다.

잘못된 계수는 서버 시작 시점에 바로 실패하도록 했습니다.

if cfg.radix_eviction_policy == "gds":
    from sglang.srt.mem_cache.evict_policy import GDSStrategy

    # Building the strategy checks its cost coefficients.
    GDSStrategy(**(cfg.radix_eviction_policy_config or {}))

전략 객체를 한 번 만들어 보는 것으로 검증을 대신합니다. 설정 오류가 첫 eviction 순간이 아니라 startup에서 드러나므로 운영 관점에서 좋은 패턴입니다.

3. MambaComponent / components/mamba.rs — lazy heap walk

가장 흥미로운 부분입니다. GDS의 gain은 "가장 가까운 얕은 state"에 의존합니다. 그래서 같은 walk 안에서 조상 state가 evict되면 자손 state의 재구성 비용이 커집니다. 매번 전체 priority를 다시 계산하면 비쌉니다. PR은 이를 lazy re-evaluation으로 처리합니다.

• 마지막으로 알려진 priority로 heap을 seed합니다.
• pop한 state의 gain을 현재 기준(가장 가까운 얕은 state)으로 다시 계산합니다.
• gain이 올랐다면 heap에 다시 push하고, 아니면 evict합니다.
• state가 어떤 노드에 새로 set되면 그 노드와 바로 아래의 state들은 cached gain을 버립니다. 그래서 모든 seed 값이 항상 lower bound로 유지됩니다.

lower bound 불변식이 핵심입니다. heap에 들어 있는 값이 실제 priority보다 작거나 같다고 보장되면, pop 후 재계산한 값이 여전히 최소일 때만 evict하면 됩니다. 이렇게 해도 정확성이 깨지지 않습니다. Lazy Dijkstra나 lazy deletion heap과 같은 원리입니다.

또한 walk의 계약은 LRU walk와 같게 유지됩니다. device leaf는 caller에게 돌려주고, internal state는 한 step으로 처리하며, write-back backup 경로는 두 walk가 공유하는 helper로 옮겼습니다. 이 변경이 Python core와 Rust core 양쪽에 모두 들어가서, Linux 기본 Rust tree core를 그대로 쓸 수 있습니다.

watermark stamp는 state가 cache될 때, match될 때, host에서 load될 때, 마지막 요청이 release할 때 찍힙니다.

4. 방어적 기본값: gds = None

리뷰 과정에서 CPU CI가 실패했습니다. 원인은 일부 테스트와 MLX auxiliary-state component가 MambaComponent.__init__을 거치지 않고 객체를 만든다는 점이었습니다. 작성자는 class-level에 gds = None을 두는 방식으로 해결했습니다. 이 경우 __init__을 건너뛴 subclass는 자연스럽게 LRU walk로 동작합니다. 작은 수정이지만, 상속 구조가 복잡한 코드에 새 attribute를 추가할 때 참고할 만한 교훈입니다.

5. 기본값 해석 (overrides.py, memory.py)

# Before
radix_eviction_policy: A[str, Arg(...)] = "lru"
_radix_eviction_policy_explicitly_set: A[bool, Arg(no_cli=True)] = False

# After
radix_eviction_policy: A[Optional[str], Arg(..., fallback="lru")] = None
_radix_eviction_policy_explicitly_set: A[Optional[bool], Arg(no_cli=True)] = None

if (
    view.radix_eviction_policy == "lru"
    and not view._radix_eviction_policy_explicitly_set
    and not view.enable_session_radix_cache
):
    # gds rejects the session radix cache on recurrent-state models.
    declared["radix_eviction_policy"] = "gds"

기존에는 argv를 스캔해서 사용자가 lru를 명시했는지 판단했습니다. PR은 이 방식을 버리고 기본값을 None으로 두었습니다. handle_missing_default_values에서 "명시 여부"를 기록한 뒤 lru로 채웁니다. 덕분에 CLI든 config 파일이든 Engine(...) 키워드든 ServerArgs(...) 직접 생성이든 같은 의미를 가집니다. argv 스캔이 놓치던 경로까지 커버하는 깔끔한 개선입니다.

왜 이게 좋은가

성능 수치 (GLM-5.3-Flash NVFP4, 2x B200, lru → gds)

Sessions	Hit rate	TTFT p90 (s)	Output tok/min/GPU
16	0.913 → 0.931	0.69-1.01 → 0.56-0.62	25.6k → 26.0k
24	0.904 → 0.926	1.14-1.26 → 0.73-0.75	35.8k → 37.0k
32	0.845 → 0.857	2.17-2.21 → 1.82-1.94	37.5k → 38.7k
40	0.36-0.43 → 0.45-0.47	범위 겹침	19.9-22.0k → 22.5-23.0k

1632 세션에서 p90 TTFT 평균이 28%, 38%, 14% 줄었고, run 간 범위가 겹치지 않습니다. CPU 시뮬레이션에서는 "KV는 남았는데 state가 없어서 놓친 토큰"이 64세션 기준 요청당 3,819개에서 281개로 줄었습니다. 반면 median uncached token은 줄지 않습니다. 개선이 mean과 tail에 집중된다는 뜻입니다. 15%의 불운한 요청이 60~85K 토큰을 다시 prefill하던 상황을 막는 효과입니다.

64세션처럼 KV pool 자체가 포화된 구간에서는 차이가 없습니다. state 축출 순서로는 이미 잃어버린 KV를 되살릴 수 없기 때문입니다. 작성자가 이 한계를 정직하게 보고한 점도 눈에 띕니다.

리뷰 논의: 기본값 변경의 무게

리뷰어 BBuf는 처음에 두 가지를 요청했습니다. 하나는 반복 측정입니다. 40세션에서 p90이 오히려 나빠진 단일 run이 있었기 때문입니다. 다른 하나는 기본값을 바꿀 계획입니다. 작성자는 16~40세션 반복 run을 추가했고, hybrid 모델의 기본값을 gds로 바꿨습니다. 그런데 병합 후 BBuf가 직접 revert(#50118)했습니다. 이유는 다음과 같습니다.

switching the default eviction policy of every hybrid (Mamba / GDN / KDA) model on one model and one trace is too big a step.

결국 #50122에서 opt-in(--radix-eviction-policy gds)으로만 재제출되었고, eviction 로직 자체는 그대로 유지되었습니다. 좋은 최적화라도 기본값 변경에는 더 넓은 모델과 워크로드 검증이 필요하다는, 대규모 서빙 프레임워크다운 판단입니다.

일반적 교훈

1. 축출 비용은 균일하지 않습니다. 캐시 항목마다 재구성 비용이 다르면 LRU보다 cost-aware 정책(GDS 계열)이 유리합니다.
2. 의존 관계가 있는 비용은 lazy heap과 lower bound 불변식으로 관리합니다. 전체를 재계산하지 않고도 정확한 순서를 얻을 수 있습니다.
3. watermark aging은 비싼 항목이 영원히 남는 cache pollution을 막아 줍니다.
4. "명시적으로 설정했는가"는 값과 별도로 기록하세요. None default와 provenance flag 조합은 argv 스캔보다 견고합니다.
5. 새 정책은 opt-in으로 시작하세요. 기본값 변경은 별도의 근거가 필요한 결정입니다.

참고 자료

• heapq — Heap queue algorithm — GDS walk의 lazy priority heap 구현 원리
• SGLang Radix Eviction Policy docs — gds 정책이 추가된 공식 문서
• FlashState code — 이 PR이 따른 GreedyDual-Size recurrent state 축출 구현

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.