본문으로 건너뛰기

[sglang] DeepSeek-V4.1 모델 성능 향상: DSpark 검증 및 MoE 커널 최적화

PR 링크: sgl-project/sglang#38879 상태: Merged | 변경: +1524 / -136

들어가며

최근 sglang 레포지토리의 PR #38879는 DeepSeek-V4.1 모델의 추론 성능을 획기적으로 개선하는 중요한 변경사항을 포함하고 있습니다. 이 PR은 특히 DSpark (Speculative Decoding)의 타겟 검증(target verification) 및 드래프트(draft) 단계, 그리고 Mixture-of-Experts (MoE) 커널의 병목 현상을 해결하는 데 중점을 둡니다. 기존 코드에서는 디코딩 과정의 여러 빠른 경로를 놓치거나, MoE 라우팅, 입력 양자화, 최종 처리 단계에서 불필요한 연산이 발생하여 성능 저하의 원인이 되었습니다. 본 글에서는 이 PR이 어떻게 이러한 문제들을 해결하고, 실제 코드 변경사항을 통해 성능 향상을 달성했는지 심층적으로 분석합니다.

코드 분석

이번 PR의 핵심은 DSpark의 타겟 검증 및 드래프트 단계의 효율성을 높이고, MoE 연산의 병목 구간을 제거하는 데 있습니다. 변경 사항은 주로 sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh 파일과 관련 로직에서 찾아볼 수 있습니다.

1. DSpark 타겟 검증 및 드래프트 단계 최적화

DSpark는 추론 속도를 높이기 위해 모델의 예측을 미리 수행하고(draft), 이를 실제 모델의 예측(target)과 비교하여 일치하는 경우 여러 토큰을 한 번에 생성하는 기법입니다. 이 과정에서 타겟 검증 및 드래프트 단계의 효율성이 전체 성능에 큰 영향을 미칩니다.

변경 전 (개념적): 기존에는 타겟 검증 및 드래프트 단계에서 통계 오버랩(statistics overlap)을 충분히 활용하지 못했고, 검증할 수 있는 형태(shapes)의 범위가 제한적이었습니다.

변경 후:

  • mHC statistics overlap 활용 및 검증 형태 확장: PR 설명에 따르면, mHC statistics overlap을 타겟 검증 및 모든 드래프트 단계에서 활용하고, 검증할 수 있는 형태의 범위를 넓혔습니다. 이는 GPU 활용률을 높이고 불필요한 계산을 줄여 성능을 향상시킵니다.

    - Use mHC statistics overlap in target verify and all three draft stages, and widen the validated verify shapes.
+ Use mHC statistics overlap in target verify and all three draft stages, and widen the validated verify shapes.

    (참고: 위 diff는 개념적인 설명이며, 실제 코드 diff는 더 복잡합니다. 이 부분은 주로 알고리즘 레벨의 개선을 의미합니다.)

2. MoE 커널 최적화

MoE 모델은 여러 전문가(expert) 네트워크를 사용하여 입력에 따라 동적으로 활성화되는 전문가를 선택합니다. 이 과정에서 라우팅, 양자화, 최종 결과 합산 등의 연산이 성능에 중요한 영향을 미칩니다.

변경 전 (개념적): MoE 라우팅, 입력 양자화, 최종 처리(finalize) 단계에서 병목 현상이 존재했습니다.

