[sglang] PD 시나리오에서 상세 캐시 히트 분류 수정

PR 링크: sgl-project/sglang#21764 상태: Merged | 변경: +16 / -8

들어가며

SGLang의 Prefill-Decode(PD) disaggregation 환경에서 캐시 히트 통계가 device/host/storage 레벨로 세분화되어 보고되어야 한다. 그러나 기존 코드에서는 decode 노드가 prefill 노드로부터 cached_tokens 메타데이터를 전달받을 때 device/host/storage 세부 정보가 누락되었다. 또한 HiCache가 비활성화된 환경에서도 기본 캐시 히트 정보를 보여줄 수 있도록 개선되었다.

핵심 코드 분석

1. PD 메타데이터 전달 수정

Before:

# decode.py - transfer commit 시
decode_req.req.cached_tokens = cached_tokens[0].item()
# device, host, storage 정보 누락

After:

decode_req.req.cached_tokens = cached_tokens[0].item()
decode_req.req.cached_tokens_device = cached_tokens[1].item()
decode_req.req.cached_tokens_host = cached_tokens[2].item()
decode_req.req.cached_tokens_storage = cached_tokens[3].item()

cached_tokens 텐서의 인덱스 1~3에 저장된 device/host/storage 정보를 개별 필드에 저장한다.

2. 비-HiCache 환경에서도 상세 정보 제공

Before:

def _get_cached_tokens_details(self, req):
    # HiCache가 비활성화되면 None 반환
    if not getattr(self, "enable_hierarchical_cache", False):
        return None

After:

def _get_cached_tokens_details(self, req):
    if req.cached_tokens_device > 0 or req.cached_tokens_host > 0 or ...:
        details = {"device": req.cached_tokens_device, "host": req.cached_tokens_host}
        if req.cached_tokens_storage > 0:
            details["storage"] = req.cached_tokens_storage
        return details

    # HiCache 없이도 기본 캐시 히트 정보 제공
    if req.cached_tokens > 0:
        return {"device": req.cached_tokens, "host": 0}
    return None

HiCache 활성화 여부와 무관하게, 실제 캐시 히트 데이터가 있으면 상세 분류를 반환한다.

왜 이게 좋은가

• PD 환경 모니터링: disaggregation 환경에서 어느 레벨(device/host/storage)의 캐시가 효과적인지 파악 가능
• 범용성 향상: HiCache 비활성화 환경에서도 기본 캐시 정보를 제공
• 데이터 무결성: 메타데이터 전달 누락으로 인한 0 보고 문제 해결

정리

PD disaggregation 환경에서 캐시 히트 통계의 device/host/storage 세분화 정보가 올바르게 전달되지 않던 버그를 수정하고, 비-HiCache 환경에서도 기본적인 캐시 히트 정보를 제공하도록 개선했다.

참고 자료

• sgl-project/sglang#21764 — 원본 PR

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