[sglang] SGLang MoE All-Reduce 최적화: NCCL Symmetric Memory 활용으로 지연 시간 50% 단축

PR 링크: sgl-project/sglang#29007 상태: Merged | 변경: +115 / -35

들어가며

대규모 언어 모델(LLM)의 추론 성능은 모델의 복잡성과 배치 크기가 증가함에 따라 더욱 중요해지고 있습니다. 특히 MoE(Mixture-of-Experts) 모델은 여러 전문가(Expert) 네트워크의 출력을 효율적으로 결합해야 하므로, 분산 환경에서의 통신 오버헤드가 전체 성능에 큰 영향을 미칩니다. SGLang은 이러한 LLM 추론을 위한 고성능 프레임워크로, --enable-symm-mem 옵션을 통해 NCCL symmetric memory를 활용하여 tensor_model_parallel_all_reduce 연산의 지연 시간을 줄이는 기능을 제공합니다. 그러나 이 PR 이전에는 MoE 레이어에서 이 최적화 경로가 제대로 활성화되지 않는 문제가 있었습니다.

이 글에서는 SGLang의 MoE 레이어에서 NCCL symmetric memory를 효과적으로 활용하도록 개선한 GitHub PR의 코드 변경사항을 분석하고, 이 최적화가 왜 중요한지, 그리고 어떤 성능 향상을 가져왔는지 설명합니다.

문제점 분석: 왜 Symmetric Memory가 활용되지 못했나?

NCCL symmetric memory는 GPU 간 통신 시 데이터가 특정 메모리 풀에 할당되어 있을 때, 불필요한 데이터 복사 없이 직접 접근하여 all_reduce와 같은 집단 통신 연산의 지연 시간을 크게 줄일 수 있는 기술입니다. SGLang의 MoE 레이어는 이 기능을 활용하도록 설계되었지만, 다음 두 가지 주요 문제로 인해 실제로는 빠른 경로를 타지 못했습니다.

1. MoE Runner의 출력 버퍼 할당 방식: DeepGemm, Triton 등 MoE 레이어의 개별 Runner들이 자체적으로 출력 버퍼를 할당할 때, 이 버퍼가 symmetric memory 풀 외부에 생성되었습니다. forward_impl에서 use_symmetric_memory 컨텍스트를 사용하더라도, 내부 Runner들이 torch.empty 등을 호출하여 새로운 메모리를 할당하면 미리 할당된 symmetric 버퍼를 재사용하지 못했습니다.
2. 업스트림 hidden_states의 메모리 위치: MoE 레이어의 입력으로 들어오는 hidden_states (예: DeepSeek-V4 모델의 hc_pre 레이어 출력) 자체가 symmetric memory 풀에 있지 않은 경우가 있었습니다. 이 경우, MoE 레이어의 전체 파이프라인이 symmetric memory 외부의 출력을 생성하게 되어, 최종 all_reduce 연산이 느린 경로를 사용하게 됩니다.

이러한 문제들은 tensor_model_parallel_all_reduce 연산이 NCCL symmetric memory의 이점을 누리지 못하게 하여, MoE 레이어의 성능 병목으로 작용했습니다.

코드 변경사항 분석: Symmetric Memory 직접 할당

이 PR은 위에서 언급된 문제들을 해결하기 위해, MoE 레이어의 핵심 출력 버퍼와 업스트림 입력 버퍼를 use_symmetric_memory 컨텍스트 내에서 직접 할당하도록 변경했습니다. 이는 불필요한 복사 없이 처음부터 symmetric memory를 사용하게 하여, all_reduce 성능을 극대화합니다.

1. python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

DeepGemm Runner에서 MoE 레이어의 최종 출력인 down_output 및 output 텐서가 symmetric memory 풀에 할당되도록 변경되었습니다. 중간 버퍼는 기본 할당자를 사용하여 풀 점유율을 제한합니다.

Before:

down_output = torch.empty(
    (all_tokens, K),
    device=hidden_states_device,
    dtype=torch.bfloat16,
)

