본문으로 건너뛰기

[sglang] SGLang에서 DeepSeek V3.2를 위한 IndexCache 최적화 구현

PR 링크: sgl-project/sglang#21405 상태: Merged | 변경: +None / -None

들어가며

최근 대규모 언어 모델(LLM)의 추론 성능 최적화는 엔지니어링의 핵심 과제입니다. 특히 DeepSeek V3.2와 같은 모델은 복잡한 구조를 가지고 있어, 연산 효율을 높이기 위한 다양한 기법이 요구됩니다. 이번 SGLang의 PR(#21286)은 THUDM의 IndexCache 기법을 포팅하여, 특정 레이어에서 Top-K 인덱스 계산을 건너뛰고 이전 레이어의 결과를 재사용함으로써 추론 속도를 획기적으로 개선했습니다.

코드 분석

1. forward_mla.py의 로직 변경

핵심 변경 사항은 forward_absorb_prepare 함수 내에서 topk_indices를 계산하는 방식의 조건부 분기입니다. 기존에는 모든 레이어에서 인덱서를 호출했으나, 이제는 prev_topk_indices를 전달받아 재사용할 수 있게 되었습니다.

# Before
topk_indices = self.indexer(x=hidden_states, q_lora=q_lora, ...)

# After
if not self.skip_topk or prev_topk_indices is None:
    topk_indices = self.indexer(x=hidden_states, q_lora=q_lora, ...)
else:
    topk_indices = prev_topk_indices

2. deepseek_v2.py에서의 레이어 제어

모델의 각 레이어는 skip_topk와 next_skip_topk 속성을 통해 인덱스 캐싱 여부를 결정합니다. index_topk_freq 파라미터를 통해 특정 빈도로만 인덱스를 계산하도록 설정하여 연산량을 줄였습니다.

# 레이어별 인덱스 계산 스킵 로직
if self.index_topk_pattern is None:
    self.skip_topk = max(layer_id - 1, 0) % self.index_topk_freq != 0
    self.next_skip_topk = layer_id % self.index_topk_freq != 0

왜 이게 좋은가

이번 최적화는 연산 집약적인 Top-K 인덱싱 과정을 레이어 간에 공유함으로써 성능 이득을 얻습니다. 벤치마크 결과, 처리량(Throughput)은 약 6.4% 향상되었으며, TTFT(Time to First Token)와 TPOT(Time per Output Token) 또한 각각 5.4%, 5.5% 개선되었습니다.

특히 요청 속도(Request Rate)가 높아질수록 그 효과는 더욱 두드러집니다. 요청 속도가 20인 환경에서 처리량은 최대 11.39%까지 향상되는 모습을 보였습니다. 이는 고부하 환경에서 모델 서빙의 안정성을 크게 높여줍니다.

일반적 교훈

  1. 중복 연산 제거: 모델의 모든 레이어에서 동일한 연산을 수행할 필요가 없는 경우, 캐싱을 통해 연산 그래프를 최적화할 수 있습니다.
  2. 데이터 흐름 최적화: 레이어 간에 중간 상태(Top-K indices)를 전달하는 구조를 설계함으로써 전체 추론 파이프라인의 병목을 해소할 수 있습니다.
  3. 벤치마크의 중요성: 다양한 요청 속도(Request Rate) 환경에서 테스트를 수행하여, 실제 프로덕션 환경에서의 성능 변화를 예측하는 것이 필수적입니다.

결론

이번 PR은 단순히 코드를 포팅하는 것을 넘어, SGLang의 유연한 모델 서빙 구조를 활용하여 최신 모델의 성능을 극대화한 좋은 사례입니다. 특히 리뷰 과정에서 제기된 레이어 스킵 로직의 정교화(Modulo 연산 적용 등)는 코드의 완성도를 높이는 데 기여했습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글