[vllm] vLLM ROCm 환경에서 FlyDSL을 활용한 MXFP8 MoE 성능 최적화

PR 링크: vllm-project/vllm#46184 상태: Merged | 변경: +374 / -4

들어가며

vLLM은 다양한 하드웨어 가속기를 지원하기 위해 노력하고 있습니다. 특히 AMD ROCm 환경에서 대규모 언어 모델(LLM)의 Mixture-of-Experts(MoE) 추론 성능은 매우 중요한 과제입니다. 이번 업데이트는 AMD의 최신 GPU 아키텍처인 gfx950에서 Minimax-M3 mxfp8 가중치를 활용할 때, AITER의 FlyDSL(Fly-DSL) 커널을 사용하여 추론 성능을 극대화하는 것을 목표로 합니다. 기존 Triton 기반의 dot-scaled 구현체와 비교하여 더 높은 처리량(throughput)과 낮은 지연 시간(TPOT)을 달성했습니다.

코드 분석

1. AITER FlyDSL 커널 통합 (vllm/_aiter_ops.py)

기존 AITER 연산 인터페이스에 swiglu_limit 파라미터를 추가하고, MXFP8 가중치를 FlyDSL 레이아웃으로 변환하는 shuffle_mxfp8_moe_weights 함수를 구현했습니다.

# Before: swiglu_limit 파라미터 부재
def _rocm_aiter_fused_moe_impl(..., moe_sorting_dispatch_policy: int = 0) -> torch.Tensor:

# After: swiglu_limit 추가 및 가중치 셔플 로직 구현
def _rocm_aiter_fused_moe_impl(..., swiglu_limit: float = 0.0) -> torch.Tensor:
    # ...
    @staticmethod
    def shuffle_mxfp8_moe_weights(w13, w2, w13_scale, w2_scale):
        w13 = shuffle_weight(w13, is_guinterleave=True, gate_up=True)
        # ...

2. 백엔드 선택 로직 및 검증 (vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py)

is_aiter_mxfp8_moe_available 함수를 통해 FlyDSL 커널이 사용 가능한 환경(gfx950 및 필수 패키지 존재 여부)인지 확인합니다. 이를 통해 실패 시 Triton 백엔드로 안전하게 폴백(fallback)할 수 있습니다.

# FlyDSL 사용 가능 여부 확인
def is_aiter_mxfp8_moe_available() -> bool:
    if not (current_platform.is_rocm() and current_platform.supports_mx()):
        return False
    # ... aiter 설정 파일 존재 여부로 검증

왜 이게 좋은가

이번 최적화는 실제 벤치마크에서 유의미한 성능 향상을 보여주었습니다. 특히 TP8(Tensor Parallel 8) 설정에서 1k/1k 토큰 처리 시 처리량이 1316% 증가했으며, TPOT(Time Per Output Token)은 1214% 감소했습니다.

핵심 교훈

1. 자동 백엔드 선택: 하드웨어와 라이브러리 지원 여부를 런타임에 체크하여 최적의 커널을 자동으로 선택함으로써 사용자 경험을 개선했습니다.
2. 안전한 폴백(Fallback): 새로운 고성능 커널을 도입할 때, 지원되지 않는 환경에서 명확한 에러 메시지와 함께 기존 구현체로 돌아가는 설계는 시스템 안정성을 높입니다.
3. 레이아웃 최적화: 가중치를 커널이 선호하는 메모리 레이아웃으로 미리 셔플(shuffle)하여 런타임 오버헤드를 최소화했습니다.

리뷰어 피드백 반영

리뷰 과정에서 swiglu_limit 파라미터가 구버전 AITER API와 호환되지 않는 문제가 발견되었습니다. 이를 해결하기 위해 AITER 버전을 v0.1.16.post2로 범핑하여 호환성을 확보했습니다. 또한, 백엔드 선택 로직에서 flydsl이라는 별도 키워드 대신 aiter 백엔드 내부에서 적절한 커널을 선택하도록 구조를 개선했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://github.com/ROCm/aiter

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