[sglang] SGLang PD-Disaggregation 최적화: Mori 백엔드에서의 증분 KV 전송 구현

PR 링크: sgl-project/sglang#26288 상태: Merged | 변경: +64 / -9

들어가며

LLM 서빙에서 Prefill과 Decode 단계를 분리하는 PD-Disaggregation(Prefill-Decode Disaggregation) 아키텍처는 리소스 효율성을 극대화할 수 있는 강력한 도구입니다. 하지만 분산 환경에서 Prefill 노드가 Decode 노드로 KV 캐시를 전송할 때, 매 턴마다 전체 프리픽스(prefix)를 재전송하는 비효율이 발생하곤 합니다. 본 PR은 SGLang의 mori 전송 백엔드에 Radix Cache 기능을 통합하여, 이미 Decode 노드가 보유하고 있는 프리픽스를 제외한 '증분(incremental)' 데이터만 전송하도록 최적화했습니다.

코드 분석

1. python/sglang/srt/disaggregation/mori/conn.py: 전송 메타데이터 확장

핵심 변경 사항은 TransferInfo 클래스에 decode_prefix_len 필드를 추가하여, Decode 노드가 현재 보유 중인 KV 캐시의 길이를 Prefill 노드에 알리는 것입니다.

# Before
class TransferInfo:
    dst_state_indices: List[npt.NDArray[np.int32]]
    required_dst_info_num: int
    is_dummy: bool

# After
class TransferInfo:
    # ... 기존 필드 ...
    decode_prefix_len: Optional[int] = None

send_metadata 함수에서는 이 정보를 ZMQ 페이로드에 인코딩하여 전송합니다. 이를 통해 Prefill 노드는 Decode 노드가 어디까지 데이터를 가지고 있는지 파악하고, 필요한 페이지만 선택적으로 보낼 수 있게 됩니다.

2. python/sglang/srt/arg_groups/pd_disaggregation_hook.py: 설정 검증 로직 업데이트

--disaggregation-decode-enable-radix-cache 플래그가 mori 백엔드와 함께 사용될 수 있도록 검증 로직을 수정했습니다.

# After
if server_args.disaggregation_transfer_backend not in (
    "nixl",
    "mooncake",
    "mori",
):
    raise ValueError(
        "... requires --disaggregation-transfer-backend in "
        "('nixl', 'mooncake', 'mori'), but got "
        f"{server_args.disaggregation_transfer_backend!r}"
    )

왜 이게 좋은가

이 최적화는 네트워크 대역폭 사용량을 획기적으로 줄여줍니다. 특히 다중 턴(multi-turn) 대화형 워크로드에서 프리픽스가 반복적으로 전송되는 문제를 해결함으로써 TTFT(Time To First Token) 성능을 크게 향상시켰습니다.

성능 수치 (DeepSeek-R1 MXFP4 기준)

Round	Baseline	+ flag (이 PR)	개선율
0	2.39 s	2.15 s	10%
1	2.44 s	1.06 s	56%
3	1.24 s	0.78 s	37%
Avg	1.73 s	1.37 s	약 20.8%

교훈

1. 데이터 중복 제거: 분산 시스템에서 상태(State)를 동기화할 때, 전체 상태를 보내기보다 '차분(delta)'을 계산하는 것이 네트워크 병목을 해결하는 핵심입니다.
2. 백엔드 추상화: Radix Cache와 같은 상위 레벨의 캐싱 로직을 특정 전송 백엔드에 종속시키지 않고, 메타데이터 교환 프로토콜을 통해 범용적으로 확장 가능하게 설계하는 것이 중요합니다.

리뷰어 피드백

리뷰 과정에서 mori 백엔드 사용 시의 호환성 검증과 린트(lint) 수정이 이루어졌습니다. 특히 CI 단계에서 AMD MI355X 환경의 로그 404 오류가 있었으나, 해당 코드가 mori 백엔드 활성화 시에만 동작하므로 기존 테스트에 영향을 주지 않음이 확인되었습니다.

참고 자료

• https://github.com/sgl-project/sglang

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