Architektur im Überblick

Quelldateien

Diese Seite wurde aus den folgenden Quelldateien erstellt:

• nanovllm/engine/llm_engine.py
• nanovllm/engine/scheduler.py
• nanovllm/engine/model_runner.py
• nanovllm/engine/block_manager.py
• nanovllm/engine/sequence.py
• nanovllm/config.py
• nanovllm/sampling_params.py
• nanovllm/llm.py
• nanovllm/init.py
• nanovllm/utils/loader.py
• pyproject.toml
• bench.py
• example.py
• nanovllm/models/qwen3.py
• nanovllm/layers/linear.py
• nanovllm/utils/context.py
• nanovllm/layers/sampler.py
• nanovllm/layers/attention.py
• nanovllm/layers/layernorm.py
• nanovllm/layers/activation.py

nano-vllm ist eine minimalistische, aber funktionsfähige Implementierung eines LLM-Inferenzsystems, das die Kernkonzepte von vLLM in einer überschaubaren Codebasis demonstriert. Das Projekt implementiert Paged Attention mit automatischem KV-Cache-Management, Prefix Caching zur Deduplizierung wiederkehrender Prompt-Präfixe, sowie CUDA Graph Optimierung für effiziente Decode-Phasen. Die Architektur folgt einem modularen Design mit klarer Trennung zwischen Engine-Koordination, Scheduling, Modell-Ausführung und Speicherverwaltung.

Systemarchitektur

Die Systemarchitektur von nano-vllm besteht aus vier zentralen Komponenten, die über definierte Schnittstellen miteinander interagieren. Die LLMEngine fungiert als zentrale Koordinationseinheit und verwaltet den gesamten Lebenszyklus von Inferenzanfragen. Der Scheduler entscheidet über die Reihenfolge und Batching-Strategie für ausstehende Sequenzen. Der ModelRunner übernimmt die eigentliche Modellinferenz auf der GPU, einschließlich KV-Cache-Verwaltung und CUDA Graph Optimierung. Der BlockManager implementiert das Paged Attention Speichermodell mit Hash-basiertem Prefix Caching.

正在加载图表渲染器...

Erklärung der Architekturkomponenten:

• LLM / LLMEngine: Haupteinstiegspunkt, der Tokenizer, Scheduler und ModelRunner koordiniert. Verwaltet den Lebenszyklus von Worker-Prozessen für Tensor-Parallelität (nanovllm/engine/llm_engine.py:15-35).
• Config: Zentrale Konfiguration mit Parametern wie max_num_seqs, tensor_parallel_size und gpu_memory_utilization (nanovllm/config.py:7-26).
• Scheduler: Implementiert Prefill- und Decode-Scheduling mit Preemption-Unterstützung bei Speichermangel (nanovllm/engine/scheduler.py:8-42).
• BlockManager: Verwaltet physische KV-Cache-Blöcke mit Hash-basiertem Prefix Caching zur Deduplizierung (nanovllm/engine/block_manager.py:26-33).

Engine-Komponenten und Initialisierung

Die LLMEngine-Klasse bildet das Herzstück des Systems und orchestriert alle Komponenten. Bei der Initialisierung werden mehrere kritische Schritte ausgeführt: Konfigurationserstellung, Prozess-Spawning für Tensor-Parallelität, Modell-Laden und Scheduler-Initialisierung.

Initialisierungsablauf

python
1def __init__(self, model, **kwargs):
2    config_fields = {field.name for field in fields(Config)}
3    config_kwargs = {k: v for kwargs.items() if k in config_fields}
4    config = Config(model, **config_kwargs)
5    # Tensor-Parallel-Prozesse spawnen
6    for i in range(1, config.tensor_parallel_size):
7        process = ctx.Process(target=ModelRunner, args=(config, i, event))
8        process.start()
9    self.model_runner = ModelRunner(config, 0, self.events)
10    self.scheduler = Scheduler(config)

