[vllm] vLLM CI 속도 개선: 70분 걸리던 MoE 테스트를 5분으로 단축하기

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

들어가며

대규모 언어 모델(LLM) 서빙 엔진인 vLLM 프로젝트에서 CI(지속적 통합) 시간은 개발 생산성에 직결되는 중요한 요소입니다. 특히 Mixture of Experts(MoE)와 Marlin 양자화 커널은 다양한 파라미터 조합(M, N, K 크기, Expert 수, 양자화 타입 등)을 가지고 있어, 이를 단순한 itertools.product로 전수 조사할 경우 테스트 케이스가 기하급수적으로 늘어나는 문제가 발생합니다.

기존 test_fused_marlin_moe는 약 2,880개의 테스트 케이스를 실행하며 완료까지 약 70분이 소요되었습니다. 이는 전체 CI 파이프라인의 병목 구간이었습니다. 이번 PR은 이러한 조합 폭발(Combinatorial Explosion) 문제를 해결하여 테스트 케이스를 153개로 줄이고, 실행 시간을 5분 내외로 단축한 최적화 사례입니다.

코드 분석: 무차별 대입에서 전략적 샘플링으로

1. 테스트 시나리오의 명시적 정의

가장 큰 변화는 모든 변수의 조합을 생성하는 대신, 실제 엣지 케이스와 일반적인 워크로드를 대변하는 MARLIN_MOE_SCENARIOS를 명시적으로 정의한 것입니다.

Before:

m_list = [1, 123, 666]
n_list = [128, 1024]
k_list = [256, 2048]
e_list = [5, 12]
topk_list = [2, 3]
ep_size_list = [1, 4]
act_order_list = [True, False]
is_k_full_list = [True, False]

all_combinations = itertools.product(
    MOE_MARLIN_QUANT_TEST_CONFIGS,
    m_list, n_list, k_list, e_list, topk_list, ep_size_list, act_order_list, is_k_full_list,
)

위 코드는 모든 리스트의 원소를 곱하여 수천 개의 조합을 만듭니다. 하지만 실제로는 m=1(Single token)인 경우와 m=128(Batch)인 경우, 그리고 Unaligned m(133 등)인 경우만 확인해도 커널의 안정성을 충분히 검증할 수 있습니다.

After:

MARLIN_MOE_SCENARIOS = [
    # (m, n, k, e, topk, ep_size, act_order, is_k_full)
    (1, 128, 256, 5, 2, 1, False, True),      # Single token, small
    (1, 1024, 2048, 5, 2, 1, False, True),    # Single token, large
    (133, 256, 256, 5, 2, 1, False, True),    # Unaligned m, small
    (128, 1024, 2048, 12, 3, 1, False, True), # Aligned batch, large
    (64, 1024, 2048, 12, 3, 4, False, True),  # Expert parallelism
    # ... (중략)
]

이렇게 유의미한 시나리오 9가지만 선별함으로써, 불필요하게 중복되는 계산 리소스를 획기적으로 절감했습니다.

2. 필터링 로직의 개선 (is_invalid -> is_valid)

기존에는 잘못된 조합을 걸러내는 is_invalid 함수를 사용했으나, 로직이 복잡해질수록 가독성이 떨어졌습니다. 이를 is_valid로 전환하여 테스트 가능한 조건만 명확히 기술하도록 변경되었습니다.

Before/After 비교:

-    def is_invalid(
+    def is_valid(
         a_type,
         b_type,
         c_type,
@@ -715,39 +721,42 @@ def is_invalid(
         group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
         if group_size > 0 and k % group_size != 0:
             return False
-        if act_order and group_size in [-1, k, n]:
-            return False
-        if group_size in [k, n]:
-            return False
-        if not act_order and is_k_full:
+        if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
             return False
-        return a_type.size_bits < 16 or a_type is c_type
+        return a_type.size_bits < 16 or a_type is c_type

특히 b_type이 FP8(float8_e4m3fn)인 경우 특정 group_size와 is_k_full 조건에서 발생할 수 있는 하드웨어 제약 사항을 명확히 반영하여 테스트의 정확도를 높였습니다.

왜 이게 좋은가?

1. CI 비용 및 시간의 극적인 감소

• 테이스 케이스 수: 2,880개 → 153개 (약 95% 감소)
• 실행 시간: 70분 → 5분 (약 14배 속도 향상)
• H100 노드 점유 시간: 80분 → 20분

CI 시간의 단축은 단순한 수치 이상의 의미를 갖습니다. 개발자가 PR을 올린 후 결과를 확인하기까지의 피드백 루프가 짧아지며, 이는 전체 팀의 개발 속도 향상으로 이어집니다.

2. 테스트 커버리지의 질적 유지

단순히 개수만 줄인 것이 아니라, Unaligned m, Expert Parallelism, Act order 등 커널 로직에서 버그가 발생하기 쉬운 지점들을 MARLIN_MOE_SCENARIOS에 포함시켰습니다. 양보다 질에 집중한 전략입니다.

3. 유지보수 용이성

새로운 양자화 방식이나 하드웨어 지원이 추가될 때, 무거운 itertools.product 결과값을 뒤지는 대신 MARLIN_MOE_SCENARIOS 리스트에 한 줄의 시나리오만 추가하면 되므로 관리가 훨씬 쉬워집니다.

결론

소프트웨어 엔지니어링에서 테스트는 필수적이지만, 그 비용이 개발을 저해할 정도로 커진다면 최적화의 대상이 됩니다. 이번 vLLM의 개선 사례는 "모든 조합을 테스트하는 것보다, 의미 있는 경계값(Boundary Value)을 테스트하는 것이 훨씬 효율적이다"라는 테스팅의 기본 원칙을 잘 보여줍니다.

참고 자료

• https://github.com/vllm-project/vllm
• https://docs.python.org/3/library/itertools.html#itertools.product

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