wehub-resource-sync
635 lines
26 KiB
Python
635 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from collections import defaultdict
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from itertools import product
|
|
from typing import Any, NamedTuple, Protocol
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from tqdm import tqdm
|
|
|
|
from vllm.compilation.breakable_cudagraph import (
|
|
BreakableCUDAGraphWrapper,
|
|
is_breakable_cudagraph_enabled,
|
|
)
|
|
from vllm.compilation.counter import compilation_counter
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.distributed.parallel_state import (
|
|
get_pp_group,
|
|
graph_capture,
|
|
is_global_first_rank,
|
|
)
|
|
from vllm.forward_context import BatchDescriptor, set_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.offloader.base import get_offloader
|
|
from vllm.platforms import current_platform
|
|
from vllm.sequence import IntermediateTensors
|
|
from vllm.utils.math_utils import round_up
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
|
|
from vllm.v1.worker.gpu.block_table import BlockTables
|
|
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
|
|
from vllm.v1.worker.gpu.model_states.interface import ModelState
|
|
from vllm.v1.worker.utils import AttentionGroup
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class AttentionState(NamedTuple):
|
|
attn_metadata: dict[str, Any] | None
|
|
slot_mappings: dict[str, torch.Tensor]
|
|
|
|
|
|
class AttentionStatePair(NamedTuple):
|
|
warmup: AttentionState
|
|
captured: AttentionState
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BatchExecutionDescriptor:
|
|
"""Describes the shape of the batch and CG mode to run; this is used to make shape
|
|
matches between the capture and runtime."""
|
|
|
|
cg_mode: CUDAGraphMode
|
|
num_tokens: int
|
|
num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs)
|
|
uniform_token_count: int | None = None
|
|
num_active_loras: int = 0
|
|
|
|
|
|
class CreateForwardFn(Protocol):
|
|
"""Factory that prepares inputs (OUTSIDE the graph) and returns a tuple of
|
|
(forward_fn, attn_state). Called with warmup=True for the warmup pass and
|
|
warmup=False for the captured pass."""
|
|
|
|
def __call__(
|
|
self,
|
|
desc: BatchExecutionDescriptor,
|
|
warmup: bool,
|
|
) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]: ...
|
|
|
|
|
|
def _is_compatible(
|
|
desc: BatchExecutionDescriptor,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
uniform_token_count: int | None,
|
|
num_active_loras: int,
|
|
) -> bool:
|
|
# desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
|
|
# desc.num_reqs=None means no request padding needed (PIECEWISE)
|
|
return (
|
|
(
|
|
desc.uniform_token_count is None
|
|
or desc.uniform_token_count == uniform_token_count
|
|
)
|
|
and (desc.num_reqs is None or desc.num_reqs >= num_reqs)
|
|
and desc.num_tokens >= num_tokens
|
|
and desc.num_active_loras == num_active_loras
|
|
)
|
|
|
|
|
|
def get_uniform_token_count(
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
max_query_len: int,
|
|
) -> int | None:
|
|
"""
|
|
Return the uniform token count if batch is uniform, else None.
|
|
A batch is uniform if all requests have the same number of tokens.
|
|
"""
|
|
if (max_query_len == num_tokens // num_reqs) and (
|
|
num_tokens == max_query_len * num_reqs
|
|
):
|
|
return max_query_len
|
|
return None
|
|
|
|
|
|
class CudaGraphManager:
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
cudagraph_mode: CUDAGraphMode,
|
|
decode_query_len: int,
|
|
lora_capture_cases: list[int] | None = None,
|
|
):
|
|
self.vllm_config = vllm_config
|
|
self.device = device
|
|
self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
|
|
self.compilation_config = vllm_config.compilation_config
|
|
assert self.compilation_config is not None
|
|
self.cudagraph_mode = cudagraph_mode
|
|
self.decode_query_len = decode_query_len
|
|
|
|
self.dp_size = vllm_config.parallel_config.data_parallel_size
|
|
self.tp_size = vllm_config.parallel_config.tensor_parallel_size
|
|
self.is_first_pp_rank = get_pp_group().is_first_rank
|
|
self.is_last_pp_rank = get_pp_group().is_last_rank
|
|
self.lora_capture_cases = lora_capture_cases or [0]
|
|
# Precompute actual num_active_loras -> captured case mapping so that
|
|
# dispatch() is a plain dict lookup instead of a per-call bisect.
|
|
self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()
|
|
|
|
self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
|
|
self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
|
|
|
|
self._graphs_captured = False
|
|
|
|
self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
|
|
self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
|
|
|
|
self._init_candidates()
|
|
|
|
# Breakable CUDA graph (PW CUDA graph without torch.compile)
|
|
self.use_breakable_cg = (
|
|
is_breakable_cudagraph_enabled()
|
|
and self.cudagraph_mode.has_piecewise_cudagraphs()
|
|
)
|
|
self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
|
|
|
|
def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:
|
|
"""Precompute actual num_active_loras -> effective captured case.
|
|
|
|
Mirrors the num_tokens candidate expansion in ``_init_candidates``:
|
|
every possible active-LoRA count is mapped ahead of time to the
|
|
smallest captured case that can serve it, so ``dispatch`` is a plain
|
|
dict lookup instead of a per-call bisect.
|
|
"""
|
|
captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0)
|
|
if not captured_with_lora:
|
|
return {}, 0
|
|
dispatch_map: dict[int, int] = {}
|
|
case_idx = 0
|
|
for n in range(1, captured_with_lora[-1] + 1):
|
|
while captured_with_lora[case_idx] < n:
|
|
case_idx += 1
|
|
dispatch_map[n] = captured_with_lora[case_idx]
|
|
return dispatch_map, captured_with_lora[-1]
|
|
|
|
def _resolve_effective_loras(self, num_active_loras: int) -> int:
|
|
"""Map an actual active-LoRA count to its captured graph case."""
|
|
if num_active_loras <= 0 or not self._lora_dispatch_map:
|
|
return num_active_loras
|
|
# Counts above the largest captured case clamp to it.
|
|
return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case)
|
|
|
|
def _init_candidates(self) -> None:
|
|
"""Build priority-ordered candidate lists for each token count."""
|
|
capture_sizes = self.compilation_config.cudagraph_capture_sizes
|
|
if not (self.cudagraph_mode and capture_sizes):
|
|
return
|
|
|
|
capture_sizes = sorted(capture_sizes)
|
|
max_decode_tokens = self.max_num_reqs * self.decode_query_len
|
|
decode_mode = self.cudagraph_mode.decode_mode()
|
|
mixed_mode = self.cudagraph_mode.mixed_mode()
|
|
separate_decode_routine = self.cudagraph_mode.separate_routine()
|
|
max_cg_capture_size = self.compilation_config.max_cudagraph_capture_size
|
|
|
|
descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = (
|
|
defaultdict(list)
|
|
)
|
|
descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = (
|
|
defaultdict(list)
|
|
)
|
|
|
|
# When using Dynamic SD, num_speculative_tokens is the max number of
|
|
# draft tokens. The scheduler might use a smaller number so we need
|
|
# to capture graphs for all possible values during decode.
|
|
speculative_config = self.vllm_config.speculative_config
|
|
if (
|
|
speculative_config
|
|
and speculative_config.uses_dynamic_speculative_decoding()
|
|
):
|
|
num_spec_per_batch_size = (
|
|
speculative_config.num_speculative_tokens_per_batch_size
|
|
)
|
|
# uses_dynamic_speculative_decoding() guarantees this is set.
|
|
assert num_spec_per_batch_size is not None
|
|
# decode_query_len = num_speculative_steps + num_new_sampled_tokens
|
|
# _per_step. Recover num_new_sampled_tokens_per_step
|
|
# from the values the manager already has.
|
|
num_new_sampled_tokens_per_step = (
|
|
self.decode_query_len - self.vllm_config.num_speculative_tokens
|
|
)
|
|
# Each entry is (range_start, range_end, num_speculative_tokens).
|
|
decode_query_lens = [
|
|
x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size
|
|
]
|
|
else:
|
|
decode_query_lens = [self.decode_query_len]
|
|
|
|
for num_tokens, num_active_loras in product(
|
|
capture_sizes, self.lora_capture_cases
|
|
):
|
|
# Capture uniform decode specfifc graphs if required
|
|
# (i.e. separate decode routine)
|
|
if separate_decode_routine and decode_mode:
|
|
for decode_query_len in decode_query_lens:
|
|
rounded_num_tokens = round_up(num_tokens, decode_query_len)
|
|
rounded_num_reqs = rounded_num_tokens // decode_query_len
|
|
|
|
if (
|
|
rounded_num_tokens > max_decode_tokens
|
|
or rounded_num_tokens > max_cg_capture_size
|
|
or rounded_num_reqs > self.max_num_reqs
|
|
):
|
|
continue
|
|
|
|
desc = BatchExecutionDescriptor(
|
|
cg_mode=decode_mode,
|
|
num_tokens=rounded_num_tokens,
|
|
num_reqs=rounded_num_reqs,
|
|
uniform_token_count=decode_query_len,
|
|
num_active_loras=num_active_loras,
|
|
)
|
|
|
|
# avoid duplicate graphs
|
|
if desc not in descs_by_mode[decode_mode]:
|
|
descs_by_mode[decode_mode].append(desc)
|
|
descs_by_token_lora[
|
|
(rounded_num_tokens, num_active_loras)
|
|
].append(desc)
|
|
|
|
if mixed_mode:
|
|
# for PIECEWISE graphs there is no limit on requests when replaying
|
|
# i.e. no request padding is needed
|
|
# so we leave it as None
|
|
num_reqs = (
|
|
min(num_tokens, self.max_num_reqs)
|
|
if mixed_mode == CUDAGraphMode.FULL
|
|
else None
|
|
)
|
|
desc = BatchExecutionDescriptor(
|
|
cg_mode=mixed_mode,
|
|
num_tokens=num_tokens,
|
|
num_reqs=num_reqs,
|
|
num_active_loras=num_active_loras,
|
|
)
|
|
descs_by_mode[mixed_mode].append(desc)
|
|
descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
|
|
|
|
if not descs_by_token_lora:
|
|
return
|
|
|
|
all_token_counts = sorted({k[0] for k in descs_by_token_lora})
|
|
current_range_start = 0
|
|
for token_cg_size in all_token_counts:
|
|
for i in range(current_range_start, token_cg_size + 1):
|
|
for num_active_loras in self.lora_capture_cases:
|
|
staging_key = (token_cg_size, num_active_loras)
|
|
if staging_key in descs_by_token_lora:
|
|
self._candidates[(i, num_active_loras)] = descs_by_token_lora[
|
|
staging_key
|
|
]
|
|
current_range_start = token_cg_size + 1
|
|
|
|
for mode, descs in descs_by_mode.items():
|
|
descs.sort(key=lambda d: d.num_tokens, reverse=True)
|
|
self._capture_descs[mode] = descs
|
|
|
|
def needs_capture(self) -> bool:
|
|
return len(self._capture_descs) > 0
|
|
|
|
@torch.inference_mode()
|
|
def capture(
|
|
self,
|
|
create_forward_fn: CreateForwardFn,
|
|
progress_bar_desc: str = "Capturing CUDA graphs",
|
|
) -> dict[BatchExecutionDescriptor, AttentionStatePair]:
|
|
"""Capture CUDA graphs.
|
|
|
|
Args:
|
|
create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
|
|
returns a tuple of (forward_fn, attn_state). For FULL cudagraph
|
|
mode, it is invoked once with warmup=True for the warmup pass,
|
|
and again with warmup=False for the captured pass. For attention
|
|
backends that perform lazy metadata initialization (e.g. FlashMLA),
|
|
FULL cudagraph capture requires distinct metadatas for warmup and
|
|
capture.
|
|
"""
|
|
attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {}
|
|
with graph_capture(device=self.device):
|
|
# Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
|
|
# activations so FULL activations should fit in already allocated
|
|
# buffers in the graph pool.
|
|
for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]:
|
|
if mode not in self._capture_descs:
|
|
continue
|
|
|
|
descs = self._capture_descs[mode]
|
|
if is_global_first_rank():
|
|
descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})")
|
|
for desc in descs:
|
|
# Prepare inputs and get forward function
|
|
forward_fn, warmup_attn_state = create_forward_fn(desc, warmup=True)
|
|
|
|
# Warmup
|
|
forward_fn(CUDAGraphMode.NONE)
|
|
|
|
# Capture
|
|
logger.debug(
|
|
"CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
|
|
)
|
|
if desc.cg_mode == CUDAGraphMode.PIECEWISE:
|
|
attn_states[desc] = AttentionStatePair(
|
|
warmup_attn_state, warmup_attn_state
|
|
)
|
|
forward_fn(CUDAGraphMode.PIECEWISE)
|
|
else:
|
|
# Capture with fresh attention state.
|
|
forward_fn, capture_attn_state = create_forward_fn(
|
|
desc, warmup=False
|
|
)
|
|
attn_states[desc] = AttentionStatePair(
|
|
warmup_attn_state, capture_attn_state
|
|
)
|
|
assert desc not in self.graphs, (
|
|
f"Graph already captured for {desc}"
|
|
)
|
|
graph = torch.cuda.CUDAGraph()
|
|
# Sync offloader's copy stream before capture.
|
|
# Ensure any pre-capture prefetches from offloader are complete.
|
|
get_offloader().sync_prev_onload()
|
|
with torch.cuda.graph(graph, self.pool):
|
|
forward_fn(CUDAGraphMode.NONE)
|
|
# Join offloader's copy stream after forward to avoid
|
|
# unjoined stream error. The last layer's start_prefetch
|
|
# forks copy_stream, but wait_prefetch only happens in
|
|
# the next forward pass.
|
|
get_offloader().join_after_forward()
|
|
self.graphs[desc] = graph
|
|
compilation_counter.num_cudagraph_captured += 1
|
|
self._graphs_captured = True
|
|
return attn_states
|
|
|
|
def dispatch(
|
|
self,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
uniform_token_count: int | None,
|
|
num_active_loras: int,
|
|
) -> BatchExecutionDescriptor:
|
|
"""Find matching cudagraph descriptor from priority-ordered candidates."""
|
|
|
|
effective_loras = self._resolve_effective_loras(num_active_loras)
|
|
key = (num_tokens, effective_loras)
|
|
if self._graphs_captured and num_tokens > 0 and key in self._candidates:
|
|
for desc in self._candidates[key]:
|
|
if _is_compatible(
|
|
desc,
|
|
num_reqs,
|
|
num_tokens,
|
|
uniform_token_count,
|
|
effective_loras,
|
|
):
|
|
return desc
|
|
return BatchExecutionDescriptor(
|
|
cg_mode=CUDAGraphMode.NONE,
|
|
num_tokens=num_tokens,
|
|
num_reqs=num_reqs,
|
|
num_active_loras=effective_loras,
|
|
)
|
|
|
|
def run_fullgraph(self, desc: BatchExecutionDescriptor):
|
|
"""Replay a captured FULL cudagraph."""
|
|
assert desc.cg_mode == CUDAGraphMode.FULL, (
|
|
f"Expected FULL mode, got {desc.cg_mode}"
|
|
)
|
|
assert desc in self.graphs, f"No cudagraph for {desc}"
|
|
# Sync offloader before replay - needed when transitioning from
|
|
# eager/piecewise to full cudagraph (e.g., prefill → decode).
|
|
# The previous eager iteration's start_prefetch may have queued
|
|
# H2D copies on copy_stream that the graph's captured events
|
|
# cannot see. Without this, replay could overwrite static buffers
|
|
# while those copies are still in flight.
|
|
get_offloader().sync_prev_onload()
|
|
self.graphs[desc].replay()
|
|
|
|
def init_breakable_cg_runner(self, model: nn.Module) -> None:
|
|
if self.breakable_cg_runner is None:
|
|
self.breakable_cg_runner = BreakableCUDAGraphWrapper(
|
|
model, self.vllm_config
|
|
)
|
|
|
|
def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any:
|
|
if not self.use_breakable_cg:
|
|
# Default: Use torch-compiled piecewise cudagraph.
|
|
return model(**model_inputs)
|
|
assert self.breakable_cg_runner is not None
|
|
return self.breakable_cg_runner(**model_inputs)
|
|
|
|
|
|
class ModelCudaGraphManager(CudaGraphManager):
|
|
"""CudaGraphManager with model-specific capture and hidden state management."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
cudagraph_mode: CUDAGraphMode,
|
|
decode_query_len: int,
|
|
lora_capture_cases: list[int] | None = None,
|
|
):
|
|
super().__init__(
|
|
vllm_config,
|
|
device,
|
|
cudagraph_mode,
|
|
decode_query_len,
|
|
lora_capture_cases=lora_capture_cases,
|
|
)
|
|
self.hidden_states: torch.Tensor | None = None
|
|
self.aux_hidden_states: list[torch.Tensor] = []
|
|
self.use_aux_hidden_state_outputs = False
|
|
self.intermediate_tensors: IntermediateTensors | None = None
|
|
|
|
def capture(
|
|
self,
|
|
model: nn.Module,
|
|
model_state: ModelState,
|
|
input_buffers: InputBuffers,
|
|
intermediate_tensors: IntermediateTensors | None,
|
|
block_tables: BlockTables,
|
|
attn_groups: list[list[AttentionGroup]],
|
|
kv_cache_config: KVCacheConfig,
|
|
has_lora: bool = False,
|
|
use_aux_hidden_state_outputs: bool = False,
|
|
lora_capture_hook: Callable[[int, int, int], None] | None = None,
|
|
progress_bar_desc: str = "Capturing CUDA graphs",
|
|
) -> dict[BatchExecutionDescriptor, AttentionStatePair]:
|
|
"""Capture CUDA graphs for model forward pass."""
|
|
self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
|
|
if self.use_breakable_cg:
|
|
self.init_breakable_cg_runner(model)
|
|
|
|
def create_forward_fn(
|
|
desc: BatchExecutionDescriptor,
|
|
warmup: bool,
|
|
) -> tuple[
|
|
Callable[[CUDAGraphMode], None],
|
|
AttentionState,
|
|
]:
|
|
num_tokens = desc.num_tokens
|
|
num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
|
|
|
|
# Set LoRA state before capture so kernels see correct adapters.
|
|
if lora_capture_hook is not None:
|
|
lora_capture_hook(desc.num_active_loras, num_reqs, num_tokens)
|
|
|
|
num_tokens_across_dp = (
|
|
torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
|
|
if self.dp_size > 1
|
|
else None
|
|
)
|
|
|
|
model_inputs = {
|
|
"input_ids": input_buffers.input_ids[:num_tokens],
|
|
"positions": input_buffers.positions[:num_tokens],
|
|
**model_state.prepare_dummy_inputs(num_reqs, num_tokens),
|
|
}
|
|
if not self.is_first_pp_rank:
|
|
# Update for non-first PP ranks.
|
|
model_inputs["input_ids"] = None
|
|
model_inputs["inputs_embeds"] = None
|
|
assert intermediate_tensors is not None
|
|
model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]
|
|
|
|
attn_metadata, slot_mappings = prepare_inputs_to_capture(
|
|
num_reqs,
|
|
num_tokens,
|
|
model_state,
|
|
input_buffers,
|
|
block_tables,
|
|
attn_groups,
|
|
kv_cache_config,
|
|
skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
|
|
)
|
|
|
|
# Capture with dummy rows marked as padding.
|
|
input_buffers.is_padding.fill_(True)
|
|
|
|
def forward_fn(cg_mode: CUDAGraphMode) -> None:
|
|
batch_descriptor = None
|
|
if cg_mode == CUDAGraphMode.PIECEWISE:
|
|
assert attn_metadata is None
|
|
batch_descriptor = BatchDescriptor(
|
|
num_tokens=num_tokens,
|
|
has_lora=has_lora,
|
|
num_active_loras=desc.num_active_loras,
|
|
)
|
|
with set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens,
|
|
cudagraph_runtime_mode=cg_mode,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
slot_mapping=slot_mappings,
|
|
batch_descriptor=batch_descriptor,
|
|
is_padding=input_buffers.is_padding[:num_tokens],
|
|
):
|
|
if cg_mode == CUDAGraphMode.PIECEWISE:
|
|
# PIECEWISE graph (compiled PW or breakable, chosen inside
|
|
# run_pw_graph).
|
|
model_output = self.run_pw_graph(model, model_inputs)
|
|
else:
|
|
model_output = model(**model_inputs)
|
|
|
|
if cg_mode == CUDAGraphMode.PIECEWISE:
|
|
# PW CUDA graph (compiled or breakable) internally handles the
|
|
# model outputs. No need to keep track of the hidden states.
|
|
return None
|
|
|
|
if self.is_last_pp_rank:
|
|
# Last PP rank (common case).
|
|
if self.use_aux_hidden_state_outputs:
|
|
hidden_states, aux_hidden_states = model_output
|
|
else:
|
|
hidden_states = model_output
|
|
aux_hidden_states = []
|
|
if self.hidden_states is None:
|
|
self.hidden_states = torch.empty_like(hidden_states)
|
|
self.hidden_states[:num_tokens] = hidden_states
|
|
if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
|
|
self.aux_hidden_states = [
|
|
torch.empty_like(x) for x in aux_hidden_states
|
|
]
|
|
for i, aux in enumerate(aux_hidden_states):
|
|
self.aux_hidden_states[i][:num_tokens] = aux
|
|
else:
|
|
# Non-last PP rank.
|
|
assert isinstance(model_output, IntermediateTensors)
|
|
intermediate_tensors = model_output
|
|
if self.intermediate_tensors is None:
|
|
self.intermediate_tensors = IntermediateTensors.empty_like(
|
|
intermediate_tensors
|
|
)
|
|
for k, v in intermediate_tensors.tensors.items():
|
|
self.intermediate_tensors[k][:num_tokens] = v
|
|
|
|
return forward_fn, AttentionState(attn_metadata, slot_mappings)
|
|
|
|
return super().capture(create_forward_fn, progress_bar_desc)
|
|
|
|
def run_fullgraph(
|
|
self, desc: BatchExecutionDescriptor
|
|
) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors:
|
|
"""Replay a captured FULL cudagraph and return hidden states."""
|
|
super().run_fullgraph(desc)
|
|
if not self.is_last_pp_rank:
|
|
assert self.intermediate_tensors is not None
|
|
return self.intermediate_tensors[: desc.num_tokens]
|
|
|
|
assert self.hidden_states is not None
|
|
hidden_states = self.hidden_states[: desc.num_tokens]
|
|
if not self.use_aux_hidden_state_outputs:
|
|
return hidden_states
|
|
return hidden_states, [x[: desc.num_tokens] for x in self.aux_hidden_states]
|
|
|
|
|
|
def prepare_inputs_to_capture(
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
model_state: ModelState,
|
|
input_buffers: InputBuffers,
|
|
block_tables: BlockTables,
|
|
attn_groups: list[list[AttentionGroup]],
|
|
kv_cache_config: KVCacheConfig,
|
|
skip_attn: bool = False,
|
|
) -> AttentionState:
|
|
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
|
|
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
|
|
slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
|
|
slot_mappings_by_layer = build_slot_mappings_by_layer(
|
|
slot_mappings, kv_cache_config
|
|
)
|
|
|
|
# HACK(woosuk): Special handling for DCP.
|
|
if block_tables.cp_size > 1:
|
|
prepare_dcp_local_seq_lens(
|
|
input_buffers.dcp_local_seq_lens,
|
|
input_batch.seq_lens,
|
|
num_reqs,
|
|
block_tables.cp_size,
|
|
block_tables.cp_rank,
|
|
block_tables.cp_interleave,
|
|
)
|
|
input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
|
|
|
|
attn_metadata = None
|
|
if not skip_attn:
|
|
attn_metadata = model_state.prepare_attn(
|
|
input_batch,
|
|
CUDAGraphMode.NONE,
|
|
input_block_tables,
|
|
slot_mappings,
|
|
attn_groups,
|
|
kv_cache_config,
|
|
for_capture=True,
|
|
)
|
|
return AttentionState(attn_metadata, slot_mappings_by_layer)
|