[sglang] SGLang NIXL 이기종 TP 환경에서 디스어그리게이션 KV 캐시 전송 버그 수정 및 성능 개선
PR 링크: sgl-project/sglang#22145 상태: Merged | 변경: +None / -None
들어가며
대규모 언어 모델(LLM) 서빙에서 효율적인 KV 캐시 관리는 성능과 직결되는 중요한 문제입니다. 특히, 추론 시 GPU 자원을 효율적으로 활용하기 위해 모델 병렬화(Tensor Parallelism, TP)를 적용할 때, 사전 생성(Prefill) 단계와 디코딩(Decode) 단계에서 서로 다른 TP 크기를 사용하는 이기종(Heterogeneous) TP 설정은 복잡성을 더합니다.
SGLang은 이러한 이기종 TP 환경을 지원하기 위해 디스어그리게이션(Disaggregation) 아키텍처와 NIXL 백엔드를 도입했습니다. 하지만 최근 sgl-project/sglang 레포지토리의 PR #774 에서는 NIXL 백엔드에서 이기종 TP 설정 시, 특히 비-MLA(Non-Multi-Layer Attention) 모델에서 KV 캐시 전송 중 발생하는 심각한 버그 두 가지를 발견하고 이를 수정했습니다. 이 버그들은 디코딩 단계에서 무한정 행(hang)을 유발하여 서빙이 불가능한 상태로 만들었습니다.
본 글에서는 해당 PR에서 어떤 문제들이 발견되었고, 어떻게 해결되었는지, 그리고 이 변경이 왜 중요한 최적화인지 코드 변경 사항을 중심으로 상세히 분석해보겠습니다.
코드 분석
이번 PR은 주로 python/sglang/srt/disaggregation/nixl/conn.py 파일의 두 함수, send_kvcache_slice와 _process_kvcache_transfer를 수정했습니다. 각각의 변경 사항을 살펴보겠습니다.
1. send_kvcache_slice 함수의 헤드 분산 로직 개선
send_kvcache_slice 함수는 사전 생성 TP에서 디코딩 TP로 KV 캐시 데이터를 전송하는 역할을 담당합니다. 이 함수는 특히 GQA(Grouped-Query Attention)와 같이 헤드 수가 TP 크기보다 작을 수 있는 경우, KV 캐시 헤드의 정확한 분산 및 복제 처리가 중요합니다.
변경 전 (Before)
- num_kv_heads = self.kv_args.kv_head_num
-
- # Calculate head distribution
- src_heads_per_rank = num_kv_heads
- dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size
-
src_kv_item_len = self.kv_args.kv_item_lens[0]
page_size = self.kv_args.page_size
bytes_per_head_slice_to_send = (
dst_kv_item_len // page_size // dst_heads_per_rank
)
# Determine which heads to send
if prefill_tp_size > decode_tp_size:
# Multiple prefill ranks to one decode rank
src_head_start_offset = 0
num_heads_to_send = src_heads_per_rank
- dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank
+ dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank
변경 전 코드에서는 self.kv_args.kv_head_num을 기준으로 헤드 수를 계산했습니다. 이 값은 종종 각 TP 랭크별로 할당된 헤드 수(total_kv_heads // tp_size)이며, 특히 total_kv_heads < tp_size인 경우 정보 손실이 발생할 수 있습니다. 또한, GQA 복제(replication)를 고려하지 않아 dst_head_start_offset 계산 시 부정확성이 있었습니다.
변경 후 (After)
+ # Use total KV head count (not per-rank) for correct head distribution.
+ # Per-rank kv_head_num is max(1, total//tp) which loses info when total < tp.
+ total_kv_heads = getattr(self.kv_args, "total_kv_head_num", 0)
+ if total_kv_heads <= 0:
+ total_kv_heads = self.kv_args.kv_head_num * prefill_tp_size
+
+ src_heads_per_rank = max(1, total_kv_heads // prefill_tp_size)
+ dst_heads_per_rank = max(1, total_kv_heads // decode_tp_size)
+
bytes_per_head_slice_to_send = (
dst_kv_item_len // page_size // dst_heads_per_rank
)
+ # GQA replication: how many prefill ranks share the same KV head
+ src_replication = max(1, prefill_tp_size // total_kv_heads)
+
# Determine which heads to send
if prefill_tp_size > decode_tp_size:
# Multiple prefill ranks to one decode rank
src_head_start_offset = 0
num_heads_to_send = src_heads_per_rank
- dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank
+ unique_head_idx = local_tp_rank_in_group // src_replication
+ dst_head_start_offset = (
+ unique_head_idx * src_heads_per_rank
+ ) % dst_heads_per_rank
else:
# Send KVCache from 1 prefill instance to multiple decode instances
src_head_start_offset = (
변경 후에는 total_kv_head_num을 사용하여 헤드 수를 정확하게 계산합니다. getattr을 사용하여 해당 값이 없을 경우 이전 로직(kv_head_num * prefill_tp_size)으로 폴백(fallback)하도록 안전장치를 마련했습니다. 또한, GQA 복제를 고려하여 src_replication 값을 계산하고, 이를 바탕으로 unique_head_idx를 구하여 dst_head_start_offset을 정확하게 계산합니다. 이는 여러 사전 생성 랭크가 동일한 KV 헤드를 공유하는 경우에도 올바른 헤드 범위를 지정할 수 있게 합니다.
리뷰어 YAMY1234의 코멘트에서도 언급되었듯이, 이 로직은 Mooncake 백엔드의 구현과 동일하게 유지되었습니다. 이는 SGLang 프로젝트 내에서 백엔드 간 일관성을 유지하려는 노력의 일환으로 보입니다.
2. _process_kvcache_transfer 함수의 알림 태그(Notification Tag) 수정
_process_kvcache_transfer 함수는 KV 캐시 전송 완료를 알리는 알림(notification)을 처리합니다. 이 알림 태그는 각 전송 작업의 고유성을 보장하는 데 사용됩니다. 이전에는 pp_rank (Pipeline Parallelism rank)를 태그에 포함시켰으나, PP=1인 경우 모든 사전 생성 랭크가 동일한 pp_rank=0을 공유하게 되어 문제가 발생했습니다.
변경 전 (Before)
- notif = f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.pp_rank}"
+
notif = (
f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.engine_rank}"
)
변경 후 (After)
- f"{req.room}_state_{self.kv_args.pp_rank}",
+ f"{req.room}_state_{self.kv_args.engine_rank}",
변경 후에는 pp_rank 대신 engine_rank를 알림 태그에 사용하도록 수정되었습니다. engine_rank는 각 엔진 인스턴스의 고유한 랭크 ID를 나타내므로, PP=1 상황에서도 각 사전 생성 랭크를 명확히 구분할 수 있습니다. 이를 통해 TransferStatus.received_kvs_per_pp에 올바른 키가 기록되고, is_done() 함수가 정상적으로 True를 반환하여 디코딩 단계로의 진행이 가능해집니다. 이 수정은 특히 PP=1 구성에서 발생하는 무한 행(hang) 문제를 직접적으로 해결합니다.
왜 이게 좋은가
이번 PR의 변경 사항은 다음과 같은 이유로 중요한 최적화 및 버그 수정입니다.
- 안정성 확보: 가장 큰 이점은 이기종 TP 환경에서 발생하던 치명적인 무한 행(hang) 버그를 해결했다는 점입니다. 이전에는 특정 설정에서 디코딩 단계로 진행 자체가 불가능하여 서빙이 불가능했습니다. 수정 후에는 Qwen3-32B 모델을 GB200에서
1P4D(Prefill TP4 → Decode TP1x4) 구성으로 테스트했을 때, 이전에는 무한정 행(hang)되던 것이 정상적으로 완료되었습니다. GSM8K 8-shot 평가에서 0.961의 높은 점수와 568,960 token/s의 처리량을 기록하며 기능적으로 문제가 없음을 입증했습니다. - 정확성 향상:
send_kvcache_slice함수의 헤드 분산 로직 개선은 GQA와 같이 헤드 수가 TP 크기보다 작은 모델에서 KV 캐시 전송의 정확성을 높였습니다. 이는 모델의 추론 정확도를 유지하는 데 필수적입니다. - 이기종 TP 지원 강화: 이 PR은 SGLang이 NIXL 백엔드에서 이기종 TP 설정을 더 안정적으로 지원할 수 있게 합니다. 이는 다양한 하드웨어 구성과 모델 아키텍처에 맞춰 최적의 성능을 추구하는 데 중요한 기반이 됩니다.
- 코드 일관성: Mooncake 백엔드의 검증된 로직을 참고하여
send_kvcache_slice함수를 수정함으로써, SGLang 프로젝트 내 여러 백엔드 간의 코드 일관성을 유지하고 잠재적인 버그 재발을 방지했습니다.
일반적인 교훈
- 분산 시스템에서의 고유 식별자: 분산 시스템에서는 각 컴포넌트나 통신 단위를 고유하게 식별할 수 있는 식별자(예:
engine_rank)를 사용하는 것이 매우 중요합니다. 공유되는 식별자(pp_rank)는 특정 조건에서 충돌을 일으켜 전체 시스템의 오작동을 유발할 수 있습니다. - 헤드 분산 로직의 복잡성: 특히 GQA와 같은 어텐션 메커니즘에서는 헤드 수와 TP 크기 간의 관계가 복잡할 수 있습니다.
total_kv_heads와 같이 전체적인 정보를 기반으로 로직을 설계하고,max(1, ...)와 같은 안전장치를 활용하여 엣지 케이스를 처리해야 합니다. - 이기종 설정의 어려움: 사전 생성과 디코딩 단계에서 다른 TP 설정을 사용하는 것은 성능 최적화에 유리할 수 있지만, 데이터 전송 및 동기화 로직이 훨씬 복잡해지므로 철저한 테스트와 검증이 필요합니다.
결론
PR #774는 SGLang의 NIXL 백엔드에서 이기종 TP 환경 설정 시 발생하던 심각한 버그를 성공적으로 해결했습니다. 알림 태그의 고유 식별자 문제와 KV 캐시 헤드 분산 로직의 부정확성을 수정함으로써, 디스어그리게이션 서빙의 안정성과 정확성을 크게 향상시켰습니다. 이는 다양한 하드웨어 및 모델 구성에서 SGLang의 활용도를 높이는 중요한 진전이라고 할 수 있습니다.
참고 자료
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [sglang] sglang, Qwen3.5-397B FP8 모델 성능 벤치마크 추가 및 CI 개선
- 현재글 : [sglang] SGLang NIXL 이기종 TP 환경에서 디스어그리게이션 KV 캐시 전송 버그 수정 및 성능 개선
- 다음글 [llm-compressor] Gemma4 MoE 모델 양자화를 위한 llm-compressor 지원 추가 분석
댓글