[sglang] SGLang의 긴 문맥 처리 최적화: fill_ids 재구성 오버헤드 줄이기

PR 링크: sgl-project/sglang#27965 상태: Merged | 변경: +363 / -66

들어가며

LLM 추론 엔진인 SGLang에서 긴 문맥(Long Context)을 처리할 때, 매 단계마다 fill_ids 리스트를 전체 재구성(reconstruction)하고 참조 카운트를 관리하는 작업은 상당한 오버헤드를 발생시킵니다. 특히 요청이 길어질수록 이 비용은 선형적으로 증가하여 전체 추론 속도에 병목이 됩니다. 본 PR은 fill_ids를 매번 새로 생성하는 대신, 필요한 부분만 점진적으로 업데이트하고 RadixKey에 limit 개념을 도입하여 메모리 복사를 최소화하는 최적화를 수행했습니다.

코드 분석

1. RadixKey의 최적화 (python/sglang/srt/mem_cache/radix_cache.py)

기존에는 token_ids를 슬라이싱하여 새로운 배열을 생성하는 방식이었습니다. 이를 limit 필드를 추가하여 원본 배열을 유지하면서 논리적으로만 길이를 제한하도록 변경했습니다.

Before:

# 슬라이싱을 통해 매번 새로운 객체 생성
token_ids_to_match = self.full_untruncated_fill_ids[: self._compute_max_prefix_len(input_len)]

After:

# limit을 사용하여 원본 배열 참조 유지
key_limit: Optional[int] = self._compute_max_prefix_len(input_len)
# ...
RadixKey(token_ids=token_ids_to_match, limit=key_limit)

2. fill_ids 점진적 업데이트 (python/sglang/srt/managers/schedule_batch.py)

_refresh_fill_ids 메서드를 도입하여, 전체를 다시 합치는 대신 새로운 토큰만 extend하도록 로직을 개선했습니다.

Before:

self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids

After:

def _refresh_fill_ids(self) -> None:
    n_have_output = len(self.full_untruncated_fill_ids) - len(self.origin_input_ids)
    if self.full_untruncated_fill_ids is not self.origin_input_ids and 0 <= n_have_output <= len(self.output_ids):
        self.full_untruncated_fill_ids.extend(self.output_ids[n_have_output:])
    else:
        self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids

왜 이게 좋은가

이 최적화의 핵심은 O(N) 복사 연산을 O(1) 혹은 O(추가된 토큰 수) 연산으로 줄였다는 점입니다. 긴 문맥에서 매번 수천 개의 토큰을 복사하는 것은 CPU 사이클 낭비와 메모리 할당 오버헤드를 유발합니다. 특히 RadixKey에서 limit을 사용함으로써, 캐시 매칭 시점에 불필요한 데이터 복사 없이 원본 데이터의 일부만 안전하게 참조할 수 있게 되었습니다.

일반적 교훈

1. 불변성 vs 성능: 데이터의 불변성을 유지하기 위해 매번 복사본을 만드는 것보다, 논리적 경계(limit)를 두어 원본을 공유하는 것이 고성능 시스템에서는 유리할 수 있습니다.
2. 점진적 업데이트: 상태를 전체 재계산하는 대신, 변경된 차이(delta)만을 반영하는 구조로 설계하면 대규모 데이터 처리 시 성능 이득이 큽니다.

리뷰 과정에서 ispobock님은 aliasing 문제와 길이 불일치 상황에 대해 예리한 질문을 던졌고, 이에 대해 Qiaolin-Yu님은 _refresh_fill_ids 내에서 안전하게 폴백(fallback) 로직을 구현하여 데이터 무결성을 보장했습니다.

참고 자료

• https://docs.python.org/3/library/array.html
• https://github.com/sgl-project/sglang

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