본문으로 건너뛰기

[vllm] vLLM, ROCm 환경에서 AITER MoE 연산 성능 최적화를 위한 환경 변수 노출

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

들어가며

대규모 언어 모델(LLM)의 발전과 함께 Mixture-of-Experts (MoE) 아키텍처가 주목받고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시 활성화되는 파라미터 수를 제한하여 효율성을 높이는 장점이 있습니다. vLLM은 이러한 MoE 모델을 효율적으로 서빙하기 위한 다양한 최적화를 적용하고 있으며, 특히 AMD GPU 환경에서의 성능 향상을 위해 ROCm 및 AITER 라이브러리를 활용하고 있습니다.

이번 PR은 vLLM에서 ROCm 환경을 사용하는 사용자들이 AITER 라이브러리의 MoE 연산에 대한 디스패치 정책을 직접 제어할 수 있도록 하는 새로운 환경 변수를 도입합니다. 이를 통해 특정 모델이나 워크로드에 더 적합한 디스패치 전략을 선택하여 성능을 최적화할 수 있는 가능성을 열었습니다. 기존에는 이 설정이 vLLM 코드 내에서 고정되어 있어 사용자가 임의로 변경하기 어려웠습니다.

코드 변경 분석

이번 PR은 주로 세 개의 파일에서 변경이 이루어졌으며, AITER 라이브러리의 MoE 연산에 moe_sorting_dispatch_policy라는 새로운 파라미터를 노출하고 이를 환경 변수로 제어할 수 있도록 하는 데 중점을 둡니다.

1. vllm/envs.py: 환경 변수 정의

새로운 환경 변수인 VLLM_ROCM_AITER_MOE_DISPATCH_POLICY가 정의되었습니다. 이 변수는 정수형(int) 값을 가지며, 기본값은 0으로 설정되어 기존의 동작 방식을 유지합니다.

--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -116,6 +116,7 @@
     VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
     VLLM_ROCM_USE_AITER_LINEAR: bool = True
     VLLM_ROCM_USE_AITER_MOE: bool = True
+    VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
     VLLM_ROCM_USE_AITER_RMSNORM: bool = True
     VLLM_ROCM_USE_AITER_MLA: bool = True
     VLLM_ROCM_USE_AITER_MHA: bool = True
@@ -1103,6 +1104,17 @@ def _resolve_rust_frontend_path() -> str | None:
     "VLLM_ROCM_USE_AITER_MOE": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")
     ),
+    # MoE sorting dispatch policy for AITER fused MoE kernels.
+    #   0 = auto (default): single-pass for small batches, multi-pass
+    #       for large batches
+    #   1 = always single-pass: one kernel launch, no workspace,
+    #       may be preferred for low-concurrency decode workloads
+    #   2 = always multi-pass: can be faster for MoE-heavy models
+    #       (e.g., +2-5% on Qwen3-Next, +1.5% on DeepSeek-V3 at TP4,
+    #       see PR #39177 for benchmarks)
+    "VLLM_ROCM_AITER_MOE_DISPATCH_POLICY": lambda: int(
+        os.getenv("VLLM_ROCM_AITER_MOE_DISPATCH_POLICY", "0")
+    ),
     # use aiter rms norm op if aiter ops are enabled.
     "VLLM_ROCM_USE_AITER_RMSNORM": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in ("true", "1")

또한, 이 환경 변수에 대한 상세한 설명이 추가되었습니다. 각 정책(0, 1, 2)의 의미와 예상되는 성능 특성이 명시되어 있어 사용자가 쉽게 이해하고 선택할 수 있도록 돕습니다.

2. vllm/_aiter_ops.py: 연산 시그니처 및 캐싱 로직 추가

_rocm_aiter_fused_moe_impl 함수와 _rocm_aiter_fused_moe_fake 함수에 moe_sorting_dispatch_policy 파라미터가 추가되었습니다. 이는 AITER 연산이 이 새로운 정책을 인지하고 사용할 수 있도록 하는 핵심적인 변경입니다.

