[vllm] MORI KV Connector - ROCm 기반 Prefill-Decode Disaggregation

PR 링크: vllm-project/vllm#29304 상태: Merged | 변경: +3,369/-3

들어가며

Prefill-Decode(PD) Disaggregation은 prefill과 decode를 별도 GPU 인스턴스에서 실행하는 아키텍처다. prefill 인스턴스가 생성한 KV cache를 decode 인스턴스로 전송해야 하는데, 이 PR은 ROCm 환경에서 MORI 라이브러리를 사용하는 KV 전송 커넥터를 구현한다.

핵심 코드 분석

Proxy 서버 기반 라우팅

# 서비스 디스커버리
def _listen_for_register(hostname, port):
    router_socket = context.socket(zmq.ROUTER)
    # prefill/decode 인스턴스 등록 수신
    if data["type"] == "register" and data["role"] == "P":
        _append_whole_dict_unique(prefill_instances, data)
    elif data["type"] == "register" and data["role"] == "D":
        _append_whole_dict_unique(decode_instances, data)

ZeroMQ 기반의 프록시 서버가 prefill과 decode 인스턴스를 관리하고, 요청을 적절한 인스턴스로 라우팅한다.

Transfer Mode 일관성 검증

if TRANSFER_TYPE is None:
    TRANSFER_TYPE = transfer_mode
elif transfer_mode != TRANSFER_TYPE:
    raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")

모든 인스턴스가 동일한 전송 모드를 사용하는지 검증하여, 설정 불일치로 인한 오류를 방지한다.

왜 이게 좋은가

1. ROCm PD Disaggregation: AMD GPU에서도 prefill-decode 분리 아키텍처 사용 가능
2. MORI 통합: ROCm의 네이티브 통신 라이브러리를 활용한 고효율 KV 전송
3. 서비스 디스커버리: 동적으로 인스턴스를 등록/해제하여 탄력적 스케일링 지원
4. 프로덕션 준비: 프록시 서버, 헬스 체크, 에러 처리가 포함된 완성도 높은 구현

정리

3,369줄의 대규모 PR로, ROCm 환경에서의 PD disaggregation을 위한 완전한 인프라를 구축한다. AMD GPU 클러스터에서 대규모 LLM 서빙을 위한 핵심 기능이다.

참고 자료

• PR #29304
• MORI (ROCm)

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.