diff --git a/docs/recipes/models.md b/docs/recipes/models.md index 2be2a886e..ecc652650 100644 --- a/docs/recipes/models.md +++ b/docs/recipes/models.md @@ -203,6 +203,9 @@ speculative decoding and paged-cache groups are both active — and prefix cachi stays on by default. Add `--enable-metrics` to read `Decoded Tok/Iter` and the speculative accept rate from the run summary. +DeepSeek V4 MTP currently requires `--disable-kvstore`; grouped paged-cache +loadback does not yet restore the independent draft KV pool. + ## Tuning Order 1. Set model ID, trust policy, tokenizer mode, and served model name. diff --git a/python/tokenspeed/runtime/cache/deepseek_v4_cache_host.py b/python/tokenspeed/runtime/cache/deepseek_v4_cache_host.py new file mode 100644 index 000000000..8349830c4 --- /dev/null +++ b/python/tokenspeed/runtime/cache/deepseek_v4_cache_host.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import math + +import torch +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + V4_SWA_KV_GROUP_ID, + v4_compressed_kv_group_id, + v4_compressor_state_group_id, +) +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, + _deepseek_v4_cache_group_page_bytes, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +def _allocate_host_group_pages( + *, + device_counts: dict[str, int], + page_bytes: dict[str, int], + host_ratio: float, + host_size_gb: int, + host_budget_bytes: int | None = None, +) -> dict[str, int]: + active_groups = [ + group_id + for group_id, bytes_per_page in page_bytes.items() + if bytes_per_page > 0 and int(device_counts.get(group_id, 0)) > 0 + ] + if not active_groups: + return {} + + if host_size_gb <= 0: + ratio_counts = { + group_id: max(2, int(math.ceil(device_counts[group_id] * host_ratio))) + for group_id in active_groups + } + if host_budget_bytes is None: + return ratio_counts + + requested_bytes = sum( + ratio_counts[group_id] * page_bytes[group_id] for group_id in active_groups + ) + if requested_bytes <= host_budget_bytes: + return ratio_counts + budget_bytes = int(host_budget_bytes) + else: + budget_bytes = int(host_size_gb * 1e9) + + min_pages_per_group = 2 # page 0 is the dummy page; keep one usable page. + min_bytes = sum( + min_pages_per_group * page_bytes[group_id] for group_id in active_groups + ) + if budget_bytes < min_bytes: + raise ValueError( + "DeepSeek V4 KVStore host budget is too small to allocate one usable " + "page plus the dummy page per group: " + f"budget={budget_bytes} bytes minimum={min_bytes} bytes" + ) + + counts = {group_id: min_pages_per_group for group_id in active_groups} + remaining_budget = budget_bytes - min_bytes + desired_bytes = { + group_id: int(device_counts[group_id]) * int(page_bytes[group_id]) + for group_id in active_groups + } + total_desired = max(sum(desired_bytes.values()), 1) + for group_id in active_groups: + extra_budget = remaining_budget * desired_bytes[group_id] // total_desired + counts[group_id] += extra_budget // int(page_bytes[group_id]) + return counts + + +class DeepseekV4TokenToKVPoolHost: + """Registered host mirror for DeepSeek V4 paged-cache groups.""" + + def __init__( + self, + device_pool: DeepseekV4TokenToKVPool, + host_to_device_ratio: float, + host_size_gb: int, + layout: str = "layer_first", + device: str = "cpu", + register_host: bool = True, + host_budget_bytes: int | None = None, + ) -> None: + if layout != "layer_first": + raise ValueError("DeepSeek V4 KVStore host pool only supports layer_first") + if host_to_device_ratio <= 0 and host_size_gb <= 0: + raise ValueError("host_to_device_ratio must be positive") + + self.device_pool = device_pool + self.layout = layout + self.device = device + self.layer_num = int(device_pool.layer_num) + + self.paged_cache_group_page_bytes = _deepseek_v4_cache_group_page_bytes( + device_pool.layout, + device_pool.paged_cache_group_specs, + self.layer_num, + ) + self.paged_cache_group_page_counts = _allocate_host_group_pages( + device_counts=device_pool.paged_cache_group_page_counts, + page_bytes=self.paged_cache_group_page_bytes, + host_ratio=float(host_to_device_ratio), + host_size_gb=int(host_size_gb), + host_budget_bytes=host_budget_bytes if host_size_gb <= 0 else None, + ) + self._shadow_page_num = self._compute_shadow_page_num(device_pool) + self.total_bytes = int( + sum( + self.paged_cache_group_page_counts[group_id] + * self.paged_cache_group_page_bytes[group_id] + for group_id in self.paged_cache_group_page_counts + ) + ) + self._check_host_memory(self.total_bytes, host_budget_bytes) + + def alloc_like_group( + source: torch.Tensor | None, + group_id: str, + ) -> torch.Tensor | None: + if source is None: + return None + pages = int(self.paged_cache_group_page_counts[group_id]) + use_pin_memory = bool( + device == "cpu" and not register_host and torch.cuda.is_available() + ) + tensor = torch.empty( + (pages, *source.shape[1:]), + dtype=source.dtype, + device=device, + pin_memory=use_pin_memory, + ) + if register_host: + current_platform().register_host_tensor_for_gpu_access(tensor) + return tensor + + self.swa_kv_buffer = [ + alloc_like_group(buf, V4_SWA_KV_GROUP_ID) + for buf in device_pool.swa_kv_buffer + ] + self.compressed_kv_buffer: list[torch.Tensor | None] = [] + self.compressor_state_buffer: list[torch.Tensor | None] = [] + self.indexer_kv_buffer: list[torch.Tensor | None] = [] + self.indexer_state_buffer: list[torch.Tensor | None] = [] + + for layer_id, ratio in enumerate(device_pool.layout.layer_ratio): + compressed_group_id = v4_compressed_kv_group_id(ratio) + state_group_id = v4_compressor_state_group_id(ratio) + self.compressed_kv_buffer.append( + alloc_like_group( + device_pool.compressed_kv_buffer[layer_id], + compressed_group_id, + ) + if ratio > 1 + else None + ) + self.compressor_state_buffer.append( + alloc_like_group( + device_pool.compressor_state_buffer[layer_id], + state_group_id, + ) + if ratio > 1 + else None + ) + self.indexer_kv_buffer.append( + alloc_like_group( + device_pool.indexer_kv_buffer[layer_id], + compressed_group_id, + ) + if ratio == 4 + else None + ) + self.indexer_state_buffer.append( + alloc_like_group( + device_pool.indexer_state_buffer[layer_id], + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + ) + if ratio == 4 + else None + ) + + logger.info( + "Allocating %.2f GB host memory for DeepSeek V4 KVStore. " + "group_pages=%s group_page_bytes=%s shadow_page_num=%s layout=%s", + self.total_bytes / 1e9, + self.paged_cache_group_page_counts, + self.paged_cache_group_page_bytes, + self._shadow_page_num, + self.layout, + ) + + def _compute_shadow_page_num(self, device_pool: DeepseekV4TokenToKVPool) -> int: + """Token-page capacity used only for scheduler HostResource pinning.""" + usable_token_pages_by_history_group: list[int] = [] + token_page_size = max(1, int(device_pool.page_size)) + for spec in device_pool.paged_cache_group_specs: + if getattr(spec, "family", "history") != "history": + continue + group_pages = int(self.paged_cache_group_page_counts.get(spec.group_id, 0)) + usable_group_pages = max(0, group_pages - 1) # page 0 is the dummy page. + raw_tokens = ( + usable_group_pages + * int(spec.rows_per_page) + * int(spec.entry_stride_tokens) + ) + usable_token_pages_by_history_group.append( + int(math.ceil(raw_tokens / token_page_size)) + ) + if not usable_token_pages_by_history_group: + return 0 + usable_token_pages = min(usable_token_pages_by_history_group) + return usable_token_pages + 1 if usable_token_pages > 0 else 0 + + @staticmethod + def _check_host_memory(requested_bytes: int, host_budget_bytes: int | None) -> None: + if host_budget_bytes is None: + return + if requested_bytes > host_budget_bytes: + raise ValueError( + "Not enough host memory available for DeepSeek V4 KVStore. " + f"Requesting {requested_bytes / 1e9:.2f} GB but only have " + f"{host_budget_bytes / 1e9:.2f} GB per-rank budget after " + "cgroup-aware host reservation. Please reduce kvstore size." + ) + + @property + def page_num(self) -> int: + return self._shadow_page_num diff --git a/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py b/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py index b9e100d3a..aa5a03406 100644 --- a/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py +++ b/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py @@ -101,7 +101,7 @@ class FlatMemoryExecutor: """ # EventLoop keys per-op inflight accounting off this: flat loadbacks are - # acked (LoadBackDoneEvent), radix loadbacks are not. + # acknowledged with one LoadBackDoneEvent per op. emits_loadback_acks = True def __init__(self, device_pool, *, host_ratio: float, host_size_gb: float): diff --git a/python/tokenspeed/runtime/cache/executor/host_executor.py b/python/tokenspeed/runtime/cache/executor/host_executor.py index 2830dafde..638a3f1d6 100644 --- a/python/tokenspeed/runtime/cache/executor/host_executor.py +++ b/python/tokenspeed/runtime/cache/executor/host_executor.py @@ -22,6 +22,7 @@ from __future__ import annotations +import logging from collections import OrderedDict from collections.abc import Iterable from typing import NamedTuple @@ -31,11 +32,18 @@ from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool from tokenspeed.runtime.cache.transfer.pool import CachePool -from tokenspeed.runtime.cache.transfer.types import CacheKind, Location, TransferUnit +from tokenspeed.runtime.cache.transfer.types import ( + PAGED_CACHE_KIND, + CacheKind, + Location, + PagedCacheTransferUnit, + TransferUnit, +) from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode from tokenspeed.runtime.utils import get_colorful_logger, get_device_module logger = get_colorful_logger(__name__) +_DEBUG = logging.DEBUG device_module = get_device_module() CONCURRENT_WRITEBACK_BLOCK_QUOTA = 2 @@ -101,6 +109,28 @@ def _ordered_unique(values: Iterable[int]) -> list[int]: return result +def _paged_queue_debug_summary(units: list[PagedCacheTransferUnit]): + pages = 0 + groups: dict[str, int] = {} + transfers = 0 + for unit in units: + for transfer in unit.transfers: + transfers += 1 + group_id = str(getattr(transfer, "group_id", "unknown")) + transfer_pages = len( + { + (int(src), int(dst)) + for src, dst in zip( + getattr(transfer, "src_pages", []), + getattr(transfer, "dst_pages", []), + ) + } + ) + pages += transfer_pages + groups[group_id] = groups.get(group_id, 0) + transfer_pages + return pages, transfers, groups + + class _Ack(NamedTuple): finish_event: object # device_module.Event op_ids: list[int] @@ -118,32 +148,46 @@ def __init__( draft_host_pool=None, draft_layer_num: int = 0, pools: list[CachePool] | None = None, + paged_pool=None, ): self.io_backend = io_backend if pools is None: - if ( - page_size is None - or device_pool is None - or host_pool is None - or layer_num is None - ): - raise ValueError("HostExecutor requires either pools or KV pool inputs") - pools = [ - KVCachePool( - device_pool=device_pool, - host_pool=host_pool, - io_backend=io_backend, - layer_num=layer_num, - draft_device_pool=draft_device_pool, - draft_host_pool=draft_host_pool, - draft_layer_num=draft_layer_num, - ) - ] - if not pools: + has_kv_inputs = ( + page_size is not None + and device_pool is not None + and host_pool is not None + and layer_num is not None + ) + if not has_kv_inputs: + if paged_pool is None: + raise ValueError( + "HostExecutor requires either pools, KV pool inputs, " + "or a paged_pool" + ) + pools = [] + else: + pools = [ + KVCachePool( + device_pool=device_pool, + host_pool=host_pool, + io_backend=io_backend, + layer_num=layer_num, + draft_device_pool=draft_device_pool, + draft_host_pool=draft_host_pool, + draft_layer_num=draft_layer_num, + ) + ] + if not pools and paged_pool is None: raise ValueError("HostExecutor requires at least one cache pool") self.pools = {CacheKind(pool.kind): pool for pool in pools} - self.device = next(iter(self.pools.values())).device + self.paged_pool = paged_pool + self.emits_loadback_acks = paged_pool is not None + self.device = ( + next(iter(self.pools.values())).device + if self.pools + else self.paged_pool.device + ) write_priority, load_priority = _cache_stream_priorities() self.write_stream = _new_cache_stream(write_priority) @@ -156,6 +200,8 @@ def __init__( self.load_queues: dict[CacheKind, list[TransferUnit]] = { kind: [] for kind in self.pools } + self.paged_write_queue: list[PagedCacheTransferUnit] = [] + self.paged_load_queue: list[PagedCacheTransferUnit] = [] self.ack_write_queue: list[_Ack] = [] self.ack_load_queue: list[_Ack] = [] @@ -164,9 +210,13 @@ def __init__( self._counters = { kind: pool.get_layer_done_counter() for kind, pool in self.pools.items() } + self._paged_counter = ( + paged_pool.get_layer_done_counter() if paged_pool is not None else None + ) self._producer_map: dict[CacheKind, OrderedDict[int, int]] = { kind: OrderedDict() for kind in self.pools } + self._paged_producer_map: OrderedDict[int, int] = OrderedDict() self._producer_map_limit = 1024 def enqueue_writeback( @@ -241,10 +291,45 @@ def enqueue_loadback( ) ) + def enqueue_paged_cache_writeback( + self, + op_id: int, + transfers: list, + is_retract: bool = False, + ) -> None: + if self.paged_pool is None: + if transfers: + raise ValueError("paged-cache writeback requires a paged_pool") + self.completed_writebacks.append(op_id) + return + if not transfers: + self.completed_writebacks.append(op_id) + return + self.paged_write_queue.append( + PagedCacheTransferUnit( + op_id=int(op_id), + transfers=list(transfers), + is_retract=is_retract, + ) + ) + + def enqueue_paged_cache_loadback(self, op_id: int, transfers: list) -> None: + if self.paged_pool is None: + if transfers: + raise ValueError("paged-cache loadback requires a paged_pool") + return + if not transfers: + return + self.paged_load_queue.append( + PagedCacheTransferUnit(op_id=int(op_id), transfers=list(transfers)) + ) + def flush(self) -> None: - throttle_writeback = self._has_work(self.load_queues) and not any( + has_load_work = self._has_work(self.load_queues) or bool(self.paged_load_queue) + has_retract_writeback = any( unit.is_retract for units in self.write_queues.values() for unit in units - ) + ) or any(unit.is_retract for unit in self.paged_write_queue) + throttle_writeback = has_load_work and not has_retract_writeback writeback_block_quota = ( CONCURRENT_WRITEBACK_BLOCK_QUOTA if throttle_writeback else None ) @@ -257,7 +342,7 @@ def flush(self) -> None: self._writeback_block_quota = previous_writeback_block_quota def _start_writing(self) -> None: - if not self._has_work(self.write_queues): + if not self._has_work(self.write_queues) and not self.paged_write_queue: return start_event = device_module.Event() @@ -279,13 +364,68 @@ def _start_writing(self) -> None: self._record_if_cuda(src_indices, self.write_stream) self._record_if_cuda(dst_indices, self.write_stream) op_ids.extend(unit.op_id for unit in units) + if self.paged_write_queue: + assert self.paged_pool is not None + transfers = [ + transfer + for unit in self.paged_write_queue + for transfer in unit.transfers + ] + prepare_transfers = getattr( + self.paged_pool, + "prepare_paged_transfers", + None, + ) + write_prepared = getattr( + self.paged_pool, + "writeback_prepared_paged", + None, + ) + prepared_transfers = ( + prepare_transfers(transfers) + if prepare_transfers is not None and write_prepared is not None + else None + ) + if logger.isEnabledFor(_DEBUG): + paged_pages, paged_transfers, paged_groups = ( + _paged_queue_debug_summary(self.paged_write_queue) + ) + prepared_count = ( + len(prepared_transfers) + if prepared_transfers is not None + else "n/a" + ) + prepared_spans = ( + sum( + int(getattr(transfer, "span_count", 0)) + for transfer in prepared_transfers + ) + if prepared_transfers is not None + else "n/a" + ) + logger.debug( + "[cache_op][paged_l2] writeback submit units=%s " + "transfers=%s coalesced=%s spans=%s pages=%s groups=%s", + len(self.paged_write_queue), + paged_transfers, + prepared_count, + prepared_spans, + paged_pages, + paged_groups, + ) + if prepared_transfers is not None and write_prepared is not None: + write_prepared(prepared_transfers) + else: + self.paged_pool.writeback_paged(transfers) + op_ids.extend(unit.op_id for unit in self.paged_write_queue) finish_event.record() self._clear_queues(self.write_queues) + self.paged_write_queue.clear() self.ack_write_queue.append(_Ack(finish_event, _ordered_unique(op_ids))) def _start_loading(self) -> None: - if not self._has_work(self.load_queues): + if not self._has_work(self.load_queues) and not self.paged_load_queue: return assert ( not get_is_capture_mode() @@ -339,7 +479,103 @@ def _start_loading(self) -> None: while len(producer_map) > self._producer_map_limit: producer_map.popitem(last=False) + if self.paged_load_queue: + assert self.paged_pool is not None + assert self._paged_counter is not None + transfers = [ + transfer + for unit in self.paged_load_queue + for transfer in unit.transfers + ] + prepare_transfers = getattr( + self.paged_pool, + "prepare_paged_transfers", + None, + ) + load_prepared = getattr( + self.paged_pool, + "loadback_prepared_paged", + None, + ) + load_prepared_range = getattr( + self.paged_pool, + "loadback_prepared_paged_range", + None, + ) + prepared_transfers = ( + prepare_transfers(transfers) + if prepare_transfers is not None and load_prepared is not None + else None + ) + num_layers = self.paged_pool.num_layers() + layer_chunk_size = 1 + if prepared_transfers is not None and load_prepared_range is not None: + layer_chunk_size = max( + 1, + int(getattr(self.paged_pool, "loadback_layer_chunk_size", 1)), + ) + layer_chunks = (num_layers + layer_chunk_size - 1) // layer_chunk_size + if logger.isEnabledFor(_DEBUG): + paged_pages, paged_transfers, paged_groups = ( + _paged_queue_debug_summary(self.paged_load_queue) + ) + prepared_count = ( + len(prepared_transfers) + if prepared_transfers is not None + else "n/a" + ) + prepared_spans = ( + sum( + int(getattr(transfer, "span_count", 0)) + for transfer in prepared_transfers + ) + if prepared_transfers is not None + else "n/a" + ) + logger.debug( + "[cache_op][paged_l2] loadback submit units=%s " + "transfers=%s coalesced=%s spans=%s pages=%s groups=%s " + "layers=%s layer_chunk=%s layer_chunks=%s", + len(self.paged_load_queue), + paged_transfers, + prepared_count, + prepared_spans, + paged_pages, + paged_groups, + num_layers, + layer_chunk_size, + layer_chunks, + ) + producer_id = self._paged_counter.update_producer() + producer_event = self._paged_counter.events[producer_id] + producer_event.start_event.record() + producer_event.start_event.wait(self.load_stream) + if prepared_transfers is not None and load_prepared_range is not None: + for layer_start in range(0, num_layers, layer_chunk_size): + layer_end = min(layer_start + layer_chunk_size, num_layers) + load_prepared_range( + prepared_transfers, + layer_start, + layer_end, + ) + for layer_index in range(layer_start, layer_end): + producer_event.complete(layer_index) + else: + for layer_index in range(num_layers): + if prepared_transfers is not None and load_prepared is not None: + load_prepared(prepared_transfers, layer_index) + else: + self.paged_pool.loadback_paged(transfers, layer_index) + producer_event.complete(layer_index) + op_ids = _ordered_unique(unit.op_id for unit in self.paged_load_queue) + self.ack_load_queue.append(_Ack(producer_event.finish_event, op_ids)) + for op_id in op_ids: + self._paged_producer_map[op_id] = producer_id + while len(self._paged_producer_map) > self._producer_map_limit: + self._paged_producer_map.popitem(last=False) + self._clear_queues(self.load_queues) + self.paged_load_queue.clear() @staticmethod def _has_work(queues: dict[CacheKind, list[TransferUnit]]) -> bool: @@ -470,6 +706,12 @@ def _poll_load_acks(self) -> list: for ack in self.ack_load_queue: if not ack.finish_event.query(): remaining.append(ack) + elif self.emits_loadback_acks: + for op_id in ack.op_ids: + evt = Cache.LoadBackDoneEvent() + evt.op_id = op_id + evt.success = True + results.append(evt) self.ack_load_queue[:] = remaining return results @@ -480,6 +722,8 @@ def get_producer_index( kind = CacheKind.KV op_id = int(kind_or_op_id) else: + if self._is_paged_kind(kind_or_op_id): + return self._paged_producer_map.pop(int(op_id), None) kind = CacheKind(kind_or_op_id) return self._producer_map[kind].pop(int(op_id), None) @@ -492,9 +736,18 @@ def set_consumer( kind = CacheKind.KV producer_index = kind_or_producer_index else: + if self._is_paged_kind(kind_or_producer_index): + assert self._paged_counter is not None + self._paged_counter.set_consumer(producer_index) + return kind = CacheKind(kind_or_producer_index) self._counters[kind].set_consumer(producer_index) + def _is_paged_kind(self, kind: CacheKind | str | int | Iterable[int]) -> bool: + if self.paged_pool is None: + return False + return str(kind) == PAGED_CACHE_KIND + def shutdown(self) -> None: self.write_stream.synchronize() self.load_stream.synchronize() @@ -502,15 +755,24 @@ def shutdown(self) -> None: shutdown = getattr(pool, "shutdown", None) if shutdown is not None: shutdown() + if self.paged_pool is not None: + shutdown = getattr(self.paged_pool, "shutdown", None) + if shutdown is not None: + shutdown() def reset(self) -> None: self.write_stream.synchronize() self.load_stream.synchronize() self._clear_queues(self.write_queues) self._clear_queues(self.load_queues) + self.paged_write_queue.clear() + self.paged_load_queue.clear() self.ack_write_queue.clear() self.ack_load_queue.clear() for producer_map in self._producer_map.values(): producer_map.clear() + self._paged_producer_map.clear() for counter in self._counters.values(): counter.reset() + if self._paged_counter is not None: + self._paged_counter.reset() diff --git a/python/tokenspeed/runtime/cache/executor/memory_executor.py b/python/tokenspeed/runtime/cache/executor/memory_executor.py index 98c037bc6..5ea6fc77f 100644 --- a/python/tokenspeed/runtime/cache/executor/memory_executor.py +++ b/python/tokenspeed/runtime/cache/executor/memory_executor.py @@ -22,11 +22,15 @@ from __future__ import annotations +import logging from collections.abc import Iterable from dataclasses import dataclass from tokenspeed_scheduler import Cache +from tokenspeed.runtime.cache.deepseek_v4_cache_host import ( + DeepseekV4TokenToKVPoolHost, +) from tokenspeed.runtime.cache.executor.host_executor import HostExecutor from tokenspeed.runtime.cache.executor.storage_executor import StorageExecutor from tokenspeed.runtime.cache.kv_cache_host import ( @@ -36,15 +40,49 @@ get_available_host_memory_bytes, ) from tokenspeed.runtime.cache.mamba_cache_host import MambaPoolHost +from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool from tokenspeed.runtime.cache.transfer.mamba_pool import MambaCachePool from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, +) from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool from tokenspeed.runtime.utils import get_colorful_logger logger = get_colorful_logger(__name__) +_DEBUG = logging.DEBUG + + +def _count_page_pair_spans(src_pages, dst_pages) -> tuple[int, int]: + pairs = sorted({(int(src), int(dst)) for src, dst in zip(src_pages, dst_pages)}) + if not pairs: + return 0, 0 + spans = 1 + prev_src, prev_dst = pairs[0] + for src, dst in pairs[1:]: + if src != prev_src + 1 or dst != prev_dst + 1: + spans += 1 + prev_src, prev_dst = src, dst + return len(pairs), spans + + +def _paged_transfer_debug_summary(paged_transfers) -> tuple[int, int, dict[str, int]]: + pages = 0 + spans = 0 + groups: dict[str, int] = {} + for transfer in paged_transfers or []: + group_id = str(getattr(transfer, "group_id", "unknown")) + transfer_pages, transfer_spans = _count_page_pair_spans( + getattr(transfer, "src_pages", []), + getattr(transfer, "dst_pages", []), + ) + pages += transfer_pages + spans += transfer_spans + groups[group_id] = groups.get(group_id, 0) + transfer_pages + return pages, spans, groups @dataclass(slots=True) @@ -124,107 +162,140 @@ def __init__( mamba_pool=None, ): self.page_size = config.page_size - kv_pool_types = (DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool) + kv_pool_types = ( + DSATokenToKVPool, + DeepseekV4TokenToKVPool, + MHATokenToKVPool, + MLATokenToKVPool, + ) # Unwrap LayerMappedKVPool (hybrid GDN models) to get the inner MHA pool. actual_pool = device_pool if hasattr(device_pool, "inner") and not isinstance(device_pool, kv_pool_types): actual_pool = device_pool.inner + self.paged_cache_pool = None actual_draft_pool = None - if draft_device_pool is not None: - actual_draft_pool = draft_device_pool - if hasattr(draft_device_pool, "inner") and not isinstance( - draft_device_pool, kv_pool_types - ): - actual_draft_pool = draft_device_pool.inner - if not isinstance(actual_draft_pool, kv_pool_types): - raise ValueError( - f"draft_device_pool only supports DSA, MHA and MLA, " - f"got {type(actual_draft_pool)}" + if isinstance(actual_pool, DeepseekV4TokenToKVPool): + if config.storage_backend is not None: + raise NotImplementedError( + "DeepSeek V4 KVStore currently supports L2 host memory only; " + "L3 storage backends are out of scope." ) - - host_size_tokens = 0 - if config.host_size_gb == 0: - target_size_per_token = _pool_size_per_token(actual_pool) - draft_size_per_token = ( - _pool_size_per_token(actual_draft_pool) - if actual_draft_pool is not None - else 0 - ) - combined_size_per_token = target_size_per_token + draft_size_per_token reserve_bytes = int(config.host_reserve_gb * (1024**3)) - available_bytes, _, cgroup_available = get_available_host_memory_bytes( - reserve_bytes - ) - requested_tokens = int(actual_pool.size * config.host_ratio) - host_size_tokens = _auto_capped_host_size_tokens( - requested_tokens=requested_tokens, - page_size=config.page_size, - size_per_token=combined_size_per_token, - available_host_memory_bytes=available_bytes, - host_parallel_count=config.host_parallel_count, - ) - if host_size_tokens > 0: - capped_tokens = _aligned_token_count(host_size_tokens, config.page_size) - requested_tokens_aligned = _aligned_token_count( - requested_tokens, config.page_size - ) - logger.warning( - "Capping KVStore host pool for cgroup budget: " - "tokens %s -> %s, total bytes %.2f GB -> %.2f GB " - "(parallel_count=%s, available=%.2f GB, cgroup_available=%s)", - requested_tokens_aligned, - capped_tokens, - requested_tokens_aligned * combined_size_per_token / 1e9, - capped_tokens * combined_size_per_token / 1e9, - config.host_parallel_count, - available_bytes / 1e9, - ( - f"{cgroup_available / 1e9:.2f} GB" - if cgroup_available is not None - else "unlimited" - ), - ) - - # DSA subclasses MLA, so it must be matched before the MLA branch. - if isinstance(actual_pool, DSATokenToKVPool): - self.host_pool = DSATokenToKVPoolHost( + available_bytes, _, _ = get_available_host_memory_bytes(reserve_bytes) + host_budget_bytes = available_bytes // max(config.host_parallel_count, 1) + self.host_pool = DeepseekV4TokenToKVPoolHost( actual_pool, config.host_ratio, config.host_size_gb, - config.page_size, config.host_layout, - host_size_tokens=host_size_tokens, + host_budget_bytes=host_budget_bytes, ) - elif isinstance(actual_pool, MHATokenToKVPool): - self.host_pool = MHATokenToKVPoolHost( - actual_pool, - config.host_ratio, - config.host_size_gb, - config.page_size, - config.host_layout, - host_size_tokens=host_size_tokens, - ) - elif isinstance(actual_pool, MLATokenToKVPool): - self.host_pool = MLATokenToKVPoolHost( - actual_pool, - config.host_ratio, - config.host_size_gb, - config.page_size, - config.host_layout, - host_size_tokens=host_size_tokens, + self.paged_cache_pool = DeepseekV4CachePool( + device_pool=actual_pool, + host_pool=self.host_pool, + io_backend=config.io_backend, ) else: - raise ValueError( - f"host_pool only supports DSA, MHA and MLA, got {type(actual_pool)} " - f"from module {type(actual_pool).__module__}" - ) + if draft_device_pool is not None: + actual_draft_pool = draft_device_pool + if hasattr(draft_device_pool, "inner") and not isinstance( + draft_device_pool, kv_pool_types + ): + actual_draft_pool = draft_device_pool.inner + if not isinstance( + actual_draft_pool, + (DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool), + ): + raise ValueError( + f"draft_device_pool only supports DSA, MHA and MLA, " + f"got {type(actual_draft_pool)}" + ) + + host_size_tokens = 0 + if config.host_size_gb == 0: + target_size_per_token = _pool_size_per_token(actual_pool) + draft_size_per_token = ( + _pool_size_per_token(actual_draft_pool) + if actual_draft_pool is not None + else 0 + ) + combined_size_per_token = target_size_per_token + draft_size_per_token + reserve_bytes = int(config.host_reserve_gb * (1024**3)) + available_bytes, _, cgroup_available = get_available_host_memory_bytes( + reserve_bytes + ) + requested_tokens = int(actual_pool.size * config.host_ratio) + host_size_tokens = _auto_capped_host_size_tokens( + requested_tokens=requested_tokens, + page_size=config.page_size, + size_per_token=combined_size_per_token, + available_host_memory_bytes=available_bytes, + host_parallel_count=config.host_parallel_count, + ) + if host_size_tokens > 0: + capped_tokens = _aligned_token_count( + host_size_tokens, config.page_size + ) + requested_tokens_aligned = _aligned_token_count( + requested_tokens, config.page_size + ) + logger.warning( + "Capping KVStore host pool for cgroup budget: " + "tokens %s -> %s, total bytes %.2f GB -> %.2f GB " + "(parallel_count=%s, available=%.2f GB, cgroup_available=%s)", + requested_tokens_aligned, + capped_tokens, + requested_tokens_aligned * combined_size_per_token / 1e9, + capped_tokens * combined_size_per_token / 1e9, + config.host_parallel_count, + available_bytes / 1e9, + ( + f"{cgroup_available / 1e9:.2f} GB" + if cgroup_available is not None + else "unlimited" + ), + ) + + # DSA subclasses MLA, so it must be matched before the MLA branch. + if isinstance(actual_pool, DSATokenToKVPool): + self.host_pool = DSATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + elif isinstance(actual_pool, MHATokenToKVPool): + self.host_pool = MHATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + elif isinstance(actual_pool, MLATokenToKVPool): + self.host_pool = MLATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + else: + raise ValueError( + "host_pool only supports DSA, MHA, MLA, and DeepSeek V4, " + f"got {type(actual_pool)} from module {type(actual_pool).__module__}" + ) # Draft model L2 cache: draft shares the same page mapping as the base # model, so its host pool must hold exactly the same number of tokens. # Pass host_size_tokens directly to bypass ratio/GB recalculation. - if actual_draft_pool is not None: + if actual_draft_pool is not None and self.paged_cache_pool is None: if isinstance(actual_draft_pool, DSATokenToKVPool): self.draft_host_pool = DSATokenToKVPoolHost( actual_draft_pool, @@ -276,7 +347,8 @@ def __init__( pools = None self.mamba_host_pool = None if ( - config.enable_mamba_l2 + self.paged_cache_pool is None + and config.enable_mamba_l2 and mamba_pool is not None and config.mamba_l2_host_slots > 0 ): @@ -313,7 +385,13 @@ def __init__( config.host_layout, ) - if pools is not None: + if self.paged_cache_pool is not None: + self.host_exec = HostExecutor( + pools=pools or [], + paged_pool=self.paged_cache_pool, + io_backend=config.io_backend, + ) + elif pools is not None: self.host_exec = HostExecutor(pools=pools, io_backend=config.io_backend) else: self.host_exec = HostExecutor( @@ -328,6 +406,7 @@ def __init__( draft_host_pool=self.draft_host_pool, draft_layer_num=draft_layer_num, ) + self.emits_loadback_acks = self.host_exec.emits_loadback_acks self.storage_exec = StorageExecutor( page_size=config.page_size, device_pool=device_pool, @@ -377,9 +456,14 @@ def submit(self, op) -> None: len(op.dst_pages), ) groups = self._page_groups_by_kind(op) + paged_transfers_by_op = getattr(op, "paged_cache_transfers", []) + is_retract_flags = getattr(op, "is_retract", []) for i in range(len(op.op_ids)): op_id = op.op_ids[i] - is_retract = bool(getattr(op, "is_retract", [False])[i]) + is_retract = ( + bool(is_retract_flags[i]) if i < len(is_retract_flags) else False + ) + submitted = False for kind, (src_groups, dst_groups) in groups.items(): if kind not in self.host_exec.pools: continue @@ -405,7 +489,32 @@ def submit(self, op) -> None: is_retract=is_retract, kind=kind, ) - if all( + submitted = True + paged_transfers = ( + paged_transfers_by_op[i] if i < len(paged_transfers_by_op) else [] + ) + if paged_transfers: + if logger.isEnabledFor(_DEBUG): + pages, spans, debug_groups = _paged_transfer_debug_summary( + paged_transfers + ) + logger.debug( + "[cache_op][paged_l2] writeback schedule op_id=%s " + "pages=%s spans=%s groups=%s transfers=%s is_retract=%s", + op_id, + pages, + spans, + debug_groups, + len(paged_transfers), + is_retract, + ) + self.host_exec.enqueue_paged_cache_writeback( + op_id, + paged_transfers, + is_retract=is_retract, + ) + submitted = True + if not submitted and all( i >= len(src_groups) or not src_groups[i] for kind, (src_groups, _) in groups.items() if kind in self.host_exec.pools @@ -419,6 +528,7 @@ def submit(self, op) -> None: len(op.dst_pages), ) groups = self._page_groups_by_kind(op) + paged_transfers_by_op = getattr(op, "paged_cache_transfers", []) for i in range(len(op.op_ids)): op_id = op.op_ids[i] for kind, (src_groups, dst_groups) in groups.items(): @@ -448,6 +558,27 @@ def submit(self, op) -> None: self.host_exec.enqueue_loadback( op_id, src_pages, dst_pages, kind=kind, **loadback_kwargs ) + paged_transfers = ( + paged_transfers_by_op[i] if i < len(paged_transfers_by_op) else [] + ) + if paged_transfers: + if logger.isEnabledFor(_DEBUG): + pages, spans, debug_groups = _paged_transfer_debug_summary( + paged_transfers + ) + logger.debug( + "[cache_op][paged_l2] loadback schedule op_id=%s " + "pages=%s spans=%s groups=%s transfers=%s", + op_id, + pages, + spans, + debug_groups, + len(paged_transfers), + ) + self.host_exec.enqueue_paged_cache_loadback( + op_id, + paged_transfers, + ) elif isinstance(op, Cache.PrefetchOp): logger.debug( diff --git a/python/tokenspeed/runtime/cache/transfer/__init__.py b/python/tokenspeed/runtime/cache/transfer/__init__.py index 7c97ee8ba..6a730bac5 100644 --- a/python/tokenspeed/runtime/cache/transfer/__init__.py +++ b/python/tokenspeed/runtime/cache/transfer/__init__.py @@ -22,8 +22,10 @@ from tokenspeed.runtime.cache.transfer.mamba_pool import MambaCachePool from tokenspeed.runtime.cache.transfer.pool import CachePool from tokenspeed.runtime.cache.transfer.types import ( + PAGED_CACHE_KIND, CacheKind, Location, + PagedCacheTransferUnit, TransferBatch, TransferUnit, ) @@ -34,6 +36,8 @@ "KVCachePool", "Location", "MambaCachePool", + "PAGED_CACHE_KIND", + "PagedCacheTransferUnit", "TransferBatch", "TransferUnit", ] diff --git a/python/tokenspeed/runtime/cache/transfer/deepseek_v4_pool.py b/python/tokenspeed/runtime/cache/transfer/deepseek_v4_pool.py new file mode 100644 index 000000000..c04f76ddc --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/deepseek_v4_pool.py @@ -0,0 +1,510 @@ +from __future__ import annotations + +import logging +from collections.abc import Iterable +from dataclasses import dataclass + +import torch +from tokenspeed_kernel.ops.kvcache.cuda import ( + DirectH2DScatterPlan, + prepare_kv_direct_h2d_scatter_plan, + transfer_kv_direct, + transfer_kv_direct_h2d_scatter_prepared, +) + +from tokenspeed.runtime.cache.deepseek_v4_cache_host import ( + DeepseekV4TokenToKVPoolHost, +) +from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter +from tokenspeed.runtime.cache.transfer.types import PAGED_CACHE_KIND +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + V4_SWA_KV_GROUP_ID, + parse_v4_compressor_state_group_id, +) +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) +_DEBUG = logging.DEBUG +_H2D_SCATTER_MIN_COPY_CALLS = 4096 + + +@dataclass(frozen=True, slots=True) +class PagedCacheTensorRef: + group_id: str + layer_id: int + device_tensor: torch.Tensor + host_tensor: torch.Tensor + page_bytes: int + + +@dataclass(slots=True) +class _PreparedPagedCacheTransfer: + group_id: str + src_indices: torch.Tensor + dst_indices: torch.Tensor + page_count: int + span_count: int + src_indices_device: torch.Tensor | None = None + dst_indices_device: torch.Tensor | None = None + + def h2d_indices(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + if self.src_indices_device is None or self.src_indices_device.device != device: + if device.type == "cuda": + if not self.src_indices.is_pinned(): + self.src_indices = self.src_indices.pin_memory() + if not self.dst_indices.is_pinned(): + self.dst_indices = self.dst_indices.pin_memory() + self.src_indices_device = self.src_indices.to(device, non_blocking=True) + self.dst_indices_device = self.dst_indices.to(device, non_blocking=True) + assert self.dst_indices_device is not None + return self.src_indices_device, self.dst_indices_device + + +@dataclass(frozen=True, slots=True) +class _PagedCopySummary: + backend: str + fallback_reason: str + pages: int + spans: int + tensor_refs: int + effective_copy_calls: int + total_bytes: int + buckets: int + kernel_launches: int + + +def _parse_v4_compressed_kv_group_id(group_id: str) -> int | None: + prefix = "v4.c" + suffix = "a.compressed_kv" + if not group_id.startswith(prefix) or not group_id.endswith(suffix): + return None + try: + return int(group_id[len(prefix) : -len(suffix)]) + except ValueError: + return None + + +def _ordered_page_pairs(src_pages: Iterable[int], dst_pages: Iterable[int]): + seen = set() + pairs = [] + for src_page, dst_page in zip(src_pages, dst_pages): + pair = (int(src_page), int(dst_page)) + if pair in seen: + continue + seen.add(pair) + pairs.append(pair) + pairs.sort() + return pairs + + +def _count_contiguous_spans(pairs: list[tuple[int, int]]) -> int: + if not pairs: + return 0 + spans = 1 + prev_src, prev_dst = pairs[0] + for src_page, dst_page in pairs[1:]: + if src_page != prev_src + 1 or dst_page != prev_dst + 1: + spans += 1 + prev_src, prev_dst = src_page, dst_page + return spans + + +def _coalesce_page_pairs_by_group(transfers: list) -> dict[str, list[tuple[int, int]]]: + pairs_by_group: dict[str, list[tuple[int, int]]] = {} + for transfer in transfers: + src_pages = list(getattr(transfer, "src_pages")) + dst_pages = list(getattr(transfer, "dst_pages")) + if len(src_pages) != len(dst_pages): + raise ValueError( + "DeepSeek V4 paged transfer page count mismatch for " + f"group={transfer.group_id!r}: {len(src_pages)} src vs " + f"{len(dst_pages)} dst" + ) + pairs = _ordered_page_pairs(src_pages, dst_pages) + if not pairs: + continue + group_id = str(transfer.group_id) + pairs_by_group.setdefault(group_id, []).extend(pairs) + + return { + group_id: sorted(set(pairs)) + for group_id, pairs in pairs_by_group.items() + if pairs + } + + +class DeepseekV4CachePool: + """Group-paged DeepSeek V4 L2 transfer pool. + + This pool deliberately does not implement the CacheKind protocol. It is + driven by typed paged-cache transfers whose group ids carry the page space. + """ + + kind = PAGED_CACHE_KIND + loadback_layer_chunk_size = 4 + + def __init__( + self, + device_pool: DeepseekV4TokenToKVPool, + host_pool: DeepseekV4TokenToKVPoolHost, + io_backend: str, + ) -> None: + if io_backend not in ("kernel", "direct"): + raise ValueError( + f"Unsupported DeepSeek V4 paged-cache io_backend={io_backend}" + ) + self.device_pool = device_pool + self.host_pool = host_pool + self.io_backend = io_backend + self._counter = LayerDoneCounter(self.num_layers()) + device_pool.register_layer_transfer_counter(self._counter) + self._h2d_scatter_plans: dict[ + str, + tuple[int, tuple[tuple[int, int, int, int], ...], DirectH2DScatterPlan], + ] = {} + + @property + def device(self): + return self.device_pool.device + + @property + def host_layout(self) -> str: + return self.host_pool.layout + + def num_layers(self) -> int: + return int(self.device_pool.layer_num) + + def supports_layerwise_loadback(self) -> bool: + return True + + def get_layer_done_counter(self) -> LayerDoneCounter: + return self._counter + + def local_layer_idx(self, global_layer_id: int) -> int: + return global_layer_id + + def tensor_refs_for_group( + self, + group_id: str, + layer_idx: int | None = None, + ) -> list[PagedCacheTensorRef]: + group_id = str(group_id) + layer_ids = ( + range(self.num_layers()) + if layer_idx is None + else range(layer_idx, layer_idx + 1) + ) + refs: list[PagedCacheTensorRef] = [] + for layer_id in layer_ids: + ratio = int(self.device_pool.layout.layer_ratio[layer_id]) + if group_id == V4_SWA_KV_GROUP_ID: + refs.append( + self._ref( + group_id, + layer_id, + self.device_pool.swa_kv_buffer[layer_id], + self.host_pool.swa_kv_buffer[layer_id], + ) + ) + continue + + state_ratio = parse_v4_compressor_state_group_id(group_id) + if state_ratio is not None: + if ratio == state_ratio: + refs.append( + self._ref( + group_id, + layer_id, + self.device_pool.compressor_state_buffer[layer_id], + self.host_pool.compressor_state_buffer[layer_id], + ) + ) + continue + + compressed_ratio = _parse_v4_compressed_kv_group_id(group_id) + if compressed_ratio is not None: + if ratio == compressed_ratio: + refs.append( + self._ref( + group_id, + layer_id, + self.device_pool.compressed_kv_buffer[layer_id], + self.host_pool.compressed_kv_buffer[layer_id], + ) + ) + if ratio == 4: + refs.append( + self._ref( + group_id, + layer_id, + self.device_pool.indexer_kv_buffer[layer_id], + self.host_pool.indexer_kv_buffer[layer_id], + ) + ) + continue + + if group_id == V4_INDEXER_COMPRESSOR_STATE_GROUP_ID: + if ratio == 4: + refs.append( + self._ref( + group_id, + layer_id, + self.device_pool.indexer_state_buffer[layer_id], + self.host_pool.indexer_state_buffer[layer_id], + ) + ) + continue + + raise KeyError(f"unknown DeepSeek V4 paged-cache group {group_id!r}") + return refs + + @staticmethod + def _ref( + group_id: str, + layer_id: int, + device_tensor: torch.Tensor | None, + host_tensor: torch.Tensor | None, + ) -> PagedCacheTensorRef: + if device_tensor is None or host_tensor is None: + raise ValueError( + f"DeepSeek V4 group {group_id!r} has no tensor for layer {layer_id}" + ) + return PagedCacheTensorRef( + group_id=group_id, + layer_id=layer_id, + device_tensor=device_tensor, + host_tensor=host_tensor, + page_bytes=int(device_tensor[0].nbytes), + ) + + def writeback_paged(self, transfers: list) -> None: + prepared = self.prepare_paged_transfers(transfers) + self.writeback_prepared_paged(prepared) + + def loadback_paged(self, transfers: list, layer_idx: int) -> None: + prepared = self.prepare_paged_transfers(transfers) + self.loadback_prepared_paged(prepared, layer_idx) + + def prepare_paged_transfers( + self, transfers: list + ) -> list[_PreparedPagedCacheTransfer]: + """Coalesce scheduler transfer fragments into group-level page copies.""" + prepared: list[_PreparedPagedCacheTransfer] = [] + for group_id, pairs in _coalesce_page_pairs_by_group(transfers).items(): + prepared.append( + _PreparedPagedCacheTransfer( + group_id=group_id, + src_indices=torch.tensor( + [src for src, _ in pairs], + dtype=torch.int64, + device="cpu", + ), + dst_indices=torch.tensor( + [dst for _, dst in pairs], + dtype=torch.int64, + device="cpu", + ), + page_count=len(pairs), + span_count=_count_contiguous_spans(pairs), + ) + ) + return prepared + + def writeback_prepared_paged( + self, transfers: list[_PreparedPagedCacheTransfer] + ) -> None: + """Submit D2H copies using coalesced page index tensors.""" + for transfer in transfers: + self._copy_prepared_transfer( + transfer, + self.tensor_refs_for_group(transfer.group_id), + host_to_device=False, + ) + + def loadback_prepared_paged( + self, transfers: list[_PreparedPagedCacheTransfer], layer_idx: int + ) -> None: + """Submit one layer's H2D copy using precomputed page index tensors.""" + self.loadback_prepared_paged_range(transfers, layer_idx, layer_idx + 1) + + def loadback_prepared_paged_range( + self, + transfers: list[_PreparedPagedCacheTransfer], + layer_start: int, + layer_end: int, + ) -> None: + """Submit a small layer range's H2D copies with shared page indices.""" + if layer_start < 0 or layer_end < layer_start or layer_end > self.num_layers(): + raise ValueError( + "invalid DeepSeek V4 paged loadback layer range: " + f"[{layer_start}, {layer_end})" + ) + summaries: list[_PagedCopySummary] = [] + for transfer in transfers: + refs: list[PagedCacheTensorRef] = [] + for layer_idx in range(layer_start, layer_end): + refs.extend( + self.tensor_refs_for_group( + transfer.group_id, + layer_idx=layer_idx, + ) + ) + summary = self._copy_prepared_transfer( + transfer, + refs, + host_to_device=True, + layer_start=layer_start, + layer_end=layer_end, + ) + if summary is not None: + summaries.append(summary) + self._log_h2d_chunk_summary(layer_start, layer_end, summaries) + + def _copy_prepared_transfer( + self, + transfer: _PreparedPagedCacheTransfer, + refs: list[PagedCacheTensorRef], + *, + host_to_device: bool, + layer_start: int | None = None, + layer_end: int | None = None, + ) -> _PagedCopySummary | None: + if transfer.page_count == 0 or not refs: + return None + transfer_bytes = transfer.page_count * sum(ref.page_bytes for ref in refs) + effective_copy_calls = int(transfer.span_count) * len(refs) + buckets = len({ref.page_bytes for ref in refs}) + if host_to_device: + src_layers = [ref.host_tensor for ref in refs] + dst_layers = [ref.device_tensor for ref in refs] + else: + src_layers = [ref.device_tensor for ref in refs] + dst_layers = [ref.host_tensor for ref in refs] + + backend = "direct" + fallback_reason = "" + kernel_launches = 0 + if host_to_device: + assert layer_start is not None and layer_end is not None + full_refs = self.tensor_refs_for_group(transfer.group_id) + threshold_copy_calls = max( + effective_copy_calls, + int(transfer.span_count) * len(full_refs), + ) + scatter_result = None + if threshold_copy_calls >= _H2D_SCATTER_MIN_COPY_CALLS: + plan, fallback_reason = self._h2d_scatter_plan( + transfer.group_id, + full_refs, + ) + if plan is not None: + src_indices_device, dst_indices_device = transfer.h2d_indices( + torch.device(self.device) + ) + scatter_result = transfer_kv_direct_h2d_scatter_prepared( + plan, + src_indices_device, + dst_indices_device, + 1, + layer_start, + layer_end, + ) + else: + fallback_reason = "below_threshold" + if scatter_result is not None and scatter_result.used: + backend = "scatter" + buckets = int(scatter_result.buckets) + kernel_launches = int(scatter_result.kernel_launches) + elif scatter_result is not None: + fallback_reason = fallback_reason or str(scatter_result.fallback_reason) + + if backend == "direct": + transfer_kv_direct( + src_layers=src_layers, + dst_layers=dst_layers, + src_indices=transfer.src_indices, + dst_indices=transfer.dst_indices, + page_size=1, + ) + return _PagedCopySummary( + backend=backend, + fallback_reason=fallback_reason, + pages=int(transfer.page_count), + spans=int(transfer.span_count), + tensor_refs=len(refs), + effective_copy_calls=effective_copy_calls, + total_bytes=transfer_bytes, + buckets=buckets, + kernel_launches=kernel_launches, + ) + + @staticmethod + def _log_h2d_chunk_summary( + layer_start: int, + layer_end: int, + summaries: list[_PagedCopySummary], + ) -> None: + if not summaries or not logger.isEnabledFor(_DEBUG): + return + backends = {summary.backend for summary in summaries} + backend = next(iter(backends)) if len(backends) == 1 else "mixed" + fallback_reasons = sorted( + { + summary.fallback_reason + for summary in summaries + if summary.fallback_reason + } + ) + logger.debug( + "[cache_op][paged_l2] h2d_chunk layers=[%s,%s) backend=%s " + "pages=%s spans=%s tensor_refs=%s effective_copy_calls=%s " + "total_bytes=%s buckets=%s kernel_launches=%s fallback_reason=%s", + layer_start, + layer_end, + backend, + sum(summary.pages for summary in summaries), + sum(summary.spans for summary in summaries), + sum(summary.tensor_refs for summary in summaries), + sum(summary.effective_copy_calls for summary in summaries), + sum(summary.total_bytes for summary in summaries), + sum(summary.buckets for summary in summaries), + sum(summary.kernel_launches for summary in summaries), + ",".join(fallback_reasons) if fallback_reasons else "none", + ) + + def _h2d_scatter_plan( + self, + group_id: str, + refs: list[PagedCacheTensorRef], + ) -> tuple[DirectH2DScatterPlan | None, str]: + device = torch.device(self.device) + stream_id = ( + int(torch.cuda.current_stream(device).cuda_stream) + if device.type == "cuda" + else 0 + ) + signature = tuple( + ( + ref.layer_id, + ref.page_bytes, + ref.host_tensor.data_ptr(), + ref.device_tensor.data_ptr(), + ) + for ref in refs + ) + cached = self._h2d_scatter_plans.get(group_id) + if cached is not None and cached[0] == stream_id and cached[1] == signature: + return cached[2], "" + + plan, reason = prepare_kv_direct_h2d_scatter_plan( + [ref.host_tensor for ref in refs], + [ref.device_tensor for ref in refs], + [ref.layer_id for ref in refs], + ) + if plan is not None: + self._h2d_scatter_plans[group_id] = (stream_id, signature, plan) + return plan, reason diff --git a/python/tokenspeed/runtime/cache/transfer/types.py b/python/tokenspeed/runtime/cache/transfer/types.py index 251090f6f..410620978 100644 --- a/python/tokenspeed/runtime/cache/transfer/types.py +++ b/python/tokenspeed/runtime/cache/transfer/types.py @@ -22,9 +22,12 @@ from dataclasses import dataclass from enum import Enum +from typing import Any import torch +PAGED_CACHE_KIND = "paged_cache" + class CacheKind(str, Enum): KV = "kv" @@ -58,3 +61,10 @@ def direction(self) -> tuple[Location, Location]: class TransferBatch: units: list[TransferUnit] op_ids: list[int] + + +@dataclass(slots=True) +class PagedCacheTransferUnit: + op_id: int + transfers: list[Any] + is_retract: bool = False diff --git a/python/tokenspeed/runtime/engine/event_loop.py b/python/tokenspeed/runtime/engine/event_loop.py index cbe4c5e55..78615871f 100644 --- a/python/tokenspeed/runtime/engine/event_loop.py +++ b/python/tokenspeed/runtime/engine/event_loop.py @@ -38,7 +38,7 @@ MemoryExecutor, MemoryExecutorConfig, ) -from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.cache.transfer.types import PAGED_CACHE_KIND, CacheKind from tokenspeed.runtime.configs.model_config import ModelConfig from tokenspeed.runtime.configs.paged_cache_spec import ( scheduler_ext_flat_kvcache, @@ -106,6 +106,24 @@ logger = get_colorful_logger(__name__) +def _paged_cache_host_group_pages_for_scheduler( + enable_kvstore: bool, host_pool +) -> dict: + if not enable_kvstore: + return {} + return dict(getattr(host_pool, "paged_cache_group_page_counts", {})) + + +def _validate_grouped_kvstore_draft_pool( + enable_kvstore: bool, paged_cache_groups: list, draft_token_to_kv_pool +) -> None: + if enable_kvstore and paged_cache_groups and draft_token_to_kv_pool is not None: + raise NotImplementedError( + "KVStore does not support speculative draft KV for grouped paged " + "caches; pass --disable-kvstore." + ) + + def calc_l3_query_hashes(scheduler, tokens: list[int]) -> list[str]: return scheduler.calc_rolling_hash(tokens, apply_match=True) @@ -211,6 +229,12 @@ def __init__( decode_input_tokens=decode_input_tokens, overlap_schedule_depth=self.overlap_schedule_depth, ) + paged_cache_groups = pool_to_paged_cache_groups(token_to_kv_pool) + _validate_grouped_kvstore_draft_pool( + server_args.enable_kvstore, + paged_cache_groups, + draft_token_to_kv_pool, + ) num_total_pages = self.max_total_num_tokens // server_args.block_size hf_config = getattr(self.model_config, "hf_config", None) @@ -273,11 +297,6 @@ def __init__( ) self._dp_local_info = torch.zeros(1, 3, dtype=torch.int32) self._dp_global_info = torch.zeros(mapping.world_size, 3, dtype=torch.int32) - if not server_args.enable_kvstore: - logger.warning( - "KVStore L2 cache will not be used during normal execution, but it will still be used when retraction happens." - ) - mamba_l2_host_slots = 0 if has_mamba and server_args.enable_mamba_l2: if server_args.mamba_l2_host_slots > 0: @@ -298,6 +317,9 @@ def __init__( int(mamba_pool_total_chunks * server_args.mamba_l2_ratio), 1 ) + enable_mamba_l2 = bool( + has_mamba and server_args.enable_mamba_l2 and mamba_l2_host_slots > 0 + ) mem_cfg = MemoryExecutorConfig( layer_num=self.model_config.num_hidden_layers, page_size=server_args.block_size, @@ -316,6 +338,7 @@ def __init__( mamba_l2_layout=server_args.mamba_l2_layout, mamba_l2_io_backend=server_args.mamba_l2_io_backend, ) + paged_cache_host_group_pages = {} if scheduler_ext_flat_kvcache() and server_args.enable_kvstore: if server_args.kvstore_storage_backend is not None: raise NotImplementedError( @@ -345,10 +368,13 @@ def __init__( draft_device_pool=draft_token_to_kv_pool, mamba_pool=mamba_pool, ) - num_host_pages = self.memory_executor.host_pool.page_num + num_host_pages = int(getattr(self.memory_executor.host_pool, "page_num", 0)) + paged_cache_host_group_pages = _paged_cache_host_group_pages_for_scheduler( + server_args.enable_kvstore, self.memory_executor.host_pool + ) - # Flat host tier acks loadbacks (LoadBackDoneEvent), so they join the - # inflight accounting in _submit_cache_ops; radix loadbacks never ack. + # Executors that pin loadback resources acknowledge completion, so their + # operations join the inflight accounting in _submit_cache_ops. self._loadback_acks_expected = getattr( self.memory_executor, "emits_loadback_acks", False ) @@ -365,8 +391,8 @@ def __init__( f"(ratio={server_args.mamba_full_memory_ratio})." ) - # Adjunct enabled only when pool opts in AND prefix-caching switch is on. - paged_cache_groups = pool_to_paged_cache_groups(token_to_kv_pool) + # Adjunct is device-side prefix-cache metadata and must stay enabled + # under --disable-kvstore to match the non-L2 DeepSeek V4 path. validate_flat_scheduler_config( flat_kvcache_ext=scheduler_ext_flat_kvcache(), paged_cache_groups=paged_cache_groups, @@ -396,21 +422,24 @@ def __init__( enable_mamba=has_mamba, mamba_cache_chunk_size=server_args.mamba_cache_chunk_size, mamba_pool_total_chunks=mamba_pool_total_chunks, - enable_mamba_l2=server_args.enable_mamba_l2, + enable_mamba_l2=enable_mamba_l2, mamba_l2_host_slots=mamba_l2_host_slots, paged_cache_groups=paged_cache_groups, + paged_cache_host_group_pages=paged_cache_host_group_pages, enable_mixed_prefill_decode=server_args.enable_mixed_batch, prefix_cache_adjunct=prefix_cache_adjunct, ) logger.info( - "Scheduler config: block_size=%s num_device_pages=%s " + "Scheduler config: block_size=%s num_device_pages=%s num_host_pages=%s " "max_scheduled_tokens=%s decode_input_tokens=%s " "overlap_schedule_depth=%s disable_l2_cache=%s " "max_batch_size=%s (global max_num_seqs=%s, dp_size=%s) " "mamba_pool_total_chunks=%s enable_mamba=%s " - "disable_prefix_cache=%s paged_cache_groups=%s", + "disable_prefix_cache=%s paged_cache_groups=%s " + "paged_cache_host_group_pages=%s prefix_cache_adjunct_groups=%s", scheduler_cfg.block_size, scheduler_cfg.num_device_pages, + scheduler_cfg.num_host_pages, scheduler_cfg.max_scheduled_tokens, scheduler_cfg.decode_input_tokens, scheduler_cfg.overlap_schedule_depth, @@ -422,6 +451,12 @@ def __init__( has_mamba, scheduler_cfg.disable_prefix_cache, [group.group_id for group in paged_cache_groups], + paged_cache_host_group_pages, + ( + list(prefix_cache_adjunct.required_groups) + if prefix_cache_adjunct is not None + else None + ), ) self.scheduler = Scheduler(scheduler_cfg) token_to_kv_pool.bind_paged_cache_scheduler(self.scheduler) @@ -750,6 +785,10 @@ def _publish_scheduler_kv_events(self) -> None: KVEventBatch(ts=time.time(), events=events, attn_dp_rank=self.dp_rank) ) + def _publish_scheduler_aborts(self, execution_plan) -> None: + for request_id, message in getattr(execution_plan, "scheduler_aborts", ()): + self.output_processor.publish_scheduler_abort(request_id, message) + def _cache_group_has_work(self, local_has_work: bool) -> bool: """Whether ANY attn-tp rank has cache work this step (unanimous via a single-int MAX all_reduce, far cheaper than the payload gather it @@ -993,8 +1032,7 @@ def _submit_cache_ops(self, execution_plan) -> None: if isinstance(op, Cache.WriteBackOp): self._num_inflight_cache_ops += len(op.op_ids) elif isinstance(op, Cache.LoadBackOp): - # Radix loadbacks are fire-and-forget (no ack, nothing in - # flight); the flat host tier acks one LoadBackDone per op_id. + # Resource-pinning host tiers ack one LoadBackDone per op_id. if self._loadback_acks_expected: self._num_inflight_cache_ops += len(op.op_ids) elif isinstance(op, (Cache.PrefetchOp, Cache.BackUpOp)): @@ -1011,9 +1049,12 @@ def _setup_layerwise_loadback(self, execution_plan) -> None: consumer_indices_by_kind: dict[CacheKind, list[int]] = { kind: [] for kind in available_pools } + consumer_indices_by_paged_cache: list[int] = [] + has_paged_pool = getattr(host_exec, "paged_pool", None) is not None for cache_op in execution_plan.cache: if isinstance(cache_op, Cache.LoadBackOp): - for op_id in cache_op.op_ids: + paged_transfers_by_op = getattr(cache_op, "paged_cache_transfers", []) + for i, op_id in enumerate(cache_op.op_ids): for kind in consumer_indices_by_kind: producer_idx = self.memory_executor.get_producer_index( kind, op_id @@ -1023,10 +1064,34 @@ def _setup_layerwise_loadback(self, execution_plan) -> None: and producer_idx not in consumer_indices_by_kind[kind] ): consumer_indices_by_kind[kind].append(producer_idx) + paged_transfers = ( + paged_transfers_by_op[i] + if i < len(paged_transfers_by_op) + else [] + ) + if has_paged_pool and paged_transfers: + producer_idx = self.memory_executor.get_producer_index( + PAGED_CACHE_KIND, + op_id, + ) + if ( + producer_idx is not None + and producer_idx not in consumer_indices_by_paged_cache + ): + consumer_indices_by_paged_cache.append(producer_idx) for kind, consumer_indices in consumer_indices_by_kind.items(): self.memory_executor.set_consumer( kind, consumer_indices if consumer_indices else -1 ) + if has_paged_pool: + self.memory_executor.set_consumer( + PAGED_CACHE_KIND, + ( + consumer_indices_by_paged_cache + if consumer_indices_by_paged_cache + else -1 + ), + ) def _flush_mamba_retract_states(self, forward_op) -> None: """Copy draft->working mamba states when retract occurred (no forward scheduled).""" @@ -1231,7 +1296,11 @@ def _process_new_requests(self): if self.kv_transfer is not None and not is_epd: self.kv_transfer.register(spec.request_id, bootstrap) - if self.memory_executor is not None: + if ( + self.memory_executor is not None + and self.server_args.enable_kvstore + and self.server_args.kvstore_storage_backend is not None + ): hashes = calc_l3_query_hashes(self.scheduler, spec.tokens) if hashes and len(hashes) > self.prefetch_threshold: hit_pages = self.memory_executor.query_l3_pages(hashes) @@ -1584,6 +1653,7 @@ def event_loop(self): self._paused_idle_step() continue execution_plan = self.scheduler.next_execution_plan() + self._publish_scheduler_aborts(execution_plan) self._publish_scheduler_kv_events() self._handle_flat_oom_terminals(execution_plan) self._submit_cache_ops(execution_plan) @@ -1736,6 +1806,7 @@ def event_loop_overlap(self): prev_forward_op = None continue execution_plan = self.scheduler.next_execution_plan() + self._publish_scheduler_aborts(execution_plan) self._publish_scheduler_kv_events() self._handle_flat_oom_terminals(execution_plan) diff --git a/python/tokenspeed/runtime/engine/generation_output_processor.py b/python/tokenspeed/runtime/engine/generation_output_processor.py index 69162c09f..475a13963 100644 --- a/python/tokenspeed/runtime/engine/generation_output_processor.py +++ b/python/tokenspeed/runtime/engine/generation_output_processor.py @@ -422,11 +422,10 @@ def register(self, rid, state): state.set_finish_with_abort("AbortReq from client") def publish_finished_at_admission(self, rid: str, state: RequestState) -> None: - """Stream a finish for a request that was finished before admission. + """Stream and release a request finished outside normal forward output. - Used for grammar-aborted requests (invalid/timed-out compile, missing - backend) so the client gets a finish_reason without us wasting a - scheduler slot or a forward step on them. + Used for grammar admission failures and scheduler-internal aborts so + the client receives a finish reason without another forward step. """ self.rid_to_state[rid] = state try: @@ -435,9 +434,8 @@ def publish_finished_at_admission(self, rid: str, state: RequestState) -> None: finally: state.release_pending_multimodal_features() self.rid_to_state.pop(rid, None) - # This path replaces register() for grammar-aborted rids — - # drop any queued abort marker so pending_aborts doesn't leak - # and a reused rid isn't instantly re-aborted on next register. + # Drop any queued abort marker so pending_aborts does not leak and + # a reused rid is not immediately re-aborted on the next register. self.pending_aborts.pop(rid, None) def reap_finished_orphan(self, rid: str, state: RequestState) -> None: @@ -453,6 +451,15 @@ def reap_finished_orphan(self, rid: str, state: RequestState) -> None: else: self.rid_to_state.pop(rid, None) + def publish_scheduler_abort(self, rid: str, message: str) -> None: + """Publish a terminal error decided internally by the scheduler.""" + state = self.rid_to_state.get(rid) + if state is None: + return + state.set_finish_with_abort(message, notify_client=True) + self._log_request_stats(rid, state, time.time()) + self.publish_finished_at_admission(rid, state) + def _host_advance_matcher(self, completion, model_execution_results): """Host-side fallback for the grammar matcher advance. diff --git a/python/tokenspeed/runtime/engine/scheduler_utils.py b/python/tokenspeed/runtime/engine/scheduler_utils.py index 11889ebc9..83c66c075 100644 --- a/python/tokenspeed/runtime/engine/scheduler_utils.py +++ b/python/tokenspeed/runtime/engine/scheduler_utils.py @@ -87,6 +87,7 @@ def make_config( enable_mamba_l2: bool = False, mamba_l2_host_slots: int = 0, paged_cache_groups: Sequence["PagedCacheGroupConfig"] | None = None, + paged_cache_host_group_pages: Mapping[str, int] | None = None, enable_mixed_prefill_decode: bool = False, prefix_cache_adjunct: "PrefixCacheAdjunctSpec | None" = None, ) -> SchedulerConfig: @@ -120,6 +121,11 @@ def make_config( cfg.enable_mixed_prefill_decode = enable_mixed_prefill_decode if paged_cache_groups: cfg.paged_cache_groups = list(paged_cache_groups) + if paged_cache_host_group_pages: + cfg.paged_cache_host_group_pages = { + str(group_id): int(page_count) + for group_id, page_count in paged_cache_host_group_pages.items() + } # Opt-in; unset means paged-cache groups are transport-only. if prefix_cache_adjunct is not None: cfg.prefix_cache_adjunct = prefix_cache_adjunct diff --git a/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py b/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py index f86ed47d6..1a6c58993 100644 --- a/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py +++ b/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py @@ -43,6 +43,9 @@ from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( deepseek_v4_fused_csa_indexer_mxfp4_cache_insert as _triton_fused_csa_indexer_mxfp4_cache_insert, ) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_fused_hca_direct_compress_cache_insert as _triton_fused_hca_direct_compress_cache_insert, +) from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( deepseek_v4_fused_indexer_q_rope_hadamard_mxfp4 as _triton_fused_indexer_q_rope_hadamard_mxfp4, ) @@ -87,6 +90,7 @@ "deepseek_v4_csa_compress_kv_cache_insert", "deepseek_v4_csa_indexer_cache_insert", "deepseek_v4_hca_compress_kv_cache_insert", + "deepseek_v4_hca_direct_compress_kv_cache_insert", "deepseek_v4_prepare_indexer_q_mxfp4", "dequantize_deepseek_v4_fp8_ds_mla_cache", "fused_qnorm_rope_kv_insert", @@ -910,6 +914,118 @@ def deepseek_v4_hca_compress_kv_cache_insert( ) +def deepseek_v4_hca_direct_compress_kv_cache_insert( + state_cache: torch.Tensor, + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + block_table: torch.Tensor, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + active_token_indices: torch.Tensor, + block_table_base_offsets: torch.Tensor | None = None, +) -> None: + """Compress HCA active rows without falling back to paged row gathering. + + The kernel reads current-step tokens directly from ``kv``/``score`` and + prefix-boundary tokens from the loaded paged ``state_cache``. This keeps HCA + cached-prefill rows on the direct path even when their 128-token window + crosses the prefix/tail boundary. + + Args: + state_cache: Paged HCA compressor state with shape + ``[blocks, block_size, 2 * head_dim]``. + kv: Current forward HCA value states with shape ``[tokens, head_dim]``. + score: Current forward HCA score states with shape ``[tokens, head_dim]``. + ape: HCA APE table with shape ``[128, head_dim]``. + token_to_req_indices: Request index for each current-step token. + query_start_loc: Prefix-sum token offsets for current-step requests. + positions: Absolute token positions for the current forward step. + block_table: Paged compressor-state block table. + compressor_block_size: Compressor-state page size. + rms_norm_weight: RMSNorm weight applied after weighted reduction. + rms_norm_eps: RMSNorm epsilon. + cos_sin_cache: RoPE cache indexed by compressed positions. + kv_cache_2d: Destination compressed KV cache byte storage. + kv_slot_mapping: Destination compressed KV slot per source token. + kv_cache_block_size: Destination compressed KV cache page size. + active_token_indices: Source token offsets for HCA compressed rows. + block_table_base_offsets: Optional logical-page offset per request. + """ + + if kv.shape != score.shape: + raise ValueError( + f"kv and score shapes must match, got {kv.shape} vs {score.shape}" + ) + if kv.dim() != 2: + raise ValueError(f"kv/score must be [tokens, head_dim], got {kv.shape}") + if state_cache.dim() != 3: + raise ValueError(f"state_cache must be 3D, got {tuple(state_cache.shape)}") + state_width = kv.shape[-1] + head_dim = int(rms_norm_weight.numel()) + if state_width != head_dim: + raise ValueError(f"HCA state width must be {head_dim}, got {state_width}") + if compressor_block_size != state_cache.shape[1]: + raise ValueError( + "compressor_block_size must match state_cache page size, " + f"got {compressor_block_size} vs {state_cache.shape[1]}" + ) + if state_cache.shape[-1] != state_width * 2: + raise ValueError( + f"state_cache last dim must be {state_width * 2}, " + f"got {state_cache.shape[-1]}" + ) + rope_dim = int(cos_sin_cache.shape[-1]) + min_block_stride = kv_cache_block_size * deepseek_v4_swa_row_bytes( + state_width, rope_dim + ) + if kv_cache_2d.dim() != 2 or kv_cache_2d.shape[1] < min_block_stride: + raise ValueError( + f"kv_cache_2d must be [blocks, >= {min_block_stride}] uint8, " + f"got {tuple(kv_cache_2d.shape)}" + ) + if kv_cache_2d.dtype != torch.uint8: + raise TypeError(f"kv_cache_2d must be uint8, got {kv_cache_2d.dtype}") + if not kv.is_cuda: + raise ValueError( + "deepseek_v4_hca_direct_compress_kv_cache_insert only supports CUDA tensors." + ) + if token_to_req_indices.device != positions.device: + token_to_req_indices = token_to_req_indices.to( + positions.device, non_blocking=True + ) + if query_start_loc.device != positions.device: + query_start_loc = query_start_loc.to(positions.device, non_blocking=True) + + _triton_fused_hca_direct_compress_cache_insert( + state_cache=state_cache, + kv=kv, + score=score, + ape=ape, + token_to_req_indices=token_to_req_indices, + query_start_loc=query_start_loc, + positions=positions, + block_table=block_table, + compressor_block_size=compressor_block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=kv_cache_2d, + kv_slot_mapping=kv_slot_mapping, + kv_cache_block_size=kv_cache_block_size, + active_token_indices=active_token_indices, + block_table_base_offsets=block_table_base_offsets, + ) + + def deepseek_v4_csa_compress_kv_cache_insert( state_cache: torch.Tensor, token_to_req_indices: torch.Tensor, diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py b/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py index dd4785814..9009dd19e 100644 --- a/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py @@ -13,7 +13,6 @@ from __future__ import annotations -import logging from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from fractions import Fraction @@ -763,7 +762,7 @@ class DeepseekV4TokenToKVPool(BaseTokenToKVPool): group rather than owning a separate group of its own. """ - supports_hierarchical_kv_cache = False + supports_hierarchical_kv_cache = True def __init__( self, @@ -815,12 +814,6 @@ def __init__( self._paged_cache_group_specs_by_id = { spec.group_id: spec for spec in self.paged_cache_group_specs } - self._paged_cache_scheduler: object | None = None - self._paged_cache_state_group_ids = tuple( - str(spec.group_id) - for spec in self.paged_cache_group_specs - if spec.family == "state" - ) self.paged_cache_group_page_counts = compute_paged_cache_group_page_counts( self.paged_cache_group_specs, max_live_requests=max_batch_size, @@ -986,36 +979,22 @@ def prefix_cache_required_group_ids(self) -> tuple[str, ...]: if spec.family == "history" ) - def bind_paged_cache_scheduler(self, scheduler: object) -> None: - self._paged_cache_scheduler = scheduler - - def maybe_log_paged_cache_group_pages(self) -> None: - scheduler = self._paged_cache_scheduler - if self.rank != 0 or scheduler is None or not self._paged_cache_state_group_ids: - return - if not logger.isEnabledFor(logging.DEBUG): - return - - parts = [] - for group_id in self._paged_cache_state_group_ids: - total = scheduler.paged_cache_group_total_pages(group_id) - available = scheduler.paged_cache_group_available_pages(group_id) - failed = scheduler.paged_cache_group_failed_alloc_count(group_id) - parts.append( - f"{group_id}: used={total - available}/{total}, " - f"available={available}, failed_alloc={failed}" - ) - logger.debug("DeepSeek V4 paged-cache state group pages. %s", "; ".join(parts)) - def _require( self, buffers: list[torch.Tensor | None], layer_id: int, name: str ) -> torch.Tensor: + self._wait_for_layer_loadback(layer_id) buf = buffers[layer_id] if buf is None: raise ValueError(f"DeepSeek V4 layer {layer_id} has no {name} cache") return buf + def _wait_for_layer_loadback(self, layer_id: int) -> None: + counter = self.layer_transfer_counter + if counter is not None: + counter.wait_until(layer_id) + def get_swa_kv_buffer(self, layer_id: int) -> torch.Tensor: + self._wait_for_layer_loadback(layer_id) return self.swa_kv_buffer[layer_id] @property @@ -1269,13 +1248,13 @@ def get_layerwise_buf_info_offsets(self, start_idx=0): def get_cpu_copy(self, token_indices: list[int]) -> list[torch.Tensor]: del token_indices raise NotImplementedError( - "DeepSeek V4 KV cache offload is not implemented; the compressed-MQA " - "and indexer buffers are page-shaped and require page-aware indexing." + "DeepSeek V4 does not support legacy token-indexed KV cache offload; " + "use the group-paged L2 KVStore path instead." ) def load_cpu_copy(self, kv_cache_cpu, token_indices: list[int]) -> None: del kv_cache_cpu, token_indices raise NotImplementedError( - "DeepSeek V4 KV cache reload is not implemented; the compressed-MQA " - "and indexer buffers are page-shaped and require page-aware indexing." + "DeepSeek V4 does not support legacy token-indexed KV cache reload; " + "use the group-paged L2 KVStore path instead." ) diff --git a/python/tokenspeed/runtime/models/deepseek_v4.py b/python/tokenspeed/runtime/models/deepseek_v4.py index 18b09b5dd..a409f44a8 100644 --- a/python/tokenspeed/runtime/models/deepseek_v4.py +++ b/python/tokenspeed/runtime/models/deepseek_v4.py @@ -102,6 +102,7 @@ deepseek_v4_csa_indexer_cache_insert, deepseek_v4_fused_inv_rope_fp8_quant, deepseek_v4_hca_compress_kv_cache_insert, + deepseek_v4_hca_direct_compress_kv_cache_insert, deepseek_v4_prepare_indexer_q_mxfp4, fused_qnorm_rope_kv_insert, save_deepseek_v4_compressor_state, @@ -197,6 +198,77 @@ def _deepseek_v4_forward_metadata(ctx: ForwardContext): return metadata +def _deepseek_v4_hca_active_token_offsets( + ctx: ForwardContext, + metadata: DeepseekV4ForwardMetadata, + positions: torch.Tensor, + compress_ratio: int, +) -> list[int] | None: + if compress_ratio != 128 or get_is_capture_mode(): + return None + forward_mode = getattr(ctx, "forward_mode", None) + if forward_mode is None or not forward_mode.is_extend_or_mixed(): + return None + num_tokens = int(positions.numel()) + if num_tokens <= 0: + return None + seq_lens_cpu = getattr(metadata, "seq_lens_cpu", None) + query_lens_cpu = getattr(metadata, "query_lens_cpu", None) + if seq_lens_cpu is None or query_lens_cpu is None: + return None + if seq_lens_cpu.device.type != "cpu" or query_lens_cpu.device.type != "cpu": + return None + num_reqs = min(int(seq_lens_cpu.numel()), int(query_lens_cpu.numel())) + if num_reqs <= 0: + return None + + active_offsets: list[int] = [] + token_start = 0 + seq_lens = seq_lens_cpu[:num_reqs].tolist() + query_lens = query_lens_cpu[:num_reqs].tolist() + for seq_len_value, query_len_value in zip(seq_lens, query_lens, strict=True): + seq_len = int(seq_len_value) + query_len = int(query_len_value) + if query_len < 0: + return None + if query_len == 0: + continue + start_pos = seq_len - query_len + if start_pos < 0: + return None + first_pos = ( + (start_pos + compress_ratio) // compress_ratio + ) * compress_ratio - 1 + for pos in range(first_pos, seq_len, compress_ratio): + token_offset = token_start + pos - start_pos + if 0 <= token_offset < num_tokens: + active_offsets.append(token_offset) + token_start += query_len + + if token_start != num_tokens: + return None + return active_offsets + + +def _deepseek_v4_hca_active_token_indices( + ctx: ForwardContext, + metadata: DeepseekV4ForwardMetadata, + positions: torch.Tensor, + compress_ratio: int, +) -> torch.Tensor | None: + active_offsets = _deepseek_v4_hca_active_token_offsets( + ctx, + metadata, + positions, + compress_ratio, + ) + if active_offsets is None: + return None + if not active_offsets: + return torch.empty(0, dtype=torch.int64, device=positions.device) + return torch.tensor(active_offsets, dtype=torch.int64, device=positions.device) + + def _dequant_fp8_weight(layer: nn.Module, shape: tuple[int, ...]) -> torch.Tensor: weight = layer.weight.view(*shape) scale = getattr(layer, "weight_scale_inv", None) @@ -2864,6 +2936,42 @@ def forward( return kv, score kv_cache_block_size = pool.get_compressed_block_size(layer_index) + insert_token_to_req_indices = metadata.token_to_req_indices[: positions.numel()] + direct_active_token_indices = None + if self.compress_ratio == 128: + active_hit = ( + memo.get(("hca_active", self.compress_ratio)) + if memo is not None + else None + ) + if active_hit is not None: + direct_active_token_indices = active_hit + else: + active_metadata = metadata + forward_mode = getattr(ctx, "forward_mode", None) + if forward_mode is not None and forward_mode.is_mixed(): + full_metadata = getattr(ctx.attn_backend, "forward_metadata", None) + if full_metadata is not None: + active_metadata = full_metadata + active_indices = _deepseek_v4_hca_active_token_indices( + ctx, + active_metadata, + positions, + self.compress_ratio, + ) + if active_indices is None and active_metadata is not metadata: + active_indices = _deepseek_v4_hca_active_token_indices( + ctx, + metadata, + positions, + self.compress_ratio, + ) + if active_indices is not None: + direct_active_token_indices = active_indices + if memo is not None: + memo[("hca_active", self.compress_ratio)] = ( + direct_active_token_indices + ) compressed_hit = ( memo.get(("compressed", self.compress_ratio)) if memo is not None else None ) @@ -2874,9 +2982,7 @@ def forward( compressed_slots = cache_metadata.compressed_slot_mapping( positions, self.compress_ratio, - token_to_req_indices=metadata.token_to_req_indices[ - : positions.numel() - ], + token_to_req_indices=insert_token_to_req_indices, query_start_loc=metadata.query_start_loc, seq_lens=metadata.seq_lens, kv_cache_block_size=kv_cache_block_size, @@ -2888,27 +2994,64 @@ def forward( if memo is not None: memo[("compressed", self.compress_ratio)] = compressed_slots with nvtx_range(f"{profile_prefix}_cache_insert"): - insert = ( - deepseek_v4_csa_compress_kv_cache_insert - if self.compress_ratio == 4 - else deepseek_v4_hca_compress_kv_cache_insert - ) - insert( - state_cache=state_cache, - token_to_req_indices=metadata.token_to_req_indices[: positions.numel()], - positions=positions, - compressor_slot_mapping=state_slot_mapping, - block_table=state_block_table, - block_table_base_offsets=state_base_logical_page, - compressor_block_size=state_block_size, - rms_norm_weight=self.norm.weight, - rms_norm_eps=self.norm.variance_epsilon, - cos_sin_cache=cos_sin_cache, - kv_cache_2d=pool.get_compressed_kv_buffer_2d(layer_index), - kv_slot_mapping=compressed_slots, - kv_cache_block_size=kv_cache_block_size, - compress_ratio=self.compress_ratio, - ) + if self.compress_ratio == 4: + deepseek_v4_csa_compress_kv_cache_insert( + state_cache=state_cache, + token_to_req_indices=insert_token_to_req_indices, + positions=positions, + compressor_slot_mapping=state_slot_mapping, + block_table=state_block_table, + block_table_base_offsets=state_base_logical_page, + compressor_block_size=state_block_size, + rms_norm_weight=self.norm.weight, + rms_norm_eps=self.norm.variance_epsilon, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=pool.get_compressed_kv_buffer_2d(layer_index), + kv_slot_mapping=compressed_slots, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=self.compress_ratio, + ) + else: + if ( + direct_active_token_indices is not None + and direct_active_token_indices.numel() > 0 + ): + deepseek_v4_hca_direct_compress_kv_cache_insert( + state_cache=state_cache, + kv=kv, + score=score, + ape=self.ape, + token_to_req_indices=insert_token_to_req_indices, + query_start_loc=metadata.query_start_loc, + positions=positions, + block_table=state_block_table, + compressor_block_size=state_block_size, + rms_norm_weight=self.norm.weight, + rms_norm_eps=self.norm.variance_epsilon, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=pool.get_compressed_kv_buffer_2d(layer_index), + kv_slot_mapping=compressed_slots, + kv_cache_block_size=kv_cache_block_size, + active_token_indices=direct_active_token_indices, + block_table_base_offsets=state_base_logical_page, + ) + if direct_active_token_indices is None: + deepseek_v4_hca_compress_kv_cache_insert( + state_cache=state_cache, + token_to_req_indices=insert_token_to_req_indices, + positions=positions, + compressor_slot_mapping=state_slot_mapping, + block_table=state_block_table, + block_table_base_offsets=state_base_logical_page, + compressor_block_size=state_block_size, + rms_norm_weight=self.norm.weight, + rms_norm_eps=self.norm.variance_epsilon, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=pool.get_compressed_kv_buffer_2d(layer_index), + kv_slot_mapping=compressed_slots, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=self.compress_ratio, + ) return kv, score diff --git a/test/runtime/cache/test_deepseek_v4_l2_offload.py b/test/runtime/cache/test_deepseek_v4_l2_offload.py new file mode 100644 index 000000000..61c94d188 --- /dev/null +++ b/test/runtime/cache/test_deepseek_v4_l2_offload.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.parametrize( + ("enable_kvstore", "paged_cache_groups", "draft_pool", "should_raise"), + [ + (True, [object()], object(), True), + (False, [object()], object(), False), + (True, [object()], None, False), + (True, [], object(), False), + ], +) +def test_grouped_kvstore_rejects_only_paged_draft_pool( + enable_kvstore, paged_cache_groups, draft_pool, should_raise +): + from tokenspeed.runtime.engine.event_loop import ( + _validate_grouped_kvstore_draft_pool, + ) + + if should_raise: + with pytest.raises(NotImplementedError, match="--disable-kvstore"): + _validate_grouped_kvstore_draft_pool( + enable_kvstore, paged_cache_groups, draft_pool + ) + else: + _validate_grouped_kvstore_draft_pool( + enable_kvstore, paged_cache_groups, draft_pool + ) + + +def test_deepseek_v4_scheduler_host_pages_follow_kvstore_mode(): + from tokenspeed.runtime.engine.event_loop import ( + _paged_cache_host_group_pages_for_scheduler, + ) + from tokenspeed.runtime.engine.scheduler_utils import make_config + + host_pool = SimpleNamespace(paged_cache_group_page_counts={"v4.swa_kv": 1}) + + hidden_host_groups = _paged_cache_host_group_pages_for_scheduler(False, host_pool) + visible_host_groups = _paged_cache_host_group_pages_for_scheduler(True, host_pool) + assert hidden_host_groups == {} + assert visible_host_groups == {"v4.swa_kv": 1} + + common_config = dict( + num_device_pages=8, + max_scheduled_tokens=4, + max_batch_size=2, + page_size=64, + enable_l3_storage=False, + prefetch_threshold=4, + role="null", + ) + disabled_config = make_config( + **common_config, + num_host_pages=32, + disable_l2_cache=True, + paged_cache_host_group_pages=hidden_host_groups, + ) + enabled_config = make_config( + **common_config, + num_host_pages=0, + disable_l2_cache=False, + paged_cache_host_group_pages=visible_host_groups, + ) + + assert disabled_config.disable_l2_cache + assert disabled_config.num_host_pages == 32 + assert disabled_config.paged_cache_host_group_pages == {} + assert enabled_config.paged_cache_host_group_pages == {"v4.swa_kv": 1} + + +def _make_v4_pool(): + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, + deepseek_v4_cache_layout_from_config, + ) + + hf_config = SimpleNamespace( + compress_ratios=(1, 4, 128), + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + ) + layout = deepseek_v4_cache_layout_from_config( + hf_config, + page_size=64, + use_fp4_indexer_cache=True, + ) + pool = DeepseekV4TokenToKVPool( + size=512, + model_dtype=torch.bfloat16, + layout=layout, + layer_num=3, + device="cpu", + enable_memory_saver=False, + max_batch_size=2, + max_context_len=512, + page_size=64, + rank=0, + hf_config=hf_config, + max_scheduled_tokens=64, + ) + return pool + + +def _make_host_pool(device_pool, ratio: float = 1.5): + from tokenspeed.runtime.cache.deepseek_v4_cache_host import ( + DeepseekV4TokenToKVPoolHost, + ) + + return DeepseekV4TokenToKVPoolHost( + device_pool, + host_to_device_ratio=ratio, + host_size_gb=0, + register_host=False, + ) + + +def _make_ratio4_paged_transfers(): + return [ + SimpleNamespace( + group_id="v4.c4a.compressed_kv", + src_pages=[3, 1, 3], + dst_pages=[9, 7, 9], + ), + SimpleNamespace( + group_id="v4.c4a.compressed_kv", + src_pages=[2, 1], + dst_pages=[8, 7], + ), + ] + + +def _patch_direct_transfer_recorder(monkeypatch): + from tokenspeed.runtime.cache.transfer import deepseek_v4_pool + + calls = [] + + def fake_transfer_kv_direct( + *, + src_layers, + dst_layers, + src_indices, + dst_indices, + page_size, + ): + calls.append( + ( + src_layers, + dst_layers, + src_indices.tolist(), + dst_indices.tolist(), + page_size, + ) + ) + + monkeypatch.setattr( + deepseek_v4_pool, + "transfer_kv_direct", + fake_transfer_kv_direct, + ) + return calls + + +def _assert_ratio4_direct_call(call, src_tensors, dst_tensors): + src_layers, dst_layers, src_indices, dst_indices, page_size = call + assert src_indices == [1, 2, 3] + assert dst_indices == [7, 8, 9] + assert page_size == 1 + assert src_layers[0] is src_tensors[0] + assert src_layers[1] is src_tensors[1] + assert dst_layers[0] is dst_tensors[0] + assert dst_layers[1] is dst_tensors[1] + + +def test_deepseek_v4_host_pool_shapes_and_group_counts(): + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool, ratio=1.5) + + for group_id, device_pages in device_pool.paged_cache_group_page_counts.items(): + assert host_pool.paged_cache_group_page_counts[group_id] >= device_pages + + assert ( + host_pool.swa_kv_buffer[0].shape[1:] == device_pool.swa_kv_buffer[0].shape[1:] + ) + assert host_pool.compressed_kv_buffer[1].shape[1:] == ( + device_pool.compressed_kv_buffer[1].shape[1:] + ) + assert host_pool.indexer_kv_buffer[1].shape[0] == ( + host_pool.compressed_kv_buffer[1].shape[0] + ) + assert host_pool.indexer_state_buffer[1].shape[1:] == ( + device_pool.indexer_state_buffer[1].shape[1:] + ) + assert host_pool.page_num > 0 + assert host_pool.total_bytes > 0 + + +def test_deepseek_v4_host_group_page_sizing_keeps_one_usable_page_per_group(): + from tokenspeed.runtime.cache.deepseek_v4_cache_host import ( + _allocate_host_group_pages, + ) + + ratio_counts = _allocate_host_group_pages( + device_counts={"a": 3, "b": 1}, + page_bytes={"a": 100, "b": 200}, + host_ratio=0.1, + host_size_gb=0, + ) + assert ratio_counts == {"a": 2, "b": 2} + + size_counts = _allocate_host_group_pages( + device_counts={"a": 3, "b": 1}, + page_bytes={"a": 100, "b": 200}, + host_ratio=1.0, + host_size_gb=1, + ) + assert size_counts["a"] >= 2 + assert size_counts["b"] >= 2 + + capped_counts = _allocate_host_group_pages( + device_counts={"a": 100, "b": 100}, + page_bytes={"a": 10, "b": 30}, + host_ratio=10.0, + host_size_gb=0, + host_budget_bytes=120, + ) + assert capped_counts["a"] >= 2 + assert capped_counts["b"] >= 2 + assert ( + sum( + capped_counts[group_id] * {"a": 10, "b": 30}[group_id] + for group_id in capped_counts + ) + <= 120 + ) + + with pytest.raises(ValueError, match="too small"): + _allocate_host_group_pages( + device_counts={"a": 100, "b": 100}, + page_bytes={"a": 10, "b": 30}, + host_ratio=10.0, + host_size_gb=0, + host_budget_bytes=79, + ) + + +def test_deepseek_v4_shadow_capacity_uses_complete_history_limit(): + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + host_pool.paged_cache_group_page_counts.update( + { + "v4.c4a.compressed_kv": 2, + "v4.c128a.compressed_kv": 32, + } + ) + + expected_usable_pages = min( + ( + (host_pool.paged_cache_group_page_counts[spec.group_id] - 1) + * spec.rows_per_page + * spec.entry_stride_tokens + + device_pool.page_size + - 1 + ) + // device_pool.page_size + for spec in device_pool.paged_cache_group_specs + if spec.family == "history" + ) + + assert host_pool._compute_shadow_page_num(device_pool) == expected_usable_pages + 1 + + +def test_deepseek_v4_descriptor_expansion_maps_paged_groups(): + from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool + + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + transfer_pool = DeepseekV4CachePool(device_pool, host_pool, io_backend="direct") + + checks = [ + ( + "v4.swa_kv", + None, + [("swa_kv_buffer", 0), ("swa_kv_buffer", 1), ("swa_kv_buffer", 2)], + ), + ( + "v4.c4a.compressed_kv", + 1, + [("compressed_kv_buffer", 1), ("indexer_kv_buffer", 1)], + ), + ("v4.c4a.compressor_state", 1, [("compressor_state_buffer", 1)]), + ("v4.c128a.compressor_state", 2, [("compressor_state_buffer", 2)]), + ("v4.c4a.indexer_compressor_state", 1, [("indexer_state_buffer", 1)]), + ] + for group_id, layer_idx, expected in checks: + refs = transfer_pool.tensor_refs_for_group(group_id, layer_idx=layer_idx) + assert len(refs) == len(expected) + for ref, (buffer_name, layer_id) in zip(refs, expected): + assert ref.layer_id == layer_id + assert ref.device_tensor is getattr(device_pool, buffer_name)[layer_id] + assert ref.host_tensor is getattr(host_pool, buffer_name)[layer_id] + assert ref.page_bytes == ref.device_tensor[0].nbytes + + +def test_deepseek_v4_paged_pool_prepares_coalesced_transfers(): + from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool + + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + transfer_pool = DeepseekV4CachePool(device_pool, host_pool, io_backend="kernel") + transfers = _make_ratio4_paged_transfers() + + prepared = transfer_pool.prepare_paged_transfers(transfers) + + assert len(prepared) == 1 + assert prepared[0].page_count == 3 + assert prepared[0].span_count == 1 + assert prepared[0].src_indices.tolist() == [1, 2, 3] + assert prepared[0].dst_indices.tolist() == [7, 8, 9] + + +def test_deepseek_v4_paged_pool_writeback_prepared_maps_group_tensors(monkeypatch): + from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool + + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + transfer_pool = DeepseekV4CachePool(device_pool, host_pool, io_backend="kernel") + calls = _patch_direct_transfer_recorder(monkeypatch) + prepared = transfer_pool.prepare_paged_transfers(_make_ratio4_paged_transfers()) + transfer_pool.writeback_prepared_paged(prepared) + + assert len(calls) == 1 + _assert_ratio4_direct_call( + calls[-1], + (device_pool.compressed_kv_buffer[1], device_pool.indexer_kv_buffer[1]), + (host_pool.compressed_kv_buffer[1], host_pool.indexer_kv_buffer[1]), + ) + + +def test_deepseek_v4_paged_pool_loadback_prepared_maps_group_tensors(monkeypatch): + from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool + + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + transfer_pool = DeepseekV4CachePool(device_pool, host_pool, io_backend="kernel") + calls = _patch_direct_transfer_recorder(monkeypatch) + prepared = transfer_pool.prepare_paged_transfers(_make_ratio4_paged_transfers()) + transfer_pool.loadback_prepared_paged(prepared, layer_idx=1) + + assert len(calls) == 1 + _assert_ratio4_direct_call( + calls[-1], + (host_pool.compressed_kv_buffer[1], host_pool.indexer_kv_buffer[1]), + (device_pool.compressed_kv_buffer[1], device_pool.indexer_kv_buffer[1]), + ) + + +def test_deepseek_v4_paged_pool_reuses_and_invalidates_h2d_scatter_plan( + monkeypatch, +): + from tokenspeed.runtime.cache.transfer import deepseek_v4_pool + from tokenspeed.runtime.cache.transfer.deepseek_v4_pool import DeepseekV4CachePool + + device_pool = _make_v4_pool() + host_pool = _make_host_pool(device_pool) + transfer_pool = DeepseekV4CachePool(device_pool, host_pool, io_backend="kernel") + monkeypatch.setattr(deepseek_v4_pool, "_H2D_SCATTER_MIN_COPY_CALLS", 0) + prepared = transfer_pool.prepare_paged_transfers(_make_ratio4_paged_transfers()) + prepared_plans = [] + launches = [] + + def prepare_plan(src_layers, dst_layers, entry_ids): + plan = object() + prepared_plans.append((plan, src_layers, dst_layers, entry_ids)) + return plan, "" + + def launch(plan, src_indices, dst_indices, page_size, entry_begin, entry_end): + launches.append( + (plan, src_indices, dst_indices, page_size, entry_begin, entry_end) + ) + return SimpleNamespace( + used=True, + buckets=1, + kernel_launches=1, + fallback_reason="", + ) + + monkeypatch.setattr( + deepseek_v4_pool, + "prepare_kv_direct_h2d_scatter_plan", + prepare_plan, + ) + monkeypatch.setattr( + deepseek_v4_pool, + "transfer_kv_direct_h2d_scatter_prepared", + launch, + ) + + transfer_pool.loadback_prepared_paged(prepared, layer_idx=1) + transfer_pool.loadback_prepared_paged(prepared, layer_idx=1) + + assert len(prepared_plans) == 1 + assert len(launches) == 2 + assert launches[0][0] is launches[1][0] is prepared_plans[0][0] + assert launches[0][1] is launches[1][1] + assert launches[0][2] is launches[1][2] + + host_pool.compressed_kv_buffer[1] = host_pool.compressed_kv_buffer[1].clone() + transfer_pool.loadback_prepared_paged(prepared, layer_idx=1) + + assert len(prepared_plans) == 2 + assert launches[-1][0] is prepared_plans[-1][0] + + +def test_deepseek_v4_layer_getters_wait_on_registered_counter(): + device_pool = _make_v4_pool() + waits = [] + + class FakeCounter: + def wait_until(self, layer_id: int): + waits.append(layer_id) + + device_pool.register_layer_transfer_counter(FakeCounter()) + + device_pool.get_swa_kv_buffer(0) + device_pool.get_compressed_kv_buffer_2d(1) + device_pool.get_compressor_state_buffer(1) + device_pool.get_indexer_kv_buffer_2d(1) + device_pool.get_indexer_state_buffer(1) + device_pool.get_kv_buffer(2) + + assert waits == [0, 1, 1, 1, 1, 2] diff --git a/test/runtime/cache/test_transfer_host_executor.py b/test/runtime/cache/test_transfer_host_executor.py index 339f66b16..3f2fffb19 100644 --- a/test/runtime/cache/test_transfer_host_executor.py +++ b/test/runtime/cache/test_transfer_host_executor.py @@ -1,11 +1,17 @@ from __future__ import annotations from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch -from tokenspeed.runtime.cache.transfer.types import CacheKind, Location, TransferUnit +from tokenspeed.runtime.cache.transfer.types import ( + PAGED_CACHE_KIND, + CacheKind, + Location, + TransferUnit, +) class FakeEvent: @@ -109,6 +115,48 @@ def reset(self): self.counter.reset() +class FakePreparedPagedPool: + kind = PAGED_CACHE_KIND + loadback_layer_chunk_size = 2 + + def __init__(self, num_layers: int): + self._num_layers = num_layers + self.device = torch.device("cpu") + self.counter = FakeCounter(num_layers) + self.prepares = [] + self.prepared_writebacks = [] + self.prepared_loadbacks = [] + self.prepared_range_loadbacks = [] + + def num_layers(self): + return self._num_layers + + def get_layer_done_counter(self): + return self.counter + + def prepare_paged_transfers(self, transfers): + prepared = tuple(("prepared", id(transfer)) for transfer in transfers) + self.prepares.append(list(transfers)) + return prepared + + def writeback_prepared_paged(self, prepared): + self.prepared_writebacks.append(tuple(prepared)) + + def loadback_prepared_paged(self, prepared, layer_idx: int): + self.prepared_loadbacks.append((layer_idx, tuple(prepared))) + + def loadback_prepared_paged_range(self, prepared, layer_start: int, layer_end: int): + self.prepared_range_loadbacks.append((layer_start, layer_end, tuple(prepared))) + + +class FakeCacheOps: + class WriteBackOp: + pass + + class LoadBackOp: + pass + + def _patch_host_executor_device(monkeypatch): import tokenspeed.runtime.cache.executor.host_executor as host_executor @@ -225,23 +273,55 @@ def test_host_executor_loadback_uses_independent_layer_counters(monkeypatch): executor.set_consumer(CacheKind.KV, [0]) executor.set_consumer(CacheKind.MAMBA, [0]) assert kv_pool.counter.consumer == [0] + assert executor.drain() == [] -def test_memory_executor_submit_dispatches_flat_op_by_cache_kind(monkeypatch): - import tokenspeed.runtime.cache.executor.memory_executor as memory_executor +def test_host_executor_paged_cache_loadback_reuses_prepared_transfers(monkeypatch): + HostExecutor = _patch_host_executor_device(monkeypatch) + paged_pool = FakePreparedPagedPool(num_layers=3) + executor = HostExecutor(pools=[], paged_pool=paged_pool, io_backend="kernel") + transfers = [object(), object()] - class FakeCache: - class WriteBackOp: - pass + executor.enqueue_paged_cache_loadback(33, transfers) + executor.flush() - class LoadBackOp: - pass + assert paged_pool.prepares == [transfers] + expected_prepared = tuple(("prepared", id(transfer)) for transfer in transfers) + assert paged_pool.prepared_range_loadbacks == [ + (0, 2, expected_prepared), + (2, 3, expected_prepared), + ] + assert paged_pool.prepared_loadbacks == [] + assert all(event.recorded for event in paged_pool.counter.events[0].load_events) + assert executor.get_producer_index(PAGED_CACHE_KIND, 33) == 0 + executor.set_consumer(PAGED_CACHE_KIND, [0]) + assert paged_pool.counter.consumer == [0] + results = executor.drain() + assert len(results) == 1 + assert results[0].op_id == 33 + assert results[0].success is True + assert type(results[0]).__name__ == "LoadBackDoneEvent" - class PrefetchOp: - pass - class BackUpOp: - pass +def test_host_executor_paged_cache_writeback_uses_prepared_transfers(monkeypatch): + HostExecutor = _patch_host_executor_device(monkeypatch) + paged_pool = FakePreparedPagedPool(num_layers=3) + executor = HostExecutor(pools=[], paged_pool=paged_pool, io_backend="kernel") + transfers = [object(), object()] + + executor.enqueue_paged_cache_writeback(44, transfers) + executor.flush() + + assert paged_pool.prepares == [transfers] + expected_prepared = tuple(("prepared", id(transfer)) for transfer in transfers) + assert paged_pool.prepared_writebacks == [expected_prepared] + results = executor.drain() + assert [event.op_id for event in results] == [44] + assert all(event.success for event in results) + + +def test_memory_executor_submit_dispatches_flat_op_by_cache_kind(monkeypatch): + import tokenspeed.runtime.cache.executor.memory_executor as memory_executor class FakeHostExec: def __init__(self): @@ -264,12 +344,12 @@ def enqueue_loadback(self, op_id, src_pages, dst_pages, kind=CacheKind.KV): def flush(self): self.order.append(("flush",)) - monkeypatch.setattr(memory_executor, "Cache", FakeCache) + monkeypatch.setattr(memory_executor, "Cache", FakeCacheOps) executor = object.__new__(memory_executor.MemoryExecutor) executor.host_exec = FakeHostExec() executor.storage_exec = None - wb = FakeCache.WriteBackOp() + wb = FakeCacheOps.WriteBackOp() wb.op_ids = [7] wb.src_pages = [[1]] wb.dst_pages = [[11]] @@ -284,7 +364,7 @@ def flush(self): ] assert executor.host_exec.completed_writebacks == [] - lb = FakeCache.LoadBackOp() + lb = FakeCacheOps.LoadBackOp() lb.op_ids = [9] lb.src_pages = [[10]] lb.dst_pages = [[20]] @@ -298,27 +378,81 @@ def flush(self): ] -def test_memory_executor_submit_plan_keeps_generic_submit_signature(monkeypatch): +def test_memory_executor_submit_dispatches_paged_cache_transfers(monkeypatch): import tokenspeed.runtime.cache.executor.memory_executor as memory_executor - class FakeCache: - class WriteBackOp: - pass + class FakeHostExec: + def __init__(self): + self.pools = {CacheKind.KV: object()} + self.completed_writebacks = [] + self.writebacks = [] + self.loadbacks = [] + self.paged_writebacks = [] + self.paged_loadbacks = [] - class LoadBackOp: - pass + def enqueue_writeback(self, op_id, src_pages, dst_pages, **kwargs): + self.writebacks.append((op_id, src_pages, dst_pages, kwargs)) - class PrefetchOp: - pass + def enqueue_loadback(self, op_id, src_pages, dst_pages, **kwargs): + self.loadbacks.append((op_id, src_pages, dst_pages, kwargs)) - class BackUpOp: + def enqueue_paged_cache_writeback(self, op_id, transfers, is_retract=False): + self.paged_writebacks.append((op_id, transfers, is_retract)) + + def enqueue_paged_cache_loadback(self, op_id, transfers): + self.paged_loadbacks.append((op_id, transfers)) + + def flush(self): pass - monkeypatch.setattr(memory_executor, "Cache", FakeCache) + monkeypatch.setattr(memory_executor, "Cache", FakeCacheOps) + monkeypatch.setattr(memory_executor.logger, "isEnabledFor", lambda _level: True) + executor = object.__new__(memory_executor.MemoryExecutor) + executor.host_exec = FakeHostExec() + executor.storage_exec = None + + transfer = SimpleNamespace( + group_id="v4.test", + src_pages=[1], + dst_pages=[2], + ) + wb = FakeCacheOps.WriteBackOp() + wb.op_ids = [7, 8] + wb.src_pages = [[], [3]] + wb.dst_pages = [[], [4]] + wb.src_pages_by_kind = {"kv": [[], [3]], "mamba": [[], []]} + wb.dst_pages_by_kind = {"kv": [[], [4]], "mamba": [[], []]} + wb.paged_cache_transfers = [[transfer], []] + wb.is_retract = [True, False] + executor.submit(wb) + + assert executor.host_exec.paged_writebacks == [(7, [transfer], True)] + assert executor.host_exec.writebacks == [ + (8, [3], [4], {"is_retract": False, "kind": CacheKind.KV}) + ] + assert executor.host_exec.completed_writebacks == [] + + lb = FakeCacheOps.LoadBackOp() + lb.op_ids = [9, 10] + lb.src_pages = [[], [5]] + lb.dst_pages = [[], [6]] + lb.src_pages_by_kind = {"kv": [[], [5]], "mamba": [[], []]} + lb.dst_pages_by_kind = {"kv": [[], [6]], "mamba": [[], []]} + lb.paged_cache_transfers = [[transfer], []] + executor.submit(lb) + + assert executor.host_exec.paged_loadbacks == [(9, [transfer])] + assert executor.host_exec.loadbacks == [(10, [5], [6], {"kind": CacheKind.KV})] + + +def test_memory_executor_submit_plan_keeps_generic_submit_signature(monkeypatch): + import tokenspeed.runtime.cache.executor.memory_executor as memory_executor + + monkeypatch.setattr(memory_executor, "Cache", FakeCacheOps) executor = object.__new__(memory_executor.MemoryExecutor) executor.seen = [] - wb = FakeCache.WriteBackOp() + wb = FakeCacheOps.WriteBackOp() plan = type("Plan", (), {"cache": [wb]})() def submit(self, op): @@ -335,19 +469,6 @@ def submit(self, op): def test_memory_executor_mamba_layerwise_cow_uses_dedicated_context(monkeypatch): import tokenspeed.runtime.cache.executor.memory_executor as memory_executor - class FakeCache: - class WriteBackOp: - pass - - class LoadBackOp: - pass - - class PrefetchOp: - pass - - class BackUpOp: - pass - class FakeHostExec: def __init__(self): self.pools = {CacheKind.KV: object(), CacheKind.MAMBA: object()} @@ -369,13 +490,13 @@ def enqueue_loadback( def flush(self): pass - monkeypatch.setattr(memory_executor, "Cache", FakeCache) + monkeypatch.setattr(memory_executor, "Cache", FakeCacheOps) executor = object.__new__(memory_executor.MemoryExecutor) executor.host_exec = FakeHostExec() executor.storage_exec = None executor.set_mamba_layerwise_cow({40: [400]}) - lb = FakeCache.LoadBackOp() + lb = FakeCacheOps.LoadBackOp() lb.op_ids = [9] lb.src_pages = [[10]] lb.dst_pages = [[20]] diff --git a/test/runtime/test_deepseek_v4_attention_ops.py b/test/runtime/test_deepseek_v4_attention_ops.py index e0c364247..0d388d991 100644 --- a/test/runtime/test_deepseek_v4_attention_ops.py +++ b/test/runtime/test_deepseek_v4_attention_ops.py @@ -15,6 +15,7 @@ import os import sys import unittest +from types import SimpleNamespace import torch @@ -36,6 +37,7 @@ deepseek_v4_swa_scale_dim, deepseek_v4_swa_token_stride, ) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode from tokenspeed.runtime.layers.attention.deepseek_v4_ops import ( deepseek_v4_combine_dense_swa_indices, deepseek_v4_combine_topk_swa_indices, @@ -46,6 +48,7 @@ deepseek_v4_decode_swa_indices_and_lens, deepseek_v4_dequantize_and_gather_k_cache, deepseek_v4_hca_compress_kv_cache_insert, + deepseek_v4_hca_direct_compress_kv_cache_insert, deepseek_v4_prepare_indexer_q_mxfp4, dequantize_deepseek_v4_fp8_ds_mla_cache, fused_qnorm_rope_kv_insert, @@ -59,6 +62,7 @@ _mask_invalid_graph_tokens, ) from tokenspeed.runtime.models.deepseek_v4 import ( + _deepseek_v4_hca_active_token_indices, _deepseek_v4_sanitize_swa_slot_mapping, ) @@ -70,6 +74,123 @@ SWA_SCALE_DIM = deepseek_v4_swa_scale_dim(HEAD_DIM, ROPE_DIM) +def _make_hca_compressor_fixture(seed: int): + torch.manual_seed(seed) + device = torch.device("cuda") + dtype = torch.bfloat16 + compress_ratio = 128 + state_block_size = 8 + kv_cache_block_size = 2 + num_tokens = compress_ratio * 2 + prefix_len = 232 + logical_start = 128 + full_tokens = prefix_len + num_tokens - logical_start + num_state_blocks = full_tokens // state_block_size + eps = 1.0e-6 + + full_kv = torch.randn(full_tokens, HEAD_DIM, device=device, dtype=dtype) + full_score = torch.randn(full_tokens, HEAD_DIM, device=device, dtype=dtype) * 0.1 + tail_start = prefix_len - logical_start + kv = full_kv[tail_start:] + score = full_score[tail_start:] + ape = ( + torch.randn( + compress_ratio, + HEAD_DIM, + device=device, + dtype=torch.float32, + ) + * 0.01 + ) + state_cache = torch.zeros( + num_state_blocks, + state_block_size, + HEAD_DIM * 2, + device=device, + dtype=torch.float32, + ) + full_positions = torch.arange( + logical_start, + prefix_len + num_tokens, + device=device, + dtype=torch.int64, + ) + positions = torch.arange( + prefix_len, + prefix_len + num_tokens, + device=device, + dtype=torch.int64, + ) + full_state_slots = torch.arange(full_tokens, device=device, dtype=torch.int64) + save_deepseek_v4_compressor_state( + kv=full_kv, + score=full_score, + ape=ape, + state_cache=state_cache, + slot_mapping=full_state_slots, + positions=full_positions, + block_size=state_block_size, + compress_ratio=compress_ratio, + ) + + token_to_req_indices = torch.zeros(num_tokens, device=device, dtype=torch.int32) + block_table = torch.arange(num_state_blocks, device=device, dtype=torch.int32).view( + 1, + -1, + ) + block_table_base_offsets = torch.tensor( + [logical_start // state_block_size], + device=device, + dtype=torch.int32, + ) + kv_slots = torch.full((num_tokens,), -1, device=device, dtype=torch.int64) + active_token_indices = torch.tensor( + [255 - prefix_len, 383 - prefix_len], + device=device, + dtype=torch.int64, + ) + kv_slots[active_token_indices[0]] = 0 + kv_slots[active_token_indices[1]] = 1 + cos_sin = torch.randn(512, ROPE_DIM, device=device, dtype=torch.float32) * 0.05 + rms_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.float32) * 0.1 + 1.0 + cache_shape = ( + 1, + kv_cache_block_size * (SWA_TOKEN_STRIDE + SWA_SCALE_DIM), + ) + common_kwargs = dict( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + block_table=block_table, + block_table_base_offsets=block_table_base_offsets, + compressor_block_size=state_block_size, + rms_norm_weight=rms_weight, + rms_norm_eps=eps, + cos_sin_cache=cos_sin, + kv_slot_mapping=kv_slots, + kv_cache_block_size=kv_cache_block_size, + ) + paged_kwargs = dict( + common_kwargs, + compressor_slot_mapping=full_state_slots[tail_start:], + compress_ratio=compress_ratio, + ) + direct_kwargs = dict( + common_kwargs, + kv=kv, + score=score, + ape=ape, + query_start_loc=torch.tensor([0, num_tokens], device=device, dtype=torch.int32), + active_token_indices=active_token_indices, + ) + return SimpleNamespace( + active_token_indices=active_token_indices, + cache_shape=cache_shape, + direct_kwargs=direct_kwargs, + paged_kwargs=paged_kwargs, + ) + + def _apply_gptj_rope_with_nope( x: torch.Tensor, positions: torch.Tensor, @@ -335,6 +456,35 @@ def test_slot_mapping_guard_masks_invalid_graph_tokens(self): torch.tensor([0, -1, -1, -1], dtype=torch.int64), ) + def test_hca_active_token_indices_cover_cached_and_plain_prefill(self): + ctx = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + cases = ( + ("cached", [300, 512], [200, 130], [27, 155, 201, 329]), + ("plain", [256], [256], [127, 255]), + ) + for name, seq_lens, query_lens, expected in cases: + with self.subTest(name=name): + num_tokens = sum(query_lens) + metadata = SimpleNamespace( + num_prefill_reqs=len(seq_lens), + num_prefill_tokens=num_tokens, + seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int32), + query_lens_cpu=torch.tensor(query_lens, dtype=torch.int32), + ) + + indices = _deepseek_v4_hca_active_token_indices( + ctx, + metadata, + torch.empty(num_tokens, dtype=torch.int64), + compress_ratio=128, + ) + + self.assertIsNotNone(indices) + torch.testing.assert_close( + indices, + torch.tensor(expected, dtype=torch.int64), + ) + def test_slot_mapping_guard_expands_per_request_validity(self): slots = torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int64) is_valid_token = torch.tensor([True, False]) @@ -774,6 +924,38 @@ def test_hca_compressor_state_insert_matches_reference(self): 0, ) + def test_hca_direct_compressor_matches_paged_insert(self): + fixture = _make_hca_compressor_fixture(seed=2468) + paged_cache = torch.zeros(fixture.cache_shape, device="cuda", dtype=torch.uint8) + direct_cache = torch.zeros_like(paged_cache) + + deepseek_v4_hca_compress_kv_cache_insert( + kv_cache_2d=paged_cache, + **fixture.paged_kwargs, + ) + deepseek_v4_hca_direct_compress_kv_cache_insert( + kv_cache_2d=direct_cache, + **{ + **fixture.direct_kwargs, + "active_token_indices": torch.cat( + ( + torch.tensor([-1], device="cuda", dtype=torch.int64), + fixture.active_token_indices, + torch.tensor( + [fixture.direct_kwargs["kv"].shape[0]], + device="cuda", + dtype=torch.int64, + ), + ) + ), + }, + ) + + torch.cuda.synchronize() + torch.testing.assert_close( + direct_cache.cpu(), paged_cache.cpu(), atol=0, rtol=0 + ) + def test_csa_compressor_state_insert_matches_reference(self): torch.manual_seed(5678) device = torch.device("cuda") diff --git a/test/runtime/test_generation_output_processor.py b/test/runtime/test_generation_output_processor.py index d050c7309..b96a6bdca 100644 --- a/test/runtime/test_generation_output_processor.py +++ b/test/runtime/test_generation_output_processor.py @@ -134,6 +134,26 @@ def test_mark_abort_notify_client_flag(): assert not client_state.abort_notify_client +def test_scheduler_abort_publishes_once_and_releases_request_state(): + from tokenspeed.runtime.engine.request_types import FINISH_ABORT + + sender = _Sender() + processor = OutputProcesser(sender, attn_tp_rank=0, metrics=_Metrics()) + state = _state([1, 2, 3]) + processor.rid_to_state["r"] = state + processor.pending_aborts["r"] = 0.0 + + processor.publish_scheduler_abort("r", "paged cache exhausted") + processor.publish_scheduler_abort("r", "duplicate") + + assert isinstance(state.finished_reason, FINISH_ABORT) + assert state.finished_reason.message == "paged cache exhausted" + assert "r" not in processor.rid_to_state + assert "r" not in processor.pending_aborts + assert len(sender.items) == 1 + assert sender.items[0].rids == ["r"] + + def test_nan_flag_finishes_request_with_numerical_error(): """A request flagged by the NaN guard is finished with ABORT_CODE.NumericalError while the rest of the batch continues.""" diff --git a/test/runtime/test_scheduler_utils_cache_events.py b/test/runtime/test_scheduler_utils_cache_events.py new file mode 100644 index 000000000..a74a7f013 --- /dev/null +++ b/test/runtime/test_scheduler_utils_cache_events.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from tokenspeed.runtime.engine.scheduler_utils import pop_common_cache_event_payloads + + +def test_cache_event_payloads_commit_only_common_success(): + rank0 = [ + {"kind": "WriteBackDoneEvent", "op_id": 3, "success": True}, + {"kind": "WriteBackDoneEvent", "op_id": 5, "success": True}, + ] + rank1 = [ + {"kind": "WriteBackDoneEvent", "op_id": 3, "success": False}, + ] + + ready = pop_common_cache_event_payloads([rank0, rank1]) + + assert ready == [{"kind": "WriteBackDoneEvent", "op_id": 3, "success": False}] + assert pop_common_cache_event_payloads([rank0, []]) == [] diff --git a/test/runtime/test_v4_prefix_cache_metadata.py b/test/runtime/test_v4_prefix_cache_metadata.py index 5814ec2b9..d33e781a6 100644 --- a/test/runtime/test_v4_prefix_cache_metadata.py +++ b/test/runtime/test_v4_prefix_cache_metadata.py @@ -30,6 +30,7 @@ ExecutionEvent, ForwardEvent, PagedCacheGroupConfig, + PagedCacheGroupFamily, PagedCacheRetention, PrefixCacheAdjunctSpec, RequestSpec, @@ -83,6 +84,22 @@ def _make_two_group_config() -> "SchedulerConfig": return cfg +def _make_l2_binding_config() -> "SchedulerConfig": + cfg = _make_two_group_config() + cfg.disable_l2_cache = False + cfg.paged_cache_host_group_pages = {"fh": 32, "swa": 32} + groups = list(cfg.paged_cache_groups) + for group in groups: + group.family = ( + PagedCacheGroupFamily.History + if group.group_id == "fh" + else PagedCacheGroupFamily.State + ) + cfg.paged_cache_groups = groups + cfg.prefix_cache_adjunct.required_groups = ["fh"] + return cfg + + def _post(sched: "Scheduler", payload) -> None: ev = ExecutionEvent() ev.add_event(payload) @@ -123,8 +140,74 @@ def _submit_r2_same_prefix(sched: "Scheduler") -> int: return int(forward.extend_prefix_lens[0]) -class TestV4PrefixCacheMetadata(unittest.TestCase): +def test_paged_l2_transfer_binding_shape() -> None: + sched = Scheduler(_make_l2_binding_config()) + spec = RequestSpec() + spec.request_id = "r1" + spec.tokens = list(range(1, 13)) + sched.submit_requests([spec]) + assert len(sched.next_execution_plan().forward) == 1 + result = ForwardEvent.ExtendResult() + result.request_id = "r1" + result.tokens = [99] + _post(sched, result) + sched.next_execution_plan() + + finish = ForwardEvent.Finish() + finish.request_id = "r1" + _post(sched, finish) + writeback_plan = sched.next_execution_plan() + writeback = next( + (op for op in writeback_plan.cache if hasattr(op, "is_retract")), + None, + ) + assert writeback is not None + assert len(writeback.op_ids) == 1 + assert len(writeback.paged_cache_transfers) == 1 + assert writeback.paged_cache_transfers[0] + + writeback_groups = { + transfer.group_id for transfer in writeback.paged_cache_transfers[0] + } + assert writeback_groups == {"fh", "swa"} + for transfer in writeback.paged_cache_transfers[0]: + assert isinstance(transfer.group_id, str) + assert len(transfer.src_pages) == len(transfer.dst_pages) + assert len(transfer.src_pages) > 0 + + +def test_scheduler_abort_binding_shape() -> None: + config = _make_l2_binding_config() + groups = list(config.paged_cache_groups) + for group in groups: + group.total_pages = 3 if group.group_id == "fh" else 5 + config.paged_cache_groups = groups + sched = Scheduler(config) + specs = [] + for request_id, start in (("z-request", 1), ("a-request", 101)): + spec = RequestSpec() + spec.request_id = request_id + spec.tokens = list(range(start, start + 4)) + specs.append(spec) + sched.submit_requests(specs) + sched.next_execution_plan() + + for request_id, token in (("z-request", 42), ("a-request", 43)): + result = ForwardEvent.ExtendResult() + result.request_id = request_id + result.tokens = [token] + _post(sched, result) + + plan = sched.next_execution_plan() + + assert len(plan.scheduler_aborts) == 1 + request_id, message = plan.scheduler_aborts[0] + assert request_id == "a-request" + assert "capacity" in message + + +class TestV4PrefixCacheMetadata(unittest.TestCase): def setUp(self) -> None: config = _make_two_group_config() fh_config = next( diff --git a/test/runtime/test_v4_sliding_window_groups_smoke.py b/test/runtime/test_v4_sliding_window_groups_smoke.py index fa1e2e230..406b80960 100644 --- a/test/runtime/test_v4_sliding_window_groups_smoke.py +++ b/test/runtime/test_v4_sliding_window_groups_smoke.py @@ -19,7 +19,6 @@ import sys import unittest from types import SimpleNamespace -from unittest.mock import patch import torch @@ -345,9 +344,6 @@ def test_page_counts_positive_finite_and_under_total_times_live(self): self.assertLess(n, bound, spec.group_id) def test_deepseek_v4_pool_exposes_scheduler_cache_groups(self): - from tokenspeed.runtime.layers.attention.kv_cache import ( - deepseek_v4 as deepseek_v4_kv, - ) from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( DeepseekV4TokenToKVPool, deepseek_v4_cache_layout_from_config, @@ -391,7 +387,7 @@ def test_deepseek_v4_pool_exposes_scheduler_cache_groups(self): pool.paged_cache_group_page_counts["v4.c4a.compressed_kv"], 1 ) self.assertFalse(hasattr(pool, "prefix_cache_state_policy")) - self.assertFalse(pool.supports_hierarchical_kv_cache) + self.assertTrue(pool.supports_hierarchical_kv_cache) self.assertEqual( pool.prefix_cache_required_group_ids, ( @@ -400,31 +396,6 @@ def test_deepseek_v4_pool_exposes_scheduler_cache_groups(self): ), ) - class FakePagedCacheScheduler: - @staticmethod - def paged_cache_group_total_pages(group_id: str) -> int: - return 11 - - @staticmethod - def paged_cache_group_available_pages(group_id: str) -> int: - return 4 - - @staticmethod - def paged_cache_group_failed_alloc_count(group_id: str) -> int: - return 2 - - pool.bind_paged_cache_scheduler(FakePagedCacheScheduler()) - with ( - patch.object(deepseek_v4_kv.logger, "isEnabledFor", return_value=True), - patch.object(deepseek_v4_kv.logger, "debug") as log_debug, - ): - pool.maybe_log_paged_cache_group_pages() - log_debug.assert_called_once() - logged_groups = log_debug.call_args.args[1] - self.assertIn("v4.swa_kv: used=7/11", logged_groups) - self.assertIn("v4.c4a.indexer_compressor_state", logged_groups) - self.assertIn("failed_alloc=2", logged_groups) - def test_deepseek_v4_capacity_profile_matches_pool_buffers(self): from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( DeepseekV4TokenToKVPool, diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/deepseek_v4.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/deepseek_v4.py index a5650f7fd..a9f8f63fe 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/deepseek_v4.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/deepseek_v4.py @@ -41,7 +41,7 @@ DEEPSEEK_V4_INDEXER_DIM // DEEPSEEK_V4_MXFP4_BLOCK_SIZE ) DEEPSEEK_V4_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 - +_HCA_DIRECT_LAUNCH_BUCKET = 64 __all__ = [ "deepseek_v4_build_dense_prefill_local_compressed_indices", "deepseek_v4_combine_dense_swa_indices", @@ -51,6 +51,7 @@ "deepseek_v4_decode_swa_indices_and_lens", "deepseek_v4_dequantize_and_gather_k_cache", "deepseek_v4_fused_csa_indexer_mxfp4_cache_insert", + "deepseek_v4_fused_hca_direct_compress_cache_insert", "deepseek_v4_fused_indexer_q_rope_hadamard_mxfp4", "deepseek_v4_fused_sparse_compress_cache_insert", "deepseek_v4_gather_indexer_mxfp4_cache", @@ -60,6 +61,14 @@ ] +def _deepseek_v4_hca_direct_launch_rows(num_actual: int) -> int: + return ( + (num_actual + _HCA_DIRECT_LAUNCH_BUCKET - 1) + // _HCA_DIRECT_LAUNCH_BUCKET + * _HCA_DIRECT_LAUNCH_BUCKET + ) + + @triton.jit def _deepseek_v4_mxfp4_e2m1_nibble(x): abs_x = tl.minimum(tl.abs(x), 6.0) @@ -254,7 +263,18 @@ def deepseek_v4_fused_indexer_q_rope_hadamard_mxfp4( ), weights_out -@triton.jit +# ``block_table_width`` varies per request with cache depth, so keeping it as a +# ``tl.constexpr`` compiles a fresh kernel variant for every distinct width and +# stalls the compute stream on first use of each width. It only bounds +# ``table_idx``, so it is a runtime scalar excluded from alignment specialization +# (see the HCA direct kernel below for the full rationale). +@triton.jit( + do_not_specialize_on_alignment=[ + "block_table_width", + "state_block_size", + "kv_cache_block_size", + ] +) def _deepseek_v4_fused_sparse_compress_cache_kernel( state_cache_ptr, state_cache_stride0, @@ -265,7 +285,7 @@ def _deepseek_v4_fused_sparse_compress_cache_kernel( block_table_ptr, block_table_base_offsets_ptr, block_table_stride, - block_table_width: tl.constexpr, + block_table_width, state_block_size, rms_norm_weight_ptr, rms_norm_eps, @@ -424,6 +444,7 @@ def deepseek_v4_fused_sparse_compress_cache_insert( compressor_slot_mapping.numel(), positions.numel(), kv_slot_mapping.numel(), + token_to_req_indices.numel(), ) if num_actual == 0: return @@ -465,7 +486,345 @@ def deepseek_v4_fused_sparse_compress_cache_insert( ) -@triton.jit +# ``block_table_width`` and ``num_source_tokens`` vary per request with cache +# depth (e.g. 704 vs 691), so their 16-divisibility flips across cached-prefill +# calls. Triton specializes non-constexpr ints on 16-alignment by default +# (``kwargs["align"]`` is ``True`` for positional ints), which turns each new +# alignment class into a fresh JIT compile that stalls the compute stream for +# seconds and shows up as a single-point cache-insert outlier. These scalars +# only drive control flow and masks, never a contiguous load base, so excluding +# them from alignment specialization removes the hidden compile dimension +# without losing any vectorization hint. +@triton.jit( + do_not_specialize_on_alignment=[ + "block_table_width", + "state_block_size", + "kv_cache_block_size", + "num_active_tokens", + "num_source_tokens", + ] +) +def _deepseek_v4_fused_hca_direct_compress_cache_kernel( + active_token_indices_ptr, + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + kv_ptr, + kv_stride, + score_ptr, + score_stride, + ape_ptr, + token_to_req_indices_ptr, + query_start_loc_ptr, + positions_ptr, + block_table_ptr, + block_table_base_offsets_ptr, + block_table_stride, + block_table_width, + state_block_size, + rms_norm_weight_ptr, + rms_norm_eps, + cos_sin_cache_ptr, + cos_sin_stride, + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + num_active_tokens, + num_source_tokens, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + SCALE_DIM: tl.constexpr, + KV_BLOCK_STRIDE: tl.constexpr, +): + row_idx = tl.program_id(0) + if row_idx >= num_active_tokens: + return + token_idx = tl.load(active_token_indices_ptr + row_idx) + if token_idx < 0 or token_idx >= num_source_tokens: + return + + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + kv_slot = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot < 0: + return + + tokens = tl.arange(0, COMPRESS_RATIO) + source_positions = position - COMPRESS_RATIO + 1 + tokens + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + req_start = tl.load(query_start_loc_ptr + req_idx) + req_end = tl.load(query_start_loc_ptr + req_idx + 1) + start_pos = tl.load(positions_ptr + req_start) + source_idx = req_start + source_positions - start_pos + tail_bounds = ( + (source_idx >= req_start) + & (source_idx < req_end) + & (source_idx >= 0) + & (source_idx < num_source_tokens) + ) + safe_source_idx = tl.minimum(tl.maximum(source_idx, 0), num_source_tokens - 1) + source_req = tl.load( + token_to_req_indices_ptr + safe_source_idx, + mask=tail_bounds, + other=-1, + ) + in_tail = tail_bounds & (source_req == req_idx) + + if block_table_base_offsets_ptr is not None: + base_logical_page = tl.load(block_table_base_offsets_ptr + req_idx) + else: + base_logical_page = tl.full((), 0, tl.int32) + table_idx = source_positions // state_block_size - base_logical_page + valid_paged = ( + (~in_tail) + & (source_positions >= 0) + & (table_idx >= 0) + & (table_idx < block_table_width) + ) + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + table_idx, + mask=valid_paged, + other=-1, + ).to(tl.int64) + valid_paged = valid_paged & (block_numbers >= 0) + pos_in_block = source_positions % state_block_size + safe_pos_in_block = tl.maximum(pos_in_block, 0) + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + safe_block_numbers = tl.maximum(block_numbers, 0) + state_row_base = ( + state_cache_ptr + + safe_block_numbers[:, None] * state_cache_stride0 + + safe_pos_in_block[:, None] * state_cache_stride1 + ) + paged_score = tl.load( + state_row_base + STATE_WIDTH + block[None, :], + mask=valid_paged[:, None] & mask[None, :], + other=float("-inf"), + ).to(tl.float32) + ape_row = tl.maximum(source_positions % COMPRESS_RATIO, 0) + ape = tl.load( + ape_ptr + ape_row[:, None] * HEAD_SIZE + block[None, :], + mask=in_tail[:, None] & mask[None, :], + other=0.0, + ).to(tl.float32) + tail_score = tl.load( + score_ptr + safe_source_idx[:, None] * score_stride + block[None, :], + mask=in_tail[:, None] & mask[None, :], + other=0.0, + ).to(tl.float32) + score = tl.where(in_tail[:, None], tail_score + ape, paged_score) + valid_source = in_tail | valid_paged + score = tl.where(valid_source[:, None], score, float("-inf")) + score = tl.softmax(score, dim=0) + + paged_kv = tl.load( + state_row_base + block[None, :], + mask=valid_paged[:, None] & mask[None, :], + other=0.0, + ).to(tl.float32) + tail_kv = tl.load( + kv_ptr + safe_source_idx[:, None] * kv_stride + block[None, :], + mask=in_tail[:, None] & mask[None, :], + other=0.0, + ).to(tl.float32) + kv = tl.where(in_tail[:, None], tail_kv, paged_kv) + kv = tl.where(valid_source[:, None], kv, 0.0) + compressed = tl.sum(kv * score, axis=0) + + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed * compressed, axis=0) / HEAD_SIZE + normed = compressed * tl.rsqrt(variance + rms_norm_eps) * rms_w + + kv_block = kv_slot // kv_cache_block_size + kv_pos = kv_slot % kv_cache_block_size + cache_block_ptr = k_cache_ptr + kv_block.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + kv_cache_block_size * TOKEN_STRIDE + kv_pos * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + quant_input = normed.to(tl.bfloat16).to(tl.float32) + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) + block_absmax = tl.max(tl.abs(quant_2d), axis=1) + block_absmax = tl.maximum(block_absmax, 1.0e-4) + exponents = tl.ceil(tl.log2(block_absmax * INV_FP8_MAX)) + inv_scales = tl.exp2(-exponents) + x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) + x_fp8 = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX).to(tl.float8e4nv) + x_uint8 = tl.reshape(x_fp8.to(tl.uint8, bitcast=True), (TRITON_BLOCK_SIZE,)) + + tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) + scale_idx = tl.arange(0, N_QUANT_BLOCKS) + encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + tl.store( + scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS + ) + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) + + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + pair_2d = tl.reshape(normed, (NUM_PAIRS, 2)) + even, odd = tl.split(pair_2d) + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair = pair_idx - NOPE_PAIRS + is_rope = rope_pair >= 0 + cs_idx = tl.maximum(rope_pair, 0) + + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cs_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cs_base + cs_idx, mask=is_rope, other=1.0) + sin_v = tl.load(cs_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0) + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + rotated = tl.interleave(new_even, new_odd) + + rope_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) + rope_local = block - NOPE_HEAD_DIM + tl.store( + rope_ptr + rope_local, + rotated.to(tl.bfloat16), + mask=(block >= NOPE_HEAD_DIM) & mask, + ) + + +def deepseek_v4_fused_hca_direct_compress_cache_insert( + *, + state_cache: torch.Tensor, + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + block_table: torch.Tensor, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + active_token_indices: torch.Tensor, + block_table_base_offsets: torch.Tensor | None = None, +) -> None: + """Compress HCA active rows from current KV/score plus paged prefix state. + + Args: + state_cache: Paged HCA compressor state for prefix tokens. + kv: Current forward HCA value states with shape ``[tokens, 512]``. + score: Current forward HCA score states with shape ``[tokens, 512]``. + ape: HCA APE table with shape ``[128, 512]``. + token_to_req_indices: Request id per current-step token. + query_start_loc: Prefix-sum token offsets for current-step requests. + positions: Absolute token positions for the current forward step. + block_table: Paged compressor-state block table. + compressor_block_size: Compressor-state page size. + rms_norm_weight: RMSNorm weight applied after weighted reduction. + rms_norm_eps: RMSNorm epsilon. + cos_sin_cache: RoPE cache indexed by compressed positions. + kv_cache_2d: Destination compressed KV cache byte storage. + kv_slot_mapping: Destination compressed KV slot per source token. + kv_cache_block_size: Destination compressed KV cache page size. + active_token_indices: Source token offsets for HCA compressed rows. + block_table_base_offsets: Optional logical-page offset per request. + """ + + num_source_tokens = min( + kv.shape[0], + score.shape[0], + positions.numel(), + kv_slot_mapping.numel(), + token_to_req_indices.numel(), + ) + if active_token_indices.dim() != 1: + raise ValueError("active_token_indices must be a 1D tensor") + if active_token_indices.dtype != torch.int64: + active_token_indices = active_token_indices.to(torch.int64) + if active_token_indices.device != positions.device: + active_token_indices = active_token_indices.to( + positions.device, + non_blocking=True, + ) + if not active_token_indices.is_contiguous(): + active_token_indices = active_token_indices.contiguous() + num_actual = active_token_indices.numel() + if num_actual == 0 or num_source_tokens == 0: + return + if ( + block_table_base_offsets is not None + and block_table_base_offsets.dtype != torch.int32 + ): + block_table_base_offsets = block_table_base_offsets.to(torch.int32) + launch_rows = _deepseek_v4_hca_direct_launch_rows(num_actual) + _deepseek_v4_fused_hca_direct_compress_cache_kernel[(launch_rows,)]( + active_token_indices[:num_actual], + state_cache, + state_cache.stride(0), + state_cache.stride(1), + kv[:num_source_tokens], + kv.stride(0), + score[:num_source_tokens], + score.stride(0), + ape, + token_to_req_indices[:num_source_tokens], + query_start_loc, + positions[:num_source_tokens], + block_table, + block_table_base_offsets, + block_table.stride(0), + block_table.shape[-1], + compressor_block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + cos_sin_cache.stride(0), + kv_cache_2d, + kv_slot_mapping[:num_source_tokens], + kv_cache_block_size, + num_actual, + num_source_tokens, + HEAD_SIZE=DEEPSEEK_V4_HEAD_DIM, + TRITON_BLOCK_SIZE=triton.next_power_of_2(DEEPSEEK_V4_HEAD_DIM), + STATE_WIDTH=state_cache.shape[-1] // 2, + COMPRESS_RATIO=128, + ROPE_HEAD_DIM=DEEPSEEK_V4_ROPE_DIM, + FP8_MAX=DEEPSEEK_V4_FP8_MAX, + QUANT_BLOCK=DEEPSEEK_V4_FP8_QUANT_BLOCK, + TOKEN_STRIDE=DEEPSEEK_V4_SWA_TOKEN_STRIDE, + SCALE_DIM=DEEPSEEK_V4_SWA_SCALE_DIM, + KV_BLOCK_STRIDE=kv_cache_2d.stride(0), + num_warps=4, + ) + + +# Same fix as the sparse/HCA compress kernels: ``block_table_width`` is a +# per-request runtime bound on ``table_idx``, not a compile-time constant, so +# keeping it ``tl.constexpr`` compiled a distinct CSA variant per width and +# charged that JIT cost to the CSA cache-insert scope. +@triton.jit( + do_not_specialize_on_alignment=[ + "block_table_width", + "state_block_size", + "kv_cache_block_size", + ] +) def _deepseek_v4_fused_csa_indexer_mxfp4_cache_kernel( state_cache_ptr, state_cache_stride0, @@ -476,7 +835,7 @@ def _deepseek_v4_fused_csa_indexer_mxfp4_cache_kernel( block_table_ptr, block_table_base_offsets_ptr, block_table_stride, - block_table_width: tl.constexpr, + block_table_width, state_block_size, rms_norm_weight_ptr, rms_norm_eps, diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/kvcache/cuda.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/kvcache/cuda.py index 4395f7a62..c117f5c39 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/kvcache/cuda.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/kvcache/cuda.py @@ -20,37 +20,48 @@ """CUDA KV cache transfer kernels.""" +from typing import Any + from tokenspeed_kernel.registry import error_fn try: from tokenspeed_kernel.thirdparty.cuda.kvcacheio import ( + DirectH2DScatterPlan, + prepare_kv_direct_h2d_scatter_plan, transfer_kv_all_layer_lf_pf, transfer_kv_all_layer_lf_ph, transfer_kv_all_layer_mla, transfer_kv_all_layer_mla_lf_pf, transfer_kv_direct, + transfer_kv_direct_h2d_scatter_prepared, transfer_kv_per_layer_mla, transfer_kv_per_layer_mla_pf_lf, transfer_kv_per_layer_pf_lf, transfer_kv_per_layer_ph_lf, ) except ImportError: + DirectH2DScatterPlan = Any + prepare_kv_direct_h2d_scatter_plan = error_fn transfer_kv_all_layer_lf_pf = error_fn transfer_kv_all_layer_lf_ph = error_fn transfer_kv_all_layer_mla = error_fn transfer_kv_all_layer_mla_lf_pf = error_fn transfer_kv_direct = error_fn + transfer_kv_direct_h2d_scatter_prepared = error_fn transfer_kv_per_layer_mla = error_fn transfer_kv_per_layer_mla_pf_lf = error_fn transfer_kv_per_layer_pf_lf = error_fn transfer_kv_per_layer_ph_lf = error_fn __all__ = [ + "DirectH2DScatterPlan", + "prepare_kv_direct_h2d_scatter_plan", "transfer_kv_all_layer_lf_pf", "transfer_kv_all_layer_lf_ph", "transfer_kv_all_layer_mla", "transfer_kv_all_layer_mla_lf_pf", "transfer_kv_direct", + "transfer_kv_direct_h2d_scatter_prepared", "transfer_kv_per_layer_mla", "transfer_kv_per_layer_mla_pf_lf", "transfer_kv_per_layer_pf_lf", diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/flashinfer_kvcacheio_binding.cu b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/flashinfer_kvcacheio_binding.cu index 38eb09583..523880416 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/flashinfer_kvcacheio_binding.cu +++ b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/flashinfer_kvcacheio_binding.cu @@ -140,6 +140,22 @@ void transfer_kv_all_layer_mla_lf_pf( int64_t block_quota, int64_t num_warps_per_block); +void transfer_kv_direct( + TensorView src_layer_ptrs, + TensorView dst_layer_ptrs, + TensorView src_indices, + TensorView dst_indices, + int64_t item_size, + int64_t page_size); + +void transfer_kv_direct_ptr_table_scatter_h2d( + TensorView src_layer_ptrs, + TensorView dst_layer_ptrs, + TensorView src_indices, + TensorView dst_indices, + int64_t item_size, + int64_t page_size); + TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_per_layer, transfer_kv_per_layer); TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_per_layer_pf_lf, transfer_kv_per_layer_pf_lf); TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_per_layer_ph_lf, transfer_kv_per_layer_ph_lf); @@ -150,3 +166,5 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_per_layer_mla, transfer_kv_per_layer_m TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_per_layer_mla_pf_lf, transfer_kv_per_layer_mla_pf_lf); TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_all_layer_mla, transfer_kv_all_layer_mla); TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_all_layer_mla_lf_pf, transfer_kv_all_layer_mla_lf_pf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_direct, transfer_kv_direct); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(transfer_kv_direct_ptr_table_scatter_h2d, transfer_kv_direct_ptr_table_scatter_h2d); diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/kvcacheio_transfer.cu b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/kvcacheio_transfer.cu index 76ef50bdf..4416de62b 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/kvcacheio_transfer.cu +++ b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/csrc/kvcacheio_transfer.cu @@ -26,6 +26,7 @@ #include #include +#include #include #include "tvm_ffi_utils.h" @@ -80,6 +81,14 @@ inline std::vector indices_to_host(TensorView indices) { return host; } +inline void check_cpu_uint64_vector(TensorView values, const char* name) { + TVM_FFI_ICHECK_EQ(values.ndim(), 1) << name << " must be a 1D tensor"; + TVM_FFI_ICHECK_EQ(values.dtype(), dl_uint64) + << name << " must be a uint64 tensor"; + TVM_FFI_ICHECK_EQ(values.device().device_type, kDLCPU) + << name << " must be a CPU tensor"; +} + inline void copy_async_bytes(void* dst, const void* src, size_t num_bytes, cudaStream_t stream) { if (num_bytes == 0) { return; @@ -87,25 +96,55 @@ inline void copy_async_bytes(void* dst, const void* src, size_t num_bytes, cudaS check_cuda(cudaMemcpyAsync(dst, src, num_bytes, cudaMemcpyDefault, stream), "cudaMemcpyAsync failed"); } -inline void copy_token_span( - TensorView src, - TensorView dst, - int64_t src_index, - int64_t dst_index, - int64_t num_tokens, - cudaStream_t stream) { - TVM_FFI_ICHECK_GE(src.dim(), 1); - TVM_FFI_ICHECK_GE(dst.dim(), 1); - - const int64_t src_stride_bytes = src.stride(0) * get_element_size(src); - const int64_t dst_stride_bytes = dst.stride(0) * get_element_size(dst); - TVM_FFI_ICHECK_EQ(src_stride_bytes, dst_stride_bytes) - << "Source and destination token stride bytes must match for direct copy"; - - const size_t copy_bytes = static_cast(num_tokens * src_stride_bytes); - const char* src_ptr = static_cast(src.data_ptr()) + src_index * src_stride_bytes; - char* dst_ptr = static_cast(dst.data_ptr()) + dst_index * dst_stride_bytes; - copy_async_bytes(dst_ptr, src_ptr, copy_bytes, stream); +inline void check_cuda_uint64_vector(TensorView values, const char* name) { + CHECK_CUDA(values); + CHECK_DIM(1, values); + CHECK_INPUT_TYPE(values, dl_uint64); +} + +__global__ void scatter_h2d_pages_kernel( + const uint64_t* __restrict__ src_layer_ptrs, + const uint64_t* __restrict__ dst_layer_ptrs, + const int64_t* __restrict__ src_indices, + const int64_t* __restrict__ dst_indices, + int64_t num_layers, + int64_t num_pages, + int64_t item_size, + int64_t page_size, + int64_t copy_bytes) { + const int64_t task_id = static_cast(blockIdx.x); + if (task_id >= num_layers * num_pages) { + return; + } + + const int64_t page_id = task_id % num_pages; + const int64_t layer_id = task_id / num_pages; + const int64_t index_pos = page_id * page_size; + const int64_t src_index = src_indices[index_pos]; + const int64_t dst_index = dst_indices[index_pos]; + + const uintptr_t src_addr = static_cast(src_layer_ptrs[layer_id]) + + static_cast(src_index * item_size); + const uintptr_t dst_addr = static_cast(dst_layer_ptrs[layer_id]) + + static_cast(dst_index * item_size); + const char* src = reinterpret_cast(src_addr); + char* dst = reinterpret_cast(dst_addr); + + const bool aligned16 = + ((src_addr | dst_addr | static_cast(copy_bytes)) & 0xF) == 0; + if (aligned16) { + const int64_t vec_count = copy_bytes / static_cast(sizeof(uint4)); + const uint4* src_vec = reinterpret_cast(src); + uint4* dst_vec = reinterpret_cast(dst); + for (int64_t i = threadIdx.x; i < vec_count; i += blockDim.x) { + dst_vec[i] = src_vec[i]; + } + return; + } + + for (int64_t i = threadIdx.x; i < copy_bytes; i += blockDim.x) { + dst[i] = src[i]; + } } template @@ -756,23 +795,29 @@ void transfer_kv_all_layer_mla_lf_pf( } void transfer_kv_direct( - const std::vector& src_layers, - std::vector dst_layers, + TensorView src_layer_ptrs, + TensorView dst_layer_ptrs, TensorView src_indices, TensorView dst_indices, + int64_t item_size, int64_t page_size) { - TVM_FFI_ICHECK_EQ(src_layers.size(), dst_layers.size()) - << "Source and destination layers must have the same number of layers"; + check_cpu_uint64_vector(src_layer_ptrs, "src_layer_ptrs"); + check_cpu_uint64_vector(dst_layer_ptrs, "dst_layer_ptrs"); + TVM_FFI_ICHECK_EQ(src_layer_ptrs.numel(), dst_layer_ptrs.numel()) + << "Source and destination pointer tables must have the same length"; TVM_FFI_ICHECK_EQ(src_indices.numel(), dst_indices.numel()) << "Source and destination indices must have the same length"; + TVM_FFI_ICHECK_GT(item_size, 0) << "Item size must be positive"; TVM_FFI_ICHECK_GT(page_size, 0) << "Page size must be positive"; TVM_FFI_ICHECK_EQ(src_indices.numel() % page_size, 0) << "Source indices size must be divisible by page size"; + const auto* src_ptrs = static_cast(src_layer_ptrs.data_ptr()); + const auto* dst_ptrs = static_cast(dst_layer_ptrs.data_ptr()); const auto src_indices_host = indices_to_host(src_indices); const auto dst_indices_host = indices_to_host(dst_indices); const int64_t num_indices = static_cast(src_indices_host.size()); - const int64_t num_layers = static_cast(src_layers.size()); + const int64_t num_layers = static_cast(src_layer_ptrs.numel()); const cudaStream_t stream = get_current_stream(); int64_t start_index = 0; @@ -791,11 +836,64 @@ void transfer_kv_direct( const int64_t src_index = src_indices_host[start_index]; const int64_t dst_index = dst_indices_host[start_index]; - const int64_t num_tokens = end_index - start_index; + const int64_t num_items = end_index - start_index; + const size_t copy_bytes = static_cast(num_items * item_size); for (int64_t j = 0; j < num_layers; ++j) { - copy_token_span(src_layers[j], dst_layers[j], src_index, dst_index, num_tokens, stream); + const char* src_ptr = + reinterpret_cast(static_cast(src_ptrs[j])) + + src_index * item_size; + char* dst_ptr = + reinterpret_cast(static_cast(dst_ptrs[j])) + + dst_index * item_size; + copy_async_bytes(dst_ptr, src_ptr, copy_bytes, stream); } start_index = end_index; } } + +void transfer_kv_direct_ptr_table_scatter_h2d( + TensorView src_layer_ptrs, + TensorView dst_layer_ptrs, + TensorView src_indices, + TensorView dst_indices, + int64_t item_size, + int64_t page_size) { + check_cuda_uint64_vector(src_layer_ptrs, "src_layer_ptrs"); + check_cuda_uint64_vector(dst_layer_ptrs, "dst_layer_ptrs"); + check_indices(src_indices, dst_indices); + TVM_FFI_ICHECK_EQ(src_layer_ptrs.numel(), dst_layer_ptrs.numel()) + << "Source and destination pointer tables must have the same length"; + TVM_FFI_ICHECK_GT(item_size, 0) << "Item size must be positive"; + TVM_FFI_ICHECK_EQ(page_size, 1) + << "Scatter H2D only supports one cache row per page"; + TVM_FFI_ICHECK_EQ(src_indices.numel() % page_size, 0) + << "Source indices size must be divisible by page size"; + + const int64_t num_layers = static_cast(src_layer_ptrs.numel()); + const int64_t num_pages = + static_cast(src_indices.numel()) / page_size; + if (num_layers == 0 || num_pages == 0) { + return; + } + const int64_t num_tasks = num_layers * num_pages; + TVM_FFI_ICHECK_LE(num_tasks, + static_cast(std::numeric_limits::max())) + << "Too many scatter H2D copy tasks for one launch"; + + const int threads = 256; + const int64_t copy_bytes = item_size * page_size; + const cudaStream_t stream = get_current_stream(); + scatter_h2d_pages_kernel<<(num_tasks), threads, 0, stream>>>( + static_cast(src_layer_ptrs.data_ptr()), + static_cast(dst_layer_ptrs.data_ptr()), + static_cast(src_indices.data_ptr()), + static_cast(dst_indices.data_ptr()), + num_layers, + num_pages, + item_size, + page_size, + copy_bytes); + check_cuda(cudaGetLastError(), + "transfer_kv_direct_ptr_table_scatter_h2d launch failed"); +} diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/kvcacheio.py b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/kvcacheio.py index 2d1d58c76..ca403a1ad 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/kvcacheio.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/kvcacheio.py @@ -19,6 +19,8 @@ # SOFTWARE. import functools +from bisect import bisect_left +from dataclasses import dataclass from pathlib import Path from typing import List @@ -43,6 +45,60 @@ def _load_kvcacheio_module(): return tvm_ffi.load_module(str(so_path)) +@functools.cache +def _load_transfer_kv_direct_func(): + try: + return getattr(_load_kvcacheio_module(), "transfer_kv_direct") + except (AttributeError, ModuleNotFoundError): + return None + except RuntimeError as exc: + if "kvcacheio library not found" in str(exc): + return None + raise + + +@functools.cache +def _load_transfer_kv_direct_scatter_h2d_func(): + try: + return getattr( + _load_kvcacheio_module(), + "transfer_kv_direct_ptr_table_scatter_h2d", + ) + except (AttributeError, ModuleNotFoundError): + return None + except RuntimeError as exc: + if "kvcacheio library not found" in str(exc): + return None + raise + + +@dataclass(frozen=True) +class _DirectH2DScatterResult: + used: bool + buckets: int = 0 + kernel_launches: int = 0 + fallback_reason: str = "" + + +@dataclass(frozen=True) +class _DirectH2DScatterBucket: + item_size: int + entry_ids: tuple[int, ...] + src_ptrs_host: torch.Tensor + dst_ptrs_host: torch.Tensor + src_ptrs_device: torch.Tensor + dst_ptrs_device: torch.Tensor + + +@dataclass(frozen=True) +class DirectH2DScatterPlan: + """Immutable H2D pointer tables prepared for one device stream.""" + + device: torch.device + stream_id: int + buckets: tuple[_DirectH2DScatterBucket, ...] + + _is_amd = current_platform().is_amd @@ -77,6 +133,278 @@ def _transfer_page_direct( ) +def _has_cuda_layer( + src_layers: List[torch.Tensor], + dst_layers: List[torch.Tensor], +) -> bool: + return any(layer.device.type == "cuda" for layer in src_layers) or any( + layer.device.type == "cuda" for layer in dst_layers + ) + + +def _leading_stride_bytes(tensor: torch.Tensor) -> int: + if tensor.dim() < 1: + raise ValueError("Direct KV transfer buffers must have at least one dimension") + return int(tensor.stride(0) * tensor.element_size()) + + +def _h2d_scatter_device( + src_layers: List[torch.Tensor], + dst_layers: List[torch.Tensor], +) -> tuple[torch.device | None, str]: + if _is_amd: + return None, "amd_fallback" + if not src_layers: + return None, "empty_layers" + if any(layer.device.type != "cpu" for layer in src_layers): + return None, "src_not_cpu" + dst_device = dst_layers[0].device + if dst_device.type != "cuda": + return None, "dst_not_cuda" + if any(layer.device != dst_device for layer in dst_layers): + return None, "mixed_dst_devices" + return dst_device, "" + + +def _to_i64_contiguous_on_device( + indices: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + if indices.dtype != torch.int64: + indices = indices.to(torch.int64) + if indices.device != device: + indices = indices.to(device, non_blocking=True) + if not indices.is_contiguous(): + indices = indices.contiguous() + return indices + + +def _record_cuda_metadata_stream(*tensors: torch.Tensor) -> None: + cuda_tensors = [tensor for tensor in tensors if tensor.device.type == "cuda"] + if not cuda_tensors: + return + stream = torch.cuda.current_stream(cuda_tensors[0].device) + for tensor in cuda_tensors: + tensor.record_stream(stream) + + +def _current_stream_id(device: torch.device) -> int: + if device.type != "cuda": + return 0 + return int(torch.cuda.current_stream(device).cuda_stream) + + +def _pinned_pointer_tensor(values: list[int], device: torch.device) -> torch.Tensor: + return torch.tensor( + values, + dtype=torch.uint64, + device="cpu", + pin_memory=device.type == "cuda" and torch.cuda.is_available(), + ) + + +def prepare_kv_direct_h2d_scatter_plan( + src_layers: List[torch.Tensor], + dst_layers: List[torch.Tensor], + entry_ids: List[int], +) -> tuple[DirectH2DScatterPlan | None, str]: + """Prepare full-group pointer tables for repeated H2D scatter launches.""" + + if len(src_layers) != len(dst_layers) or len(src_layers) != len(entry_ids): + raise ValueError( + "Source layers, destination layers, and entry ids must have equal lengths" + ) + scatter_func = _load_transfer_kv_direct_scatter_h2d_func() + if scatter_func is None: + return None, "symbol_missing" + + device, reason = _h2d_scatter_device(src_layers, dst_layers) + if device is None: + return None, reason + + grouped: dict[int, list[tuple[int, torch.Tensor, torch.Tensor]]] = {} + for entry_id, src_layer, dst_layer in zip(entry_ids, src_layers, dst_layers): + src_item_size = _leading_stride_bytes(src_layer) + dst_item_size = _leading_stride_bytes(dst_layer) + if src_item_size != dst_item_size: + return None, "item_size_mismatch" + grouped.setdefault(src_item_size, []).append( + (int(entry_id), src_layer, dst_layer) + ) + + platform = current_platform() + buckets: list[_DirectH2DScatterBucket] = [] + for item_size, entries in grouped.items(): + entries.sort(key=lambda entry: entry[0]) + bucket_entry_ids = tuple(entry_id for entry_id, _, _ in entries) + src_ptrs_host = _pinned_pointer_tensor( + [platform.device_visible_data_ptr(src) for _, src, _ in entries], + device, + ) + dst_ptrs_host = _pinned_pointer_tensor( + [platform.device_visible_data_ptr(dst) for _, _, dst in entries], + device, + ) + buckets.append( + _DirectH2DScatterBucket( + item_size=item_size, + entry_ids=bucket_entry_ids, + src_ptrs_host=src_ptrs_host, + dst_ptrs_host=dst_ptrs_host, + src_ptrs_device=src_ptrs_host.to(device, non_blocking=True), + dst_ptrs_device=dst_ptrs_host.to(device, non_blocking=True), + ) + ) + + return ( + DirectH2DScatterPlan( + device=device, + stream_id=_current_stream_id(device), + buckets=tuple(buckets), + ), + "", + ) + + +def transfer_kv_direct_h2d_scatter_prepared( + plan: DirectH2DScatterPlan, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + page_size: int, + entry_begin: int, + entry_end: int, +) -> _DirectH2DScatterResult: + """Launch a layer-range H2D scatter from a prepared full-group plan.""" + + _check_direct_copy_args(src_indices, dst_indices, page_size) + if page_size != 1: + return _DirectH2DScatterResult( + used=False, + fallback_reason="unsupported_page_size", + ) + if _current_stream_id(plan.device) != plan.stream_id: + return _DirectH2DScatterResult( + used=False, + fallback_reason="stream_mismatch", + ) + + scatter_func = _load_transfer_kv_direct_scatter_h2d_func() + if scatter_func is None: + return _DirectH2DScatterResult( + used=False, + fallback_reason="symbol_missing", + ) + + src_indices_device = _to_i64_contiguous_on_device(src_indices, plan.device) + dst_indices_device = _to_i64_contiguous_on_device(dst_indices, plan.device) + kernel_launches = 0 + for bucket in plan.buckets: + first = bisect_left(bucket.entry_ids, entry_begin) + last = bisect_left(bucket.entry_ids, entry_end) + if first == last: + continue + src_ptrs = bucket.src_ptrs_device[first:last] + dst_ptrs = bucket.dst_ptrs_device[first:last] + scatter_func( + src_ptrs, + dst_ptrs, + src_indices_device, + dst_indices_device, + bucket.item_size, + page_size, + ) + _record_cuda_metadata_stream( + src_ptrs, + dst_ptrs, + src_indices_device, + dst_indices_device, + ) + kernel_launches += 1 + + return _DirectH2DScatterResult( + used=True, + buckets=kernel_launches, + kernel_launches=kernel_launches, + ) + + +def _transfer_kv_direct_cpp( + direct_func, + src_layers: List[torch.Tensor], + dst_layers: List[torch.Tensor], + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + page_size: int, +) -> bool: + buckets: dict[int, tuple[list[torch.Tensor], list[torch.Tensor]]] = {} + for src_layer, dst_layer in zip(src_layers, dst_layers): + src_item_size = _leading_stride_bytes(src_layer) + dst_item_size = _leading_stride_bytes(dst_layer) + if src_item_size != dst_item_size: + return False + src_bucket, dst_bucket = buckets.setdefault(src_item_size, ([], [])) + src_bucket.append(src_layer) + dst_bucket.append(dst_layer) + + platform = current_platform() + for item_size, (bucket_src_layers, bucket_dst_layers) in buckets.items(): + src_ptrs = torch.tensor( + [platform.device_visible_data_ptr(layer) for layer in bucket_src_layers], + dtype=torch.uint64, + device="cpu", + ) + dst_ptrs = torch.tensor( + [platform.device_visible_data_ptr(layer) for layer in bucket_dst_layers], + dtype=torch.uint64, + device="cpu", + ) + direct_func( + src_ptrs, + dst_ptrs, + src_indices, + dst_indices, + item_size, + page_size, + ) + return True + + +def _transfer_kv_direct_python( + src_layers: List[torch.Tensor], + dst_layers: List[torch.Tensor], + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + page_size: int, +): + src_indices_host = _indices_to_host_list(src_indices) + dst_indices_host = _indices_to_host_list(dst_indices) + + start_index = 0 + end_index = 0 + num_indices = len(src_indices_host) + + for i in range(num_indices): + if i < num_indices - 1: + src_diff = src_indices_host[i + 1] - src_indices_host[i] + dst_diff = dst_indices_host[i + 1] - dst_indices_host[i] + if src_diff == 1 and dst_diff == 1: + continue + end_index = i + 1 + else: + end_index = num_indices + + src_index = src_indices_host[start_index] + dst_index = dst_indices_host[start_index] + num_tokens = end_index - start_index + + for src_layer, dst_layer in zip(src_layers, dst_layers): + _transfer_page_direct( + src_layer, dst_layer, src_index, dst_index, num_tokens + ) + + start_index = end_index + + def transfer_kv_per_layer( src_k: torch.Tensor, dst_k: torch.Tensor, @@ -260,33 +588,32 @@ def transfer_kv_direct( ) _check_direct_copy_args(src_indices, dst_indices, page_size) - src_indices_host = _indices_to_host_list(src_indices) - dst_indices_host = _indices_to_host_list(dst_indices) - - start_index = 0 - end_index = 0 - num_indices = len(src_indices_host) - - for i in range(num_indices): - if i < num_indices - 1: - src_diff = src_indices_host[i + 1] - src_indices_host[i] - dst_diff = dst_indices_host[i + 1] - dst_indices_host[i] - if src_diff == 1 and dst_diff == 1: - continue - end_index = i + 1 - else: - end_index = num_indices - - src_index = src_indices_host[start_index] - dst_index = dst_indices_host[start_index] - num_tokens = end_index - start_index - - for src_layer, dst_layer in zip(src_layers, dst_layers): - _transfer_page_direct( - src_layer, dst_layer, src_index, dst_index, num_tokens - ) - - start_index = end_index + direct_func = ( + _load_transfer_kv_direct_func() + if _has_cuda_layer( + src_layers, + dst_layers, + ) + else None + ) + if direct_func is not None: + if _transfer_kv_direct_cpp( + direct_func, + src_layers, + dst_layers, + src_indices, + dst_indices, + page_size, + ): + return + + _transfer_kv_direct_python( + src_layers, + dst_layers, + src_indices, + dst_indices, + page_size, + ) def transfer_kv_per_layer_mla( @@ -380,11 +707,14 @@ def transfer_kv_all_layer_mla_lf_pf( __all__ = [ + "DirectH2DScatterPlan", + "prepare_kv_direct_h2d_scatter_plan", "transfer_kv_all_layer_lf_pf", "transfer_kv_all_layer_lf_ph", "transfer_kv_all_layer_mla", "transfer_kv_all_layer_mla_lf_pf", "transfer_kv_direct", + "transfer_kv_direct_h2d_scatter_prepared", "transfer_kv_per_layer_mla", "transfer_kv_per_layer_mla_pf_lf", "transfer_kv_per_layer_pf_lf", diff --git a/tokenspeed-kernel/test/test_kvcacheio_direct.py b/tokenspeed-kernel/test/test_kvcacheio_direct.py new file mode 100644 index 000000000..ab9dcb9c2 --- /dev/null +++ b/tokenspeed-kernel/test/test_kvcacheio_direct.py @@ -0,0 +1,232 @@ +import pytest +import torch + + +class _FakePointerPlatform: + def device_visible_data_ptr(self, tensor): + return tensor.data_ptr() + + +def test_transfer_kv_direct_prefers_cpp_binding(monkeypatch): + from tokenspeed_kernel.thirdparty.cuda import kvcacheio + + calls = [] + + def fake_direct( + src_layer_ptrs, + dst_layer_ptrs, + src_indices, + dst_indices, + item_size, + page_size, + ): + calls.append( + ( + src_layer_ptrs, + dst_layer_ptrs, + src_indices, + dst_indices, + item_size, + page_size, + ) + ) + + monkeypatch.setattr(kvcacheio, "_load_transfer_kv_direct_func", lambda: fake_direct) + monkeypatch.setattr(kvcacheio, "_has_cuda_layer", lambda *_args: True) + monkeypatch.setattr(kvcacheio, "current_platform", lambda: _FakePointerPlatform()) + + src = torch.arange(4, dtype=torch.float32).reshape(2, 2) + dst = torch.zeros_like(src) + indices = torch.tensor([0, 1], dtype=torch.int64) + + kvcacheio.transfer_kv_direct([src], [dst], indices, indices, page_size=1) + + assert len(calls) == 1 + ( + called_src_ptrs, + called_dst_ptrs, + called_src_indices, + called_dst_indices, + item_size, + page_size, + ) = calls[0] + assert called_src_ptrs.tolist() == [src.data_ptr()] + assert called_dst_ptrs.tolist() == [dst.data_ptr()] + assert called_src_indices is indices + assert called_dst_indices is indices + assert item_size == src.stride(0) * src.element_size() + assert page_size == 1 + assert torch.equal(dst, torch.zeros_like(dst)) + + +def test_transfer_kv_direct_python_fallback(monkeypatch): + from tokenspeed_kernel.thirdparty.cuda import kvcacheio + + monkeypatch.setattr(kvcacheio, "_load_transfer_kv_direct_func", lambda: None) + + src = torch.arange(6, dtype=torch.float32).reshape(3, 2) + dst = torch.zeros_like(src) + src_indices = torch.tensor([0, 2], dtype=torch.int64) + dst_indices = torch.tensor([1, 2], dtype=torch.int64) + + kvcacheio.transfer_kv_direct([src], [dst], src_indices, dst_indices, page_size=1) + + expected = torch.zeros_like(src) + expected[1] = src[0] + expected[2] = src[2] + assert torch.equal(dst, expected) + + +def test_prepared_h2d_scatter_reuses_bucketed_pointer_tables(monkeypatch): + from tokenspeed_kernel.thirdparty.cuda import kvcacheio + + calls = [] + pointer_lookups = [] + + def fake_scatter( + src_layer_ptrs, + dst_layer_ptrs, + src_indices, + dst_indices, + item_size, + page_size, + ): + calls.append( + ( + src_layer_ptrs, + dst_layer_ptrs, + src_indices, + dst_indices, + item_size, + page_size, + ) + ) + + monkeypatch.setattr( + kvcacheio, + "_load_transfer_kv_direct_scatter_h2d_func", + lambda: fake_scatter, + ) + monkeypatch.setattr( + kvcacheio, + "_h2d_scatter_device", + lambda *_args: (torch.device("cpu"), ""), + ) + platform = _FakePointerPlatform() + raw_data_ptr = platform.device_visible_data_ptr + + def data_ptr(tensor): + pointer_lookups.append(tensor) + return raw_data_ptr(tensor) + + platform.device_visible_data_ptr = data_ptr + monkeypatch.setattr(kvcacheio, "current_platform", lambda: platform) + + src_a0 = torch.arange(8, dtype=torch.float32).reshape(4, 2) + dst_a0 = torch.zeros_like(src_a0) + src_b0 = torch.arange(12, dtype=torch.float32).reshape(4, 3) + dst_b0 = torch.zeros_like(src_b0) + src_a1 = src_a0.clone() + dst_a1 = torch.zeros_like(src_a1) + src_b1 = src_b0.clone() + dst_b1 = torch.zeros_like(src_b1) + src_indices = torch.tensor([0, 2], dtype=torch.int64) + dst_indices = torch.tensor([1, 3], dtype=torch.int64) + + plan, reason = kvcacheio.prepare_kv_direct_h2d_scatter_plan( + [src_a0, src_b0, src_a1, src_b1], + [dst_a0, dst_b0, dst_a1, dst_b1], + [0, 0, 1, 1], + ) + assert reason == "" + assert plan is not None + assert len(pointer_lookups) == 8 + + first = kvcacheio.transfer_kv_direct_h2d_scatter_prepared( + plan, + src_indices, + dst_indices, + page_size=1, + entry_begin=0, + entry_end=1, + ) + second = kvcacheio.transfer_kv_direct_h2d_scatter_prepared( + plan, + src_indices, + dst_indices, + page_size=1, + entry_begin=1, + entry_end=2, + ) + + assert first.used and second.used + assert first.kernel_launches == second.kernel_launches == 2 + assert len(pointer_lookups) == 8 + assert [call[4] for call in calls] == [ + src_a0.stride(0) * src_a0.element_size(), + src_b0.stride(0) * src_b0.element_size(), + src_a1.stride(0) * src_a1.element_size(), + src_b1.stride(0) * src_b1.element_size(), + ] + assert calls[0][0].tolist() == [src_a0.data_ptr()] + assert calls[0][1].tolist() == [dst_a0.data_ptr()] + assert calls[1][0].tolist() == [src_b0.data_ptr()] + assert calls[1][1].tolist() == [dst_b0.data_ptr()] + assert calls[2][0].tolist() == [src_a1.data_ptr()] + assert calls[3][0].tolist() == [src_b1.data_ptr()] + assert calls[0][2] is src_indices + assert calls[0][3] is dst_indices + assert calls[0][5] == 1 + + monkeypatch.setattr(kvcacheio, "_current_stream_id", lambda _device: 1) + mismatch = kvcacheio.transfer_kv_direct_h2d_scatter_prepared( + plan, + src_indices, + dst_indices, + page_size=1, + entry_begin=0, + entry_end=1, + ) + assert not mismatch.used + assert mismatch.fallback_reason == "stream_mismatch" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_prepared_h2d_scatter_cuda_smoke(): + from tokenspeed_kernel.platform import current_platform + from tokenspeed_kernel.thirdparty.cuda import kvcacheio + + if kvcacheio._load_transfer_kv_direct_scatter_h2d_func() is None: + pytest.skip("scatter H2D kvcacheio symbol is not built") + + src = torch.arange(24, dtype=torch.float32).reshape(6, 4) + current_platform().register_host_tensor_for_gpu_access(src) + dst = torch.zeros_like(src, device="cuda") + src_indices = torch.tensor([0, 2, 4], dtype=torch.int64) + dst_indices = torch.tensor([1, 3, 5], dtype=torch.int64) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + plan, reason = kvcacheio.prepare_kv_direct_h2d_scatter_plan( + [src], + [dst], + [0], + ) + assert reason == "" + assert plan is not None + result = kvcacheio.transfer_kv_direct_h2d_scatter_prepared( + plan, + src_indices, + dst_indices, + page_size=1, + entry_begin=0, + entry_end=1, + ) + + assert result.used + stream.synchronize() + expected = torch.zeros_like(src) + expected[1] = src[0] + expected[3] = src[2] + expected[5] = src[4] + assert torch.equal(dst.cpu(), expected) diff --git a/tokenspeed-scheduler/CMakeLists.txt b/tokenspeed-scheduler/CMakeLists.txt index 90d3fa6b7..45ebfddb7 100644 --- a/tokenspeed-scheduler/CMakeLists.txt +++ b/tokenspeed-scheduler/CMakeLists.txt @@ -150,6 +150,7 @@ if(TOKENSPEED_SCHEDULER_BUILD_TESTS) tests/cpp/test_paged_cache_attach_loop.cpp tests/cpp/test_paged_cache_eviction.cpp tests/cpp/test_paged_cache_family_split.cpp + tests/cpp/test_paged_cache_l2_offload.cpp tests/cpp/test_paged_cache_prefix_hit_commit.cpp tests/cpp/test_paged_cache_replay.cpp tests/cpp/test_retract_abort_pages.cpp diff --git a/tokenspeed-scheduler/bindings/python_module.cpp b/tokenspeed-scheduler/bindings/python_module.cpp index c97d25513..7d10c002d 100644 --- a/tokenspeed-scheduler/bindings/python_module.cpp +++ b/tokenspeed-scheduler/bindings/python_module.cpp @@ -243,6 +243,7 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { "num_host_pages", [](const tokenspeed::SchedulerConfig& c) { return c.host_allocator.total_pages; }, [](tokenspeed::SchedulerConfig& c, std::int32_t v) { c.host_allocator.total_pages = v; }) .def_rw("paged_cache_groups", &tokenspeed::SchedulerConfig::paged_cache_groups) + .def_rw("paged_cache_host_group_pages", &tokenspeed::SchedulerConfig::paged_cache_host_group_pages) .def_rw("prefix_cache_adjunct", &tokenspeed::SchedulerConfig::prefix_cache_adjunct) .def_rw("disable_l2_cache", &tokenspeed::SchedulerConfig::disable_l2_cache) .def_rw("enable_l3_storage", &tokenspeed::SchedulerConfig::enable_l3_storage) @@ -378,6 +379,11 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { .value("KV", tokenspeed::CacheKind::kKV) .value("MAMBA", tokenspeed::CacheKind::kMamba); + nb::class_(cache, "PagedCacheTransferPair") + .def_ro("group_id", &tokenspeed::PagedCacheTransferPair::group_id) + .def_ro("src_pages", &tokenspeed::PagedCacheTransferPair::src_pages) + .def_ro("dst_pages", &tokenspeed::PagedCacheTransferPair::dst_pages); + auto prefetch_op = nb::class_(cache, "PrefetchOp"); BindCacheCommonFields(prefetch_op); prefetch_op.def(nb::init<>()) @@ -393,7 +399,8 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { .def_ro("src_pages", &tokenspeed::FlatLoadBackOperation::src_pages) .def_ro("dst_pages", &tokenspeed::FlatLoadBackOperation::dst_pages) .def_ro("src_pages_by_kind", &tokenspeed::FlatLoadBackOperation::src_pages_by_kind) - .def_ro("dst_pages_by_kind", &tokenspeed::FlatLoadBackOperation::dst_pages_by_kind); + .def_ro("dst_pages_by_kind", &tokenspeed::FlatLoadBackOperation::dst_pages_by_kind) + .def_ro("paged_cache_transfers", &tokenspeed::FlatLoadBackOperation::paged_cache_transfers); nb::class_(cache, "WriteBackOp") .def_ro("op_ids", &tokenspeed::FlatWriteBackOperation::op_ids) @@ -401,6 +408,7 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { .def_ro("dst_pages", &tokenspeed::FlatWriteBackOperation::dst_pages) .def_ro("src_pages_by_kind", &tokenspeed::FlatWriteBackOperation::src_pages_by_kind) .def_ro("dst_pages_by_kind", &tokenspeed::FlatWriteBackOperation::dst_pages_by_kind) + .def_ro("paged_cache_transfers", &tokenspeed::FlatWriteBackOperation::paged_cache_transfers) .def_ro("is_retract", &tokenspeed::FlatWriteBackOperation::is_retract); auto collect_forward = [](const tokenspeed::ExecutionPlan& plan) -> nb::list { @@ -423,11 +431,20 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { return result; }; + auto collect_scheduler_aborts = [](const tokenspeed::ExecutionPlan& plan) -> nb::list { + nb::list result; + for (const auto& abort : plan.SchedulerAborts()) { + result.append(nb::make_tuple(abort.request_id, abort.message)); + } + return result; + }; + nb::class_(m, "ExecutionPlan") .def(nb::init<>()) .def_prop_ro("forward", collect_forward) .def_prop_ro("cache", collect_cache) - .def_ro("flat_oom_request_ids", &tokenspeed::ExecutionPlan::flat_oom_request_ids); + .def_ro("flat_oom_request_ids", &tokenspeed::ExecutionPlan::flat_oom_request_ids) + .def_prop_ro("scheduler_aborts", collect_scheduler_aborts); nb::class_(m, "Scheduler") .def(nb::init(), nb::arg("config") = tokenspeed::SchedulerConfig{}) @@ -461,6 +478,12 @@ NB_MODULE(tokenspeed_scheduler_ext, m) { nb::arg("group_id")) .def("paged_cache_group_failed_alloc_count", &tokenspeed::Scheduler::PagedCacheGroupFailedAllocCount, nb::arg("group_id")) + .def("paged_cache_host_group_total_pages", &tokenspeed::Scheduler::PagedCacheHostGroupTotalPages, + nb::arg("group_id")) + .def("paged_cache_host_group_available_pages", &tokenspeed::Scheduler::PagedCacheHostGroupAvailablePages, + nb::arg("group_id")) + .def("paged_cache_host_group_failed_alloc_count", &tokenspeed::Scheduler::PagedCacheHostGroupFailedAllocCount, + nb::arg("group_id")) .def("get_request_paged_cache_page_ids", &tokenspeed::Scheduler::GetRequestPagedCachePageIds, nb::arg("request_id"), nb::arg("group_id")) .def("get_request_paged_cache_base_logical_page", &tokenspeed::Scheduler::GetRequestPagedCacheBaseLogicalPage, diff --git a/tokenspeed-scheduler/csrc/fsm/forward_events.cpp b/tokenspeed-scheduler/csrc/fsm/forward_events.cpp index 6e0755946..faf857b7d 100644 --- a/tokenspeed-scheduler/csrc/fsm/forward_events.cpp +++ b/tokenspeed-scheduler/csrc/fsm/forward_events.cpp @@ -60,6 +60,54 @@ std::vector BuildWriteBackPairs(const std::vector +tokenspeed::TreeNode* LastNodeWithResourceOrRoot(tokenspeed::TreeNode* node) { + for (tokenspeed::TreeNode* candidate = node; candidate != nullptr; candidate = candidate->Parent()) { + if (candidate->IsRoot()) return candidate; + if constexpr (RType == tokenspeed::ResourceType::Device) { + if (candidate->OnDevice()) { + return candidate; + } + } else { + if (candidate->OnHost()) { + return candidate; + } + } + } + return nullptr; +} + +struct KVResourceNodes { + tokenspeed::TreeNode* device_node{nullptr}; + tokenspeed::TreeNode* host_node{nullptr}; + std::int32_t device_page_size{0}; + std::int32_t host_page_size{0}; + + std::int32_t DeviceDepthInPage() const { + return device_node == nullptr ? 0 : device_node->DepthInPage(device_page_size); + } + + std::int32_t HostDepthInPage() const { return host_node == nullptr ? 0 : host_node->DepthInPage(host_page_size); } + + std::vector NodesWithoutDevice() const { + std::vector result; + for (tokenspeed::TreeNode* node : tokenspeed::LeafToRoot(host_node)) { + if (node->OnDevice()) break; + result.push_back(node); + } + return result; + } +}; + +KVResourceNodes GetKVResourceNodes(const tokenspeed::MatchResult& match) { + return KVResourceNodes{ + LastNodeWithResourceOrRoot(match.device.last_node), + LastNodeWithResourceOrRoot(match.host.last_node), + match.device.page_size, + match.host.page_size, + }; +} + std::vector MambaNodesForTransferPairs(const std::vector& candidates, const std::vector& transfers) { std::unordered_set src_slots; @@ -78,6 +126,18 @@ std::vector MambaNodesForTransferPairs(const std::vector< return nodes; } +std::vector PagedCacheSnapshotNodesForWriteBack(tokenspeed::TreeNode* last_node) { + std::vector nodes; + for (tokenspeed::TreeNode* node : tokenspeed::LeafToRoot(last_node)) { + if (node == nullptr) continue; + if (node->HasPagedCacheSnapshot() && !node->HasPagedCacheHostSnapshot() && + !node->HasPagedCachePendingHostSnapshot()) { + nodes.push_back(node); + } + } + return nodes; +} + void DemoteWrittenBackDevice(tokenspeed::KVPrefixCache* kv_prefix_cache, tokenspeed::HybridPrefixCache* hybrid_prefix_cache, tokenspeed::TreeNode* device_node) { if (kv_prefix_cache == nullptr || device_node == nullptr) return; @@ -88,6 +148,28 @@ void DemoteWrittenBackDevice(tokenspeed::KVPrefixCache* kv_prefix_cache, }); } +void ReleaseFailedHostWriteBack(tokenspeed::KVPrefixCache* kv_prefix_cache, + tokenspeed::HybridPrefixCache* hybrid_prefix_cache, + const std::vector& host_writeback_nodes) { + if (kv_prefix_cache == nullptr || host_writeback_nodes.empty()) return; + kv_prefix_cache->ReleaseHostResources(host_writeback_nodes, [hybrid_prefix_cache](tokenspeed::TreeNode* n) { + if (hybrid_prefix_cache != nullptr) { + hybrid_prefix_cache->OnKVHostEvict(n); + } + }); +} + +void RollbackHostWriteBack(tokenspeed::KVPrefixCache* kv_prefix_cache, + tokenspeed::HybridPrefixCache* hybrid_prefix_cache, tokenspeed::fsm::WritingBack& state) { + if (hybrid_prefix_cache != nullptr) { + hybrid_prefix_cache->OnMambaHostWriteBackDone(state.MambaWriteBackNodes(), false); + hybrid_prefix_cache->OnPagedCacheHostWriteBackDone(state.PagedCacheWriteBackNodes(), false); + } + state.DropHostNodeRef(); + ReleaseFailedHostWriteBack(kv_prefix_cache, hybrid_prefix_cache, state.HostWriteBackNodes()); + state.DropDeviceNodeRef(); +} + bool ShouldPublishMambaCheckpoint(tokenspeed::HybridPrefixCache* hybrid_cache, std::int32_t chunk_begin, std::int32_t chunk_size, std::int32_t page_size) { if (hybrid_cache == nullptr || chunk_size <= 0 || page_size <= 0) return false; @@ -220,13 +302,13 @@ std::variant SchedulePrefillFirstChunkEvent::operator() #else std::unique_ptr host_node_ref{nullptr}; std::unique_ptr device_node_ref{nullptr}; - if (!disable_l2_cache_ && (match_result_.host.DepthInPage() > match_result_.device.DepthInPage())) { - host_node_ref = std::make_unique(match_result_.host.last_node); - kv_prefix_cache_->AllocateResourceOfType( - match_result_.NodesWithout()); - device_node_ref = std::make_unique(match_result_.host.last_node); + const KVResourceNodes kv_nodes = GetKVResourceNodes(match_result_); + if (!disable_l2_cache_ && (kv_nodes.HostDepthInPage() > kv_nodes.DeviceDepthInPage())) { + host_node_ref = std::make_unique(kv_nodes.host_node); + kv_prefix_cache_->AllocateResourceOfType(kv_nodes.NodesWithoutDevice()); + device_node_ref = std::make_unique(kv_nodes.host_node); } else { - device_node_ref = std::make_unique(match_result_.device.last_node); + device_node_ref = std::make_unique(kv_nodes.device_node); } auto local_kv_allocator = std::make_unique(device_allocator_, tokens_this_round_); @@ -248,10 +330,13 @@ std::variant SchedulePrefillFirstChunkEvent::operator() TokenContainer* token_container = state.GetTokenContainer(); - std::int32_t max_matched_pages = - disable_l2_cache_ ? match_result_.device.DepthInPage() - : std::max(match_result_.device.DepthInPage(), match_result_.host.DepthInPage()); - std::int32_t window_begin = max_matched_pages * state.GetPageSize(); + std::int32_t window_begin = matched_prefix_len_tokens_; + if (window_begin < 0) { + const std::int32_t max_matched_pages = + disable_l2_cache_ ? match_result_.device.DepthInPage() + : std::max(match_result_.device.DepthInPage(), match_result_.host.DepthInPage()); + window_begin = max_matched_pages * state.GetPageSize(); + } TokenContainer::Window window{.begin = window_begin, .size = tokens_this_round_}; bool is_last_chunk = (window.begin + window.size) == token_container->PrefillSize(); @@ -468,17 +553,18 @@ Decoding ScheduleDecodeFromRetractedEvent::operator()(Retracted&& state) { #else std::unique_ptr host_node_ref{nullptr}; std::unique_ptr device_node_ref{nullptr}; - if (match_result_.host.DepthInPage() > match_result_.device.DepthInPage()) { - host_node_ref = std::make_unique(match_result_.host.last_node); - if (!kv_prefix_cache_->AllocateResourceOfType( - match_result_.NodesWithout())) { + const KVResourceNodes kv_nodes = GetKVResourceNodes(match_result_); + if (kv_nodes.HostDepthInPage() > kv_nodes.DeviceDepthInPage()) { + host_node_ref = std::make_unique(kv_nodes.host_node); + if (!kv_prefix_cache_->AllocateResourceOfType(kv_nodes.NodesWithoutDevice())) { + // Device allocation failed (race between capacity check and actual alloc). throw std::logic_error( "ScheduleDecodeFromRetractedEvent: failed to allocate device pages for host cache recovery"); } - // Device pages were just attached along the host-matched chain: pinning the HOST last node is not a typo. - device_node_ref = std::make_unique(match_result_.host.last_node); + // This is not a typo + device_node_ref = std::make_unique(kv_nodes.host_node); } else { - device_node_ref = std::make_unique(match_result_.device.last_node); + device_node_ref = std::make_unique(kv_nodes.device_node); } TokenContainer* token_container = state.GetTokenContainer(); std::int32_t page_size = state.GetPageSize(); @@ -578,14 +664,24 @@ std::variant FinishEvent::apply(ForwardStateT&& state) { auto pages_to_transfer = BuildWriteBackPairs(write_diff); std::vector mamba_writeback_nodes; + std::vector paged_cache_writeback_transfers; + std::vector paged_cache_writeback_nodes; if (hybrid_prefix_cache_ != nullptr) { auto mamba_pairs = hybrid_prefix_cache_->PrepareMambaHostWriteBack(write_diff); mamba_writeback_nodes = MambaNodesForTransferPairs(write_diff, mamba_pairs); pages_to_transfer.insert(pages_to_transfer.end(), std::make_move_iterator(mamba_pairs.begin()), std::make_move_iterator(mamba_pairs.end())); + paged_cache_writeback_nodes = PagedCacheSnapshotNodesForWriteBack(match.device.last_node); + paged_cache_writeback_transfers = + hybrid_prefix_cache_->PreparePagedCacheHostWriteBack(paged_cache_writeback_nodes); } - return Draining{std::move(pages_to_transfer), std::move(device_node_ref), std::move(host_node_ref), - std::move(mamba_writeback_nodes)}; + return Draining{std::move(pages_to_transfer), + std::move(device_node_ref), + std::move(host_node_ref), + std::move(write_diff), + std::move(mamba_writeback_nodes), + std::move(paged_cache_writeback_transfers), + std::move(paged_cache_writeback_nodes)}; } return Finished{}; #endif @@ -608,15 +704,23 @@ WritingBack FinishEvent::operator()(Retracting&& state) { WritingBack CommitDrainingEvent::operator()(Draining&& state) { auto device_node_ref = std::move(state).TakeDeviceNodeRef(); auto host_node_ref = std::move(state).TakeHostNodeRef(); + auto host_writeback_nodes = std::move(state).TakeHostWriteBackNodes(); auto mamba_writeback_nodes = std::move(state).TakeMambaWriteBackNodes(); - return WritingBack{std::move(device_node_ref), std::move(host_node_ref), std::move(mamba_writeback_nodes)}; + auto paged_cache_writeback_nodes = std::move(state).TakePagedCacheWriteBackNodes(); + return WritingBack{std::move(device_node_ref), std::move(host_node_ref), std::move(host_writeback_nodes), + std::move(mamba_writeback_nodes), std::move(paged_cache_writeback_nodes)}; } // WritingBack -> Finished: written-back cache demotes to host-only, so the next hit must load back. Finished WriteBackDoneEvent::operator()(WritingBack&& state) { TreeNode* device_node = state.DeviceNode(); + if (!success_) { + RollbackHostWriteBack(kv_prefix_cache_, hybrid_prefix_cache_, state); + return Finished{}; + } if (hybrid_prefix_cache_ != nullptr) { - hybrid_prefix_cache_->OnMambaHostWriteBackDone(state.MambaWriteBackNodes()); + hybrid_prefix_cache_->OnMambaHostWriteBackDone(state.MambaWriteBackNodes(), true); + hybrid_prefix_cache_->OnPagedCacheHostWriteBackDone(state.PagedCacheWriteBackNodes(), true); } state.DropDeviceNodeRef(); DemoteWrittenBackDevice(kv_prefix_cache_, hybrid_prefix_cache_, device_node); @@ -626,15 +730,20 @@ Finished WriteBackDoneEvent::operator()(WritingBack&& state) { return Finished{}; } -Retracted WriteBackDoneEvent::operator()(Retracting&& state) { +std::variant WriteBackDoneEvent::operator()(Retracting&& state) { #if TOKENSPEED_FLAT_KVCACHE FlatRetractUnsupported(); #else + if (!success_) { + RollbackHostWriteBack(kv_prefix_cache_, hybrid_prefix_cache_, state); + return Finished{}; + } TokenContainer* token_container = state.GetTokenContainer(); std::int32_t page_size = state.GetPageSize(); TreeNode* device_node = state.DeviceNode(); if (hybrid_prefix_cache_ != nullptr) { - hybrid_prefix_cache_->OnMambaHostWriteBackDone(state.MambaWriteBackNodes()); + hybrid_prefix_cache_->OnMambaHostWriteBackDone(state.MambaWriteBackNodes(), true); + hybrid_prefix_cache_->OnPagedCacheHostWriteBackDone(state.PagedCacheWriteBackNodes(), true); } state.DropDeviceNodeRef(); DemoteWrittenBackDevice(kv_prefix_cache_, hybrid_prefix_cache_, device_node); @@ -658,7 +767,15 @@ Aborting AbortEvent::operator()(Prefetching&& state) { return Aborting{std::move(state).TakeHostPages()}; } -Finished AbortEvent::operator()(Draining&&) { +Finished AbortEvent::operator()(Draining&& state) { + auto device_node_ref = std::move(state).TakeDeviceNodeRef(); + auto host_node_ref = std::move(state).TakeHostNodeRef(); + auto host_writeback_nodes = std::move(state).TakeHostWriteBackNodes(); + auto mamba_writeback_nodes = std::move(state).TakeMambaWriteBackNodes(); + auto paged_cache_writeback_nodes = std::move(state).TakePagedCacheWriteBackNodes(); + auto writeback = WritingBack{std::move(device_node_ref), std::move(host_node_ref), std::move(host_writeback_nodes), + std::move(mamba_writeback_nodes), std::move(paged_cache_writeback_nodes)}; + RollbackHostWriteBack(kv_prefix_cache_, hybrid_prefix_cache_, writeback); return Finished{}; } @@ -703,10 +820,6 @@ Finished AbortEvent::operator()(Decoding&& state) { return Finished{}; } -Finished AbortEvent::operator()(Retracting&&) { - return Finished{}; -} - Finished AbortEvent::operator()(Retracted&&) { return Finished{}; } @@ -741,7 +854,10 @@ Retracting ScheduleRetractEvent::applyRetract(ForwardStateT&& state) { std::unique_ptr device_node_ref = nullptr; std::unique_ptr host_node_ref = nullptr; std::vector pages_to_transfer; + std::vector host_writeback_nodes; std::vector mamba_writeback_nodes; + std::vector paged_cache_writeback_transfers; + std::vector paged_cache_writeback_nodes; if (match_result_.device.DepthInPage() > match_result_.host.DepthInPage()) { std::vector write_diff = match_result_.NodesWithout(); @@ -749,12 +865,16 @@ Retracting ScheduleRetractEvent::applyRetract(ForwardStateT&& state) { if (!kv_prefix_cache_->AllocateResourceOfType(write_diff)) { throw std::logic_error("ScheduleRetractEvent: failed to allocate host pages for device cache writeback"); } + host_writeback_nodes = write_diff; pages_to_transfer = BuildWriteBackPairs(write_diff); if (hybrid_prefix_cache_ != nullptr) { auto mamba_pairs = hybrid_prefix_cache_->PrepareMambaHostWriteBack(write_diff); mamba_writeback_nodes = MambaNodesForTransferPairs(write_diff, mamba_pairs); pages_to_transfer.insert(pages_to_transfer.end(), std::make_move_iterator(mamba_pairs.begin()), std::make_move_iterator(mamba_pairs.end())); + paged_cache_writeback_nodes = PagedCacheSnapshotNodesForWriteBack(match_result_.device.last_node); + paged_cache_writeback_transfers = + hybrid_prefix_cache_->PreparePagedCacheHostWriteBack(paged_cache_writeback_nodes); } host_node_ref = std::make_unique(match_result_.device.last_node); } else { @@ -788,7 +908,10 @@ Retracting ScheduleRetractEvent::applyRetract(ForwardStateT&& state) { std::move(local_allocator), std::move(pages_to_transfer), std::move(mamba_writeback_nodes), - std::move(local_mamba_allocator)}; + std::move(local_mamba_allocator), + std::move(paged_cache_writeback_transfers), + std::move(paged_cache_writeback_nodes), + std::move(host_writeback_nodes)}; #endif } diff --git a/tokenspeed-scheduler/csrc/fsm/forward_events.h b/tokenspeed-scheduler/csrc/fsm/forward_events.h index 811b5148a..0f46d125f 100644 --- a/tokenspeed-scheduler/csrc/fsm/forward_events.h +++ b/tokenspeed-scheduler/csrc/fsm/forward_events.h @@ -71,30 +71,32 @@ struct SchedulePrefillFirstChunkEvent : InvalidTransitionHandler loadback_diff, HybridPrefixCache* hybrid_prefix_cache = nullptr, MambaChunkAllocator* mamba_allocator = nullptr, - std::vector mamba_loadback_nodes = {} + std::vector mamba_loadback_nodes = {}, #if TOKENSPEED_FLAT_KVCACHE // coordinator defaults to nullptr because radix-only call // sites (production and tests) compile in flat builds too; // every flat transition body asserts coordinator_ != nullptr. - , KvCacheCoordinator* coordinator = nullptr, // Admission-layer prefix match, threaded from the scheduler; // default {} is the zero hit for call sites that never match. CoordinatorMatch flat_hit = {}, // Host-tier match above flat_hit's boundary (read-only; the // load emission pins both sides when it builds the ticket). - CoordinatorMatch flat_host = {}, std::vector flat_ext_hashes = {} + CoordinatorMatch flat_host = {}, std::vector flat_ext_hashes = {}, #endif - ) + std::vector paged_cache_loadback_transfers = {}, + std::int32_t matched_prefix_len_tokens = -1) : tokens_this_round_(tokens_this_round), decode_input_tokens_(decode_input_tokens), device_allocator_(device_allocator), req_pool_allocator_(req_pool_allocator), - match_result_(match_result), + match_result_(std::move(match_result)), role_{role}, disable_l2_cache_{disable_l2_cache}, loadback_diff_(std::move(loadback_diff)), mamba_loadback_nodes_(std::move(mamba_loadback_nodes)), + paged_cache_loadback_transfers_(std::move(paged_cache_loadback_transfers)), + matched_prefix_len_tokens_(matched_prefix_len_tokens), kv_prefix_cache_(kv_prefix_cache), hybrid_prefix_cache_(hybrid_prefix_cache), mamba_allocator_(mamba_allocator) @@ -111,10 +113,13 @@ struct SchedulePrefillFirstChunkEvent : InvalidTransitionHandler operator()(Submitted&& state); - const MatchResult GetMatchResult() const { return match_result_; } + const MatchResult& GetMatchResult() const { return match_result_; } - const std::vector& GetLoadbackDiff() const { return loadback_diff_; } - const std::vector& GetMambaLoadbackNodes() const { return mamba_loadback_nodes_; } + std::vector TakeLoadbackDiff() { return std::move(loadback_diff_); } + std::vector TakeMambaLoadbackNodes() { return std::move(mamba_loadback_nodes_); } + std::vector TakePagedCacheLoadbackTransfers() { + return std::move(paged_cache_loadback_transfers_); + } #if TOKENSPEED_FLAT_KVCACHE // Post-apply channel for the scheduler's LoadBack emission (transition fills the pairs). @@ -126,11 +131,13 @@ struct SchedulePrefillFirstChunkEvent : InvalidTransitionHandler loadback_diff_; std::vector mamba_loadback_nodes_; + std::vector paged_cache_loadback_transfers_; + std::int32_t matched_prefix_len_tokens_{-1}; KVPrefixCache* kv_prefix_cache_; HybridPrefixCache* hybrid_prefix_cache_{}; MambaChunkAllocator* mamba_allocator_{}; @@ -211,7 +218,8 @@ struct ScheduleDecodeFromRetractedEvent : InvalidTransitionHandler loadback_diff, MambaChunkAllocator* mamba_allocator = nullptr, - std::vector mamba_loadback_nodes = {}) + std::vector mamba_loadback_nodes = {}, + std::vector paged_cache_loadback_transfers = {}) : decode_input_tokens_(decode_input_tokens), device_allocator_(device_allocator), req_pool_allocator_(req_pool_allocator), @@ -219,6 +227,7 @@ struct ScheduleDecodeFromRetractedEvent : InvalidTransitionHandler& GetLoadbackDiff() const { return loadback_diff_; } const std::vector& GetMambaLoadbackNodes() const { return mamba_loadback_nodes_; } + const std::vector& GetPagedCacheLoadbackTransfers() const { + return paged_cache_loadback_transfers_; + } + std::vector TakeLoadbackDiff() { return std::move(loadback_diff_); } + std::vector TakeMambaLoadbackNodes() { return std::move(mamba_loadback_nodes_); } + std::vector TakePagedCacheLoadbackTransfers() { + return std::move(paged_cache_loadback_transfers_); + } private: std::int32_t decode_input_tokens_{}; @@ -236,6 +253,7 @@ struct ScheduleDecodeFromRetractedEvent : InvalidTransitionHandler loadback_diff_; std::vector mamba_loadback_nodes_; + std::vector paged_cache_loadback_transfers_; MambaChunkAllocator* mamba_allocator_{}; }; @@ -288,8 +306,14 @@ struct FinishEvent : InvalidTransitionHandler { struct AbortEvent : InvalidTransitionHandler { using InvalidTransitionHandler::operator(); + AbortEvent() = default; + AbortEvent(KVPrefixCache* kv_prefix_cache, HybridPrefixCache* hybrid_prefix_cache) + : kv_prefix_cache_(kv_prefix_cache), hybrid_prefix_cache_(hybrid_prefix_cache) {} + #if TOKENSPEED_FLAT_KVCACHE - explicit AbortEvent(KvCacheCoordinator* coordinator = nullptr) : coordinator_(coordinator) {} + explicit AbortEvent(KvCacheCoordinator* coordinator) : coordinator_(coordinator) {} + AbortEvent(KVPrefixCache* kv_prefix_cache, HybridPrefixCache* hybrid_prefix_cache, KvCacheCoordinator* coordinator) + : kv_prefix_cache_(kv_prefix_cache), hybrid_prefix_cache_(hybrid_prefix_cache), coordinator_(coordinator) {} #endif Finished operator()(Submitted&& state); @@ -298,15 +322,16 @@ struct AbortEvent : InvalidTransitionHandler { Finished operator()(Prefilling&&); Finished operator()(PrefillDone&&); Finished operator()(Decoding&&); - Finished operator()(Retracting&&); Finished operator()(Retracted&&); Finished operator()(Draining&&); // Defensive: late or duplicate abort after terminalization, stay Finished. Finished operator()(Finished&& state) { return std::move(state); } Aborting operator()(Aborting&& state); // Defensive: duplicate abort, stay Aborting -#if TOKENSPEED_FLAT_KVCACHE private: + KVPrefixCache* kv_prefix_cache_{}; + HybridPrefixCache* hybrid_prefix_cache_{}; +#if TOKENSPEED_FLAT_KVCACHE KvCacheCoordinator* coordinator_{}; #endif }; @@ -336,21 +361,19 @@ struct ScheduleRetractEvent : InvalidTransitionHandler { HybridPrefixCache* hybrid_prefix_cache = nullptr) : kv_prefix_cache_(kv_prefix_cache), host_allocator_(host_allocator), - match_result_(match_result), + match_result_(std::move(match_result)), hybrid_prefix_cache_(hybrid_prefix_cache) {} Retracting operator()(Decoding&& state); Retracting operator()(PrefillDone&& state); - MatchResult GetMatchResult() { return match_result_; } - private: template Retracting applyRetract(ForwardStateT&& state); KVPrefixCache* kv_prefix_cache_{}; PageAllocator* host_allocator_{}; - const MatchResult match_result_{}; + MatchResult match_result_{}; HybridPrefixCache* hybrid_prefix_cache_{}; }; @@ -362,20 +385,20 @@ struct CommitDrainingEvent : InvalidTransitionHandler { }; // WritingBack → Finished: async Device→Host transfer complete; node-ref locks released. -// Retracting → Retracted: same transfer path for preempted requests; -// device_node_ref drops (frees GPU pages), host_node_ref moves into Retracted. +// Retracting → Retracted on success, Finished on failure; the device pin is released in either case. struct WriteBackDoneEvent : InvalidTransitionHandler { explicit WriteBackDoneEvent(KVPrefixCache* kv_prefix_cache = nullptr, - HybridPrefixCache* hybrid_prefix_cache = nullptr) - : kv_prefix_cache_(kv_prefix_cache), hybrid_prefix_cache_(hybrid_prefix_cache) {} + HybridPrefixCache* hybrid_prefix_cache = nullptr, bool success = true) + : kv_prefix_cache_(kv_prefix_cache), hybrid_prefix_cache_(hybrid_prefix_cache), success_(success) {} using InvalidTransitionHandler::operator(); Finished operator()(WritingBack&& state); - Retracted operator()(Retracting&& state); + std::variant operator()(Retracting&& state); private: KVPrefixCache* kv_prefix_cache_{}; HybridPrefixCache* hybrid_prefix_cache_{}; + bool success_{true}; }; struct UpdateReserveNumTokensEvent : InvalidTransitionHandler { diff --git a/tokenspeed-scheduler/csrc/fsm/forward_states.h b/tokenspeed-scheduler/csrc/fsm/forward_states.h index da8ce4b53..5c9c5afce 100644 --- a/tokenspeed-scheduler/csrc/fsm/forward_states.h +++ b/tokenspeed-scheduler/csrc/fsm/forward_states.h @@ -322,46 +322,69 @@ struct Draining { // newWriteBackOperation split-safe (no re-walk after splitChild redistributes pages). using PagePair = TransferPair; Draining(std::vector pages_to_transfer, std::unique_ptr&& device_node_ref, - std::unique_ptr&& host_node_ref, std::vector mamba_writeback_nodes = {}) + std::unique_ptr&& host_node_ref, std::vector host_writeback_nodes = {}, + std::vector mamba_writeback_nodes = {}, + std::vector paged_cache_writeback_transfers = {}, + std::vector paged_cache_writeback_nodes = {}) : pages_to_transfer_(std::move(pages_to_transfer)), device_node_ref_(std::move(device_node_ref)), host_node_ref_(std::move(host_node_ref)), - mamba_writeback_nodes_(std::move(mamba_writeback_nodes)) {} + host_writeback_nodes_(std::move(host_writeback_nodes)), + mamba_writeback_nodes_(std::move(mamba_writeback_nodes)), + paged_cache_writeback_transfers_(std::move(paged_cache_writeback_transfers)), + paged_cache_writeback_nodes_(std::move(paged_cache_writeback_nodes)) {} public: const std::vector& GetPagesToTransfer() const { return pages_to_transfer_; } std::unique_ptr TakeDeviceNodeRef() && { return std::move(device_node_ref_); } std::unique_ptr TakeHostNodeRef() && { return std::move(host_node_ref_); } + std::vector TakeHostWriteBackNodes() && { return std::move(host_writeback_nodes_); } std::vector TakeMambaWriteBackNodes() && { return std::move(mamba_writeback_nodes_); } + const std::vector& GetPagedCacheWriteBackTransfers() const { + return paged_cache_writeback_transfers_; + } + const std::vector& PagedCacheWriteBackNodes() const { return paged_cache_writeback_nodes_; } + std::vector TakePagedCacheWriteBackNodes() && { return std::move(paged_cache_writeback_nodes_); } private: std::vector pages_to_transfer_; // concrete mixed-kind pairs to copy std::unique_ptr device_node_ref_; // keeps matched Device node alive until WritingBack std::unique_ptr host_node_ref_; // keeps pre-allocated Host node alive until WritingBack + std::vector host_writeback_nodes_; // exact KV host resources newly allocated by this op std::vector mamba_writeback_nodes_; // exact Mamba nodes covered by this writeback op + std::vector paged_cache_writeback_transfers_; + std::vector paged_cache_writeback_nodes_; }; // Writeback op executing; both node refs are RAII locks pinning the pages while the transfer is in flight. struct WritingBack { WritingBack(std::unique_ptr&& device_node_ref, std::unique_ptr&& host_node_ref, - std::vector mamba_writeback_nodes = {}) + std::vector host_writeback_nodes = {}, std::vector mamba_writeback_nodes = {}, + std::vector paged_cache_writeback_nodes = {}) : device_node_ref_(std::move(device_node_ref)), host_node_ref_(std::move(host_node_ref)), - mamba_writeback_nodes_(std::move(mamba_writeback_nodes)) {} + host_writeback_nodes_(std::move(host_writeback_nodes)), + mamba_writeback_nodes_(std::move(mamba_writeback_nodes)), + paged_cache_writeback_nodes_(std::move(paged_cache_writeback_nodes)) {} WritingBack(WritingBack&&) noexcept = default; WritingBack& operator=(WritingBack&&) noexcept = default; std::unique_ptr TakeHostNodeRef() && { return std::move(host_node_ref_); } TreeNode* DeviceNode() const { return device_node_ref_ ? device_node_ref_->Node() : nullptr; } + const std::vector& HostWriteBackNodes() const { return host_writeback_nodes_; } const std::vector& MambaWriteBackNodes() const { return mamba_writeback_nodes_; } + const std::vector& PagedCacheWriteBackNodes() const { return paged_cache_writeback_nodes_; } void DropDeviceNodeRef() { device_node_ref_.reset(); } + void DropHostNodeRef() { host_node_ref_.reset(); } private: std::unique_ptr device_node_ref_; // released after WriteBackDone std::unique_ptr host_node_ref_; // released after WriteBackDone + std::vector host_writeback_nodes_; // newly attached KV host resources std::vector mamba_writeback_nodes_; // pending host Mamba slots published by this op ack + std::vector paged_cache_writeback_nodes_; }; // Keeps the local KV allocator (tail-page info) and token container for recovery. @@ -371,13 +394,18 @@ struct Retracting : public WritingBack { Retracting(TokenContainer* token_container, std::int32_t page_size, std::unique_ptr&& host_node_ref, std::unique_ptr&& device_node_ref, std::unique_ptr&& local_kv_allocator, std::vector pages_to_transfer, std::vector mamba_writeback_nodes = {}, - std::unique_ptr&& local_mamba_allocator = nullptr) - : WritingBack(std::move(device_node_ref), std::move(host_node_ref), std::move(mamba_writeback_nodes)), + std::unique_ptr&& local_mamba_allocator = nullptr, + std::vector paged_cache_writeback_transfers = {}, + std::vector paged_cache_writeback_nodes = {}, + std::vector host_writeback_nodes = {}) + : WritingBack(std::move(device_node_ref), std::move(host_node_ref), std::move(host_writeback_nodes), + std::move(mamba_writeback_nodes), std::move(paged_cache_writeback_nodes)), token_container_{token_container}, page_size_{page_size}, local_kv_allocator_(std::move(local_kv_allocator)), pages_to_transfer_(std::move(pages_to_transfer)), - local_mamba_allocator_(std::move(local_mamba_allocator)) {} + local_mamba_allocator_(std::move(local_mamba_allocator)), + paged_cache_writeback_transfers_(std::move(paged_cache_writeback_transfers)) {} Retracting(Retracting&&) noexcept = default; Retracting& operator=(Retracting&&) noexcept = default; @@ -388,6 +416,9 @@ struct Retracting : public WritingBack { std::int32_t GetPageSize() const { return page_size_; } const std::vector& GetPagesToTransfer() const { return pages_to_transfer_; } + const std::vector& GetPagedCacheWriteBackTransfers() const { + return paged_cache_writeback_transfers_; + } void ExtendResultTokens(const std::vector result_tokens) { token_container_->Extend(result_tokens); } std::vector GetLocalAllocatorPages() const { @@ -400,6 +431,7 @@ struct Retracting : public WritingBack { std::unique_ptr local_kv_allocator_{}; std::vector pages_to_transfer_{}; std::unique_ptr local_mamba_allocator_{}; + std::vector paged_cache_writeback_transfers_{}; }; struct Retracted { diff --git a/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.cpp b/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.cpp index 84905647f..7d5439fff 100644 --- a/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.cpp +++ b/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.cpp @@ -383,6 +383,45 @@ void PagedCacheGroupTable::ImportPrefixBorrowed(std::vector ids, s RefreshPageIdsView(); } +void PagedCacheGroupTable::ImportPrefixOwned(OwnedPages pages, std::int32_t base_logical_page, + std::int32_t raw_tokens_covered) { + if (allocator_ == nullptr) { + throw std::logic_error("PagedCacheGroupTable::ImportPrefixOwned: no allocator bound"); + } + if (!(borrowed_page_ids_.empty() && owned_pages_.Empty() && raw_token_cursor_ == 0 && base_logical_page_ == 0 && + committed_prefix_len_tokens_ == 0)) { + throw std::logic_error("PagedCacheGroupTable::ImportPrefixOwned: only legal on a fresh-empty table"); + } + if (base_logical_page < 0) { + throw std::invalid_argument("PagedCacheGroupTable::ImportPrefixOwned: base_logical_page must be >= 0"); + } + if (raw_tokens_covered < 0) { + throw std::invalid_argument("PagedCacheGroupTable::ImportPrefixOwned: raw_tokens_covered must be >= 0"); + } + const std::int32_t raw_per_page = allocator_->Config().RawTokensPerPage(); + if (raw_per_page <= 0) { + throw std::logic_error("PagedCacheGroupTable::ImportPrefixOwned: invalid group config"); + } + if (raw_tokens_covered % raw_per_page != 0) { + throw std::invalid_argument("PagedCacheGroupTable::ImportPrefixOwned: raw_tokens_covered must be page-aligned"); + } + if (!pages.Empty()) { + const std::int32_t covered_pages = raw_tokens_covered / raw_per_page; + if (base_logical_page + pages.Size() != covered_pages) { + throw std::invalid_argument( + "PagedCacheGroupTable::ImportPrefixOwned: page ids do not end at covered raw " + "token cursor"); + } + } + owned_pages_ = std::move(pages); + base_logical_page_ = base_logical_page; + raw_token_cursor_ = raw_tokens_covered; + committed_prefix_len_tokens_ = allocator_->Config().family == PagedCacheGroupFamily::History + ? base_logical_page_ * raw_per_page + : raw_tokens_covered; + RefreshPageIdsView(); +} + std::vector PagedCacheGroupTable::ReleaseSkipped(std::int32_t window_lower_bound) { if (allocator_ == nullptr || window_lower_bound <= 0) { return {}; diff --git a/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.h b/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.h index 306c1b6dd..a860fd3e9 100644 --- a/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.h +++ b/tokenspeed-scheduler/csrc/resource/allocator/paged_cache_group.h @@ -30,6 +30,8 @@ namespace tokenspeed { +class HybridPrefixCache; + // Positive-only ceiling division; returns 0 for non-positive numerators. // Lives here because paged-cache admission/table math is its only caller. inline std::int32_t CeilDivPositive(std::int32_t numer, std::int32_t denom) { @@ -91,6 +93,7 @@ class PagedCacheGroupAllocator { private: friend class PagedCacheGroupTable; + friend class HybridPrefixCache; // Empty OwnedPages on insufficient capacity (bumps failed_alloc_count_). OwnedPages AcquireOwned(std::int32_t num_pages); @@ -155,6 +158,13 @@ class PagedCacheGroupTable { void ImportPrefixBorrowed(std::vector ids, std::int32_t base_logical_page, std::int32_t raw_tokens_covered); + // Adopt owned device pages that were freshly materialized from host L2. + // Legal only on a fresh-empty table. History-family pages remain + // uncommitted so CommitChunk can publish them into device snapshots before + // extending the chain; State-family pages represent terminal state already + // covered by the host hit. + void ImportPrefixOwned(OwnedPages pages, std::int32_t base_logical_page, std::int32_t raw_tokens_covered); + // Sliding-only: drop front pages strictly below `window_lower_bound`. // On an empty table, advances base_logical_page_ so first allocation starts // at the live sliding window. Commit cursor untouched. Idempotent. diff --git a/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.cpp b/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.cpp index 80969f7b3..8741b3d55 100644 --- a/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.cpp +++ b/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.cpp @@ -34,9 +34,11 @@ #include #include #include +#include #include #include #include +#include #include namespace tokenspeed { @@ -112,12 +114,20 @@ HybridPrefixCache::~HybridPrefixCache() { TreeNode* node = *paged_cache_snapshot_nodes_.begin(); DetachPagedCacheSnapshotFromNode(node); } + while (!paged_cache_host_snapshot_nodes_.empty()) { + TreeNode* node = *paged_cache_host_snapshot_nodes_.begin(); + DetachPagedCacheHostSnapshotFromNode(node); + } + while (!paged_cache_pending_host_snapshot_nodes_.empty()) { + TreeNode* node = *paged_cache_pending_host_snapshot_nodes_.begin(); + DetachPagedCachePendingHostSnapshotFromNode(node); + } } MatchResult HybridPrefixCache::Match(const token_vec_t& token_ids, MatchIntent intent) { auto match = kv_prefix_cache_.Match(token_ids, intent); augmentMatch(match); - augmentMatchPagedCache(match); + augmentMatchPagedCache(match, intent); return match; } @@ -125,7 +135,7 @@ MatchResult HybridPrefixCache::Match(const std::vector snapshot) { + if (node == nullptr || snapshot == nullptr) return false; + if (!node->OnHost()) return false; + RefreshPagedCacheSnapshotCompleteness(*snapshot); + node->AttachPagedCacheHostSnapshot(std::move(snapshot)); + paged_cache_host_snapshot_nodes_.insert(node); + return true; +} + std::unique_ptr HybridPrefixCache::DetachPagedCacheSnapshotFromNode(TreeNode* node) { if (node == nullptr) return nullptr; paged_cache_snapshot_nodes_.erase(node); @@ -518,6 +538,18 @@ bool HybridPrefixCache::isPagedCacheSnapshotBorrowed(const TreeNode* node, return false; } +std::unique_ptr HybridPrefixCache::DetachPagedCacheHostSnapshotFromNode(TreeNode* node) { + if (node == nullptr) return nullptr; + paged_cache_host_snapshot_nodes_.erase(node); + return node->DetachPagedCacheHostSnapshot(); +} + +std::unique_ptr HybridPrefixCache::DetachPagedCachePendingHostSnapshotFromNode(TreeNode* node) { + if (node == nullptr) return nullptr; + paged_cache_pending_host_snapshot_nodes_.erase(node); + return node->DetachPagedCachePendingHostSnapshot(); +} + void HybridPrefixCache::OnKVEvict(TreeNode* node) { if (node == nullptr) return; if (mamba_allocator_ != nullptr && node->HasMamba()) { @@ -557,6 +589,12 @@ void HybridPrefixCache::OnNodeDestroyed(TreeNode* node) { "HybridPrefixCache::OnNodeDestroyed: paged snapshot still has a request-table borrower"); DetachPagedCacheSnapshotFromNode(node); } + if (node->HasPagedCacheHostSnapshot()) { + DetachPagedCacheHostSnapshotFromNode(node); + } + if (node->HasPagedCachePendingHostSnapshot()) { + DetachPagedCachePendingHostSnapshotFromNode(node); + } // Host-side Mamba L2 bookkeeping (no-op when the L2 pool is disabled, since // these sets are only populated when mamba_host_allocator_ is present). pending_mamba_host_writebacks_.erase(node); @@ -568,7 +606,14 @@ void HybridPrefixCache::OnNodeDestroyed(TreeNode* node) { } void HybridPrefixCache::OnKVHostEvict(TreeNode* node) { - if (node == nullptr || mamba_host_allocator_ == nullptr) return; + if (node == nullptr) return; + if (node->HasPagedCacheHostSnapshot()) { + DetachPagedCacheHostSnapshotFromNode(node); + } + if (node->HasPagedCachePendingHostSnapshot()) { + DetachPagedCachePendingHostSnapshotFromNode(node); + } + if (mamba_host_allocator_ == nullptr) return; pending_mamba_host_writebacks_.erase(node); if (node->HasMambaOnHost()) { node->DetachMambaHost(); @@ -603,17 +648,17 @@ void HybridPrefixCache::DemoteIdleMambaDeviceCopiesPresentOnHost() { } } -void HybridPrefixCache::OnMambaHostWriteBackDone(TreeNode* last_node) { +void HybridPrefixCache::OnMambaHostWriteBackDone(TreeNode* last_node, bool success) { if (last_node == nullptr) return; std::vector nodes; for (TreeNode* node : LeafToRoot(last_node)) { if (node == nullptr || !node->OnHost()) break; nodes.push_back(node); } - OnMambaHostWriteBackDone(nodes); + OnMambaHostWriteBackDone(nodes, success); } -void HybridPrefixCache::OnMambaHostWriteBackDone(const std::vector& nodes) { +void HybridPrefixCache::OnMambaHostWriteBackDone(const std::vector& nodes, bool success) { if (mamba_allocator_ == nullptr || mamba_host_allocator_ == nullptr) return; std::int32_t attached = 0; @@ -622,6 +667,10 @@ void HybridPrefixCache::OnMambaHostWriteBackDone(const std::vector& n if (node == nullptr || !node->OnHost()) continue; auto pending = pending_mamba_host_writebacks_.find(node); if (pending != pending_mamba_host_writebacks_.end()) { + if (!success) { + pending_mamba_host_writebacks_.erase(pending); + continue; + } node->AttachMambaHost(std::move(pending->second)); pending_mamba_host_writebacks_.erase(pending); mamba_host_nodes_.insert(node); @@ -636,11 +685,17 @@ void HybridPrefixCache::OnMambaHostWriteBackDone(const std::vector& n spdlog::debug("[HybridPrefixCache][mamba_l2] host writeback done attach_count={} completed_nodes={}", attached, completed); } - DemoteIdleMambaDeviceCopiesPresentOnHost(); + if (success) { + DemoteIdleMambaDeviceCopiesPresentOnHost(); + } } void HybridPrefixCache::OnKVDeviceDemote(TreeNode* node) { - if (node == nullptr || mamba_allocator_ == nullptr) return; + if (node == nullptr) return; + if (node->HasPagedCacheSnapshot() && node->HasPagedCacheHostSnapshot() && !isPagedCacheSnapshotBorrowed(node)) { + DetachPagedCacheSnapshotFromNode(node); + } + if (mamba_allocator_ == nullptr) return; if (node->HasMamba() && node->HasMambaOnHost()) { mamba_eviction_manager_.UntrackNode(node); node->DetachMamba(); @@ -666,6 +721,29 @@ void HybridPrefixCache::RegisterPagedCacheGroup(std::unique_ptr allocator) { + if (allocator == nullptr) { + throw std::invalid_argument("HybridPrefixCache::RegisterPagedCacheHostGroup: null allocator"); + } + const std::string gid = allocator->Config().group_id; + if (paged_cache_host_allocators_.find(gid) != paged_cache_host_allocators_.end()) { + throw std::invalid_argument("HybridPrefixCache::RegisterPagedCacheHostGroup: duplicate group_id: " + gid); + } + auto device_it = paged_cache_allocators_.find(gid); + if (device_it == paged_cache_allocators_.end() || device_it->second == nullptr) { + throw std::invalid_argument("HybridPrefixCache::RegisterPagedCacheHostGroup: device group missing: " + gid); + } + const auto& host_cfg = allocator->Config(); + const auto& device_cfg = device_it->second->Config(); + if (host_cfg.rows_per_page != device_cfg.rows_per_page || + host_cfg.entry_stride_tokens != device_cfg.entry_stride_tokens || host_cfg.retention != device_cfg.retention || + host_cfg.sliding_window_tokens != device_cfg.sliding_window_tokens || host_cfg.family != device_cfg.family) { + throw std::invalid_argument("HybridPrefixCache::RegisterPagedCacheHostGroup: host/device config mismatch for " + + gid); + } + paged_cache_host_allocators_.emplace(gid, std::move(allocator)); +} + void HybridPrefixCache::EnablePagedCacheAdjunct( std::vector required_groups, std::unordered_map sliding_window_per_group) { if (required_groups.empty()) { @@ -776,6 +854,21 @@ TreeNode* CapNodeToDepth(TreeNode* from, std::int32_t depth) { return node; } +enum class PagedCacheResidency { kDevice, kHost }; + +const PagedCacheSnapshot* GetPagedCacheSnapshotForResidency(TreeNode* node, PagedCacheResidency residency) { + if (node == nullptr) return nullptr; + switch (residency) { + case PagedCacheResidency::kDevice: + if (!node->OnDevice()) return nullptr; + return node->GetPagedCacheSnapshot(); + case PagedCacheResidency::kHost: + if (!node->OnHost()) return nullptr; + return node->GetPagedCacheHostSnapshot(); + } + return nullptr; +} + // Ancestor path (excluding root), reversed so element 0 is closest to root. std::vector CollectAncestorPathRootToLeaf(TreeNode* from) { std::vector path; @@ -787,11 +880,11 @@ std::vector CollectAncestorPathRootToLeaf(TreeNode* from) { } void AssemblePagedCacheGroupPages(MatchResult::PagedCache& out, const std::string& gid, - std::span chain, bool is_sliding) { + std::span chain, bool is_sliding, PagedCacheResidency residency) { std::vector page_ids; std::int32_t base_logical_page = 0; if (!chain.empty()) { - const PagedCacheSnapshot* earliest_snap = chain.front()->GetPagedCacheSnapshot(); + const PagedCacheSnapshot* earliest_snap = GetPagedCacheSnapshotForResidency(chain.front(), residency); if (earliest_snap != nullptr && is_sliding) { auto git = earliest_snap->groups.find(gid); if (git != earliest_snap->groups.end()) { @@ -799,7 +892,7 @@ void AssemblePagedCacheGroupPages(MatchResult::PagedCache& out, const std::strin } } for (TreeNode* anc : chain) { - const PagedCacheSnapshot* snap = anc->GetPagedCacheSnapshot(); + const PagedCacheSnapshot* snap = GetPagedCacheSnapshotForResidency(anc, residency); if (snap == nullptr) continue; auto git = snap->groups.find(gid); if (git == snap->groups.end()) continue; @@ -827,7 +920,8 @@ bool ImportPagedCacheGroupSnapshot(MatchResult::PagedCache& out, const std::stri bool AssemblePagedCacheStateGroupPagesToTarget(MatchResult::PagedCache& out, const std::string& gid, std::span chain, const PagedCacheGroupAllocator& allocator, - std::int32_t target_raw_tokens, std::int32_t retained_tokens) { + std::int32_t target_raw_tokens, std::int32_t retained_tokens, + PagedCacheResidency residency) { const std::int32_t raw_per_page = allocator.Config().RawTokensPerPage(); if (raw_per_page <= 0 || target_raw_tokens <= 0 || target_raw_tokens % raw_per_page != 0) return false; if (retained_tokens <= 0 || retained_tokens % raw_per_page != 0) return false; @@ -842,7 +936,7 @@ bool AssemblePagedCacheStateGroupPagesToTarget(MatchResult::PagedCache& out, con bool started = false; for (TreeNode* anc : chain) { - const PagedCacheSnapshot* snap = anc != nullptr ? anc->GetPagedCacheSnapshot() : nullptr; + const PagedCacheSnapshot* snap = GetPagedCacheSnapshotForResidency(anc, residency); if (snap == nullptr) continue; auto git = snap->groups.find(gid); if (git == snap->groups.end()) continue; @@ -881,151 +975,182 @@ bool AssemblePagedCacheStateGroupPagesToTarget(MatchResult::PagedCache& out, con } // namespace -void HybridPrefixCache::augmentMatchPagedCache(MatchResult& match) const { +void HybridPrefixCache::augmentMatchPagedCache(MatchResult& match, MatchIntent intent) const { if (!HasPagedCacheAdjunct()) return; - if (match.device.last_node == nullptr) return; + TreeNode* token_terminal = match.token_terminal != nullptr ? match.token_terminal : match.device.last_node; + if (token_terminal == nullptr) return; const std::int32_t align = paged_cache_history_alignment_tokens_; auto cap_to_root = [&]() { - TreeNode* root = RootOf(match.device.last_node); + TreeNode* root = RootOf(token_terminal); match.device.last_node = root; match.host.last_node = RootOf(match.host.last_node); match.paged_cache = MatchResult::PagedCache{}; + match.paged_cache_host = MatchResult::PagedCache{}; }; - std::vector path = CollectAncestorPathRootToLeaf(match.device.last_node); - - TreeNode* deepest_history = nullptr; - std::vector history_chain; - std::int32_t expected_depth = align; - for (TreeNode* n : path) { - const std::int32_t d = static_cast(n->DepthInTokens()); - if (d < expected_depth) continue; - if (d > expected_depth) break; - const auto* snap = n->GetPagedCacheSnapshot(); - if (snap == nullptr) break; - if (!snap->IsCompleteFor(PagedCacheGroupFamily::History)) break; - deepest_history = n; - history_chain.push_back(n); - expected_depth += align; - } - if (deepest_history == nullptr) { - cap_to_root(); - return; - } - - match.paged_cache.per_group_page_ids.clear(); - match.paged_cache.per_group_base_logical_page.clear(); - match.paged_cache.history_hit_tokens = static_cast(deepest_history->DepthInTokens()); + auto build_hit = [&](PagedCacheResidency residency) -> MatchResult::PagedCache { + MatchResult::PagedCache hit{}; + std::vector path = CollectAncestorPathRootToLeaf(token_terminal); + + TreeNode* deepest_history = nullptr; + std::vector history_chain; + std::int32_t expected_depth = align; + for (TreeNode* n : path) { + const std::int32_t d = static_cast(n->DepthInTokens()); + if (d < expected_depth) continue; + if (d > expected_depth) break; + const auto* snap = GetPagedCacheSnapshotForResidency(n, residency); + if (snap == nullptr) break; + if (!snap->IsCompleteFor(PagedCacheGroupFamily::History)) break; + deepest_history = n; + history_chain.push_back(n); + expected_depth += align; + } + if (deepest_history == nullptr) { + return hit; + } + + hit.history_hit_tokens = static_cast(deepest_history->DepthInTokens()); + + const std::int32_t history_hit = hit.history_hit_tokens; + auto build_exact_continuation_hit = [&](TreeNode* terminal, + std::span history_span) -> MatchResult::PagedCache { + const std::int32_t terminal_depth = static_cast(terminal->DepthInTokens()); + const auto* terminal_snap = GetPagedCacheSnapshotForResidency(terminal, residency); + const bool terminal_state_complete = + paged_cache_continuation_state_groups_.empty() || + (terminal_snap != nullptr && terminal_snap->continuation_state_complete); + if (terminal_snap == nullptr || terminal_snap->prefix_len_tokens != terminal_depth || + !terminal_state_complete) { + return {}; + } - const std::int32_t history_hit = match.paged_cache.history_hit_tokens; - if (!paged_cache_continuation_state_groups_.empty()) { - const auto* terminal_snap = deepest_history->GetPagedCacheSnapshot(); - if (terminal_snap != nullptr && terminal_snap->continuation_state_complete && - terminal_snap->prefix_len_tokens == history_hit) { MatchResult::PagedCache terminal_hit{}; - terminal_hit.last_node = deepest_history; - terminal_hit.prefix_len_tokens = history_hit; - terminal_hit.history_hit_tokens = history_hit; + terminal_hit.last_node = terminal; + terminal_hit.prefix_len_tokens = terminal_depth; + terminal_hit.history_hit_tokens = terminal_depth; - const std::span history_span{history_chain}; for (const auto& gid : paged_cache_history_groups_) { const bool is_sliding = paged_cache_sliding_window_per_group_.find(gid) != paged_cache_sliding_window_per_group_.end(); - AssemblePagedCacheGroupPages(terminal_hit, gid, history_span, is_sliding); + AssemblePagedCacheGroupPages(terminal_hit, gid, history_span, is_sliding, residency); } - bool continuation_ok = true; for (const auto& gid : paged_cache_continuation_state_groups_) { auto alloc_it = paged_cache_allocators_.find(gid); - if (alloc_it == paged_cache_allocators_.end()) { - continuation_ok = false; - break; - } + if (alloc_it == paged_cache_allocators_.end()) return {}; const std::int32_t retained_tokens = alloc_it->second->Config().sliding_window_tokens.value_or(0); if (!AssemblePagedCacheStateGroupPagesToTarget(terminal_hit, gid, history_span, *alloc_it->second, - history_hit, retained_tokens)) { - continuation_ok = false; - break; + terminal_depth, retained_tokens, residency)) { + return {}; } } - if (continuation_ok) { - match.paged_cache = std::move(terminal_hit); - match.device.last_node = deepest_history; - match.host.last_node = CapNodeToDepth(match.host.last_node, history_hit); - return; - } + return terminal_hit; + }; + + if (intent == MatchIntent::StateRecovery) { + if (deepest_history != token_terminal) return {}; + const std::span history_span{history_chain}; + return build_exact_continuation_hit(token_terminal, history_span); } - } - const bool has_transport_only_state = - paged_cache_continuation_state_group_set_.size() != paged_cache_state_group_set_.size(); - if (has_transport_only_state) { - cap_to_root(); - return; - } + if (!paged_cache_continuation_state_groups_.empty()) { + const std::span history_span{history_chain}; + MatchResult::PagedCache terminal_hit = build_exact_continuation_hit(deepest_history, history_span); + if (terminal_hit.last_node != nullptr) return terminal_hit; - std::int32_t worst_window = 0; - for (const auto& gid : paged_cache_state_groups_) { - auto it = paged_cache_sliding_window_per_group_.find(gid); - if (it != paged_cache_sliding_window_per_group_.end()) { - worst_window = std::max(worst_window, it->second); + const bool has_transport_only_state = std::any_of( + paged_cache_continuation_state_groups_.begin(), paged_cache_continuation_state_groups_.end(), + [&](const std::string& gid) { return !paged_cache_state_group_set_.contains(gid); }); + if (has_transport_only_state) return {}; } - } - const std::int32_t segments_needed = worst_window > 0 ? (worst_window + align - 1) / align : 1; - TreeNode* usable_node = nullptr; - if (paged_cache_state_groups_.empty()) { - usable_node = deepest_history; - } else { - for (std::int32_t end_idx = static_cast(history_chain.size()) - 1; end_idx >= 0; --end_idx) { - const std::int32_t start_idx = std::max(0, end_idx - segments_needed + 1); - bool ok = true; - for (std::int32_t i = start_idx; i <= end_idx; ++i) { - const auto* snap = history_chain[i]->GetPagedCacheSnapshot(); - if (snap == nullptr || !snap->IsCompleteFor(PagedCacheGroupFamily::State)) { - ok = false; + // Configurations whose State groups are all required can rebuild them from + // checkpoint segments. Pure History configurations also use this path. + std::int32_t worst_window = 0; + for (const auto& gid : paged_cache_state_groups_) { + auto it = paged_cache_sliding_window_per_group_.find(gid); + if (it != paged_cache_sliding_window_per_group_.end()) { + worst_window = std::max(worst_window, it->second); + } + } + const std::int32_t segments_needed = worst_window > 0 ? (worst_window + align - 1) / align : 1; + + TreeNode* usable_node = nullptr; + if (paged_cache_state_groups_.empty()) { + usable_node = deepest_history; + } else { + for (std::int32_t end_idx = static_cast(history_chain.size()) - 1; end_idx >= 0; --end_idx) { + const std::int32_t start_idx = std::max(0, end_idx - segments_needed + 1); + bool ok = true; + for (std::int32_t i = start_idx; i <= end_idx; ++i) { + const auto* snap = GetPagedCacheSnapshotForResidency(history_chain[i], residency); + if (snap == nullptr || !snap->IsCompleteFor(PagedCacheGroupFamily::State)) { + ok = false; + break; + } + } + if (ok) { + usable_node = history_chain[end_idx]; break; } } - if (ok) { - usable_node = history_chain[end_idx]; - break; - } } - } - if (usable_node == nullptr) { - cap_to_root(); - return; - } + if (usable_node == nullptr) { + return MatchResult::PagedCache{}; + } - const std::int32_t usable = static_cast(usable_node->DepthInTokens()); - while (!history_chain.empty() && static_cast(history_chain.back()->DepthInTokens()) > usable) { - history_chain.pop_back(); - } + const std::int32_t usable = static_cast(usable_node->DepthInTokens()); + while (!history_chain.empty() && static_cast(history_chain.back()->DepthInTokens()) > usable) { + history_chain.pop_back(); + } - match.paged_cache.last_node = usable_node; - match.paged_cache.prefix_len_tokens = usable; + hit.last_node = usable_node; + hit.prefix_len_tokens = usable; - const std::span history_span{history_chain}; - for (const auto& gid : paged_cache_history_groups_) { - const bool is_sliding = - paged_cache_sliding_window_per_group_.find(gid) != paged_cache_sliding_window_per_group_.end(); - AssemblePagedCacheGroupPages(match.paged_cache, gid, history_span, is_sliding); - } - if (!paged_cache_state_groups_.empty()) { - const std::size_t take = std::min(history_chain.size(), static_cast(segments_needed)); - const std::span state_span = history_span.last(take); - for (const auto& gid : paged_cache_state_groups_) { + const std::span history_span{history_chain}; + for (const auto& gid : paged_cache_history_groups_) { const bool is_sliding = paged_cache_sliding_window_per_group_.find(gid) != paged_cache_sliding_window_per_group_.end(); - AssemblePagedCacheGroupPages(match.paged_cache, gid, state_span, is_sliding); + AssemblePagedCacheGroupPages(hit, gid, history_span, is_sliding, residency); } + if (!paged_cache_state_groups_.empty()) { + const std::size_t take = + std::min(history_chain.size(), static_cast(segments_needed)); + const std::span state_span = history_span.last(take); + for (const auto& gid : paged_cache_state_groups_) { + const bool is_sliding = + paged_cache_sliding_window_per_group_.find(gid) != paged_cache_sliding_window_per_group_.end(); + AssemblePagedCacheGroupPages(hit, gid, state_span, is_sliding, residency); + } + } + return hit; + }; + + match.paged_cache = build_hit(PagedCacheResidency::kDevice); + match.paged_cache_host = build_hit(PagedCacheResidency::kHost); + + const bool has_device_hit = match.paged_cache.last_node != nullptr && match.paged_cache.prefix_len_tokens > 0; + const bool has_host_hit = + match.paged_cache_host.last_node != nullptr && match.paged_cache_host.prefix_len_tokens > 0; + if (!has_device_hit && !has_host_hit) { + cap_to_root(); + return; } - match.device.last_node = usable_node; + TreeNode* root = RootOf(token_terminal); + if (has_device_hit) { + match.device.last_node = match.paged_cache.last_node; + } else { + match.device.last_node = root; + } + const std::int32_t usable = std::max(match.paged_cache.prefix_len_tokens, match.paged_cache_host.prefix_len_tokens); match.host.last_node = CapNodeToDepth(match.host.last_node, usable); + if (match.host.last_node == nullptr) { + match.host.last_node = RootOf(token_terminal); + } } std::vector HybridPrefixCache::PagedCacheGroupIds() const { @@ -1061,6 +1186,30 @@ std::int64_t HybridPrefixCache::PagedCacheGroupFailedAllocCount(const std::strin return it->second->FailedAllocCount(); } +std::int32_t HybridPrefixCache::PagedCacheHostGroupTotalPages(const std::string& group_id) const { + auto it = paged_cache_host_allocators_.find(group_id); + if (it == paged_cache_host_allocators_.end()) { + throw std::out_of_range("HybridPrefixCache::PagedCacheHostGroupTotalPages: group_id not configured"); + } + return it->second->TotalPages(); +} + +std::int32_t HybridPrefixCache::PagedCacheHostGroupAvailablePages(const std::string& group_id) const { + auto it = paged_cache_host_allocators_.find(group_id); + if (it == paged_cache_host_allocators_.end()) { + throw std::out_of_range("HybridPrefixCache::PagedCacheHostGroupAvailablePages: group_id not configured"); + } + return it->second->AvailablePages(); +} + +std::int64_t HybridPrefixCache::PagedCacheHostGroupFailedAllocCount(const std::string& group_id) const { + auto it = paged_cache_host_allocators_.find(group_id); + if (it == paged_cache_host_allocators_.end()) { + throw std::out_of_range("HybridPrefixCache::PagedCacheHostGroupFailedAllocCount: group_id not configured"); + } + return it->second->FailedAllocCount(); +} + std::vector HybridPrefixCache::GetRequestPagedCachePageIds(const std::string& request_id, const std::string& group_id) const { if (paged_cache_allocators_.find(group_id) == paged_cache_allocators_.end()) { @@ -1093,6 +1242,11 @@ std::int32_t HybridPrefixCache::GetRequestPagedCacheBaseLogicalPage(const std::s return group_it->second.BaseLogicalPage(); } +bool HybridPrefixCache::HasRequestPagedCacheTables(const std::string& request_id) const { + auto it = request_paged_cache_tables_.find(request_id); + return it != request_paged_cache_tables_.end() && !it->second.empty(); +} + std::map HybridPrefixCache::InitialSimulatedFree() const { std::map out; for (const auto& [gid, allocator] : paged_cache_allocators_) { @@ -1169,6 +1323,179 @@ void HybridPrefixCache::RewindRequest(const std::string& request_id, std::int32_ } } +std::vector HybridPrefixCache::PreparePagedCacheDeviceLoadBack( + const std::string& request_id, const MatchResult::PagedCache& host_hit) { + std::vector transfers; + if (!HasPagedCacheHostAdjunct() || host_hit.last_node == nullptr || host_hit.prefix_len_tokens <= 0) { + return transfers; + } + auto& tables = request_paged_cache_tables_[request_id]; + if (!tables.empty()) { + throw std::logic_error("HybridPrefixCache::PreparePagedCacheDeviceLoadBack: request already has tables"); + } + + for (const auto& [gid, host_pages] : host_hit.per_group_page_ids) { + if (host_pages.empty()) continue; + auto device_alloc_it = paged_cache_allocators_.find(gid); + if (device_alloc_it == paged_cache_allocators_.end() || device_alloc_it->second == nullptr) { + ReleaseRequest(request_id); + return {}; + } + OwnedPages device_pages = device_alloc_it->second->AcquireOwned(static_cast(host_pages.size())); + if (device_pages.Size() != static_cast(host_pages.size())) { + ReleaseRequest(request_id); + return {}; + } + std::int32_t base_logical_page = 0; + auto base_it = host_hit.per_group_base_logical_page.find(gid); + if (base_it != host_hit.per_group_base_logical_page.end()) { + base_logical_page = base_it->second; + } + + PagedCacheTransferPair transfer{gid, host_pages, device_pages.Ids()}; + auto [table_it, inserted] = tables.emplace(gid, PagedCacheGroupTable(device_alloc_it->second.get())); + if (!inserted) { + ReleaseRequest(request_id); + throw std::logic_error("HybridPrefixCache::PreparePagedCacheDeviceLoadBack: duplicate table for group " + + gid); + } + table_it->second.ImportPrefixOwned(std::move(device_pages), base_logical_page, host_hit.prefix_len_tokens); + transfers.push_back(std::move(transfer)); + } + if (transfers.empty()) { + ReleaseRequest(request_id); + } + return transfers; +} + +std::vector HybridPrefixCache::PreparePagedCacheHostWriteBack( + const std::vector& nodes) { + std::vector transfers; + if (!HasPagedCacheHostAdjunct()) return transfers; + + auto is_pinned = [](TreeNode* node) { + for (TreeNode* cur = node; cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { + if (!cur->OnHost()) return true; + if (cur->Host().RefCount() > 0) return true; + } + return false; + }; + std::vector prune_candidates; + std::size_t prune_cursor = 0; + bool prune_candidates_initialized = false; + auto prune_one = [&]() { + if (!prune_candidates_initialized) { + prune_candidates.reserve(paged_cache_host_snapshot_nodes_.size()); + for (TreeNode* candidate : paged_cache_host_snapshot_nodes_) { + if (candidate != nullptr && candidate->HasPagedCacheHostSnapshot()) { + prune_candidates.push_back(candidate); + } + } + std::sort(prune_candidates.begin(), prune_candidates.end(), [](TreeNode* a, TreeNode* b) { + if (a->Time() != b->Time()) return a->Time() < b->Time(); + if (a->DepthInTokens() != b->DepthInTokens()) return a->DepthInTokens() > b->DepthInTokens(); + return a->SeqId() < b->SeqId(); + }); + prune_candidates_initialized = true; + } + + while (prune_cursor < prune_candidates.size()) { + TreeNode* victim = prune_candidates[prune_cursor++]; + if (victim == nullptr || !victim->HasPagedCacheHostSnapshot() || is_pinned(victim)) continue; + + const std::size_t victim_depth = victim->DepthInTokens(); + auto primary = DetachPagedCacheHostSnapshotFromNode(victim); + (void)primary; + + std::vector descendants; + for (TreeNode* candidate : paged_cache_host_snapshot_nodes_) { + if (candidate == nullptr || candidate == victim || !candidate->HasPagedCacheHostSnapshot()) continue; + if (candidate->DepthInTokens() <= victim_depth) continue; + for (TreeNode* cur = candidate->Parent(); cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { + if (cur == victim) { + descendants.push_back(candidate); + break; + } + } + } + for (TreeNode* descendant : descendants) { + if (is_pinned(descendant)) continue; + auto cascaded = DetachPagedCacheHostSnapshotFromNode(descendant); + (void)cascaded; + } + return true; + } + return false; + }; + + for (TreeNode* node : nodes) { + if (node == nullptr || !node->HasPagedCacheSnapshot() || node->HasPagedCacheHostSnapshot() || + node->HasPagedCachePendingHostSnapshot() || !node->OnHost()) { + continue; + } + const PagedCacheSnapshot* device_snapshot = node->GetPagedCacheSnapshot(); + if (device_snapshot == nullptr || device_snapshot->groups.empty()) continue; + + auto pending = std::make_unique(); + pending->prefix_len_tokens = device_snapshot->prefix_len_tokens; + std::vector node_transfers; + bool ok = true; + for (const auto& [gid, device_group] : device_snapshot->groups) { + auto host_alloc_it = paged_cache_host_allocators_.find(gid); + if (host_alloc_it == paged_cache_host_allocators_.end() || host_alloc_it->second == nullptr) { + ok = false; + break; + } + const auto& src_ids = device_group.pages.Ids(); + if (src_ids.empty()) continue; + OwnedPages host_pages = host_alloc_it->second->AcquireOwned(static_cast(src_ids.size())); + while (host_pages.Size() != static_cast(src_ids.size()) && prune_one()) { + host_pages = host_alloc_it->second->AcquireOwned(static_cast(src_ids.size())); + } + if (host_pages.Size() != static_cast(src_ids.size())) { + ok = false; + break; + } + + PagedCacheGroupSnapshot host_group{}; + host_group.pages = std::move(host_pages); + host_group.base_logical_page = device_group.base_logical_page; + host_group.raw_token_cursor = device_group.raw_token_cursor; + host_group.sliding = device_group.sliding; + + node_transfers.push_back(PagedCacheTransferPair{gid, src_ids, host_group.pages.Ids()}); + pending->groups.emplace(gid, std::move(host_group)); + } + if (!ok || pending->groups.empty()) { + continue; + } + RefreshPagedCacheSnapshotCompleteness(*pending); + node->AttachPagedCachePendingHostSnapshot(std::move(pending)); + paged_cache_pending_host_snapshot_nodes_.insert(node); + transfers.insert(transfers.end(), std::make_move_iterator(node_transfers.begin()), + std::make_move_iterator(node_transfers.end())); + } + + return transfers; +} + +void HybridPrefixCache::OnPagedCacheHostWriteBackDone(const std::vector& nodes, bool success) { + if (!HasPagedCacheHostAdjunct()) return; + + for (TreeNode* node : nodes) { + if (node == nullptr || !node->HasPagedCachePendingHostSnapshot()) continue; + if (!success) { + DetachPagedCachePendingHostSnapshotFromNode(node); + continue; + } + paged_cache_pending_host_snapshot_nodes_.erase(node); + node->PromotePagedCachePendingHostSnapshot(); + if (node->HasPagedCacheHostSnapshot()) { + paged_cache_host_snapshot_nodes_.insert(node); + } + } +} + void HybridPrefixCache::PopulateOp(ForwardOperationBase& op_base) const { if (paged_cache_allocators_.empty()) return; auto req_it = request_paged_cache_tables_.find(op_base.request_id); @@ -1390,11 +1717,107 @@ bool HybridPrefixCache::admitPagedCacheChunk(const std::string& request_id, std: const PagedCacheAdmissionContext& context) { PagedCacheGroupAdmission admission = checkPagedCacheGroupAdmission( request_id, first_raw_position_of_op, target_raw_tokens_exclusive, simulated_free, paged_cache_hit, context); - const std::size_t prune_budget = paged_cache_snapshot_nodes_.size(); - for (std::size_t pruned = 0; !admission.ok && pruned < prune_budget; ++pruned) { + if (admission.ok) { + applyPagedCacheGroupAdmissionDebit(simulated_free, admission); + return true; + } + + std::vector candidates; + candidates.reserve(paged_cache_snapshot_nodes_.size()); + for (TreeNode* node : paged_cache_snapshot_nodes_) { + if (node != nullptr && node->HasPagedCacheSnapshot()) { + candidates.push_back(node); + } + } + std::sort(candidates.begin(), candidates.end(), [](TreeNode* a, TreeNode* b) { + if (a->Time() != b->Time()) return a->Time() < b->Time(); + if (a->DepthInTokens() != b->DepthInTokens()) return a->DepthInTokens() > b->DepthInTokens(); + return a->SeqId() < b->SeqId(); + }); + + std::unordered_map> borrowed_pages_by_group; + for (const auto& [_, tables] : request_paged_cache_tables_) { + for (const auto& [group_id, table] : tables) { + auto& borrowed = borrowed_pages_by_group[group_id]; + borrowed.insert(table.BorrowedPageIds().begin(), table.BorrowedPageIds().end()); + } + } + auto is_borrowed = [&](const TreeNode* node, std::optional family = std::nullopt) { + if (node == nullptr || node->GetPagedCacheSnapshot() == nullptr) return false; + for (const auto& [group_id, group_snapshot] : node->GetPagedCacheSnapshot()->groups) { + if (family.has_value()) { + auto allocator_it = paged_cache_allocators_.find(group_id); + if (allocator_it == paged_cache_allocators_.end() || allocator_it->second == nullptr || + allocator_it->second->Config().family != *family) { + continue; + } + } + auto borrowed_it = borrowed_pages_by_group.find(group_id); + if (borrowed_it == borrowed_pages_by_group.end()) continue; + for (std::int32_t page_id : group_snapshot.pages.Ids()) { + if (borrowed_it->second.contains(page_id)) return true; + } + } + return false; + }; + auto is_pinned = [&](TreeNode* node, std::optional family = std::nullopt) { + if (is_borrowed(node, family)) return true; + for (TreeNode* cur = node; cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { + if (cur->OnDevice() && cur->Device().RefCount() > 0) return true; + } + return false; + }; + + std::size_t state_cursor = 0; + std::size_t full_cursor = 0; + auto prune_state = [&]() { + while (state_cursor < candidates.size()) { + TreeNode* node = candidates[state_cursor++]; + if (node == nullptr || !node->HasPagedCacheSnapshot() || is_pinned(node, PagedCacheGroupFamily::State)) { + continue; + } + const auto* snapshot = node->GetPagedCacheSnapshot(); + if (!snapshot->IsCompleteFor(PagedCacheGroupFamily::State) && !snapshot->continuation_state_complete) { + continue; + } + if (DetachStateSnapshotFromNode(node)) return true; + } + return false; + }; + auto prune_full = [&]() { + while (full_cursor < candidates.size()) { + TreeNode* victim = candidates[full_cursor++]; + if (victim == nullptr || !victim->HasPagedCacheSnapshot() || is_pinned(victim)) continue; + + const std::size_t victim_depth = victim->DepthInTokens(); + auto primary = DetachPagedCacheSnapshotFromNode(victim); + (void)primary; + std::vector descendants; + for (TreeNode* node : paged_cache_snapshot_nodes_) { + if (node == nullptr || node == victim || !node->HasPagedCacheSnapshot()) continue; + if (node->DepthInTokens() <= victim_depth) continue; + for (TreeNode* cur = node->Parent(); cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { + if (cur == victim) { + descendants.push_back(node); + break; + } + } + } + for (TreeNode* descendant : descendants) { + if (is_pinned(descendant)) continue; + auto cascaded = DetachPagedCacheSnapshotFromNode(descendant); + (void)cascaded; + } + return true; + } + return false; + }; + + while (!admission.ok) { AdmissionFailureKind kind = ClassifyAdmissionFailure(admission); if (kind == AdmissionFailureKind::kNone) break; - if (!tryPrunePagedCacheSnapshot(kind)) break; + const bool pruned = kind == AdmissionFailureKind::kStateStarved ? prune_state() : prune_full(); + if (!pruned) break; refreshPagedCacheSimulatedFree(simulated_free); admission = checkPagedCacheGroupAdmission(request_id, first_raw_position_of_op, target_raw_tokens_exclusive, simulated_free, paged_cache_hit, context); @@ -1426,89 +1849,6 @@ bool HybridPrefixCache::DetachStateSnapshotFromNode(TreeNode* node) { return true; } -bool HybridPrefixCache::tryPrunePagedCacheSnapshot(AdmissionFailureKind kind) { - if (!HasPagedCacheAdjunct()) return false; - if (kind == AdmissionFailureKind::kNone) return false; - - auto is_pinned = [this](TreeNode* node, std::optional borrowed_family = std::nullopt) { - if (isPagedCacheSnapshotBorrowed(node, borrowed_family)) return true; - for (TreeNode* cur = node; cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { - if (!cur->OnDevice()) continue; - if (cur->Device().RefCount() > 0) return true; - } - return false; - }; - - // Sort once and share between branches: oldest first, then deepest within - // same Time(). Both try_state_only and try_full walk this same order. - std::vector candidates; - candidates.reserve(paged_cache_snapshot_nodes_.size()); - for (TreeNode* node : paged_cache_snapshot_nodes_) { - if (node == nullptr) continue; - if (!node->HasPagedCacheSnapshot()) continue; - candidates.push_back(node); - } - std::sort(candidates.begin(), candidates.end(), [](TreeNode* a, TreeNode* b) { - if (a->Time() != b->Time()) return a->Time() < b->Time(); - return a->DepthInTokens() > b->DepthInTokens(); - }); - - auto try_state_only = [&]() { - for (TreeNode* node : candidates) { - if (is_pinned(node, PagedCacheGroupFamily::State)) continue; - const auto* snap = node->GetPagedCacheSnapshot(); - if (snap == nullptr) continue; - if (!snap->IsCompleteFor(PagedCacheGroupFamily::State) && !snap->continuation_state_complete) continue; - if (DetachStateSnapshotFromNode(node)) return true; - } - return false; - }; - - auto try_full = [&]() { - TreeNode* victim = nullptr; - for (TreeNode* node : candidates) { - if (is_pinned(node)) continue; - victim = node; - break; - } - if (victim == nullptr) return false; - const std::size_t victim_depth = victim->DepthInTokens(); - auto primary = DetachPagedCacheSnapshotFromNode(victim); - (void)primary; - std::vector descendants; - for (TreeNode* node : paged_cache_snapshot_nodes_) { - if (node == nullptr || node == victim) continue; - if (!node->HasPagedCacheSnapshot()) continue; - if (node->DepthInTokens() <= victim_depth) continue; - for (TreeNode* cur = node->Parent(); cur != nullptr && !cur->IsRoot(); cur = cur->Parent()) { - if (cur == victim) { - descendants.push_back(node); - break; - } - } - } - for (TreeNode* d : descendants) { - if (is_pinned(d)) continue; - auto cascaded = DetachPagedCacheSnapshotFromNode(d); - (void)cascaded; - } - return true; - }; - - // kBothStarved: state-only cannot solve history shortage; go straight to - // full. The outer admit loop will re-classify if state still needs more. - switch (kind) { - case AdmissionFailureKind::kStateStarved: - return try_state_only(); - case AdmissionFailureKind::kHistoryStarved: - case AdmissionFailureKind::kBothStarved: - return try_full(); - case AdmissionFailureKind::kNone: - return false; - } - return false; -} - bool HybridPrefixCache::AdmitChunk(const std::string& request_id, std::int32_t first_raw_position_of_op, std::int32_t target_raw_tokens_exclusive, std::map& simulated_free, @@ -1535,8 +1875,9 @@ void HybridPrefixCache::CommitChunk(const std::string& request_id, TreeNode* ter if (lcm <= 0) return; const auto& required_groups = paged_cache_required_groups_; if (required_groups.empty()) return; + const auto& canonical_groups = paged_cache_history_groups_.empty() ? required_groups : paged_cache_history_groups_; - auto canonical_it = tables.find(required_groups.front()); + auto canonical_it = tables.find(canonical_groups.front()); if (canonical_it == tables.end()) return; std::int32_t last_committed = canonical_it->second.CommittedPrefixLenTokens(); @@ -1550,37 +1891,63 @@ void HybridPrefixCache::CommitChunk(const std::string& request_id, TreeNode* ter if (attach_node == nullptr) break; bool preflight_ok = true; + const char* preflight_fail_reason = "unknown"; + std::string preflight_fail_group; + std::int32_t fail_raw_per_page = -1; + std::int32_t fail_committed = -1; + std::int32_t fail_raw_cursor = -1; + std::int32_t fail_base = -1; + std::int32_t fail_borrowed = -1; + std::int32_t fail_owned = -1; + std::int32_t fail_committed_page = -1; + std::int32_t fail_owned_base_page = -1; + std::int32_t fail_pages_to_commit = -1; + + auto mark_preflight_failure = [&](const char* reason, const std::string& gid, + const PagedCacheGroupTable* table = nullptr, std::int32_t raw_per_page = -1) { + preflight_ok = false; + preflight_fail_reason = reason; + preflight_fail_group = gid; + fail_raw_per_page = raw_per_page; + if (table == nullptr) return; + fail_committed = table->CommittedPrefixLenTokens(); + fail_raw_cursor = table->RawTokenCursor(); + fail_base = table->BaseLogicalPage(); + fail_borrowed = table->BorrowedPagesCount(); + fail_owned = table->OwnedPagesCount(); + }; + for (const auto& gid : required_groups) { auto t_it = tables.find(gid); if (t_it == tables.end()) { - preflight_ok = false; + mark_preflight_failure("missing_table", gid); break; } const auto& table = t_it->second; const std::int32_t raw_per_page = table.RawTokensPerPage(); if (raw_per_page <= 0) { - preflight_ok = false; + mark_preflight_failure("invalid_raw_per_page", gid, &table, raw_per_page); break; } if (table.CommittedPrefixLenTokens() % raw_per_page != 0) { - preflight_ok = false; + mark_preflight_failure("committed_unaligned", gid, &table, raw_per_page); break; } if (target % raw_per_page != 0) { - preflight_ok = false; + mark_preflight_failure("target_unaligned", gid, &table, raw_per_page); break; } if (target <= table.CommittedPrefixLenTokens()) { - preflight_ok = false; + mark_preflight_failure("target_not_advanced", gid, &table, raw_per_page); break; } if (target > table.RawTokenCursor()) { - preflight_ok = false; + mark_preflight_failure("target_exceeds_raw_cursor", gid, &table, raw_per_page); break; } auto group_alloc_it = paged_cache_allocators_.find(gid); if (group_alloc_it == paged_cache_allocators_.end() || group_alloc_it->second == nullptr) { - preflight_ok = false; + mark_preflight_failure("missing_allocator", gid, &table, raw_per_page); break; } const auto& cfg = group_alloc_it->second->Config(); @@ -1588,17 +1955,30 @@ void HybridPrefixCache::CommitChunk(const std::string& request_id, TreeNode* ter const std::int32_t committed_page = table.CommittedPrefixLenTokens() / raw_per_page; const std::int32_t owned_base_page = table.BaseLogicalPage() + table.BorrowedPagesCount(); const std::int32_t pages_to_commit = (target - table.CommittedPrefixLenTokens()) / raw_per_page; - if (owned_base_page != committed_page || pages_to_commit > table.OwnedPagesCount()) { - preflight_ok = false; + if (owned_base_page != committed_page) { + mark_preflight_failure("history_owned_base_mismatch", gid, &table, raw_per_page); + fail_committed_page = committed_page; + fail_owned_base_page = owned_base_page; + fail_pages_to_commit = pages_to_commit; + break; + } + if (pages_to_commit > table.OwnedPagesCount()) { + mark_preflight_failure("history_owned_pages_short", gid, &table, raw_per_page); + fail_committed_page = committed_page; + fail_owned_base_page = owned_base_page; + fail_pages_to_commit = pages_to_commit; break; } } } if (!preflight_ok) { spdlog::warn( - "[HybridPrefixCache] CommitChunk: preflight failed for request {} at target " - "depth {}; leaving prior commits intact", - request_id, target); + "[HybridPrefixCache] CommitChunk: preflight failed request={} target={} chunk_depth={} lcm={} " + "reason={} group={} committed={} raw_cursor={} raw_per_page={} base={} borrowed={} owned={} " + "committed_page={} owned_base_page={} pages_to_commit={}; leaving prior commits intact", + request_id, target, chunk_depth, lcm, preflight_fail_reason, preflight_fail_group, fail_committed, + fail_raw_cursor, fail_raw_per_page, fail_base, fail_borrowed, fail_owned, fail_committed_page, + fail_owned_base_page, fail_pages_to_commit); break; } diff --git a/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.h b/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.h index 24f13a2f4..ded28d8d2 100644 --- a/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.h +++ b/tokenspeed-scheduler/csrc/resource/hybrid_prefix_cache/hybrid_prefix_cache.h @@ -67,21 +67,23 @@ class HybridPrefixCache { std::vector PrepareMambaDeviceLoadBack(const std::vector& nodes); void OnKVHostEvict(TreeNode* node); void OnKVDeviceDemote(TreeNode* node); - void OnMambaHostWriteBackDone(TreeNode* last_node); - void OnMambaHostWriteBackDone(const std::vector& nodes); + void OnMambaHostWriteBackDone(TreeNode* last_node, bool success = true); + void OnMambaHostWriteBackDone(const std::vector& nodes, bool success = true); void DemoteIdleMambaDeviceCopiesPresentOnHost(); // Takes ownership. Duplicate group_id throws std::invalid_argument. void RegisterPagedCacheGroup(std::unique_ptr allocator); + void RegisterPagedCacheHostGroup(std::unique_ptr allocator); - // History alignment is the LCM of RawTokensPerPage() over the History-family - // groups; state groups only need the trailing window. Sliding groups must - // have a window entry; full-history groups must not. + // History alignment is the LCM of RawTokensPerPage() over the required + // History-family groups. Continuation state groups only need their trailing + // window and are intentionally outside prefix-reuse admission. void EnablePagedCacheAdjunct(std::vector required_groups, std::unordered_map sliding_window_per_group); bool HasMambaAdjunct() const { return mamba_allocator_ != nullptr; } bool HasPagedCacheAdjunct() const { return paged_cache_history_alignment_tokens_ > 0; } + bool HasPagedCacheHostAdjunct() const { return !paged_cache_host_allocators_.empty(); } std::int32_t PagedCacheHistoryAlignmentTokens() const { return paged_cache_history_alignment_tokens_; } const std::vector& PagedCacheRequiredGroups() const { return paged_cache_required_groups_; } @@ -90,11 +92,15 @@ class HybridPrefixCache { std::int32_t PagedCacheGroupTotalPages(const std::string& group_id) const; std::int32_t PagedCacheGroupAvailablePages(const std::string& group_id) const; std::int64_t PagedCacheGroupFailedAllocCount(const std::string& group_id) const; + std::int32_t PagedCacheHostGroupTotalPages(const std::string& group_id) const; + std::int32_t PagedCacheHostGroupAvailablePages(const std::string& group_id) const; + std::int64_t PagedCacheHostGroupFailedAllocCount(const std::string& group_id) const; // Per-request introspection: unknown group_id throws; unknown request_id returns empty. std::vector GetRequestPagedCachePageIds(const std::string& request_id, const std::string& group_id) const; std::int32_t GetRequestPagedCacheBaseLogicalPage(const std::string& request_id, const std::string& group_id) const; + bool HasRequestPagedCacheTables(const std::string& request_id) const; // Unified paged-cache lifecycle surface used by the Scheduler. All methods // below are no-ops when no paged-cache groups are registered. @@ -117,6 +123,16 @@ class HybridPrefixCache { void RewindRequest(const std::string& request_id, std::int32_t accepted_raw_tokens, std::int32_t protected_tail_tokens = 0); + // Allocate request-local device pages for a host paged-cache hit and import + // them as owned prefix pages. Returns the HOST->DEVICE transfer specs. + std::vector PreparePagedCacheDeviceLoadBack(const std::string& request_id, + const MatchResult::PagedCache& host_hit); + + // Allocate pending host snapshots from device snapshots. Host metadata is + // invisible to Match() until OnPagedCacheHostWriteBackDone(..., true). + std::vector PreparePagedCacheHostWriteBack(const std::vector& nodes); + void OnPagedCacheHostWriteBackDone(const std::vector& nodes, bool success = true); + // Fill op.paged_cache_pages / op.paged_cache_page_base_offsets from the tables. void PopulateOp(ForwardOperationBase& op_base) const; @@ -142,9 +158,11 @@ class HybridPrefixCache { // null (defensive no-op). Accepts partial snapshots; the per-policy // "snapshot must be full" invariant is enforced upstream by CommitChunk. bool AttachPagedCacheSnapshotToNode(TreeNode* node, std::unique_ptr snapshot); + bool AttachPagedCacheHostSnapshotToNode(TreeNode* node, std::unique_ptr snapshot); // Drops `node` from the membership set, then detaches and returns the snapshot. std::unique_ptr DetachPagedCacheSnapshotFromNode(TreeNode* node); + std::unique_ptr DetachPagedCacheHostSnapshotFromNode(TreeNode* node); // Callback from KV prefix-cache eviction. void OnKVEvict(TreeNode* node); @@ -183,6 +201,7 @@ class HybridPrefixCache { // Drop only state-family groups from `node`'s snapshot; history portion // remains and the node stays registered. Returns true iff state groups removed. bool DetachStateSnapshotFromNode(TreeNode* node); + std::unique_ptr DetachPagedCachePendingHostSnapshotFromNode(TreeNode* node); void RefreshPagedCacheSnapshotCompleteness(PagedCacheSnapshot& snapshot) const; bool isPagedCacheSnapshotBorrowed(const TreeNode* node, @@ -193,11 +212,7 @@ class HybridPrefixCache { std::int32_t target); void augmentMatch(MatchResult& match) const; - void augmentMatchPagedCache(MatchResult& match) const; - - // Detach oldest evictable snapshot to free pool pages. State-only path is - // used only on kStateStarved; history/both go to full cascade. - bool tryPrunePagedCacheSnapshot(AdmissionFailureKind kind); + void augmentMatchPagedCache(MatchResult& match, MatchIntent intent) const; bool admitPagedCacheChunk(const std::string& request_id, std::int32_t first_raw_position_of_op, std::int32_t target_raw_tokens_exclusive, @@ -229,6 +244,7 @@ class HybridPrefixCache { // `paged_cache_history_alignment_tokens_ == 0` means adjunct disabled; tables still work. std::map> paged_cache_allocators_; + std::map> paged_cache_host_allocators_; std::unordered_map> request_paged_cache_tables_; std::int32_t paged_cache_history_alignment_tokens_{0}; std::vector paged_cache_required_groups_; @@ -242,8 +258,9 @@ class HybridPrefixCache { std::unordered_set paged_cache_state_group_set_; std::unordered_set paged_cache_continuation_state_group_set_; - // TODO(snapshot-lru-perf): O(N log N) per prune; swap in LRU index if profiling shows it matters. std::unordered_set paged_cache_snapshot_nodes_; + std::unordered_set paged_cache_host_snapshot_nodes_; + std::unordered_set paged_cache_pending_host_snapshot_nodes_; }; } // namespace tokenspeed diff --git a/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.cpp b/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.cpp index 6272e0fd8..b6c315580 100644 --- a/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.cpp +++ b/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.cpp @@ -191,6 +191,7 @@ MatchResult KVPrefixCache::RootMatch() const { TreeNode* root = tree_.Root(); const std::int32_t page_size = tree_.PageSize(); return MatchResult{ + .token_terminal = root, .device = {.last_node = root, .page_size = page_size}, .host = {.last_node = root, .page_size = page_size}, }; @@ -385,6 +386,42 @@ std::vector KVPrefixCache::ReleaseDeviceResourcesPresentOnHost(TreeNo return released; } +std::vector KVPrefixCache::ReleaseHostResources(const std::vector& nodes, + std::function on_release) { + std::vector candidates; + candidates.reserve(nodes.size()); + std::unordered_set seen; + for (TreeNode* node : nodes) { + if (node == nullptr || node->IsRoot() || !node->OnHost()) { + continue; + } + if (seen.insert(node).second) { + candidates.push_back(node); + } + } + + std::sort(candidates.begin(), candidates.end(), + [](const TreeNode* lhs, const TreeNode* rhs) { return lhs->DepthInTokens() > rhs->DepthInTokens(); }); + + std::vector released; + for (TreeNode* node : candidates) { + if (node == nullptr || !node->OnHost() || node->Host().RefCount() != 0) { + continue; + } + if (on_release) { + on_release(node); + } + host_.RemoveLeaf(node); + auto detached = node->DetachResource(); + if (detached == nullptr) { + continue; + } + released.push_back(node); + host_.UpdateLeaves(node->Parent()); + } + return released; +} + cache_op_id KVPrefixCache::AllocateCacheOpId() { return next_op_id_++; } diff --git a/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.h b/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.h index 1027d5e42..cc18a10ad 100644 --- a/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.h +++ b/tokenspeed-scheduler/csrc/resource/kv_prefix_cache/kv_prefix_cache.h @@ -68,6 +68,8 @@ class KVPrefixCache { std::vector ReleaseDeviceResourcesPresentOnHost(TreeNode* last_node, std::function on_release = {}); + std::vector ReleaseHostResources(const std::vector& nodes, + std::function on_release = {}); void EnqueueTransfer(TreeNode* last_node); diff --git a/tokenspeed-scheduler/csrc/resource/radix_tree/radix_tree.cpp b/tokenspeed-scheduler/csrc/resource/radix_tree/radix_tree.cpp index 68bbc6e53..8be3486ea 100644 --- a/tokenspeed-scheduler/csrc/resource/radix_tree/radix_tree.cpp +++ b/tokenspeed-scheduler/csrc/resource/radix_tree/radix_tree.cpp @@ -78,7 +78,9 @@ SplitResult RadixTree::splitChild(TreeNode* parent, const token_vec_t& child_key TreeNode* RadixTree::PruneEmptyByNode(TreeNode* node) { TreeNode* current = node; while (current != nullptr && !current->IsRoot()) { - if (current->NumChildren() != 0 || current->OnDevice() || current->OnHost()) { + if (current->NumChildren() != 0 || current->OnDevice() || current->OnHost() || + current->HasPagedCacheSnapshot() || current->HasPagedCacheHostSnapshot() || + current->HasPagedCachePendingHostSnapshot()) { break; } @@ -116,10 +118,6 @@ TreeNode* RadixTree::SplitAt(TreeNode* descendant, std::int32_t depth_in_tokens) return current; } if (depth_in_tokens > parent_depth && depth_in_tokens < this_depth) { - // Refuse to split a snapshot-bearing node (would dangle borrowed ids). - if (current->HasPagedCacheSnapshot()) { - return nullptr; - } TreeNode* parent = current->Parent(); const token_vec_t child_key = getFirstPage(current->Tokens(), page_size_); const std::size_t prefix_pages = static_cast((depth_in_tokens - parent_depth) / page_size_); @@ -140,6 +138,7 @@ WalkResult RadixTree::WalkDownUtilMismatch(token_slice aligned_tokens, TreeNode: .remaining_tokens = aligned_tokens, .match = { + .token_terminal = current, .device = {.last_node = current}, .host = {.last_node = current}, }, @@ -173,10 +172,6 @@ WalkResult RadixTree::WalkDownUtilMismatch(token_slice aligned_tokens, TreeNode: break; } if (matched_num_pages != static_cast(child->Tokens().size() / page_size_)) { - // Refuse to split a snapshot-bearing node; borrowed ids rely on it. - if (child->HasPagedCacheSnapshot()) { - break; - } SplitResult split = splitChild(current, walk_key_cache, matched_num_pages); child = split.prefix; } @@ -188,6 +183,7 @@ WalkResult RadixTree::WalkDownUtilMismatch(token_slice aligned_tokens, TreeNode: current = child; result.terminal = child; + result.match.token_terminal = child; result.remaining_tokens = result.remaining_tokens.subspan(matched_num_pages * page_size_); } diff --git a/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.cpp b/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.cpp index 56111d943..65701b8e9 100644 --- a/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.cpp +++ b/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "resource/allocator/owned_pages.h" @@ -40,9 +41,12 @@ void TreeNode::AddChild(const token_vec_t& key, std::unique_ptr&& chil if (child == nullptr) [[unlikely]] { return; } - child->parent_ = this; - child->depth_in_tokens_ = depth_in_tokens_ + child->tokens_.size(); - children_[key] = std::move(child); + auto [it, inserted] = children_.try_emplace(key, std::move(child)); + if (!inserted) { + throw std::logic_error("TreeNode::AddChild: duplicate child key"); + } + it->second->parent_ = this; + it->second->depth_in_tokens_ = depth_in_tokens_ + it->second->tokens_.size(); } std::unique_ptr TreeNode::RemoveChild(const token_vec_t& key) { @@ -96,12 +100,8 @@ void TreeNode::SplitSelfInto(TreeNode& prefix, std::size_t prefix_pages, std::in std::int32_t ref_count = host_resource_->RefCount(); prefix.AttachResource(std::make_unique(host_resource_->SplitFirst(prefix_pages), ref_count)); } - // Mamba stays in suffix. - // Invariant: snapshot-bearing nodes are never split (RadixTree refuses). - // A split here would dangle borrowed ids in active requests. - _assert(paged_cache_snapshot_ == nullptr, - "TreeNode::SplitSelfInto called on a node with an attached paged-cache snapshot; " - "splitting would invalidate borrowed page id references in active requests"); + // Mamba and paged-cache snapshots stay on this suffix object. Its pointer and + // absolute depth are unchanged, so adjunct membership and borrowed ids remain valid. } void TreeNode::AttachPagedCacheSnapshot(std::unique_ptr snapshot) { @@ -112,6 +112,29 @@ std::unique_ptr TreeNode::DetachPagedCacheSnapshot() { return std::move(paged_cache_snapshot_); } +void TreeNode::AttachPagedCacheHostSnapshot(std::unique_ptr snapshot) { + paged_cache_host_snapshot_ = std::move(snapshot); +} + +std::unique_ptr TreeNode::DetachPagedCacheHostSnapshot() { + return std::move(paged_cache_host_snapshot_); +} + +void TreeNode::AttachPagedCachePendingHostSnapshot(std::unique_ptr snapshot) { + paged_cache_pending_host_snapshot_ = std::move(snapshot); +} + +std::unique_ptr TreeNode::DetachPagedCachePendingHostSnapshot() { + return std::move(paged_cache_pending_host_snapshot_); +} + +void TreeNode::PromotePagedCachePendingHostSnapshot() { + if (paged_cache_pending_host_snapshot_ == nullptr) { + return; + } + paged_cache_host_snapshot_ = std::move(paged_cache_pending_host_snapshot_); +} + void TreeNode::SetPersisted(bool persisted) { storage_persisted_ = persisted; } diff --git a/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.h b/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.h index cc5b66a15..ae8a9aa99 100644 --- a/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.h +++ b/tokenspeed-scheduler/csrc/resource/radix_tree/tree_node.h @@ -122,6 +122,9 @@ class TreeNode { // snapshot itself (see `PagedCacheSnapshot::IsCompleteFor`). bool HasPagedCacheSnapshot() const { return paged_cache_snapshot_ != nullptr; } const PagedCacheSnapshot* GetPagedCacheSnapshot() const { return paged_cache_snapshot_.get(); } + bool HasPagedCacheHostSnapshot() const { return paged_cache_host_snapshot_ != nullptr; } + bool HasPagedCachePendingHostSnapshot() const { return paged_cache_pending_host_snapshot_ != nullptr; } + const PagedCacheSnapshot* GetPagedCacheHostSnapshot() const { return paged_cache_host_snapshot_.get(); } std::optional CacheOpId() const; @@ -141,6 +144,12 @@ class TreeNode { void AttachPagedCacheSnapshot(std::unique_ptr snapshot); std::unique_ptr DetachPagedCacheSnapshot(); PagedCacheSnapshot* GetPagedCacheSnapshotMut() { return paged_cache_snapshot_.get(); } + void AttachPagedCacheHostSnapshot(std::unique_ptr snapshot); + std::unique_ptr DetachPagedCacheHostSnapshot(); + void AttachPagedCachePendingHostSnapshot(std::unique_ptr snapshot); + std::unique_ptr DetachPagedCachePendingHostSnapshot(); + void PromotePagedCachePendingHostSnapshot(); + PagedCacheSnapshot* GetPagedCacheHostSnapshotMut() { return paged_cache_host_snapshot_.get(); } private: TreeNode* parent_{}; @@ -157,6 +166,8 @@ class TreeNode { std::unique_ptr mamba_slot_{}; std::unique_ptr mamba_host_slot_{}; std::unique_ptr paged_cache_snapshot_{}; + std::unique_ptr paged_cache_host_snapshot_{}; + std::unique_ptr paged_cache_pending_host_snapshot_{}; static std::atomic next_seq_id_; }; diff --git a/tokenspeed-scheduler/csrc/resource/types.h b/tokenspeed-scheduler/csrc/resource/types.h index 92e939bd6..6f1365172 100644 --- a/tokenspeed-scheduler/csrc/resource/types.h +++ b/tokenspeed-scheduler/csrc/resource/types.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -52,6 +53,9 @@ using DeviceNodeRef = NodeRef; using HostNodeRef = NodeRef; struct MatchResult { + // Deepest token-radix node reached before tier-specific resource capping. + TreeNode* token_terminal{nullptr}; + struct Device { TreeNode* last_node; std::int32_t page_size{0}; @@ -85,6 +89,9 @@ struct MatchResult { std::map> per_group_page_ids; std::map per_group_base_logical_page; } paged_cache; + // Host-tier paged-cache hit. These page ids are host page ids and must be + // materialized into device pages before a forward op is built. + PagedCache paged_cache_host; }; struct InsertResult { @@ -108,6 +115,8 @@ struct CacheOpSpec { std::string request_id; TreeNode* last_node{nullptr}; std::vector nodes; + std::vector paged_cache_nodes; + std::unique_ptr paged_loadback_host_pin; CacheOpSpec(); ~CacheOpSpec(); diff --git a/tokenspeed-scheduler/csrc/scheduler/execution_plan.h b/tokenspeed-scheduler/csrc/scheduler/execution_plan.h index 581a1d653..c22fdb545 100644 --- a/tokenspeed-scheduler/csrc/scheduler/execution_plan.h +++ b/tokenspeed-scheduler/csrc/scheduler/execution_plan.h @@ -20,6 +20,8 @@ #pragma once +#include +#include #include #include @@ -27,11 +29,16 @@ namespace tokenspeed { +struct SchedulerAbort { + std::string request_id; + std::string message; +}; + class ExecutionPlan { public: template ExecutionPlan& With(OperationType operation) { - operations_.emplace_back(operation); + operations_.emplace_back(std::move(operation)); return *this; } @@ -43,7 +50,13 @@ class ExecutionPlan { return *this; } + ExecutionPlan& WithSchedulerAborts(std::vector aborts) { + scheduler_aborts_ = std::move(aborts); + return *this; + } + const std::vector& Operations() const { return operations_; } + const std::vector& SchedulerAborts() const { return scheduler_aborts_; } // Flat KV-cache: requests terminalized this round as OOM -- the pool was wedged by // unretractable mid-prefill holders (possibly the request itself, or a mutual wedge) @@ -52,6 +65,7 @@ class ExecutionPlan { private: std::vector operations_; + std::vector scheduler_aborts_; }; } // namespace tokenspeed diff --git a/tokenspeed-scheduler/csrc/scheduler/operations/cache.h b/tokenspeed-scheduler/csrc/scheduler/operations/cache.h index a1876ca57..08b2285d6 100644 --- a/tokenspeed-scheduler/csrc/scheduler/operations/cache.h +++ b/tokenspeed-scheduler/csrc/scheduler/operations/cache.h @@ -90,9 +90,20 @@ struct TransferPairHash { } }; +struct PagedCacheTransferPair { + std::string group_id; + std::vector src_pages; + std::vector dst_pages; + + bool operator==(const PagedCacheTransferPair& other) const { + return group_id == other.group_id && src_pages == other.src_pages && dst_pages == other.dst_pages; + } +}; + struct WriteBackOperation { cache_op_id op_id{0}; - std::vector transfers; // DEVICE→HOST by cache kind. + std::vector transfers; // DEVICE→HOST by cache kind. + std::vector paged_cache_transfers; // DEVICE→HOST by paged-cache group. bool is_retract{false}; WriteBackOperation() = default; @@ -101,6 +112,12 @@ struct WriteBackOperation { : op_id{op_id}, transfers{ToTransferPairs(CacheKind::kKV, pages_to_transfer)}, is_retract{is_retract} {} WriteBackOperation(cache_op_id op_id, std::vector transfers, bool is_retract = false) : op_id{op_id}, transfers{std::move(transfers)}, is_retract{is_retract} {} + WriteBackOperation(cache_op_id op_id, std::vector transfers, + std::vector paged_cache_transfers, bool is_retract = false) + : op_id{op_id}, + transfers{std::move(transfers)}, + paged_cache_transfers{std::move(paged_cache_transfers)}, + is_retract{is_retract} {} }; struct FlatWriteBackOperation { @@ -111,11 +128,13 @@ struct FlatWriteBackOperation { // Generic view keyed by CacheKindName(kind), currently "kv" and "mamba". std::map>> src_pages_by_kind; std::map>> dst_pages_by_kind; + // Indexed by op_ids; each item preserves the op's paged-cache group transfer specs. + std::vector> paged_cache_transfers; std::vector is_retract; - explicit FlatWriteBackOperation(const std::vector& ops) { + explicit FlatWriteBackOperation(std::vector ops) { std::unordered_set seen; - for (const auto& op : ops) { + for (auto& op : ops) { std::map> src_this_op; std::map> dst_this_op; src_this_op[CacheKindName(CacheKind::kKV)]; @@ -140,6 +159,7 @@ struct FlatWriteBackOperation { for (auto& [kind, pages] : dst_this_op) { dst_pages_by_kind[kind].push_back(std::move(pages)); } + paged_cache_transfers.push_back(std::move(op.paged_cache_transfers)); is_retract.push_back(op.is_retract); } } @@ -147,13 +167,17 @@ struct FlatWriteBackOperation { struct LoadBackOperation { cache_op_id op_id{0}; - std::vector transfers; // HOST→DEVICE by cache kind. + std::vector transfers; // HOST→DEVICE by cache kind. + std::vector paged_cache_transfers; // HOST→DEVICE by paged-cache group. LoadBackOperation() = default; LoadBackOperation(cache_op_id op_id, std::vector> pages_to_transfer) : op_id{op_id}, transfers{ToTransferPairs(CacheKind::kKV, pages_to_transfer)} {} LoadBackOperation(cache_op_id op_id, std::vector transfers) : op_id{op_id}, transfers{std::move(transfers)} {} + LoadBackOperation(cache_op_id op_id, std::vector transfers, + std::vector paged_cache_transfers) + : op_id{op_id}, transfers{std::move(transfers)}, paged_cache_transfers{std::move(paged_cache_transfers)} {} }; struct FlatLoadBackOperation { @@ -164,10 +188,12 @@ struct FlatLoadBackOperation { // Generic view keyed by CacheKindName(kind), currently "kv" and "mamba". std::map>> src_pages_by_kind; std::map>> dst_pages_by_kind; + // Indexed by op_ids; each item preserves the op's paged-cache group transfer specs. + std::vector> paged_cache_transfers; - explicit FlatLoadBackOperation(const std::vector& ops) { + explicit FlatLoadBackOperation(std::vector ops) { std::unordered_set seen; - for (const auto& op : ops) { + for (auto& op : ops) { std::map> src_this_op; std::map> dst_this_op; src_this_op[CacheKindName(CacheKind::kKV)]; @@ -192,6 +218,7 @@ struct FlatLoadBackOperation { for (auto& [kind, pages] : dst_this_op) { dst_pages_by_kind[kind].push_back(std::move(pages)); } + paged_cache_transfers.push_back(std::move(op.paged_cache_transfers)); } } }; diff --git a/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp b/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp index bf643e331..7c052a1df 100644 --- a/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp +++ b/tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp @@ -100,6 +100,23 @@ static void MaybeFillFlatBlockTables(Op& op, Request* request, std::span +TreeNode* LastNodeWithResourceOrRoot(TreeNode* node) { + for (TreeNode* candidate = node; candidate != nullptr; candidate = candidate->Parent()) { + if (candidate->IsRoot()) return candidate; + if constexpr (RType == ResourceType::Device) { + if (candidate->OnDevice()) { + return candidate; + } + } else { + if (candidate->OnHost()) { + return candidate; + } + } + } + return nullptr; +} + } // namespace #if TOKENSPEED_FLAT_KVCACHE @@ -275,17 +292,29 @@ std::optional Scheduler::schedulePrefillFir std::int32_t unscheduled = 0; std::vector loadback_diff; std::vector mamba_loadback_nodes; + std::vector paged_cache_loadback_transfers; const std::int32_t device_matched = match_result.device.DepthInPage(); const std::int32_t host_matched = match_result.host.DepthInPage(); + const bool has_paged_cache = hybrid_prefix_cache_ && hybrid_prefix_cache_->HasPagedCacheAdjunct(); + const std::int32_t paged_device_matched = match_result.paged_cache.prefix_len_tokens; + const std::int32_t paged_host_matched = disable_l2_cache ? 0 : match_result.paged_cache_host.prefix_len_tokens; + const bool use_paged_host_hit = has_paged_cache && paged_host_matched > paged_device_matched; + std::int32_t matched_prefix_len_tokens = 0; if (disable_l2_cache) { - unscheduled = request->PrefillSize() - device_matched * config_.block_size; + matched_prefix_len_tokens = device_matched * config_.block_size; + unscheduled = request->PrefillSize() - matched_prefix_len_tokens; } else { loadback_diff = match_result.NodesWithout(); if (host_matched > device_matched) { loadback_tokens = config_.block_size * (host_matched - device_matched); } - unscheduled = request->PrefillSize() - std::max(device_matched, host_matched) * config_.block_size; + matched_prefix_len_tokens = has_paged_cache ? std::max(paged_device_matched, paged_host_matched) + : std::max(device_matched, host_matched) * config_.block_size; + unscheduled = request->PrefillSize() - matched_prefix_len_tokens; + } + if (unscheduled < 0) { + unscheduled = 0; } std::int32_t tokens_this_round = std::min(remaining, unscheduled); @@ -299,7 +328,8 @@ std::optional Scheduler::schedulePrefillFir std::int32_t num_tokens = loadback_tokens + tokens_this_round + decode_input_tokens; std::int32_t device_pages_needed = (num_tokens + config_.block_size - 1) / config_.block_size; - std::unique_ptr temp_lock = std::make_unique(match_result.device.last_node); + std::unique_ptr temp_lock = std::make_unique( + LastNodeWithResourceOrRoot(match_result.device.last_node)); // Evict unlocked prefix-cache nodes before allocating request-local pages. if (!(kv_prefix_cache_.EnsureCapacityByEvict(device_pages_needed))) { @@ -344,12 +374,22 @@ std::optional Scheduler::schedulePrefillFir return {}; } - const std::int32_t first_pos = request->PrefillSize() - unscheduled; + const std::int32_t first_pos = disable_l2_cache ? request->PrefillSize() - unscheduled : matched_prefix_len_tokens; const std::int32_t target = first_pos + tokens_this_round; + const MatchResult::PagedCache empty_paged_hit{}; + const MatchResult::PagedCache& paged_hit_for_admission = + use_paged_host_hit ? empty_paged_hit : match_result.paged_cache; if (hybrid_prefix_cache_ && - !hybrid_prefix_cache_->AdmitChunk(request->Id(), first_pos, target, simulated_free, match_result.paged_cache)) { + !hybrid_prefix_cache_->AdmitChunk(request->Id(), first_pos, target, simulated_free, paged_hit_for_admission)) { return {}; } + if (use_paged_host_hit) { + paged_cache_loadback_transfers = + hybrid_prefix_cache_->PreparePagedCacheDeviceLoadBack(request->Id(), match_result.paged_cache_host); + if (paged_cache_loadback_transfers.empty()) { + return {}; + } + } if (needs_mamba_loadback) { hybrid_prefix_cache_->PrepareMambaDeviceLoadBack(mamba_loadback_nodes); TreeNode* mamba_node = hybrid_prefix_cache_->FindLastMambaNode(match_result.host.last_node); @@ -372,7 +412,7 @@ std::optional Scheduler::schedulePrefillFir decode_input_tokens, &device_allocator_, &req_pool_allocator_, - match_result, + std::move(match_result), config_.role, &kv_prefix_cache_, disable_l2_cache, @@ -386,6 +426,8 @@ std::optional Scheduler::schedulePrefillFir std::move(flat_match.host), std::move(flat_match.ext_hashes), #endif + std::move(paged_cache_loadback_transfers), + matched_prefix_len_tokens, }; } @@ -444,8 +486,8 @@ std::optional Scheduler::schedulePrefill( }; } -std::optional Scheduler::scheduleDecode(Request* request, - std::map& simulated_free) { +Scheduler::ScheduleAttempt Scheduler::scheduleDecode( + Request* request, std::map& simulated_free) { std::int32_t tail_available = request->TailPageAvailableTokens(); std::int32_t extra_tokens = std::max(0, request->GetReserveNumTokensInNextScheduleEvent() - tail_available); std::int32_t pages_needed = (extra_tokens + config_.block_size - 1) / config_.block_size; @@ -478,19 +520,21 @@ std::optional Scheduler::scheduleDecode(Request* reque } if (!hybrid_prefix_cache_->AdmitChunk(request->Id(), first_pos, target, simulated_free, {}, commit_target, commit_token_pages)) { - return {}; + return {.event = std::nullopt, .failure = ScheduleFailure::kPagedCache}; } } - return fsm::ScheduleDecodeEvent{config_.decode_input_tokens, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr + return {.event = fsm::ScheduleDecodeEvent{config_.decode_input_tokens, + hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr #if TOKENSPEED_FLAT_KVCACHE - , - &coordinator_ + , + &coordinator_ #endif - }; + }, + .failure = ScheduleFailure::kNone}; } -std::optional Scheduler::scheduleDecodeFromRetracted( +Scheduler::ScheduleAttempt Scheduler::scheduleDecodeFromRetracted( Request* request, std::map& simulated_free) { if (req_pool_allocator_.AvailableSlots() == 0) return {}; @@ -500,6 +544,12 @@ std::optional Scheduler::scheduleDecodeFr : kv_prefix_cache_.Match(request->GetFullPagedTokens(true), MatchIntent::StateRecovery); std::vector loadback_diff = match_result.NodesWithout(); std::vector mamba_loadback_nodes; + std::vector paged_cache_loadback_transfers; + const bool has_paged_cache = hybrid_prefix_cache_ && hybrid_prefix_cache_->HasPagedCacheAdjunct(); + const std::int32_t paged_device_matched = match_result.paged_cache.prefix_len_tokens; + const std::int32_t paged_host_matched = + config_.disable_l2_cache ? 0 : match_result.paged_cache_host.prefix_len_tokens; + const bool use_paged_host_hit = has_paged_cache && paged_host_matched > paged_device_matched; TreeNode* mamba_recovery_node = nullptr; bool needs_mamba_loadback = false; if (hybrid_prefix_cache_ && mamba_allocator_) { @@ -512,13 +562,7 @@ std::optional Scheduler::scheduleDecodeFr } } if (mamba_recovery_node == nullptr) { - spdlog::warn("[Scheduler] Retracted request {} lost tree-owned Mamba state, aborting request", - request->Id()); - request->Apply(fsm::AbortEvent{ -#if TOKENSPEED_FLAT_KVCACHE - &coordinator_ -#endif - }); + abortRequest(request, "Retracted request lost tree-owned Mamba state"); return {}; } if (!needs_mamba_loadback) { @@ -536,7 +580,8 @@ std::optional Scheduler::scheduleDecodeFr } std::int32_t device_pages_needed = (num_tokens + config_.block_size - 1) / config_.block_size; - std::unique_ptr temp_lock = std::make_unique(match_result.device.last_node); + std::unique_ptr temp_lock = std::make_unique( + LastNodeWithResourceOrRoot(match_result.device.last_node)); if (!kv_prefix_cache_.EnsureCapacityByEvict(device_pages_needed)) { return {}; } @@ -552,9 +597,21 @@ std::optional Scheduler::scheduleDecodeFr const std::int32_t target = std::max( request->TokenSize(), DecodePagedCacheReservationEnd(first_pos, config_.decode_input_tokens, config_.overlap_schedule_depth)); + const MatchResult::PagedCache empty_paged_hit{}; + const MatchResult::PagedCache& paged_hit_for_admission = + use_paged_host_hit ? empty_paged_hit : match_result.paged_cache; if (hybrid_prefix_cache_ && - !hybrid_prefix_cache_->AdmitChunk(request->Id(), first_pos, target, simulated_free, match_result.paged_cache)) { - return {}; + !hybrid_prefix_cache_->AdmitChunk(request->Id(), first_pos, target, simulated_free, paged_hit_for_admission)) { + return {.event = std::nullopt, .failure = ScheduleFailure::kPagedCache}; + } + if (use_paged_host_hit) { + if (!hybrid_prefix_cache_->HasRequestPagedCacheTables(request->Id())) { + paged_cache_loadback_transfers = + hybrid_prefix_cache_->PreparePagedCacheDeviceLoadBack(request->Id(), match_result.paged_cache_host); + if (paged_cache_loadback_transfers.empty()) { + return {}; + } + } } if (needs_mamba_loadback) { hybrid_prefix_cache_->PrepareMambaDeviceLoadBack(mamba_loadback_nodes); @@ -566,16 +623,19 @@ std::optional Scheduler::scheduleDecodeFr return {}; } - return fsm::ScheduleDecodeFromRetractedEvent{ - config_.decode_input_tokens, - &device_allocator_, - &req_pool_allocator_, - &kv_prefix_cache_, - std::move(match_result), - loadback_diff, - mamba_allocator_ ? &*mamba_allocator_ : nullptr, - std::move(mamba_loadback_nodes), - }; + return {.event = + fsm::ScheduleDecodeFromRetractedEvent{ + config_.decode_input_tokens, + &device_allocator_, + &req_pool_allocator_, + &kv_prefix_cache_, + std::move(match_result), + std::move(loadback_diff), + mamba_allocator_ ? &*mamba_allocator_ : nullptr, + std::move(mamba_loadback_nodes), + std::move(paged_cache_loadback_transfers), + }, + .failure = ScheduleFailure::kNone}; } std::optional Scheduler::scheduleRetract(Request* request) { @@ -610,12 +670,14 @@ std::optional Scheduler::scheduleRetract(Request* req if (!kv_prefix_cache_.EnsureCapacityByEvict(host_pages_needed)) { return {}; } - return fsm::ScheduleRetractEvent{&kv_prefix_cache_, &host_allocator_, match_result, + return fsm::ScheduleRetractEvent{&kv_prefix_cache_, &host_allocator_, std::move(match_result), hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr}; } -LoadBackOperation GenerateLoadBackOp(const std::vector& diff, const std::vector& mamba_nodes, - cache_op_id op_id) { +LoadBackOperation Scheduler::newLoadBackOperation(const std::string& request_id, const std::vector& diff, + const std::vector& mamba_nodes, + std::vector paged_cache_transfers, + TreeNode* paged_cache_host_node) { std::vector transfers; for (TreeNode* node : diff) { @@ -630,7 +692,18 @@ LoadBackOperation GenerateLoadBackOp(const std::vector& diff, const s transfers.push_back(TransferPair{CacheKind::kMamba, node->MambaHostSlotIndex(), node->MambaSlotIndex()}); } } - return LoadBackOperation{op_id, std::move(transfers)}; + + cache_op_id op_id = kv_prefix_cache_.AllocateCacheOpId(); + if (!paged_cache_transfers.empty()) { + if (paged_cache_host_node == nullptr || !paged_cache_host_node->OnHost()) { + throw std::logic_error("paged-cache loadback requires a host-resident snapshot node"); + } + CacheOpSpec spec; + spec.request_id = request_id; + spec.paged_loadback_host_pin = std::make_unique(paged_cache_host_node); + cache_op_tracker_[op_id] = std::move(spec); + } + return LoadBackOperation{op_id, std::move(transfers), std::move(paged_cache_transfers)}; } std::optional Scheduler::applyEventAndGenerateOp(Request* request, @@ -638,7 +711,8 @@ std::optional Scheduler::applyEventAndGenerateOp(Request* re request->Apply(std::move(event)); const auto& pages_to_transfer = request->GetPagesToTransfer(); - if (pages_to_transfer.empty()) { + const auto& paged_cache_transfers = request->GetPagedCacheWriteBackTransfers(); + if (pages_to_transfer.empty() && paged_cache_transfers.empty()) { // No copy needed; advance Retracting to Retracted without an op_id. request->Apply( fsm::WriteBackDoneEvent{&kv_prefix_cache_, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr}); @@ -647,9 +721,65 @@ std::optional Scheduler::applyEventAndGenerateOp(Request* re cache_op_id op_id = kv_prefix_cache_.AllocateCacheOpId(); CacheOpSpec spec; spec.request_id = request->Id(); + spec.paged_cache_nodes = request->GetPagedCacheWriteBackNodes(); cache_op_tracker_[op_id] = std::move(spec); - return WriteBackOperation{op_id, std::vector(pages_to_transfer.begin(), pages_to_transfer.end()), - true}; + return WriteBackOperation{ + op_id, std::vector(pages_to_transfer.begin(), pages_to_transfer.end()), + std::vector(paged_cache_transfers.begin(), paged_cache_transfers.end()), true}; +} + +bool Scheduler::hasInFlightCacheOp(const std::string& request_id) const { + return std::any_of(cache_op_tracker_.begin(), cache_op_tracker_.end(), + [&](const auto& item) { return item.second.request_id == request_id; }); +} + +void Scheduler::deferAbort(const std::string& request_id, bool discard_writeback, std::string scheduler_message) { + auto& deferred = deferred_aborts_[request_id]; + if (discard_writeback) { + // A runtime-originated abort is already terminal to the client. It wins over + // an internal abort that had not yet become observable. + deferred.discard_writeback = true; + deferred.scheduler_message.clear(); + } else if (!deferred.discard_writeback && deferred.scheduler_message.empty()) { + deferred.scheduler_message = std::move(scheduler_message); + } +} + +void Scheduler::tryFinalizeDeferredAbort(const std::string& request_id) { + if (deferred_aborts_.empty()) return; + + auto deferred_it = deferred_aborts_.find(request_id); + if (deferred_it == deferred_aborts_.end() || pending_forward_results_.contains(request_id) || + hasInFlightCacheOp(request_id)) { + return; + } + + if (Request* request = find_request(request_id); request != nullptr && !request->Is()) { + request->Apply(fsm::AbortEvent{&kv_prefix_cache_, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr +#if TOKENSPEED_FLAT_KVCACHE + , + &coordinator_ +#endif + }); + } + if (hybrid_prefix_cache_) { + hybrid_prefix_cache_->ReleaseRequest(request_id); + } + + std::string scheduler_message = std::move(deferred_it->second.scheduler_message); + deferred_aborts_.erase(deferred_it); + if (!scheduler_message.empty()) { + spdlog::warn("[Scheduler] Aborting request {}: {}", request_id, scheduler_message); + scheduler_aborts_.push_back(SchedulerAbort{request_id, std::move(scheduler_message)}); + } +} + +void Scheduler::abortRequest(Request* request, std::string message) { + if (request == nullptr || request->Is()) return; + + const std::string request_id = request->Id(); + deferAbort(request_id, /*discard_writeback=*/false, std::move(message)); + tryFinalizeDeferredAbort(request_id); } std::optional Scheduler::newRetractOperation(Request* retract_request) { @@ -658,13 +788,7 @@ std::optional Scheduler::newRetractOperation(Request* retrac return std::move(*op); } } else { - spdlog::warn("[Scheduler] Retract failed for request {}: host capacity exhausted, aborting request", - retract_request->Id()); - retract_request->Apply(fsm::AbortEvent{ -#if TOKENSPEED_FLAT_KVCACHE - &coordinator_ -#endif - }); + abortRequest(retract_request, "Retract failed because host cache capacity is exhausted"); } return std::nullopt; } @@ -833,12 +957,13 @@ DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Schedu return op; } -DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::ScheduleDecodeFromRetractedEvent event) { - const std::int32_t mamba_cow_src_index = event.GetMatchResult().mamba_cow_src_index; +DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::ScheduleDecodeFromRetractedEvent& event) { + const MatchResult& match = event.GetMatchResult(); + const std::int32_t mamba_cow_src_index = match.mamba_cow_src_index; #if !TOKENSPEED_FLAT_KVCACHE - auto paged_cache_hit = event.GetMatchResult().paged_cache; + const bool has_paged_cache_loadback = !event.GetPagedCacheLoadbackTransfers().empty(); #endif - request->Apply(std::move(event)); + request->Apply(event); if (!request->Is()) { throw std::logic_error( "Scheduler::applyEventAndGenerateOp: expected state=Decoding after loadback recovery; got state=" + @@ -871,10 +996,13 @@ DecodeOperation Scheduler::applyEventAndGenerateOp(Request* request, fsm::Schedu const std::int32_t target = std::max( request->TokenSize(), DecodePagedCacheReservationEnd(op.hist_token_len, op.input_length, config_.overlap_schedule_depth)); + if (!has_paged_cache_loadback && !hybrid_prefix_cache_->HasRequestPagedCacheTables(op.request_id)) { + hybrid_prefix_cache_->ReleaseRequest(op.request_id); + } // Preserve the existing table across retraction. Its request-local // tail contains state after the last published prefix checkpoint and // cannot be reconstructed by importing that older snapshot alone. - hybrid_prefix_cache_->AcquireForRequest(op.request_id, op.hist_token_len, target, paged_cache_hit); + hybrid_prefix_cache_->AcquireForRequest(op.request_id, op.hist_token_len, target, match.paged_cache); hybrid_prefix_cache_->PopulateOp(op); } #endif @@ -914,17 +1042,22 @@ Scheduler::newForwardOperation(std::vector candidates) { } ops.push_back(std::move(op)); }; -#if TOKENSPEED_FLAT_KVCACHE // Mid-prefill chunk ops emit no ExtendResult; only decode and prefill-completing ops owe one. auto note_result_owed = [&](Request* request) { if (!request->Is()) { ++pending_forward_results_[request->Id()]; } }; -#else - auto note_result_owed = [](Request*) {}; -#endif std::vector loadback_ops; + std::vector paged_cache_blocked; + bool all_decode_failures_are_paged_cache = true; + auto record_decode_failure = [&](Request* request, ScheduleFailure failure) { + if (request->Is() || failure != ScheduleFailure::kPagedCache) { + all_decode_failures_are_paged_cache = false; + return; + } + paged_cache_blocked.push_back(request); + }; auto simulated_free = hybrid_prefix_cache_ ? hybrid_prefix_cache_->InitialSimulatedFree() : std::map{}; for (Request* request : candidates) { @@ -942,36 +1075,52 @@ Scheduler::newForwardOperation(std::vector candidates) { if (auto ev = schedulePrefillFirstChunk(request, token_budget, decode_input_tokens, config_.disable_l2_cache, simulated_free)) { - std::vector loadback_diff = ev->GetLoadbackDiff(); - std::vector mamba_loadback_nodes = ev->GetMambaLoadbackNodes(); + TreeNode* paged_cache_host_node = ev->GetMatchResult().paged_cache_host.last_node; + std::vector loadback_diff = ev->TakeLoadbackDiff(); + std::vector mamba_loadback_nodes = ev->TakeMambaLoadbackNodes(); + std::vector paged_cache_loadback_transfers = + ev->TakePagedCacheLoadbackTransfers(); push_op(applyEventAndGenerateOp(request, std::move(*ev), loadback_ops)); note_result_owed(request); // will be empty when disable_l2_cache - if (!loadback_diff.empty() || !mamba_loadback_nodes.empty()) { - cache_op_id op_id = kv_prefix_cache_.AllocateCacheOpId(); - loadback_ops.push_back(GenerateLoadBackOp(loadback_diff, mamba_loadback_nodes, op_id)); + if (!loadback_diff.empty() || !mamba_loadback_nodes.empty() || + !paged_cache_loadback_transfers.empty()) { + loadback_ops.push_back(newLoadBackOperation(request->Id(), loadback_diff, mamba_loadback_nodes, + std::move(paged_cache_loadback_transfers), + paged_cache_host_node)); } } } else if (request->Is() || (request->Is() && config_.role != Role::kP)) { // Mixed-batch disabled: skip ALL decode once a prefill was scheduled. if (!config_.enable_mixed_prefill_decode && pushed_prefill) break; - if (auto ev = scheduleDecode(request, simulated_free)) { - push_op(applyEventAndGenerateOp(request, *ev)); + auto attempt = scheduleDecode(request, simulated_free); + if (attempt.event) { + push_op(applyEventAndGenerateOp(request, *attempt.event)); note_result_owed(request); + } else { + record_decode_failure(request, attempt.failure); } } else if (request->Is() && config_.role != Role::kP) { if (!config_.enable_mixed_prefill_decode && pushed_prefill) break; - if (auto ev = scheduleDecodeFromRetracted(request, simulated_free)) { - std::vector loadback_diff = ev->GetLoadbackDiff(); - std::vector mamba_loadback_nodes = ev->GetMambaLoadbackNodes(); - push_op(applyEventAndGenerateOp(request, std::move(*ev))); + auto attempt = scheduleDecodeFromRetracted(request, simulated_free); + if (attempt.event) { + TreeNode* paged_cache_host_node = attempt.event->GetMatchResult().paged_cache_host.last_node; + push_op(applyEventAndGenerateOp(request, *attempt.event)); note_result_owed(request); - if (!loadback_diff.empty() || !mamba_loadback_nodes.empty()) { - cache_op_id op_id = kv_prefix_cache_.AllocateCacheOpId(); - loadback_ops.push_back(GenerateLoadBackOp(loadback_diff, mamba_loadback_nodes, op_id)); + std::vector loadback_diff = attempt.event->TakeLoadbackDiff(); + std::vector mamba_loadback_nodes = attempt.event->TakeMambaLoadbackNodes(); + std::vector paged_cache_loadback_transfers = + attempt.event->TakePagedCacheLoadbackTransfers(); + if (!loadback_diff.empty() || !mamba_loadback_nodes.empty() || + !paged_cache_loadback_transfers.empty()) { + loadback_ops.push_back(newLoadBackOperation(request->Id(), loadback_diff, mamba_loadback_nodes, + std::move(paged_cache_loadback_transfers), + paged_cache_host_node)); } + } else { + record_decode_failure(request, attempt.failure); } } } @@ -979,6 +1128,20 @@ Scheduler::newForwardOperation(std::vector candidates) { #if TOKENSPEED_FLAT_KVCACHE resolveFlatStarvation(candidates, /*made_progress=*/!ops.empty()); #else + const bool no_async_progress_pending = pending_forward_results_.empty() && cache_op_tracker_.empty(); + if (ops.empty() && no_async_progress_pending && all_decode_failures_are_paged_cache && + !paged_cache_blocked.empty()) { + Request* victim = paged_cache_blocked.front(); + for (Request* request : paged_cache_blocked) { + if (request->TokenSize() > victim->TokenSize() || + (request->TokenSize() == victim->TokenSize() && request->Id() < victim->Id())) { + victim = request; + } + } + abortRequest(victim, "Paged cache group capacity is exhausted"); + return {std::vector{}, std::move(loadback_ops)}; + } + // If all active decode requests failed, device memory is exhausted: retract the longest one. if (ops.empty() && !candidates.empty()) { std::vector retract_candidates; diff --git a/tokenspeed-scheduler/csrc/scheduler/outside_event_handler.cpp b/tokenspeed-scheduler/csrc/scheduler/outside_event_handler.cpp index 03b221a7f..0a4d11a0c 100644 --- a/tokenspeed-scheduler/csrc/scheduler/outside_event_handler.cpp +++ b/tokenspeed-scheduler/csrc/scheduler/outside_event_handler.cpp @@ -107,8 +107,8 @@ void Scheduler::handleEvent(const pd::BootstrappedEvent& event) { void Scheduler::handleEvent(const pd::FailedEvent& event) {} void Scheduler::handleEvent(const pd::SucceededEvent& event) { -#if TOKENSPEED_FLAT_KVCACHE pending_forward_results_.erase(event.request_id); +#if TOKENSPEED_FLAT_KVCACHE flat_reserved_pages_.erase(event.request_id); #endif std::vector page_hashes; @@ -120,6 +120,7 @@ void Scheduler::handleEvent(const pd::SucceededEvent& event) { &coordinator_ #endif }); + tryFinalizeDeferredAbort(event.request_id); } void Scheduler::handleEvent(const pd::RemotePrefillDoneEvent& event) { @@ -127,8 +128,8 @@ void Scheduler::handleEvent(const pd::RemotePrefillDoneEvent& event) { } void Scheduler::handleEvent(const forward::Finish& event) { -#if TOKENSPEED_FLAT_KVCACHE pending_forward_results_.erase(event.request_id); +#if TOKENSPEED_FLAT_KVCACHE flat_reserved_pages_.erase(event.request_id); #endif if (auto req = find_request(event.request_id)) { @@ -152,6 +153,7 @@ void Scheduler::handleEvent(const forward::Finish& event) { #endif }); } + tryFinalizeDeferredAbort(event.request_id); } void Scheduler::handleEvent(const forward::UpdateReserveNumTokens& event) { @@ -160,42 +162,48 @@ void Scheduler::handleEvent(const forward::UpdateReserveNumTokens& event) { } } void Scheduler::handleEvent(const forward::ExtendResult& event) { -#if TOKENSPEED_FLAT_KVCACHE // One owed forward result delivered (see pending_forward_results_). if (auto it = pending_forward_results_.find(event.request_id); it != pending_forward_results_.end()) { if (--it->second <= 0) { pending_forward_results_.erase(it); } } -#endif if (auto req = find_request(event.request_id)) { const std::int32_t protected_tail_tokens = config_.overlap_schedule_depth * config_.decode_input_tokens; req->Apply(fsm::ExtendResultEvent{event.request_id, event.tokens, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr, protected_tail_tokens}); } + tryFinalizeDeferredAbort(event.request_id); } void Scheduler::handleEvent(const forward::Abort& event) { + // Terminal for this request's forward stream: late results are ignored. + pending_forward_results_.erase(event.request_id); #if TOKENSPEED_FLAT_KVCACHE - // Terminal for this request's forward stream: drop any remaining result debt - // and any decode reservation it never consumed -- an abort between the + // Drop any decode reservation it never consumed -- an abort between the // prefill-completing admission and the PrefillDone->Decoding transition must // not leave a permanent phantom reservation deflating every later gate. - pending_forward_results_.erase(event.request_id); flat_reserved_pages_.erase(event.request_id); #endif - auto iter = requests_.find(event.request_id); - if (iter == requests_.end()) { + Request* req = find_request(event.request_id); + if (req == nullptr) { return; } - Request* req = iter->second.get(); - req->Apply(fsm::AbortEvent{ + // L3 prefetch owns a separate abort handoff: the request remains Aborting + // until PrefetchDone releases its host-page transfer. + if (req->Is() || req->Is()) { + req->Apply(fsm::AbortEvent{ #if TOKENSPEED_FLAT_KVCACHE - &coordinator_ + &coordinator_ #endif - }); + }); + return; + } + + deferAbort(event.request_id, /*discard_writeback=*/true); + tryFinalizeDeferredAbort(event.request_id); } void Scheduler::handleEvent(const cache::WriteBackDone& event) { @@ -226,15 +234,38 @@ void Scheduler::handleEvent(const cache::WriteBackDone& event) { auto spec = std::move(it->second); cache_op_tracker_.erase(it); + const auto deferred_it = deferred_aborts_.find(spec.request_id); + const bool discard_writeback = deferred_it != deferred_aborts_.end() && deferred_it->second.discard_writeback; + const bool effective_success = event.success && !discard_writeback; + auto now = std::chrono::steady_clock::now(); for (TreeNode* n : spec.nodes) n->Touch(now); if (!spec.request_id.empty()) { if (auto* req = find_request(spec.request_id)) { - req->Apply( - fsm::WriteBackDoneEvent{&kv_prefix_cache_, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr}); + const bool was_writing_back = req->Is(); + const bool was_retracting = req->Is(); + // A finishing request no longer needs its paged-cache table once D2H completes. + // A successful retract keeps its table for recovery; the deferred-abort finalizer + // releases failed or explicitly aborted retracts. + if (hybrid_prefix_cache_ && was_writing_back) { + hybrid_prefix_cache_->ReleaseRequest(spec.request_id); + } + req->Apply(fsm::WriteBackDoneEvent{ + &kv_prefix_cache_, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr, effective_success}); + if (was_retracting && !event.success && !discard_writeback) { + deferAbort(spec.request_id, /*discard_writeback=*/false, "L2 retract write-back failed"); + } + tryFinalizeDeferredAbort(spec.request_id); + return; } } + if (hybrid_prefix_cache_ && !spec.paged_cache_nodes.empty()) { + hybrid_prefix_cache_->OnPagedCacheHostWriteBackDone(spec.paged_cache_nodes, effective_success); + } + if (!spec.request_id.empty()) { + tryFinalizeDeferredAbort(spec.request_id); + } } void Scheduler::handleEvent(const cache::LoadBackDone& event) { @@ -249,7 +280,14 @@ void Scheduler::handleEvent(const cache::LoadBackDone& event) { return; } #endif - // Radix loadbacks emit no LoadBackDone today: unknown op_ids are silently ignored. + auto it = cache_op_tracker_.find(event.op_id); + if (it == cache_op_tracker_.end() || !it->second.paged_loadback_host_pin) { + return; + } + _assert(event.success, "radix paged-cache host loadback failed"); + const std::string request_id = it->second.request_id; + cache_op_tracker_.erase(it); + tryFinalizeDeferredAbort(request_id); } } // namespace tokenspeed diff --git a/tokenspeed-scheduler/csrc/scheduler/request.h b/tokenspeed-scheduler/csrc/scheduler/request.h index 2157be683..9e44bfd06 100644 --- a/tokenspeed-scheduler/csrc/scheduler/request.h +++ b/tokenspeed-scheduler/csrc/scheduler/request.h @@ -299,6 +299,36 @@ class Request { state_); } + template + requires(std::same_as || std::same_as) + const std::vector& GetPagedCacheWriteBackTransfers() const { + return std::visit(Overloaded{ + [](const T& s) -> const std::vector& + requires(std::same_as) + { return s.GetPagedCacheWriteBackTransfers(); }, + [this](const auto&) -> const std::vector& { + throw std::logic_error("Request::GetPagedCacheWriteBackTransfers: expected state=" + + std::string(detail::TypeName()) + "; got state=" + StateName()); + }, + }, + state_); + } + + template + requires(std::same_as || std::same_as) + const std::vector& GetPagedCacheWriteBackNodes() const { + return std::visit(Overloaded{ + [](const T& s) -> const std::vector& + requires(std::same_as) + { return s.PagedCacheWriteBackNodes(); }, + [this](const auto&) -> const std::vector& { + throw std::logic_error("Request::GetPagedCacheWriteBackNodes: expected state=" + + std::string(detail::TypeName()) + "; got state=" + StateName()); + }, + }, + state_); + } + private: std::string id_; TokenContainer token_container_; diff --git a/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp b/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp index d518d1d9b..21e4f59c9 100644 --- a/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp +++ b/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp @@ -129,6 +129,25 @@ Scheduler::Scheduler(SchedulerConfig config) copy.Validate(); hybrid_prefix_cache_->RegisterPagedCacheGroup(std::make_unique(std::move(copy))); } + std::unordered_set registered_paged_group_ids; + for (const auto& cfg : config_.paged_cache_groups) { + registered_paged_group_ids.insert(cfg.group_id); + auto host_it = config_.paged_cache_host_group_pages.find(cfg.group_id); + if (host_it == config_.paged_cache_host_group_pages.end() || host_it->second <= 0) { + continue; + } + PagedCacheGroupConfig host_copy = cfg; + host_copy.total_pages = host_it->second; + host_copy.Validate(); + hybrid_prefix_cache_->RegisterPagedCacheHostGroup( + std::make_unique(std::move(host_copy))); + } + for (const auto& [gid, _] : config_.paged_cache_host_group_pages) { + if (registered_paged_group_ids.find(gid) == registered_paged_group_ids.end()) { + throw std::invalid_argument("Scheduler: paged_cache_host_group_pages references unknown group_id '" + + gid + "'"); + } + } if (has_prefix_cache_adjunct) { const auto& spec = *config_.prefix_cache_adjunct; @@ -259,6 +278,10 @@ std::size_t Scheduler::AvailableKvPages() const { #endif } +std::size_t Scheduler::AvailableHostKvPages() const { + return host_allocator_.AvailablePages(); +} + std::size_t Scheduler::ActiveKvPages() const { std::unordered_set active_pages; for (const auto& [_, req] : requests_) { @@ -297,6 +320,27 @@ std::int64_t Scheduler::PagedCacheGroupFailedAllocCount(const std::string& group return hybrid_prefix_cache_->PagedCacheGroupFailedAllocCount(group_id); } +std::int32_t Scheduler::PagedCacheHostGroupTotalPages(const std::string& group_id) const { + if (!hybrid_prefix_cache_) { + throw std::out_of_range("Scheduler::PagedCacheHostGroupTotalPages: group_id not configured"); + } + return hybrid_prefix_cache_->PagedCacheHostGroupTotalPages(group_id); +} + +std::int32_t Scheduler::PagedCacheHostGroupAvailablePages(const std::string& group_id) const { + if (!hybrid_prefix_cache_) { + throw std::out_of_range("Scheduler::PagedCacheHostGroupAvailablePages: group_id not configured"); + } + return hybrid_prefix_cache_->PagedCacheHostGroupAvailablePages(group_id); +} + +std::int64_t Scheduler::PagedCacheHostGroupFailedAllocCount(const std::string& group_id) const { + if (!hybrid_prefix_cache_) { + throw std::out_of_range("Scheduler::PagedCacheHostGroupFailedAllocCount: group_id not configured"); + } + return hybrid_prefix_cache_->PagedCacheHostGroupFailedAllocCount(group_id); +} + std::vector Scheduler::GetRequestPagedCachePageIds(const std::string& request_id, const std::string& group_id) const { if (!hybrid_prefix_cache_) { @@ -328,21 +372,25 @@ std::vector Scheduler::newWriteBackOperation( return ops; } for (auto& [id, req] : requests) { - if (!req->Is()) continue; + if (!req->Is() || (!deferred_aborts_.empty() && deferred_aborts_.contains(id))) continue; const auto& pages_to_transfer = req->GetPagesToTransfer(); + const auto& paged_cache_transfers = req->GetPagedCacheWriteBackTransfers(); - if (!pages_to_transfer.empty()) { + if (!pages_to_transfer.empty() || !paged_cache_transfers.empty()) { cache_op_id op_id = kv_prefix_cache_.AllocateCacheOpId(); CacheOpSpec spec; spec.request_id = id; + spec.paged_cache_nodes = req->GetPagedCacheWriteBackNodes(); cache_op_tracker_[op_id] = std::move(spec); ops.push_back(WriteBackOperation{ - op_id, std::vector(pages_to_transfer.begin(), pages_to_transfer.end())}); + op_id, std::vector(pages_to_transfer.begin(), pages_to_transfer.end()), + std::vector(paged_cache_transfers.begin(), paged_cache_transfers.end())}); req->Apply(fsm::CommitDrainingEvent{}); } else { - req->Apply(fsm::AbortEvent{ + req->Apply(fsm::AbortEvent{&kv_prefix_cache_, hybrid_prefix_cache_ ? &*hybrid_prefix_cache_ : nullptr #if TOKENSPEED_FLAT_KVCACHE - &coordinator_ + , + &coordinator_ #endif }); } @@ -356,19 +404,23 @@ ExecutionPlan Scheduler::NextExecutionPlan() { std::vector write_back_ops; write_back_ops = std::move(newWriteBackOperation(requests_)); + const bool has_deferred_aborts = !deferred_aborts_.empty(); if (hybrid_prefix_cache_) { for (const auto& [id, req] : requests_) { - if (req->Is()) { + if (req->Is() && (!has_deferred_aborts || !deferred_aborts_.contains(id))) { hybrid_prefix_cache_->ReleaseRequest(id); } } } - std::erase_if(requests_, [](const auto& req) { return req.second->template Is(); }); + std::erase_if(requests_, [this, has_deferred_aborts](const auto& req) { + return req.second->template Is() && + (!has_deferred_aborts || !deferred_aborts_.contains(req.first)); + }); std::vector candidates; for (auto& [id, req] : requests_) { - if (!req->Is() && !req->Is() && !req->Is() && - !req->Is()) { + if ((!has_deferred_aborts || !deferred_aborts_.contains(id)) && !req->Is() && + !req->Is() && !req->Is() && !req->Is()) { candidates.push_back(req.get()); } } @@ -418,16 +470,17 @@ ExecutionPlan Scheduler::NextExecutionPlan() { } #endif if (!write_back_ops.empty()) { - plan.With(CacheOperation{FlatWriteBackOperation{write_back_ops}}); + plan.With(CacheOperation{FlatWriteBackOperation{std::move(write_back_ops)}}); } if (auto* lb = std::get_if>(&cache_ops)) { if (!lb->empty()) { - plan.With(CacheOperation{FlatLoadBackOperation{*lb}}); + plan.With(CacheOperation{FlatLoadBackOperation{std::move(*lb)}}); } } if (std::getenv("DEBUG_MEM")) { check_device_mem(); } + plan.WithSchedulerAborts(std::exchange(scheduler_aborts_, {})); return plan; } diff --git a/tokenspeed-scheduler/csrc/scheduler/scheduler.h b/tokenspeed-scheduler/csrc/scheduler/scheduler.h index 0975e1d8f..0e7a8c2b9 100644 --- a/tokenspeed-scheduler/csrc/scheduler/scheduler.h +++ b/tokenspeed-scheduler/csrc/scheduler/scheduler.h @@ -74,6 +74,7 @@ class Scheduler { std::size_t DecodingSize() const; std::size_t RetractedSize() const; std::size_t AvailableKvPages() const; + std::size_t AvailableHostKvPages() const; std::size_t ActiveKvPages() const; std::size_t PrefillSize() const; std::int32_t GetRequestTokenSize(const std::string& id) const; @@ -81,6 +82,9 @@ class Scheduler { std::int32_t PagedCacheGroupTotalPages(const std::string& group_id) const; std::int32_t PagedCacheGroupAvailablePages(const std::string& group_id) const; std::int64_t PagedCacheGroupFailedAllocCount(const std::string& group_id) const; + std::int32_t PagedCacheHostGroupTotalPages(const std::string& group_id) const; + std::int32_t PagedCacheHostGroupAvailablePages(const std::string& group_id) const; + std::int64_t PagedCacheHostGroupFailedAllocCount(const std::string& group_id) const; std::vector GetRequestPagedCachePageIds(const std::string& request_id, const std::string& group_id) const; // Compact-view base logical-page offset; 0 for full-history / unseen. @@ -94,6 +98,18 @@ class Scheduler { #endif private: + enum class ScheduleFailure { + kNone, + kGenericResource, + kPagedCache, + }; + + template + struct ScheduleAttempt { + std::optional event; + ScheduleFailure failure{ScheduleFailure::kGenericResource}; + }; + // Second element is LoadBackOperation list (normal path) or WriteBackOperation list (retract triggered). std::tuple, std::variant, std::vector>> @@ -106,7 +122,7 @@ class Scheduler { std::vector& loadback_ops); PrefillOperation applyEventAndGenerateOp(Request* request, fsm::SchedulePrefillEvent event); DecodeOperation applyEventAndGenerateOp(Request* request, fsm::ScheduleDecodeEvent event); - DecodeOperation applyEventAndGenerateOp(Request* request, fsm::ScheduleDecodeFromRetractedEvent event); + DecodeOperation applyEventAndGenerateOp(Request* request, fsm::ScheduleDecodeFromRetractedEvent& event); std::optional applyEventAndGenerateOp(Request* request, fsm::ScheduleRetractEvent event); PrefetchOperation applyEventAndGenerateOp(Request* request, fsm::SchedulePrefetchEvent event); @@ -118,11 +134,15 @@ class Scheduler { std::optional schedulePrefill(Request* request, std::int32_t remaining, std::int32_t reserve_num_tokens_in_next_schedule_event, std::map& simulated_free); - std::optional scheduleDecode(Request* request, - std::map& simulated_free); - std::optional scheduleDecodeFromRetracted( + ScheduleAttempt scheduleDecode(Request* request, + std::map& simulated_free); + ScheduleAttempt scheduleDecodeFromRetracted( Request* request, std::map& simulated_free); std::optional scheduleRetract(Request* request); + LoadBackOperation newLoadBackOperation(const std::string& request_id, const std::vector& diff, + const std::vector& mamba_nodes, + std::vector paged_cache_transfers, + TreeNode* paged_cache_host_node); #if TOKENSPEED_FLAT_KVCACHE // One hash pass at admission: the device match, the read-only host-tier match above its @@ -144,9 +164,19 @@ class Scheduler { void resolveFlatStarvation(const std::vector& candidates, bool made_progress); #endif + void abortRequest(Request* request, std::string message); void check_device_mem(); private: + struct DeferredAbort { + bool discard_writeback{false}; + std::string scheduler_message; + }; + + bool hasInFlightCacheOp(const std::string& request_id) const; + void deferAbort(const std::string& request_id, bool discard_writeback, std::string scheduler_message = {}); + void tryFinalizeDeferredAbort(const std::string& request_id); + void handleEvent(const cache::PrefetchDone& event); void handleEvent(const cache::WriteBackDone& event); void handleEvent(const cache::LoadBackDone& event); @@ -185,6 +215,9 @@ class Scheduler { KVPrefixCache kv_prefix_cache_; ReqPoolAllocator req_pool_allocator_; std::optional hybrid_prefix_cache_{}; + // Prefill-completing and decode operations still owed by the executor. + // Starvation recovery must not release their request-owned cache pages. + std::unordered_map pending_forward_results_; #if TOKENSPEED_FLAT_KVCACHE BlockPool block_pool_; @@ -193,9 +226,6 @@ class Scheduler { BlockPool flat_host_pool_; KvCacheCoordinator coordinator_; std::vector flat_group_ids_; // group_id per cache group, index-aligned to coordinator groups - // ExtendResults the executor still owes per request (erased on Finish/Abort/PD-success); non-empty means - // an in-flight forward can still free pool pages, which flatPoolWedged keys off. - std::unordered_map pending_forward_results_; // Reserve ledger: decode pages promised at admission but Acquired only at PrefillDone->Decoding; until // then they sit in the free count, so every flat gate subtracts OTHER requests' entries. std::unordered_map flat_reserved_pages_; @@ -273,7 +303,9 @@ class Scheduler { private: std::unordered_map> requests_; std::unordered_map cache_op_tracker_; + std::unordered_map deferred_aborts_; std::vector kv_events_; + std::vector scheduler_aborts_; SchedulerStats stats_; }; diff --git a/tokenspeed-scheduler/csrc/scheduler/types.h b/tokenspeed-scheduler/csrc/scheduler/types.h index 538b49adf..d0333c5eb 100644 --- a/tokenspeed-scheduler/csrc/scheduler/types.h +++ b/tokenspeed-scheduler/csrc/scheduler/types.h @@ -21,6 +21,7 @@ #pragma once #include +#include #include #include #include @@ -86,6 +87,7 @@ struct SchedulerConfig { } device_allocator; std::vector paged_cache_groups{}; + std::map paged_cache_host_group_pages{}; // GCD of every group's effective block_size (per-group override, else the global // block_size): the base page granularity all group block sizes are multiples of. diff --git a/tokenspeed-scheduler/python/tests/test_paged_cache_group_admission.py b/tokenspeed-scheduler/python/tests/test_paged_cache_group_admission.py index c9b237424..acb142848 100644 --- a/tokenspeed-scheduler/python/tests/test_paged_cache_group_admission.py +++ b/tokenspeed-scheduler/python/tests/test_paged_cache_group_admission.py @@ -380,13 +380,17 @@ def test_existing_terminal_state_snapshot_credit_is_branch_specific(): assert scheduler.paged_cache_group_failed_alloc_count("c4.test") == 0 -def test_transport_state_decode_checkpoint_credit_is_not_overcounted(): +# Cover both failed generic retract and available-host fallback paths. +@pytest.mark.parametrize("num_host_pages", [0, 256]) +def test_transport_state_decode_checkpoint_credit_is_not_overcounted( + num_host_pages: int, +): scheduler = _transport_state_checkpoint_scheduler( state_total_pages=20, max_scheduled_tokens=72, decode_input_tokens=64, enable_mixed_prefill_decode=True, - num_host_pages=256, + num_host_pages=num_host_pages, requests=[ ("A", list(range(68))), ("B", list(range(1_000, 1_004))), @@ -401,6 +405,7 @@ def test_transport_state_decode_checkpoint_credit_is_not_overcounted(): # would otherwise return, so A cannot safely fit at this capacity. second_plan = scheduler.next_execution_plan() assert _request_input_lengths(second_plan) == {} + assert not second_plan.scheduler_aborts assert scheduler.paged_cache_group_available_pages("c4.test") == 1 assert scheduler.paged_cache_group_failed_alloc_count("c4.test") == 0 diff --git a/tokenspeed-scheduler/tests/cpp/flat_known_radix_failures.txt b/tokenspeed-scheduler/tests/cpp/flat_known_radix_failures.txt index d3debac36..ae00186e2 100644 --- a/tokenspeed-scheduler/tests/cpp/flat_known_radix_failures.txt +++ b/tokenspeed-scheduler/tests/cpp/flat_known_radix_failures.txt @@ -43,6 +43,7 @@ ChunkedPrefillTestSuite.ChunkedPrefill_SplitsAcrossPlans ChunkedPrefillTestSuite.ExtendPrefixLen_GrowsPerChunk ChunkedPrefillTestSuite.InputIds_CorrectPerChunk ChunkedPrefillTestSuite.PrefillFirst_ContinuesPrefillBeforeNewSubmitted +DisableL2RetractTestSuite.RetractStillEmitsHostWriteBackWhenL2Disabled DisablePrefixCacheMambaRetractTest.RetractedRequestRecoversFromTreeOwnedMambaState DisablePrefixCacheTestSuite.PrefetchNotGeneratedForStorageHit DisablePrefixCacheTestSuite.SamePromptDoesNotReuseDevicePrefix @@ -76,6 +77,9 @@ OverlapDepthsAndAcceptLengths/PagedCacheOverlapRetractTest.LateResultRecoveryReb OverlapDepthsAndAcceptLengths/PagedCacheOverlapRetractTest.LateResultRecoveryRebuildsDynamicHorizon/6 PagedCacheAttachLoopTest.FullyCachedPrefillBorrowedPrefixReimported PagedCacheDecodePublishTest.ContinuingDecodePublishesAcceptedPagesOnly +PagedCacheExhaustionSchedulerTest.AbortsOneDeterministicVictimAndSurvivorAdvancesNextRound +PagedCacheL2SchedulerTest.FailedWriteBackDoesNotPublishHostSnapshotOrDemoteDeviceSnapshot +PagedCacheL2SchedulerTest.WriteBackAckDemotesDeviceSnapshotAndNextHitLoadsFromHost PagedCacheTerminalMixedSchedulerTest.MixedPrefillDecodePagedTablesCoverScheduledTokens PrefetchTestSuite.NoPrefetch_BelowThreshold PrefetchTestSuite.NoPrefetch_WhenL3Disabled @@ -89,6 +93,7 @@ RetractAbortPagesSuite.Retract_ThenAbort_NoDoubleFree RetractAbortPagesSuite.Retract_WriteBackDone_PagesAccounted RetractFromPrefillDoneTestSuite.Retract_FromPrefillDone_FullCycle RetractFromPrefillDoneTestSuite.Retract_FromPrefillDone_TriggeredWhenDeviceFull +RetractHostExhaustionSuite.InternalAbortIsPublishedInExecutionPlan RetractTailPageTestSuite.Retract_TailPageZero_PreservesBoundaryPage RetractTestSuite.Retract_RequestNotInForwardWhileRetracting RetractTestSuite.Retract_RetractedRequestRecoversToDecoding diff --git a/tokenspeed-scheduler/tests/cpp/paged_cache_test_fixture.h b/tokenspeed-scheduler/tests/cpp/paged_cache_test_fixture.h index 981cfcd6b..94782cdb8 100644 --- a/tokenspeed-scheduler/tests/cpp/paged_cache_test_fixture.h +++ b/tokenspeed-scheduler/tests/cpp/paged_cache_test_fixture.h @@ -145,9 +145,10 @@ class PagedCacheTestFixtureT : public ::testing::Test { PagedCacheGroupSnapshot BuildGroupSnap(PagedCacheGroupAllocator* alloc, std::int32_t prefix_len_tokens, std::int32_t base_logical_page, bool sliding) { PagedCacheGroupTable t{alloc}; - t.Acquire(kLcm); + t.Acquire(prefix_len_tokens); // Caller chooses absolute base; fresh table commits at 0. - auto committed = sliding ? t.CheckpointStateToSnapshot(kLcm) : t.CommitHistoryToSnapshot(kLcm); + auto committed = + sliding ? t.CheckpointStateToSnapshot(prefix_len_tokens) : t.CommitHistoryToSnapshot(prefix_len_tokens); PagedCacheGroupSnapshot g{}; g.pages = std::move(committed.pages); g.base_logical_page = base_logical_page; diff --git a/tokenspeed-scheduler/tests/cpp/test_cache_operation_kinds.cpp b/tokenspeed-scheduler/tests/cpp/test_cache_operation_kinds.cpp index 7771fa2b0..8d3c96894 100644 --- a/tokenspeed-scheduler/tests/cpp/test_cache_operation_kinds.cpp +++ b/tokenspeed-scheduler/tests/cpp/test_cache_operation_kinds.cpp @@ -4,7 +4,7 @@ namespace tokenspeed::test { -TEST(CacheOperationKindTest, FlatWriteBackBucketsTransfersByKind) { +TEST(CacheOperationKindTest, FlatWriteBackPreservesGenericAndPagedTransfers) { WriteBackOperation op; op.op_id = 7; op.transfers = { @@ -13,35 +13,65 @@ TEST(CacheOperationKindTest, FlatWriteBackBucketsTransfersByKind) { TransferPair{CacheKind::kKV, 1, 11}, TransferPair{CacheKind::kMamba, 3, 23}, }; + op.paged_cache_transfers = { + PagedCacheTransferPair{"v4.c4a.compressed_kv", {1, 2}, {101, 102}}, + PagedCacheTransferPair{"v4.swa_kv", {3}, {103}}, + }; + WriteBackOperation second; + second.op_id = 8; + second.paged_cache_transfers = { + PagedCacheTransferPair{"v4.c128a.compressed_kv", {4}, {104}}, + }; - FlatWriteBackOperation flat({op}); + FlatWriteBackOperation flat({op, second}); - ASSERT_EQ(flat.op_ids, std::vector({7})); + ASSERT_EQ(flat.op_ids, std::vector({7, 8})); EXPECT_EQ(flat.src_pages[0], std::vector({1})); EXPECT_EQ(flat.dst_pages[0], std::vector({11})); EXPECT_EQ(flat.src_pages_by_kind.at("kv")[0], std::vector({1})); EXPECT_EQ(flat.dst_pages_by_kind.at("kv")[0], std::vector({11})); EXPECT_EQ(flat.src_pages_by_kind.at("mamba")[0], std::vector({2, 3})); EXPECT_EQ(flat.dst_pages_by_kind.at("mamba")[0], std::vector({22, 23})); + ASSERT_EQ(flat.paged_cache_transfers.size(), 2u); + ASSERT_EQ(flat.paged_cache_transfers[0].size(), 2u); + EXPECT_EQ(flat.paged_cache_transfers[0][0].group_id, "v4.c4a.compressed_kv"); + EXPECT_EQ(flat.paged_cache_transfers[0][0].src_pages, std::vector({1, 2})); + EXPECT_EQ(flat.paged_cache_transfers[0][0].dst_pages, std::vector({101, 102})); + EXPECT_EQ(flat.paged_cache_transfers[0][1].group_id, "v4.swa_kv"); + ASSERT_EQ(flat.paged_cache_transfers[1].size(), 1u); + EXPECT_EQ(flat.paged_cache_transfers[1][0].group_id, "v4.c128a.compressed_kv"); } -TEST(CacheOperationKindTest, FlatLoadBackBucketsTransfersByKind) { +TEST(CacheOperationKindTest, FlatLoadBackPreservesGenericAndPagedTransfers) { LoadBackOperation op; op.op_id = 9; op.transfers = { TransferPair{CacheKind::kKV, 10, 20}, TransferPair{CacheKind::kMamba, 30, 40}, }; + LoadBackOperation paged_only; + paged_only.op_id = 10; + paged_only.paged_cache_transfers = { + PagedCacheTransferPair{"v4.c4a.compressed_kv", {101, 102}, {1, 2}}, + PagedCacheTransferPair{"v4.c4a.indexer_compressor_state", {103}, {3}}, + }; - FlatLoadBackOperation flat({op}); + FlatLoadBackOperation flat({op, paged_only}); - ASSERT_EQ(flat.op_ids, std::vector({9})); + ASSERT_EQ(flat.op_ids, std::vector({9, 10})); EXPECT_EQ(flat.src_pages[0], std::vector({10})); EXPECT_EQ(flat.dst_pages[0], std::vector({20})); EXPECT_EQ(flat.src_pages_by_kind.at("kv")[0], std::vector({10})); EXPECT_EQ(flat.dst_pages_by_kind.at("kv")[0], std::vector({20})); EXPECT_EQ(flat.src_pages_by_kind.at("mamba")[0], std::vector({30})); EXPECT_EQ(flat.dst_pages_by_kind.at("mamba")[0], std::vector({40})); + ASSERT_EQ(flat.paged_cache_transfers.size(), 2u); + EXPECT_TRUE(flat.paged_cache_transfers[0].empty()); + ASSERT_EQ(flat.paged_cache_transfers[1].size(), 2u); + EXPECT_EQ(flat.paged_cache_transfers[1][0].group_id, "v4.c4a.compressed_kv"); + EXPECT_EQ(flat.paged_cache_transfers[1][0].src_pages, std::vector({101, 102})); + EXPECT_EQ(flat.paged_cache_transfers[1][0].dst_pages, std::vector({1, 2})); + EXPECT_EQ(flat.paged_cache_transfers[1][1].group_id, "v4.c4a.indexer_compressor_state"); } } // namespace tokenspeed::test diff --git a/tokenspeed-scheduler/tests/cpp/test_paged_cache_l2_offload.cpp b/tokenspeed-scheduler/tests/cpp/test_paged_cache_l2_offload.cpp new file mode 100644 index 000000000..56b8912cc --- /dev/null +++ b/tokenspeed-scheduler/tests/cpp/test_paged_cache_l2_offload.cpp @@ -0,0 +1,839 @@ +// Copyright (c) 2026 LightSeek Foundation +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include "paged_cache_test_fixture.h" +#include "integration_test_helper.h" +#include "resource/radix_tree/tree_resource.h" +#include "scheduler/operations/forward.h" + +#include +#include +#include + +namespace tokenspeed::test { + +class PagedCacheL2OffloadTest : public PagedCacheSmallFixture { +protected: + void SetUp() override { + PagedCacheSmallFixture::SetUp(); + + auto fh_host_cfg = MakeGroupConfig("fh", kSmallFixtureParams.fh_rows_per_page, kSmallFixtureParams.fh_stride, + PagedCacheGroupConfig::Retention::FullHistory, /*window=*/0, + PagedCacheGroupFamily::History); + fh_host_cfg.total_pages = 3; + auto swa_host_cfg = MakeGroupConfig( + "swa", kSmallFixtureParams.swa_rows_per_page, kSmallFixtureParams.swa_stride, + PagedCacheGroupConfig::Retention::SlidingWindow, kSlidingWindow, PagedCacheGroupFamily::State); + swa_host_cfg.total_pages = 3; + auto fh_host = std::make_unique(std::move(fh_host_cfg)); + auto swa_host = std::make_unique(std::move(swa_host_cfg)); + hybrid_->RegisterPagedCacheHostGroup(std::move(fh_host)); + hybrid_->RegisterPagedCacheHostGroup(std::move(swa_host)); + } + + TreeNode* SeedCompleteDeviceSnapshot(token_t token_start = 1) { + TreeNode* terminal = InsertDevicePages(/*num_pages=*/2, token_start); + EXPECT_NE(terminal, nullptr); + hybrid_->AttachPagedCacheSnapshotToNode(terminal, MakeCompleteSnapshot(kLcm)); + return terminal; + } + + void AttachHostResource(TreeNode* node) { + ASSERT_NE(node, nullptr); + node->AttachResource(std::make_unique>(OwnedPages{})); + } +}; + +class PagedCacheL2SchedulerTest : public SchedulerTestSuite { +protected: + SchedulerConfig MakeConfig() override { + auto cfg = SchedulerTestSuite::MakeConfig(); + cfg.block_size = kSmallFixtureParams.page_size; + cfg.device_allocator.total_pages = 32; + cfg.host_allocator.total_pages = 32; + cfg.max_scheduled_tokens = 64; + cfg.max_batch_size = 8; + cfg.enable_l3_storage = false; + cfg.disable_l2_cache = false; + + cfg.paged_cache_groups = { + MakeSchedulerGroup("fh", kSmallFixtureParams.fh_rows_per_page, kSmallFixtureParams.fh_stride, + PagedCacheGroupConfig::Retention::FullHistory, /*window=*/0, + PagedCacheGroupFamily::History), + MakeSchedulerGroup("swa", kSmallFixtureParams.swa_rows_per_page, kSmallFixtureParams.swa_stride, + PagedCacheGroupConfig::Retention::SlidingWindow, + kSmallFixtureParams.sliding_window_tokens, PagedCacheGroupFamily::State), + }; + cfg.paged_cache_host_group_pages = {{"fh", 16}, {"swa", 16}}; + PrefixCacheAdjunctSpec adjunct; + adjunct.required_groups = {"fh"}; + cfg.prefix_cache_adjunct = adjunct; + return cfg; + } + + static PagedCacheGroupConfig MakeSchedulerGroup(std::string group_id, std::int32_t rows_per_page, + std::int32_t stride, PagedCacheGroupConfig::Retention retention, + std::int32_t window, PagedCacheGroupFamily family) { + PagedCacheGroupConfig cfg{}; + cfg.group_id = std::move(group_id); + cfg.rows_per_page = rows_per_page; + cfg.entry_stride_tokens = stride; + cfg.total_pages = 16; + cfg.retention = retention; + if (retention == PagedCacheGroupConfig::Retention::SlidingWindow) { + cfg.sliding_window_tokens = window; + } + cfg.family = family; + return cfg; + } + + void BringToDecoding(const std::string& id, token_t start = 1) { + Submit(MakeRequestSpec(id, /*num_pages=*/2, start)); + PlanOnce(); + SendForwardDone(id, {42}); + PlanOnce(); + } + + static const FlatWriteBackOperation* GetWriteBack(const ExecutionPlan& plan) { + for (const auto& op : plan.Operations()) { + if (auto* cache_op = std::get_if(&op)) { + if (auto* wb = std::get_if(cache_op)) { + return wb; + } + } + } + return nullptr; + } + + static const FlatLoadBackOperation* GetLoadBack(const ExecutionPlan& plan) { + for (const auto& op : plan.Operations()) { + if (auto* cache_op = std::get_if(&op)) { + if (auto* lb = std::get_if(cache_op)) { + return lb; + } + } + } + return nullptr; + } + + static const FlatForwardOperation* GetForward(const ExecutionPlan& plan) { + for (const auto& op : plan.Operations()) { + if (auto* fwd = std::get_if(&op)) { + return fwd; + } + } + return nullptr; + } + + static std::map UniqueDstPagesByGroup(const FlatWriteBackOperation& writeback) { + std::map> pages_by_group; + for (const auto& transfers : writeback.paged_cache_transfers) { + for (const auto& transfer : transfers) { + auto& pages = pages_by_group[transfer.group_id]; + pages.insert(transfer.dst_pages.begin(), transfer.dst_pages.end()); + } + } + std::map counts; + for (const auto& [group_id, pages] : pages_by_group) { + counts[group_id] = static_cast(pages.size()); + } + return counts; + } + + static std::int32_t UniqueKvDstPages(const FlatWriteBackOperation& writeback) { + std::set pages; + for (const auto& op_pages : writeback.dst_pages) { + pages.insert(op_pages.begin(), op_pages.end()); + } + return static_cast(pages.size()); + } +}; + +class PagedCacheExhaustionSchedulerTest : public PagedCacheL2SchedulerTest { +protected: + SchedulerConfig MakeConfig() override { + auto cfg = PagedCacheL2SchedulerTest::MakeConfig(); + for (auto& group : cfg.paged_cache_groups) { + group.total_pages = group.group_id == "fh" ? 3 : 5; + } + return cfg; + } +}; + +TEST(TreeNodeTest, DuplicateChildPreservesExistingAndIncomingOwnership) { + TreeNode root; + const token_vec_t key{1, 2}; + auto existing = std::make_unique(key); + auto incoming = std::make_unique(key); + TreeNode* existing_ptr = existing.get(); + TreeNode* incoming_ptr = incoming.get(); + + root.AddChild(key, std::move(existing)); + EXPECT_THROW(root.AddChild(key, std::move(incoming)), std::logic_error); + + ASSERT_NE(incoming, nullptr); + EXPECT_EQ(incoming.get(), incoming_ptr); + auto retained = root.RemoveChild(key); + ASSERT_NE(retained, nullptr); + EXPECT_EQ(retained.get(), existing_ptr); +} + +TEST_F(PagedCacheExhaustionSchedulerTest, AbortsOneDeterministicVictimAndSurvivorAdvancesNextRound) { + Submit({MakeRequestSpec("z-request", /*num_pages=*/2, /*start=*/1), + MakeRequestSpec("a-request", /*num_pages=*/2, /*start=*/101)}); + + auto prefill_plan = PlanOnce(); + const auto* prefill = GetForward(prefill_plan); + ASSERT_NE(prefill, nullptr); + ASSERT_EQ(prefill->request_ids.size(), 2u); + SendForwardDone("z-request", {42}); + SendForwardDone("a-request", {43}); + + auto abort_plan = PlanOnce(); + ASSERT_EQ(abort_plan.SchedulerAborts().size(), 1u); + EXPECT_EQ(abort_plan.SchedulerAborts().front().request_id, "a-request"); + const auto* blocked = GetForward(abort_plan); + ASSERT_NE(blocked, nullptr); + EXPECT_TRUE(blocked->request_ids.empty()); + + auto survivor_plan = PlanOnce(); + EXPECT_TRUE(survivor_plan.SchedulerAborts().empty()); + const auto* survivor = GetForward(survivor_plan); + ASSERT_NE(survivor, nullptr); + ASSERT_EQ(survivor->request_ids.size(), 1u); + EXPECT_EQ(survivor->request_ids.front(), "z-request"); +} + +class PagedCacheHistoryOnlyL2OffloadTest : public ::testing::Test { +protected: + static constexpr std::int32_t kPageSize = 2; + static constexpr std::int32_t kLcm = 4; + + void SetUp() override { + device_alloc_ = std::make_unique(kPageSize, /*total_pages=*/64); + kv_cache_ = std::make_unique(device_alloc_.get(), /*host=*/nullptr); + + auto history_cfg = MakeHistoryGroup(/*total_pages=*/32); + auto history_owner = std::make_unique(history_cfg); + history_alloc_ = history_owner.get(); + hybrid_ = std::make_unique(*kv_cache_, /*mamba=*/nullptr, /*mamba_chunk_size=*/0); + hybrid_->RegisterPagedCacheGroup(std::move(history_owner)); + + auto host_history_cfg = MakeHistoryGroup(/*total_pages=*/32); + auto host_history = std::make_unique(host_history_cfg); + hybrid_->RegisterPagedCacheHostGroup(std::move(host_history)); + hybrid_->EnablePagedCacheAdjunct({"fh"}, {}); + kv_cache_->GetDeviceManager().SetEvictionCallback([this](TreeNode* node) { hybrid_->OnKVEvict(node); }); + } + + TreeNode* InsertDeviceTokens(std::int32_t num_pages, token_t token_start = 1) { + auto tokens = MakeAlignedTokens(num_pages, kPageSize, token_start); + OwnedPages pages = device_alloc_->Allocate(num_pages); + auto inserted = + kv_cache_->Insert(tokens, /*prefix_pages=*/{}, std::move(pages), /*page_hashes=*/{}); + return inserted.last_node; + } + + std::unique_ptr MakeHistorySnapshot(std::int32_t prefix_len_tokens) { + PagedCacheGroupTable table{history_alloc_}; + table.Acquire(prefix_len_tokens); + auto committed = table.CommitHistoryToSnapshot(prefix_len_tokens); + + PagedCacheGroupSnapshot group{}; + group.pages = std::move(committed.pages); + group.base_logical_page = committed.segment_base_logical_page; + group.raw_token_cursor = prefix_len_tokens; + group.sliding = false; + + auto snapshot = std::make_unique(); + snapshot->prefix_len_tokens = prefix_len_tokens; + snapshot->groups.emplace("fh", std::move(group)); + return snapshot; + } + + static PagedCacheGroupConfig MakeHistoryGroup(std::int32_t total_pages) { + PagedCacheGroupConfig cfg{}; + cfg.group_id = "fh"; + cfg.rows_per_page = 4; + cfg.entry_stride_tokens = 1; + cfg.total_pages = total_pages; + cfg.retention = PagedCacheGroupConfig::Retention::FullHistory; + cfg.family = PagedCacheGroupFamily::History; + return cfg; + } + + std::unique_ptr device_alloc_; + std::unique_ptr kv_cache_; + PagedCacheGroupAllocator* history_alloc_{nullptr}; + std::unique_ptr hybrid_; +}; + +TEST_F(PagedCacheL2OffloadTest, PendingHostSnapshotInvisibleUntilWriteBackAck) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + + auto transfers = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_EQ(transfers.size(), 2u); + EXPECT_TRUE(terminal->HasPagedCachePendingHostSnapshot()); + EXPECT_FALSE(terminal->HasPagedCacheHostSnapshot()); + + auto pending_match = hybrid_->Match(tokens); + EXPECT_EQ(pending_match.paged_cache_host.last_node, nullptr); + EXPECT_EQ(pending_match.paged_cache_host.prefix_len_tokens, 0); + + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + EXPECT_FALSE(terminal->HasPagedCachePendingHostSnapshot()); + ASSERT_TRUE(terminal->HasPagedCacheHostSnapshot()); + + auto visible_match = hybrid_->Match(tokens); + ASSERT_NE(visible_match.paged_cache_host.last_node, nullptr); + EXPECT_EQ(visible_match.paged_cache_host.last_node, terminal); + EXPECT_EQ(visible_match.paged_cache_host.prefix_len_tokens, kLcm); +} + +TEST_F(PagedCacheL2OffloadTest, FailedHostWriteBackReleasesPendingSnapshotWithoutPublishing) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + const std::int32_t fh_before = hybrid_->PagedCacheHostGroupAvailablePages("fh"); + const std::int32_t swa_before = hybrid_->PagedCacheHostGroupAvailablePages("swa"); + + auto transfers = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(transfers.empty()); + ASSERT_TRUE(terminal->HasPagedCachePendingHostSnapshot()); + ASSERT_LT(hybrid_->PagedCacheHostGroupAvailablePages("fh"), fh_before); + ASSERT_LT(hybrid_->PagedCacheHostGroupAvailablePages("swa"), swa_before); + + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/false); + + EXPECT_FALSE(terminal->HasPagedCachePendingHostSnapshot()); + EXPECT_FALSE(terminal->HasPagedCacheHostSnapshot()); + EXPECT_TRUE(terminal->HasPagedCacheSnapshot()); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("fh"), fh_before); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("swa"), swa_before); + auto match = hybrid_->Match(tokens); + EXPECT_EQ(match.paged_cache_host.last_node, nullptr); + EXPECT_EQ(match.paged_cache_host.prefix_len_tokens, 0); +} + +TEST_F(PagedCacheL2OffloadTest, DemotedDeviceSnapshotDoesNotMatchAsDeviceHit) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + + auto device_match = hybrid_->Match(tokens); + ASSERT_EQ(device_match.paged_cache.last_node, terminal); + ASSERT_EQ(device_match.paged_cache.prefix_len_tokens, kLcm); + + auto demoted = terminal->DetachResource(); + ASSERT_NE(demoted, nullptr); + hybrid_->OnKVDeviceDemote(terminal); + + EXPECT_TRUE(terminal->HasPagedCacheSnapshot()); + EXPECT_FALSE(terminal->OnDevice()); + auto demoted_match = hybrid_->Match(tokens); + EXPECT_EQ(demoted_match.paged_cache.last_node, nullptr); + EXPECT_EQ(demoted_match.paged_cache.prefix_len_tokens, 0); + EXPECT_EQ(demoted_match.paged_cache_host.last_node, nullptr); +} + +TEST_F(PagedCacheL2OffloadTest, DeviceDemotePreservesSnapshotBorrowedByRequestTable) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens); + ASSERT_EQ(match.paged_cache.last_node, terminal); + hybrid_->AcquireForRequest("borrower", /*first_raw_position_of_op=*/0, + /*target_raw_tokens_exclusive=*/kLcm, match.paged_cache); + const auto borrowed_fh = hybrid_->GetRequestPagedCachePageIds("borrower", "fh"); + ASSERT_FALSE(borrowed_fh.empty()); + + auto demoted = terminal->DetachResource(); + ASSERT_NE(demoted, nullptr); + hybrid_->OnKVDeviceDemote(terminal); + + EXPECT_TRUE(terminal->HasPagedCacheSnapshot()); + EXPECT_EQ(hybrid_->GetRequestPagedCachePageIds("borrower", "fh"), borrowed_fh); + + hybrid_->ReleaseRequest("borrower"); + hybrid_->OnKVDeviceDemote(terminal); + EXPECT_FALSE(terminal->HasPagedCacheSnapshot()); +} + +TEST_F(PagedCacheL2OffloadTest, ReleasingBorrowedOnlyTableAllowsAdmissionToPruneSnapshot) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens); + ASSERT_EQ(match.paged_cache.last_node, terminal); + + hybrid_->AcquireForRequest("borrower", /*first_raw_position_of_op=*/0, + /*target_raw_tokens_exclusive=*/kLcm, match.paged_cache); + auto blocked_free = hybrid_->InitialSimulatedFree(); + blocked_free["fh"] = 0; + blocked_free["swa"] = 0; + EXPECT_FALSE(hybrid_->AdmitChunk("pressure", /*first_raw_position_of_op=*/0, + /*target_raw_tokens_exclusive=*/kLcm, blocked_free)); + EXPECT_TRUE(terminal->HasPagedCacheSnapshot()); + + hybrid_->ReleaseRequest("borrower"); + auto released_free = hybrid_->InitialSimulatedFree(); + released_free["fh"] = 0; + released_free["swa"] = 0; + EXPECT_TRUE(hybrid_->AdmitChunk("pressure", /*first_raw_position_of_op=*/0, + /*target_raw_tokens_exclusive=*/kLcm, released_free)); + EXPECT_FALSE(terminal->HasPagedCacheSnapshot()); +} + +TEST_F(PagedCacheL2OffloadTest, PendingSnapshotPartialSplitPreservesSuffixIdentityAndBorrowedPages) { + PageAllocator split_host_alloc{kPageSize, /*total_pages=*/3}; + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + const PagedCacheSnapshot* device_snapshot = terminal->GetPagedCacheSnapshot(); + const std::size_t terminal_depth = terminal->DepthInTokens(); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens); + ASSERT_EQ(match.paged_cache.last_node, terminal); + hybrid_->AcquireForRequest("borrower", /*first_raw_position_of_op=*/0, + /*target_raw_tokens_exclusive=*/kLcm, match.paged_cache); + const auto borrowed_fh = hybrid_->GetRequestPagedCachePageIds("borrower", "fh"); + const auto borrowed_swa = hybrid_->GetRequestPagedCachePageIds("borrower", "swa"); + + terminal->AttachResource( + std::make_unique>(split_host_alloc.Allocate(/*count=*/2))); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + ASSERT_TRUE(terminal->HasPagedCachePendingHostSnapshot()); + + TreeNode* prefix = kv_cache_->GetRadixTree().SplitAt(terminal, kPageSize); + ASSERT_NE(prefix, nullptr); + EXPECT_NE(prefix, terminal); + EXPECT_EQ(terminal->Parent(), prefix); + EXPECT_EQ(terminal->DepthInTokens(), terminal_depth); + EXPECT_EQ(terminal->GetPagedCacheSnapshot(), device_snapshot); + EXPECT_TRUE(terminal->HasPagedCachePendingHostSnapshot()); + EXPECT_EQ(hybrid_->GetRequestPagedCachePageIds("borrower", "fh"), borrowed_fh); + EXPECT_EQ(hybrid_->GetRequestPagedCachePageIds("borrower", "swa"), borrowed_swa); + + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + EXPECT_FALSE(terminal->HasPagedCachePendingHostSnapshot()); + EXPECT_TRUE(terminal->HasPagedCacheHostSnapshot()); + EXPECT_EQ(terminal->GetPagedCacheSnapshot(), device_snapshot); + hybrid_->ReleaseRequest("borrower"); + + auto prefix_host = prefix->DetachResource(); + auto suffix_host = terminal->DetachResource(); + ASSERT_NE(prefix_host, nullptr); + ASSERT_NE(suffix_host, nullptr); +} + +TEST_F(PagedCacheL2OffloadTest, HostSnapshotPartialSplitStaysOnSuffix) { + PageAllocator split_host_alloc{kPageSize, /*total_pages=*/3}; + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + terminal->AttachResource( + std::make_unique>(split_host_alloc.Allocate(/*count=*/2))); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + const PagedCacheSnapshot* device_snapshot = terminal->GetPagedCacheSnapshot(); + const PagedCacheSnapshot* host_snapshot = terminal->GetPagedCacheHostSnapshot(); + ASSERT_NE(device_snapshot, nullptr); + ASSERT_NE(host_snapshot, nullptr); + const std::size_t terminal_depth = terminal->DepthInTokens(); + + TreeNode* prefix = kv_cache_->GetRadixTree().SplitAt(terminal, kPageSize); + + ASSERT_NE(prefix, nullptr); + EXPECT_EQ(terminal->Parent(), prefix); + EXPECT_EQ(terminal->DepthInTokens(), terminal_depth); + EXPECT_EQ(terminal->GetPagedCacheSnapshot(), device_snapshot); + EXPECT_EQ(terminal->GetPagedCacheHostSnapshot(), host_snapshot); + + auto prefix_host = prefix->DetachResource(); + auto suffix_host = terminal->DetachResource(); + ASSERT_NE(prefix_host, nullptr); + ASSERT_NE(suffix_host, nullptr); +} + +TEST_F(PagedCacheL2OffloadTest, HostHitMaterializesDeviceDestinationPagesForForwardMetadata) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + auto detached_device = hybrid_->DetachPagedCacheSnapshotFromNode(terminal); + ASSERT_NE(detached_device, nullptr); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens); + ASSERT_EQ(match.paged_cache.last_node, nullptr); + ASSERT_NE(match.paged_cache_host.last_node, nullptr); + + auto loadback = hybrid_->PreparePagedCacheDeviceLoadBack("r-load", match.paged_cache_host); + ASSERT_EQ(loadback.size(), 2u); + + PrefillOperation op{{ + .request_id = "r-load", + .request_pool_index = 0, + .input_length = 0, + .occupied_pages = {}, + .begin = 0, + .size = 0, + .prefill_length = kLcm, + }}; + hybrid_->AcquireForRequest("r-load", /*first_raw_position_of_op=*/0, /*target_raw_tokens_exclusive=*/kLcm); + hybrid_->PopulateOp(op); + + for (const auto& transfer : loadback) { + auto pages_it = op.paged_cache_pages.find(transfer.group_id); + ASSERT_NE(pages_it, op.paged_cache_pages.end()) << transfer.group_id; + EXPECT_EQ(pages_it->second, transfer.dst_pages) << transfer.group_id; + } + + hybrid_->ReleaseRequest("r-load"); +} + +TEST_F(PagedCacheHistoryOnlyL2OffloadTest, HostLoadedHistoryPrefixCanPublishContinuationSnapshot) { + TreeNode* prefix = InsertDeviceTokens(/*num_pages=*/2); + ASSERT_NE(prefix, nullptr); + ASSERT_TRUE(hybrid_->AttachPagedCacheSnapshotToNode(prefix, MakeHistorySnapshot(kLcm))); + prefix->AttachResource(std::make_unique>(OwnedPages{})); + + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({prefix}); + ASSERT_EQ(writeback.size(), 1u); + hybrid_->OnPagedCacheHostWriteBackDone({prefix}, /*success=*/true); + auto detached_device = hybrid_->DetachPagedCacheSnapshotFromNode(prefix); + ASSERT_NE(detached_device, nullptr); + + auto host_match = hybrid_->Match(MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1)); + ASSERT_EQ(host_match.paged_cache.last_node, nullptr); + ASSERT_NE(host_match.paged_cache_host.last_node, nullptr); + EXPECT_EQ(host_match.paged_cache_host.prefix_len_tokens, kLcm); + + auto loadback = hybrid_->PreparePagedCacheDeviceLoadBack("r-load", host_match.paged_cache_host); + ASSERT_EQ(loadback.size(), 1u); + + TreeNode* extended = InsertDeviceTokens(/*num_pages=*/4); + ASSERT_NE(extended, nullptr); + hybrid_->AcquireForRequest("r-load", /*first_raw_position_of_op=*/kLcm, + /*target_raw_tokens_exclusive=*/2 * kLcm); + hybrid_->CommitChunk("r-load", extended); + + ASSERT_TRUE(extended->HasPagedCacheSnapshot()); + auto device_match = hybrid_->Match(MakeAlignedTokens(/*num_pages=*/4, kPageSize, /*start=*/1)); + ASSERT_NE(device_match.paged_cache.last_node, nullptr); + EXPECT_EQ(device_match.paged_cache.prefix_len_tokens, 2 * kLcm); + ASSERT_EQ(device_match.paged_cache.per_group_page_ids.count("fh"), 1u); + EXPECT_EQ(device_match.paged_cache.per_group_page_ids.at("fh").size(), 2u); +} + +TEST_F(PagedCacheL2OffloadTest, HostStateRecoveryRequiresExactTerminalStateCompleteness) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + auto detached_device = hybrid_->DetachPagedCacheSnapshotFromNode(terminal); + ASSERT_NE(detached_device, nullptr); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens, MatchIntent::StateRecovery); + + ASSERT_EQ(match.paged_cache.last_node, nullptr); + ASSERT_NE(match.paged_cache_host.last_node, nullptr); + EXPECT_EQ(match.paged_cache_host.last_node, terminal); + EXPECT_EQ(match.paged_cache_host.prefix_len_tokens, kLcm); +} + +TEST_F(PagedCacheL2OffloadTest, HostStateRecoveryCapsWhenTerminalStateGroupMissing) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + auto detached_device = hybrid_->DetachPagedCacheSnapshotFromNode(terminal); + ASSERT_NE(detached_device, nullptr); + + auto host_snapshot = hybrid_->DetachPagedCacheHostSnapshotFromNode(terminal); + ASSERT_NE(host_snapshot, nullptr); + host_snapshot->groups.erase("swa"); + ASSERT_TRUE(hybrid_->AttachPagedCacheHostSnapshotToNode(terminal, std::move(host_snapshot))); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens, MatchIntent::StateRecovery); + + EXPECT_EQ(match.paged_cache.last_node, nullptr); + EXPECT_EQ(match.paged_cache_host.last_node, nullptr); + ASSERT_NE(match.device.last_node, nullptr); + EXPECT_TRUE(match.device.last_node->IsRoot()); + ASSERT_NE(match.host.last_node, nullptr); + EXPECT_TRUE(match.host.last_node->IsRoot()); +} + +TEST_F(PagedCacheL2OffloadTest, PrefixReuseMissesWhenContinuationStateGroupMissing) { + hybrid_->EnablePagedCacheAdjunct({"fh"}, {}); + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + + auto device_snapshot = hybrid_->DetachPagedCacheSnapshotFromNode(terminal); + ASSERT_NE(device_snapshot, nullptr); + device_snapshot->groups.erase("swa"); + ASSERT_TRUE(hybrid_->AttachPagedCacheSnapshotToNode(terminal, std::move(device_snapshot))); + + const auto tokens = MakeAlignedTokens(/*num_pages=*/2, kPageSize, /*start=*/1); + auto device_match = hybrid_->Match(tokens); + EXPECT_EQ(device_match.paged_cache.last_node, nullptr); + EXPECT_EQ(device_match.paged_cache_host.last_node, nullptr); + ASSERT_NE(device_match.device.last_node, nullptr); + EXPECT_TRUE(device_match.device.last_node->IsRoot()); + + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + auto detached_device = hybrid_->DetachPagedCacheSnapshotFromNode(terminal); + ASSERT_NE(detached_device, nullptr); + + auto host_snapshot = hybrid_->DetachPagedCacheHostSnapshotFromNode(terminal); + ASSERT_NE(host_snapshot, nullptr); + host_snapshot->groups.erase("swa"); + ASSERT_TRUE(hybrid_->AttachPagedCacheHostSnapshotToNode(terminal, std::move(host_snapshot))); + + auto host_match = hybrid_->Match(tokens); + + EXPECT_EQ(host_match.paged_cache.last_node, nullptr); + EXPECT_EQ(host_match.paged_cache_host.last_node, nullptr); + ASSERT_NE(host_match.device.last_node, nullptr); + EXPECT_TRUE(host_match.device.last_node->IsRoot()); + ASSERT_NE(host_match.host.last_node, nullptr); + EXPECT_TRUE(host_match.host.last_node->IsRoot()); +} + +TEST_F(PagedCacheL2OffloadTest, KVHostEvictReleasesPagedCacheHostSnapshotPages) { + TreeNode* terminal = SeedCompleteDeviceSnapshot(); + ASSERT_NE(terminal, nullptr); + AttachHostResource(terminal); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({terminal}); + ASSERT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({terminal}, /*success=*/true); + ASSERT_TRUE(terminal->HasPagedCacheHostSnapshot()); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("fh"), 1); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("swa"), 0); + + hybrid_->OnKVHostEvict(terminal); + + EXPECT_FALSE(terminal->HasPagedCacheHostSnapshot()); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("fh"), 2); + EXPECT_EQ(hybrid_->PagedCacheHostGroupAvailablePages("swa"), 2); +} + +TEST_F(PagedCacheL2OffloadTest, HostWriteBackPrunesOnlyUnpinnedPagedCacheHostSnapshots) { + TreeNode* pinned = SeedCompleteDeviceSnapshot(/*token_start=*/1); + ASSERT_NE(pinned, nullptr); + AttachHostResource(pinned); + auto first_writeback = hybrid_->PreparePagedCacheHostWriteBack({pinned}); + ASSERT_FALSE(first_writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({pinned}, /*success=*/true); + ASSERT_TRUE(pinned->HasPagedCacheHostSnapshot()); + + TreeNode* next = SeedCompleteDeviceSnapshot(/*token_start=*/101); + ASSERT_NE(next, nullptr); + AttachHostResource(next); + { + HostNodeRef lock(pinned); + auto blocked = hybrid_->PreparePagedCacheHostWriteBack({next}); + EXPECT_TRUE(blocked.empty()); + EXPECT_TRUE(pinned->HasPagedCacheHostSnapshot()); + EXPECT_FALSE(next->HasPagedCachePendingHostSnapshot()); + } + + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({next}); + ASSERT_FALSE(writeback.empty()); + EXPECT_FALSE(pinned->HasPagedCacheHostSnapshot()); + EXPECT_TRUE(next->HasPagedCachePendingHostSnapshot()); +} + +TEST_F(PagedCacheL2OffloadTest, HostWriteBackPrunesMultipleSnapshotsForOneAllocation) { + auto publish_history_snapshot = [&](token_t token_start) { + TreeNode* node = InsertDevicePages(/*num_pages=*/2, token_start); + EXPECT_NE(node, nullptr); + EXPECT_TRUE(hybrid_->AttachPagedCacheSnapshotToNode(node, MakeHistoryOnlySnapshot(kLcm))); + AttachHostResource(node); + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({node}); + EXPECT_FALSE(writeback.empty()); + hybrid_->OnPagedCacheHostWriteBackDone({node}, /*success=*/true); + return node; + }; + + TreeNode* first = publish_history_snapshot(/*token_start=*/1); + TreeNode* second = publish_history_snapshot(/*token_start=*/101); + ASSERT_TRUE(first->HasPagedCacheHostSnapshot()); + ASSERT_TRUE(second->HasPagedCacheHostSnapshot()); + + constexpr std::int32_t kTargetTokens = 2 * kLcm; + TreeNode* target = InsertDevicePages(/*num_pages=*/2, /*token_start=*/201, first); + ASSERT_NE(target, nullptr); + ASSERT_TRUE(hybrid_->AttachPagedCacheSnapshotToNode(target, MakeHistoryOnlySnapshot(kTargetTokens))); + AttachHostResource(target); + + auto writeback = hybrid_->PreparePagedCacheHostWriteBack({target}); + + ASSERT_FALSE(writeback.empty()); + EXPECT_FALSE(first->HasPagedCacheHostSnapshot()); + EXPECT_FALSE(second->HasPagedCacheHostSnapshot()); + EXPECT_TRUE(target->HasPagedCachePendingHostSnapshot()); +} + +#if !TOKENSPEED_FLAT_KVCACHE +TEST_F(PagedCacheL2SchedulerTest, WriteBackAckDemotesDeviceSnapshotAndLoadbackPinsHostSnapshot) { + const auto request_tokens = MakeAlignedTokens(/*num_pages=*/3, kSmallFixtureParams.page_size, /*start=*/1); + BringToDecoding("r1"); + SendFinish("r1"); + + auto writeback_plan = PlanOnce(); + const auto* writeback = GetWriteBack(writeback_plan); + ASSERT_NE(writeback, nullptr); + ASSERT_EQ(writeback->op_ids.size(), 1u); + ASSERT_EQ(writeback->paged_cache_transfers.size(), 1u); + ASSERT_FALSE(writeback->paged_cache_transfers[0].empty()); + + SendWriteBackDone(writeback->op_ids[0]); + PlanOnce(); + EXPECT_EQ(scheduler_->PagedCacheGroupAvailablePages("fh"), 15); + EXPECT_EQ(scheduler_->PagedCacheGroupAvailablePages("swa"), 15); + EXPECT_EQ(scheduler_->PagedCacheHostGroupTotalPages("fh"), 16); + EXPECT_EQ(scheduler_->PagedCacheHostGroupTotalPages("swa"), 16); + const std::int32_t fh_host_after_writeback = scheduler_->PagedCacheHostGroupAvailablePages("fh"); + const std::int32_t swa_host_after_writeback = scheduler_->PagedCacheHostGroupAvailablePages("swa"); + EXPECT_LT(fh_host_after_writeback, scheduler_->PagedCacheHostGroupTotalPages("fh")); + EXPECT_LT(swa_host_after_writeback, scheduler_->PagedCacheHostGroupTotalPages("swa")); + + Submit(RequestSpec{ + .request_id = "r2", + .tokens = request_tokens, + }); + auto loadback_plan = PlanOnce(); + const auto* loadback = GetLoadBack(loadback_plan); + ASSERT_NE(loadback, nullptr); + ASSERT_EQ(loadback->paged_cache_transfers.size(), 1u); + ASSERT_FALSE(loadback->paged_cache_transfers[0].empty()); + + std::set groups; + for (const auto& transfer : loadback->paged_cache_transfers[0]) { + groups.insert(transfer.group_id); + EXPECT_EQ(transfer.src_pages.size(), transfer.dst_pages.size()); + EXPECT_FALSE(transfer.src_pages.empty()); + } + EXPECT_EQ(groups, (std::set{"fh", "swa"})); + EXPECT_EQ(scheduler_->PagedCacheHostGroupAvailablePages("fh"), fh_host_after_writeback); + EXPECT_EQ(scheduler_->PagedCacheHostGroupAvailablePages("swa"), swa_host_after_writeback); + + ASSERT_EQ(loadback->op_ids.size(), 1u); + EXPECT_THROW(SendLoadBackDone(loadback->op_ids[0], /*success=*/false), std::runtime_error); + EXPECT_NO_THROW(SendLoadBackDone(loadback->op_ids[0])); + EXPECT_NO_THROW(SendLoadBackDone(loadback->op_ids[0], /*success=*/false)); +} + +TEST_F(PagedCacheL2SchedulerTest, AbortDuringLoadbackWaitsForAck) { + const auto request_tokens = MakeAlignedTokens(/*num_pages=*/3, kSmallFixtureParams.page_size, /*start=*/1); + BringToDecoding("r1"); + SendFinish("r1"); + + auto writeback_plan = PlanOnce(); + const auto* writeback = GetWriteBack(writeback_plan); + ASSERT_NE(writeback, nullptr); + ASSERT_EQ(writeback->op_ids.size(), 1u); + SendWriteBackDone(writeback->op_ids[0]); + PlanOnce(); + + Submit(RequestSpec{.request_id = "r2", .tokens = request_tokens}); + auto loadback_plan = PlanOnce(); + const auto* loadback = GetLoadBack(loadback_plan); + ASSERT_NE(loadback, nullptr); + ASSERT_EQ(loadback->op_ids.size(), 1u); + + ExecutionEvent abort; + abort.With(ForwardEvent{forward::Abort{.request_id = "r2"}}); + scheduler_->Advance(std::move(abort)); + EXPECT_EQ(scheduler_->PrefillSize(), 1u); + auto waiting_plan = PlanOnce(); + ASSERT_NE(GetForward(waiting_plan), nullptr); + EXPECT_TRUE(GetForward(waiting_plan)->request_ids.empty()); + EXPECT_EQ(GetLoadBack(waiting_plan), nullptr); + + SendLoadBackDone(loadback->op_ids[0]); + EXPECT_EQ(scheduler_->PrefillSize(), 0u); + EXPECT_TRUE(PlanOnce().SchedulerAborts().empty()); +} +#endif + +TEST_F(PagedCacheL2SchedulerTest, FailedWriteBackDoesNotPublishHostSnapshotOrDemoteDeviceSnapshot) { + const auto request_tokens = MakeAlignedTokens(/*num_pages=*/3, kSmallFixtureParams.page_size, /*start=*/1); + BringToDecoding("r1"); + SendFinish("r1"); + + auto writeback_plan = PlanOnce(); + const auto* writeback = GetWriteBack(writeback_plan); + ASSERT_NE(writeback, nullptr); + ASSERT_EQ(writeback->op_ids.size(), 1u); + ASSERT_EQ(writeback->paged_cache_transfers.size(), 1u); + ASSERT_FALSE(writeback->paged_cache_transfers[0].empty()); + const auto transferred_host_pages = UniqueDstPagesByGroup(*writeback); + const std::int32_t transferred_kv_host_pages = UniqueKvDstPages(*writeback); + const std::size_t kv_host_before_ack = scheduler_->AvailableHostKvPages(); + const std::int32_t fh_host_before_ack = scheduler_->PagedCacheHostGroupAvailablePages("fh"); + const std::int32_t swa_host_before_ack = scheduler_->PagedCacheHostGroupAvailablePages("swa"); + ASSERT_GT(transferred_kv_host_pages, 0); + ASSERT_GT(transferred_host_pages.at("fh"), 0); + ASSERT_GT(transferred_host_pages.at("swa"), 0); + + SendWriteBackDone(writeback->op_ids[0], /*success=*/false); + PlanOnce(); + + EXPECT_EQ(scheduler_->AvailableHostKvPages(), kv_host_before_ack + transferred_kv_host_pages); + EXPECT_EQ(scheduler_->PagedCacheHostGroupAvailablePages("fh"), + fh_host_before_ack + transferred_host_pages.at("fh")); + EXPECT_EQ(scheduler_->PagedCacheHostGroupAvailablePages("swa"), + swa_host_before_ack + transferred_host_pages.at("swa")); + + Submit(RequestSpec{ + .request_id = "r2", + .tokens = request_tokens, + }); + auto reuse_plan = PlanOnce(); + EXPECT_EQ(GetLoadBack(reuse_plan), nullptr); + ASSERT_NE(GetForward(reuse_plan), nullptr); +} + +} // namespace tokenspeed::test diff --git a/tokenspeed-scheduler/tests/cpp/test_paged_cache_replay.cpp b/tokenspeed-scheduler/tests/cpp/test_paged_cache_replay.cpp index 64b667a9a..9e225c47f 100644 --- a/tokenspeed-scheduler/tests/cpp/test_paged_cache_replay.cpp +++ b/tokenspeed-scheduler/tests/cpp/test_paged_cache_replay.cpp @@ -392,6 +392,25 @@ TEST_F(PagedCacheTerminalContinuationTest, ExactTerminalHitUsesContinuationState EXPECT_EQ(second_match.paged_cache.per_group_page_ids.at(kRequiredStateGroup).size(), 2u); } +TEST_F(PagedCacheTerminalContinuationTest, StateRecoveryDoesNotUseEarlierContinuationCheckpoint) { + TreeNode* checkpoint = InsertDeviceTokens(256); + ASSERT_NE(checkpoint, nullptr); + CommitRequest("seed", /*first_token=*/0, /*target=*/256, checkpoint); + hybrid_->ReleaseRequest("seed"); + + TreeNode* terminal = InsertDeviceTokens(320); + ASSERT_NE(terminal, nullptr); + ASSERT_FALSE(terminal->HasPagedCacheSnapshot()); + + auto tokens = MakeAlignedTokens(/*num_pages=*/5, kPageSize, /*start=*/1); + auto match = hybrid_->Match(tokens, MatchIntent::StateRecovery); + + EXPECT_EQ(match.paged_cache.last_node, nullptr); + EXPECT_EQ(match.paged_cache_host.last_node, nullptr); + ASSERT_NE(match.device.last_node, nullptr); + EXPECT_TRUE(match.device.last_node->IsRoot()); +} + TEST_F(PagedCacheTerminalContinuationTest, SnapshotCursorStopsAtCheckpointBeforeReservedTail) { constexpr std::int32_t kCheckpoint = 256; constexpr std::int32_t kVerifyWidth = 4; @@ -530,7 +549,7 @@ TEST_F(PagedCacheTerminalContinuationTest, CurrentBorrowPinsAncestorWhenTerminal hybrid_->ReleaseRequest("current"); } -TEST_F(PagedCacheTerminalContinuationTest, StatePruneDropsContinuationAndFallsBackToColdPrefill) { +TEST_F(PagedCacheTerminalContinuationTest, StatePruneDropsContinuationAndDisablesPrefixReuse) { TreeNode* n256 = InsertDeviceTokens(256); ASSERT_NE(n256, nullptr); CommitRequest("r1", /*first_token=*/0, /*target=*/256, n256); @@ -549,9 +568,10 @@ TEST_F(PagedCacheTerminalContinuationTest, StatePruneDropsContinuationAndFallsBa n256->GetPagedCacheSnapshot()->groups.end()); auto match = MatchTokens(256); - EXPECT_EQ(match.paged_cache.history_hit_tokens, 0); + EXPECT_EQ(match.paged_cache.last_node, nullptr); EXPECT_EQ(match.paged_cache.prefix_len_tokens, 0); - EXPECT_TRUE(match.paged_cache.per_group_page_ids.empty()); + ASSERT_NE(match.device.last_node, nullptr); + EXPECT_TRUE(match.device.last_node->IsRoot()); } TEST_F(PagedCacheTerminalContinuationTest, StateOnlyPruneIgnoresHistoryOnlySideTableBorrow) { @@ -822,6 +842,7 @@ TEST_P(PagedCacheOverlapRetractTest, LateResultRecoveryRebuildsDynamicHorizon) { const auto state_pages_before = scheduler_->GetRequestPagedCachePageIds("r", "retract.state"); SendReserveNumTokens(/*value=*/100); auto retract_plan = PlanOnce(); + EXPECT_TRUE(retract_plan.SchedulerAborts().empty()); const auto* writeback = GetWriteBack(retract_plan); ASSERT_NE(writeback, nullptr); ASSERT_FALSE(writeback->op_ids.empty()); @@ -851,6 +872,7 @@ TEST_P(PagedCacheOverlapRetractTest, LateResultRecoveryRebuildsDynamicHorizon) { ASSERT_EQ(scheduler_->RetractedSize(), 1u); auto recovery_plan = PlanOnce(); + EXPECT_TRUE(recovery_plan.SchedulerAborts().empty()); const auto* recovery = GetForwardOp(recovery_plan); ASSERT_NE(recovery, nullptr); ASSERT_EQ(recovery->request_ids.size(), 1u); diff --git a/tokenspeed-scheduler/tests/cpp/test_retract.cpp b/tokenspeed-scheduler/tests/cpp/test_retract.cpp index 5b5830b8e..5e3307f74 100644 --- a/tokenspeed-scheduler/tests/cpp/test_retract.cpp +++ b/tokenspeed-scheduler/tests/cpp/test_retract.cpp @@ -68,6 +68,15 @@ class RetractTestSuite : public SchedulerTestSuite { } }; +class DisableL2RetractTestSuite : public RetractTestSuite { +protected: + SchedulerConfig MakeConfig() override { + auto cfg = RetractTestSuite::MakeConfig(); + cfg.disable_l2_cache = true; + return cfg; + } +}; + // Device full + extra page needed → WriteBack (Retract) op emitted. TEST_F(RetractTestSuite, Retract_TriggeredWhenDeviceFull) { BringToDecoding("r1"); @@ -88,6 +97,17 @@ TEST_F(RetractTestSuite, Retract_TriggeredWhenDeviceFull) { EXPECT_TRUE(any_pages); } +TEST_F(DisableL2RetractTestSuite, RetractStillEmitsHostWriteBackWhenL2Disabled) { + BringToDecoding("r1"); + SendReserveNumTokens("r1", 3); + + auto plan = PlanOnce(); + const auto* wb = GetWriteBack(plan); + ASSERT_NE(wb, nullptr); + EXPECT_FALSE(wb->op_ids.empty()); + EXPECT_EQ(scheduler_->AvailableHostKvPages(), Config().host_allocator.total_pages - 2); +} + // Retracting request must not appear in the forward batch. TEST_F(RetractTestSuite, Retract_RequestNotInForwardWhileRetracting) { BringToDecoding("r1"); @@ -129,6 +149,28 @@ TEST_F(RetractTestSuite, Retract_WriteBackDoneTransitionsToRetracted) { EXPECT_EQ(scheduler_->DecodingSize(), 0u); } +#if !TOKENSPEED_FLAT_KVCACHE +TEST_F(RetractTestSuite, FailedWriteBackAbortsAfterOutstandingForwardResult) { + BringToDecoding("r1"); + SendReserveNumTokens("r1", 3); + + auto retract_plan = PlanOnce(); + const auto* writeback = GetWriteBack(retract_plan); + ASSERT_NE(writeback, nullptr); + ASSERT_EQ(writeback->op_ids.size(), 1u); + + SendWriteBackDone(writeback->op_ids[0], /*success=*/false); + EXPECT_EQ(scheduler_->RetractedSize(), 0u); + EXPECT_TRUE(PlanOnce().SchedulerAborts().empty()); + + SendForwardDone("r1", {}); + auto abort_plan = PlanOnce(); + ASSERT_EQ(abort_plan.SchedulerAborts().size(), 1u); + EXPECT_EQ(abort_plan.SchedulerAborts().front().request_id, "r1"); + EXPECT_EQ(abort_plan.SchedulerAborts().front().message, "L2 retract write-back failed"); +} +#endif + // Retracted request recovers to Decoding on next PlanOnce. TEST_F(RetractTestSuite, Retract_RetractedRequestRecoversToDecoding) { BringToDecoding("r1"); diff --git a/tokenspeed-scheduler/tests/cpp/test_retract_abort_pages.cpp b/tokenspeed-scheduler/tests/cpp/test_retract_abort_pages.cpp index 4d1d0ba26..dc30295ed 100644 --- a/tokenspeed-scheduler/tests/cpp/test_retract_abort_pages.cpp +++ b/tokenspeed-scheduler/tests/cpp/test_retract_abort_pages.cpp @@ -89,6 +89,32 @@ class RetractAbortPagesSuite : public SchedulerTestSuite { } }; +class RetractHostExhaustionSuite : public RetractAbortPagesSuite { +protected: + SchedulerConfig MakeConfig() override { + auto cfg = RetractAbortPagesSuite::MakeConfig(); + cfg.host_allocator.total_pages = 1; + return cfg; + } +}; + +TEST_F(RetractHostExhaustionSuite, InternalAbortIsPublishedInExecutionPlan) { + // BringToDecoding schedules zero-token decode ops; settle their result debt before testing terminal abort. + BringToDecoding("r1", /*num_pages=*/2, /*start=*/1); + SendForwardDone("r1", {}); + BringToDecoding("r2", /*num_pages=*/1, /*start=*/5); + SendForwardDone("r1", {}); + SendForwardDone("r2", {}); + SendReserveNumTokens("r1", 2); + SendReserveNumTokens("r2", 1); + + auto plan = PlanOnce(); + + ASSERT_EQ(plan.SchedulerAborts().size(), 1u); + EXPECT_EQ(plan.SchedulerAborts().front().request_id, "r1"); + EXPECT_NE(plan.SchedulerAborts().front().message.find("host cache capacity"), std::string::npos); +} + // Core scenario: Decoding request gets retracted, then aborted. // Before OwnedPages: ~LocalKVAllocator would double-free pages already in the tree. // After OwnedPages: TakeFullPages moved pages out before Insert; abort is clean. diff --git a/tokenspeed-scheduler/tests/cpp/test_write_back.cpp b/tokenspeed-scheduler/tests/cpp/test_write_back.cpp index 100963741..8bb3605d6 100644 --- a/tokenspeed-scheduler/tests/cpp/test_write_back.cpp +++ b/tokenspeed-scheduler/tests/cpp/test_write_back.cpp @@ -46,6 +46,12 @@ class WriteBackTestSuite : public SchedulerTestSuite { scheduler_->Advance(std::move(event)); } + void SendAbort(const std::string& id) { + ExecutionEvent event; + event.With(ForwardEvent{forward::Abort{.request_id = id}}); + scheduler_->Advance(std::move(event)); + } + // Submitted → PrefillDone → Decoding. // decoding_peers: requests already in Decoding that need reserve set before the next PlanOnce. void BringToDecoding(const std::string& id, std::int32_t num_pages, token_t start = 1, @@ -129,6 +135,41 @@ TEST_F(WriteBackTestSuite, WriteBack_FinishedAfterWriteBackDone) { EXPECT_EQ(scheduler_->DecodingSize(), 0u); } +#if !TOKENSPEED_FLAT_KVCACHE +TEST_F(WriteBackTestSuite, AbortBeforeWriteBackRollsBackPreparedHostState) { + BringToDecoding("r1", /*num_pages=*/2); + const std::size_t host_pages_before_finish = scheduler_->AvailableHostKvPages(); + + SendFinish("r1"); + ASSERT_LT(scheduler_->AvailableHostKvPages(), host_pages_before_finish); + + SendAbort("r1"); + EXPECT_EQ(scheduler_->AvailableHostKvPages(), host_pages_before_finish); + EXPECT_EQ(GetWriteBack(PlanOnce()), nullptr); +} + +TEST_F(WriteBackTestSuite, AbortDuringWriteBackWaitsForAckAndDiscardsHostState) { + BringToDecoding("r1", /*num_pages=*/2); + const std::size_t host_pages_before_finish = scheduler_->AvailableHostKvPages(); + SendFinish("r1"); + + auto writeback_plan = PlanOnce(); + const auto* writeback = GetWriteBack(writeback_plan); + ASSERT_NE(writeback, nullptr); + ASSERT_EQ(writeback->op_ids.size(), 1u); + const std::size_t host_pages_during_writeback = scheduler_->AvailableHostKvPages(); + ASSERT_LT(host_pages_during_writeback, host_pages_before_finish); + + SendAbort("r1"); + EXPECT_EQ(scheduler_->AvailableHostKvPages(), host_pages_during_writeback); + EXPECT_EQ(GetWriteBack(PlanOnce()), nullptr); + + SendWriteBackDone(writeback->op_ids[0]); + EXPECT_EQ(scheduler_->AvailableHostKvPages(), host_pages_before_finish); + EXPECT_TRUE(PlanOnce().SchedulerAborts().empty()); +} +#endif + // Two requests finish in the same round → two separate WriteBack ops. TEST_F(WriteBackTestSuite, WriteBack_MultipleRequestsGetSeparateOps) { BringToDecoding("r1", /*num_pages=*/2);