wehub-resource-sync
866 lines
33 KiB
Python
866 lines
33 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""CUDA graph manager for vision encoder budget-batch execution."""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.distributed import (
|
|
get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size,
|
|
tensor_model_parallel_all_gather,
|
|
)
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.interfaces import (
|
|
SupportsEncoderCudaGraph,
|
|
)
|
|
from vllm.model_executor.models.utils import scatter_output_slices
|
|
from vllm.model_executor.models.vision import get_load_balance_assignment
|
|
from vllm.v1.worker.encoder_cudagraph_defs import (
|
|
EncoderCudaGraphConfig,
|
|
EncoderItemSpec,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class BudgetGraphMetadata:
|
|
"""Metadata for a single budget graph.
|
|
|
|
CUDA graph replay pattern:
|
|
* Copy precomputed values into input_buffers
|
|
* Replay graph
|
|
* Read encoder outputs from output_buffer
|
|
"""
|
|
|
|
token_budget: int
|
|
max_batch_size: int # Max number of images/videos per batch
|
|
max_frames_per_batch: int # Max total frames per batch (for video)
|
|
graph: torch.cuda.CUDAGraph
|
|
# Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
|
|
# Before replay the manager updates these in-place. By default buffers are
|
|
# zeroed before slice-copying the actual values; model-specific padding
|
|
# behavior is provided by EncoderCudaGraphConfig.padding_logics.
|
|
input_buffers: dict[str, torch.Tensor]
|
|
# Output written by graph, read after replay
|
|
output_buffer: torch.Tensor
|
|
|
|
|
|
class EncoderCudaGraphManager:
|
|
"""Budget-based CUDA graph capture/replay for vision encoders."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
model: SupportsEncoderCudaGraph,
|
|
):
|
|
"""Initialize CUDA graph manager with provided token budgets
|
|
and max batch size."""
|
|
self.vllm_config = vllm_config
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.model = model
|
|
self.config: EncoderCudaGraphConfig = model.get_encoder_cudagraph_config()
|
|
|
|
comp_config = vllm_config.compilation_config
|
|
user_budgets = comp_config.encoder_cudagraph_token_budgets
|
|
user_max_vision_items = comp_config.encoder_cudagraph_max_vision_items_per_batch
|
|
user_max_frames = comp_config.encoder_cudagraph_max_frames_per_batch
|
|
|
|
multimodal_config = vllm_config.model_config.multimodal_config
|
|
|
|
# Invariant: max_batch_size <= min_token_budget.
|
|
# This ensures per_image_output = budget // max_batch_size >= 1
|
|
# for every captured budget, preventing reshape crashes on empty
|
|
# tensors during CUDA graph capture. Validated/enforced below for
|
|
# each configuration path.
|
|
if user_budgets and user_max_vision_items > 0:
|
|
# Fully user-specified: validate the invariant.
|
|
self.token_budgets = sorted(user_budgets)
|
|
self.max_batch_size = user_max_vision_items
|
|
min_tok = min(self.token_budgets)
|
|
if self.max_batch_size > min_tok:
|
|
raise ValueError(
|
|
f"encoder_cudagraph_max_vision_items_per_batch "
|
|
f"({self.max_batch_size}) must be <= smallest token "
|
|
f"budget ({min_tok}). With budgets="
|
|
f"{self.token_budgets}, per_image_output = "
|
|
f"{min_tok} // {self.max_batch_size} = "
|
|
f"{min_tok // self.max_batch_size}, which would cause "
|
|
f"a capture failure. Either increase the smallest "
|
|
f"budget or decrease max_vision_items_per_batch."
|
|
)
|
|
else:
|
|
# Auto-infer missing values from model.
|
|
min_budget, max_budget = model.get_encoder_cudagraph_budget_range(
|
|
vllm_config
|
|
)
|
|
if min_budget <= 0 or max_budget <= 0:
|
|
raise ValueError(
|
|
f"Invalid encoder cudagraph budget range: "
|
|
f"min_budget={min_budget}, max_budget={max_budget}. "
|
|
f"Both must be positive."
|
|
)
|
|
if min_budget > max_budget:
|
|
raise ValueError(
|
|
f"Invalid encoder cudagraph budget range: "
|
|
f"min_budget={min_budget} > max_budget={max_budget}."
|
|
)
|
|
|
|
if user_max_vision_items > 0:
|
|
# User provided max_vision_items only; adjust auto-inferred
|
|
# budgets so min(budgets) >= max_batch_size.
|
|
self.max_batch_size = user_max_vision_items
|
|
effective_min = max(min_budget, user_max_vision_items)
|
|
self.token_budgets = self._generate_budgets(effective_min, max_budget)
|
|
elif user_budgets:
|
|
# User provided budgets only; cap auto-inferred
|
|
# max_batch_size to min(user_budgets).
|
|
self.token_budgets = sorted(user_budgets)
|
|
self.max_batch_size = min(
|
|
max_budget // min_budget,
|
|
min(self.token_budgets),
|
|
)
|
|
else:
|
|
# Fully auto-inferred.
|
|
self.token_budgets = self._generate_budgets(min_budget, max_budget)
|
|
self.max_batch_size = min(
|
|
max_budget // min_budget,
|
|
min(self.token_budgets),
|
|
)
|
|
|
|
assert multimodal_config is not None
|
|
if multimodal_config.get_limit_per_prompt("video") == 0:
|
|
self.max_frames_per_batch = 0
|
|
elif user_max_frames is not None:
|
|
self.max_frames_per_batch = user_max_frames
|
|
else:
|
|
# Set it to the model-specific value from config.
|
|
max_frames_per_video = self.config.max_frames_per_video
|
|
self.max_frames_per_batch = self.max_batch_size * max_frames_per_video
|
|
|
|
mm_config = vllm_config.model_config.multimodal_config
|
|
self.use_dp = (
|
|
mm_config is not None
|
|
and mm_config.mm_encoder_tp_mode == "data"
|
|
and vllm_config.parallel_config.tensor_parallel_size > 1
|
|
)
|
|
|
|
self.budget_graphs: dict[str, dict[int, BudgetGraphMetadata]] = {}
|
|
self.graph_pool: Any | None = None
|
|
self.graph_hits = 0
|
|
self.graph_misses = 0
|
|
self.log_stats_interval = 100
|
|
|
|
if self.config.enable_dual_path_graph:
|
|
max_budget = self.token_budgets[-1]
|
|
self.global_token_budgets = self._generate_budgets(
|
|
self.config.global_token_per_image,
|
|
max_budget,
|
|
)
|
|
self.local_token_budgets = self._generate_budgets(
|
|
self.config.local_token_per_patch,
|
|
max_budget,
|
|
)
|
|
# When `image_width <= 640 and image_height <= 640`, the mm inputs
|
|
# will only contain global image, without generating local patches.
|
|
self.local_token_budgets.insert(0, 0)
|
|
logger.info(
|
|
"EncoderCudaGraphManager dual-path mode: "
|
|
"global_budgets=%s, local_budgets=%s",
|
|
self.global_token_budgets,
|
|
self.local_token_budgets,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"EncoderCudaGraphManager initialized with "
|
|
"budgets=%s, max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
|
|
self.token_budgets,
|
|
self.max_batch_size,
|
|
self.max_frames_per_batch,
|
|
self.use_dp,
|
|
)
|
|
|
|
@staticmethod
|
|
def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
|
|
"""Generate power-of-2 token budgets from min_budget to max_budget."""
|
|
budgets: list[int] = []
|
|
b = min_budget
|
|
while b <= max_budget:
|
|
budgets.append(b)
|
|
b *= 2
|
|
# Always include max_budget if it's not already a power-of-2 boundary
|
|
if not budgets or budgets[-1] < max_budget:
|
|
budgets.append(max_budget)
|
|
return budgets
|
|
|
|
def supports_modality(self, modality: str) -> bool:
|
|
"""Check if a modality is supported by this manager."""
|
|
return modality in self.config.modalities
|
|
|
|
def clear(self) -> None:
|
|
"""Release captured encoder CUDA graphs and the manager-local pool."""
|
|
for graph_set in self.budget_graphs.values():
|
|
graph_set.clear()
|
|
self.graph_pool = None
|
|
|
|
def capture(self, graph_pool: Any):
|
|
"""Capture CUDA graphs for all token budgets."""
|
|
self.graph_pool = graph_pool
|
|
|
|
if self.config.enable_dual_path_graph:
|
|
for token_budget in sorted(self.global_token_budgets, reverse=True):
|
|
self._capture_budget_graph(token_budget, path="global")
|
|
for token_budget in sorted(self.local_token_budgets, reverse=True):
|
|
if token_budget == 0:
|
|
continue
|
|
self._capture_budget_graph(token_budget, path="local")
|
|
logger.info(
|
|
"Encoder CUDA graph capture complete. "
|
|
"Captured %d global + %d local budget graphs.",
|
|
len(self.budget_graphs["global"]),
|
|
len(self.budget_graphs["local"]),
|
|
)
|
|
return
|
|
|
|
for token_budget in sorted(self.token_budgets, reverse=True):
|
|
self._capture_budget_graph(token_budget)
|
|
|
|
logger.info(
|
|
"Encoder CUDA graph capture complete. Captured %d budget graphs.",
|
|
len(self.budget_graphs["default"]),
|
|
)
|
|
|
|
def get_num_graphs_to_capture(self) -> int:
|
|
if self.config.enable_dual_path_graph:
|
|
return len(self.global_token_budgets) + len(self.local_token_budgets)
|
|
return len(self.token_budgets)
|
|
|
|
def _get_graph_set(self, path: str = "default") -> dict[int, BudgetGraphMetadata]:
|
|
# Lazy init global/local graph sets for dual-path models, or default graph
|
|
# set for single-path models.
|
|
if path not in self.budget_graphs:
|
|
self.budget_graphs[path] = {}
|
|
return self.budget_graphs[path]
|
|
|
|
def _capture_budget_graph(self, token_budget: int, path: str = "default"):
|
|
"""Capture CUDA graph for a single token budget."""
|
|
logger.debug(
|
|
"Capturing encoder cudagraph for budget=%d, max_batch_size=%d, "
|
|
"max_frames_per_batch=%d",
|
|
token_budget,
|
|
self.max_batch_size,
|
|
self.max_frames_per_batch,
|
|
)
|
|
|
|
graph_set = self._get_graph_set(path)
|
|
|
|
capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(
|
|
token_budget,
|
|
self.max_batch_size,
|
|
self.max_frames_per_batch,
|
|
self.device,
|
|
self.dtype,
|
|
path,
|
|
)
|
|
|
|
values = capture_inputs.values
|
|
|
|
with torch.inference_mode():
|
|
output = self.model.encoder_cudagraph_forward({**values}, path=path)
|
|
output_buffer = torch.empty_like(output)
|
|
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool):
|
|
output = self.model.encoder_cudagraph_forward({**values}, path=path)
|
|
output_buffer.copy_(output)
|
|
|
|
graph_set[token_budget] = BudgetGraphMetadata(
|
|
token_budget=token_budget,
|
|
max_batch_size=self.max_batch_size,
|
|
max_frames_per_batch=self.max_frames_per_batch,
|
|
graph=graph,
|
|
input_buffers=values,
|
|
output_buffer=output_buffer,
|
|
)
|
|
|
|
def _find_smallest_fitting_budget_given_tokens(
|
|
self, total_tokens: int, budgets: list[int] | None = None
|
|
) -> int | None:
|
|
"""Find smallest budget >= total_tokens.
|
|
|
|
Returns:
|
|
Token budget if found, None if no fitting budget.
|
|
"""
|
|
budgets = budgets if budgets is not None else self.token_budgets
|
|
for budget in budgets:
|
|
if budget >= total_tokens:
|
|
return budget
|
|
return None
|
|
|
|
def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:
|
|
"""Get item specs from the model."""
|
|
return self.model.get_encoder_cudagraph_item_specs(mm_kwargs)
|
|
|
|
def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:
|
|
"""Get per-item output token counts as plain ints."""
|
|
return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)]
|
|
|
|
@staticmethod
|
|
def _copy_padded_buffer(
|
|
dst: torch.Tensor,
|
|
src: torch.Tensor,
|
|
) -> None:
|
|
dst.zero_()
|
|
dst[: src.shape[0]].copy_(src)
|
|
|
|
def _run_budget_graph(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
token_budget: int,
|
|
path: str = "default",
|
|
) -> torch.Tensor | None:
|
|
"""Execute budget graph.
|
|
|
|
Args:
|
|
mm_kwargs: Multimodal inputs for the batch.
|
|
token_budget: Token budget to use.
|
|
path: Path for the graph. Should be one of ["default", "global", "local"].
|
|
Returns:
|
|
Encoder outputs, or None if graph not captured.
|
|
"""
|
|
graph_set = self._get_graph_set(path)
|
|
num_items = len(self._get_item_specs(mm_kwargs))
|
|
|
|
if token_budget not in graph_set:
|
|
self.graph_misses += num_items
|
|
return None
|
|
|
|
graph_meta = graph_set[token_budget]
|
|
|
|
replay = self.model.prepare_encoder_cudagraph_replay_buffers(
|
|
mm_kwargs,
|
|
self.max_batch_size,
|
|
self.max_frames_per_batch,
|
|
path,
|
|
)
|
|
|
|
# Copy replay buffers into graph input buffers. Iterate over the
|
|
# graph's own buffer keys (which may differ per path for dual-path
|
|
# models) rather than the global config.buffer_keys.
|
|
for key, buf in graph_meta.input_buffers.items():
|
|
src = replay.values.get(key)
|
|
if src is None:
|
|
continue
|
|
if src.ndim == 0:
|
|
buf.copy_(src)
|
|
else:
|
|
padding_logic = self.config.padding_logics.get(
|
|
key, self._copy_padded_buffer
|
|
)
|
|
padding_logic(buf, src)
|
|
|
|
graph_meta.graph.replay()
|
|
|
|
self.graph_hits += num_items
|
|
return graph_meta.output_buffer
|
|
|
|
def _execute_local(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
) -> list[torch.Tensor]:
|
|
"""Execute encoder on local inputs using greedy-packed CUDA graphs.
|
|
|
|
Sort images by output token count (smallest first), then greedily pack
|
|
as many images as possible into each batch while staying within
|
|
max_budget tokens and max_batch_size. Once a batch is finalised (next
|
|
image would overflow either constraint), find the smallest fitting
|
|
budget once for that batch.
|
|
|
|
For dual-path models (``enable_dual_path_graph=True``), two independent
|
|
graph sets are used: one for global images, one for local patches.
|
|
Budgets are found independently per path; if only one path fits, the
|
|
other falls back to eager via partial fallback.
|
|
|
|
By exchange argument, greedy smallest-first packing minimises eager
|
|
fallbacks -- any other ordering yields a higher token sum in some batch,
|
|
making that batch more likely to exceed the budget.
|
|
|
|
Stats note:
|
|
graph_hits -- counted inside _run_budget_graph after successful replay.
|
|
graph_misses -- counted here for single-image batches where the image
|
|
exceeds max_budget. Batches split due to max_batch_size
|
|
always satisfy total_tokens <= max_budget and therefore
|
|
always find a valid budget (no miss).
|
|
"""
|
|
if self.config.enable_dual_path_graph:
|
|
return self._execute_local_dual_path(mm_kwargs)
|
|
return self._execute_local_single_path(mm_kwargs)
|
|
|
|
def _execute_local_single_path(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
) -> list[torch.Tensor]:
|
|
"""Single-path greedy-packing execution (original behaviour)."""
|
|
item_specs = self._get_item_specs(mm_kwargs)
|
|
num_items = len(item_specs)
|
|
max_budget = self.token_budgets[-1]
|
|
|
|
per_item_out_tokens = [spec.output_tokens for spec in item_specs]
|
|
|
|
# Sort ascending by output token count (smallest first)
|
|
sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])
|
|
|
|
# Greedy pack against max_budget and max_batch_size.
|
|
# _find_smallest_fitting_budget_given_tokens is called once per
|
|
# finalised batch, not per image.
|
|
batches: list[tuple[list[int], int | None]] = []
|
|
current_batch: list[int] = []
|
|
current_batch_tokens = 0
|
|
|
|
for orig_idx in sorted_indices:
|
|
item_tokens = per_item_out_tokens[orig_idx]
|
|
if (
|
|
current_batch_tokens + item_tokens <= max_budget
|
|
and len(current_batch) < self.max_batch_size
|
|
):
|
|
current_batch.append(orig_idx)
|
|
current_batch_tokens += item_tokens
|
|
else:
|
|
if current_batch:
|
|
batches.append(
|
|
(
|
|
current_batch,
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_batch_tokens
|
|
),
|
|
)
|
|
)
|
|
current_batch = [orig_idx]
|
|
current_batch_tokens = item_tokens
|
|
|
|
if current_batch:
|
|
batches.append(
|
|
(
|
|
current_batch,
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_batch_tokens
|
|
),
|
|
)
|
|
)
|
|
|
|
# outputs_by_orig_idx maps each original image index to its output
|
|
# tensor. Needed because greedy packing reorders images; we restore
|
|
# the original order before returning.
|
|
outputs_by_orig_idx: dict[int, torch.Tensor] = {}
|
|
|
|
for batch_orig_indices, token_budget in batches:
|
|
batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
|
|
mm_kwargs, batch_orig_indices
|
|
)
|
|
batch_out_tokens = sum(per_item_out_tokens[i] for i in batch_orig_indices)
|
|
|
|
if token_budget is None:
|
|
# Single oversized image: item_tokens > max_budget.
|
|
# graph_misses counted here for this eager fallback.
|
|
logger.debug(
|
|
"Encoder CUDA graph fallback to eager: no budget for "
|
|
"%d tokens from %d images",
|
|
batch_out_tokens,
|
|
len(batch_orig_indices),
|
|
)
|
|
self.graph_misses += len(batch_orig_indices)
|
|
with torch.inference_mode():
|
|
raw = self.model.encoder_eager_forward(batch_mm_kwargs)
|
|
scatter_output_slices(
|
|
raw,
|
|
batch_orig_indices,
|
|
per_item_out_tokens,
|
|
outputs_by_orig_idx,
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"Encoder CUDA graph: batch_size=%d, tokens=%d, "
|
|
"budget=%d, waste=%.1f%%",
|
|
len(batch_orig_indices),
|
|
batch_out_tokens,
|
|
token_budget,
|
|
(token_budget - batch_out_tokens) / token_budget * 100,
|
|
)
|
|
|
|
# graph_hits counted inside _run_budget_graph after replay.
|
|
output = self._run_budget_graph(batch_mm_kwargs, token_budget)
|
|
assert output is not None
|
|
self.model.postprocess_encoder_output(
|
|
output,
|
|
batch_orig_indices,
|
|
per_item_out_tokens,
|
|
outputs_by_orig_idx,
|
|
clone=True,
|
|
batch_mm_kwargs=batch_mm_kwargs,
|
|
)
|
|
|
|
# Return in original batch order (caller maps outputs to token positions)
|
|
return [outputs_by_orig_idx[i] for i in range(num_items)]
|
|
|
|
def _execute_local_dual_path(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
) -> list[torch.Tensor]:
|
|
"""Dual-path greedy-packing execution.
|
|
|
|
Each image contributes both global tokens (constant per image)
|
|
and local tokens (patches * patch_tokens). Greedy packing
|
|
respects both budgets independently, then selects the smallest
|
|
fitting budget per path with partial eager fallback.
|
|
"""
|
|
item_specs = self._get_item_specs(mm_kwargs)
|
|
num_items = len(item_specs)
|
|
|
|
max_global_budget = self.global_token_budgets[-1]
|
|
max_local_budget = self.local_token_budgets[-1]
|
|
|
|
per_item_global_tokens = [spec.global_output_tokens for spec in item_specs]
|
|
per_item_local_tokens = [spec.local_output_tokens for spec in item_specs]
|
|
per_item_total_tokens = [spec.output_tokens for spec in item_specs]
|
|
|
|
# Sort ascending by total output tokens
|
|
sorted_indices = sorted(
|
|
range(num_items), key=lambda i: per_item_total_tokens[i]
|
|
)
|
|
|
|
# Each batch is a tuple of (indices, global_budget, local_budget).
|
|
batches: list[tuple[list[int], int | None, int | None]] = []
|
|
current_batch: list[int] = []
|
|
current_global_tokens = 0
|
|
current_local_tokens = 0
|
|
|
|
for orig_idx in sorted_indices:
|
|
global_token = per_item_global_tokens[orig_idx]
|
|
local_token = per_item_local_tokens[orig_idx]
|
|
if (
|
|
current_global_tokens + global_token <= max_global_budget
|
|
and current_local_tokens + local_token <= max_local_budget
|
|
and len(current_batch) < self.max_batch_size
|
|
):
|
|
current_batch.append(orig_idx)
|
|
current_global_tokens += global_token
|
|
current_local_tokens += local_token
|
|
else:
|
|
if current_batch:
|
|
batches.append(
|
|
(
|
|
current_batch,
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_global_tokens, self.global_token_budgets
|
|
),
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_local_tokens, self.local_token_budgets
|
|
),
|
|
)
|
|
)
|
|
current_batch = [orig_idx]
|
|
current_global_tokens = global_token
|
|
current_local_tokens = local_token
|
|
|
|
if current_batch:
|
|
batches.append(
|
|
(
|
|
current_batch,
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_global_tokens, self.global_token_budgets
|
|
),
|
|
self._find_smallest_fitting_budget_given_tokens(
|
|
current_local_tokens, self.local_token_budgets
|
|
),
|
|
)
|
|
)
|
|
|
|
outputs_by_orig_idx: dict[int, torch.Tensor] = {}
|
|
|
|
for batch_orig_indices, global_budget, local_budget in batches:
|
|
batch_mm_kwargs = self.model.select_encoder_cudagraph_items(
|
|
mm_kwargs, batch_orig_indices
|
|
)
|
|
batch_global_tokens = sum(
|
|
per_item_global_tokens[i] for i in batch_orig_indices
|
|
)
|
|
batch_local_tokens = sum(
|
|
per_item_local_tokens[i] for i in batch_orig_indices
|
|
)
|
|
|
|
both_eager = global_budget is None and local_budget is None
|
|
|
|
if both_eager:
|
|
logger.debug(
|
|
"Encoder CUDA graph dual-path full eager fallback: "
|
|
"%d global + %d local tokens from %d images",
|
|
batch_global_tokens,
|
|
batch_local_tokens,
|
|
len(batch_orig_indices),
|
|
)
|
|
self.graph_misses += len(batch_orig_indices)
|
|
with torch.inference_mode():
|
|
raw = self.model.encoder_eager_forward(batch_mm_kwargs)
|
|
per_item_total = [
|
|
per_item_global_tokens[i] + per_item_local_tokens[i] + 1
|
|
for i in batch_orig_indices
|
|
]
|
|
scatter_output_slices(
|
|
raw, batch_orig_indices, per_item_total, outputs_by_orig_idx
|
|
)
|
|
continue
|
|
|
|
logger.debug(
|
|
"Encoder CUDA graph dual-path: batch_size=%d, "
|
|
"global=%d (budget=%s), local=%d (budget=%s)",
|
|
len(batch_orig_indices),
|
|
batch_global_tokens,
|
|
global_budget,
|
|
batch_local_tokens,
|
|
local_budget,
|
|
)
|
|
|
|
# Execute global path: graph or eager fallback
|
|
if global_budget is not None:
|
|
global_output = self._run_budget_graph(
|
|
batch_mm_kwargs,
|
|
global_budget,
|
|
path="global",
|
|
)
|
|
assert global_output is not None
|
|
else:
|
|
with torch.inference_mode():
|
|
global_output = self.model.encoder_eager_forward(
|
|
batch_mm_kwargs, path="global"
|
|
)
|
|
|
|
# Execute local path: graph or eager fallback
|
|
if local_budget is not None and batch_local_tokens > 0:
|
|
local_output = self._run_budget_graph(
|
|
batch_mm_kwargs,
|
|
local_budget,
|
|
path="local",
|
|
)
|
|
assert local_output is not None
|
|
elif batch_local_tokens > 0:
|
|
with torch.inference_mode():
|
|
local_output = self.model.encoder_eager_forward(
|
|
batch_mm_kwargs, path="local"
|
|
)
|
|
else:
|
|
local_output = None
|
|
|
|
self.model.postprocess_encoder_output(
|
|
global_output,
|
|
batch_orig_indices,
|
|
per_item_global_tokens,
|
|
outputs_by_orig_idx,
|
|
clone=True,
|
|
batch_mm_kwargs=batch_mm_kwargs,
|
|
local_output=local_output,
|
|
)
|
|
|
|
return [outputs_by_orig_idx[i] for i in range(num_items)]
|
|
|
|
def _dp_shard(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
per_item_out_tokens: list[int],
|
|
) -> tuple[dict[str, Any], list[int], list[int], int]:
|
|
"""Distribute items across TP ranks for data-parallel execution.
|
|
|
|
Uses get_load_balance_assignment() to balance load by input size,
|
|
then select_encoder_cudagraph_items() to extract each rank's inputs.
|
|
|
|
Returns:
|
|
local_mm_kwargs: Inputs for this rank.
|
|
image_rank_assignment: Flattened assignment order across all ranks.
|
|
images_per_rank: Number of items per rank.
|
|
max_output_tokens_per_rank: Max output tokens across all ranks
|
|
(for padding during all_gather).
|
|
"""
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
current_rank = get_tensor_model_parallel_rank()
|
|
|
|
item_specs = self._get_item_specs(mm_kwargs)
|
|
per_item_input_sizes = [spec.input_size for spec in item_specs]
|
|
|
|
(image_rank_assignment, images_per_rank, input_patches_per_rank) = (
|
|
get_load_balance_assignment(per_item_input_sizes, tp_size)
|
|
)
|
|
|
|
# Extract local indices for this rank
|
|
cum_images_per_rank = [0]
|
|
for count in images_per_rank:
|
|
cum_images_per_rank.append(cum_images_per_rank[-1] + count)
|
|
|
|
local_indices = image_rank_assignment[
|
|
cum_images_per_rank[current_rank] : cum_images_per_rank[current_rank + 1]
|
|
]
|
|
|
|
if len(local_indices) > 0:
|
|
local_mm_kwargs = self.model.select_encoder_cudagraph_items(
|
|
mm_kwargs, local_indices
|
|
)
|
|
else:
|
|
local_mm_kwargs = self.model.select_encoder_cudagraph_items(mm_kwargs, [])
|
|
|
|
max_output_tokens_per_rank = (
|
|
max(
|
|
sum(
|
|
per_item_out_tokens[i]
|
|
for i in image_rank_assignment[
|
|
cum_images_per_rank[r] : cum_images_per_rank[r + 1]
|
|
]
|
|
)
|
|
for r in range(tp_size)
|
|
)
|
|
if len(per_item_out_tokens) > 0
|
|
else 0
|
|
)
|
|
|
|
return (
|
|
local_mm_kwargs,
|
|
image_rank_assignment,
|
|
images_per_rank,
|
|
max_output_tokens_per_rank,
|
|
)
|
|
|
|
def _dp_gather(
|
|
self,
|
|
local_outputs: list[torch.Tensor],
|
|
per_item_out_tokens: list[int],
|
|
image_rank_assignment: list[int],
|
|
images_per_rank: list[int],
|
|
max_output_tokens_per_rank: int,
|
|
) -> list[torch.Tensor]:
|
|
"""Gather outputs from all TP ranks and reorder to original sequence.
|
|
|
|
Assumes 2D output tensors [tokens, hidden]. Follows the same
|
|
pad -> all_gather -> unpad -> reorder algorithm as
|
|
run_dp_sharded_mrope_vision_model() in the eager path.
|
|
"""
|
|
hidden_size = self.config.out_hidden_size
|
|
tp_size = len(images_per_rank)
|
|
|
|
if len(local_outputs) > 0:
|
|
local_concat = torch.cat(local_outputs, dim=0)
|
|
else:
|
|
local_concat = torch.empty(
|
|
(0, hidden_size), device=self.device, dtype=self.dtype
|
|
)
|
|
|
|
# Pad to max_output_tokens_per_rank for all_gather
|
|
current_len = local_concat.shape[0]
|
|
if current_len < max_output_tokens_per_rank:
|
|
padding = torch.empty(
|
|
(max_output_tokens_per_rank - current_len, hidden_size),
|
|
dtype=self.dtype,
|
|
device=self.device,
|
|
)
|
|
local_padded = torch.cat([local_concat, padding], dim=0)
|
|
else:
|
|
local_padded = local_concat
|
|
|
|
gathered = tensor_model_parallel_all_gather(local_padded, dim=0)
|
|
|
|
# Unpad each rank's contribution
|
|
rank_outputs: list[torch.Tensor] = []
|
|
current_idx = 0
|
|
for rank in range(tp_size):
|
|
start = rank * max_output_tokens_per_rank
|
|
rank_count = images_per_rank[rank]
|
|
rank_indices = image_rank_assignment[current_idx : current_idx + rank_count]
|
|
rank_tokens = sum(per_item_out_tokens[i] for i in rank_indices)
|
|
current_idx += rank_count
|
|
rank_outputs.append(gathered[start : start + rank_tokens])
|
|
|
|
# Reorder to original sequence
|
|
total_items = len(per_item_out_tokens)
|
|
result: list[torch.Tensor | None] = [None] * total_items
|
|
current_idx = 0
|
|
for rank in range(tp_size):
|
|
count = images_per_rank[rank]
|
|
if count > 0:
|
|
rank_items = image_rank_assignment[current_idx : current_idx + count]
|
|
scatter_output_slices(
|
|
rank_outputs[rank],
|
|
rank_items,
|
|
per_item_out_tokens,
|
|
result,
|
|
)
|
|
current_idx += count
|
|
|
|
return [t for t in result if t is not None]
|
|
|
|
def execute(
|
|
self,
|
|
mm_kwargs: dict[str, Any],
|
|
) -> list[torch.Tensor]:
|
|
"""Execute encoder using CUDA graph with optional DP.
|
|
|
|
Args:
|
|
mm_kwargs: Multimodal keyword arguments containing the
|
|
input tensor and grid dimensions.
|
|
|
|
Returns:
|
|
List of encoder outputs (one per item).
|
|
"""
|
|
if self.use_dp:
|
|
per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)
|
|
|
|
(
|
|
local_mm_kwargs,
|
|
image_rank_assignment,
|
|
images_per_rank,
|
|
max_output_tokens_per_rank,
|
|
) = self._dp_shard(mm_kwargs, per_item_out_tokens)
|
|
|
|
local_outputs = self._execute_local(local_mm_kwargs)
|
|
|
|
result = self._dp_gather(
|
|
local_outputs,
|
|
per_item_out_tokens,
|
|
image_rank_assignment,
|
|
images_per_rank,
|
|
max_output_tokens_per_rank,
|
|
)
|
|
else:
|
|
result = self._execute_local(mm_kwargs)
|
|
|
|
# Log cumulative stats periodically
|
|
stats = self.get_cumulative_stats()
|
|
total_requests = self.graph_hits + self.graph_misses
|
|
if total_requests > 0 and total_requests % self.log_stats_interval == 0:
|
|
logger.debug(
|
|
"Encoder CUDA graph cumulative stats: "
|
|
"hits=%d, misses=%d, hit_rate=%.1f%%",
|
|
stats["graph_hits"],
|
|
stats["graph_misses"],
|
|
stats["hit_rate"] * 100,
|
|
)
|
|
|
|
return result
|
|
|
|
def get_cumulative_stats(self) -> dict[str, Any]:
|
|
"""Get cumulative CUDA graph statistics."""
|
|
total_requests = self.graph_hits + self.graph_misses
|
|
hit_rate = self.graph_hits / total_requests if total_requests > 0 else 0.0
|
|
|
|
num_budgets = sum(len(g) for g in self.budget_graphs.values())
|
|
|
|
return {
|
|
"graph_hits": self.graph_hits,
|
|
"graph_misses": self.graph_misses,
|
|
"hit_rate": hit_rate,
|
|
"num_budgets": num_budgets,
|
|
"token_budgets": self.token_budgets,
|
|
}
|