262 lines
10 KiB
Python
262 lines
10 KiB
Python
|
from collections import deque
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from vllm.logger import logger
|
|||
|
|
from vllm.utils.platform_utils import is_pin_memory_available
|
|||
|
|
from vllm.v1.attention.backend import AttentionBackend # type: ignore
|
|||
|
|
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
|
|||
|
|
from vllm.v1.kv_offload.worker.worker import OffloadingHandler, TransferResult, TransferSpec
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class Transfer:
|
|||
|
|
job_id: int
|
|||
|
|
stream: torch.npu.Stream
|
|||
|
|
start_event: torch.npu.Event
|
|||
|
|
end_event: torch.npu.Event
|
|||
|
|
num_bytes: int
|
|||
|
|
|
|||
|
|
|
|||
|
|
def expand_block_ids(
|
|||
|
|
block_ids: np.ndarray,
|
|||
|
|
block_size_factor: int,
|
|||
|
|
output: np.ndarray,
|
|||
|
|
skip_count: int = 0,
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
Convert a list of block IDs to a list of matching block ids,
|
|||
|
|
assuming each block is composed of actual block_size_factor blocks.
|
|||
|
|
Outputs to output tensor.
|
|||
|
|
The first skip_count blocks will be skipped.
|
|||
|
|
Note that skip_count must be less than block_size_factor.
|
|||
|
|
|
|||
|
|
For example, if block_ids = [0, 1, 3] and block_size_factor = 4,
|
|||
|
|
then it yields [0, 1, 2, 3, 4, 5, 6, 7, 12, 13, 14, 15]
|
|||
|
|
since 0 maps to [0, 1, 2, 3]
|
|||
|
|
1 maps to [4, 5, 6, 7]
|
|||
|
|
and 3 maps to [12, 13, 14, 15]
|
|||
|
|
"""
|
|||
|
|
assert skip_count < block_size_factor
|
|||
|
|
|
|||
|
|
# Vectorized: compute all sub-block IDs at once
|
|||
|
|
bases = block_ids * block_size_factor
|
|||
|
|
offsets = np.arange(block_size_factor)
|
|||
|
|
# shape: (num_blocks, block_size_factor) -> ravel to 1D
|
|||
|
|
all_ids = (bases[:, None] + offsets[None, :]).ravel()
|
|||
|
|
# Skip the first skip_count elements (only affects first block)
|
|||
|
|
if skip_count > 0:
|
|||
|
|
all_ids = all_ids[skip_count:]
|
|||
|
|
output[: len(all_ids)] = all_ids
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CpuNpuOffloadingHandler(OffloadingHandler):
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
gpu_block_size: int,
|
|||
|
|
cpu_block_size: int,
|
|||
|
|
num_cpu_blocks: int,
|
|||
|
|
gpu_caches: dict[str, torch.Tensor],
|
|||
|
|
attn_backends: dict[str, type[AttentionBackend]],
|
|||
|
|
):
|
|||
|
|
assert cpu_block_size % gpu_block_size == 0
|
|||
|
|
self.block_size_factor = cpu_block_size // gpu_block_size
|
|||
|
|
|
|||
|
|
# npu streams for npu->cpu and cpu->npu
|
|||
|
|
self.d2h_stream = torch.npu.Stream()
|
|||
|
|
self.h2d_stream = torch.npu.Stream()
|
|||
|
|
|
|||
|
|
# Ordered queue of in-flight transfers per direction
|
|||
|
|
self._d2h_transfers: deque[Transfer] = deque()
|
|||
|
|
self._h2d_transfers: deque[Transfer] = deque()
|
|||
|
|
|
|||
|
|
# Reusable event pool to avoid allocation overhead
|
|||
|
|
self._event_pool: list[torch.npu.Event] = []
|
|||
|
|
|
|||
|
|
pin_memory = is_pin_memory_available()
|
|||
|
|
|
|||
|
|
# allocate cpu tensors
|
|||
|
|
logger.info("Allocating %d CPU tensors...", len(gpu_caches))
|
|||
|
|
self.npu_tensors: list[torch.Tensor] = []
|
|||
|
|
self.cpu_tensors: list[torch.Tensor] = []
|
|||
|
|
for layer_name, gpu_tensor in gpu_caches.items():
|
|||
|
|
self.npu_tensors.append(gpu_tensor)
|
|||
|
|
|
|||
|
|
gpu_shape = gpu_tensor[0].shape
|
|||
|
|
|
|||
|
|
num_blocks_idx = 0
|
|||
|
|
cpu_shape = list(gpu_shape)
|
|||
|
|
cpu_shape[num_blocks_idx] = num_cpu_blocks * self.block_size_factor
|
|||
|
|
|
|||
|
|
logger.debug("Allocating CPU tensor of shape %r", cpu_shape)
|
|||
|
|
self.cpu_tensors.append(
|
|||
|
|
(
|
|||
|
|
torch.zeros(
|
|||
|
|
cpu_shape,
|
|||
|
|
dtype=gpu_tensor[0].dtype,
|
|||
|
|
device="cpu",
|
|||
|
|
pin_memory=pin_memory,
|
|||
|
|
),
|
|||
|
|
torch.zeros(
|
|||
|
|
cpu_shape,
|
|||
|
|
dtype=gpu_tensor[0].dtype,
|
|||
|
|
device="cpu",
|
|||
|
|
pin_memory=pin_memory,
|
|||
|
|
),
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Pre-compute base pointers and block sizes for batch copies.
|
|||
|
|
# In vllm-ascend, each layer's KV cache is stored as a tuple
|
|||
|
|
# (key_cache, value_cache), so we flatten them into individual
|
|||
|
|
# sub-tensors for batching: [layer0_key, layer0_value,
|
|||
|
|
# layer1_key, layer1_value, ...].
|
|||
|
|
npu_base_ptrs = []
|
|||
|
|
cpu_base_ptrs = []
|
|||
|
|
block_sizes_in_bytes = []
|
|||
|
|
|
|||
|
|
for npu_tensor, cpu_tensor in zip(self.npu_tensors, self.cpu_tensors):
|
|||
|
|
for kv_idx in range(2): # 0=key, 1=value
|
|||
|
|
npu_t = npu_tensor[kv_idx]
|
|||
|
|
cpu_t = cpu_tensor[kv_idx]
|
|||
|
|
npu_base_ptrs.append(npu_t.data_ptr())
|
|||
|
|
cpu_base_ptrs.append(cpu_t.data_ptr())
|
|||
|
|
# block size in bytes = stride of dim 0 (elements) * element size
|
|||
|
|
block_sizes_in_bytes.append(npu_t.stride(0) * npu_t.element_size())
|
|||
|
|
|
|||
|
|
self._npu_base_ptrs = np.array(npu_base_ptrs, dtype=np.int64)
|
|||
|
|
self._cpu_base_ptrs = np.array(cpu_base_ptrs, dtype=np.int64)
|
|||
|
|
self._block_size_in_bytes_arr = np.array(block_sizes_in_bytes, dtype=np.int64)
|
|||
|
|
# Total bytes per block across all sub-tensors (for transfer stats)
|
|||
|
|
self._total_bytes_per_block = int(self._block_size_in_bytes_arr.sum())
|
|||
|
|
|
|||
|
|
def _get_event(self) -> torch.npu.Event:
|
|||
|
|
if self._event_pool:
|
|||
|
|
return self._event_pool.pop()
|
|||
|
|
return torch.npu.Event(enable_timing=True)
|
|||
|
|
|
|||
|
|
def _recycle_event(self, event: torch.npu.Event) -> None:
|
|||
|
|
self._event_pool.append(event)
|
|||
|
|
|
|||
|
|
def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:
|
|||
|
|
src_spec, dst_spec = spec
|
|||
|
|
if isinstance(src_spec, CPULoadStoreSpec):
|
|||
|
|
assert isinstance(dst_spec, GPULoadStoreSpec)
|
|||
|
|
stream = self.h2d_stream
|
|||
|
|
src_base_ptrs = self._cpu_base_ptrs
|
|||
|
|
dst_base_ptrs = self._npu_base_ptrs
|
|||
|
|
src_block_size_factor = self.block_size_factor
|
|||
|
|
dst_block_size_factor = 1
|
|||
|
|
is_d2h = False
|
|||
|
|
transfers = self._h2d_transfers
|
|||
|
|
else:
|
|||
|
|
assert isinstance(src_spec, GPULoadStoreSpec)
|
|||
|
|
assert isinstance(dst_spec, CPULoadStoreSpec)
|
|||
|
|
stream = self.d2h_stream
|
|||
|
|
src_base_ptrs = self._npu_base_ptrs
|
|||
|
|
dst_base_ptrs = self._cpu_base_ptrs
|
|||
|
|
src_block_size_factor = 1
|
|||
|
|
dst_block_size_factor = self.block_size_factor
|
|||
|
|
is_d2h = True
|
|||
|
|
transfers = self._d2h_transfers
|
|||
|
|
|
|||
|
|
src_blocks = src_spec.block_ids
|
|||
|
|
dst_blocks = dst_spec.block_ids
|
|||
|
|
assert src_blocks.ndim == 1
|
|||
|
|
assert dst_blocks.ndim == 1
|
|||
|
|
|
|||
|
|
dst_sub_blocks_to_skip = -src_blocks.size % dst_block_size_factor
|
|||
|
|
src_sub_block_count = src_blocks.size * src_block_size_factor
|
|||
|
|
|
|||
|
|
assert src_sub_block_count == dst_blocks.size * dst_block_size_factor - dst_sub_blocks_to_skip
|
|||
|
|
|
|||
|
|
# Expand block IDs into sub-block IDs
|
|||
|
|
src_block_ids = np.empty(src_sub_block_count, dtype=np.int64)
|
|||
|
|
dst_block_ids = np.empty(src_sub_block_count, dtype=np.int64)
|
|||
|
|
expand_block_ids(src_blocks, src_block_size_factor, src_block_ids)
|
|||
|
|
expand_block_ids(
|
|||
|
|
dst_blocks,
|
|||
|
|
dst_block_size_factor,
|
|||
|
|
dst_block_ids,
|
|||
|
|
skip_count=dst_sub_blocks_to_skip,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Build flat pointer arrays for all sub-tensors × all block pairs.
|
|||
|
|
# sub-tensors = [layer0_key, layer0_value, layer1_key, layer1_value, ...]
|
|||
|
|
# Fully vectorized via numpy broadcasting (no Python loop).
|
|||
|
|
num_pairs = src_sub_block_count
|
|||
|
|
num_sub_tensors = len(self._block_size_in_bytes_arr)
|
|||
|
|
total = num_pairs * num_sub_tensors
|
|||
|
|
|
|||
|
|
# (num_sub_tensors, 1) + (1, num_pairs) * (num_sub_tensors, 1) -> (num_sub_tensors, num_pairs)
|
|||
|
|
bsz_col = self._block_size_in_bytes_arr[:, None] # (T, 1)
|
|||
|
|
all_src = (src_base_ptrs[:, None] + src_block_ids[None, :] * bsz_col).ravel()
|
|||
|
|
all_dst = (dst_base_ptrs[:, None] + dst_block_ids[None, :] * bsz_col).ravel()
|
|||
|
|
all_sizes = np.broadcast_to(bsz_col, (num_sub_tensors, num_pairs)).ravel().copy()
|
|||
|
|
|
|||
|
|
batch_src = torch.from_numpy(all_src)
|
|||
|
|
batch_dst = torch.from_numpy(all_dst)
|
|||
|
|
batch_sizes = torch.from_numpy(all_sizes)
|
|||
|
|
|
|||
|
|
start_event = self._get_event()
|
|||
|
|
end_event = self._get_event()
|
|||
|
|
|
|||
|
|
if is_d2h:
|
|||
|
|
# Wait for model computation to finish before reading NPU data
|
|||
|
|
stream.wait_stream(torch.npu.current_stream())
|
|||
|
|
if transfers:
|
|||
|
|
# Ensure this transfer starts only after the previous one completes
|
|||
|
|
last_transfer = transfers[-1]
|
|||
|
|
stream.wait_event(last_transfer.end_event)
|
|||
|
|
|
|||
|
|
with torch.npu.stream(stream):
|
|||
|
|
start_event.record(stream)
|
|||
|
|
if total > 0:
|
|||
|
|
direction = 0 if not is_d2h else 1
|
|||
|
|
torch.ops._C_ascend.swap_blocks_batch(batch_src, batch_dst, batch_sizes, direction)
|
|||
|
|
end_event.record(stream)
|
|||
|
|
|
|||
|
|
transfers.append(
|
|||
|
|
Transfer(
|
|||
|
|
job_id=job_id,
|
|||
|
|
stream=stream,
|
|||
|
|
start_event=start_event,
|
|||
|
|
end_event=end_event,
|
|||
|
|
num_bytes=src_sub_block_count * self._total_bytes_per_block,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def get_finished(self) -> list[TransferResult]:
|
|||
|
|
results: list[TransferResult] = []
|
|||
|
|
for transfers, transfer_type in [
|
|||
|
|
(self._d2h_transfers, ("NPU", "CPU")),
|
|||
|
|
(self._h2d_transfers, ("CPU", "NPU")),
|
|||
|
|
]:
|
|||
|
|
while transfers and transfers[0].end_event.query():
|
|||
|
|
transfer = transfers.popleft()
|
|||
|
|
transfer_time = transfer.start_event.elapsed_time(transfer.end_event) * 1e-3
|
|||
|
|
results.append(
|
|||
|
|
TransferResult(
|
|||
|
|
job_id=transfer.job_id,
|
|||
|
|
success=True,
|
|||
|
|
transfer_size=transfer.num_bytes,
|
|||
|
|
transfer_time=transfer_time,
|
|||
|
|
transfer_type=transfer_type,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
self._recycle_event(transfer.start_event)
|
|||
|
|
self._recycle_event(transfer.end_event)
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
def wait(self, job_ids: set[int]) -> None:
|
|||
|
|
"""
|
|||
|
|
Wait (block) until all specified transfer jobs are completed.
|
|||
|
|
"""
|
|||
|
|
for transfers in (self._d2h_transfers, self._h2d_transfers):
|
|||
|
|
for transfer in transfers:
|
|||
|
|
if transfer.job_id in job_ids:
|
|||
|
|
transfer.end_event.synchronize()
|