Die Engine erstellt für jede GPU einen separaten Prozess, wobei Rank 0 als Master fungiert. Die Kommunikation zwischen Prozessen erfolgt über Shared Memory und Event-Synchronisation (nanovllm/engine/llm_engine.py:15-35).

Konfigurationsparameter

Die Config-Klasse definiert alle steuerbaren Parameter mit sinnvollen Standardwerten:

Die Konfiguration validiert kritische Constraints wie kvcache_block_size % 256 == 0 und begrenzt max_model_len auf die Modellkapazität (nanovllm/config.py:7-26).

Anfrageverarbeitung

Die add_request-Methode transformiert Text-Prompts in Token-Sequenzen und registriert diese beim Scheduler:

python
1def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
2    if isinstance(prompt, str):
3        prompt = self.tokenizer.encode(prompt)
4    seq = Sequence(prompt, sampling_params)
5    self.scheduler.add(seq)

Die step-Methode führt einen Scheduling-Zyklus aus, ruft das Modell auf und verarbeitet die generierten Tokens (nanovllm/engine/llm_engine.py:42-54).

Scheduler und Anfrageverwaltung

Der Scheduler implementiert die zentrale Logik für Batching-Entscheidungen und Ressourcenverwaltung. Er verwaltet zwei Warteschlangen: waiting für neue Sequenzen und running für aktive Sequenzen.

Prefill-Scheduling

Während des Prefill-Phasen-Scheduling werden Sequenzen aus der Warteschlange basierend auf zwei Constraints ausgewählt:

1. Token-Limit: num_batched_tokens + len(seq) <= max_num_batched_tokens
2. Block-Verfügbarkeit: block_manager.can_allocate(seq) muss True zurückgeben

python
1def schedule(self) -> tuple[list[Sequence], bool]:
2    # Prefill-Phase
3    while self.waiting and num_seqs < self.max_num_seqs:
4        seq = self.waiting[0]
5        if num_batched_tokens + len(seq) > self.max_num_batched_tokens:
6            break
7        if not self.block_manager.can_allocate(seq):
8            break
9        self.block_manager.allocate(seq)
10        seq.status = SequenceStatus.RUNNING
11        self.running.append(seq)

Der Algorithmus priorisiert First-Come-First-Served und bricht ab, sobald ein Constraint verletzt wird (nanovllm/engine/scheduler.py:8-42).

Decode-Scheduling mit Preemption

In der Decode-Phase werden laufende Sequenzen verarbeitet. Wenn kein Platz für neue Tokens verfügbar ist, greift die Preemption-Logik:

python
1# Decode-Phase
2while self.running and num_seqs < self.max_num_seqs:
3    seq = self.running.popleft()
4    while not self.block_manager.can_append(seq):
5        if self.running:
6            self.preempt(self.running.pop())  # Andere Sequenz preempten
7        else:
8            self.preempt(seq)  # Sich selbst preempten
9            break

Die preempt-Methode setzt die Sequenz zurück in den WAITING-Status und gibt alle Blöcke frei (nanovllm/engine/scheduler.py:43-71).

Postprocessing und Abschluss-Erkennung

Nach der Token-Generierung prüft postprocess zwei Abschlussbedingungen:

1. EOS-Token: token_id == self.eos (sofern ignore_eos == False)
2. Maximale Länge: num_completion_tokens == max_tokens

python
1def postprocess(self, seqs: list[Sequence], token_ids: list[int]):
2    for seq, token_id in zip(seqs, token_ids):
3        seq.append_token(token_id)
4        if (not seq.ignore_eos and token_id == self.eos) or \
5           seq.num_completion_tokens == seq.max_tokens:
6            seq.status = SequenceStatus.FINISHED
7            self.block_manager.deallocate(seq)

Beende Sequenzen werden aus der Running-Queue entfernt und ihre Blöcke freigegeben (nanovllm/engine/scheduler.py:65-71).

