처음부터 미니 vLLM 만들기: LLM 추론 최적화 심층 탐구

vLLM 같은 LLM 추론 엔진을 돌릴 때 내부에서 무슨 일이 일어나는지 궁금했던 적 있나요? 저는 정말 궁금했습니다. 그래서 nano-vllm을 만들었습니다. 고성능 LLM 추론 엔진을 처음부터 구현한, 미니멀하고 교육용인 프로젝트입니다. 한마디로 “초보자를 위한 vLLM”(저 자신 포함)이라고 생각하면 됩니다.

이 글은 제가 배운 것들과 각 최적화 기법이 어떻게 작동하는지 정리한 것입니다. 준비하세요 — 깊게 들어갑니다!

문제: 왜 LLM 추론은 까다로운가

LLM을 돌릴 때 핵심은 단순히 행렬곱만 하는 게 아닙니다. 요청을 하나씩 순차적으로 처리하는 나이브한 방식은 GPU 메모리와 연산을 민망할 정도로 낭비합니다. 왜 그런지 설명해 보겠습니다.

LLM이 텍스트를 생성할 때는 두 단계로 동작합니다:

1. 프리필(prefill): 프롬프트 전체를 한 번에 처리 (연산(Compute) 병목)
2. 디코드(decode): 토큰을 한 개씩 생성 (메모리 병목)

흥미로운 건 디코드 단계입니다. 새 토큰을 만들 때마다 모델은 어텐션 메커니즘을 통해 이전의 모든 토큰을 다시 참조해야 합니다. 캐시가 없다면 같은 계산을 계속 반복하게 됩니다. 여기서 등장하는 것이 KV 캐시(KV cache)입니다.

하지만 문제가 하나 더 있습니다. 최대 시퀀스 길이를 기준으로 KV 캐시 메모리를 미리 할당해 두면 엄청난 낭비가 발생합니다. 예를 들어 최대 길이가 2048인데 실제 시퀀스가 100 토큰이라면, 메모리의 95%를 버리는 셈이죠!

vLLM은 PagedAttention으로 이 문제를 해결했고, nano-vllm에서도 이를 구현했습니다.

아키텍처 개요

nano-vllm은 다음과 같이 구성되어 있습니다:

nano_vllm/
├── engine.py              # 메인 추론 엔진
├── config.py              # 모델 설정
├── cache.py               # KV 캐시 구현
├── sampler.py             # 토큰 샘플링
├── core/
│   ├── sequence.py        # 요청 추적
│   ├── scheduler.py       # 우선순위 기반 배치 스케줄링
│   ├── block.py           # PagedAttention용 메모리 블록
│   └── block_manager.py   # 블록 할당 (OS 메모리 관리자처럼)
├── attention/
│   ├── paged_attention.py # PagedAttention 커널
│   └── flash_attention.py # FlashAttention 통합
├── speculative/
│   └── speculative_decoding.py  # 드래프트 모델 기반 추측
├── educational/           # 관찰하며 배우는 모드
│   ├── narrator.py        # 쉬운 설명
│   ├── xray.py            # 텐서 시각화
│   └── dashboard.py       # 실시간 터미널 UI
└── model/
    ├── loader.py          # HuggingFace 모델 로딩
    └── llama.py           # Llama 구현 (RMSNorm, RoPE, GQA, SwiGLU)

이제 주요 최적화들을 하나씩 살펴보겠습니다!

---

1. PagedAttention: vLLM의 심장

문제는 무엇인가?

전통적인 KV 캐시 할당 방식은 “혹시 1999명의 친구를 데려올 수도 있으니 영화관 전체를 한 사람에게 예약해 두는 것”과 같습니다. 낭비죠.

나이브한 방식에서는 각 시퀀스에 대해 최대 길이를 기준으로 연속(contiguous)된 메모리 덩어리를 미리 할당합니다. 그러면:

• 메모리 단편화: 시퀀스들이 서로 다른 시점에 끝나면서 구멍이 생김
• 메모리 낭비: 대부분의 시퀀스는 최대 길이까지 안 감
• 제한된 배치 크기: GPU 메모리에 더 많은 요청을 올릴 수 없음

PagedAttention은 어떻게 해결하나

PagedAttention은 운영체제의 가상 메모리에서 아이디어를 가져옵니다. 연속 할당 대신 KV 캐시를 고정 크기의 블록(메모리 페이지처럼)으로 나눕니다:

python
# core/block.py
@dataclass
class Block:
    """KV 캐시 메모리의 고정 크기 청크.

    각 블록은 `block_size` 토큰에 대한 KV 상태를 저장한다.
    """
    block_id: int
    block_size: int = 16  # 블록당 16 토큰
    ref_count: int = 1    # 공유(프리픽스 캐싱)를 위한 참조 카운트
    prefix_hash: Optional[int] = None

각 시퀀스는 BlockTable을 가지며, 논리적 위치를 물리 블록으로 매핑합니다:

python
# core/block.py
@dataclass
class BlockTable:
    """논리적 위치를 물리 블록으로 매핑.

    가상 메모리의 페이지 테이블과 유사:
    - 위치 p의 토큰은 논리 블록: p // block_size
    - 블록 내 슬롯: p % block_size
    - 물리 블록: block_ids[p // block_size]

    예시 (block_size=16, 시퀀스가 35 토큰):
        block_table.block_ids = [5, 12, 3]  # 물리 블록 3개
        Token 0-15  -> block 5
        Token 16-31 -> block 12
        Token 32-34 -> block 3 (slots 0-2)
    """
    block_ids: List[int]
    block_size: int = 16

BlockManager는 OS처럼 할당/해제를 관리합니다:

python
# core/block_manager.py
class BlockManager:
    """KV 캐시 블록 할당을 관리.

    간단한 프리 리스트(스택)를 사용해 O(1) 할당/해제.
    """
    def allocate_block(self) -> int:
        if not self.free_blocks:
            raise RuntimeError("KV 캐시 블록이 부족합니다!")
        return self.free_blocks.pop()

    def free_block(self, block_id: int) -> None:
        block = self.blocks[block_id]
        if block.decrement_ref() <= 0:
            self.free_blocks.append(block_id)

Paged Attention 계산

어텐션 계산 시, 비연속(non-contiguous) 블록들에서 K와 V를 모아서(gather) 사용합니다:

python
# attention/paged_attention.py
def paged_attention(
    query: torch.Tensor,
    key_cache: torch.Tensor,      # [num_blocks, block_size, num_kv_heads, head_dim]
    value_cache: torch.Tensor,
    block_tables: List[BlockTable],
    context_lens: List[int],
    block_size: int,
    num_kv_heads: int,
) -> torch.Tensor:
    # 각 시퀀스에 대해 블록에서 gather
    for batch_idx in range(batch_size):
        block_table = block_tables[batch_idx]

        for pos in range(context_len):
            logical_block = pos // block_size
            slot_in_block = pos % block_size
            physical_block = block_table.block_ids[logical_block]

            # 캐시에서 복사
            gathered_keys[batch_idx, :, pos, :] = key_cache[physical_block, slot_in_block]
            gathered_values[batch_idx, :, pos, :] = value_cache[physical_block, slot_in_block]

    # 표준 어텐션 계산
    attn_weights = torch.matmul(query, gathered_keys.transpose(-2, -1)) * scale
    # ... 마스킹, 소프트맥스 적용 후 출력 계산

왜 중요한가

PagedAttention은 다음을 가능하게 합니다:

• 거의 0에 가까운 메모리 낭비: 필요한 만큼만 할당
• 메모리 공유: 공통 프리픽스가 블록을 공유 가능(프리픽스 캐싱)
• 더 높은 처리량: 더 많은 요청을 메모리에 올릴 수 있어 병렬성이 증가

---

2. 연속 배칭(Continuous Batching): 더는 기다리지 않기

전통적 배칭의 문제

구식 배칭은 한 배치의 모든 시퀀스가 끝날 때까지 새 요청을 받지 않습니다. 예를 들어:

• 요청 A: 50 토큰 생성
• 요청 B: 5 토큰 생성

B는 빨리 끝나도 A가 끝날 때까지 기다려야 합니다. GPU가 놀게 되죠.

연속 배칭의 해결

nano-vllm은 이터레이션 단위로 스케줄링합니다:

• 새 요청이 배치 중간에 합류
• 끝난 요청은 즉시 퇴장
• GPU를 계속 바쁘게 유지

스케줄러 동작은 다음과 같습니다:

python
# core/scheduler.py
class Scheduler:
    """시퀀스를 생명주기 전반에서 관리:
    - WAITING: 큐 대기
    - RUNNING: 처리 중
    - SWAPPED: 선점됨
    - FINISHED: 완료
    """

    def schedule(self) -> SchedulerOutputs:
        outputs = SchedulerOutputs()

        # 1. 높은 우선순위 요청 대기 시 선점 처리
        if self.enable_preemption and self.block_manager:
            self._handle_preemption(outputs)

        # 2. 실행 중 시퀀스 계속 처리(디코드)
        for seq in self.running:
            if seq.is_chunked_prefill():
                outputs.chunked_prefill_sequences.append(seq)
            else:
                outputs.decode_sequences.append(seq)

        # 3. 대기 큐에서 새 시퀀스 입장
        while can_admit_more():
            seq = self._pop_waiting()
            seq.status = SequenceStatus.RUNNING
            outputs.prefill_sequences.append(seq)

        return outputs

엔진은 한 이터레이션에서 다음을 처리합니다:

python
# engine.py
def step(self) -> List[GenerationOutput]:
    """연속 배칭의 한 이터레이션."""
    scheduler_outputs = self.scheduler.schedule()

    # 청크 프리필 처리
    for seq, num_tokens in zip(chunked_prefill_seqs, chunked_prefill_tokens):
        self._run_chunked_prefill(seq, num_tokens)

    # 전체 프리필 처리(새 시퀀스)
    for seq in prefill_sequences:
        self._run_prefill(seq)

    # 디코드 처리(함께 배칭!)
    if decode_sequences:
        self._run_decode(decode_sequences)

    # 완료된 시퀀스 반환
    return newly_finished

---

3. 우선순위 스케줄링 & 선점(Preemption)

어떤 요청에는 VIP 대우가 필요할 수 있습니다. nano-vllm은 이를 지원합니다.

우선순위 기반 스케줄링

요청은 우선순위를 가지며, 우선순위가 높을수록 먼저 처리됩니다:

python
# core/scheduler.py
def _get_priority_key(self, seq: Sequence) -> Tuple[int, float, int]:
    """힙 정렬을 위한 우선순위 키. 튜플이 작을수록 우선순위가 높다."""
    # 큰 priority가 먼저 오도록 음수로 변환
    return (-seq.priority, seq.arrival_time, seq.seq_id)

# O(log n) 스케줄링을 위한 힙 사용
heapq.heappush(self._waiting_heap, (priority_key, sequence))

선점: 낮은 우선순위 요청 내보내기

높은 우선순위 요청이 들어왔는데 메모리가 부족하면, 실행 중인 낮은 우선순위 요청을 선점할 수 있습니다:

python
# core/scheduler.py
def _handle_preemption(self, outputs):
    """높은 우선순위 대기를 위해 낮은 우선순위 실행 시퀀스를 선점."""
    highest_waiting = self._peek_waiting()

    while not self.block_manager.can_allocate(blocks_needed) and self.running:
        # 실행 중인 시퀀스 중 가장 낮은 우선순위 선택
        lowest_running = min(self.running, key=lambda s: s.priority)

        if highest_waiting.priority > lowest_running.priority:
            # 선점! 블록을 해제하고 재계산을 위해 리셋
            self.running.remove(lowest_running)
            self.block_manager.free_sequence_blocks(lowest_running.block_table)
            lowest_running.reset_for_recompute()
            self._push_waiting(lowest_running)

선점된 시퀀스는 다시 대기 상태로 돌아가며 나중에 다시 프리필됩니다. CPU로 스왑하는 대신 재계산(recompute) 기반 선점을 택했는데, 구현이 단순하고 실전에서도 꽤 잘 동작합니다.

---

4. 프리픽스 캐싱(Prefix Caching): 공통 프리픽스 공유

많은 요청은 같은 시스템 프롬프트로 시작합니다. 같은 KV 캐시를 왜 매번 다시 계산할까요?

동작 방식

블록은 토큰 내용과 시퀀스 내 위치를 기반으로 해시됩니다:

python
# core/block.py
def hash_token_block(token_ids: Tuple[int, ...], parent_hash: Optional[int] = None) -> int:
    """전체 프리픽스 체인을 포함하는 누적 해시.

    이를 통해 '전체 프리픽스가 완전히 같을 때만' 블록을 공유하게 된다.
    """
    if parent_hash is None:
        return hash(token_ids)
    return hash((parent_hash, token_ids))

새 시퀀스가 들어오면 프리픽스 블록이 이미 존재하는지 확인합니다:

