Model Runner V2 설계 문서

Model Runner V2 설계 문서 (Model Runner V2 Design Document)

소개

vLLM V1이 처음 구현된 이후 여러 근본적인 설계 실수가 발견됐고 상당한 기술 부채가 쌓였습니다. 원래 설계에서 고려되지 않았던 많은 기능이 얹혀졌습니다. 또한 샘플링 기법(예: Gumbel-max 샘플링), 도구(예: Triton), CUDA 기능(예: UVA)에 대한 귀중한 통찰도 얻었습니다. 이런 지식을 바탕으로 Model Runner V2(MRV2)를 원칙에 입각해(first principles) 더 깔끔하고, 효율적이며, 모듈화된 방식으로 구현했습니다.

돌이켜 보면 V1의 많은 설계 선택은 차선이었습니다. MRV2는 아직 기능이 완전하지 않고, 엄격한 테스트를 거치지 않았으며, 열려 있는 설계 결정도 있지만, V1보다 상당히 개선된 것이라고 생각합니다.

이 문서는 MRV2의 설계를 설명합니다.

출처: 문서

본문

1. 영속 배치 (Persistent Batch)

V1의 마찰(moving friction) 요인 중 하나는 영속 배치 구현입니다.

배경 (Background)

V1은 입력 준비 과정에서 CPU 오버헤드를 최소화하기 위해 영속 배치를 도입했습니다. 스텝에 요청이 스케줄링되면 모델 러너는 모델에 넣을 연속 입력 텐서(예: 블록 테이블, 요청별 temperature 값)를 구성해야 합니다. 이런 텐서를 매 스텝 처음부터 만드는 것은 특히 블록 테이블 같은 큰 텐서에서 Python으로 매우 느립니다.

영속 배치 최적화는 연속된 스텝의 요청 배치가 대부분 동일하다는 사실을 활용합니다. 스텝마다 몇 개의 요청만(있다면) 들어오거나 끝납니다. 영속 상태 텐서를 유지하고 입력을 처음부터 재구성하는 대신 증분 diff를 적용하면 CPU 오버헤드를 크게 줄일 수 있습니다.

