본문으로 건너뛰기

[vllm] Blackwell을 위한 새로운 MLA 백엔드: TOKENSPEED_MLA 분석 (DeepSeek R1 최적화)

PR 링크: vllm-project/vllm#41778 상태: Merged | 변경: +0 / -0

들어가며

최근 LLM 추론 분야에서 가장 뜨거운 주제 중 하나는 DeepSeek-R1이 채택한 MLA(Multi-head Latent Attention) 구조입니다. MLA는 KV 캐시의 크기를 획기적으로 줄이면서도 성능을 유지하지만, 그 복잡한 구조로 인해 하드웨어 가속기에서 최적의 성능을 내기가 매우 까다롭습니다.

이번에 vLLM 프로젝트에 병합된 PR은 NVIDIA의 최신 아키텍처인 Blackwell(SM100) 환경에서 DeepSeek-R1 및 Kimi K25 모델의 추론 속도를 극대화하기 위해 TOKENSPEED_MLA라는 새로운 백엔드를 도입했습니다. 이 백엔드는 CuTe DSL(Domain Specific Language)로 작성된 커널을 사용하여, 특히 MTP(Multi-Token Prediction) 시나리오와 FP8 KV 캐시 환경에서 기존의 TRT-LLM이나 FlashInfer를 압도하는 성능을 보여줍니다.

본 글에서는 이 PR이 구체적으로 어떤 코드를 변경했으며, 왜 Blackwell 환경에서 중요한 최적화인지 분석해 보겠습니다.

코드 분석: 백엔드 통합과 구조적 개선

1. 백엔드 선택 로직의 현대화 (Registry-based Enum)

기존 vLLM의 MLA 백엔드 선택 로직은 여러 개의 불리언(boolean) 플래그를 사용하여 다소 복잡하고 확장이 어려운 구조였습니다. 이번 PR에서는 이를 MLAPrefillBackendEnum 기반의 레지스트리 방식으로 전환하여 코드의 가독성과 유지보수성을 높였습니다.

Before:

# mla_runner.py (기존 방식: 개별 플래그 설정)
if prefill_backend is not None:
    prefill_cfg = get_prefill_backend_config(prefill_backend)
    if prefill_cfg["flash_attn_version"] is not None:
        vllm_config.attention_config.flash_attn_version = prefill_cfg["flash_attn_version"]
    vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg["disable_flashinfer_prefill"]
    vllm_config.attention_config.use_cudnn_prefill = prefill_cfg["use_cudnn_prefill"]
    vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg["use_trtllm_ragged_deepseek_prefill"]

After:

# mla_runner.py (개선된 방식: Enum 기반 레지스트리 활용)
if prefill_backend is not None:
    prefill_cfg = get_prefill_backend_config(prefill_backend)
    if prefill_cfg.get("mla_prefill_backend_enum") is not None:
        from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackendEnum
        vllm_config.attention_config.mla_prefill_backend = MLAPrefillBackendEnum[
            prefill_cfg["mla_prefill_backend_enum"]
        ]
    else:
        # 하위 호환성을 위한 기존 로직 유지
        ...

이 변경은 단순히 코드가 깔끔해진 것을 넘어, TOKENSPEED_MLA와 같은 새로운 하드웨어 특화 백엔드를 추가할 때 기존 로직을 건드리지 않고 Enum에 등록만 하면 되도록 설계되었습니다.

2. Blackwell 아키텍처에서의 우선순위 조정

PR은 Blackwell(SM100) 환경에서 TOKENSPEED_MLA를 매우 높은 우선순위로 배치했습니다. 이는 해당 커널이 Blackwell의 하드웨어 특성(FP8 연산 가속 등)을 가장 잘 활용한다는 자신감을 보여줍니다.

docs/design/attention_backends.md 변경사항:

 | Priority | Backend |
 | -------- | ------- |
 | 1 | `FLASHINFER_MLA` |
+| 2 | `TOKENSPEED_MLA` |
 | 3 | `CUTLASS_MLA` |
 | 4 | `FLASH_ATTN_MLA` |

Blackwell 환경에서 FLASHINFER_MLA 다음으로 높은 2순위(Priority 2)를 부여받았습니다. 이는 범용적인 FlashInfer가 대응하지 못하는 특정 워크로드(예: 대규모 배치 사이즈의 Decode)에서 TOKENSPEED_MLA가 주력으로 사용될 것임을 시사합니다.

왜 이게 좋은 최적화인가?

1. MTP(Multi-Token Prediction) 시나리오에서의 압도적 성능

최근 DeepSeek-V3 등에서 사용되는 MTP 기술은 한 번에 여러 토큰을 예측합니다. 이 경우 q_len이 1보다 커지게 되는데, tokenspeed_mla_decode 커널은 이 영역에서 놀라운 효율을 보여줍니다.

  • MTP-7 (q_len=8) 환경에서 Batch Size 12일 때, 기존 TRT-LLM 대비 4.23배의 속도 향상을 기록했습니다.
  • 이는 CuTe DSL을 통해 Blackwell의 공유 메모리(Shared Memory)와 레지스터 레이아웃을 정밀하게 제어하여, 여러 토큰의 쿼리가 KV 캐시를 스캔할 때 발생하는 오버헤드를 최소화했기 때문입니다.

2. KV Length 스케일링의 효율성

벤치마크 결과에 따르면, TOKENSPEED_MLA는 KV 길이가 길어질수록 성능 우위가 더 뚜렷해집니다.

bs kv=1024 kv=65536
64 0.65× 3.88×

TRT-LLM은 KV 스캔량에 따라 성능이 선형적으로 하락하는 반면, TOKENSPEED_MLA는 하위 선형적(Sub-linear)으로 스케일링됩니다. 이는 긴 문맥(Long Context)을 처리해야 하는 최신 LLM 서비스에서 매우 중요한 경쟁력입니다.

3. 고정 오버헤드의 상쇄 (Amortization)

Prefill 단계에서 TOKENSPEED_MLA는 약 120µs의 고정된 런치 오버헤드(Launch Overhead)를 가집니다. 이 때문에 짧은 문맥에서는 TRT-LLM보다 느릴 수 있지만, 문맥이 64k 이상으로 길어지면 이 오버헤드가 전체 연산 시간에서 차지하는 비중이 줄어들며 성능이 대등해집니다.

결론 및 교훈

이번 PR은 특정 하드웨어(Blackwell)와 특정 모델 구조(MLA)에 특화된 커널이 범용 커널을 어떻게 압도할 수 있는지를 잘 보여줍니다.

  1. 하드웨어 특화 DSL의 힘: CuTe DSL과 같은 도구를 사용하여 하드웨어의 잠재력을 끝까지 끌어올리는 것이 중요합니다.
  2. 워크로드에 따른 백엔드 전략: 모든 상황에서 최고의 백엔드는 없습니다. vLLM은 이번 PR을 통해 Prefill/Decode, 문맥 길이, 배치 사이즈에 따라 최적의 백엔드를 선택할 수 있는 더 넓은 선택지를 갖게 되었습니다.
  3. MTP 최적화의 중요성: 향후 추론 엔진의 핵심 경쟁력은 단일 토큰 생성을 넘어 MTP와 같은 병렬 생성 기법을 얼마나 잘 가속하느냐에 달려 있음을 알 수 있습니다.

Blackwell 서버를 운영하거나 DeepSeek-R1 급의 대형 모델을 서비스하려는 엔지니어라면, 이번에 추가된 TOKENSPEED_MLA 백엔드 설정을 반드시 검토해 보시기 바랍니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글