python
# core/block_manager.py
def allocate_blocks_with_prefix_caching(self, token_ids: List[int]):
    """가능하면 캐시된 프리픽스 블록을 재사용하면서 블록을 할당."""
    parent_hash = None

    for block_idx in range(num_full_blocks):
        block_tokens = tuple(token_ids[start:end])
        cache_key = (parent_hash, block_tokens)

        if cache_key in self.prefix_cache:
            # 캐시 히트! 기존 블록 재사용
            cached_block_id = self.prefix_cache[cache_key]
            self.blocks[cached_block_id].increment_ref()  # 참조 카운팅
            block_table.append_block(cached_block_id)
        else:
            # 캐시 미스 - 새 블록 할당
            block_id = self.allocate_block()
            self.prefix_cache[cache_key] = block_id
            block_table.append_block(block_id)

        parent_hash = self.blocks[block_id].prefix_hash

    return block_table, shared_prefix_len

참조 카운팅(reference counting) 덕분에 다른 시퀀스가 사용 중인 블록은 해제되지 않습니다.

---

5. 청크 프리필(Chunked Prefill): 긴 프롬프트로 배치를 막지 말자

긴 프롬프트(예: 4000 토큰)는 프리필 단계에서 배치 전체를 오래 잡아먹을 수 있습니다. 청크 프리필은 이를 더 작은 조각으로 나눕니다:

python
# engine.py
def _run_chunked_prefill_paged(self, seq: Sequence, num_tokens: int):
    """프롬프트 토큰의 일부 청크를 처리."""
    start_pos = seq.num_prefilled_tokens
    end_pos = start_pos + num_tokens
    chunk_tokens = seq.prompt_token_ids[start_pos:end_pos]

    # 이 청크에 대한 블록 할당
    # ...

    # 이 청크만 forward
    logits = self.model(input_ids, block_kv_cache=..., start_positions=[start_pos])

    # 진행 업데이트
    seq.num_prefilled_tokens = end_pos

    # 전체 프롬프트 토큰 처리가 끝난 뒤에만 샘플링
    if seq.num_prefilled_tokens >= len(seq.prompt_token_ids):
        next_token = self.sampler.sample(logits)
        seq.append_token(next_token.item())

스케줄러는 이터레이션당 프리필 토큰 수를 조절합니다:

python
# max_prefill_tokens가 이터레이션당 연산량을 제한
if prompt_len <= prefill_budget:
    outputs.prefill_sequences.append(seq)  # 전체 프리필
else:
    outputs.chunked_prefill_sequences.append(seq)  # 부분 프리필
    outputs.chunked_prefill_tokens.append(prefill_budget)

---

6. FlashAttention: 메모리 효율적인 어텐션

표준 어텐션은 전체 N×N 어텐션 행렬을 물리화(materialize)합니다. 2048 토큰이면 400만 원소입니다! FlashAttention은 **타일링(tiling)**을 사용해 이를 피합니다.

nano-vllm에서의 통합

python
# attention/flash_attention.py
def flash_attention(query, key, value, causal=True):
    """O(N^2) 대신 O(N) 메모리를 사용하도록 FlashAttention 사용."""
    # FlashAttention 입력: [batch, seq_len, num_heads, head_dim]
    query = query.transpose(1, 2)
    key = key.transpose(1, 2)
    value = value.transpose(1, 2)

    output = flash_attn_func(query, key, value, causal=causal)

    return output.transpose(1, 2)

# 폴백 포함 통합 인터페이스
def attention(query, key, value, use_flash_attn=True, causal=True):
    if use_flash_attn and FLASH_ATTN_AVAILABLE:
        return flash_attention(query, key, value, causal)

    # PyTorch SDPA로 폴백(이것도 최적화되어 있음!)
    return F.scaled_dot_product_attention(query, key, value, is_causal=causal)

FlashAttention은 모델의 어텐션 레이어에서 사용됩니다:

python
# model/llama.py
class LlamaAttention(nn.Module):
    def __init__(self, config, layer_idx, use_flash_attn=True):
        self.use_flash_attn = use_flash_attn and is_flash_attn_available()

    def forward(self, hidden_states, ...):
        # ... Q, K, V 계산 및 RoPE 적용 ...

        # 통합 어텐션 사용(가능하면 FlashAttention)
        attn_output = unified_attention(
            query=query_states,
            key=key_states,
            value=value_states,
            use_flash_attn=self.use_flash_attn,
            causal=True,
        )

