본문으로 건너뛰기

[sglang] [DeepSeek-V4.1] Blackwell(SM100) 성능의 한계를 끌어올리는 Fused WO-A 커널 최적화 분석

PR 링크: sgl-project/sglang#39957 상태: Merged | 변경: +1070 / -27

들어가며

LLM 서빙에서 작은 배치 사이즈(Small Batch)와 높은 병렬성(TP4 등)을 사용하는 환경은 커널 실행 오버헤드(Kernel Launch Overhead)에 매우 민감합니다. 특히 DeepSeek-V4.1과 같은 최신 모델의 디코드 및 검증(Verify) 경로에서는 Inverse-RoPE, WO-A 프로젝션, 그리고 MXFP8 양자화가 각각 별도의 커널로 실행되어 왔습니다.

이 PR은 NVIDIA의 최신 아키텍처인 Blackwell(SM100/SM103)의 기능을 활용하여, 이 세 가지 단계를 하나의 Fused Kernel로 통합했습니다. 이를 통해 불필요한 중간 메모리 접근을 제거하고 커널 런칭 횟수를 줄여, 실제 서비스 환경에서 유의미한 성능 향상을 이끌어냈습니다.

핵심 변경 사항 분석

1. 기존의 3단계 워크플로우 (Before)

기존에는 DeepSeek-V4의 추론 과정에서 다음과 같이 세 번의 커널이 순차적으로 호출되었습니다.

  1. fused_rope_inplace: RoPE 역변환 수행
  2. _wo_a_partial: WO-A 부분 GEMM 수행
  3. _wo_a_reduce_quant: 리덕션 및 MXFP8 양자화 수행

이 방식은 각 단계마다 GPU 메모리(VRAM)에 중간 결과를 썼다 읽어야 하며, 작은 배치에서는 커널을 띄우는 시간 자체가 연산 시간보다 길어지는 문제가 있었습니다.

2. Blackwell 전용 Fused 커널 (After)

이번 PR의 핵심은 python/sglang/kernels/jit/csrc/deepseek_v4/wo_a_fused.cuh 파일에 구현된 새로운 퓨즈드 커널입니다. 이 커널은 Blackwell 아키텍처의 TMA(Tensor Memory Accelerator)와 tcgen05(Blackwell MMA) 명령어를 직접 제어합니다.

TMA를 활용한 데이터 로드

// After: TMA를 이용한 2D 텐서 비동기 복사
SGL_DEVICE void
cp_async_bulk_tensor_2d(uint32_t dst_smem, const CUtensorMap* tmap, int32_t x, int32_t y, uint64_t* bar) {
  asm volatile(
      "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes"
      " [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_smem),
      "l"(tmap),
      "r"(x),
      "r"(y),
      "r"(to_shared(bar))
      : "memory");
}

기존의 일반적인 로드 방식 대신, 하드웨어 가속기인 TMA를 사용하여 Global Memory에서 Shared Memory로 데이터를 비동기적으로 가져옵니다. 이는 CPU나 CUDA 코어의 개입을 최소화합니다.

Blackwell MMA (tcgen05) 활용

Blackwell 아키텍처에서 도입된 새로운 행렬 연산 명령어인 tcgen05를 직접 호출하여 연산 효율을 극대화했습니다.

// After: Blackwell 전용 tcgen05 MMA 명령어 사용
SGL_DEVICE void tcgen05_mma_f16(uint32_t d, uint64_t desc_a, uint64_t desc_b, uint32_t idesc, bool accumulate) {
  asm volatile(
      "{ .reg .pred p; setp.ne.b32 p, %4, 0;"
      "  tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p; }" ::"r"(d),
      "l"(desc_a),
      "l"(desc_b),
      "r"(idesc),
      "r"(uint32_t(accumulate)));
}

이 코드는 BF16 정밀도의 행렬 곱셈을 수행하며, accumulate 플래그를 통해 이전 연산 결과에 누적할지 여부를 결정합니다. 퓨즈드 커널 내에서 RoPE 연산 직후 이 MMA가 실행되므로 데이터가 레지스터나 Shared Memory 내에서 유지됩니다.

3. Cluster-level 동기화 및 mbarrier

Blackwell의 특징 중 하나인 Thread Block Cluster를 효율적으로 관리하기 위해 mbarrier와 cluster_arrive 등의 PTX 명령어를 사용했습니다.

// After: Cluster 내 CTA 간 동기화
SGL_DEVICE void cluster_arrive_relaxed() {
  asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory");
}
SGL_DEVICE void cluster_wait_acquire() {
  asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
}

이를 통해 여러 CTA(Thread Block)가 협력하여 데이터를 처리할 때 발생하는 동기화 비용을 하드웨어 수준에서 최적화했습니다.

왜 이게 좋은가?

1. 성능 수치 (Performance Gains)

리뷰어 BBuf가 4× GB300 (SM103) 환경에서 테스트한 결과에 따르면:

  • 전체 디코드 처리량(Throughput): 1167.22 → 1202.65 tokens/s (+3.03% 향상)
  • 커널 체인 지연시간(Latency): T=8(Token rows) 기준 9.226 µs → 4.983 µs (약 46% 감소)
  • 오버헤드 제거: 20단계의 디코드 스텝당 1600개의 커널 호출을 제거했습니다.

2. 메모리 대역폭 절약

Inverse-RoPE의 결과물과 WO-A의 중간 부분합(Partial Sum)은 매우 큰 메모리 대역폭을 소모합니다. 이를 퓨즈드 커널 내부의 Shared Memory와 Register에서 처리함으로써, VRAM으로의 불필요한 Write/Read를 완전히 제거했습니다.

3. 정확도 유지

PR 설명에 따르면 GSM8K, AIME, GPQA 등 주요 벤치마크에서 기존 모델과 거의 동일한 정확도를 유지함을 확인했습니다. 미세한 차이는 부동 소수점 연산 순서 변경에 따른 일반적인 변동 범위 내에 있습니다.

일반적인 교훈

  1. 작은 연산일수록 퓨전이 답이다: 연산 자체의 시간보다 커널을 실행하고 메모리를 관리하는 오버헤드가 더 큰 경우, 하드웨어 특화 기능을 활용한 커널 퓨전은 드라마틱한 성능 향상을 가져옵니다.
  2. 아키텍처 특화 최적화: Blackwell과 같은 최신 GPU 아키텍처는 TMA, tcgen05와 같은 강력한 전용 명령어를 제공합니다. 이를 활용하기 위해 PTX 수준의 최적화가 필요할 수 있습니다.
  3. JIT의 활용: 다양한 텐서 모양(Shape)에 대응하기 위해 JIT(Just-In-Time) 컴파일러를 통해 최적화된 커널을 동적으로 생성하는 전략이 유효합니다.

결론

이번 PR은 DeepSeek-V4.1 모델을 Blackwell 하드웨어에서 최적으로 구동하기 위한 정교한 엔지니어링의 결과물입니다. 단순한 코드 정리를 넘어, 하드웨어의 로우레벨 기능을 직접 제어하여 LLM 추론 엔진의 효율성을 한 단계 끌어올렸습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글