[vllm] vLLM, DFlash 도입으로 추론 속도 1.2배 향상: MRV2와 CUDAGraph의 시너지

2026년 6월 10일수정: 2026년 6월 10일

PR 링크: vllm-project/vllm#44586 상태: Merged | 변경: +868 / -47

들어가며

최근 대규모 언어 모델(LLM)의 추론 속도 최적화는 실시간 서비스 제공에 있어 핵심적인 과제가 되었습니다. 특히, vLLM은 KV 캐시 최적화와 같은 혁신적인 기술로 주목받아 왔습니다. 이번 PR은 vLLM에 DFlash (Block Diffusion Speculative Decoding)를 도입하고, 이를 Model Runner V2 (MRV2) 및 CUDA Graph (CUDAGraph)와 통합하여 추론 성능을 크게 향상시키는 것을 목표로 합니다.

본 글에서는 이 PR이 해결하고자 하는 문제점, 코드 변경의 핵심 내용, 그리고 이러한 변경이 왜 성능 향상으로 이어지는지에 대해 실제 코드 변경 사항을 중심으로 심층적으로 분석합니다. 특히, DFlash와 MRV2의 결합이 가져오는 시너지 효과와 CUDAGraph 지원을 통한 최적화 방안을 중점적으로 다룰 것입니다.

코드 분석

이번 PR은 DFlash 알고리즘을 vLLM의 MRV2 아키텍처에 통합하고, CUDAGraph와의 호환성을 확보하는 데 중점을 둡니다. 주요 변경 사항은 다음과 같습니다.

1. DFlash Speculator 도입 및 MRV2 통합 (`vllm/v1/worker/gpu/spec_decode/init.py`, `vllm/v1/worker/gpu/spec_decode/dflash/speculator.py` 등)

DFlash는 기존의 자기회귀적 추론 방식과 달리, 여러 개의 토큰을 한 번에 예측하는 '블록' 단위의 추론을 통해 속도를 높이는 기술입니다. 이 PR에서는 DFlash를 위한 새로운 DFlashSpeculator 클래스를 도입하고, 이를 vllm/v1/worker/gpu/spec_decode/__init__.py에서 초기화 로직에 포함시켰습니다.

Before:

def init_speculator(vllm_config: VllmConfig, device: torch.device):
    speculative_config = vllm_config.speculative_config
    assert speculative_config is not None
    if speculative_config.use_gemma4_mtp():
        from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import \
            Gemma4Speculator

        return Gemma4Speculator(vllm_config, device)
    # ... other speculators

After:

def init_speculator(vllm_config: VllmConfig, device: torch.device):
    speculative_config = vllm_config.speculative_config
    assert speculative_config is not None
    if speculative_config.method == "dflash":
        from vllm.v1.worker.gpu.spec_decode.dflash.speculator import \
            DFlashSpeculator

        return DFlashSpeculator(vllm_config, device)
    elif speculative_config.use_gemma4_mtp():
        from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import \
            Gemma4Speculator

        return Gemma4Speculator(vllm_config, device)
    # ... other speculators

speculative_config.method에 "dflash"가 추가되어, 해당 설정 시 DFlashSpeculator가 인스턴스화됩니다. 이는 DFlash를 vLLM의 다양한 추론 전략 중 하나로 선택할 수 있게 합니다.

또한, vllm/v1/worker/gpu/model_runner.py에서는 DFlash와 같은 특정 투기적 디코딩 방법이 파이프라인 병렬 처리(Pipeline Parallelism)와 호환되지 않음을 명시했습니다.

Before:

            if self.speculative_config.method == "eagle3":
                # EAGLE3 may require auxiliary hidden states from target model outputs.
                self.use_aux_hidden_state_outputs = True
                if self.use_pp:
                    raise ValueError("EAGLE3 with pipeline parallel is not supported.")

After:

            if self.speculative_config.method in ("eagle3", "dflash"):
                # Drafting may require auxiliary hidden states from target model outputs
                self.use_aux_hidden_state_outputs = True
                if self.use_pp:
                    raise ValueError(
                        f"{self.speculative_config.method} with pipeline parallel "
                        "is not supported."
                    )

eagle3뿐만 아니라 dflash도 파이프라인 병렬 처리와 함께 사용할 수 없도록 예외 처리가 추가되었습니다. 이는 DFlash의 내부 동작 방식이 파이프라인 병렬 처리의 단계별 실행과 충돌할 수 있기 때문입니다.

2. CUDAGraph 지원 강화 (`vllm/v1/worker/gpu/attn_utils.py`, `vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py` 등)

CUDAGraph는 CUDA 커널 실행을 캡처하여 재사용함으로써 오버헤드를 줄이고 성능을 향상시키는 기술입니다. DFlash를 CUDAGraph와 함께 사용하기 위해 build_attn_metadata 함수에 causal 파라미터가 추가되었습니다.

