본문으로 건너뛰기

[sglang] SGLang, Ulysses 어텐션의 IPC 기반 파이프라이닝으로 성능 극대화

PR 링크: sgl-project/sglang#43164 상태: Merged | 변경: +914 / -10

들어가며

최근 SGLang 프로젝트의 GitHub 저장소에는 Ulysses 모델의 MiniMax-H3 블록에서 발생하는 성능 병목 현상을 해결하기 위한 중요한 PR이 제출되었습니다. 이 PR의 핵심은 기존의 순차적인 쿼리(q), 키(k), 밸류(v) 교환 및 출력 처리 방식을 개선하여, 특히 FlashAttention과 같이 연산 집약적인 커널이 SM(Streaming Multiprocessor)을 최대한 활용하는 상황에서 발생하는 오버헤드를 줄이는 것입니다. 기존 방식에서는 q/k/v 교환 시 NCCL을 사용하고, 출력 처리 시에는 SM 복사를 사용했는데, 이 과정들이 연산 커널 사이에 끼어들어 전체 성능을 저하시키는 요인이었습니다. 이 PR은 이러한 문제를 해결하기 위해 Copy Engine을 활용한 IPC(Inter-Process Communication) 기반의 올투올(all-to-all) 통신과 파이프라이닝 기법을 도입하여, Ulysses의 어텐션 메커니즘을 혁신적으로 최적화했습니다.

코드 분석

이 PR은 총 네 개의 커밋으로 구성되어 있으며, 각 커밋은 점진적으로 성능을 개선하고 새로운 기능을 도입합니다. 모든 변경 사항은 기존의 순차적 교환 방식과 비트 단위로 동일한 결과를 보장하면서도 성능 향상을 목표로 합니다.

1. Copy-engine all-to-all with head-group pipelining (ipc_a2a_multi.py, _minimax_h3_pipelined_dense_attention)

이 변경의 핵심은 여러 개의 동일 호스트 내 랭크(rank)를 대상으로 하는 IPC 기반의 올투올 통신 설계입니다. 각 랭크는 자신의 데이터를 피어 랭크의 스테이징 버퍼에 직접 기록하며, 이 과정에서 Copy Engine을 활용한 디바이스-투-디바이스 복사가 수행됩니다. 이는 기존의 NCCL 기반 올투올 통신이 SM을 사용했던 것과 달리, SM 자원을 전혀 소모하지 않아 연산 커널의 성능을 온전히 활용할 수 있게 합니다.

또한, 헤드 그룹(head group)별 파이프라이닝을 도입했습니다. 이는 그룹 g의 데이터 복사가 그룹 g-1의 어텐션 연산과 동시에 진행되도록 하여, 전체 처리 시간을 단축시킵니다. 각 그룹은 별도의 스트림에서 처리되며, cuStreamWriteValue32 및 cuStreamWaitValue32와 같은 스트림 메모리 연산을 사용하여 복사 작업과 커널 실행 간의 순서를 효율적으로 관리합니다. 이를 통해 SM을 불필요하게 점유하지 않으면서도 데이터 종속성을 정확하게 관리할 수 있습니다.

이 새로운 IPC 기반 통신은 특정 조건 하에서만 활성화됩니다. 예를 들어, 모든 랭크가 서로 접근 가능해야 하고, 스트림 메모리 연산을 지원해야 합니다. 또한, 그래프 캡처 중이거나 torch.compile 사용 시, 또는 VSA 및 링(ring) 어텐션 백엔드와 같이 다른 어텐션 백엔드를 사용할 경우에는 기존의 순차적 교환 방식으로 자동 폴백(fallback)됩니다.

Before: 기존에는 q/k/v 교환 시 NCCL을 사용하고, 출력 처리 시 SM 복사를 사용했습니다. 이 과정은 연산 집약적인 커널 사이에 위치하여 성능 저하의 원인이 되었습니다.

After:

# python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py
# ... (중략) ...
        # The transport is the two-rank `ipc_a2a` design widened to N same-host ranks:
        # each rank writes its chunk straight into its peers' staging buffers with
        # device-to-device copies on the copy engine, and a per-peer sequence flag
        # orders data before flag.
        # Pipelining: group g's copies run while group g-1 attends; each group attends
        # on its own stream; stream memory operations (`cuStreamWriteValue32` /
        # `cuStreamWaitValue32`) order copies and kernels without holding an SM.
# ... (중략) ...
        # It falls back to the sequential exchange when any rank lacks peer access
        # or stream memory operations (agreed across the group before any handle
        # exchange), during graph capture, under `torch.compile`, and for other
        # attention backends, VSA and ring.