Model Runner und Inferenz-Ausführung

Der ModelRunner kapselt alle GPU-bezogenen Operationen: Modell-Laden, KV-Cache-Allokation, CUDA Graph Capture und die eigentliche Inferenz.

Initialisierung und Prozessgruppen-Setup

python
1def __init__(self, config: Config, rank: int, event: Event | list[Event]):
2    dist.init_process_group("nccl", "tcp://localhost:2333", 
3                            world_size=self.world_size, rank=rank)
4    torch.cuda.set_device(rank)
5    self.model = Qwen3ForCausalLM(hf_config)
6    load_model(self.model, config.model)
7    self.allocate_kv_cache()
8    if not self.enforce_eager:
9        self.capture_cudagraph()

Für Tensor-Parallelität wird NCCL als Backend verwendet. Rank 0 erstellt eine Shared Memory-Region für die Kommunikation mit Worker-Prozessen (nanovllm/engine/model_runner.py:15-48).

Dynamische KV-Cache-Allokation

Die KV-Cache-Größe wird dynamisch basierend auf dem verfügbaren GPU-Speicher berechnet:

python
1def allocate_kv_cache(self):
2    free, total = torch.cuda.mem_get_info()
3    used = total - free
4    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
5    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
6    
7    block_bytes = 2 * num_hidden_layers * block_size * num_kv_heads * head_dim * dtype_size
8    config.num_kvcache_blocks = int(total * gpu_memory_utilization - used - peak + current) // block_bytes
9    
10    self.kv_cache = torch.empty(2, num_hidden_layers, num_kvcache_blocks, 
11                                 block_size, num_kv_heads, head_dim)

Die Formel berücksichtigt bereits allozierten Speicher und Peak-Memory während des Warmup. Die KV-Cache-Tensoren werden direkt in die Attention-Module injiziert (nanovllm/engine/model_runner.py:100-125).

CUDA Graph Optimierung

Für die Decode-Phase werden CUDA Graphs für verschiedene Batch-Größen aufgezeichnet:

python
1def capture_cudagraph(self):
2    self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
3    for bs in reversed(self.graph_bs):
4        graph = torch.cuda.CUDAGraph()
5        with torch.cuda.graph(graph, self.graph_pool):
6            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])
7        self.graphs[bs] = graph

Die Graphs werden in absteigender Reihenfolge aufgezeichnet, um Speicherfragmentierung zu minimieren. Ein gemeinsamer Graph-Pool ermöglicht Speicheroptimierungen zwischen verschiedenen Batch-Größen (nanovllm/engine/model_runner.py:216-251).

Prefill- und Decode-Vorbereitung

Die prepare_prefill-Methode erstellt die notwendigen Tensoren für die Prefill-Phase:

python
1def prepare_prefill(self, seqs: list[Sequence]):
2    for seq in seqs:
3        input_ids.extend(seq[seq.num_cached_tokens:])
4        positions.extend(range(seq.num_cached_tokens, seqlen))
5        # Slot-Mapping für KV-Cache-Writes
6        for i in range(seq.num_cached_blocks, seq.num_blocks):
7            slot_mapping.extend(range(start, end))

num_cached_tokens berücksichtigt Prefix Caching – nur nicht-gecachte Tokens werden verarbeitet. Die prepare_decode-Methode bereitet einzelne Token für die autoregressive Generierung vor (nanovllm/engine/model_runner.py:126-180).

Block Manager und KV-Cache-Verwaltung

Der BlockManager implementiert Paged Attention mit automatischer Speicherverwaltung und Hash-basiertem Prefix Caching.

Block-Struktur und Referenzzählung

Jeder Block verwaltet vier Attribute:

python
1class Block:
2    def __init__(self, block_id):
3        self.block_id = block_id
4        self.ref_count = 0      # Referenzzähler für Sharing
5        self.hash = -1          # Hash für Prefix Caching
6        self.token_ids = []     # Token-Inhalt für Validierung

