본문으로 건너뛰기

[vllm] vLLM의 MoE 튜닝 스크립트 최적화: RDNA 아키텍처를 위한 성능 개선

PR 링크: vllm-project/vllm#41197 상태: Merged | 변경: +38 / -6

들어가며

vLLM의 fused_moe_triton 커널은 모델 성능에 매우 중요한 역할을 합니다. 특히 GLM4.7-Flash와 같은 모델에서 최적의 튜닝 설정은 최대 15%의 속도 향상을 가져올 수 있습니다. 하지만 기존의 튜닝 스크립트는 RDNA(gfx11/gfx12) 아키텍처에서 매우 비효율적이었으며, 8개의 GPU를 사용해도 4시간 이상 소요되는 긴 시간과 함께 Ray 워커가 멈추는 크래시 문제가 빈번했습니다. 본 PR은 RDNA 아키텍처에 맞지 않는 파라미터를 제거하고, 런타임 안정성을 높여 튜닝 프로세스를 최적화합니다.

코드 분석

1. RDNA 아키텍처에 맞춘 탐색 공간(Search Space) 최적화

기존 스크립트는 CDNA 아키텍처 전용 파라미터인 matrix_instr_nonkdim과 kpack을 RDNA에서도 탐색하고 있었습니다. 이는 RDNA에서 "no matching matrix core intrinsic" 에러를 유발하는 원인이었습니다.

Before:

matrix_instr_nonkdim_range = [16, 32] if use_fp16 else []
kpack_range = [1, 2] if use_fp16 else []
# ...
if use_fp16:
    param_ranges["matrix_instr_nonkdim"] = matrix_instr_nonkdim_range
    param_ranges["kpack"] = kpack_range

After:

# RDNA doesn't use matrix_instr_nonkdim nor kpack
if use_fp16 and on_cdna():
    param_ranges["matrix_instr_nonkdim"] = [16, 32]
    param_ranges["kpack"] = [1, 2]

2. 레지스터 스필링(Register Spilling) 방지를 위한 Pruning

RDNA 아키텍처의 VGPR(Vector General Purpose Register) 제한을 고려하여, 과도한 리소스를 사용하는 설정을 사전에 필터링하도록 로직을 추가했습니다.

if on_rdna():
    accum_vgprs = (
        (BLOCK_SIZE_M // WMMA_TILE)
        * (BLOCK_SIZE_N // WMMA_TILE)
        // num_warps
        * 8
    )
    if accum_vgprs >= 256:
        continue

3. Triton 캐시 관리 및 런타임 예외 처리

Triton 3.6+ 버전에서 캐시를 올바르게 정리하고, 튜닝 중 발생하는 예외를 적절히 처리하여 전체 프로세스가 중단되지 않도록 개선했습니다.

except RuntimeError as e:
    logger.warning("Skipping config %s: %s", config, e)
    continue

왜 이게 좋은가

  1. 탐색 공간 축소: 불필요한 파라미터(matrix_instr_nonkdim, kpack)를 제거함으로써 탐색 공간을 최대 4배까지 줄였습니다. 이는 튜닝 시간을 획기적으로 단축합니다.
  2. 안정성 확보: 런타임 에러 발생 시 전체 프로세스가 중단되던 문제를 해결했습니다. 개별 설정에서 발생하는 RuntimeError를 로깅하고 건너뛰도록 하여, 긴 튜닝 작업이 중도에 실패하는 리스크를 제거했습니다.
  3. 하드웨어 최적화: RDNA의 WMMA(Wave Matrix Multiply Accumulate) 특성을 고려한 VGPR 제한 로직을 통해, 커널 크래시를 방지하고 하드웨어 리소스를 효율적으로 사용합니다.

일반적 교훈: 대규모 파라미터 탐색(Autotuning)을 수행할 때는 하드웨어 아키텍처별 제약 사항을 명확히 구분해야 합니다. 또한, 수 시간이 소요되는 작업에서는 예외 처리를 통해 부분적인 실패가 전체 작업의 실패로 이어지지 않도록 방어적인 코드를 작성하는 것이 필수적입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글