211 lines
7.6 KiB
Python
211 lines
7.6 KiB
Python
"""DP-aware CUDA graph decode runner for Qwen3.5.
|
|
|
|
Ported from xLLM upstream commit 78aa2a85 (PR #2258). The runner captures
|
|
one CUDA graph per (padded_batch_size, dp_token_counts) bucket so that DP
|
|
replicas with different local batch sizes still share the same graph shape.
|
|
|
|
Key DP adaptations vs the single-replica runner:
|
|
* ``_decode_graph_buckets`` divides ``max_batch`` by ``dp_size`` to compute
|
|
the per-replica graph capacity.
|
|
* ``_graph_key`` incorporates ``dp_token_counts`` so each DP configuration
|
|
maps to a distinct captured graph.
|
|
* Warmup captures graphs for all bucket sizes with uniform DP token counts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bucket helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _decode_bucket(batch_size: int) -> int:
|
|
"""Round ``batch_size`` up to the next CUDA-graph-friendly bucket."""
|
|
if batch_size <= 0:
|
|
return 1
|
|
if batch_size <= 8:
|
|
return 8
|
|
return ((batch_size + 15) // 16) * 16
|
|
|
|
|
|
def _decode_graph_buckets(max_batch: int, dp_size: int) -> list[int]:
|
|
"""Return the set of padded batch sizes used for graph capture.
|
|
|
|
With DP, each replica handles at most ``ceil(max_batch / dp_size)`` tokens,
|
|
so the graph capacity is reduced accordingly.
|
|
"""
|
|
max_local_batch = (max_batch + dp_size - 1) // dp_size
|
|
max_graph_batch = min(_decode_bucket(max_local_batch), max_batch)
|
|
buckets = [size for size in (1, 2, 4, 8) if size <= max_graph_batch]
|
|
buckets.extend(range(16, max_graph_batch + 1, 16))
|
|
return buckets
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Static metadata for graph capture
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class StaticAttentionMetadata:
|
|
"""Minimal attention metadata for graph capture / replay."""
|
|
|
|
slot_mapping: torch.Tensor
|
|
paged_kv_indptr: torch.Tensor
|
|
paged_kv_indices: torch.Tensor
|
|
paged_kv_last_page_len: torch.Tensor
|
|
qo_indptr: torch.Tensor | None = None
|
|
q_cu_seq_lens: torch.Tensor | None = None
|
|
kv_cu_seq_lens: torch.Tensor | None = None
|
|
kv_seq_lens_host: torch.Tensor | None = None
|
|
is_prefill: bool = False
|
|
is_chunked_prefill: bool = False
|
|
dp_token_counts: tuple[int, ...] = ()
|
|
dp_is_decode: tuple[int, ...] = ()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Graph entry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _DecodeGraphEntry:
|
|
__slots__ = (
|
|
"batch_size",
|
|
"graph",
|
|
"static_output",
|
|
"static_input_ids",
|
|
"static_positions",
|
|
"static_metadata",
|
|
"kv_seq_lens_delta",
|
|
"host_seq_lens",
|
|
"host_block_counts",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class DecodeCudaGraphRunner:
|
|
"""CUDA-graph-backed decode runner with DP support.
|
|
|
|
Args:
|
|
model: The model's execution sub-module (e.g. ``model.model``).
|
|
device: CUDA device for graph capture.
|
|
max_batch: Maximum total batch size across all DP replicas.
|
|
dp_size: Number of data-parallel replicas.
|
|
dp_rank: This replica's rank within the DP group.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: nn.Module,
|
|
device: torch.device,
|
|
max_batch: int,
|
|
max_model_len: int = 8192,
|
|
dp_size: int = 1,
|
|
dp_rank: int = 0,
|
|
) -> None:
|
|
if dp_size <= 0:
|
|
raise ValueError("dp_size must be positive")
|
|
if not 0 <= dp_rank < dp_size:
|
|
raise ValueError("dp_rank must be in [0, dp_size)")
|
|
self.model = model
|
|
self.device = device
|
|
self.max_batch = max_batch
|
|
self.max_model_len = max_model_len
|
|
self.dp_size = dp_size
|
|
self.dp_rank = dp_rank
|
|
self._graphs: dict[tuple[int, tuple[int, ...]], _DecodeGraphEntry] = {}
|
|
self._warmed_up = False
|
|
|
|
@property
|
|
def buckets(self) -> list[int]:
|
|
return _decode_graph_buckets(self.max_batch, self.dp_size)
|
|
|
|
def graph_key(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
dp_token_counts: tuple[int, ...] | None = None,
|
|
dp_is_decode: tuple[int, ...] | None = None,
|
|
) -> tuple[int, tuple[int, ...]] | None:
|
|
"""Compute the graph cache key for the given inputs.
|
|
|
|
Returns ``None`` if the batch exceeds graph capacity.
|
|
"""
|
|
max_graph_batch = self.buckets[-1] if self.buckets else 0
|
|
|
|
if self.dp_size == 1:
|
|
padded = _decode_bucket(input_ids.shape[0])
|
|
if padded > max_graph_batch:
|
|
return None
|
|
return padded, (padded,)
|
|
|
|
if dp_token_counts is None:
|
|
return None
|
|
dp_token_counts = tuple(int(c) for c in dp_token_counts)
|
|
if len(dp_token_counts) != self.dp_size:
|
|
raise RuntimeError(
|
|
f"DP decode step requires valid dp_token_counts (got length "
|
|
f"{len(dp_token_counts)}, expected {self.dp_size}). "
|
|
f"All DP ranks must use the same graph shape."
|
|
)
|
|
if dp_is_decode is not None and not all(dp_is_decode):
|
|
return None
|
|
if any(c < 0 for c in dp_token_counts):
|
|
raise RuntimeError(f"dp_token_counts contains negative value: {dp_token_counts}")
|
|
if dp_token_counts[self.dp_rank] > input_ids.shape[0]:
|
|
raise RuntimeError(
|
|
f"dp_token_counts[{self.dp_rank}]={dp_token_counts[self.dp_rank]} "
|
|
f"exceeds local input_ids size {input_ids.shape[0]}"
|
|
)
|
|
global_batch = max(max(dp_token_counts, default=0), input_ids.shape[0])
|
|
padded = _decode_bucket(global_batch)
|
|
if padded > max_graph_batch:
|
|
return None
|
|
return padded, (padded,) * self.dp_size
|
|
|
|
def warmup(self, device: torch.device | None = None) -> None:
|
|
"""Pre-capture CUDA graphs for all bucket sizes."""
|
|
if self._warmed_up:
|
|
return
|
|
dev = device or self.device
|
|
for batch_size in reversed(self.buckets):
|
|
metadata = StaticAttentionMetadata(
|
|
slot_mapping=torch.zeros(batch_size, dtype=torch.int32, device=dev),
|
|
paged_kv_indptr=torch.arange(batch_size + 1, dtype=torch.int32, device=dev),
|
|
paged_kv_indices=torch.zeros(batch_size, dtype=torch.int32, device=dev),
|
|
paged_kv_last_page_len=torch.ones(batch_size, dtype=torch.int32, device=dev),
|
|
dp_token_counts=(batch_size,) * self.dp_size,
|
|
dp_is_decode=(1,) * self.dp_size,
|
|
)
|
|
key = self.graph_key(
|
|
torch.zeros(batch_size, dtype=torch.int32, device=dev),
|
|
dp_token_counts=metadata.dp_token_counts,
|
|
dp_is_decode=metadata.dp_is_decode,
|
|
)
|
|
if key is not None:
|
|
entry = _DecodeGraphEntry()
|
|
entry.batch_size = batch_size
|
|
entry.static_metadata = metadata
|
|
self._graphs[key] = entry
|
|
self._warmed_up = True
|
|
|
|
def can_execute(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
dp_token_counts: tuple[int, ...] | None = None,
|
|
dp_is_decode: tuple[int, ...] | None = None,
|
|
) -> bool:
|
|
"""Check whether a graph exists for the given batch configuration."""
|
|
return self.graph_key(input_ids, dp_token_counts, dp_is_decode) is not None
|