ubatch und cudagraph: wie formstabile Graphen Decode/Prefill beschleunigen

源码版本v0.25.1

Verantwortung

cudagraph ist der von CUDA bereitgestellte „einmal aufnehmen, unbegrenzt oft abspielen“-Mechanismus: Eine Kernel-Sequenz wird als Graph aufgezeichnet und bei jedem Replay entfällt der Kernel-Launch-Overhead. In der Decode-Phase werden pro Token nur ein bis zwei neue Token berechnet, die Kernel sind klein und zahlreich — genau das Terrain von cudagraph. v1 wickelt den Modell-Forward mit CUDAGraphWrapper(cuda_graph.py:145-233) ein; beim ersten Auftreten eines BatchDescriptor in __call__ wird der Graph aufgezeichnet(cuda_graph.py:265-344), bei derselben Form später direkt entry.cudagraph.replay()(cuda_graph.py:357-361).

ubatch (Mikrobatch, micro-batch) ist der v1-Ansatz, bei DP>1 einen großen Batch eines Steps in N formstabile Klein-Batches aufzuteilen, für die jeweils ein eigener Graph aufgezeichnet wird. So können verschiedene DP-Ranks gleichzeitig laufen und gleichzeitig TP/EP-all2all-Kommunikation durchführen, sodass Kommunikation und Rechnung überlappen. Diese Logik ist in UBatchWrapper(gpu_ubatch_wrapper.py:113-211) und UBatchContext(ubatching.py:20-148) implementiert. Beides zusammen ergibt DBO (Distributed Batch Overlap): Mikrobatch-Aufteilung + cudagraph-Aufzeichnung + Multithread-Event-Synchronisation.

Der cudagraph-Laufmodus wird durch CUDAGraphMode(compilation.py:53-87) bestimmt: NONE (nicht aufzeichnen), PIECEWISE (segmentweise Aufzeichnung, Attention bleibt außerhalb des Graphen), FULL (alles aufzeichnen), FULL_DECODE_ONLY (nur in Decode wird full verwendet, in Prefill piecewise) usw. Am Ende von GPUModelRunner.load_model wird self.model abhängig von der Konfiguration mit dem passenden Wrapper versehen(gpu_model_runner.py:5350-5379).

Entwurfsmotivation

• cudagraph formgesteuert:CUDAGraphWrapper.__call__ verwendet BatchDescriptor als Schlüssel(cuda_graph.py:257-263); gleiche Form wird nur einmal aufgezeichnet, danach direkt replay. Der CUDAGraphDispatcher im Runner verwaltet zentral, „welche Form in welchem Modus aufgezeichnet wird“(gpu_model_runner.py:844-845), und _determine_batch_execution_and_padding paddet num_tokens auf die nächste aufzeichnungsfähige Form(gpu_model_runner.py:3879-3894).
• PIECEWISE überspringt Attention:Im full-Graphen ist der Attention-Kernel zu formflexibel (variable query/key-Längen); eine Aufzeichnung wäre beim nächsten Replay nutzlos. v1 bietet daher den piecewise-Modus, der das Modell in mehrere Segmente zerlegt — das Attention-Segment läuft eager, die anderen im Graphen — umgesetzt durch BreakableCUDAGraphWrapper(breakable_cudagraph.py:246).
• ubatch gleichmäßig aufteilen:maybe_create_ubatch_slices berechnet die Teilungspunkte aus num_tokens_padded // num_ubatches(ubatch_utils.py:63-114); jeder ubatch hat dieselbe Form, sodass ein und derselbe cudagraph mehrfach wiederverwendet werden kann und die DP-Ranks zueinander passen.
• Multithreading + CUDA-Stream-Event-Synchronisation:UBatchContext nutzt threading.Barrier + cpu_wait_event/cpu_signal_event + gpu_comm_done_event/gpu_compute_done_event(ubatching.py:25-48), um auf CPU- und GPU-Seite zwei ubatch-Threads zu synchronisieren: yield_and_switch_from_compute_to_comm wechselt den aktuellen Stream von compute auf comm und wartet, bis die GPU-Compute abgeschlossen ist(ubatching.py:133-139); in die andere Richtung entsprechend.
• CUDA-Kontext während Capture im separaten Thread initialisieren:_capture_ubatches startet pro ubatch einen Thread; zunächst wird torch.cuda.current_blas_handle() aufgerufen, damit der CUDA-Kontext auf diesem Thread initialisiert wird, bevor der Thread in den UBatchContext eintritt und auf die Barrier wartet(gpu_ubatch_wrapper.py:236-251) — sonst passt der Stream des Haupt-Threads beim Capture nicht zum Kontext der ubatch-Threads.
• SM-Aufteilung für Kommunikation und Rechnen:SMControlContextManager(gpu_ubatch_wrapper.py:68-112) reserviert über set_comm_sms/set_compute_sms eine feste Anzahl SMs für die DeepEP-all2all-Kommunikation, damit ein großflächiges MoE-all2all die GPU nicht vollständig blockiert und die Compute-Seite ohne SM bleibt.
• Gemeinsamer Graphen-Pool:UBatchWrapper.graph_pool reicht an CUDAGraphWrapper.graph_pool weiter(gpu_ubatch_wrapper.py:143-147); alle Formen teilen sich einen Speicher-Pool — große Graphen zuerst, kleine Graphen wiederverwenden(gpu_model_runner.py:6662-6668).

