[vllm] vLLM의 ROCm 환경에서 듀얼 스트림 디코드를 통한 성능 최적화

2026년 8월 30일수정: 2026년 8월 30일

PR 링크: vllm-project/vllm#52033 상태: Merged | 변경: +70 / -57

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 추론 성능 최적화는 매우 중요한 과제가 되었습니다. 특히, vLLM은 KV 캐싱과 같은 혁신적인 기법으로 LLM 추론 속도를 크게 향상시켰습니다. 이번 PR (#52033)은 vLLM의 ROCm 환경에서의 성능을 더욱 끌어올리기 위한 중요한 개선 사항을 담고 있습니다. 핵심은 Mixture-of-Experts (MoE) 모델에서 '듀얼 스트림 디코드' 기능을 최적화하고, 특정 모델에서 발생하던 정확도 문제를 해결하는 것입니다. 이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 왜 이러한 변경이 성능 및 정확도 향상에 기여하는지 자세히 살펴보겠습니다.

코드 변경 분석

이번 PR의 핵심은 vllm/model_executor/layers/fused_moe/runner/moe_runner.py와 vllm/model_executor/layers/fused_moe/runner/shared_experts.py 파일의 수정입니다. 이 변경들은 ROCm 환경에서 MoE 모델의 듀얼 스트림 디코드 동작 방식을 개선하는 데 초점을 맞추고 있습니다.

1. `moe_runner.py`: 듀얼 스트림 로직 개선

moe_runner.py 파일에서는 _apply_quant_method 함수와 _forward_impl 함수가 주로 수정되었습니다. 이전에는 듀얼 스트림 디코드 시 공유 전문가(shared experts)의 실행 순서와 동기화 방식에 일부 문제가 있었습니다.

Before:

def _apply_quant_method(
    self,
    router_logits: torch.Tensor,
    shared_experts_input: torch.Tensor | None,
    input_ids: torch.Tensor | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
    ...
    self._maybe_apply_shared_experts(
        shared_experts_input,
        SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
    )
    ...

After:

def _apply_quant_method(
    self,
    router_logits: torch.Tensor,
    shared_experts_input: torch.Tensor | None,
    input_ids: torch.Tensor | None = None,
    shared_experts_overlapping: bool = False, # New argument
) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
    ...
    if shared_experts_overlapping:
        assert self._shared_experts is not None
        self._shared_experts.wait() # New call
    ...

_apply_quant_method 함수에 shared_experts_overlapping이라는 새로운 인자가 추가되었습니다. 이 인자는 공유 전문가가 비동기적으로 실행되었는지 여부를 나타냅니다. 만약 비동기적으로 실행되었다면, self._shared_experts.wait()를 호출하여 메인 스트림에서 해당 작업이 완료될 때까지 기다립니다.

또한, _forward_impl 함수에서는 공유 전문가의 비동기 실행 로직이 _shared_experts.maybe_forward_async 호출로 변경되었습니다. 이 함수는 공유 전문가를 별도의 스트림에 비동기적으로 올리고, 실제로 듀얼 스트림 오버랩이 가능한 경우에만 True를 반환합니다.

Before (_forward_impl):

# Sync aux and main stream for shared expert multi-stream overlap.
self._maybe_sync_shared_experts_stream(shared_experts_input)

After (_forward_impl):

# If using multi-stream overlap for shared experts, we must launch it
# before routed expert dispatch.
some_overlap_possible = False
if self._shared_experts is not None:
    shared_experts_overlapping = self._shared_experts.maybe_forward_async(
        shared_experts_input
    )

이 변경은 공유 전문가의 실행을 라우팅된 전문가(routed experts)의 디스패치 이전에 비동기적으로 시작할 수 있도록 하여, ROCm 환경에서 진정한 듀얼 스트림 오버랩을 가능하게 합니다. 이전에는 _maybe_sync_shared_experts_stream과 같은 함수를 통해 동기화가 이루어졌지만, 이는 실제 오버랩을 보장하지 못하는 경우가 있었습니다.

2. `shared_experts.py`: 멀티 스트림 안전성 검사 및 로직 분리

shared_experts.py 파일에서는 SharedExperts 클래스의 __init__ 메서드와 maybe_forward_async, wait 메서드가 추가/수정되었습니다.

__init__ 메서드 변경:

SharedExperts 클래스의 생성자에 is_multistream_safe 콜백 함수가 추가되었습니다. 이 함수는 공유 전문가의 입력과 라우팅된 전문가의 입력이 동일한 텐서를 참조하는지 (alias) 여부를 검사합니다. 만약 동일한 텐서를 참조한다면, 멀티 스트림 모드에서 경쟁 상태(race condition)가 발생할 수 있으므로 듀얼 스트림 오버랩을 비활성화해야 합니다.

Before (__init__):

def __init__(
    self,
    shared_experts: torch.nn.Module,
    moe_config: FusedMoEConfig,
    enable_dbo: bool,
    mk_can_overlap_shared_experts: Callable[[], bool],
):
    ...

After (__init__):

def __init__(
    self,
    shared_experts: torch.nn.Module,
    moe_config: FusedMoEConfig,
    enable_dbo: bool,
    mk_can_overlap_shared_experts: Callable[[], bool],
    is_multistream_safe: Callable[[], bool], # New argument
):
    ...
    self._is_multistream_safe = is_multistream_safe
    ...

새로운 메서드 maybe_forward_async 및 wait:

이전에는 _run_in_aux_stream과 maybe_sync_shared_experts_stream과 같은 함수들이 혼재되어 있었으나, 이를 maybe_forward_async와 wait로 명확히 분리했습니다.

maybe_forward_async 함수는 공유 전문가 연산을 별도의 CUDA 스트림에 비동기적으로 제출합니다. 이 함수는 실제로 오버랩이 안전하고 유익할 때만 True를 반환하며, 그렇지 않으면 False를 반환합니다. 이 함수는 공유 전문가의 입력 텐서에 record_stream을 호출하여, 해당 텐서가 공유 전문가 스트림에서 사용될 것임을 명시합니다.

wait 함수는 maybe_forward_async가 True를 반환했을 경우, 메인 스트림에서 공유 전문가 연산이 완료될 때까지 기다리는 역할을 합니다. 이는 torch.cuda.Event를 사용하여 구현됩니다.

이러한 분리는 코드의 가독성을 높이고, 비동기 스트림 관리 로직을 더욱 명확하게 만듭니다.

3. 멀티 스트림 안전성 조건 강화

shared_experts.py의 _determine_shared_experts_order 함수에서도 중요한 로직이 추가되었습니다.

Before:

should_run_shared_in_aux_stream = (
    current_platform.is_cuda()
    and self._stream is not None
    and hidden_states.shape[0]
    <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
)

After:

overlap_is_beneficial = not current_platform.is_rocm() or (
    self._moe_config.moe_parallel_config.dp_size > 1
    and self._moe_config.moe_parallel_config.tp_size == 1
)

should_run_shared_in_aux_stream = (
    current_platform.is_cuda_alike()
    and self._stream is not None
    and hidden_states.shape[0]
    <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
    and overlap_is_beneficial
    and self._is_multistream_safe()
)

overlap_is_beneficial이라는 새로운 조건이 추가되었습니다. 이 조건은 다음과 같은 경우에만 멀티 스트림 오버랩이 유익하다고 판단합니다:

ROCm 환경이 아니고, 또는
ROCm 환경이지만 데이터 병렬화(DP) 크기가 1보다 크고 텐서 병렬화(TP) 크기가 1인 경우 (즉, DP-only 구성).

또한, self._is_multistream_safe() 조건이 추가되어, 공유 전문가 입력이 양자화되지 않아 경쟁 상태가 발생할 수 있는 경우에는 오버랩을 비활성화합니다.

리뷰어(shen-shanshan)의 지적대로, TP > 1일 때 DP > 1이더라도 오버랩이 성능에 부정적인 영향을 줄 수 있다는 점을 고려하여, 최종적으로 dp_size > 1 and tp_size == 1 조건으로 더욱 엄격하게 제한되었습니다. 이는 TP 환경에서는 듀얼 스트림 오버랩이 오히려 성능 저하를 유발할 수 있다는 경험적 테스트 결과에 따른 것입니다.

왜 이게 좋은가?

이번 PR의 변경 사항은 다음과 같은 측면에서 vLLM의 성능과 안정성을 향상시킵니다.

ROCm 환경에서의 듀얼 스트림 디코드 최적화:
- 이전에는 ROCm 환경에서 듀얼 스트림 디코드가 제대로 동작하지 않거나, 오히려 성능 저하를 유발하는 경우가 있었습니다. 이번 PR은 공유 전문가(shared experts)와 라우팅된 전문가(routed experts) 간의 스트림 동기화 및 실행 순서를 최적화하여, ROCm 환경에서도 효과적인 듀얼 스트림 오버랩을 가능하게 합니다.
- 특히, maybe_forward_async와 wait 메서드를 통해 공유 전문가를 비동기적으로 실행하고 필요할 때만 기다리는 방식으로 변경하여, 실제 오버랩을 통한 성능 향상을 달성했습니다.
- 테스트 결과, 1k/1k 및 8k/1k 시나리오에서 DP8+TP 구성 시 TPOT(Time Per Output Token)이 최대 11.44%까지 개선되었습니다. 이는 GPU 자원을 더 효율적으로 활용하여 초당 더 많은 토큰을 처리할 수 있음을 의미합니다.
특정 모델(Qwen3.5)의 정확도 문제 해결:
- Qwen3.5와 같이 공유 전문가 입력이 양자화되지 않는 모델의 경우, 공유 전문가와 라우팅된 전문가가 동일한 텐서를 참조하게 되어 멀티 스트림 모드에서 경쟁 상태(race condition)가 발생했습니다. 이로 인해 모델의 정확도가 저하되는 문제가 있었습니다.
- 이번 PR에서는 is_multistream_safe 검사를 통해, 이러한 경쟁 상태가 발생할 수 있는 경우 자동으로 듀얼 스트림 오버랩을 비활성화합니다. 이를 통해 Qwen3.5 모델의 정확도를 유지하면서도, 다른 모델에서는 듀얼 스트림의 이점을 누릴 수 있게 되었습니다.
- vllm serve Qwen/Qwen3.5-35B-A3B ... 명령어로 테스트한 결과, Qwen3.5 모델의 GSM8k 정확도는 이전과 동일하게 유지되었습니다.
ROCm 환경에서의 TP(Tensor Parallelism) 사용 시 주의:
- 테스트 결과, ROCm 환경에서 TP > 1인 경우 듀얼 스트림 오버랩이 오히려 성능 저하를 야기하는 것으로 나타났습니다. 따라서 이 PR에서는 TP > 1인 경우에는 듀얼 스트림 오버랩을 비활성화하도록 조건을 강화했습니다 (dp_size > 1 and tp_size == 1). 이는 불필요한 성능 저하를 방지하고, 각 병렬화 전략에 맞는 최적의 실행 방식을 선택하도록 합니다.
코드 구조 개선:
- 비동기 스트림 관리 로직을 maybe_forward_async와 wait로 분리함으로써 코드의 가독성과 유지보수성이 향상되었습니다.

일반적인 교훈

하드웨어별 최적화의 중요성: ROCm과 CUDA는 하드웨어 아키텍처 및 API 동작 방식에 차이가 있습니다. 따라서 특정 하드웨어 플랫폼에 맞춰 최적화 로직을 조정하는 것이 중요합니다. 이 PR은 ROCm 환경의 특성을 고려하여 듀얼 스트림 디코드를 개선했습니다.
경쟁 상태(Race Condition) 방지: 멀티 스트림 프로그래밍에서는 여러 스트림이 동일한 메모리나 자원에 접근할 때 경쟁 상태가 발생할 수 있습니다. 이러한 문제를 방지하기 위해 입력 텐서의 복사본 사용 여부, 양자화 상태 등을 고려한 안전성 검사가 필수적입니다.
성능 측정 및 분석의 중요성: 새로운 기능을 도입하거나 기존 기능을 수정할 때는 다양한 시나리오(다른 모델, 다른 병렬화 설정, 다른 배치 크기 등)에서 성능을 측정하고 분석해야 합니다. 이 PR은 TPOT, TTFT 등의 지표를 통해 성능 개선 효과를 명확히 입증했습니다.
점진적 개선 및 리뷰: 복잡한 시스템에서는 한 번에 완벽한 해결책을 찾기 어렵습니다. 이 PR은 이전 시도(#48223)의 실패를 바탕으로 문제의 근본 원인을 파악하고, 리뷰어(shen-shanshan, simondanielsson 등)의 피드백을 반영하여 더욱 견고한 해결책을 제시했습니다. 특히, is_multistream_safe 조건과 TP 환경에서의 제약 조건은 리뷰를 통해 개선된 부분입니다.

결론

이번 vLLM PR은 ROCm 환경에서 MoE 모델의 듀얼 스트림 디코드 기능을 정교하게 최적화하여 상당한 성능 향상을 이루었습니다. 또한, 특정 모델에서 발생하던 정확도 문제를 해결하고, 다양한 병렬화 설정에서의 안정성을 확보했습니다. 이러한 개선은 vLLM이 더 넓은 범위의 하드웨어와 모델에서 효율적으로 동작할 수 있도록 하는 데 크게 기여할 것입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] SGLang, NVIDIA Blackwell GPU를 위한 Wan2.2 모델의 ML P 연산 최적화
현재글 : [vllm] vLLM의 ROCm 환경에서 듀얼 스트림 디코드를 통한 성능 최적화
다음글 [sglang] ROCm 환경에서 BF16 All-Reduce의 수치 안정성 확보하기: QuickReduce의 FP16 Saturation 이슈 해결