Die Referenzzählung ermöglicht das Sharing von Blöcken zwischen Sequenzen mit identischem Prefix. Blöcke werden erst freigegeben, wenn ref_count auf 0 fällt (nanovllm/engine/block_manager.py:8-33).

Hash-basiertes Prefix Caching

Die allocate-Methode implementiert Cache-Lookup und Deduplizierung:

python
1def allocate(self, seq: Sequence):
2    h = -1
3    for i in range(seq.num_blocks):
4        token_ids = seq.block(i)
5        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
6        block_id = self.hash_to_block_id.get(h, -1)
7        
8        if block_id != -1 and self.blocks[block_id].token_ids == token_ids:
9            # Cache Hit - Block teilen
10            seq.num_cached_tokens += self.block_size
11            block.ref_count += 1
12        else:
13            # Cache Miss - neuen Block allokieren
14            block = self._allocate_block(free_block_ids[0])

Der Hash wird rekursiv mit dem Prefix-Hash berechnet, was eine O(1) Lookup-Zit ermöglicht. Ein zusätzlicher Token-Vergleich schützt vor Hash-Kollisionen (nanovllm/engine/block_manager.py:56-91).

Dynamische Block-Erweiterung

Während der Decode-Phase werden Blöcke bei Bedarf erweitert:

python
1def may_append(self, seq: Sequence):
2    if len(seq) % self.block_size == 1:
3        # Neuen Block allokieren
4        block_id = self.free_block_ids[0]
5        self._allocate_block(block_id)
6        block_table.append(block_id)
7    elif len(seq) % self.block_size == 0:
8        # Block ist voll - Hash berechnen
9        h = self.compute_hash(token_ids, prefix)
10        last_block.update(h, token_ids)
11        self.hash_to_block_id[h] = last_block.block_id

Die Logik unterscheidet drei Fälle basierend auf len(seq) % block_size: neuer Block benötigt, Block wird voll, oder Block noch nicht voll (nanovllm/engine/block_manager.py:93-112).

Sequenz-Repräsentation und Datenfluss

Die Sequence-Klasse repräsentiert eine einzelne Inferenzanfrage mit allen notwendigen Metadaten.

Sequenz-Status und Eigenschaften

python
1class SequenceStatus(Enum):
2    WAITING = auto()    # In Warteschlange
3    RUNNING = auto()    # Wird verarbeitet
4    FINISHED = auto()   # Abgeschlossen
5
6class Sequence:
7    def __init__(self, token_ids: list[int], sampling_params):
8        self.seq_id = next(Sequence.counter)
9        self.status = SequenceStatus.WAITING
10        self.num_cached_tokens = 0
11        self.block_table = []  # Mapping zu physischen Blöcken
12        self.temperature = sampling_params.temperature
13        self.max_tokens = sampling_params.max_tokens

Die Klasse berechnet abgeleitete Eigenschaften wie num_blocks und last_block_num_tokens dynamisch (nanovllm/engine/sequence.py:8-30).

Block-Berechnung und Token-Appending

python
1@property
2def num_blocks(self):
3    return (self.num_tokens + self.block_size - 1) // self.block_size
4
5@property
6def last_block_num_tokens(self):
7    return self.num_tokens - (self.num_blocks - 1) * self.block_size
8
9def append_token(self, token_id: int):
10    self.token_ids.append(token_id)
11    self.last_token = token_id
12    self.num_tokens += 1

Die Block-Größe ist statisch auf 256 Tokens festgelegt. Die block(i)-Methode ermöglicht den Zugriff auf Token-Bereiche pro Block (nanovllm/engine/sequence.py:53-73).

Datenfluss und End-to-End-Ablauf

Der folgende Sequenzdiagramm zeigt den vollständigen Datenfluss von der Anfrageeingabe bis zur Token-Ausgabe:

正在加载图表渲染器...

Erklärung der Datenfluss-Phasen:

1. Anfrage-Registrierung: Prompts werden tokenisiert und als Sequence-Objekte in die Waiting-Queue eingereiht (nanovllm/engine/llm_engine.py:42-46).
2. Scheduling-Entscheidung: Der Scheduler wählt Sequenzen basierend auf Token-Limit und Block-Verfügbarkeit aus (nanovllm/engine/scheduler.py:24-41).
3. Block-Allokation: Der BlockManager prüft auf Prefix-Cache-Hits und allokiert neue Blöcke bei Bedarf (nanovllm/engine/block_manager.py:59-82).
4. GPU-Inferenz: Der ModelRunner bereitet Tensoren vor und führt das Modell aus (nanovllm/engine/model_runner.py:208-214).
5. Postprocessing: Generierte Tokens werden angehängt und Abschlussbedingungen geprüft (nanovllm/engine/scheduler.py:65-71).

Modulabhängigkeiten und Schichtarchitektur

Die folgende Grafik zeigt die Abhängigkeiten zwischen den Kernmodulen:

正在加载图表渲染器...

Schicht-Erklärung:

• API-Schicht: Exportiert LLM und SamplingParams als öffentliche Schnittstellen (nanovllm/init.py:1-2).
• Engine-Schicht: Koordiniert alle Komponenten und verwaltet den Lebenszyklus.
• Kernlogik-Schicht: Implementiert Scheduling, Sequenzverwaltung und Speichermanagement.
• Ausführungs-Schicht: Führt GPU-Operationen aus und lädt Modellgewichte.
• Modell-Schicht: Definiert die neuronale Netzwerkarchitektur.

Kern-Designentscheidungen

1. Paged Attention mit Block-Granularität

Entscheidung: KV-Cache wird in Blöcken zu 256 Tokens verwaltet, nicht als zusammenhängender Tensor.

Begründung: Ermöglicht dynamische Speicherallokation und effizientes Sharing zwischen Sequenzen. Die Block-Größe von 256 ist ein Kompromiss zwischen Fragmentierung und Verwaltungsoverhead.

Nachteil: Komplexere Index-Berechnung und Block-Tabellen-Verwaltung.

2. Hash-basiertes Prefix Caching

Entscheidung: Blöcke werden mit xxHash gehasht und über ein hash_to_block_id-Dictionary nachgeschlagen.

Begründung: Dedupliziert identische Prompt-Präfixe über verschiedene Anfragen hinweg. Reduziert Rechenkosten für wiederkehrende System-Prompts signifikant.

Nachteil: Zusätzlicher Speicher für Hash-Tabellen und Token-Validierung bei Kollisionen.

3. CUDA Graph für Decode-Phase

Entscheidung: CUDA Graphs werden für Batch-Größen [1, 2, 4, 8, 16, 32, ...] aufgezeichnet.

Begründung: Eliminiert Python-Overhead und Kernel-Launch-Latenz in der Decode-Phase, die den Großteil der Inferenzzeit ausmacht.

Nachteil: Erhöhter Speicherverbrauch durch statische Graph-Puffer und längere Initialisierungszeit.

4. Preemption statt Swapping

Entscheidung: Bei Speichermangel werden Sequenzen preempted (zurück in WAITING), nicht auf CPU ausgelagert.

Begründung: Vereinfachte Implementierung ohne CPU-GPU-Transfer-Logik. Für kurzfristige Engpässe ist Neuberechnung oft schneller als Swapping.

Nachteil: Potenziell höhere Latenz für preemptete Sequenzen bei langen Prompts.

5. Tensor-Parallelität über Multiprocessing

Entscheidung: Jede GPU läuft in einem separaten Prozess mit Shared-Memory-Kommunikation.

Begründung: Umgeht Python's GIL für echte Parallelität. NCCL handles die GPU-Kommunikation effizient.

Nachteil: Komplexere Prozessverwaltung und Synchronisation.

6. Eager-Mode als Fallback