--- a/vllm/_aiter_ops.py
+++ b/vllm/_aiter_ops.py
@@ -144,6 +144,7 @@ def _rocm_aiter_fused_moe_impl(
     intermediate_pad: int = 0,
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
+    moe_sorting_dispatch_policy: int = 0,
 ) -> torch.Tensor:
     from aiter import ActivationType, QuantType
     from aiter.fused_moe import fused_moe
@@ -171,6 +172,7 @@ def _rocm_aiter_fused_moe_impl(
         intermediate_pad=intermediate_pad,
         bias1=bias1,
         bias2=bias2,
+        moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
     )
 
 
@@ -194,6 +196,7 @@ def _rocm_aiter_fused_moe_fake(
     intermediate_pad: int = 0,
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
+    moe_sorting_dispatch_policy: int = 0,
 ) -> torch.Tensor:
     if output_dtype is not None:
         return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1282,6 +1285,18 @@ class rocm_aiter_ops:
         - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
     */
 
+    _MOE_DISPATCH_POLICY: int | None = None
+
+    @classmethod
+    @if_aiter_supported
+    def get_moe_dispatch_policy(cls) -> int:
+        """Cached MoE sorting dispatch policy."""
+        if cls._MOE_DISPATCH_POLICY is None:
+            import vllm.envs as envs
+
+            cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
+        return cls._MOE_DISPATCH_POLICY
+
     # Check if the env variable is set
     _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
     _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