---

7. 추측 디코딩(Speculative Decoding): 초안 작성 후 검증

디코딩은 토큰을 한 번에 하나씩 생성하므로 느립니다. 큰 모델을 한 번 forward할 때 여러 토큰을 얻을 수 있다면 어떨까요?

아이디어

1. 작고 빠른 드래프트(draft) 모델로 K개 후보 토큰 생성
2. 큰 타깃(target) 모델이 K+1 위치를 한 번의 forward로 검증
3. 일치하면 수용, 아니면 거부하고 재샘플

python
# speculative/speculative_decoding.py
def _speculative_step(self, current_ids, target_kv_cache, draft_kv_cache, remaining_tokens):
    """추측 디코딩의 한 스텝."""
    K = self.config.num_speculative_tokens

    # 1) 드래프트 토큰 K개 생성(저렴!)
    draft_tokens, draft_probs = self._generate_draft_tokens(current_ids, draft_kv_cache, K)

    # 2) 타깃 모델로 검증(K+1 토큰을 한 번에!)
    verify_ids = [[current_ids[-1]] + draft_tokens]
    target_logits = self.target_model(verify_ids, kv_cache=target_kv_cache)
    target_probs = F.softmax(target_logits, dim=-1)

    # 3) 거부 샘플링으로 수용/거부
    accepted_tokens = []
    for i, draft_token in enumerate(draft_tokens):
        target_prob = target_probs[0, i, draft_token].item()
        draft_prob = draft_probs[i]

        # target prob >= draft prob이면 수용(타깃 분포 유지!)
        acceptance_prob = min(1.0, target_prob / draft_prob)

        if random() < acceptance_prob:
            accepted_tokens.append(draft_token)
        else:
            # 조정된 분포에서 재샘플
            resampled = sample_from_adjusted(target_probs[0, i], draft_prob, draft_token)
            accepted_tokens.append(resampled)
            break  # 첫 거부 이후 중단

    # 전부 수용됐다면 보너스 토큰 1개 더 샘플!
    if len(accepted_tokens) == len(draft_tokens):
        bonus_token = sample(target_probs[0, -1])
        accepted_tokens.append(bonus_token)

    return accepted_tokens

핵심: 품질 손실이 없다

이 방식은 거부 샘플링(rejection sampling)이며, 출력 분포가 타깃 모델과 수학적으로 동일함을 보장합니다. 근사가 아닙니다!

가속 효과는 다음에 달려 있습니다:

• 드래프트 모델 속도(타깃보다 ~10배 빠르면 좋음)
• 수용률(높을수록 타깃 forward당 더 많은 토큰)
• K 값(K가 클수록 잠재 이득 증가)

---

8. Llama 모델 구현

nano-vllm은 최신 구성요소를 포함한 Llama를 처음부터 구현했습니다.

RMSNorm (LayerNorm 대신)

python
# model/llama.py
class RMSNorm(nn.Module):
    """RMS 정규화 - LayerNorm보다 단순."""
    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

Rotary Position Embeddings (RoPE)

python
# model/llama.py
def apply_rotary_pos_emb(q, k, cos, sin):
    """Q와 K 벡터를 회전시켜 위치 정보를 인코딩.

    회전 공식: q_rotated = q * cos + rotate_half(q) * sin
    점곱을 통해 상대 위치를 학습하게 해준다.
    """
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

Grouped Query Attention (GQA)

python
# model/llama.py
class LlamaAttention(nn.Module):
    """GQA: Q 헤드보다 KV 헤드 수를 줄여 메모리 절감."""
    def __init__(self, config):
        self.num_heads = config.num_attention_heads        # 예: 32
        self.num_kv_heads = config.num_key_value_heads      # 예: 8
        self.num_kv_groups = self.num_heads // self.num_kv_heads  # = 4

        # Q 프로젝션이 K,V 프로젝션보다 큼
        self.q_proj = nn.Linear(hidden, num_heads * head_dim)
        self.k_proj = nn.Linear(hidden, num_kv_heads * head_dim)  # 더 작음!
        self.v_proj = nn.Linear(hidden, num_kv_heads * head_dim)

SwiGLU MLP

python
# model/llama.py
class LlamaMLP(nn.Module):
    """SwiGLU: output = down(silu(gate(x)) * up(x))"""
    def forward(self, x):
        gate = F.silu(self.gate_proj(x))  # Swish 활성화
        up = self.up_proj(x)
        return self.down_proj(gate * up)  # 게이트드 선형 유닛

