[vllm] vLLM의 Hopper(SM90) 아키텍처를 위한 GEMM 커널 최적화

PR 링크: vllm-project/vllm#54088 상태: Merged | 변경: +175 / -35

들어가며

vLLM은 대규모 언어 모델 추론을 위한 고성능 엔진입니다. 특히 최신 NVIDIA Hopper(SM90) 아키텍처에서 고성능을 유지하기 위해서는 연산 커널의 세밀한 튜닝이 필수적입니다. 이번 PR은 Kimi-K3 모델의 추론 성능을 극대화하기 위해, Hopper GPU 환경에 최적화된 저지연(low-latency) GEMM 커널을 도입하고, 기존의 SM100/SM103 타겟 테이블과 별도로 SM90 전용 튜닝 테이블을 구축하여 성능을 개선했습니다.

코드 분석

1. vllm/models/kimi_k3/nvidia/low_latency_gemm.py

핵심 변경 사항은 SM90 아키텍처를 위한 전용 KIMI_K3_PROJECTIONS_SM90 테이블의 추가입니다. 각 GPU 아키텍처마다 최적의 커널 구현(CuTe 또는 dsv3)이 다르기 때문에, 이를 하드웨어별로 분리하여 관리합니다.

# Before: SM103과 SM100만 지원
# After: SM90(H200) 전용 사양 추가
KIMI_K3_PROJECTIONS_SM90: dict[tuple[int, int], ProjectionSpec] = {
    (1536, 128): _sm90_spec(1536, 128, frozenset(range(1, 9))),
    (3072, 128): _sm90_spec(3072, 128, frozenset({1, 2, 5, 6, 7, 8, 9})),
    # ... (생략)
}

2. tests/kernels/test_bf16_skinny_gemm.py

새로운 하드웨어 타겟이 추가됨에 따라 테스트 스위트도 확장되었습니다. K3_GPU_TABLES를 통해 여러 아키텍처를 파라미터화하여 테스트할 수 있도록 구조를 개선했습니다.

# Before: SM100 위주의 테스트
# After: 파라미터화된 테스트를 통해 SM90 커버리지 확보
@pytest.mark.parametrize(
    "table",
    [k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90],
)
def test_device_table_is_keyed_by_shape(table: K3ProjectionTable) -> None:
    for (n, k), spec in table.items():
        assert (spec.n, spec.k) == (n, k)

왜 이게 좋은가

이번 최적화의 핵심은 '하드웨어 아키텍처별 최적의 커널 선택'입니다. 벤치마크 결과, H200(SM90) 환경에서 기존 방식 대비 기하평균(geomean) 기준 1.28배의 성능 향상을 보였으며, 특정 케이스에서는 최대 1.97배의 속도 개선을 달성했습니다.

• 일반적 교훈: 고성능 연산 라이브러리에서 하드웨어 아키텍처(SM 버전)는 커널의 레지스터 사용량, 공유 메모리 접근 패턴, 워프 스케줄링에 큰 영향을 미칩니다. 단순히 범용 커널을 사용하는 것보다, 각 아키텍처의 특성에 맞는 튜닝 테이블을 분리하여 관리하는 것이 추론 지연 시간을 줄이는 데 결정적입니다.
• 성능 수치: 36개의 테스트 포인트에서 최소 1.04x에서 최대 1.97x의 속도 향상을 기록했습니다.

리뷰 피드백

리뷰어들은 CI 환경에서의 검증을 위해 /ci run 명령을 사용하여 변경 사항이 기존 모델들에 회귀(regression)를 일으키지 않는지 확인했습니다. 하드웨어별 분기 로직이 명확하게 구현되어 있어 유지보수성이 높다는 평가입니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities
• https://github.com/vllm-project/vllm

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