Before:

def build_attn_metadata(
    # ... other params
    for_cudagraph_capture: bool = False,
) -> dict[str, Any]:
    # ...
            causal=True,
    # ...

After:

def build_attn_metadata(
    # ... other params
    for_cudagraph_capture: bool = False,
    causal: bool = True,
) -> dict[str, Any]:
    # ...
            causal=causal,
    # ...

causal 파라미터가 기본값 True에서 동적으로 설정 가능하게 변경되었습니다. 이는 DFlash가 비인과적(non-causal) 어텐션 메커니즘을 사용할 때도 CUDAGraph 캡처가 가능하도록 유연성을 제공합니다. dflash/cudagraph.py 파일은 DFlash를 위한 CUDAGraph 관련 유틸리티 함수들을 정의하며, _prepare_dflash_inputs_to_capture 함수는 CUDAGraph 캡처에 필요한 입력들을 준비하는 역할을 합니다.

3. Gumbel Sample Kernel 최적화 (`vllm/v1/worker/gpu/sample/gumbel.py`)

투기적 디코딩에서 토큰 샘플링은 중요한 부분입니다. Gumbel 샘플링 커널에 PER_TOKEN_COL이라는 새로운 tl.constexpr 파라미터가 추가되었습니다.

Before:

def _gumbel_sample_kernel(
    # ... other params
):
    # ...
    if processed_logits_ptr is not None:
        # Store the temperature-applied logits.
        if processed_logits_col_ptr is not None:
            col = tl.load(processed_logits_col_ptr)
        else:
            col = 0
        # ...

After:

def _gumbel_sample_kernel(
    # ... other params
    PER_TOKEN_COL: tl.constexpr,
):
    # ...
    if processed_logits_ptr is not None:
        # Store the temperature-applied logits.
        if processed_logits_col_ptr is not None:
            if PER_TOKEN_COL:
                col = tl.load(processed_logits_col_ptr + token_idx)
            else:
                col = tl.load(processed_logits_col_ptr)
        else:
            col = 0
        # ...

PER_TOKEN_COL이 True일 경우, processed_logits_col_ptr에서 각 토큰별로 다른 값을 로드하게 됩니다. 이는 샘플링 과정에서 각 토큰에 대해 더 세밀한 제어가 가능함을 시사하며, DFlash와 같이 복잡한 샘플링 전략에서 유용할 수 있습니다.

gumbel_sample 함수 자체에서도 output_processed_logits_col의 차원을 확인하여 per_token_col 값을 결정합니다. 이는 output_processed_logits_col이 스칼라가 아닌 텐서일 경우, 즉 각 토큰별로 다른 값을 사용할 때 해당 기능을 활성화하기 위함입니다.

4. 테스트 코드 업데이트 (`tests/v1/e2e/spec_decode/test_spec_decode.py`)

MRV1과 MRV2 모두에서 DFlash의 정확성과 수용률(acceptance rate)을 검증하기 위해 E2E 테스트가 업데이트되었습니다. @pytest.mark.parametrize("use_mrv2", [False, True]) 데코레이터를 사용하여 MRV1과 MRV2 환경 모두에서 테스트가 실행되도록 변경되었습니다.

Before:

def test_dflash_acceptance_rates(dflash_config):
    # ...

After:

@pytest.mark.parametrize("use_mrv2", [False, True])
def test_dflash_acceptance_rates(
    monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config
):
    """
    E2E test for DFlash (block diffusion) speculative decoding.
    Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval
    comparing against baseline results from the paper (Table 1).
    See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology.
    """
    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
    # ...

monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", ...)를 통해 테스트 실행 시 MRV2 사용 여부를 동적으로 설정합니다. 이는 DFlash가 MRV1과 MRV2 모두에서 일관되게 작동하는지 확인하는 데 필수적입니다.

왜 이게 좋은가?

1. 성능 향상

PR 설명에 따르면, Qwen3 8B FP8 모델과 DFlash를 1x GB200에서 실행했을 때 약 1.2배의 속도 향상을 달성했습니다. 기존 5.2ms/step에서 4.3ms/step으로 감소한 것입니다. 이는 DFlash의 블록 단위 예측이 기존의 자기회귀적 토큰 생성보다 훨씬 효율적임을 보여줍니다.

MRV1과 MRV2의 성능 비교 이미지에서도 DFlash 적용 시 MRV2가 MRV1보다 더 나은 성능을 보이는 것을 확인할 수 있습니다. MRV2는 vLLM의 차세대 모델 러너로, 더 최적화된 내부 구조와 연산 처리를 제공하여 DFlash와 같은 고급 추론 기법의 성능을 극대화합니다.