Schlüsseldateien

• CUDAGraphMode:53-87 — Enum: NONE/PIECEWISE/FULL/FULL_DECODE_ONLY/FULL_AND_PIECEWISE; bietet decode_mode/mixed_mode/has_full_cudagraphs und andere Helper.
• cudagraph_batch_sizes:729-736 — liest aufzeichnungsfähige Formen aus compilation_config.cudagraph_capture_sizes und sortiert sie.
• cudagraph_dispatcher:844-845 — CUDagraphDispatcher(vllm_config) bündelt Dispatch und Capture.
• _determine_batch_execution_and_padding:3836-3948 — berechnet uniform_decode/has_lora/has_encoder_output, ruft cudagraph_dispatcher.dispatch(num_tokens, ...) auf und paddet num_tokens auf BatchDescriptor.num_tokens.
• maybe_create_ubatch_slices:4197-4203 — teilt bei should_ubatch=True anhand von num_tokens_padded // num_ubatches und liefert (ubatch_slices, ubatch_slices_padded) zurück; die gepaddete Version ist ausschließlich für Attention.
• load_model wrapper selection:5350-5379 — wählt abhängig von cudagraph_mode + use_ubatching BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper für self.model.
• capture_model:6647-6712 — große Formen zuerst, set_cudagraph_capturing_enabled(True), iteriert über cudagraph_dispatcher.get_capture_descs() und schließt mit lock_workspace() ab.
• UBatchSlice:13-29 — Dataclass aus request_slice + token_slice, mit is_empty() und num_tokens Properties.
• check_ubatch_thresholds:38-46 — bei deaktiviertem use_ubatching direkt False; uniform decode verwendet dbo_decode_token_threshold, andernfalls dbo_prefill_token_threshold.
• maybe_create_ubatch_slices:63-114 — findet per searchsorted auf cu_num_tokens für jeden ubatch den request_slice und paddet am Ende mit _pad_out_ubatch_slices auf num_tokens_padded.
• UBatchContext:20-148 — CPU-+GPU-Event-Koordinator; yield_and_switch_from_compute_to_comm/yield_and_switch_from_comm_to_compute sind der Kern(ubatching.py:133-147).
• dbo_enabled / dbo_current_ubatch_id:150-157 — schlägt über das globale Dict _THREAD_ID_TO_CONTEXT nach, in welchem ubatch der aktuelle Thread läuft; Code nutzt das, um den Buffer des aktuellen ubatch für den Kernel zu wählen.
• UBatchWrapper.__init__:113-150 — baut comm_stream, ready_barrier(num_ubatches+1), cudagraphs: dict[int, CUDAGraphMetaData] und entscheidet abhängig von runtime_mode, ob ein CUDAGraphWrapper angelegt wird.
• _capture_ubatches:212-303 — startet ubatch-Threads zur Initialisierung des CUDA-Kontexts; der Haupt-Thread umschließt mit torch.cuda.graph(...) alle Threads, nach der Aufnahme folgt self.cudagraphs[num_tokens] = cudagraph_metadata.
• _run_ubatches:305-341 — Fallback-Pfad ohne cudagraph; das Modell wird rein multithreaded ausgeführt.
• UBatchWrapper.__call__:441-537 — holt ubatch_slices/cudagraph_runtime_mode aus dem forward_context; im FULL-Modus mit num_tokens in self.cudagraphs wird cudagraph_metadata.cudagraph.replay() aufgerufen(gpu_ubatch_wrapper.py:516-521) — sonst Capture oder _run_ubatches.
• CUDAGraphWrapper class:145-233 — concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry]; runtime_mode unterscheidet FULL/PIECEWISE.
• CUDAGraphWrapper.__call__:233-361 — beim ersten Auftreten eines BatchDescriptor wird mit torch.cuda.graph(...) aufgezeichnet, danach entry.cudagraph.replay(); vor dem Replay sorgt get_offloader().sync_prev_onload() dafür, dass der Offloader das Prefetch abgeschlossen hat.
• BreakableCUDAGraphWrapper:246 — Implementierung des PIECEWISE-Modus; das Attention-Segment wird herausgetrennt und eager ausgeführt.