2. QK-norm + RoPE writes straight into the exchange

이 변경은 QK-norm 커널의 출력을 직접 IPC 교환 버퍼에 쓰도록 최적화했습니다. 기존에는 QK-norm 연산 후 정규화된 q/k와 v를 별도의 패킹(packing) 과정을 거쳐야 했습니다. 하지만 이제 QK-norm 커널이 목적지 블록의 정규화된 q/k와 v를 직접 해당 랭크의 수신 슬롯에 기록합니다. 이는 불필요한 디바이스-투-디바이스 복사 단계를 제거하여 효율성을 높입니다. 또한, 헤드 그룹이 순서대로 생성되므로, 나중에 생성되는 그룹의 정규화 연산이 이전 그룹의 어텐션 연산과 겹쳐서 실행될 수 있어 파이프라이닝 효과를 극대화합니다.

Before: QK-norm 연산 후, 별도의 패킹 및 복사 단계를 거쳐 교환 버퍼에 데이터를 준비했습니다.

After:

# (이 변경은 특정 파일의 명시적인 diff보다는 로직 설명에 가까우므로, 코드 인용은 생략합니다. PR 설명의 내용을 기반으로 합니다.)
# The out-of-place QK-norm kernel writes each destination block's normalized q/k,
# plus v, in place of the in-place norm followed by a pack. Blocks for this rank
# go straight into its own receive slot, so they skip a device-to-device copy.
# Groups are produced in order, so the norm of later groups overlaps the
# attention of earlier ones.

3. Each head group is merged into the output as it arrives

이 최적화는 각 헤드 그룹의 출력이 도착하는 즉시 최종 텐서에 병합되도록 합니다. 각 그룹 스트림은 해당 랭크의 행(row)을 먼저 쓰고, 그 다음 피어 랭크의 행을 최종 출력 텐서의 해당 위치에 직접 기록합니다. 이는 후처리 단계에서 수행되던 후행 병합 커널과 자체 랭크 출력 복사 단계를 대체하여, 데이터 이동 및 연산 오버헤드를 줄입니다.

Before: 모든 헤드 그룹의 연산이 완료된 후, 별도의 병합 커널과 출력 복사 단계를 거쳐 최종 출력을 생성했습니다.

After:

# (이 변경 역시 특정 파일의 명시적인 diff보다는 로직 설명에 가까우므로, 코드 인용은 생략합니다. PR 설명의 내용을 기반으로 합니다.)
# Each group streams write this rank's rows, and then each peer's, directly into
# the final tensor. This replaces the trailing merge kernel and the own-rank
# output copy.

4. On by default (SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS)

이 PR은 새로운 파이프라이닝 기능을 기본적으로 활성화했습니다. 환경 변수 SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS의 기본값이 -1로 설정되어, MiniMax-H3의 56개 헤드를 Ulysses 2, 4, 8 구성에 맞춰 자동으로 최적의 그룹 수(예: 7개)를 선택합니다. 만약 계산된 버퍼 크기가 특정 랭크의 메모리에 맞지 않으면, 모든 랭크는 자동으로 기존의 순차적 교환 방식으로 돌아갑니다. 0 또는 1로 설정하면 파이프라이닝이 비활성화되고, N >= 2로 설정하면 N개의 그룹으로 강제됩니다.

Before: 파이프라이닝 기능은 명시적으로 활성화해야 했습니다.

After:

--- a/python/sglang/multimodal_gen/envs.py
+++ b/python/sglang/multimodal_gen/envs.py
@@ -61,6 +61,14 @@
     # back to NCCL when unavailable. Set 0 to force NCCL. Keep this in step with
     # the resolver below -- that is the value the runtime reads.
     SGLANG_DIFFUSION_IPC_A2A: bool = True
+    # copy-engine all-to-all for Ulysses groups of any size on one host; off by
+    # default while it is validated. Falls back to NCCL when unavailable.
+    SGLANG_DIFFUSION_IPC_A2A_MULTI: bool = False
+    # head groups for pipelining MiniMax-H3's Ulysses exchange against dense
+    # attention over the copy-engine transport: -1 picks a count that divides
+    # the heads per rank and steps aside when its buffers do not fit, 0 or 1
+    # keeps the sequential exchange, N >= 2 forces N groups
+    SGLANG_DIFFUSION_ULYSSES_PIPELINE_GROUPS: int = -1
     # a deadlock backstop, not a per-step budget: a rank can legitimately stall
     # for seconds (layerwise offload, wan2.2 expert-tower swaps), and expiry now
     # retires the transport on every rank and fails the request

