[sglang] MiniMax-H3 모델의 AdaLN 추론 최적화: Pinned-Host 캐싱과 LRU 전략

PR 링크: sgl-project/sglang#37266 상태: Merged | 변경: +1621 / -335

들어가며

SGLang의 MiniMax-H3 모델은 --minimax-h3-adaln-online 모드 사용 시, 매번 24.2 GiB에 달하는 adaln_proj 가중치를 체크포인트에서 다시 읽어와야 했습니다. 이로 인해 요청당 약 5.8~6.7초의 GPU 유휴 시간이 발생하며, 캐시 슬래브 오버플로우 시 전체 리셋이 일어나는 비효율적인 구조였습니다. 본 PR은 per-plan LRU(Least Recently Used) 캐싱과 Pinned-Host 메모리 계층을 도입하여, 이전에 빌드된 계획을 PCIe를 통해 빠르게 스왑함으로써 반복적인 스케줄 요청 시 성능을 크게 향상시켰습니다.

코드 분석

1. GPU Slab 및 Host Cache 도입 (python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py)

기존에는 GPU 메모리 내에서만 캐싱을 시도했으나, 이제는 호스트 메모리를 활용한 계층적 캐싱을 지원합니다.

# Before: GPU Slab에만 의존
# After: Host Cache 크기 설정 및 정밀도 제어 추가
init_params["adaln_max_plans"] = envs.SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS
init_params["adaln_host_cache_bytes"] = int(
    component_server_args.minimax_h3_adaln_host_cache_gb * 1e9
)
init_params["adaln_precision"] = (
    "fp32" if envs.SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32 else "match"
)

2. 가중치 업데이트 무결성 보장 (python/sglang/multimodal_gen/runtime/models/dits/base.py)

캐시 모드에서 가중치가 변경될 경우, 기존의 캐시된 AdaLN 계획이 유효하지 않게 되는 문제를 해결하기 위해 validate_weight_update_source 훅을 추가했습니다.

# After: 가중치 업데이트 전 검증 로직 추가
def validate_weight_update_source(self, *, weights_path: str | None) -> None:
    """Reject a weight update this model cannot stay coherent under."""
    return None

왜 이게 좋은가

이번 최적화의 핵심은 재구축(Rebuild) 비용의 제거입니다. 1x H200 환경에서 반복되는 스케줄 요청 시, 기존 94.8초가 소요되던 작업이 88.4초로 단축되었습니다. 이는 약 6.4초의 지연 시간을 절감한 것으로, nsys로 측정한 재구축 비용과 정확히 일치합니다.

일반적 교훈

1. 계층적 캐싱(Tiered Caching): GPU 메모리 부족 시 전체를 리셋하기보다, Host 메모리를 백업으로 활용하여 PCIe 대역폭을 통해 데이터를 복구하는 것이 훨씬 효율적입니다.
2. Fail-Closed 설계: 가중치 업데이트와 같이 캐시 무결성을 해칠 수 있는 작업은, 변경이 발생하기 전에 미리 검증(Validation)하여 일관성을 유지해야 합니다.
3. LRU 전략: 스케줄 단위의 LRU 캐싱을 통해 자주 사용되는 계획을 우선적으로 유지함으로써 캐시 히트율을 극대화할 수 있습니다.

결론

이번 개선을 통해 MiniMax-H3 모델은 더 이상 반복적인 체크포인트 읽기 작업에 얽매이지 않게 되었습니다. 특히 멀티 랭크 환경에서도 일관된 성능을 보여주며, 캐시 모드에서의 가중치 업데이트 안정성까지 확보했습니다.

참고 자료

• torch.compile
• SGLang Contribution Guide

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