[sglang] SGLang 최적화: NPU 환경을 위한 RoPE 캐싱 메모리 효율화

PR 링크: sgl-project/sglang#25399 상태: Merged | 변경: +0 / -0

들어가며

LLM 추론 엔진인 SGLang은 다양한 하드웨어 가속기를 지원하며, 그 과정에서 메모리 효율성은 매우 중요한 요소입니다. 특히 RoPE(Rotary Positional Embedding)와 같은 위치 임베딩 계산 시, 미리 계산된 cos/sin 캐시 값은 추론 속도를 높여주지만, 불필요하게 메모리를 점유할 경우 대규모 모델 구동 시 병목이 될 수 있습니다. 이번 PR은 NPU(Neural Processing Unit) 환경이 아닌 경우, 불필요한 cos_cached_total 및 sin_cached_total 생성을 조건부로 제한하여 약 230MB의 메모리를 절감하는 최적화를 수행했습니다.

코드 분석

이번 변경은 python/sglang/srt/layers/rotary_embedding/rope_variant.py 파일 내의 _compute_cos_sin_cache 메서드에서 이루어졌습니다.

Before

기존 코드에서는 하드웨어 환경과 관계없이 항상 cos_cached_total과 sin_cached_total을 계산하고 메모리에 할당했습니다.

        cos = freqs.cos() * self.mscale
        sin = freqs.sin() * self.mscale
        cache = torch.cat((cos, sin), dim=-1)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.cos_cached_total = torch.cos(emb) * self.mscale
        self.sin_cached_total = torch.sin(emb) * self.mscale
        return cache

After

_is_npu 플래그를 확인하여 NPU 환경에서만 해당 연산을 수행하도록 변경되었습니다. 이를 통해 GPU 등 다른 가속기 환경에서는 불필요한 메모리 할당을 방지합니다.

        cos = freqs.cos() * self.mscale
        sin = freqs.sin() * self.mscale
        cache = torch.cat((cos, sin), dim=-1)
        if _is_npu:
            emb = torch.cat((freqs, freqs), dim=-1)
            self.cos_cached_total = torch.cos(emb) * self.mscale
            self.sin_cached_total = torch.sin(emb) * self.mscale
        return cache

왜 이게 좋은가

1. 메모리 효율성: 약 230MB의 VRAM을 절감했습니다. 이는 모델 가중치를 올리거나 KV 캐시를 더 확보할 수 있는 공간을 제공하므로, 특히 메모리 제한이 엄격한 환경에서 큰 이점을 가집니다.
2. 조건부 연산(Lazy Evaluation): 모든 하드웨어에서 동일한 연산을 강제하는 대신, 특정 아키텍처(NPU)의 요구사항에만 맞게 로직을 분기함으로써 불필요한 연산 사이클과 메모리 할당을 제거했습니다.
3. 교훈: 하드웨어 추상화 계층을 다룰 때는 '범용적인 코드'가 항상 최선은 아닙니다. 특정 가속기에서만 필요한 데이터 구조라면, 이를 조건부로 처리하는 것만으로도 상당한 리소스 최적화가 가능함을 보여줍니다.

결론

이번 SGLang의 변경사항은 작은 조건문 하나가 대규모 모델 추론 엔진에서 어떻게 메모리 효율을 극대화할 수 있는지 보여주는 좋은 사례입니다. 하드웨어별 특성을 고려한 최적화는 추론 엔진의 범용성과 성능을 동시에 잡는 핵심 전략입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cat.html
• https://pytorch.org/docs/stable/generated/torch.cos.html

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