V1 접근 방식의 문제점 (Problems with V1's Approach)

효율적이긴 하지만 V1의 영속 배치 설계는 영속 상태와 입력 텐서를 결합해 불필요한 복잡성을 만들었습니다. V1은 영속 상태 텐서를 모델·샘플러 입력으로 직접 사용해 엄격한 레이아웃·순서 요구 사항을 부과합니다. 요청이 들어오거나 끝날 때 단순한 행 삽입·제거가 아니라 복잡한 텐서 전체 재배치가 필요한 경우가 많습니다.

또한 V1은 요청이 아직 활성 상태인 동안 영속 텐서의 행이 덮어써질 수 있기 때문에, 요청 상태의 중복 백업 사본인 CachedRequestState를 유지해야 했습니다.

결과적으로 비동기 스케줄링에서 더 어려워지는 복잡한 북키핑(bookkeeping)이 생깁니다.

MRV2의 해결책 (MRV2's Solution)

MRV2는 영속 상태 텐서와 스텝별 입력 텐서를 분리합니다. 스텝의 요청 순서(보통 attention 백엔드가 결정)가 주어지면 MRV2는 영속 상태에서 입력 텐서를 gather합니다.

• max_num_reqs개 행(대부분의 플랫폼에서 기본 1024)으로 고정 크기 텐서를 사전 할당합니다.
• 각 요청에 활성 생명주기(종료 또는 선점까지) 동안 영구 행을 할당합니다.
• 선점을 완료로 취급합니다. 재개 시 요청 데이터를 새 상태로 다시 추가합니다.

이렇게 하면 CachedRequestState가 필요 없어지고 북키핑이 단순해집니다. 큰 상태 텐서는 대부분 GPU 메모리에 저장되므로 gather는 GPU에서 병렬로 낮은 오버헤드로 실행됩니다.

2. 비동기 우선 (Async-First)

vLLM은 이제 비동기 스케줄링에 크게 의존합니다. GPU가 스텝 N을 실행하는 동안 스케줄러와 워커가 스텝 N+1의 입력을 준비해 CPU·GPU 작업을 겹쳐 활용도를 극대화합니다.

V1은 원래 비동기 스케줄링을 염두에 두고 설계되지 않아 지원하려면 개조된 동작과 해킹이 필요했습니다. MRV2는 핵심 모델 실행 루프를 CPU 동기화 지점이 없는 CUDA 스트림으로 가정합니다. CPU 진입점은 작업을 스트림에 큐잉합니다.

3. 비동기 배리어 제거 (Removing Async Barrier)

비동기 실행의 핵심 요구 사항은 CPU 연산이 비블로킹 상태를 유지하는 것입니다. 명시적 동기화(예: torch.accelerator.synchronize)와 암시적 동기화(예: pin 되지 않은 .to("cuda"))를 모두 피해야 합니다.

그러나 비동기 실행은 CPU와 GPU가 같은 메모리를 동시에 건드릴 때 경쟁 조건을 유발할 수 있습니다.

안전하지 않은 예:

class ModelRunner:
    def __init__(self, ...):
        # Pinned buffer
        self.states = torch.zeros(
            max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=True
        )

    def execute_step(self, ...):
        self.states[req_idx] = new_req.data
        states = self.states.to("cuda", non_blocking=True)

CPU가 self.states를 수정하는 동안 GPU는 비동기 복사로 아직 그것을 읽고 있을 수 있습니다.

V1은 임계 구간(critical section) 주변에 async barrier를 두어 해결합니다. 경쟁은 피하지만 단점이 있습니다:

• 보호 버퍼를 놓치기 쉽습니다(버그 유발).
• 유연하지 않은 구성(모든 CPU 작업이 배리어 안에 있어야 함).
• 동기화로 인해 겹침이 줄어들 수 있습니다.

MRV2의 해결책: 경쟁 제거 (Eliminate the Race)

MRV2는 영속 CPU 상태와 복사된 텐서를 분리합니다:

class ModelRunner:
    def __init__(self, ...):
        # Not pinned
        self.states = torch.zeros(
            max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=False
        )

    def execute_step(self, ...):
        self.states[req_idx] = new_req.data
        tmp_states = self.states.pin_memory()
        states = tmp_states.to("cuda", non_blocking=True)

이제 CPU는 self.states에 쓰고 GPU는 tmp_states에서 읽으므로 명시적 동기화 없이 경쟁이 제거됩니다.

4. StagedWriteTensor

블록 테이블 같은 큰 텐서의 경우 MRV2는 매 스텝 전체 CPU→GPU 복사를 피하기 위해 StagedWriteTensor를 사용합니다:

• 기본 텐서는 GPU에 유지합니다.
• diff를 CPU에 스테이징합니다.
• diff를 연속 버퍼로 패킹합니다.
• 패킹된 diff를 GPU로 복사합니다.
• diff를 적용하는 커널 하나를 실행합니다.

사용 예:

# Initialize state on GPU
state = StagedWriteTensor(size=(1024, 1000), dtype=torch.int32, device="cuda")

# Write [3, 1, 2] into row 2, starting at index 3
state.stage_write(row=2, start=3, value=[3, 1, 2])

# Write [-1, -2, -5] into row 0, starting at index 1
state.stage_write(row=0, start=1, value=[-1, -2, -5])

# Apply staged changes
state.apply_write()

이것은 CPU-GPU 동기화 없이 최소한의 커널 실행으로 ragged 업데이트를 지원합니다. 특히 블록 테이블과 num_computed_tokens 같은 CPU/GPU 혼합 기록 상태에 유용합니다.

5. GPU 네이티브 입력 메타데이터 준비 및 출력 처리 (GPU-Native Input Metadata Preparation and Output Processing)

MRV2는 input_ids, positions, query_start_loc, seq_lens 같은 입력을 준비하는 데 Triton 커널을 사용합니다.

이점:

• 더 나은 비동기 동작: GPU가 CPU가 아직 모를 수 있는 값(예: 추측 디코딩)을 파생할 수 있습니다.
• 더 낮은 CPU 오버헤드: 입력 준비가 GPU에서 매우 저렴하고 Python 병목을 피합니다.

UVA (Universal Virtual Addressing)

MRV2는 일부 경로에서 UVA를 사용해 GPU 커널이 prefill_token_ids 같은 대형 CPU 상주 텐서를 GPU 메모리로 중복 복사하지 않고 직접 접근하게 합니다.

6. Triton 네이티브 샘플러 (Triton-Native Sampler)

MRV2는 더 나은 수치·메모리 제어와 최적화를 위해 샘플링을 대부분 Triton으로 재구현합니다.

Gumbel 샘플링 커널 (Gumbel Sampling Kernel)

MRV2는 명시적 softmax 실체화를 피하고 시드 입력에서 무상태(stateless) 인커널 RNG를 사용하는 Triton Gumbel 샘플링 커널을 도입합니다.

효율적인 Top-K Logprobs

V1은 top-k 전에 전체 어휘 logprob를 실체화합니다. MRV2는 먼저 logits에서 top-k 토큰을 식별한 뒤 선택된 토큰에 대해서만 logprob을 계산합니다. 이는 최대 GPU 메모리 사용량을 줄입니다.

메모리 효율적인 프롬프트 Logprobs

MRV2는 단일 프롬프트 내부 청킹을 포함한 더 세밀한 청킹을 지원해 긴 프롬프트에서 메모리 급증을 피합니다.

추측 디코딩과의 더 나은 호환성

MRV2는 요청별 샘플링 상태를 logits 모양에 맞게 확장하는 대신 커널 내부에서 간접 참조(idx_mapping)를 사용해 각 logits 벡터를 올바른 요청 상태에 매핑합니다. 이는 복잡한 샘플링 파라미터와 logits 프로세서 지원을 단순화합니다.

7. 모듈성 (Modularity)

MRV2는 모듈성을 강조합니다. V1의 크고 얽힌 gpu_model_runner.py와 달리 MRV2는 기능 로직을 전용 파일(mrope_utils.py, penalties.py 등)로 나눕니다.

또한 모델 입력을 InputBatch 클래스로 통합하고 모델 러너 속성 간 직접 결합을 줄입니다.

8. dummy_run 남용 없음 (No Abuse of dummy_run)

V1에서 dummy_run은 너무 많은 책임을 졌습니다:

• 초기 메모리 프로파일링 및 torch.compile
• CUDA graph 캡처
• 워밍업
• EP+DP용 빈 DP forward pass

MRV2는 이를 단순화합니다:

• execute_model은 상태에 영향을 주지 않고 dummy 실행을 지원합니다.
• dummy_run은 프로파일링, 워밍업, 빈 DP forward pass를 위해 execute_model에 위임합니다.
• CUDA graph 캡처는 별도의 전용 경로를 사용합니다.

이것은 복잡성을 줄이고 execute_model과 dummy_run 동작의 차이로 인한 버그를 제거합니다.

9. 명시적 CUDA Graph 관리 (Explicit CUDA Graph Management)

V1의 CUDA graph 처리는 암시적이고 추론하기 어렵습니다. MRV2는 표준 PyTorch API를 통해 전체 CUDA graph를 명시적으로 캡처하고 실행하는 CUDAGraphManager를 사용합니다.

이로써 graph 수명 주기와 실행 모드 결정이 이해하기 쉽고 확장하기 쉬워집니다. 예: MRV2는 여러 draft 모델 forward pass를 하나의 CUDA graph로 캡처할 수 있습니다.

융합 멀티 스텝 Draft 디코딩 (Fused Multi-Step Draft Decoding)

자기회귀 추측 디코딩은 스케줄러 스텝당 여러 개의 종속 draft 스텝을 실행합니다. 융합 경로에서 MRV2는 draft 토큰마다 별도 graph를 재생(replay)하는 대신, post-prefill draft 스텝 전체를 하나의 전체 CUDA graph로 캡처합니다. Attention 메타데이터는 루프 전에 한 번 구축되고, 공통 스텝 의존 텐서는 안정적인 주소를 유지하며 draft 스텝 사이에 in-place로 갱신됩니다.

일부 attention 백엔드는 스케줄러 메타데이터나 sparse 인덱스 같은 파생 상태도 실체화합니다. 융합 경로를 선택하기 전에 이런 백엔드는 AttentionMetadataBuilder.update_draft_decode_metadata()를 구현해 draft 입력이 진행된 후 그 상태를 갱신하거나 무효화해야 합니다. 이 훅은 CUDA graph 캡처 중 실행되므로, 기록된 GPU 연산만 재생 중 실행되고 Python 본문은 다시 실행되지 않습니다. 따라서 구현은 캡처 안전(capture-safe) 연산을 사용하고 재생되는 모든 텐서 상태를 영속 스토리지에 유지해야 합니다.

위치를 진행하는 draft 모델의 경우, 융합 경로는 모든 draft attention 그룹이 supports_draft_decode_metadata_update를 선언할 때만 활성화됩니다. 그렇지 않으면 MRV2는 draft 스텝 사이에 attention 메타데이터를 다시 구축하는 방식으로 폴백합니다. 위치를 고정하는 draft 모델은 이 갱신이 필요 없습니다. 백엔드를 활성화하기 전에 부모 빌더에서 상속받은 상태나 보조 attention 백엔드가 소유한 상태를 포함한 모든 파생 메타데이터를 감사해야 합니다.

개발 철학 (Development Philosophy)

MRV2 변경은 더 높은 코드 품질 기준을 충족해야 합니다. V1과의 기능 격차가 채워지면서, 기능은 V1 동작을 빠르게 포팅하는 대신 MRV2 설계 맥락에서 원칙에 입각해 재고되어야 합니다.

핵심 요구 사항은 더 많은 설계 반복이 필요하더라도 모듈성과 깨끗한 추상화 경계를 보존하는 것입니다.

더 알아보기 (Learn more)

• 모델 실행 아키텍처 — vLLM 아키텍처 개요
• Attention 백엔드 — attention 백엔드 기능 지원
• 추측 디코딩 — 추측 디코딩 기능