Entscheidung: enforce_eager=True deaktiviert CUDA Graphs vollständig.

Begründung: Ermöglicht Debugging und Kompatibilität mit nicht-unterstützten Operationen. Reduziert Initialisierungszeit für Tests.

Nachteil: Signifikant niedrigere Decode-Throughput.

7. Statische Block-Größe

Entscheidung: block_size=256 ist fest kodiert und nicht konfigurierbar.

Begründung: Vereinfacht Index-Berechnungen und ermöglicht Optimierungen im Attention-Kernel.

Nachteil: Suboptimale Speichereffizienz für sehr kurze oder sehr lange Sequenzen.

Technologie-Stack und Auswahlkriterien

Sampling-Parameter und Konfiguration

Die SamplingParams-Klasse definiert die Generierungsparameter:

python
1@dataclass
2class SamplingParams:
3    temperature: float = 1.0      # Sampling-Temperatur
4    max_tokens: int = 64          # Maximale Output-Tokens
5    ignore_eos: bool = False      # EOS-Token ignorieren

Die Klasse erzwingt temperature > 1e-10, da Greedy-Sampling (temperature=0) nicht unterstützt wird. Dies ist eine bewusste Designentscheidung für deterministisches Verhalten durch niedrige Temperaturen statt exakter Greedy-Implementierung (nanovllm/sampling_params.py:5-11).

Modell-Architektur Qwen3-Integration

Das Projekt implementiert eine vollständige Qwen3-Architektur mit folgenden Komponenten:

Attention-Mechanismus

python
1class Qwen3Attention(nn.Module):
2    def __init__(self, hidden_size, num_heads, num_kv_heads, ...):
3        self.qkv_proj = QKVParallelLinear(...)  # Fused QKV-Projektion
4        self.o_proj = RowParallelLinear(...)
5        self.rotary_emb = get_rope(...)
6        self.attn = Attention(num_heads, head_dim, scaling, num_kv_heads)

Die Attention nutzt fused QKV-Projektionen für Effizienz und unterstützt Grouped Query Attention (GQA) durch separate num_kv_heads (nanovllm/models/qwen3.py:14-87).

MLP mit Gated Activation

python
1class Qwen3MLP(nn.Module):
2    def __init__(self, hidden_size, intermediate_size, hidden_act):
3        self.gate_up_proj = MergedColumnParallelLinear(...)  # Fused Gate+Up
4        self.down_proj = RowParallelLinear(...)
5        self.act_fn = SiluAndMul()  # SiLU-Gating

Die fused Gate-Up-Projektion reduziert Memory-Bandwidth durch einen einzigen Kernel (nanovllm/models/qwen3.py:90-116).

Packed Modules Mapping

Für korrektes Gewichts-Laden definiert das Modell ein Mapping:

python
1packed_modules_mapping = {
2    "q_proj": ("qkv_proj", "q"),
3    "k_proj": ("qkv_proj", "k"),
4    "v_proj": ("qkv_proj", "v"),
5    "gate_proj": ("gate_up_proj", 0),
6    "up_proj": ("gate_up_proj", 1),
7}

Dies ermöglicht das Laden von Standard-HuggingFace-Checkpoints in die fused Architektur (nanovllm/models/qwen3.py:186-192).

Bekannte Einschränkungen

1. Kein Swapping: Bei Speichermangel werden Sequenzen preempted, nicht ausgelagert.
2. Kein Speculative Decoding: Nur Standard-Autoregressive-Generierung.
3. Kein LoRA-Support: Nur Full-Model-Inferenz.
4. Nur Qwen3-Modell: Andere Architekturen erfordern manuelle Implementierung.
5. Kein Greedy-Sampling: Temperature muss > 0 sein.
6. Statische CUDA Graph-Größen: Nur vordefinierte Batch-Größen sind optimiert.
7. Single-Node-Only: Keine verteilte Inferenz über mehrere Nodes.