After:

with use_symmetric_memory(
    get_tp_group(), disabled=not is_allocation_symmetric()
):
    down_output = torch.empty(
        (all_tokens, K),
        device=hidden_states_device,
        dtype=torch.bfloat16,
    )

이와 유사하게 _run_bf16_contiguous_gemm, _run_masked_gemm, _run_masked_bf16_gemm, post_permute_deep_gemm_to_standard 함수 내의 down_output 또는 output 할당에도 동일한 use_symmetric_memory 컨텍스트가 적용되었습니다.

2. python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

Triton 기반의 _fused_moe_kernel_sequence 함수에서 inplace가 False일 때 할당되는 out_hidden_states 또한 symmetric memory 풀에 할당되도록 변경되었습니다.

Before:

else:
    out_hidden_states = torch.empty_like(hidden_states)

After:

else:
    # Allocate the MoE output in the NCCL symmetric memory pool when symmetric
    # allocation is required, so the downstream all-reduce takes the low-latency
    # symmetric path. Only this output enters the pool; the intermediate caches
    # below stay on the default allocator to bound pool occupancy.
    with use_symmetric_memory(
        get_tp_group(), disabled=not is_allocation_symmetric()
    ):
        out_hidden_states = torch.empty_like(hidden_states)

3. python/sglang/kernels/ops/layernorm/mhc.py

MoE 레이어의 입력으로 사용되는 layer_input (TileLang mHC 경로의 _mhc_pre_impl에서) 및 layer_input_cur (mhc_fused_post_pre에서) 텐서가 symmetric memory 풀에 할당되도록 변경되었습니다. 이는 inplace Triton MoE Runner가 symmetric all_reduce 입력을 생성하도록 보장합니다.

Before (mhc_pre):

layer_input = torch.empty(
    num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
)

After (mhc_pre):

with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
    layer_input = torch.empty(
        num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
    )

mhc_fused_post_pre의 layer_input_cur에도 동일한 변경이 적용되었습니다.

4. python/sglang/srt/models/deepseek_v4.py

DeepSeek-V4 모델의 비-TileLang hc_pre 경로에서 MoE 레이어의 입력으로 사용되는 y 텐서가 symmetric memory 풀에 할당되도록 변경되었습니다. 이는 mhc.py의 변경사항과 상응하는 조치입니다.

Before (hc_pre_torch_impl):

y = (pre.squeeze(1).unsqueeze(-1) * x_flat.view(shape)).sum(dim=1)
return y.to(dtype), post.squeeze(1), comb.squeeze(1), False

After (hc_pre_torch_impl):

with use_symmetric_memory(
    get_tp_group(), disabled=not is_allocation_symmetric()
):
    y = (pre.squeeze(1).unsqueeze(-1) * x_flat.view(shape)).sum(dim=1).to(dtype)
return y, post.squeeze(1), comb.squeeze(1), False

5. python/sglang/srt/layers/dp_attention.py

_DpGatheredBufferWrapper 클래스의 클래스 속성들이 AttributeError를 방지하기 위해 기본값을 갖도록 초기화되었습니다. 이는 is_allocation_symmetric()와 같은 함수가 첫 forward 호출 이전에 실행될 때 안전성을 보장합니다.

Before:

_global_dp_buffer_len: int
_local_dp_buffer_len: int
_dp_max_padding: bool
_global_num_tokens: Optional[List[int]]

After:

_global_dp_buffer_len: int = 0
_local_dp_buffer_len: int = 0
_dp_max_padding: bool = False
_global_num_tokens: Optional[List[int]] = None

6. test/registered/kernels/test_mhc_kernels.py

단일 프로세스 커널 유닛 테스트 환경에서는 TP 그룹이 초기화되지 않으므로, use_symmetric_memory를 nullcontext()로, is_allocation_symmetric을 False로, get_tp_group을 None으로 monkeypatch하여 일반 torch.empty를 통해 메모리를 할당하도록 변경되었습니다. 이는 테스트 환경에서 symmetric memory 관련 로직이 오류 없이 작동하도록 합니다.

