[sglang] SGLang NPU 최적화: MoE 모델을 위한 Dual Stream 병렬 처리 도입

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

들어가며

최근 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 구조가 널리 사용되면서, 추론 효율성을 극대화하는 것이 중요해졌습니다. 특히 NPU 환경에서는 Shared Expert와 Routed Expert가 동일한 스트림에서 순차적으로 실행될 경우, 하드웨어 자원이 충분히 활용되지 못하는 병목 현상이 발생합니다. 본 PR은 이러한 문제를 해결하기 위해 NPU의 다중 스트림(Multi-stream) 기능을 활용하여 두 연산을 병렬로 처리하는 최적화를 도입했습니다.

코드 분석

1. NPU 스트림 관리 로직 추가 (cmo.py)

Shared Expert 연산을 독립적인 스트림에서 실행하기 위해 share_stream 관리 로직을 추가했습니다. shared_expert_on_independent_stream 함수는 새로운 NPU 스트림을 생성하거나 기존 스트림을 재사용하여 비동기 연산을 수행합니다.

# Before: 없음
# After: python/sglang/srt/hardware_backend/npu/cmo.py
def shared_expert_on_independent_stream(hidden_states, forward_func):
    stream = get_share_stream()
    if stream is None:
        stream = torch.npu.Stream()
        set_share_stream(stream)
    stream.wait_stream(torch.npu.current_stream())
    with torch.npu.stream(stream):
        shared_output = forward_func(hidden_states)
        return shared_output

2. 모델 Forward 로직 수정 (qwen2_moe.py)

_forward_deepep 메서드에서 SGLANG_NPU_USE_MULTI_STREAM 환경 변수가 활성화되고 CUDA Graph 모드일 경우, Shared Expert 연산을 별도 스트림으로 오프로드합니다. 연산 종료 후에는 wait_share_stream()을 호출하여 동기화를 보장합니다.

# After: python/sglang/srt/models/qwen2_moe.py
if enable_dual_stream:
    shared_output = shared_expert_on_independent_stream(
        hidden_states.clone(), self._forward_shared_experts
    )
else:
    shared_output = self._forward_shared_experts(hidden_states)

# ... (중략) ...

if enable_dual_stream:
    wait_share_stream()

왜 이게 좋은가

이 최적화는 연산의 오버랩(Overlap)을 통해 하드웨어 가동률을 극대화합니다. 벤치마크 결과, Qwen3.6-35B-A3B 모델 기준 처리량(Throughput)이 약 11.2% 향상되었으며, E2E Latency 또한 8.2% 감소하는 성과를 거두었습니다.

• 성능 향상: Throughput 11.2% 증가, Mean TPOT 8.8% 감소
• 교훈: MoE 모델의 구조적 특성상 Shared Expert와 Routed Expert는 독립적인 연산이 가능합니다. NPU와 같은 가속기 환경에서는 스트림 분리를 통해 유휴 자원을 활용하는 것이 성능 최적화의 핵심 전략임을 보여줍니다.

리뷰 과정에서 다른 MoE 모델로의 확장 가능성에 대한 논의가 있었으며, 현재는 Qwen 시리즈에 최적화되어 있으나 구조적으로 다른 MoE 모델에도 적용 가능한 범용적인 패턴입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.npu.Stream.html
• https://pytorch.org/docs/stable/generated/torch.npu.stream.html

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