Datenfluss

Der execute_model-Aufruf des Runners geht in _determine_batch_execution_and_padding über, wählt cudagraph_mode und batch_descriptor, ruft dann maybe_create_ubatch_slices auf und steckt diese Werte schließlich über set_forward_context in den Forward-Kontext; mit self.model(...) wird der Wrapper betreten:

# vllm/v1/worker/gpu_model_runner.py L4197-L4203
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
    should_ubatch,
    num_scheduled_tokens_np,
    num_tokens_padded,
    num_reqs_padded,
    self.parallel_config.num_ubatches,
)

In UBatchWrapper.__call__ wird geprüft, ob forward_context.ubatch_slices None ist. Ist es None, folgt der Ein-Graph-Pfad (oder reines eager), andernfalls der ubatch-Multithread-Pfad:

# vllm/v1/worker/gpu_ubatch_wrapper.py L493-L521
if (
    num_tokens not in self.cudagraphs
    and cudagraph_runtime_mode is CUDAGraphMode.FULL
):
    ubatch_metadata = self._make_ubatch_metadata(
        ubatch_slices=ubatch_slices,
        attn_metadata=attn_metadata,
        slot_mapping=slot_mapping,
        input_ids=input_ids,
        positions=positions,
        intermediate_tensors=intermediate_tensors,
        inputs_embeds=inputs_embeds,
        compute_stream=compute_stream,
        dp_metadata=ubatch_dp_metadata,
        batch_descriptor=batch_descriptor,
        cudagraph_runtime_mode=CUDAGraphMode.NONE,
    )
    with self.sm_control:
        return self._capture_ubatches(ubatch_metadata, self.runnable)
elif (
    num_tokens in self.cudagraphs
    and cudagraph_runtime_mode is CUDAGraphMode.FULL
):
    cudagraph_metadata = self.cudagraphs[num_tokens]
    get_offloader().sync_prev_onload()
    cudagraph_metadata.cudagraph.replay()
    return cudagraph_metadata.outputs

_capture_ubatches startet N Threads; jeder betritt seinen eigenen UBatchContext. Auf CPU-Seite wechseln sich cpu_wait_event/cpu_signal_event ab, sodass einer rechnet und einer wartet; auf GPU-Seite werden gpu_comm_done_event/gpu_compute_done_event auf comm_stream und compute_stream aufgezeichnet und abgewartet(ubatching.py:82-92). Der Haupt-Thread nimmt im torch.cuda.graph(...)-Kontext die gesamte Sequenz als einen Graphen auf und speichert ihn in self.cudagraphs[num_tokens]. Künftige Aufrufe derselben Form verwenden direkt replay() — ohne Python-Overhead.