@@ -1890,6 +1905,7 @@ def fused_moe(
         intermediate_pad: int = 0,
         bias1: torch.Tensor | None = None,
         bias2: torch.Tensor | None = None,
+        moe_sorting_dispatch_policy: int = 0,
     ) -> torch.Tensor:
         return torch.ops.vllm.rocm_aiter_fused_moe(
             hidden_states,
@@ -1911,6 +1927,7 @@ def fused_moe(
             intermediate_pad,
             bias1,
             bias2,
+            moe_sorting_dispatch_policy,
         )
 
     @staticmethod

또한, rocm_aiter_ops 클래스에 get_moe_dispatch_policy 메서드가 추가되었습니다. 이 메서드는 환경 변수 값을 읽어와 캐싱하여 반복적인 호출 시 오버헤드를 줄입니다. 리뷰 과정에서 envs.XXX를 핫 패스(hot path)에서 직접 호출하는 것이 비효율적이라는 지적이 있었고, 이를 반영한 개선입니다.

3. vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py: 실제 연산에 정책 적용

rocm_aiter_fused_experts 함수에서 AITER 연산을 호출할 때, 새로 정의된 rocm_aiter_ops.get_moe_dispatch_policy() 메서드를 통해 읽어온 moe_sorting_dispatch_policy 값을 전달합니다.

--- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
@@ -246,6 +246,7 @@ def rocm_aiter_fused_experts(
     a1q_scale: torch.Tensor | None = None,
     num_local_tokens: torch.Tensor | None = None,
     output_dtype: torch.dtype | None = None,
+    moe_sorting_dispatch_policy: int = 0,
 ) -> torch.Tensor:
     """ROCm AITER fused MoE expert computation."""
     if quant_config is None:
@@ -361,6 +362,7 @@ def rocm_aiter_fused_experts(
             intermediate_pad=intermediate_pad // 64 * 64 * 2,
             bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
             bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
+            moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
         )
 
 
@@ -504,6 +506,7 @@ def apply(
             a1q_scale=a1q_scale,
             num_local_tokens=num_local_tokens,
             output_dtype=output.dtype,
+            moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
         )
         # avoid redundant copy when output is a view of the result
         if (

이 변경을 통해 AiterExperts.apply() 메서드는 이제 환경 변수에 설정된 값에 따라 AITER의 MoE 연산에 다른 디스패치 정책을 적용할 수 있게 됩니다.

왜 이게 좋은가?

성능 향상 가능성

이 PR의 핵심 가치는 사용자에게 AITER MoE 연산의 디스패치 정책을 선택할 수 있는 유연성을 제공한다는 점입니다. MoE 모델의 성능은 전문가(expert) 수, 토큰 분포, 배치 크기, 병렬 처리 설정 등 다양한 요인에 따라 달라질 수 있습니다. AITER 라이브러리에서 제공하는 moe_sorting_dispatch_policy는 이러한 복잡한 상호작용을 고려하여 최적의 성능을 내기 위한 여러 전략을 포함하고 있습니다.

PR 설명과 리뷰 댓글에서 제공된 테스트 결과는 이러한 성능 향상의 가능성을 뒷받침합니다.

  • Qwen3-Next-80B-A3B-Instruct-FP8 모델 테스트:

    • 기본 정책 (policy=0) 대비 policy=2 사용 시 출력 처리량(Output throughput)이 1.6% 증가하고, 평균 토큰 생성 시간(Mean TPOT)은 0.6% 감소하는 결과를 보였습니다.

  • nholmber의 추가 성능 데이터:

    • Qwen3-Next-80B 모델 계열에서 policy=2가 policy=0 대비 최대 5.3%의 TPOT 속도 향상을 보였습니다 (예: conc8_speedup% 105%).
    • 반면, DeepSeek-V3.2, MiniMax-M2.5, gpt-oss-120b와 같은 다른 모델에서는 policy=0이 더 나은 성능을 보이거나 큰 차이가 없는 경우도 있었습니다.

이러한 결과는 moe_sorting_dispatch_policy의 값이 모델 및 워크로드에 따라 최적의 성능을 내는 전략이 다르다는 것을 명확히 보여줍니다. 사용자가 이 값을 직접 제어함으로써, 자신의 환경에 맞는 최적의 설정을 찾아 성능을 극대화할 수 있습니다.

일반적인 교훈

  1. 사용자 제어권 확대: 라이브러리 내부의 민감한 성능 관련 파라미터를 환경 변수 등을 통해 사용자에게 노출하는 것은 매우 중요합니다. 이는 사용자가 특정 환경이나 모델에 맞게 라이브러리를 미세 조정할 수 있게 하여 전반적인 생태계의 성능을 향상시킵니다.
  2. 성능 튜닝의 복잡성: MoE와 같은 복잡한 아키텍처에서는 단일 최적 설정이 존재하지 않습니다. 다양한 파라미터와 환경 요인이 성능에 영향을 미치므로, 실험과 벤치마킹을 통해 최적의 설정을 찾아야 합니다.
  3. 코드 리뷰의 중요성: 리뷰 과정에서 제기된 envs.XXX 호출 오버헤드 문제는 핫 패스에서의 성능 최적화를 위해 캐싱 메커니즘을 도입하는 계기가 되었습니다. 이는 코드 품질과 성능을 동시에 향상시키는 좋은 예입니다.
  4. AITER 라이브러리 의존성 관리: 이 PR은 AITER 라이브러리의 특정 버전 또는 수정 사항(ROCm/aiter#2639)에 대한 의존성을 가지고 있었습니다. 이러한 외부 라이브러리 의존성을 명확히 하고 관리하는 것이 중요합니다.

리뷰 피드백 반영

  • 환경 변수 남발에 대한 우려: gshtras와 tpopp의 논의처럼, 새로운 환경 변수를 계속 추가하는 것에 대한 우려가 있었습니다. 이에 대한 대안으로 내부 설정을 더 체계적으로 관리하는 방안(#41159)이 제안되었으나, 해당 리팩토링이 완료되기 전까지는 이 PR을 진행하는 것이 합리적이라는 결론이 내려졌습니다.
  • 캐싱 메커니즘 도입: tjtanaa의 지적에 따라, envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY를 핫 패스에서 직접 읽는 대신 vllm/_aiter_ops.py에 get_moe_dispatch_policy 클래스 메서드를 추가하여 값을 캐싱하도록 개선되었습니다. 이는 성능 오버헤드를 줄이는 중요한 수정입니다.
  • 문서화 강화: AndreasKaratzas의 제안에 따라, 각 디스패치 정책(0, 1, 2)에 대한 설명과 예상되는 사용 사례를 vllm/envs.py의 주석에 추가하여 사용자가 더 쉽게 이해하고 활용할 수 있도록 개선되었습니다.
  • AITER 버전 의존성: 초기에는 AITER 라이브러리의 특정 수정 사항이 필요했지만, 이후 해당 수정 사항이 병합되면서 의존성 문제가 해결되었습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글