wehub-resource-sync
1289 lines
50 KiB
Python
1289 lines
50 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""LMCache-driven KV cache transfer operations for the MPCacheServer."""
|
|
|
|
# Standard
|
|
from dataclasses import dataclass
|
|
from itertools import islice
|
|
from typing import Generator, Sequence
|
|
import threading
|
|
import time
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache import torch_dev, torch_device_type
|
|
from lmcache.logging import init_logger
|
|
from lmcache.utils import (
|
|
EngineType,
|
|
_lmcache_nvtx_annotate,
|
|
check_interprocess_event_support,
|
|
)
|
|
from lmcache.v1.distributed.api import (
|
|
MemoryLayoutDesc,
|
|
ObjectKey,
|
|
)
|
|
from lmcache.v1.gpu_connector.gpu_ops import (
|
|
build_staging_copies,
|
|
lmcache_memcpy_async_d2h,
|
|
lmcache_memcpy_async_h2d,
|
|
)
|
|
from lmcache.v1.gpu_connector.utils import LayoutHints
|
|
from lmcache.v1.memory_allocators.lazy_memory_allocator import LazyMemoryAllocator
|
|
from lmcache.v1.memory_management import GDSMemoryObject, MemoryObj
|
|
from lmcache.v1.mp_observability.event import Event, EventType
|
|
from lmcache.v1.multiprocess.custom_types import (
|
|
IPCCacheServerKey,
|
|
KVCache,
|
|
)
|
|
from lmcache.v1.multiprocess.engine_context import MPCacheServerContext
|
|
from lmcache.v1.multiprocess.engine_module import (
|
|
HandlerSpec,
|
|
InstanceLivenessTarget,
|
|
ThreadPoolType,
|
|
)
|
|
from lmcache.v1.multiprocess.group_view import EngineGroupInfo
|
|
from lmcache.v1.multiprocess.native_completion import (
|
|
DeviceHostFuncDispatcher,
|
|
submit_callback_to_stream,
|
|
)
|
|
from lmcache.v1.multiprocess.protocols.base import RequestType
|
|
from lmcache.v1.platform.base_cache_context import BaseCacheContext
|
|
from lmcache.v1.platform.cache_context import create_cache_context
|
|
import lmcache.c_ops as lmc_ops
|
|
import lmcache.python_ops_fallback as _python_ops_fallback
|
|
|
|
logger = init_logger(__name__)
|
|
_HAS_NATIVE_OBJECT_GROUP_TRANSFER: bool = (
|
|
lmc_ops.execute_object_group_transfer
|
|
is not _python_ops_fallback.execute_object_group_transfer
|
|
)
|
|
|
|
|
|
def get_layout_desc(
|
|
cache_context: BaseCacheContext,
|
|
num_tokens: int,
|
|
object_group_id: int,
|
|
) -> MemoryLayoutDesc:
|
|
"""Get the memory layout description for a specific object group.
|
|
|
|
The returned layout describes the single memory object that backs
|
|
``object_group_id``: one (shape, dtype) entry per kernel group in that
|
|
object group, in the kernel groups' declared layout order. Kernel groups
|
|
may have different shapes and dtypes.
|
|
|
|
Args:
|
|
cache_context: The cache context containing the KV cache information.
|
|
num_tokens: The number of tokens to determine the layout for.
|
|
object_group_id: Index of the object group whose layout to build.
|
|
|
|
Returns:
|
|
MemoryLayoutDesc: The memory layout description containing shapes and
|
|
dtypes, one entry per kernel group in the object group.
|
|
"""
|
|
object_group = cache_context.kv_layer_groups_manager.object_groups[object_group_id]
|
|
shapes_and_dtypes = [
|
|
cache_context.get_kernel_group_shape_dtype(num_tokens, kernel_group_idx)
|
|
for kernel_group_idx in object_group.kernel_group_indices
|
|
]
|
|
shapes, dtypes = zip(*shapes_and_dtypes, strict=False)
|
|
return MemoryLayoutDesc(shapes=list(shapes), dtypes=list(dtypes))
|
|
|
|
|
|
def batched_iteration_with_skip(
|
|
lst: Sequence,
|
|
batch_size: int,
|
|
skip_count: int,
|
|
) -> Generator[tuple[int, tuple], None, None]:
|
|
"""Utility function to iterate over a list in batches with an initial skip.
|
|
|
|
Args:
|
|
lst: The list to iterate over.
|
|
batch_size: The size of each batch.
|
|
skip_count: The number of items to skip at the start of the list.
|
|
|
|
Yields:
|
|
Tuples of (batch_start_idx, batch) where batch is a tuple of items
|
|
from the list, and batch_start_idx is the "original" index of the first
|
|
item in the batch.
|
|
|
|
Raises:
|
|
ValueError: If batch_size is less than 1 or skip_count is negative.
|
|
|
|
Note:
|
|
Batch_idx is the index of the batch in the original list, accounting
|
|
for the skipped items. For example, if skip_count is 10 and batch_size
|
|
is 5, the first yielded batch will have batch_start_idx=10.
|
|
"""
|
|
if batch_size < 1:
|
|
raise ValueError("batch size must be at least one")
|
|
if skip_count < 0:
|
|
raise ValueError("skip_count must be non-negative")
|
|
|
|
it = iter(lst)
|
|
# Skip the initial items
|
|
for _ in range(skip_count):
|
|
next(it, None)
|
|
batch_start_idx = skip_count
|
|
while batch := tuple(islice(it, batch_size)):
|
|
yield batch_start_idx, batch
|
|
batch_start_idx += len(batch)
|
|
|
|
|
|
def downsample_and_stage_block_ids(
|
|
cache_context: BaseCacheContext,
|
|
block_ids: list[list[int]],
|
|
) -> list[torch.Tensor]:
|
|
"""Cut the block id lists to skip the unneeded blocks in a chunk and
|
|
stage it into GPU tensors for later use.
|
|
|
|
This mainly targets the case where a portion of the blocks are not
|
|
needed for every chunk, such as deepseek v4's swa cache.
|
|
|
|
Note that the we do NOT do any object-level skipping here.
|
|
|
|
Args:
|
|
cache_context: The cache context containing the KV cache information.
|
|
block_ids: The original block id lists, indexed by LMCache KV group index.
|
|
|
|
Returns:
|
|
The cut block id lists, indexed by LMCache KV group index.
|
|
|
|
Note:
|
|
This function has some coupled logic with transfer_kv_per_object_group below.
|
|
The caller need to make sure that the block ids seen by
|
|
transfer_kv_per_object_group are produced by this function.
|
|
|
|
Example:
|
|
If a model have 2 kernel groups, one is full attention with block size 32,
|
|
one is swa attention with block size 32 and sliding window size 64, and
|
|
LMCache has a chunk size of 128. And there are 2 chunks in total (256 tokens).
|
|
|
|
The input will be:
|
|
[
|
|
[1, 2, 3, 4, 5, 6, 7, 8], # block ids for the full attention group
|
|
[11, 12, 13, 14, 15, 16, 17, 18], # block ids for the swa attention group
|
|
]
|
|
|
|
The output will be
|
|
[
|
|
[1, 2, 3, 4, 5, 6, 7, 8], # full attention group still needs all block ids
|
|
[13, 14, 17, 18], # swa attention group only needs the last 2 block per chunk
|
|
]
|
|
"""
|
|
num_kernel_groups = cache_context.kv_layer_groups_manager.num_kernel_groups
|
|
for kernel_group_id in range(num_kernel_groups):
|
|
subchunk_sw_size_tokens = (
|
|
cache_context.kv_layer_groups_manager.get_subchunk_sw_size_tokens(
|
|
kernel_group_id
|
|
)
|
|
)
|
|
tokens_per_chunk = min(
|
|
cache_context.lmcache_tokens_per_chunk, subchunk_sw_size_tokens
|
|
)
|
|
keep_blocks_per_chunk = cache_context.calculate_num_blocks(
|
|
tokens_per_chunk, kernel_group_id
|
|
)
|
|
total_blocks_per_chunk = cache_context.calculate_num_blocks(
|
|
cache_context.lmcache_tokens_per_chunk, kernel_group_id
|
|
)
|
|
|
|
new_block_ids = []
|
|
old_block_ids = block_ids[kernel_group_id]
|
|
assert len(old_block_ids) % total_blocks_per_chunk == 0, (
|
|
f"len(block_ids[{kernel_group_id}]) should be a multiple "
|
|
f"of total_blocks_per_chunk ({total_blocks_per_chunk}), but got "
|
|
f"{len(old_block_ids)}"
|
|
)
|
|
|
|
for i in range(0, len(old_block_ids), total_blocks_per_chunk):
|
|
chunk_block_ids = old_block_ids[i : i + total_blocks_per_chunk]
|
|
new_block_ids.extend(chunk_block_ids[-keep_blocks_per_chunk:])
|
|
|
|
block_ids[kernel_group_id] = new_block_ids
|
|
|
|
# Stage the cut block ids into GPU tensors
|
|
block_ids_gpu = cache_context.stage_block_ids(block_ids)
|
|
return block_ids_gpu
|
|
|
|
|
|
def _recalculate_blocks_to_skip(
|
|
blocks_per_chunk: int,
|
|
blocks_per_window: int,
|
|
blocks_to_skip: int,
|
|
) -> int:
|
|
"""Re-calculate the number of blocks to skip for a batch of chunks based
|
|
on the blocks per chunk and blocks per sliding window WHEN the window
|
|
size is smaller than the lmcache chunk size.
|
|
|
|
Args:
|
|
blocks_per_chunk: The total number of blocks in one chunk for the
|
|
current group.
|
|
blocks_per_window: The number of blocks in the sliding window
|
|
for the current group. Should be less than or equal to
|
|
blocks_per_chunk.
|
|
blocks_to_skip: The number of blocks to skip.
|
|
|
|
Returns:
|
|
The re-calculated number of blocks to skip for the current batch of
|
|
chunks.
|
|
"""
|
|
if blocks_per_chunk == blocks_per_window:
|
|
return blocks_to_skip
|
|
|
|
full_windows_to_skip = blocks_to_skip // blocks_per_chunk
|
|
tail_blocks = blocks_to_skip % blocks_per_chunk
|
|
tail_blocks_to_skip = tail_blocks - (blocks_per_chunk - blocks_per_window)
|
|
return full_windows_to_skip * blocks_per_window + max(0, tail_blocks_to_skip)
|
|
|
|
|
|
def _run_object_group_transfer_plan(
|
|
cache_context: BaseCacheContext,
|
|
block_ids_gpu: list[torch.Tensor],
|
|
memory_objs: Sequence[MemoryObj | None],
|
|
object_group_id: int,
|
|
batch_size: int,
|
|
skip_first_n_tokens: int,
|
|
direction: "lmc_ops.TransferDirection",
|
|
) -> None:
|
|
"""Plan and execute one object group's transfer in a single native call.
|
|
|
|
This is the fast path of :func:`transfer_kv_per_object_group`: it runs the
|
|
same batched-iteration / skip logic, but instead of issuing each staging
|
|
copy and kernel launch immediately (each a GIL release/re-acquire), it
|
|
resolves every argument to plain pointers/scalars (the "planner", GIL held
|
|
throughout) and hands the whole plan to ``execute_object_group_transfer``,
|
|
which issues all of it on the stream within a single GIL release.
|
|
|
|
Requires every object to be non-GDS (staged through the lazy-allocator
|
|
path); the caller skips groups that contain any GDS-backed object.
|
|
|
|
Args:
|
|
cache_context: The GPU cache context containing the KV cache information.
|
|
block_ids_gpu: GPU block IDs, indexed by LMCache KV group index.
|
|
memory_objs: The MemoryObj instances to copy. None entries are only
|
|
valid for D2H (the batch is skipped); H2D raises.
|
|
object_group_id: Index of the object group being copied.
|
|
batch_size: Number of memory objects per batched copy.
|
|
skip_first_n_tokens: Tokens to skip writing at the start of the range.
|
|
direction: H2D (retrieve) or D2H (store).
|
|
|
|
Raises:
|
|
ValueError: If a None entry is found in memory_objs when direction is
|
|
H2D, or if an object's size does not match its GPU staging buffer.
|
|
"""
|
|
lmcache_chunk_size = cache_context.lmcache_tokens_per_chunk
|
|
kv_groups_manager = cache_context.kv_layer_groups_manager
|
|
object_group = kv_groups_manager.object_groups[object_group_id]
|
|
kernel_group_ids = object_group.kernel_group_indices
|
|
is_h2d = direction == lmc_ops.TransferDirection.H2D
|
|
max_batch_size = cache_context.max_batch_size
|
|
|
|
# --- Per-kernel-group invariants, resolved once (vs. every batch before) ---
|
|
kernel_group_specs: list["lmc_ops.KernelGroupSpec"] = []
|
|
spec_index_by_kg: dict[int, int] = {}
|
|
blocks_per_chunk_by_kg: dict[int, int] = {}
|
|
blocks_per_window_by_kg: dict[int, int] = {}
|
|
for kernel_group_id in kernel_group_ids:
|
|
blocks_per_chunk = cache_context.calculate_num_blocks(
|
|
lmcache_chunk_size, kernel_group_id
|
|
)
|
|
tokens_per_window = min(
|
|
lmcache_chunk_size,
|
|
kv_groups_manager.get_subchunk_sw_size_tokens(kernel_group_id),
|
|
)
|
|
blocks_per_window = cache_context.calculate_num_blocks(
|
|
tokens_per_window, kernel_group_id
|
|
)
|
|
blocks_per_chunk_by_kg[kernel_group_id] = blocks_per_chunk
|
|
blocks_per_window_by_kg[kernel_group_id] = blocks_per_window
|
|
|
|
paged_ptrs = cache_context.get_kernel_group_kv_pointers(kernel_group_id)
|
|
block_ids_tensor = block_ids_gpu[kernel_group_id]
|
|
temp_buffers = [
|
|
cache_context.get_temp_kernel_group_buffer(slot, kernel_group_id)
|
|
for slot in range(max_batch_size)
|
|
]
|
|
|
|
spec_index_by_kg[kernel_group_id] = len(kernel_group_specs)
|
|
kernel_group_specs.append(
|
|
lmc_ops.KernelGroupSpec(
|
|
paged_ptrs.data_ptr(),
|
|
[buffer.data_ptr() for buffer in temp_buffers],
|
|
cache_context.get_shape_desc(kernel_group_id),
|
|
cache_context.get_slots_per_chunk_in_sw(kernel_group_id),
|
|
cache_context.get_engine_kv_format(kernel_group_id),
|
|
block_ids_tensor.data_ptr(),
|
|
block_ids_tensor.numel(),
|
|
)
|
|
)
|
|
|
|
# Temp object-group staging buffers (reused per batch slot, like above).
|
|
object_group_buffers = [
|
|
cache_context.get_temp_object_group_buffer(slot, object_group_id)
|
|
for slot in range(max_batch_size)
|
|
]
|
|
|
|
attn_desc = kv_groups_manager.get_attn_desc()
|
|
num_objects_to_skip = 0
|
|
if not attn_desc.is_full_attention(object_group_id) and is_h2d:
|
|
sw_size_chunks = attn_desc.num_chunks_in_sw[object_group_id]
|
|
num_objects_to_skip = max(0, len(memory_objs) - sw_size_chunks)
|
|
logger.debug(
|
|
"Detected sliding window for object group %d: "
|
|
"skipping the first %d objects in the batch",
|
|
object_group_id,
|
|
num_objects_to_skip,
|
|
)
|
|
|
|
# --- Walk the batches in order, emitting staging + launch work per step ---
|
|
batch_steps: list["lmc_ops.BatchStep"] = []
|
|
for start_object_idx, memory_object_batch in batched_iteration_with_skip(
|
|
memory_objs, batch_size, skip_count=num_objects_to_skip
|
|
):
|
|
if any(mo is None for mo in memory_object_batch):
|
|
if is_h2d:
|
|
raise ValueError(
|
|
"MemoryObj is None for some objects in the batch, cannot "
|
|
"perform H2D copy. memory_object_batch: "
|
|
f"{memory_object_batch}"
|
|
)
|
|
else:
|
|
continue
|
|
|
|
batch_len = len(memory_object_batch)
|
|
batch_start_token = start_object_idx * lmcache_chunk_size
|
|
batch_end_token = batch_start_token + batch_len * lmcache_chunk_size
|
|
|
|
effective_start = max(batch_start_token, skip_first_n_tokens)
|
|
if effective_start >= batch_end_token:
|
|
continue
|
|
|
|
skip_tokens_in_chunk = effective_start - batch_start_token
|
|
|
|
staging = build_staging_copies(
|
|
memory_object_batch,
|
|
object_group_buffers[:batch_len],
|
|
is_h2d,
|
|
)
|
|
|
|
launches: list["lmc_ops.LaunchVar"] = []
|
|
for kernel_group_id in kernel_group_ids:
|
|
blocks_per_chunk = blocks_per_chunk_by_kg[kernel_group_id]
|
|
blocks_per_window = blocks_per_window_by_kg[kernel_group_id]
|
|
|
|
start_block_pos = start_object_idx * blocks_per_window
|
|
end_block_pos = (start_object_idx + batch_len) * blocks_per_window
|
|
|
|
orig_skip_blocks = cache_context.calculate_num_blocks(
|
|
skip_tokens_in_chunk, kernel_group_id
|
|
)
|
|
recalculated_skip_blocks = _recalculate_blocks_to_skip(
|
|
blocks_per_chunk,
|
|
blocks_per_window,
|
|
orig_skip_blocks,
|
|
)
|
|
|
|
launches.append(
|
|
lmc_ops.LaunchVar(
|
|
spec_index_by_kg[kernel_group_id],
|
|
start_block_pos,
|
|
end_block_pos - start_block_pos,
|
|
batch_len,
|
|
recalculated_skip_blocks,
|
|
)
|
|
)
|
|
|
|
batch_steps.append(lmc_ops.BatchStep(staging, launches))
|
|
|
|
if not batch_steps:
|
|
return
|
|
|
|
lmc_ops.execute_object_group_transfer(
|
|
direction,
|
|
cache_context.device,
|
|
LazyMemoryAllocator.PIN_CHUNK_SIZE,
|
|
kernel_group_specs,
|
|
batch_steps,
|
|
)
|
|
|
|
|
|
def transfer_kv_per_object_group(
|
|
cache_context: BaseCacheContext,
|
|
block_ids_gpu: list[torch.Tensor],
|
|
memory_objs: Sequence[MemoryObj | None],
|
|
object_group_id: int,
|
|
batch_size: int,
|
|
skip_first_n_tokens: int,
|
|
direction: "lmc_ops.TransferDirection",
|
|
) -> None:
|
|
"""Helper function to transfer memory objects of a single object group
|
|
to/from GPU, with batching support.
|
|
|
|
Args:
|
|
cache_context: The GPU cache context containing the KV cache information.
|
|
block_ids_gpu: GPU block IDs to retrieve into, indexed by LMCache KV group
|
|
index. It should satisfy `len(block_ids_gpu[i]) == len(memory_objs) *
|
|
blocks_per_chunk[i]` for each group `i`.
|
|
Note that the block IDs list are already on GPU.
|
|
memory_objs: The list of MemoryObj instances to copy from. It could be
|
|
None when allocation or retrieval fails. For store (D2H), it should
|
|
ignore the None entry and continue copying the rest. For retrieve
|
|
(H2D), it should raise the error and stop copying.
|
|
object_group_id: Index of the object group being copied.
|
|
batch_size: The number of memory objects to perform batched copy
|
|
skip_first_n_tokens: Number of tokens to skip writing at the start of
|
|
the retrieve range. This avoids overwriting APC-shared GPU blocks that
|
|
may be read concurrently by other requests.
|
|
direction: The transfer direction, H2D (retrieve) or D2H (store).
|
|
|
|
Raises:
|
|
ValueError: If it founds None entry in memory_objs when direction is H2D.
|
|
Note:
|
|
This function expects the caller to stage the block ids (list[list[int]])
|
|
into GPU tensors and pass them in as `block_ids_gpu`.
|
|
"""
|
|
if _HAS_NATIVE_OBJECT_GROUP_TRANSFER and not any(
|
|
isinstance(mo, GDSMemoryObject) for mo in memory_objs
|
|
):
|
|
_run_object_group_transfer_plan(
|
|
cache_context,
|
|
block_ids_gpu,
|
|
memory_objs,
|
|
object_group_id,
|
|
batch_size,
|
|
skip_first_n_tokens,
|
|
direction,
|
|
)
|
|
return
|
|
|
|
lmcache_chunk_size = cache_context.lmcache_tokens_per_chunk
|
|
kv_groups_manager = cache_context.kv_layer_groups_manager
|
|
object_group = kv_groups_manager.object_groups[object_group_id]
|
|
kernel_group_ids = object_group.kernel_group_indices
|
|
is_h2d = direction == lmc_ops.TransferDirection.H2D
|
|
|
|
attn_desc = kv_groups_manager.get_attn_desc()
|
|
num_objects_to_skip = 0
|
|
if not attn_desc.is_full_attention(object_group_id) and is_h2d:
|
|
sw_size_chunks = attn_desc.num_chunks_in_sw[object_group_id]
|
|
num_objects_to_skip = max(0, len(memory_objs) - sw_size_chunks)
|
|
logger.debug(
|
|
"Detected sliding window for object group %d: "
|
|
"skipping the first %d objects in the batch",
|
|
object_group_id,
|
|
num_objects_to_skip,
|
|
)
|
|
|
|
for start_object_idx, memory_object_batch in batched_iteration_with_skip(
|
|
memory_objs, batch_size, skip_count=num_objects_to_skip
|
|
):
|
|
if any(mo is None for mo in memory_object_batch):
|
|
if is_h2d:
|
|
raise ValueError(
|
|
"MemoryObj is None for some objects in the batch, cannot "
|
|
"perform H2D copy. memory_object_batch: "
|
|
f"{memory_object_batch}"
|
|
)
|
|
else:
|
|
continue
|
|
|
|
batch_len = len(memory_object_batch)
|
|
batch_start_token = start_object_idx * lmcache_chunk_size
|
|
batch_end_token = batch_start_token + batch_len * lmcache_chunk_size
|
|
|
|
effective_start = max(batch_start_token, skip_first_n_tokens)
|
|
if effective_start >= batch_end_token:
|
|
continue
|
|
|
|
skip_tokens_in_chunk = effective_start - batch_start_token
|
|
|
|
# For H2D, copy from CPU to GPU tmp buffers before the kernel launch
|
|
if is_h2d:
|
|
for chunk_idx, memory_obj in enumerate(memory_object_batch):
|
|
lmcache_memcpy_async_h2d(
|
|
memory_obj,
|
|
cache_context.get_temp_object_group_buffer(
|
|
chunk_idx, object_group_id
|
|
),
|
|
)
|
|
|
|
# Do paged KV copy
|
|
for kernel_group_id in kernel_group_ids:
|
|
blocks_per_chunk = cache_context.calculate_num_blocks(
|
|
lmcache_chunk_size, kernel_group_id
|
|
)
|
|
tokens_per_window = min(
|
|
lmcache_chunk_size,
|
|
kv_groups_manager.get_subchunk_sw_size_tokens(kernel_group_id),
|
|
)
|
|
blocks_per_window = cache_context.calculate_num_blocks(
|
|
tokens_per_window, kernel_group_id
|
|
)
|
|
|
|
# Get the block ids for this chunk
|
|
start_block_pos = start_object_idx * blocks_per_window
|
|
end_block_pos = (start_object_idx + batch_len) * blocks_per_window
|
|
|
|
block_ids_curr_batch = block_ids_gpu[kernel_group_id][
|
|
start_block_pos:end_block_pos
|
|
]
|
|
|
|
# Re-calculate the skip blocks for this kernel group
|
|
orig_skip_blocks = cache_context.calculate_num_blocks(
|
|
skip_tokens_in_chunk, kernel_group_id
|
|
)
|
|
recalculated_skip_blocks = _recalculate_blocks_to_skip(
|
|
blocks_per_chunk,
|
|
blocks_per_window,
|
|
orig_skip_blocks,
|
|
)
|
|
|
|
# Launch kernel
|
|
group_kv_pointers = cache_context.get_kernel_group_kv_pointers(
|
|
kernel_group_id
|
|
)
|
|
group_lmcache_chunk_size = cache_context.get_slots_per_chunk_in_sw(
|
|
kernel_group_id
|
|
)
|
|
tmp_gpu_buffers_batched = [
|
|
cache_context.get_temp_kernel_group_buffer(
|
|
i, kernel_group_id
|
|
).data_ptr()
|
|
for i in range(batch_len)
|
|
]
|
|
lmc_ops.multi_layer_block_kv_transfer(
|
|
group_kv_pointers,
|
|
tmp_gpu_buffers_batched,
|
|
block_ids_curr_batch,
|
|
cache_context.device,
|
|
direction,
|
|
cache_context.get_shape_desc(kernel_group_id),
|
|
group_lmcache_chunk_size,
|
|
cache_context.get_engine_kv_format(kernel_group_id),
|
|
recalculated_skip_blocks,
|
|
)
|
|
|
|
# For D2H, copy from GPU tmp buffers to CPU after the kernel launch
|
|
if not is_h2d:
|
|
for chunk_idx, memory_obj in enumerate(memory_object_batch):
|
|
lmcache_memcpy_async_d2h(
|
|
cache_context.get_temp_object_group_buffer(
|
|
chunk_idx, object_group_id
|
|
),
|
|
memory_obj,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ContextEntry:
|
|
"""Registered cache context metadata for a single worker instance.
|
|
|
|
The concrete type is whatever :func:`create_cache_context` returned
|
|
for the wrapper list at registration time -- a
|
|
:class:`GPUCacheContext` for CUDA-IPC wrappers, a
|
|
:class:`CPUCacheContext` for POSIX-SHM wrappers. Both expose
|
|
the same ``kv_tensors`` / ``engine_kv_format`` / ``num_layers`` / ...
|
|
duck-typed surface, so downstream consumers stay agnostic.
|
|
|
|
Args:
|
|
cache_context: Platform cache context (GPU or CPU) managing
|
|
shape and pointers to the registered KV cache tensors.
|
|
model_name: The name of the model associated with this KV cache.
|
|
world_size: The world size associated with this KV cache.
|
|
last_seen: ``time.monotonic()`` of the most recent activity from
|
|
this instance (register, PING, store, or retrieve). Drives reaping.
|
|
has_liveness_signal: True once the instance has sent at least one
|
|
PING. Selects the reap window (timeout vs registration grace).
|
|
Latched only by PING, never by traffic.
|
|
"""
|
|
|
|
cache_context: BaseCacheContext
|
|
model_name: str
|
|
world_size: int
|
|
last_seen: float = 0.0
|
|
has_liveness_signal: bool = False
|
|
|
|
|
|
class LMCacheDrivenTransferModule(InstanceLivenessTarget):
|
|
"""Handles LMCache-driven KV cache transfer operations.
|
|
|
|
Owns GPU context registrations and provides handlers for
|
|
register, unregister, store, and retrieve of GPU KV caches.
|
|
|
|
Args:
|
|
ctx: The shared engine context.
|
|
"""
|
|
|
|
def __init__(self, ctx: MPCacheServerContext) -> None:
|
|
self._ctx = ctx
|
|
self._cache_contexts: dict[int, ContextEntry] = {}
|
|
# Guards all reads/writes of _cache_contexts. The reaper mutates it
|
|
# off the MQ main loop, so register/unregister/store/retrieve and
|
|
# report_status all serialize through this lock. Held only for dict
|
|
# ops -- never across context creation, layout-registry calls, or
|
|
# empty_cache (leaf-lock invariant: no thread holds two locks).
|
|
self._lock = threading.Lock()
|
|
|
|
# Route finish_write / finish_read_prefetched through a C++ host
|
|
# callback so the driver thread doesn't acquire the GIL.
|
|
self._device_host_func_dispatcher = DeviceHostFuncDispatcher()
|
|
self._device_host_func_dispatcher.register(
|
|
"finish_write",
|
|
self._ctx.storage_manager.finish_write,
|
|
payload_type=list[ObjectKey],
|
|
)
|
|
self._device_host_func_dispatcher.register(
|
|
"finish_read_prefetched",
|
|
self._ctx.storage_manager.finish_read_prefetched,
|
|
payload_type=list[ObjectKey],
|
|
)
|
|
self._device_host_func_dispatcher.start()
|
|
|
|
@property
|
|
def context(self) -> MPCacheServerContext:
|
|
"""Return the shared engine context. Exposed for testing only."""
|
|
return self._ctx
|
|
|
|
def get_and_touch_context_entry(self, instance_id: int) -> ContextEntry | None:
|
|
"""Return the entry for ``instance_id``, refreshing its last-seen time.
|
|
|
|
The refresh keeps an actively transferring worker from being reaped
|
|
even if its PINGs are briefly delayed. Does not latch the
|
|
ping-proven flag -- only PINGs do that.
|
|
|
|
Args:
|
|
instance_id: The worker instance ID.
|
|
|
|
Returns:
|
|
The entry, or None if the instance is not (or no longer) tracked.
|
|
"""
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
entry = self._cache_contexts.get(instance_id)
|
|
if entry is not None:
|
|
entry.last_seen = now
|
|
return entry
|
|
|
|
def context_entries_snapshot(self) -> dict[int, ContextEntry]:
|
|
"""Return a shallow copy of the registry for iteration or status.
|
|
|
|
Returns:
|
|
A new dict mapping instance ID to entry; does not refresh
|
|
last-seen times.
|
|
"""
|
|
with self._lock:
|
|
return dict(self._cache_contexts)
|
|
|
|
def touch_instance(self, instance_id: int) -> None:
|
|
"""Refresh the worker's last-seen time and mark it ping-proven.
|
|
|
|
A no-op if the instance is not tracked.
|
|
|
|
Args:
|
|
instance_id: The worker instance ID.
|
|
"""
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
entry = self._cache_contexts.get(instance_id)
|
|
if entry is not None:
|
|
entry.last_seen = now
|
|
entry.has_liveness_signal = True
|
|
|
|
def tracked_instance_count(self) -> int:
|
|
"""Return the number of currently registered instances."""
|
|
with self._lock:
|
|
return len(self._cache_contexts)
|
|
|
|
def reap_stale_instances(
|
|
self, reap_timeout_s: float, registration_grace_s: float
|
|
) -> list[int]:
|
|
"""Reap GPU registrations that have gone silent.
|
|
|
|
A ping-proven instance is judged against ``reap_timeout_s``; one
|
|
that has never pinged against the larger ``registration_grace_s``.
|
|
|
|
Args:
|
|
reap_timeout_s: Silence budget for ping-proven instances.
|
|
registration_grace_s: Silence budget for never-pinged instances.
|
|
|
|
Returns:
|
|
The instance IDs reaped this scan.
|
|
"""
|
|
now = time.monotonic()
|
|
reaped: list[tuple[int, ContextEntry]] = []
|
|
with self._lock:
|
|
stale_ids = [
|
|
iid
|
|
for iid, entry in self._cache_contexts.items()
|
|
if now - entry.last_seen
|
|
> (
|
|
reap_timeout_s
|
|
if entry.has_liveness_signal
|
|
else registration_grace_s
|
|
)
|
|
]
|
|
for iid in stale_ids:
|
|
reaped.append((iid, self._cache_contexts.pop(iid)))
|
|
for iid, entry in reaped:
|
|
self._release_entry(entry)
|
|
logger.warning(
|
|
"Reaped GPU instance %d: silent for %.1fs (pinged=%s)",
|
|
iid,
|
|
now - entry.last_seen,
|
|
entry.has_liveness_signal,
|
|
)
|
|
return [iid for iid, _ in reaped]
|
|
|
|
def _release_entry(self, entry: ContextEntry) -> None:
|
|
"""Release resources for a popped entry (run outside the lock).
|
|
|
|
Args:
|
|
entry: The entry removed from the registry.
|
|
"""
|
|
entry.cache_context.close()
|
|
self._ctx.layout_desc_registry.unregister(entry.model_name, entry.world_size)
|
|
torch_dev.empty_cache()
|
|
|
|
def get_handlers(self) -> list[HandlerSpec]:
|
|
"""Return handler specs for all request types this module serves.
|
|
|
|
Returns:
|
|
A list of HandlerSpec entries mapping request types to
|
|
their handler callables and thread pool assignments.
|
|
"""
|
|
return [
|
|
HandlerSpec(
|
|
RequestType.REGISTER_KV_CACHE,
|
|
self.register_kv_cache,
|
|
ThreadPoolType.SYNC,
|
|
),
|
|
HandlerSpec(
|
|
RequestType.UNREGISTER_KV_CACHE,
|
|
self.unregister_kv_cache,
|
|
ThreadPoolType.SYNC,
|
|
),
|
|
HandlerSpec(
|
|
RequestType.STORE,
|
|
self.store,
|
|
ThreadPoolType.AFFINITY,
|
|
),
|
|
HandlerSpec(
|
|
RequestType.RETRIEVE,
|
|
self.retrieve,
|
|
ThreadPoolType.AFFINITY,
|
|
),
|
|
]
|
|
|
|
def report_status(self) -> dict:
|
|
"""Return GPU transfer module status information.
|
|
|
|
Returns:
|
|
A dict containing registered GPU instance IDs and
|
|
per-instance KV cache layout metadata.
|
|
"""
|
|
registered_gpu_ids: list[int] = []
|
|
cache_context_meta: dict[str, dict] = {}
|
|
|
|
for instance_id, entry in self.context_entries_snapshot().items():
|
|
registered_gpu_ids.append(instance_id)
|
|
ctx = entry.cache_context
|
|
cache_context_meta[str(instance_id)] = {
|
|
"model_name": entry.model_name,
|
|
"world_size": entry.world_size,
|
|
"kv_cache_layout": ctx.report_status(),
|
|
}
|
|
|
|
return {
|
|
"registered_gpu_ids": registered_gpu_ids,
|
|
"cache_context_meta": cache_context_meta,
|
|
}
|
|
|
|
def close(self) -> None:
|
|
"""Release GPU resources owned by this module."""
|
|
# Stop the drain thread before storage_manager.close() so any
|
|
# in-flight completions reach a live storage manager.
|
|
self._device_host_func_dispatcher.stop()
|
|
|
|
with self._lock:
|
|
entries = list(self._cache_contexts.values())
|
|
self._cache_contexts.clear()
|
|
for entry in entries:
|
|
entry.cache_context.close()
|
|
if entries:
|
|
torch_dev.empty_cache()
|
|
|
|
def register_kv_cache(
|
|
self,
|
|
instance_id: int,
|
|
kv_caches: KVCache,
|
|
model_name: str,
|
|
world_size: int,
|
|
engine_type: EngineType,
|
|
layout_hints: LayoutHints,
|
|
engine_group_infos: list[EngineGroupInfo],
|
|
) -> None:
|
|
"""Register the KV cache tensors for a given GPU instance ID.
|
|
|
|
Args:
|
|
instance_id: The GPU instance ID (such as PID).
|
|
kv_caches: The KV cache tensor wrappers from the
|
|
serving engine.
|
|
model_name: The name of the model associated with this KV cache.
|
|
world_size: The world size associated with this KV cache.
|
|
engine_type: Which serving engine produced the caches.
|
|
Forwarded to GPUCacheContext for format detection.
|
|
layout_hints: See LayoutHints. Forwarded to
|
|
GPUCacheContext for GPU KV format detection.
|
|
engine_group_infos: Engine-neutral KV cache group metadata
|
|
(already msgspec-decoded by the message queue).
|
|
"""
|
|
now = time.monotonic()
|
|
# NOOP-register: an already-registered instance (e.g. a recovering
|
|
# worker re-registering on its first ping) refreshes its last-seen
|
|
# time so a stale entry is not reaped right after recovery. REGISTER
|
|
# is SYNC-serialized on the MQ main loop, so it is the sole inserter.
|
|
with self._lock:
|
|
existing = self._cache_contexts.get(instance_id)
|
|
if existing is not None:
|
|
existing.last_seen = now
|
|
logger.info(
|
|
"Instance %d already registered; refreshing liveness",
|
|
instance_id,
|
|
)
|
|
return
|
|
|
|
# Build the context and layout descriptor outside the lock.
|
|
cache_context = create_cache_context(
|
|
kv_caches,
|
|
self._ctx.chunk_size,
|
|
layout_hints=layout_hints or None,
|
|
engine_group_infos=engine_group_infos,
|
|
engine_type=engine_type,
|
|
separate_object_groups=self._ctx.separate_object_groups,
|
|
full_sw_kv=self._ctx.full_sw_kv,
|
|
)
|
|
layout_desc = get_layout_desc(
|
|
cache_context, self._ctx.chunk_size, object_group_id=0
|
|
)
|
|
kv_groups_manager = cache_context.kv_layer_groups_manager
|
|
attn_desc = kv_groups_manager.get_attn_desc()
|
|
self._ctx.layout_desc_registry.register(
|
|
model_name, world_size, layout_desc, attn_desc
|
|
)
|
|
|
|
with self._lock:
|
|
self._cache_contexts[instance_id] = ContextEntry(
|
|
cache_context=cache_context,
|
|
model_name=model_name,
|
|
world_size=world_size,
|
|
last_seen=now,
|
|
has_liveness_signal=False,
|
|
)
|
|
|
|
logger.info(
|
|
"Registered KV cache for GPU ID %d with %d layers",
|
|
instance_id,
|
|
cache_context.num_layers,
|
|
)
|
|
|
|
def unregister_kv_cache(self, instance_id: int) -> None:
|
|
"""Unregister the KV cache tensors for a given GPU instance ID.
|
|
|
|
Args:
|
|
instance_id: The GPU instance ID (such as PID).
|
|
"""
|
|
with self._lock:
|
|
entry = self._cache_contexts.pop(instance_id, None)
|
|
if entry is None:
|
|
logger.warning(
|
|
"No registered GPU context found for instance ID %d", instance_id
|
|
)
|
|
return
|
|
|
|
self._release_entry(entry)
|
|
logger.info("Unregistered KV cache for GPU ID %d", instance_id)
|
|
|
|
@_lmcache_nvtx_annotate
|
|
def store(
|
|
self,
|
|
key: IPCCacheServerKey,
|
|
instance_id: int,
|
|
gpu_block_ids: list[list[int]],
|
|
event_ipc_handle: bytes,
|
|
) -> tuple[bytes, bool]:
|
|
"""Store the GPU KV cache blocks to CPU.
|
|
|
|
Args:
|
|
key: The IPC key for the KV cache blocks.
|
|
Must have worker_id != None (worker store operation).
|
|
instance_id: The GPU instance ID (such as PID).
|
|
gpu_block_ids: GPU block IDs to store, indexed by LMCache KV
|
|
group index.
|
|
event_ipc_handle: The IPC handle of the event to wait on.
|
|
|
|
Returns:
|
|
A tuple where the first element is the IPC handle of the event
|
|
that signals the completion of the store operation, and the second
|
|
element indicates whether the store operation completed without a
|
|
fatal error (not whether every requested chunk was stored; see
|
|
Notes).
|
|
|
|
Raises:
|
|
ValueError: If no GPU context is registered for the given instance ID.
|
|
RuntimeError: If the backend does not support IPC event handles.
|
|
|
|
Notes:
|
|
All-or-nothing. If ``gpu_block_ids`` do not fully cover every chunk
|
|
``key`` resolves to for every LMCache group (e.g. a caller/protocol
|
|
bug), or a copy fails, the whole store is skipped and nothing is
|
|
committed (logged at WARNING); a subsequent retrieve simply misses
|
|
and the engine recomputes. The boolean result reports whether the
|
|
store completed without such a failure.
|
|
"""
|
|
st = time.perf_counter()
|
|
|
|
entry = self.get_and_touch_context_entry(instance_id)
|
|
if entry is None:
|
|
raise ValueError(f"No GPU context registered for instance ID {instance_id}")
|
|
cache_context = entry.cache_context
|
|
model_name = entry.model_name
|
|
|
|
num_object_groups = cache_context.kv_layer_groups_manager.num_object_groups
|
|
obj_keys_per_obj_group = self._ctx.resolve_obj_keys(
|
|
key, list(range(num_object_groups))
|
|
)
|
|
num_chunks = len(obj_keys_per_obj_group[0])
|
|
|
|
# NOTE: different engine groups may have different block sizes, so
|
|
# ``blocks_per_chunk[i]`` is the number of blocks in one chunk for
|
|
# group ``i``.
|
|
blocks_per_chunk = [
|
|
cache_context.calculate_num_blocks(self._ctx.chunk_size, group_idx)
|
|
for group_idx in range(
|
|
cache_context.kv_layer_groups_manager.num_kernel_groups
|
|
)
|
|
]
|
|
|
|
with (
|
|
torch_dev.device(cache_context.device),
|
|
torch_dev.stream(cache_context.stream),
|
|
):
|
|
check_interprocess_event_support()
|
|
event = torch_dev.Event(interprocess=True)
|
|
|
|
# Fail closed: every LMCache group must have block IDs covering all
|
|
# chunks. A short list (e.g. a caller/protocol bug) would otherwise
|
|
# drive the transfer kernel to read out-of-bounds GPU memory, so skip
|
|
# the whole store and commit nothing rather than caching a partial or
|
|
# garbage entry. A later request can store it once the block IDs are
|
|
# complete. Checked on the raw block ids, before cutting drops the
|
|
# per-chunk blocks that sliding-window groups do not need.
|
|
if any(
|
|
len(group_block_ids) < num_chunks * bpc
|
|
for group_block_ids, bpc in zip(
|
|
gpu_block_ids, blocks_per_chunk, strict=True
|
|
)
|
|
):
|
|
logger.warning(
|
|
"STORE block ID underflow for request_id=%s: each group needs "
|
|
"num_chunks * blocks_per_chunk block IDs for %d chunks "
|
|
"(per-group blocks_per_chunk=%s); skipping the store.",
|
|
key.request_id,
|
|
num_chunks,
|
|
blocks_per_chunk,
|
|
)
|
|
event.record()
|
|
return event.ipc_handle(), False
|
|
|
|
block_ids_per_group_gpu = downsample_and_stage_block_ids(
|
|
cache_context, gpu_block_ids
|
|
)
|
|
|
|
if not hasattr(torch_dev.Event, "from_ipc_handle"):
|
|
raise RuntimeError(
|
|
f"Backend '{torch_device_type}' does not support IPC event "
|
|
"handles (Event.from_ipc_handle not available). "
|
|
"Multiprocess IPC requires CUDA."
|
|
)
|
|
vllm_event = torch_dev.Event.from_ipc_handle(
|
|
cache_context.device, event_ipc_handle
|
|
)
|
|
vllm_event.wait(stream=cache_context.stream)
|
|
|
|
# CPU-synchronous sentinel: a GPU store is about to be enqueued.
|
|
# Must be published via publish() (not publish_on_stream) so the
|
|
# drain thread sees it before MP_REQUEST_END can race MP_STORE_END.
|
|
self._ctx.event_bus.publish(
|
|
Event(
|
|
event_type=EventType.MP_STORE_SUBMITTED,
|
|
session_id=key.request_id,
|
|
metadata={"device": str(cache_context.device)},
|
|
)
|
|
)
|
|
|
|
self._ctx.event_bus.publish_on_stream(
|
|
cache_context.cupy_stream,
|
|
Event(
|
|
event_type=EventType.MP_STORE_START,
|
|
session_id=key.request_id,
|
|
metadata={
|
|
"device": str(cache_context.device),
|
|
"engine_id": instance_id,
|
|
"model_name": model_name,
|
|
},
|
|
),
|
|
)
|
|
|
|
reserved_dict: dict[ObjectKey, MemoryObj] = {}
|
|
all_dict: dict[ObjectKey, MemoryObj] = {}
|
|
total_bytes: int = 0
|
|
store_succeeded = False
|
|
try:
|
|
for obj_group_id in range(num_object_groups):
|
|
obj_keys = obj_keys_per_obj_group[obj_group_id]
|
|
layout_desc = get_layout_desc(
|
|
cache_context,
|
|
self._ctx.chunk_size,
|
|
object_group_id=obj_group_id,
|
|
)
|
|
reserved_dict = self._ctx.storage_manager.reserve_write(
|
|
obj_keys, layout_desc, "new"
|
|
)
|
|
all_dict.update(reserved_dict)
|
|
if reserved_dict:
|
|
total_bytes += next(
|
|
iter(reserved_dict.values())
|
|
).get_size() * len(reserved_dict)
|
|
|
|
# Keys not in reserved_dict (skipped by the storage manager)
|
|
# become None entries; the helper skips them for D2H.
|
|
memory_objs: list[MemoryObj | None] = [
|
|
reserved_dict.get(obj_key) for obj_key in obj_keys
|
|
]
|
|
|
|
# NOTE: batch_size must stay 1 for store.
|
|
transfer_kv_per_object_group(
|
|
cache_context,
|
|
block_ids_per_group_gpu,
|
|
memory_objs,
|
|
object_group_id=obj_group_id,
|
|
batch_size=1,
|
|
skip_first_n_tokens=0,
|
|
direction=lmc_ops.TransferDirection.D2H,
|
|
)
|
|
|
|
store_succeeded = True
|
|
except Exception:
|
|
logger.exception("Cannot store keys due to exception")
|
|
return event.ipc_handle(), False
|
|
finally:
|
|
event.record()
|
|
# Fail closed: commit the reserved objects only when every chunk
|
|
# copied successfully; otherwise the whole store is skipped.
|
|
stored_count = len(all_dict) if store_succeeded else 0
|
|
if stored_count:
|
|
submit_callback_to_stream(
|
|
cache_context.cupy_stream,
|
|
"finish_write",
|
|
list(all_dict.keys()),
|
|
)
|
|
else:
|
|
total_bytes = 0
|
|
self._ctx.event_bus.publish_on_stream(
|
|
cache_context.cupy_stream,
|
|
Event(
|
|
event_type=EventType.MP_STORE_END,
|
|
session_id=key.request_id,
|
|
metadata={
|
|
"stored_count": stored_count,
|
|
"device": str(cache_context.device),
|
|
"engine_id": instance_id,
|
|
"model_name": model_name,
|
|
"total_bytes": total_bytes,
|
|
},
|
|
),
|
|
)
|
|
|
|
ed = time.perf_counter()
|
|
if stored_count:
|
|
logger.info(
|
|
"Stored %d tokens in %.3f seconds",
|
|
num_chunks * self._ctx.chunk_size,
|
|
ed - st,
|
|
)
|
|
return event.ipc_handle(), True
|
|
|
|
@_lmcache_nvtx_annotate
|
|
def retrieve(
|
|
self,
|
|
key: IPCCacheServerKey,
|
|
instance_id: int,
|
|
gpu_block_ids: list[list[int]],
|
|
event_ipc_handle: bytes,
|
|
skip_first_n_tokens: int = 0,
|
|
) -> tuple[bytes, bool]:
|
|
"""Retrieve the CPU KV cache and put into GPU blocks.
|
|
|
|
Args:
|
|
key: The IPC key for the KV cache blocks.
|
|
Must have worker_id != None (worker retrieve operation).
|
|
instance_id: The GPU instance ID (such as PID).
|
|
gpu_block_ids: GPU block IDs to retrieve into, indexed by LMCache
|
|
KV group index.
|
|
event_ipc_handle: The IPC handle of the event to wait on.
|
|
skip_first_n_tokens: Number of tokens to skip writing at
|
|
the start of the retrieve range. This avoids overwriting
|
|
APC-shared GPU blocks that may be read concurrently by other
|
|
requests.
|
|
|
|
Returns:
|
|
A tuple where the first element is the IPC handle of the event
|
|
that signals the completion of the retrieve operation, and the
|
|
second element indicates whether the key was successfully retrieved.
|
|
|
|
Raises:
|
|
ValueError: If no GPU context is registered for the given instance ID.
|
|
"""
|
|
st = time.perf_counter()
|
|
|
|
entry = self.get_and_touch_context_entry(instance_id)
|
|
if entry is None:
|
|
raise ValueError(f"No GPU context registered for instance ID {instance_id}")
|
|
cache_context = entry.cache_context
|
|
model_name = entry.model_name
|
|
|
|
num_object_groups = cache_context.kv_layer_groups_manager.num_object_groups
|
|
obj_keys_per_obj_group = self._ctx.resolve_obj_keys(
|
|
key, list(range(num_object_groups))
|
|
)
|
|
num_chunks = len(obj_keys_per_obj_group[0])
|
|
|
|
# CPU-synchronous sentinel: a GPU retrieve is about to be enqueued.
|
|
# Must be published via publish() (not publish_on_stream) so the
|
|
# drain thread sees it before MP_REQUEST_END can race MP_RETRIEVE_END.
|
|
self._ctx.event_bus.publish(
|
|
Event(
|
|
event_type=EventType.MP_RETRIEVE_SUBMITTED,
|
|
session_id=key.request_id,
|
|
metadata={"device": str(cache_context.device)},
|
|
)
|
|
)
|
|
|
|
self._ctx.event_bus.publish_on_stream(
|
|
cache_context.cupy_stream,
|
|
Event(
|
|
event_type=EventType.MP_RETRIEVE_START,
|
|
session_id=key.request_id,
|
|
metadata={
|
|
"device": str(cache_context.device),
|
|
"engine_id": instance_id,
|
|
"model_name": model_name,
|
|
},
|
|
),
|
|
)
|
|
|
|
blocks_per_chunk = [
|
|
cache_context.calculate_num_blocks(self._ctx.chunk_size, group_idx)
|
|
for group_idx in range(
|
|
cache_context.kv_layer_groups_manager.num_kernel_groups
|
|
)
|
|
]
|
|
|
|
with (
|
|
torch_dev.device(cache_context.device),
|
|
torch_dev.stream(cache_context.stream),
|
|
):
|
|
check_interprocess_event_support()
|
|
event = torch_dev.Event(interprocess=True)
|
|
|
|
# Fail closed: a short block-id list would drive the transfer
|
|
# kernel to write out-of-bounds GPU memory. Checked on the raw
|
|
# block ids, before cutting drops the per-chunk blocks that
|
|
# sliding-window groups do not need.
|
|
if any(
|
|
len(group_block_ids) < num_chunks * bpc
|
|
for group_block_ids, bpc in zip(
|
|
gpu_block_ids, blocks_per_chunk, strict=True
|
|
)
|
|
):
|
|
logger.error(
|
|
"RETRIEVE block ID underflow for request_id=%s: each group "
|
|
"needs num_chunks * blocks_per_chunk block IDs for %d "
|
|
"chunks (per-group blocks_per_chunk=%s); skipping the "
|
|
"retrieve.",
|
|
key.request_id,
|
|
num_chunks,
|
|
blocks_per_chunk,
|
|
)
|
|
event.record()
|
|
return event.ipc_handle(), False
|
|
|
|
# Cut and stage all block_ids to GPU once before the transfer
|
|
block_ids_per_group_gpu = downsample_and_stage_block_ids(
|
|
cache_context, gpu_block_ids
|
|
)
|
|
|
|
prefetched_keys: list[ObjectKey] = []
|
|
total_bytes = 0
|
|
try:
|
|
for obj_group_id in range(num_object_groups):
|
|
obj_keys = obj_keys_per_obj_group[obj_group_id]
|
|
with self._ctx.storage_manager.read_prefetched_results(
|
|
obj_keys
|
|
) as memory_objs:
|
|
if not memory_objs or len(memory_objs) != len(obj_keys):
|
|
logger.error("Some keys not found during retrieve!")
|
|
return event.ipc_handle(), False
|
|
|
|
total_bytes += sum(mo.get_size() for mo in memory_objs)
|
|
|
|
transfer_kv_per_object_group(
|
|
cache_context,
|
|
block_ids_per_group_gpu,
|
|
memory_objs,
|
|
object_group_id=obj_group_id,
|
|
batch_size=cache_context.max_batch_size,
|
|
skip_first_n_tokens=skip_first_n_tokens,
|
|
direction=lmc_ops.TransferDirection.H2D,
|
|
)
|
|
# Extend only after the copy is enqueued: on exception,
|
|
# read_prefetched_results releases this group's locks
|
|
# itself, and a key must not be released twice.
|
|
prefetched_keys.extend(obj_keys)
|
|
except Exception:
|
|
logger.exception("Cannot retrieve keys due to exception")
|
|
return event.ipc_handle(), False
|
|
finally:
|
|
event.record()
|
|
if prefetched_keys:
|
|
submit_callback_to_stream(
|
|
cache_context.cupy_stream,
|
|
"finish_read_prefetched",
|
|
prefetched_keys,
|
|
)
|
|
self._ctx.event_bus.publish_on_stream(
|
|
cache_context.cupy_stream,
|
|
Event(
|
|
event_type=EventType.MP_RETRIEVE_END,
|
|
session_id=key.request_id,
|
|
metadata={
|
|
"retrieved_count": len(prefetched_keys),
|
|
"device": str(cache_context.device),
|
|
"engine_id": instance_id,
|
|
"model_name": model_name,
|
|
"cache_salt": key.cache_salt,
|
|
"total_bytes": total_bytes,
|
|
},
|
|
),
|
|
)
|
|
tokens_retrieved = num_chunks * self._ctx.chunk_size
|
|
ed = time.perf_counter()
|
|
logger.info(
|
|
"Retrieved %d tokens in %.3f seconds",
|
|
tokens_retrieved,
|
|
ed - st,
|
|
)
|
|
|
|
return event.ipc_handle(), True
|