왜 이게 좋은가

이 PR의 최적화는 여러 측면에서 매우 긍정적인 결과를 가져옵니다.

  1. 성능 향상: 벤치마크 결과에 따르면, 이 변경 사항들은 T2VA 및 Ref2VA 모델에서 상당한 지연 시간 감소를 보여줍니다. 예를 들어, 4xH200, 50 스텝, G=7 설정에서 T2VA Ulysses 2는 1.45%에서 최대 2.4%까지, Ulysses 4는 1.71%에서 최대 2.6%까지 성능이 향상되었습니다. Ref2VA의 경우 약 1.7%의 성능 개선이 관찰되었습니다.

    • T2VA Ulysses 2 (4 GPUs, 50 steps, G=7): 144.33s -> 141.33s (-2.1%) (커밋 3 적용 후)
    • T2VA Ulysses 4 (4 GPUs, 50 steps, G=7): 73.81s -> 72.23s (-2.1%) (커밋 3 적용 후)
    • 전체 적용 (No env, T2VA Ulysses 2): 144.80s -> 141.32s (-2.4%)
    • 전체 적용 (No env, T2VA Ulysses 4): 74.19s -> 72.23s (-2.6%)

  2. SM 자원 효율 극대화: Copy Engine을 활용한 IPC 통신은 SM을 전혀 사용하지 않습니다. 이는 FlashAttention과 같이 SM 집약적인 커널이 실행될 때, 해당 SM 자원을 방해 없이 온전히 사용할 수 있게 하여 전체적인 처리량을 높입니다. 실제로 Ref2VA at Ulysses 2의 레이어별 예산 분석 결과, 이전에는 0.58ms의 노출된 교환 시간이 있었으나, 최적화 후에는 이 시간이 크게 줄어들었습니다.

  3. 파이프라이닝을 통한 잠재적 성능 향상: 헤드 그룹별 파이프라이닝은 데이터 복사와 연산 간의 중첩을 통해 잠재적인 성능 향상을 제공합니다. 비록 H200 GPU가 최대 전력(700W)에서 작동하여 일부 해방된 SM 시간이 클럭 속도 저하로 상쇄될 수 있지만, 이는 하드웨어의 물리적 한계 내에서 최대한의 효율을 끌어내는 전략입니다.

  4. 자동화 및 유연성: 새로운 기능이 기본적으로 활성화되고, 버퍼 크기 제약 조건에 따라 자동으로 폴백되는 기능은 사용자가 별도의 복잡한 설정 없이도 최적의 성능을 경험할 수 있게 합니다. 또한, 그룹 수를 수동으로 지정할 수 있는 옵션은 특정 환경에 대한 미세 조정을 가능하게 합니다.

일반적 교훈:

  • 통신 오버헤드 최소화: 딥러닝 모델, 특히 대규모 모델에서는 통신(네트워크, IPC)이 주요 병목 지점이 될 수 있습니다. SM 커널이 아닌 전용 하드웨어(Copy Engine)를 활용하는 것은 이러한 오버헤드를 크게 줄일 수 있는 효과적인 방법입니다.
  • 파이프라이닝의 힘: 연산과 통신, 또는 다른 연산 단계를 중첩시키는 파이프라이닝은 전체 처리 시간을 줄이는 강력한 기법입니다. 특히 GPU 아키텍처의 다양한 자원(SM, Copy Engine, 스트림)을 활용하여 파이프라이닝을 구현하는 것이 중요합니다.
  • 자동화된 최적화 및 폴백: 복잡한 최적화 기법은 기본적으로 활성화하되, 하드웨어 제약이나 특정 조건에서는 안정적인 폴백 메커니즘을 제공하는 것이 사용자 경험과 시스템 안정성 측면에서 중요합니다.

리뷰 피드백 반영

PR 리뷰 댓글에서 mickqian님이 /tag-and-rerun-ci 태그를 사용하여 CI를 재실행하도록 요청했습니다. 이는 코드 변경이 CI 환경에서 올바르게 테스트되고 검증되었는지 확인하기 위한 일반적인 절차입니다. 이 PR은 이미 정확성 테스트와 속도 테스트 결과를 포함하고 있으며, 모든 변경 사항이 비트 단위로 동일함을 보장하므로, CI 통과는 예상되는 결과입니다. 리뷰어의 이러한 요청은 코드의 신뢰성을 확보하는 데 중요한 역할을 합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글