[SGLang] MoE 모델을 위한 Single Batch Overlap 기법

PR 링크: sgl-project/sglang#9660 상태: Merged | 변경: +226 / -43

들어가며

Mixture-of-Experts(MoE) 모델의 추론에서 가장 큰 병목은 expert 간 통신(all-to-all dispatch/combine)과 GEMM 연산이 순차적으로 실행되는 것이다. 이 PR은 Hopper GPU에서 DeepEP low-latency 통신과 down projection GEMM을 두 개의 CUDA stream으로 overlap하여 실행하는 Single Batch Overlap(SBO) 기법을 확장한다.

핵심은 Blackwell과 Hopper GPU 각각에 맞는 overlap 전략을 분기하고, dispatch 단계에서도 shared expert 연산을 overlap하는 새로운 경로를 추가하는 것이다.

핵심 코드 분석

Hopper vs Blackwell 분기 처리

Before:

@classmethod
def enable_combine_down_gemm_two_stream_overlap(cls):
    return (
        is_sbo_enabled()
        and get_moe_runner_backend().is_flashinfer_cutedsl()
    )

@classmethod
def enable_combine_shared_two_stream_overlap(cls):
    return is_sbo_enabled()

After:

@classmethod
def enable_combine_down_gemm_two_stream_overlap(cls):
    return (
        is_sbo_enabled()
        and (
            get_moe_runner_backend().is_flashinfer_cutedsl()
            or (get_moe_runner_backend().is_deep_gemm() and not is_blackwell())
        )
    )

@classmethod
def enable_combine_shared_two_stream_overlap(cls):
    return is_sbo_enabled() and not cls.enable_dispatch_shared_one_stream_overlap()

@classmethod
def enable_dispatch_shared_one_stream_overlap(cls):
    return is_sbo_enabled() and not is_blackwell()

Hopper에서는 DeepGemm 백엔드로도 combine-down GEMM overlap을 사용할 수 있도록 확장했다. 또한 enable_dispatch_shared_one_stream_overlap이라는 새 경로를 추가하여, dispatch 단계에서 shared expert 연산을 하나의 stream 안에서 overlap한다.

Combine signal 구조 차이

Before:

combine_signal = torch.zeros(
    num_local_experts, dtype=torch.uint32, device=hidden_states.device
)

After:

if is_blackwell():
    combine_signal = torch.zeros(
        num_local_experts, dtype=torch.uint32, device=hidden_states.device
    )
else:
    MIN_BLOCK_M = 64
    combine_signal_size = num_local_experts * (
        (num_tokens_static + MIN_BLOCK_M - 1) // MIN_BLOCK_M
    )
    combine_signal = torch.zeros(
        combine_signal_size, dtype=torch.int32, device=hidden_states.device
    )

Blackwell은 expert 단위 signal로 충분하지만, Hopper에서는 block 단위(64 토큰)로 더 세밀한 signal이 필요하다. 이를 통해 GEMM이 완료된 block부터 즉시 combine 통신을 시작할 수 있다.

SM 할당 최적화

communicate_num_sms = get_int_env_var(
    "SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32 if is_blackwell() else 3
)

Hopper에서는 통신에 3개의 SM만 할당하고 나머지를 연산에 사용한다. Blackwell의 32개와 대비되는 수치로, 아키텍처별 최적 분배가 다르다는 점을 보여준다.

왜 이게 좋은가

MoE 추론에서 combine 통신과 down projection GEMM은 서로 데이터 의존성이 없는 구간이 존재한다. 이를 두 CUDA stream으로 겹쳐 실행하면, 통신 latency를 GEMM 연산 뒤에 숨길 수 있다. Hopper GPU까지 지원 범위를 넓힘으로써 더 많은 배포 환경에서 성능 이득을 얻을 수 있다.

정리

• Hopper GPU에서 DeepGemm 백엔드를 사용한 SBO를 지원한다
• Dispatch 단계에서도 shared expert overlap 경로를 새로 추가했다
• Block 단위 fine-grained signal로 Hopper의 통신-연산 overlap 효율을 높인다
• GPU 아키텍처별로 SM 할당, signal 구조를 분기하여 최적화한다

참고 자료

• DeepEP -- DeepSeek의 Expert Parallelism 통신 라이브러리
• SGLang SBO 문서 -- Single/Two Batch Overlap 구현

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