---

9. 교육 모드: 보면서 배우기

제가 특히 좋아하는 기능입니다! nano-vllm은 추론 중 무슨 일이 일어나는지 설명해 주는 교육 모드를 제공합니다.

내레이터(Narrator) 모드

수술을 전문가가 설명하듯, 쉬운 말로 진행 상황을 코멘트합니다:

bash
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "The capital of France is" --narrate

출력:

═══════════════════════════════════════════════════════════════════
  INFERENCE ANATOMY - Educational Mode
═══════════════════════════════════════════════════════════════════

  Prompt: "The capital of France is"
  Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0

═════ ACT 1: TOKENIZATION ═════
Converting your prompt into numbers the model understands...

  "The capital of France is"
       ↓ Tokenizer (BPE algorithm)
  [The] [capital] [of] [France] [is] → [450, 7483, 310, 3444, 338]

═════ ACT 2: PREFILL PHASE ═════
The model reads your entire prompt at once...
  Processing 5 tokens through 22 layers
  ✓ Parallel computation (all tokens at once)
  ✓ Building the KV cache

═════ ACT 3: DECODE PHASE ═════
Now generating one token at a time...

  Step 1: Predicting token #6
  │   Top 5 predictions:
  │   Paris    ████████████████████ 82.3%
  │   the      ███                   7.1%
  │   located  ██                    4.2%
  └── Sampled: "Paris" (82.3%)

X-Ray 모드

텐서 shape와 수학 연산을 보여줍니다:

bash
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "Hello" --xray

대시보드(Dashboard) 모드

실시간 진행 상황을 보여주는 터미널 UI(rich 필요):

bash
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "Hello" --dashboard

인터랙티브 튜토리얼

단계별 학습 경험:

bash
python -m nano_vllm.cli --tutorial

---

빠른 시작

설치

bash
pip install -e .

# 선택: FlashAttention(더 빠른 추론)
pip install flash-attn --no-build-isolation

기본 사용법

bash
# 단일 프롬프트
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "Hello, world"

# 여러 프롬프트(연속 배칭)
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "The capital of France is" \
    --prompt "The largest planet is" \
    --prompt "Python is a"

# 우선순위 스케줄링
python -m nano_vllm.cli --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "Low priority task" --priority 1 \
    --prompt "High priority task" --priority 10

# 추측 디코딩
python -m nano_vllm.speculative.cli \
    --target-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
    --prompt "The future of AI is" \
    --num-speculative-tokens 5

Python API

python
from nano_vllm.engine import LLMEngine

engine = LLMEngine(
    model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    use_paged_attention=True,
    enable_prefix_caching=True,
    use_flash_attn=True,
)

# 단일 생성
output = engine.generate("What is machine learning?", max_tokens=100)

# 우선순위를 포함한 배치 생성
engine.add_request("Prompt 1", max_tokens=50, priority=1)
engine.add_request("Prompt 2", max_tokens=50, priority=10)  # 더 높은 우선순위
outputs = engine.run_to_completion()

---

배운 점

nano-vllm을 만들며 얻은 교훈:

1. 병목은 메모리다 — 대부분의 LLM 추론 최적화는 연산이 아니라 메모리 문제다
2. OS 개념이 그대로 적용된다 — PagedAttention은 KV 캐시를 위한 가상 메모리 그 자체
3. 배칭은 섬세하다 — 연속 배칭은 “그냥 배치에 넣기”보다 훨씬 복잡하다
4. 추측은 강력하다 — 비싼 forward 한 번으로 여러 토큰을 얻는 효과가 크다
5. 악마는 디테일에 있다 — causal 마스킹, position id, 참조 카운팅… 끝없는 엣지 케이스

다음 단계

TODO 목록에는 아직 이런 것들이 있습니다:

• CUDA 그래프(커널 런치 오버헤드 감소)
• 멀티 GPU를 위한 텐서 병렬화
• 양자화(AWQ, GPTQ)
• OpenAI 호환 API 서버

자료

• vLLM 논문
• vLLM 저장소
• PagedAttention 블로그
• FlashAttention 논문
• Speculative Decoding 논문

---

도움이 되었다면 저장소에 스타를 눌러 주세요! 버그를 발견하거나 제안이 있다면 PR도 환영합니다. 즐거운 추론 되세요!