[vllm] vLLM FlashInfer Allreduce 최적화: 단일 노드 성능 회복과 안정성 확보

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

들어가며

vLLM은 대규모 언어 모델(LLM) 추론을 위한 고성능 서빙 엔진으로, 그 핵심에는 효율적인 GPU 활용과 분산 처리 기술이 있습니다. 특히, 여러 GPU가 협력하여 작업을 수행하는 분산 환경에서는 Allreduce와 같은 통신 연산의 효율성이 전체 시스템 성능에 지대한 영향을 미칩니다. 이번에 분석할 GitHub PR(vllm-project/vllm#6804)은 vLLM의 FlashInfer Allreduce 백엔드 선택 로직을 최적화하여, 최근 발생했던 단일 노드(single-node) 환경에서의 성능 회귀(regression)를 해결하고 안정성을 확보한 사례입니다.

문제의 발단: FlashInfer Allreduce 백엔드 선택의 오해

vLLM v0.32.0 버전에서 약 3-5%의 디코딩 성능 저하가 발생했습니다. 이 문제는 이전 PR인 #56891에 의해 야기되었습니다. 해당 PR은 FlashInfer Allreduce의 자동 백엔드 선택 로직을 변경하여, 모든 노드 로컬(node-local) 프로세스 그룹에 대해 trtllm 백엔드를 사용하도록 했습니다. 그 전에는 단일 노드 실행 시 mnnvl을 기본으로 사용하고, NVSwitch 멀티캐스트를 지원하지 않는 노드에서만 trtllm으로 폴백(fallback)하도록 되어 있었습니다.

문제는 #56891이 성능 데이터 없이 병합되었고, 단일 노드 NVSwitch 시스템에서는 trtllm이 mnnvl보다 느리다는 사실이 간과되었다는 점입니다. 이로 인해 단일 노드, 특히 NVSwitch가 장착된 시스템에서 TP(Tensor Parallelism)가 1보다 큰 워크로드에서 성능 저하가 발생했습니다.

코드 분석: 올바른 백엔드 선택으로의 회귀

이번 PR은 단일 노드 환경에서 FlashInfer Allreduce 백엔드 선택을 다시 mnnvl로 선호하도록 변경하고, 동시에 trtllm으로의 안전한 폴백 메커니즘을 복원합니다. 핵심 변경 사항은 두 파일에서 이루어졌습니다.

1. tests/distributed/test_comm_ops.py

이 파일은 분산 통신 연산의 동작을 테스트하는 부분입니다. 단일 노드 환경에서의 FlashInfer Allreduce 백엔드 선택에 대한 기대치를 업데이트하여, 변경된 로직을 반영합니다.

Before:

@pytest.mark.parametrize(
    ("node_count", "is_blackwell", "expected"),
    [
        (1, False, ("trtllm", False)),
        (1, True, ("trtllm", False)),
        (2, False, (None, False)),
        (2, True, ("mnnvl", False)),
    ],
)

After:

@pytest.mark.parametrize(
    ("node_count", "is_blackwell", "expected"),
    [
        (1, False, ("mnnvl", True)),
        (1, True, ("mnnvl", True)),
        (2, False, (None, False)),
        (2, True, ("mnnvl", False)),
    ],
)

단일 노드(node_count=1)의 경우, expected 백엔드가 trtllm에서 mnnvl로 변경되었고, allow_trtllm_fallback 플래그가 False에서 True로 변경된 것을 확인할 수 있습니다. 이는 mnnvl을 선호하되, 문제가 발생하면 trtllm으로 전환할 수 있도록 하는 새로운 정책을 명확히 보여줍니다.

2. vllm/distributed/device_communicators/flashinfer_all_reduce.py

이 파일은 FlashInfer Allreduce 백엔드를 실제로 결정하는 핵심 로직을 담고 있습니다. _resolve_fi_ar_backend 함수 내에서 node_count == 1인 경우의 백엔드 선택 로직이 변경되었습니다.

Before:

    node_count = _node_count(group)
    if node_count == 1:
        # TRTLLM supports single-node allreduce on Hopper and Blackwell. Do not
        # create an MNNVL workspace when the process group is node-local.
        backend = "trtllm"
        allow_trtllm_fallback = False

After:

    node_count = _node_count(group)
    if node_count == 1:
        # MNNVL decodes faster than TRTLLM on single-node NVSwitch systems
        # (about 3-5% throughput at TP4/TP8 on H200 and B200). Fall back to
        # TRTLLM where the node lacks NVSwitch multicast and MNNVL cannot
        # create its workspace.
        backend = "mnnvl"
        allow_trtllm_fallback = True

주석에서 명확히 설명하듯이, 단일 노드 NVSwitch 시스템에서 mnnvl이 trtllm보다 3-5% 더 빠른 디코딩 성능을 제공한다는 점을 인지하고 mnnvl을 기본 백엔드로 설정했습니다. 또한, NVSwitch 멀티캐스트를 지원하지 않아 mnnvl 워크스페이스를 생성할 수 없는 경우를 대비하여 trtllm으로의 폴백을 allow_trtllm_fallback = True를 통해 다시 활성화했습니다. 이는 성능과 안정성을 동시에 고려한 현명한 접근 방식입니다.

왜 이 최적화가 좋은가?

이 PR은 단순한 코드 변경을 넘어, vLLM의 성능과 안정성에 여러 긍정적인 영향을 미칩니다.

1. 명확한 성능 향상

PR 설명에 포함된 A/B 테스트 결과는 이 최적화의 효과를 명확히 보여줍니다.

workload	default (trtllm)	mnnvl	개선율
Qwen3.5-397B NVFP4, 8xB200, 8k/1k conc 128	5462 tok/s, TPOT 21.78 ms	5728 tok/s, TPOT 20.22 ms	약 4.8%
GLM-5.3-Flash TP4, B200, 3 repeats	2533 / 2753 / 2771 tok/s	2581 / 2811 / 2845 tok/s	약 1.9% ~ 3.2%

Qwen3.5-397B 워크로드에서 토큰 처리량이 약 4.8% 증가하고, TPOT(Time Per Output Token)이 21.78ms에서 20.22ms로 감소한 것은 매우 유의미한 개선입니다. 이는 단일 노드에서 mnnvl이 NVSwitch의 멀티캐스트 기능을 활용하여 Allreduce 연산을 더 효율적으로 처리하기 때문입니다.

2. 올바른 백엔드 선택 로직 복원

이전 PR에서 trtllm을 기본으로 설정했던 것은 단일 노드 NVSwitch 시스템의 특성을 고려하지 않은 결정이었습니다. 이 PR은 각 백엔드의 강점과 약점을 정확히 파악하고, 특정 환경(단일 노드 NVSwitch)에서 더 나은 성능을 제공하는 mnnvl을 다시 선호하도록 함으로써, 시스템의 잠재력을 최대한 활용합니다.

3. 안정성을 위한 견고한 폴백 메커니즘

allow_trtllm_fallback = True로 설정한 것은 이 최적화의 중요한 부분입니다. 리뷰어 elvircrn의 초기 피드백에서 mnnvl 백엔드 초기화 실패(Failed to initialize FlashInfer Allreduce norm fusion workspace with backend=mnnvl) 사례가 언급되었습니다. 이는 mnnvl이 모든 환경에서 항상 성공적으로 작동하는 것은 아님을 시사합니다. 특히 NVSwitch 멀티캐스트가 없거나 특정 하드웨어/소프트웨어 구성에서 문제가 발생할 수 있습니다.

이 PR은 mnnvl을 선호하되, 만약 mnnvl이 워크스페이스를 생성할 수 없는 경우 trtllm으로 안전하게 폴백하도록 하여 시스템의 안정성을 보장합니다. elvircrn은 이후 Kimi K3 모델을 H200 노드에서 테스트하여, 이 PR이 병합된 브랜치에서 mnnvl 초기화 실패 경고 없이 벤치마크가 성공적으로 완료되었음을 확인했습니다. 이는 단일 노드 성능 개선이 멀티 노드 안정성을 해치지 않음을 입증합니다.

일반적인 교훈

이 사례는 성능 최적화에서 다음과 같은 중요한 교훈을 제공합니다.

• 벤치마크의 중요성: 새로운 변경 사항을 적용할 때는 반드시 다양한 워크로드와 하드웨어 구성에서 성능 벤치마크를 수행해야 합니다. 이론적인 개선이 실제 환경에서 항상 적용되는 것은 아닙니다.
• 특정 환경에 대한 이해: 분산 시스템에서는 하드웨어(예: NVSwitch)와 소프트웨어(예: 특정 Allreduce 백엔드)의 상호작용이 복잡합니다. 각 환경의 특성을 정확히 이해하고 최적의 솔루션을 선택하는 것이 중요합니다.
• 견고한 폴백 설계: 고성능을 추구하되, 예외 상황에 대비한 견고한 폴백 메커니즘을 마련하여 시스템의 안정성과 호환성을 유지해야 합니다.

결론

이번 vLLM PR은 FlashInfer Allreduce 백엔드 선택 로직을 재조정하여 단일 노드 NVSwitch 시스템에서의 성능 회귀를 성공적으로 해결했습니다. mnnvl 백엔드를 다시 선호하고 trtllm으로의 안전한 폴백을 활성화함으로써, vLLM은 더 높은 처리량과 안정성을 동시에 달성했습니다. 이는 미묘한 성능 병목 현상을 식별하고 해결하는 시니어 엔지니어의 통찰력과, 실제 환경에서의 철저한 검증의 중요성을 보여주는 좋은 예시입니다.

참고 자료

• https://github.com/flashinfer/flashinfer
• https://docs.vllm.ai/en/latest/runtimes/distributed_inference.html

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