본문으로 건너뛰기

[vllm] [vLLM] Hopper와 Blackwell을 위한 Router GEMM 최적화: DeepSeek-V4와 GLM5.2 성능 극대화하기

PR 링크: vllm-project/vllm#48774 상태: Merged | 변경: +93 / -20

들어가며: MoE 모델의 숨은 병목, Router GEMM

Mixture-of-Experts(MoE) 아키텍처를 사용하는 최신 LLM(DeepSeek-V3/V4, GLM5.2 등)에서 Router GEMM은 매우 빈번하게 호출되지만 최적화하기 까다로운 연산입니다. Router GEMM은 입력 토큰이 어떤 Expert로 갈지 결정하는 'Gate' 역할을 수행하는데, 보통 Batch Size($M$)는 작고 Hidden Dimension($K$)은 매우 큰 형태를 띱니다.

기존 vLLM은 이러한 경로에 대해 특정 셰이프(Shape)에 고정된 PTX 커널을 사용하거나, cuteDSL 기반의 범용 커널을 사용해 왔습니다. 하지만 새로운 모델이 등장할 때마다 Router의 셰이프가 달라지고, 특히 NVIDIA의 Hopper(SM90)와 Blackwell(SM100) 아키텍처 간의 하드웨어 특성 차이로 인해 범용적인 휴리스틱만으로는 최상의 성능을 내기 어려웠습니다.

이번 PR(#12836)은 주요 모델들의 Router 셰이프에 대해 아키텍처별로 전수 조사를 수행하고, 최적의 컴파일 설정을 적용하여 성능을 개선했습니다.

코드 분석 1: 아키텍처별 정밀 튜닝 테이블 도입

가장 큰 변화는 vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py에서 확인할 수 있습니다. 기존에는 하드웨어 구분 없이 단일 튜닝 테이블을 사용했으나, 이제는 SM90(Hopper)과 SM100F(Blackwell)를 구분하여 관리합니다.

Before

기존에는 아키텍처 구분 없이 _TUNED_CONFIGS라는 단일 딕셔너리에 의존했습니다.

_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = {
    (7168, 256): 6,
}
_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
    (7168, 384): {
        5: (4, 4),
        **{M: (5, 4) for M in range(6, 17)},
    },
}

After

변경 후에는 아키텍처별로 _SM100F_...와 _SM90_... 테이블을 분리하고, 런타임에 현재 플랫폼을 감지하여 최적의 설정을 선택합니다.

# SM100f(Blackwell) 전용 튜닝 테이블
_SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = {
    (6144, 256): {M: 256 for M in (1, 3, 4)},
}
_SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
    (4096, 256): {
        **{M: (8, 5) for M in (5, 8)},
        9: (8, 2),
    },
    # ... 생략 ...
}

# 아키텍처 감지 및 디스패치 로직
def _arch_tuned_configs():
    from vllm.platforms import current_platform
    if current_platform.is_device_capability_family(100):
        return (_SM100F_TUNED_DOTPROD_BS, _SM100F_TUNED_SPLITK_CONFIGS)
    if current_platform.is_device_capability(90):
        return (_SM90_TUNED_DOTPROD_BS, _SM90_TUNED_SPLITK_CONFIGS)
    return {}, {}

이 변경을 통해 split_k(K 차원 병렬화 정도)와 num_stages(파이프라이닝 스테이지 수)를 각 GPU의 메모리 대역폭과 연산 유닛 특성에 맞춰 세밀하게 조정할 수 있게 되었습니다.

코드 분석 2: 정적 리스트에서 동적 추출로, 스마트해진 Warmup

커널 성능을 보장하기 위해서는 모델 실행 전 'Warmup' 과정이 필수적입니다. 기존에는 Warmup 대상 셰이프가 하드코딩되어 있어, 새로운 모델이 추가될 때마다 이 리스트를 수동으로 업데이트해야 했습니다.

Before

_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = (
    (6144, 264),  # Inkling
    (7168, 256),  # DSV3
    (7168, 384),  # DSV4-Pro
    (14400, 256),  # DSV4-Flash
)

After

이제는 로드된 모델의 실제 레이어를 순회하며 GateLinear 모듈의 가중치 셰이프를 직접 추출합니다.

def _ll_bf16_router_shapes_from_model(model: torch.nn.Module) -> tuple[tuple[int, int], ...]:
    from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
    shapes: set[tuple[int, int]] = set()
    for module in model.modules():
        if not isinstance(module, GateLinear):
            continue
        # ... 가중치 체크 로직 ...
        n, k = weight.shape
        if k % 8 == 0:
            shapes.add((int(k), int(n)))
    return tuple(sorted(shapes))

이 방식은 유지보수성 측면에서 매우 뛰어납니다. 개발자가 새로운 모델을 지원할 때 Warmup 리스트를 깜빡하고 누락하는 실수를 원천 차단하며, 실제 사용되는 셰이프에 대해서만 Warmup을 수행하므로 불필요한 리소스 낭비도 줄여줍니다.

왜 이게 좋은가: E2E 성능의 중요성

이 PR의 핵심 교훈은 "마이크로 벤치마크 수치가 반드시 End-to-End(E2E) 성능으로 이어지지는 않는다"는 점입니다. PR 설명에 언급되었듯, PDL(Predictive Data Loading) 오버랩이 발생하는 상황에서는 커널 단독 실행 시간보다 전체 파이프라인에서의 조화가 더 중요합니다.

주요 성능 지표 (H200 / Hopper 기준)

  • DeepSeek-V3.2 (7168, 256): Median ITL(Inter-Token Latency) 5.8% 개선
  • GLM5.2 (6144, 256): Throughput 2.9% 향상
  • DeepSeek-V4-Flash (4096, 256): Throughput 2.4% 향상

단순히 cuBLAS를 호출하는 것보다, 특정 셰이프에 대해 split_k와 dotprod 백엔드를 적절히 섞어 쓰는 튜닝된 휴리스틱이 훨씬 강력함을 입증했습니다.

결론

이번 최적화는 vLLM이 단순히 범용 엔진을 넘어, 최신 하드웨어(Blackwell)와 최신 모델(DeepSeek-V4)에 대해 얼마나 기민하게 대응하고 있는지를 보여줍니다. 특히 아키텍처별 튜닝 테이블 분리와 동적 셰이프 추출 방식은 향후 더 많은 GPU와 모델이 등장하더라도 확장 가능한 구조를 마련했다는 점에서 높은 점수를 줄 수 있습니다.

MoE 모델을 운영하는 엔지니어라면, Router GEMM과 같은 작은 연산의 최적화가 전체 서비스의 ITL과 처리량에 얼마나 큰 영향을 미치는지 다시 한번 상기할 수 있는 좋은 사례입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글