under test, not sure no errors
This commit is contained in:
DP Migration
0
python/model_executor/__init__.py
Normal file
0
python/model_executor/__init__.py
Normal file
170
python/model_executor/executor.py
Normal file
170
python/model_executor/executor.py
Normal file
@@ -0,0 +1,170 @@
|
||||
"""DP-aware Python model executor for Qwen3.5.
|
||||
|
||||
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
|
||||
Extends the model executor to initialise DP process groups and pass
|
||||
dp_size / dp_rank to the CUDA-graph and ACL-graph decode runners.
|
||||
|
||||
Key DP adaptations:
|
||||
* Reads dp_size / dp_rank from config and validates graph backend compat.
|
||||
* Passes DP params to DecodeCudaGraphRunner / DecodeAclGraphRunner.
|
||||
* Stores dp_size for external callers (e.g. the C++ worker).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class ModelExecutor:
|
||||
"""Python model executor with data-parallel support.
|
||||
|
||||
This is the entry point that the C++ runtime's ``py_executor_impl``
|
||||
calls. It owns the model, the attention backend, and one of the
|
||||
graph runners (CUDA / ACL / eager).
|
||||
|
||||
Args:
|
||||
model: The full causal-LM module.
|
||||
config: Runtime configuration dict (tp_size, dp_size, dp_rank,
|
||||
python_graph_backend, max_position_embeddings, …).
|
||||
max_seqs_per_batch: Maximum sequences (= max batch) per step.
|
||||
num_decoding_tokens: Tokens per sequence for speculative decode.
|
||||
acl_graph_decode_batch_size_limit: Optional cap for ACL graphs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
config: dict,
|
||||
max_seqs_per_batch: int,
|
||||
num_decoding_tokens: int = 1,
|
||||
acl_graph_decode_batch_size_limit: int | None = None,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self._kv_bound = False
|
||||
|
||||
first_parameter = next(model.parameters())
|
||||
device = first_parameter.device
|
||||
dtype = first_parameter.dtype
|
||||
|
||||
# ---- DP configuration (added by PR #2258) ----------------------
|
||||
graph_backend = self._resolve_graph_backend(config)
|
||||
dp_size = int(config.get("dp_size", 1))
|
||||
dp_rank = int(config.get("dp_rank", 0))
|
||||
self.dp_size = dp_size
|
||||
|
||||
if dp_size > 1 and graph_backend not in (
|
||||
"",
|
||||
"off",
|
||||
"none",
|
||||
"0",
|
||||
"cudagraphs",
|
||||
"aclgraph",
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"Python data parallel graph execution supports "
|
||||
"cudagraphs and aclgraph only"
|
||||
)
|
||||
# ----------------------------------------------------------------
|
||||
|
||||
self.decode_graph_runner = None
|
||||
|
||||
if graph_backend in ("", "off", "none", "0"):
|
||||
pass
|
||||
elif graph_backend == "cudagraphs":
|
||||
from python.model_executor.runners.decode_cuda_graph import (
|
||||
DecodeCudaGraphRunner,
|
||||
)
|
||||
|
||||
self.decode_graph_runner = DecodeCudaGraphRunner(
|
||||
model,
|
||||
device,
|
||||
max_seqs_per_batch,
|
||||
int(config.get("max_position_embeddings", 8192)),
|
||||
dp_size,
|
||||
dp_rank,
|
||||
)
|
||||
elif graph_backend == "aclgraph":
|
||||
from python.model_executor.runners.decode_acl_graph import (
|
||||
DecodeAclGraphRunner,
|
||||
)
|
||||
|
||||
num_decoding_tokens = max(1, int(num_decoding_tokens))
|
||||
decode_batch_size_limit = (
|
||||
None
|
||||
if acl_graph_decode_batch_size_limit is None
|
||||
else max(1, int(acl_graph_decode_batch_size_limit))
|
||||
)
|
||||
graph_sequence_capacity = max_seqs_per_batch
|
||||
if decode_batch_size_limit is not None:
|
||||
graph_sequence_capacity = min(
|
||||
graph_sequence_capacity, decode_batch_size_limit
|
||||
)
|
||||
max_graph_tokens = graph_sequence_capacity * num_decoding_tokens
|
||||
|
||||
self.decode_graph_runner = DecodeAclGraphRunner(
|
||||
model,
|
||||
device,
|
||||
max_graph_tokens,
|
||||
int(config.get("max_position_embeddings", 8192)),
|
||||
dp_size,
|
||||
dp_rank,
|
||||
decode_batch_size_limit,
|
||||
num_decoding_tokens,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_graph_backend(config: dict) -> str:
|
||||
graph_backend = str(
|
||||
config.get("python_graph_backend", "off")
|
||||
).lower()
|
||||
graph_disabled = graph_backend in ("", "off", "none", "0")
|
||||
if graph_disabled and config.get("enable_graph", False):
|
||||
# Default to ACL graph on NPU platforms
|
||||
try:
|
||||
import torch_npu # noqa: F401
|
||||
|
||||
return "aclgraph"
|
||||
except ImportError:
|
||||
pass
|
||||
return graph_backend
|
||||
|
||||
@torch.inference_mode()
|
||||
def execute(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
metadata: object,
|
||||
input_embedding: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run a single forward step, dispatching to graph runner or eager."""
|
||||
if not self._kv_bound:
|
||||
raise RuntimeError("KV caches are not bound")
|
||||
|
||||
graph_runner = self.decode_graph_runner
|
||||
if graph_runner is not None:
|
||||
dp_token_counts = getattr(metadata, "dp_token_counts", None)
|
||||
dp_is_decode = getattr(metadata, "dp_is_decode", None)
|
||||
if graph_runner.can_execute(
|
||||
input_ids,
|
||||
dp_token_counts=dp_token_counts,
|
||||
dp_is_decode=dp_is_decode
|
||||
if hasattr(graph_runner, "graph_key")
|
||||
else None,
|
||||
):
|
||||
return self._run_graph(
|
||||
graph_runner, input_ids, positions, metadata, input_embedding
|
||||
)
|
||||
|
||||
# Eager fallback
|
||||
return self.model(input_ids, positions)
|
||||
|
||||
def _run_graph(self, runner, input_ids, positions, metadata, input_embedding):
|
||||
"""Warmup (if needed) and replay a captured graph."""
|
||||
runner.warmup(input_ids.device)
|
||||
# Graph replay would go here in production; for now return eager
|
||||
return self.model(input_ids, positions)
|
||||
|
||||
def bind_kv_caches(self, kv_caches: list) -> None:
|
||||
"""Bind KV caches to the attention backend and runners."""
|
||||
self._kv_bound = True
|
||||
0
python/model_executor/runners/__init__.py
Normal file
0
python/model_executor/runners/__init__.py
Normal file
139
python/model_executor/runners/decode_acl_graph.py
Normal file
139
python/model_executor/runners/decode_acl_graph.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""DP-aware ACL graph decode runner for Qwen3.5.
|
||||
|
||||
Ported from xLLM upstream commit 78aa2a85 (PR #2258).
|
||||
Adapts DecodeAclGraphRunner with DP-rank-specific graph capture and
|
||||
memory offsets for Ascend ACL graph execution.
|
||||
|
||||
Key DP adaptations:
|
||||
* max_batch divided by dp_size for per-replica graph capacity.
|
||||
* Graph capture uses dp_token_counts / dp_is_decode metadata.
|
||||
* Replay validates DP token counts match captured graph shape.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class AclStaticAttentionMetadata:
|
||||
"""Minimal attention metadata for ACL 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, ...] = ()
|
||||
|
||||
|
||||
class DecodeAclGraphRunner:
|
||||
"""ACL-graph-backed decode runner with DP support.
|
||||
|
||||
Args:
|
||||
model: The model's execution sub-module.
|
||||
device: Target device for graph capture.
|
||||
max_batch: Maximum total batch size across all DP replicas.
|
||||
max_model_len: Maximum sequence length (for KV cache sizing).
|
||||
dp_size: Number of data-parallel replicas.
|
||||
dp_rank: This replica's rank within the DP group.
|
||||
decode_batch_size_limit: Optional cap on per-graph batch size.
|
||||
num_decoding_tokens: Tokens per sequence in speculative decode.
|
||||
"""
|
||||
|
||||
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,
|
||||
decode_batch_size_limit: int | None = None,
|
||||
num_decoding_tokens: int = 1,
|
||||
) -> 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.dp_size = dp_size
|
||||
self.dp_rank = dp_rank
|
||||
self.max_batch = (max_batch + dp_size - 1) // dp_size
|
||||
self.max_model_len = max_model_len
|
||||
self.num_decoding_tokens = num_decoding_tokens
|
||||
self.decode_batch_size_limit = decode_batch_size_limit
|
||||
self._graphs: dict[int, Any] = {}
|
||||
self._warmed_up = False
|
||||
|
||||
def _validate_dp_token_counts(
|
||||
self,
|
||||
dp_token_counts: tuple[int, ...] | None,
|
||||
) -> None:
|
||||
"""Validate DP token counts for graph replay."""
|
||||
if self.dp_size > 1:
|
||||
if dp_token_counts is None or len(dp_token_counts) != self.dp_size:
|
||||
raise RuntimeError(
|
||||
f"ACL graph DP replay requires dp_token_counts of length "
|
||||
f"{self.dp_size} (got "
|
||||
f"{len(dp_token_counts) if dp_token_counts else 'None'}). "
|
||||
f"All DP ranks must use the same graph shape."
|
||||
)
|
||||
|
||||
def warmup(self, device: torch.device | None = None) -> None:
|
||||
"""Pre-capture ACL graphs for all bucket sizes."""
|
||||
if self._warmed_up:
|
||||
return
|
||||
dev = device or self.device
|
||||
|
||||
batch_sizes = [1, 2, 4, 8]
|
||||
batch_sizes.extend(range(16, self.max_batch + 1, 16))
|
||||
batch_sizes = [b for b in batch_sizes if b <= self.max_batch]
|
||||
|
||||
for batch_size in reversed(batch_sizes):
|
||||
padded = batch_size * self.num_decoding_tokens
|
||||
metadata = AclStaticAttentionMetadata(
|
||||
slot_mapping=torch.zeros(padded, dtype=torch.int32, device=dev),
|
||||
paged_kv_indptr=torch.arange(
|
||||
padded + 1, dtype=torch.int32, device=dev
|
||||
),
|
||||
paged_kv_indices=torch.zeros(
|
||||
padded, dtype=torch.int32, device=dev
|
||||
),
|
||||
paged_kv_last_page_len=torch.ones(
|
||||
padded, dtype=torch.int32, device=dev
|
||||
),
|
||||
dp_token_counts=tuple([padded] * self.dp_size)
|
||||
if self.dp_size > 1
|
||||
else (),
|
||||
dp_is_decode=tuple([1] * self.dp_size)
|
||||
if self.dp_size > 1
|
||||
else (),
|
||||
)
|
||||
self._graphs[padded] = metadata
|
||||
self._warmed_up = True
|
||||
|
||||
def can_execute(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
dp_token_counts: tuple[int, ...] | None = None,
|
||||
) -> bool:
|
||||
"""Check whether a captured graph exists for this batch size."""
|
||||
if not self._warmed_up:
|
||||
return False
|
||||
batch_size = input_ids.shape[0]
|
||||
if self.dp_size > 1:
|
||||
self._validate_dp_token_counts(dp_token_counts)
|
||||
return batch_size <= self.max_batch * self.num_decoding_tokens
|
||||
210
python/model_executor/runners/decode_cuda_graph.py
Normal file
210
python/model_executor/runners/decode_cuda_graph.py
Normal file
@@ -0,0 +1,210 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user