399 lines
14 KiB
Python
399 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from collections.abc import Mapping
|
|
|
|
import torch
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
ENABLE_ENV = "BI100_BLOCK_MAJOR_CPU_KV"
|
|
TRACE_ENV = "BI100_BLOCK_MAJOR_CPU_KV_TRACE"
|
|
CPU_OFFLOAD_ENV = "BI100_CPU_KV_OFFLOAD"
|
|
HYBRID_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
|
|
NUM_ATTENTION_LAYERS = 10
|
|
KV_PLANES = 2
|
|
ELEMENTS_PER_PLANE_BLOCK = 4096
|
|
STAGING_BLOCKS = 512
|
|
STAGING_BUFFER_COUNT = 2
|
|
BYTES_PER_BLOCK = (
|
|
NUM_ATTENTION_LAYERS * KV_PLANES * ELEMENTS_PER_PLANE_BLOCK * 2
|
|
)
|
|
GPU_STAGING_BYTES = STAGING_BLOCKS * STAGING_BUFFER_COUNT * BYTES_PER_BLOCK
|
|
|
|
|
|
def _strict_binary_selector(
|
|
name: str,
|
|
environ: Mapping[str, str] | None = None,
|
|
) -> bool:
|
|
source = os.environ if environ is None else environ
|
|
raw = source.get(name, "0")
|
|
if raw == "0":
|
|
return False
|
|
if raw == "1":
|
|
return True
|
|
raise RuntimeError(f"{name} must be exactly '0' or '1', got {raw!r}")
|
|
|
|
|
|
def block_major_cpu_kv_enabled(
|
|
environ: Mapping[str, str] | None = None,
|
|
) -> bool:
|
|
return _strict_binary_selector(ENABLE_ENV, environ)
|
|
|
|
|
|
def block_major_cpu_kv_trace_enabled(
|
|
environ: Mapping[str, str] | None = None,
|
|
) -> bool:
|
|
return _strict_binary_selector(TRACE_ENV, environ)
|
|
|
|
|
|
def _require_block_major_runtime(
|
|
environ: Mapping[str, str] | None = None,
|
|
) -> None:
|
|
source = os.environ if environ is None else environ
|
|
if source.get(CPU_OFFLOAD_ENV, "0") != "1":
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires {CPU_OFFLOAD_ENV}=1")
|
|
if source.get(HYBRID_ACCOUNTING_ENV, "legacy40") != "full_attention":
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires "
|
|
f"{HYBRID_ACCOUNTING_ENV}=full_attention")
|
|
|
|
|
|
def reserve_block_major_gpu_blocks(
|
|
num_gpu_blocks: int,
|
|
cache_block_size: int,
|
|
environ: Mapping[str, str] | None = None,
|
|
) -> int:
|
|
if (not isinstance(num_gpu_blocks, int)
|
|
or isinstance(num_gpu_blocks, bool)
|
|
or num_gpu_blocks < 0):
|
|
raise ValueError("num_gpu_blocks must be a non-negative integer")
|
|
if not block_major_cpu_kv_enabled(environ):
|
|
return num_gpu_blocks
|
|
|
|
_require_block_major_runtime(environ)
|
|
if cache_block_size != BYTES_PER_BLOCK:
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires cache block size "
|
|
f"{BYTES_PER_BLOCK}, got {cache_block_size}")
|
|
reserved_blocks = (
|
|
GPU_STAGING_BYTES + cache_block_size - 1
|
|
) // cache_block_size
|
|
remaining_blocks = num_gpu_blocks - reserved_blocks
|
|
if remaining_blocks <= 0:
|
|
raise RuntimeError(
|
|
"block-major GPU staging leaves no usable GPU KV blocks")
|
|
logger.info(
|
|
"[BI100 BLOCK KV] capacity reserve blocks=%d bytes=%d "
|
|
"profiled_blocks=%d usable_blocks=%d",
|
|
reserved_blocks,
|
|
GPU_STAGING_BYTES,
|
|
num_gpu_blocks,
|
|
remaining_blocks,
|
|
)
|
|
return remaining_blocks
|
|
|
|
|
|
def validate_block_mapping(
|
|
mapping: torch.Tensor,
|
|
source_limit: int,
|
|
destination_limit: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if not isinstance(mapping, torch.Tensor):
|
|
raise TypeError("block mapping must be a torch.Tensor")
|
|
if mapping.device.type != "cpu":
|
|
raise ValueError("block mapping must be on CPU")
|
|
if mapping.dtype != torch.int64:
|
|
raise ValueError("block mapping must use torch.int64")
|
|
if not mapping.is_contiguous():
|
|
raise ValueError("block mapping must be contiguous")
|
|
if mapping.dim() != 2 or mapping.shape[1] != 2:
|
|
raise ValueError("block mapping must have shape [N, 2]")
|
|
if source_limit <= 0 or destination_limit <= 0:
|
|
raise ValueError("block mapping limits must be positive")
|
|
|
|
sources: set[int] = set()
|
|
destinations: set[int] = set()
|
|
for row, pair in enumerate(mapping.tolist()):
|
|
source, destination = pair
|
|
if not 0 <= source < source_limit:
|
|
raise ValueError(
|
|
f"source block out of range at row {row}: {source}")
|
|
if not 0 <= destination < destination_limit:
|
|
raise ValueError(
|
|
f"destination block out of range at row {row}: "
|
|
f"{destination}")
|
|
if source in sources:
|
|
raise ValueError(f"duplicate source block: {source}")
|
|
if destination in destinations:
|
|
raise ValueError(f"duplicate destination block: {destination}")
|
|
sources.add(source)
|
|
destinations.add(destination)
|
|
|
|
return mapping[:, 0].contiguous(), mapping[:, 1].contiguous()
|
|
|
|
|
|
class BlockMajorCpuKVCache:
|
|
|
|
def __init__(
|
|
self,
|
|
gpu_cache: list[torch.Tensor],
|
|
num_cpu_blocks: int,
|
|
pin_memory: bool,
|
|
) -> None:
|
|
self._validate_gpu_cache(gpu_cache)
|
|
if block_major_cpu_kv_enabled():
|
|
_require_block_major_runtime()
|
|
if num_cpu_blocks <= 0:
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires a positive CPU block count")
|
|
if not pin_memory:
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires pinned CPU memory")
|
|
|
|
try:
|
|
from vllm import corex_block_major_kv_transfer as extension
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"block-major CoreX extension is unavailable") from exc
|
|
|
|
self.extension = extension
|
|
self.gpu_cache = gpu_cache
|
|
self.device = gpu_cache[0].device
|
|
self.dtype = gpu_cache[0].dtype
|
|
self.num_gpu_blocks = gpu_cache[0].shape[1]
|
|
self.num_cpu_blocks = num_cpu_blocks
|
|
self.trace_enabled = block_major_cpu_kv_trace_enabled()
|
|
|
|
self.cpu_pool = torch.zeros(
|
|
(
|
|
num_cpu_blocks,
|
|
NUM_ATTENTION_LAYERS,
|
|
KV_PLANES,
|
|
ELEMENTS_PER_PLANE_BLOCK,
|
|
),
|
|
dtype=self.dtype,
|
|
device="cpu",
|
|
pin_memory=True,
|
|
)
|
|
if not self.cpu_pool.is_pinned():
|
|
raise RuntimeError("block-major CPU pool is not pinned")
|
|
|
|
# Preserve the public CacheEngine shape without allocating a second
|
|
# layer-major CPU cache. Transfer methods use cpu_pool directly.
|
|
self.layer_views = [
|
|
self.cpu_pool[:, layer, :, :].permute(1, 0, 2)
|
|
for layer in range(NUM_ATTENTION_LAYERS)
|
|
]
|
|
self.cpu_staging = [
|
|
torch.empty(
|
|
(
|
|
STAGING_BLOCKS,
|
|
NUM_ATTENTION_LAYERS,
|
|
KV_PLANES,
|
|
ELEMENTS_PER_PLANE_BLOCK,
|
|
),
|
|
dtype=self.dtype,
|
|
device="cpu",
|
|
pin_memory=True,
|
|
)
|
|
for _ in range(STAGING_BUFFER_COUNT)
|
|
]
|
|
if not all(staging.is_pinned() for staging in self.cpu_staging):
|
|
raise RuntimeError("block-major CPU staging is not pinned")
|
|
|
|
with torch.cuda.device(self.device):
|
|
self.gpu_staging = [
|
|
torch.empty_like(staging, device=self.device)
|
|
for staging in self.cpu_staging
|
|
]
|
|
self.events = [
|
|
torch.cuda.Event(enable_timing=False)
|
|
for _ in range(STAGING_BUFFER_COUNT)
|
|
]
|
|
self.error_flag = torch.zeros(
|
|
1, dtype=torch.int32, device=self.device)
|
|
|
|
logger.info(
|
|
"[BI100 BLOCK KV] enabled device=%s gpu_blocks=%d cpu_blocks=%d "
|
|
"layers=%d block_bytes=%d staging_blocks=%d staging_buffers=%d",
|
|
self.device,
|
|
self.num_gpu_blocks,
|
|
self.num_cpu_blocks,
|
|
NUM_ATTENTION_LAYERS,
|
|
BYTES_PER_BLOCK,
|
|
STAGING_BLOCKS,
|
|
STAGING_BUFFER_COUNT,
|
|
)
|
|
|
|
@staticmethod
|
|
def _validate_gpu_cache(gpu_cache: list[torch.Tensor]) -> None:
|
|
if len(gpu_cache) != NUM_ATTENTION_LAYERS:
|
|
raise RuntimeError(
|
|
f"{ENABLE_ENV}=1 requires exactly "
|
|
f"{NUM_ATTENTION_LAYERS} GPU attention caches, got "
|
|
f"{len(gpu_cache)}")
|
|
first = gpu_cache[0]
|
|
if first.device.type != "cuda":
|
|
raise RuntimeError("block-major GPU cache must be on CUDA")
|
|
if first.dtype != torch.float16:
|
|
raise RuntimeError("block-major GPU cache must use float16")
|
|
if (first.dim() != 3 or first.shape[0] != KV_PLANES
|
|
or first.shape[2] != ELEMENTS_PER_PLANE_BLOCK):
|
|
raise RuntimeError(
|
|
"block-major GPU cache must have shape [2, blocks, 4096]")
|
|
if not first.is_contiguous():
|
|
raise RuntimeError("block-major GPU cache must be contiguous")
|
|
|
|
for layer, tensor in enumerate(gpu_cache):
|
|
if tensor.device != first.device:
|
|
raise RuntimeError(
|
|
f"GPU cache layer {layer} is on a different device")
|
|
if tensor.dtype != first.dtype or tensor.shape != first.shape:
|
|
raise RuntimeError(
|
|
f"GPU cache layer {layer} has inconsistent geometry")
|
|
if not tensor.is_contiguous():
|
|
raise RuntimeError(
|
|
f"GPU cache layer {layer} is not contiguous")
|
|
|
|
def _to_gpu_ids(self, block_ids: torch.Tensor) -> torch.Tensor:
|
|
return block_ids.to(
|
|
device=self.device,
|
|
dtype=torch.int32,
|
|
non_blocking=False,
|
|
)
|
|
|
|
@staticmethod
|
|
def _chunks(
|
|
source: torch.Tensor,
|
|
destination: torch.Tensor,
|
|
gpu_ids: torch.Tensor,
|
|
):
|
|
for start in range(0, source.numel(), STAGING_BLOCKS):
|
|
end = min(start + STAGING_BLOCKS, source.numel())
|
|
yield (
|
|
source[start:end],
|
|
destination[start:end],
|
|
gpu_ids[start:end],
|
|
end - start,
|
|
)
|
|
|
|
def _begin(self) -> None:
|
|
self.error_flag.zero_()
|
|
|
|
def _finish(
|
|
self,
|
|
direction: str,
|
|
block_count: int,
|
|
started: float | None,
|
|
) -> None:
|
|
# check_error performs the final stream synchronization. This also
|
|
# makes every staging slot safe to reuse in the next CacheEngine call.
|
|
self.extension.check_error(self.error_flag)
|
|
if started is not None:
|
|
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
|
logger.info(
|
|
"[BI100 BLOCK KV TRACE] direction=%s blocks=%d bytes=%d "
|
|
"elapsed_ms=%.3f",
|
|
direction,
|
|
block_count,
|
|
block_count * BYTES_PER_BLOCK,
|
|
elapsed_ms,
|
|
)
|
|
|
|
def swap_out(self, mapping: torch.Tensor) -> None:
|
|
started = time.perf_counter() if self.trace_enabled else None
|
|
source_gpu, destination_cpu = validate_block_mapping(
|
|
mapping,
|
|
source_limit=self.num_gpu_blocks,
|
|
destination_limit=self.num_cpu_blocks,
|
|
)
|
|
block_count = source_gpu.numel()
|
|
if block_count == 0:
|
|
return
|
|
source_gpu_ids = self._to_gpu_ids(source_gpu)
|
|
|
|
self._begin()
|
|
pending: tuple[int, torch.Tensor, int] | None = None
|
|
for index, (_, destination, gpu_ids, count) in enumerate(
|
|
self._chunks(
|
|
source_gpu, destination_cpu, source_gpu_ids)):
|
|
slot = index % STAGING_BUFFER_COUNT
|
|
self.extension.pack(
|
|
self.gpu_cache,
|
|
gpu_ids,
|
|
self.gpu_staging[slot],
|
|
self.error_flag,
|
|
count,
|
|
)
|
|
self.cpu_staging[slot][:count].copy_(
|
|
self.gpu_staging[slot][:count],
|
|
non_blocking=True,
|
|
)
|
|
self.events[slot].record()
|
|
if pending is not None:
|
|
pending_slot, pending_destination, pending_count = pending
|
|
self.events[pending_slot].synchronize()
|
|
self.extension.cpu_scatter(
|
|
self.cpu_staging[pending_slot],
|
|
self.cpu_pool,
|
|
pending_destination,
|
|
pending_count,
|
|
)
|
|
pending = (slot, destination, count)
|
|
|
|
if pending is not None:
|
|
pending_slot, pending_destination, pending_count = pending
|
|
self.events[pending_slot].synchronize()
|
|
self.extension.cpu_scatter(
|
|
self.cpu_staging[pending_slot],
|
|
self.cpu_pool,
|
|
pending_destination,
|
|
pending_count,
|
|
)
|
|
self._finish("d2h", block_count, started)
|
|
|
|
def swap_in(self, mapping: torch.Tensor) -> None:
|
|
started = time.perf_counter() if self.trace_enabled else None
|
|
source_cpu, destination_gpu = validate_block_mapping(
|
|
mapping,
|
|
source_limit=self.num_cpu_blocks,
|
|
destination_limit=self.num_gpu_blocks,
|
|
)
|
|
block_count = source_cpu.numel()
|
|
if block_count == 0:
|
|
return
|
|
destination_gpu_ids = self._to_gpu_ids(destination_gpu)
|
|
|
|
self._begin()
|
|
for index, (source, _, gpu_ids, count) in enumerate(
|
|
self._chunks(
|
|
source_cpu, destination_gpu, destination_gpu_ids)):
|
|
slot = index % STAGING_BUFFER_COUNT
|
|
if index >= STAGING_BUFFER_COUNT:
|
|
self.events[slot].synchronize()
|
|
self.extension.cpu_gather(
|
|
self.cpu_pool,
|
|
source,
|
|
self.cpu_staging[slot],
|
|
count,
|
|
)
|
|
self.gpu_staging[slot][:count].copy_(
|
|
self.cpu_staging[slot][:count],
|
|
non_blocking=True,
|
|
)
|
|
self.extension.scatter(
|
|
self.gpu_staging[slot],
|
|
gpu_ids,
|
|
self.gpu_cache,
|
|
self.error_flag,
|
|
count,
|
|
)
|
|
self.events[slot].record()
|
|
self._finish("h2d", block_count, started)
|