[vllm] vLLM에서 DeepSeek V4.1을 위한 Mega-mHC 커널 최적화

PR 링크: vllm-project/vllm#56962 상태: Merged | 변경: +316 / -41

들어가며

최신 대규모 언어 모델인 DeepSeek V4.1은 복잡한 연산 구조를 가지고 있어, 기존의 범용 커널로는 추론 성능을 극대화하는 데 한계가 있었습니다. 특히 Multi-Head Latent Attention(MHA)과 유사한 구조인 mHC(Multi-Head Context) 연산은 메모리 대역폭과 연산 효율성에 민감합니다. 본 PR은 기존의 TileLang 기반 구현을 대체하거나 보완하기 위해 DeepGEMM 라이브러리의 Mega-mHC 커널을 vLLM에 통합하여, 특히 고부하 상황(Concurrency 32)에서 출력 처리량을 크게 개선했습니다.

코드 분석

1. vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py (신규 추가)

이 파일은 DeepGEMM의 mega_mhc 기능을 vLLM 모델 레이어와 연결하는 브릿지 역할을 합니다. is_mega_mhc_supported 함수를 통해 하드웨어 가용성(SM100 계열)과 입력 파라미터(hidden_size, hc_mult)를 검증합니다.

@functools.cache
def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool:
    if not (is_deep_gemm_supported() and current_platform.is_device_capability_family(100) ...):
        return False
    deep_gemm = _import_deep_gemm()
    return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None))

2. vllm/models/deepseek_v41/nvidia/model.py (기존 코드 수정)

기존에 사용하던 mhc_fused_post_pre_delayed_tilelang를 새로운 mhc_shifted_post_pre로 교체했습니다. 이는 추상화된 인터페이스를 통해 최적화된 커널을 선택적으로 호출하게 합니다.

# Before
residual, post_mix, res_mix, x, attn_pre, aux = mhc_fused_post_pre_delayed_tilelang(...)

# After
residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre(...)

왜 이게 좋은가

이번 최적화의 핵심은 연산 융합(Kernel Fusion)입니다. Mega-mHC는 여러 단계의 텐서 연산을 하나의 CUDA 커널 내에서 처리함으로써, GPU 메모리 읽기/쓰기 횟수를 획기적으로 줄였습니다.

• 성능 수치: 7168 hidden size 기준, 64 토큰 디코딩 시 기존 TileLang 대비 약 1.51배(22.37µs -> 14.85µs)의 속도 향상을 보였습니다.
• 교훈: LLM 추론 최적화에서 가장 큰 병목은 연산 자체가 아니라 메모리 접근입니다. 특정 모델 구조에 최적화된 커널(DeepGEMM 등)을 도입하여 메모리 트래픽을 최소화하는 것이 대규모 모델 서빙의 핵심 전략임을 보여줍니다.

결론

이번 통합을 통해 DeepSeek V4.1 모델은 더 높은 처리량과 낮은 지연 시간을 확보하게 되었습니다. 특히 Concurrency 32 환경에서 기존 대비 약 5% 이상의 출력 처리량(Output tok/s) 향상을 확인했으며, 이는 실제 프로덕션 환경에서 비용 효율성을 크게 높여줄 것입니다.

참고 자료

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

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