왜 이 최적화가 좋은가?

이 PR의 최적화는 NCCL symmetric memory를 MoE 레이어의 all_reduce 연산에 효과적으로 적용함으로써, 다음과 같은 중요한 이점들을 제공합니다.

1. 현저한 성능 향상: DeepSeek-V4-Flash-FP8 모델에서 4K 입력 / 1.5K 출력 벤치마크 결과, TPOT(Time Per Output Token)이 6.55% (11.44ms → 10.69ms) 감소했습니다. 또한, all_reduce 호출당 평균 지연 시간은 50% (20μs → 10μs) 감소했습니다. 이는 분산 환경에서 통신 오버헤드가 큰 MoE 모델에 특히 중요한 개선입니다.

   벤치마크 결과 요약:

   • TPOT: 11.44 ms → 10.69 ms (6.55% 감소)
   • 평균 all_reduce 지연 시간: 20 μs → 10 μs (50% 감소)
   • 전체 Latency: 158.093 s → 144.824 s (약 8.4% 감소)
   • Output throughput: 560.387 token/s → 611.551 token/s (약 9.1% 증가)

2. 불필요한 메모리 복사 제거: 초기 PR 설명에서는 data_ptr 비교 후 복사하는 방식도 고려되었으나, 리뷰어 nvcastet의 피드백을 반영하여 MoE 레이어의 출력 및 업스트림 입력 텐서를 처음부터 symmetric memory에 직접 할당하는 방식으로 변경되었습니다. 이 접근 방식은 불필요한 memcpy 오버헤드를 완전히 제거하여, 더욱 효율적인 메모리 관리를 가능하게 합니다.

3. 일관된 Symmetric Memory 활용: MoE Runner의 내부 할당과 업스트림 레이어의 출력 모두 symmetric memory를 사용하도록 보장함으로써, enable-symm-mem 옵션이 의도한 대로 MoE 레이어 전반에 걸쳐 저지연 통신 경로를 활성화합니다.

이 최적화는 대규모 분산 LLM 추론 시스템에서 통신 병목 현상을 줄이고, 전체 시스템의 처리량을 향상시키는 데 중요한 기여를 합니다. 특히 MoE 모델과 같이 통신이 빈번한 아키텍처에서 그 효과가 두드러집니다.

리뷰어 피드백 반영

이 PR의 개발 과정에서 nvcastet 리뷰어의 피드백은 최종 솔루션의 품질을 높이는 데 결정적인 역할을 했습니다. 초기에는 MoE 출력 버퍼가 symmetric memory에 없는 경우 data_ptr를 비교하여 복사하는 방식이 고려되었으나, nvcastet는 복사를 피하고 필요한 텐서를 상위 스트림에서부터 symmetric memory에 등록하는 방식을 제안했습니다. wangfakang은 이 피드백을 수용하여, FusedMoE.forward_impl의 복사 로직을 제거하고 hc_pre와 같은 업스트림 레이어에서부터 symmetric allocation을 수행하도록 코드를 수정했습니다. 이는 불필요한 오버헤드를 완전히 제거하고, 더욱 깔끔하고 효율적인 최적화로 이어졌습니다.

결론

SGLang의 MoE all_reduce 최적화는 NCCL symmetric memory의 잠재력을 최대한 활용하여 분산 LLM 추론의 성능을 크게 향상시켰습니다. 메모리 할당 방식을 세심하게 조정하고, 불필요한 데이터 복사를 제거함으로써, all_reduce 지연 시간을 절반으로 줄이고 전체 추론 처리량을 개선할 수 있었습니다. 이러한 최적화는 고성능 LLM 서빙 시스템을 구축하는 데 있어 메모리 관리와 분산 통신의 중요성을 다시 한번 상기시켜 줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.empty.html
• https://pytorch.org/docs/stable/generated/torch.empty_like.html
• https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/overview.html

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