2. CUDAGraph 통합을 통한 오버헤드 감소

CUDAGraph 지원 강화는 반복적인 커널 호출에 드는 CPU 오버헤드를 크게 줄여줍니다. 특히 LLM 추론과 같이 동일한 연산이 반복되는 경우, CUDAGraph는 GPU 활용률을 높이고 전체적인 지연 시간을 단축하는 데 결정적인 역할을 합니다. DFlash의 복잡한 연산 흐름을 CUDAGraph로 캡처함으로써, 각 스텝마다 발생하는 오버헤드를 최소화할 수 있습니다.

3. 유연성 및 확장성 확보

MRV2 지원: DFlash를 최신 MRV2 아키텍처에 통합함으로써, vLLM의 최신 성능 개선 사항을 활용할 수 있게 되었습니다.
다양한 추론 전략 지원: speculative_config.method에 "dflash"를 추가하고, 테스트 코드에서 MRV1/MRV2를 모두 지원함으로써, 사용자는 다양한 추론 전략을 쉽게 비교하고 선택할 수 있습니다.
Gumbel Sample Kernel 최적화: PER_TOKEN_COL과 같은 세부적인 최적화는 복잡한 샘플링 로직을 더 효율적으로 처리할 수 있게 하여, 잠재적으로 더 높은 수용률과 정확도를 달성하는 데 기여할 수 있습니다.

4. 일반적인 교훈

알고리즘과 아키텍처의 결합: DFlash와 같은 새로운 추론 알고리즘은 기존의 고성능 아키텍처(MRV2) 및 최적화 기술(CUDAGraph)과 결합될 때 가장 큰 시너지를 발휘합니다.
테스트의 중요성: 새로운 기능을 도입할 때는 기존 기능(MRV1)과 새로운 기능(MRV2) 모두에서 철저한 E2E 테스트를 통해 정확성과 성능을 검증해야 합니다.
세밀한 최적화: Gumbel 샘플링 커널과 같이 저수준의 연산 최적화는 전체 시스템 성능에 상당한 영향을 미칠 수 있습니다.

리뷰 피드백 반영

주요 리뷰 댓글들을 살펴보면, 코드 중복 제거, CUDAGraph 관련 TODO 처리, 그리고 파라미터 사용의 정확성에 대한 논의가 있었습니다.

benchislett의 TODO 코멘트는 dflash/cudagraph.py 파일이 최신 상태가 아닐 수 있음을 지적하며, 실제 코드에서는 해당 TODO가 제거되었거나 수정되었을 가능성이 있습니다. 또한, dflash/speculator.py가 autoregressive/speculator.py와 코드 중복이 많다는 지적은, 코드 재사용성을 높이기 위한 리팩토링의 필요성을 시사합니다. 실제로 TheEpicDolphin은 DraftModelSpeculator 클래스에 헬퍼 메소드를 추가하여 재사용성을 높일 것을 제안했습니다.
TheEpicDolphin은 _build_draft_attn_metadata 함수를 재사용하고 num_query_per_req 및 causal 파라미터를 추가할 것을 제안했습니다. 이는 코드 중복을 줄이고 일관성을 유지하는 좋은 방법입니다.
benchislett의 max_tokens_per_req 사용에 대한 지적은, DFlash의 max_query_len 계산이 전체 토큰 수가 아닌 타겟 배치(target batch)의 최대 쿼리 길이(maximum query length)를 사용해야 함을 명확히 했습니다. PR에서는 이 부분을 수정하여 정확성을 높였습니다.

이러한 리뷰 피드백은 코드의 품질, 재사용성, 그리고 정확성을 향상시키는 데 중요한 역할을 했습니다.

결론

이번 PR은 vLLM에 DFlash를 성공적으로 통합하고, MRV2 및 CUDAGraph와의 시너지를 통해 상당한 추론 성능 향상을 달성했습니다. DFlash는 블록 단위 예측을 통해 LLM 추론 속도를 크게 개선할 수 있는 강력한 기법이며, vLLM의 최신 아키텍처 및 최적화 기술과 결합될 때 그 효과가 극대화됩니다. 앞으로도 vLLM은 이러한 최신 기술들을 적극적으로 도입하여 LLM 추론의 한계를 넓혀나갈 것으로 기대됩니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [loki] Grafana Loki의 Shuffle Sharding 알고리즘 최적화: 성능 향상의 비결
현재글 : [vllm] vLLM, DFlash 도입으로 추론 속도 1.2배 향상: MRV2와 CUDAGraph의 시너지
다음글 [flashinfer] FlashInfer Unified MoE API: NVFP4 백엔드 통합 및 자동 튜닝 최적화

댓글

관련 포스트

PR Analysis 의 다른글