[sglang] SGLang의 Linear-Attention 성능 최적화: int8 체크포인트 풀 도입

PR 링크: sgl-project/sglang#28185 상태: Merged | 변경: +1010 / -18

들어가며

최근 LLM 서빙 환경에서 GDN(Gated DeltaNet)이나 KDA와 같은 Linear-attention 모델이 주목받고 있습니다. 이 모델들은 기존의 KV 캐시 대신 시퀀스당 하나의 압축된 recurrent state를 유지합니다. SGLang은 이 상태를 MambaRadixCache를 통해 Radix 트리에 캐싱하여 프리픽스 재사용을 지원합니다. 하지만 기존 설계에서는 이 상태를 bf16으로 저장하며, 실행 중인 요청(active set)과 캐시된 상태가 동일한 bf16 풀을 공유했습니다. 이로 인해 캐시 용량이 실행 중인 요청 수에 강하게 종속되어, 다양한 프리픽스를 사용하는 워크로드(Multi-doc RAG 등)에서 캐시 적중률이 급격히 떨어지는 'reuse-collapse' 문제가 발생했습니다.

코드 분석

1. hybrid_linear_attn_backend.py: 캐시 적중 시 dequantize 로직 추가

기존에는 단순히 copy_from을 통해 bf16 상태를 복사했으나, 이제는 int8 체크포인트 풀이 존재할 경우 이를 dequantize하여 활성 bf16 슬롯으로 로드합니다.

# Before
self.req_to_token_pool.mamba_pool.copy_from(
    forward_batch.mamba_cow_src_indices, forward_batch.mamba_cow_dst_indices
)

# After
ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
if ckpt_pool is not None:
    ckpt_pool.load_to_active(
        self.req_to_token_pool.mamba_pool,
        forward_batch.mamba_cow_src_indices,
        forward_batch.mamba_cow_dst_indices,
    )

2. invariant_checker.py: 이중 풀 구조에 맞춘 불변성 검사

풀이 분리됨에 따라, 기존의 단일 풀 기반 불변성 검사 로직을 int8 풀을 고려한 독립적인 검사 로직으로 확장했습니다.

# After
ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
if ckpt_pool is not None:
    return self._check_mamba_pool_with_int8(ps, ckpt_pool)

왜 이게 좋은가

1. 캐시 용량 2배 증대: bf16에서 int8로 상태를 압축하여 동일한 HBM 메모리 내에서 약 2배 더 많은 프리픽스를 캐싱할 수 있습니다. 이는 RAG와 같이 다양한 문서를 다루는 환경에서 캐시 적중률을 획기적으로 높입니다.
2. 메모리 효율성: 실행 중인 요청을 위한 bf16 풀과 캐시를 위한 int8 풀을 분리함으로써, 활성 풀은 동시성(concurrency)에 맞게 작게 유지하고 캐시 풀은 독립적으로 크게 확장할 수 있습니다.
3. 품질 안전성: 상태는 저장 시점에 단 한 번 양자화되고, 적중 시점에 단 한 번 역양자화됩니다. recurrence 루프 내에서 양자화가 반복되지 않으므로 오차 누적이 발생하지 않으며, gated-delta rule의 특성상 양자화 오차가 기하급수적으로 감쇠되어 결과 품질에 미치는 영향이 매우 작습니다.

성능 수치

• Reuse-collapse knee: 기존 256에서 512로 2배 확장.
• 성능: K=384 환경에서 기존 대비 약 26% 낮은 프리필 레이턴시와 33% 높은 처리량(throughput)을 기록.

리뷰어 피드백 반영

리뷰 과정에서 int8_checkpoint_store.py를 별도 파일로 두는 대신 mamba_checkpoint_pool.py로 통합하여 관리 효율성을 높였으며, maybe_init_int8_mamba_checkpoint_pool과 같은 헬퍼 함수를 통해 초기화 로직을 깔끔하게 정리했습니다. 또한 Qwen3.5-35B 모델을 활용한 E2E 테스트를 추가하여 실제 프로덕션 환경에서의 안정성을 검증했습니다.

참고 자료

• https://docs.sglang.ai/developer_guide/contribution_guide.html

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