변경 후:

  • WO-A 출력 직접 작성 및 인덱스 후보 마스킹 융합: Write WO-A output directly in token-major layout and fuse index candidate masking. 이 변경은 메모리 레이아웃을 최적화하고, 불필요한 인덱스 계산을 줄여 연산 효율성을 높입니다. 토큰 단위의 레이아웃을 직접 사용하고, 인덱스 후보 마스킹을 융합함으로써 데이터 접근 및 처리 속도를 개선합니다.

  • 패딩된 라우터 출력 생성 및 MXFP8 양자화 분리 스트림 실행: Produce padded, packed router outputs and run routed-input MXFP8 quantization on a separate stream. 라우터 출력을 패딩하고 팩킹하여 효율성을 높였으며, MXFP8 양자화 연산을 별도의 스트림에서 실행하여 GPU의 병렬 처리 능력을 최대한 활용합니다. 이는 연산 간의 종속성을 줄이고 전체 처리 시간을 단축시킵니다.

  • 지연된 MoE 최종 처리, 공유 전문가 추가, 사용자 정의 all-reduce 융합: Fuse deferred MoE finalize, shared-expert addition, and custom push all-reduce for supported small TP4 batches. 이 변경은 여러 MoE 관련 연산을 하나로 융합하여 오버헤드를 줄이는 데 중점을 둡니다. 특히, finalize, shared-expert addition, all-reduce 연산을 융합하여 데이터 이동 및 커널 호출 비용을 절감합니다. 이는 Blackwell 아키텍처의 특정 구성(TP4 배치)에 최적화되었습니다.

    // Fused deferred-MoE finalize -> 1shot lamport push all-reduce [-> RMSNorm]
// over the CustomAllReduceV2 push plane, for decode-sized batches (bf16).
// ... (중략) ...
// The rank-local finalize never materializes in global memory: each thread
// computes one 16B vector of it and pushes it straight into every peer's push
// slot with unicast `st.relaxed.sys` stores, exactly like the generic
// `all_reduce_1shot_push_kernel`, so no multicast mapping is required.

    위 C++ 커널 코드는 MoE 최종 처리, 공유 전문가 추가, 그리고 사용자 정의 all-reduce 연산을 어떻게 융합하는지 보여줍니다. 특히, finalize 연산 결과를 전역 메모리에 저장하지 않고 직접 각 피어(peer)의 메모리 슬롯으로 푸시하는 방식을 사용하여 데이터 이동을 최소화합니다. 이는 all_reduce_1shot_push_kernel과 유사한 방식으로 동작하며, 멀티캐스트 매핑의 필요성을 제거합니다.

  • BF16 반올림 지점 보존 및 위상 카운터 재사용 동기화: Preserve the BF16 rounding points and synchronize phase-counter reuse. BF16 데이터 타입의 반올림 지점을 기존과 동일하게 유지하여 정확도를 보장하고, 위상 카운터(phase counter) 재사용을 동기화하여 커널 간의 안정적인 실행 흐름을 유지합니다.

  • 단일 노드 Blackwell V4.1 구성에 사용자 정의 all-reduce 선호: Prefer custom all-reduce for the supported single-node Blackwell V4.1 configuration. 특정 하드웨어 구성(단일 노드 Blackwell V4.1)에서는 성능이 더 좋은 사용자 정의 all-reduce 방식을 우선적으로 사용합니다. 다른 구성에서는 기존의 기본값이나 폴백(fallback) 방식을 유지합니다.

3. 정확도 테스트

성능 향상과 더불어 정확도 저하가 없는지도 중요한 검증 항목입니다. GSM8K 및 AIME 데이터셋을 사용한 테스트 결과, 성능 향상에도 불구하고 정확도는 거의 동일하거나 오히려 소폭 개선되었습니다.

  • GSM8K: 정확도가 97.00%에서 98.00%로 소폭 상승했습니다.
  • AIME: 정확도는 유사하게 유지되었습니다.

