[sglang] SGLang에서 NPU를 위한 LTX-2/2.3 추론 성능 최적화 및 호환성 개선

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

들어가며

최근 생성형 AI 모델인 LTX-2 및 LTX-2.3을 NPU(Ascend) 환경에서 구동하려는 시도가 늘고 있습니다. 하지만 기존 SGLang 코드는 주로 CUDA 환경에 최적화되어 있어, NPU 환경에서 float64 연산 지원 미비로 인한 크래시나 성능 저하 문제가 발생했습니다. 본 PR은 이러한 호환성 문제를 해결하고 NPU 환경에 맞는 최적화된 연산을 적용하여 추론 성능을 개선하는 것을 목표로 합니다.

코드 분석

1. 플랫폼 추상화 및 float64 제약 해결

NPU는 float64 연산을 지원하지 않으므로, 이를 감지하고 float32로 폴백(fallback)하는 로직이 필요했습니다. current_platform 인터페이스를 확장하여 플랫폼별 기능을 추상화했습니다.

Before:

noise = torch.randn(reference_tensor.shape, generator=generator, dtype=torch.float64, device=reference_tensor.device)

After:

noise = torch.randn(reference_tensor.shape, generator=generator, dtype=(torch.float32 if not current_platform.is_float64_supported() else torch.float64), device=reference_tensor.device)

2. NPU 전용 커널 및 레이어 최적화

torch.nn.RMSNorm 대신 SGLang의 NPU 최적화 RMSNorm을 사용하도록 변경하고, Triton 커널 호출 시 is_cuda 체크를 current_platform.tensor_on_device(t)로 추상화하여 NPU 호환성을 확보했습니다.

Before:

self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)

After:

if _is_npu:
    self.q_norm = RMSNorm(self.inner_dim, eps=self.norm_eps)
else:
    self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)

왜 이게 좋은가

이번 최적화는 단순히 크래시를 방지하는 것을 넘어, NPU 환경에 최적화된 연산 경로를 사용함으로써 실질적인 성능 향상을 이끌어냈습니다. 벤치마크 결과에 따르면 T2V(Text-to-Video) 작업에서 약 7.6%에서 10.9%의 레이턴시 개선을 보였으며, I2V(Image-to-Video) 작업에서도 유사한 수준의 성능 향상을 확인했습니다.

일반적 교훈

1. 하드웨어 추상화의 중요성: 특정 하드웨어(CUDA)에 종속된 is_cuda와 같은 코드를 플랫폼 인터페이스로 추상화함으로써, 향후 다른 가속기(NPU, ROCm 등)로의 확장이 훨씬 용이해집니다.
2. 정밀도 관리: NPU와 같이 특정 정밀도를 지원하지 않는 하드웨어에서는 연산 파이프라인 전체에서 float64 사용을 엄격히 통제하고, 필요한 경우 float32로의 자동 변환 로직을 포함하는 것이 안정성 확보의 핵심입니다.
3. Triton 커널 재사용: NPU에서도 기존 CUDA용으로 작성된 고성능 Triton 커널을 적절히 활용할 수 있도록 인터페이스를 맞추는 것이 개발 효율성 측면에서 매우 효과적입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.RMSNorm.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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