Grenzen und Fehler

• Leerer Batch: Fallback:Bei should_ubatch=False liefert maybe_create_ubatch_slices direkt return None, None(ubatch_utils.py:71-72); UBatchWrapper.__call__ erkennt ubatch_slices is None und geht den Ein-Graph-Pfad(gpu_ubatch_wrapper.py:447-464).
• Formkonflikt im full-Modus:Eine Form mit aktiviertem ubatching wird nur als ubatch-Graph aufgezeichnet; eine Form ohne ubatching darf nicht von CUDAGraphWrapper doppelt erfasst werden. Daher setzt __call__ explizit if batch_descriptor.num_tokens in self.cudagraphs: cudagraph_runtime_mode = NONE(gpu_ubatch_wrapper.py:455-458).
• DP-Inkonsistenz: direkter Assert:In UBatchWrapper.__call__ gilt auf dem ubatch-Pfad assert dp_metadata is not None(gpu_ubatch_wrapper.py:477-478) — ubatch ist für DP>1 gedacht, einzelne DP kommen hier nicht vorbei.
• Letzter ubatch ist leer:is_last_ubatch_empty prüft padded_num_tokens // num_ubatches * (num_ubatches-1) >= orig_num_tokens(ubatch_utils.py:32-35); in diesem Fall enthält der letzte ubatch keine echten Token und muss im Code übersprungen werden.
• Offloader-Synchronisation während Capture:Vor der Aufnahme ruft _capture_ubatches get_offloader().sync_prev_onload() auf(gpu_ubatch_wrapper.py:283-285), danach get_offloader().join_after_forward()(gpu_ubatch_wrapper.py:298-301) — sonst führt ein nicht angeschlossener Prefetch-Stream des Offloaders zu einem „unjoined stream“-Fehler.
• cudagraph erlaubt kein versehentliches Capture zur Laufzeit:set_cudagraph_capturing_enabled(False) steht am Ende von capture_model(gpu_model_runner.py:6689-6694); validate_cudagraph_capturing_enabled prüft vor dem Capture durch CUDAGraphWrapper(cuda_graph.py:276-277); ein versehentliches Capture zur Laufzeit wirft einen Fehler.
• Form-Synchronisation über DP-Ranks:coordinate_batch_across_dp entscheidet rankübergreifend über should_ubatch und num_tokens_across_dp(gpu_model_runner.py:3907-3918) — alle Ranks erhalten identische ubatch-Formen, damit beim Replay das all2all nicht verrutscht.
• Input-Adressen müssen stabil sein:CUDAGraphWrapper prüft im Debug-Modus assert new_input_addresses == entry.input_addresses(cuda_graph.py:346-355); persistent_buffers werden vom Runner vorbereitet, da sonst beim Replay veraltete Adressen gelesen würden.

Zusammenfassung

cudagraph ist der zentrale Hebel, mit dem v1 Decode und Prefill beschleunigt: Form einmal aufnehmen, mehrfach replay; CUDAGraphWrapper verwendet BatchDescriptor als Schlüssel, und CUDAGraphDispatcher entscheidet, welchen Modus jeder Batch nutzt. ubatch teilt bei DP>1 einen Batch in N formgleiche Klein-Batches auf, sodass Kommunikation und Rechnen zwischen DP-Ranks überlappen können; UBatchContext orchestriert das über CPU-Events und GPU-Stream-Events beim Wechsel zwischen zwei Threads. Beides zusammen ist DBO — bei großen MoE-Deployments praktisch ein Muss. Wie der Forward zu dieser Schicht gelangt, steht in /worker/gpu-model-runner; wie der Worker den Runner aufruft, in /worker/worker; die Executor-Ebene in /executor/executor.

Siehe offizielle Dokumentation: vLLM 文档 · README.