이는 최적화 과정에서 모델의 핵심적인 추론 능력이 손상되지 않았음을 시사합니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 훌륭한 최적화 및 개선이라고 할 수 있습니다.

  1. 획기적인 성능 향상: 실제 속도 테스트 결과는 놀랍습니다. 특히 BS64 워크로드에서 안정적인 디코드 출력 토큰/초가 6,555.95에서 13,472.51로 105.5% 증가했으며, 전체 배치(prefill 포함) 처리량도 56.4% 향상되었습니다. BS1 워크로드에서도 34.1%의 상당한 성능 향상을 보였습니다.

    Workload Baseline Candidate Change
    BS1, input 4096/output 1024, post-first-event output tok/s 569.88 764.29 +34.1%
    BS64, input 4096/output 2048, stable decode output tok/s 6,555.95 13,472.51 +105.5%
    BS64, same workload, full batch output tok/s including prefill 4,114.78 6,434.55 +56.4%

  2. 정확도 유지: 성능 향상이 정확도 희생을 동반하지 않았습니다. 오히려 GSM8K와 같은 벤치마크에서는 정확도가 소폭 개선되는 모습을 보였습니다. 이는 최적화가 모델의 근본적인 동작을 변경하지 않으면서 연산 효율성만 높였음을 의미합니다.

  3. 하드웨어 특성 활용: Blackwell 아키텍처와 같은 최신 GPU 하드웨어의 특성을 고려하여 사용자 정의 all-reduce 커널을 도입하고, MXFP8 양자화와 같은 최신 기술을 활용하여 성능을 극대화했습니다. 이는 특정 하드웨어에 대한 깊은 이해를 바탕으로 한 최적화입니다.

  4. 코드 복잡성 관리 및 융합: 여러 개의 작은 커널 호출과 데이터 이동을 하나의 융합된 커널로 통합함으로써, 커널 호출 오버헤드를 줄이고 데이터 지역성(locality)을 개선했습니다. 이는 코드의 복잡성을 증가시킬 수 있지만, 성능 향상이라는 명확한 이득을 가져옵니다.

  5. 일반적인 교훈:

    • 병목 지점 식별 및 집중: LLM 추론에서 가장 성능에 영향을 미치는 부분을 정확히 식별하고, 해당 부분을 집중적으로 최적화하는 것이 중요합니다.
    • 연산 융합 (Kernel Fusion): 여러 단계의 연산을 하나의 커널로 융합하는 것은 데이터 이동을 줄이고 GPU 활용률을 높이는 강력한 기법입니다.
    • 하드웨어 특화 최적화: 최신 하드웨어의 기능을 최대한 활용하는 커스텀 커널 개발은 상당한 성능 향상을 가져올 수 있습니다.
    • 정확도와 성능의 균형: 성능 최적화 시 정확도 저하를 면밀히 검증하고, 필요한 경우 정확도를 보장하는 메커니즘을 함께 도입해야 합니다.

리뷰 피드백 반영

PR 설명에 따르면, 이 PR은 @DarkSharpness (Ziyi Xu)와의 공동 작업으로 진행되었습니다. MoE 라우터 및 입력 사전 양자화, 최종 처리/공유 추가/all-reduce 융합, 그리고 사용자 정의 all-reduce 백엔드 선택 등 핵심적인 기여가 있었습니다. 이는 복잡한 최적화 작업이 여러 개발자의 전문성을 결합하여 이루어졌음을 보여줍니다. PR 설명에 구체적인 리뷰 코멘트가 명시되어 있지는 않지만, 공동 작업 및 상세한 성능/정확도 테스트 결과 제시는 코드의 완성도와 신뢰도를 높이는 데 기여했을 것입니다.

결론

PR #38879는 DeepSeek-V4.1 모델의 추론 성능을 극대화하기 위한 매우 성공적인 최적화 사례입니다. DSpark 및 MoE 커널의 핵심 병목 지점을 정밀하게 파악하고, 커널 융합, 하드웨어 특화 최적화, 효율적인 데이터 처리 기법을 적용하여 놀라운 성능 향상을 달성했습니다. 특히, 100% 이상의 성능 향상과 정확도 유지라는 두 마리 토끼를 모두 잡았다는 점에서 높은 평가를 받을 만합니다. 이러한 최적화는 sglang 라이브러리의 전반적인 성능을 향상시키고, 더 많은 사용자가 LLM을 효율적으로 활용할 수 있도록 기여할 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글