vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker ¶

class KVCacheStoreRecvingThread(KVTransferThread):
    """Background thread for loading KV cache blocks from the store."""

    def __init__(
        self,
        store: Any,
        coord: MooncakeStoreCoordinator,
        token_databases: list[ChunkedTokenDatabase],
        block_size: int,
        tp_rank: int,
        ready_event: threading.Event,
        disk_offload_buffer_budget_bytes: int | None = None,
        record_operation: Callable[..., None] | None = None,
        request_queue: queue.Queue[Any] | None = None,
    ):
        super().__init__(
            store,
            token_databases,
            block_size,
            tp_rank,
            ready_event,
            name="KVCacheStoreRecvingThread",
            record_operation=record_operation,
            request_queue=request_queue,
        )
        # _invalid_block_ids can be access by both the Worker and RecvingThread
        self._invalid_block_ids_lock = threading.Lock()
        self._invalid_block_ids: set[int] = set()
        self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
        self.usable_disk_offload_buffer_budget_bytes = (
            None
            if disk_offload_buffer_budget_bytes is None
            else _get_usable_disk_offload_buffer_budget_bytes(
                disk_offload_buffer_budget_bytes
            )
        )
        self.coord = coord

    def _add_load_error_block_ids(self, block_ids: list[int]) -> None:
        with self._invalid_block_ids_lock:
            self._invalid_block_ids.update(block_ids)

    def get_and_clear_block_ids_with_load_errors(self) -> set[int]:
        with self._invalid_block_ids_lock:
            invalid_block_ids = self._invalid_block_ids.copy()
            self._invalid_block_ids.clear()
        return invalid_block_ids

    def _handle_request(self, req_meta: ReqMeta):
        token_len = req_meta.load_spec.token_len  # type: ignore[union-attr]
        req_id = req_meta.req_id
        mask_num = (
            req_meta.load_spec.vllm_cached_tokens  # type: ignore[union-attr]
            // self.block_size
            * self.block_size
        )

        # Skip chunks the consumer's per-group spec wouldn't populate
        # locally (e.g. SWA pre-window) even if the producer stored them.
        load_mask_per_group = self.coord.load_mask(req_meta.block_hashes, token_len)

        addr_list: list[list[int]] = []
        size_list: list[list[int]] = []
        key_list: list[str] = []
        block_id_list: list[int] = []
        for g_idx, db in enumerate(self.token_databases):
            mask = load_mask_per_group[g_idx]
            chunks: list[tuple[int, int]] = []
            for start, end, block_hash in db.process_tokens(
                token_len, req_meta.block_hashes, mask_num
            ):
                chunk_idx = start // db.block_size
                if chunk_idx >= len(mask) or not mask[chunk_idx]:
                    continue
                key_list.append(db.key_for(block_hash))
                chunks.append((start, end))
            g_addrs, g_sizes, g_block_ids = db.prepare_values(
                chunks, req_meta.block_ids[g_idx]
            )
            addr_list.extend(g_addrs)
            size_list.extend(g_sizes)
            block_id_list.extend(g_block_ids)

        # Rotate aligned lists by tp_rank for load balancing.
        rotation = self.tp_rank % len(key_list)
        key_list_c = _rotate_list(key_list, rotation)
        addr_list_c = _rotate_list(addr_list, rotation)
        size_list_c = _rotate_list(size_list, rotation)
        block_id_list_c = _rotate_list(block_id_list, rotation)

        load_batches = [(key_list_c, addr_list_c, size_list_c, block_id_list_c)]
        if self.usable_disk_offload_buffer_budget_bytes is not None:
            total_staging_bytes = sum(
                _estimate_disk_offload_staging_bytes(size) for size in size_list_c
            )
            if total_staging_bytes > self.usable_disk_offload_buffer_budget_bytes:
                assert self.disk_offload_buffer_budget_bytes is not None
                split_batches, oversized_key = _split_disk_offload_load_batches(
                    key_list_c,
                    addr_list_c,
                    size_list_c,
                    self.usable_disk_offload_buffer_budget_bytes,
                    self.disk_offload_buffer_budget_bytes,
                )
                if oversized_key is not None:
                    oversized_key_index = key_list_c.index(oversized_key)
                    # Mark every block: we skip the whole request, and the
                    # tp_rank rotation means oversized_key isn't necessarily
                    # the first block in the request's original order.
                    self._add_load_error_block_ids(block_id_list_c)
                    oversized_key_bytes = _estimate_disk_offload_staging_bytes(
                        size_list_c[oversized_key_index]
                    )
                    logger.warning(
                        "Skipping Mooncake load for request %s because key %s "
                        "requires %d staging bytes, exceeding budget %d",
                        req_id,
                        oversized_key,
                        oversized_key_bytes,
                        self.disk_offload_buffer_budget_bytes,
                    )
                    self.set_finished_request(req_id)
                    self.request_queue.task_done()
                    return
                load_batches = []
                block_id_offset = 0
                for batch_keys, batch_addrs, batch_sizes in split_batches:
                    next_block_id_offset = block_id_offset + len(batch_keys)
                    batch_block_ids = block_id_list_c[
                        block_id_offset:next_block_id_offset
                    ]
                    load_batches.append(
                        (batch_keys, batch_addrs, batch_sizes, batch_block_ids)
                    )
                    block_id_offset = next_block_id_offset

        current_batch_keys: list[str] = key_list_c
        current_batch_block_ids: list[int] = block_id_list_c
        batch_bytes = 0
        try:
            for batch_keys, batch_addrs, batch_sizes, batch_block_ids in load_batches:
                current_batch_keys = batch_keys
                current_batch_block_ids = batch_block_ids
                batch_bytes = _sum_batch_bytes(batch_sizes)
                tiers_by_key: dict[str, str] | None = None
                if envs.VLLM_MOONCAKE_STORE_TIER_LOG:
                    tiers_by_key = _get_replica_tiers_by_key(self.store, batch_keys)
                # Reset so the recorded RPC duration excludes tier lookup.
                load_get_start = time.perf_counter()
                res = self.store.batch_get_into_multi_buffers(
                    batch_keys, batch_addrs, batch_sizes
                )
                if tiers_by_key is not None:
                    _log_mooncake_load_tier_summary(
                        req_id, batch_keys, res, tiers_by_key
                    )
                failed = [
                    (key, value, block_id)
                    for key, value, block_id in zip(
                        batch_keys, res, batch_block_ids, strict=True
                    )
                    if value < 0
                ]
                self._record_operation(
                    "load_get",
                    load_get_start,
                    len(batch_keys),
                    num_bytes=batch_bytes,
                    status="partial_failure" if failed else "ok",
                    num_failed_keys=len(failed),
                )
                if failed:
                    self._add_load_error_block_ids(
                        [block_id for _, _, block_id in failed]
                    )
                    logger.warning(
                        "Failed to get %d Mooncake keys from sub-batch "
                        "(batch_keys=%d, first_failures=%s)",
                        len(failed),
                        len(batch_keys),
                        [(key, value) for key, value, _ in failed[:3]],
                    )
                    break
        except Exception as e:
            self._add_load_error_block_ids(current_batch_block_ids)
            self._record_operation(
                "load_get",
                load_get_start,
                len(current_batch_keys),
                num_bytes=batch_bytes,
                status="error",
                num_failed_keys=len(current_batch_keys),
            )
            logger.warning(
                "Failed to get Mooncake sub-batch %s, error: %s",
                current_batch_keys[:3],
                e,
            )

        self.set_finished_request(req_id)
        self.request_queue.task_done()

class KVCacheStoreSendingThread(KVTransferThread):
    """Background thread for storing KV cache blocks to the store."""

    def __init__(
        self,
        store: Any,
        coord: MooncakeStoreCoordinator,
        token_databases: list[ChunkedTokenDatabase],
        block_size: int,
        tp_rank: int,
        group_put_steps: Sequence[int],
        kv_role: str,
        ready_event: threading.Event,
        enable_kv_event: bool = False,
        replicate_config: Any = None,
        enable_group_semantics: bool = False,
        supports_group_ids: bool = False,
        record_operation: Callable[..., None] | None = None,
    ):
        super().__init__(
            store,
            token_databases,
            block_size,
            tp_rank,
            ready_event,
            name="KVCacheStoreSendingThread",
            record_operation=record_operation,
        )
        # Only ranks with identical group bytes may stripe PUTs (e.g., MLA).
        self.group_put_steps = group_put_steps
        self.coord = coord
        self.kv_role = kv_role
        self.stored_requests: defaultdict[str, int] = defaultdict(int)
        self.enable_kv_event = enable_kv_event
        # Caller always passes a non-None ReplicateConfig — see
        # MooncakeStoreWorker.__init__ where store_replicate_config is built.
        self.replicate_config = replicate_config
        self.enable_group_semantics = enable_group_semantics
        self.supports_group_ids = supports_group_ids

        # Pause store requests when CPU/disk offloading is under pressure.
        self._store_pressure_active = False
        self._skip_store_requests: set[str] = set()

        # Per-request high-water mark of tokens actually persisted; the next
        # batch resumes here, so pressure-skipped or failed ranges are retried.
        self._saved_offset: dict[str, int] = {}

    def add_stored_request(self, req_id: str):
        with self.done_task_lock:
            self.stored_requests[req_id] += 1

    def dec_stored_request(self, req_id: str):
        with self.done_task_lock:
            if req_id in self.stored_requests:
                self.stored_requests[req_id] -= 1

    def delete_finished_stored_request(self, req_id: str):
        with self.done_task_lock:
            if req_id in self.stored_requests:
                del self.stored_requests[req_id]
            self._skip_store_requests.discard(req_id)
            self._saved_offset.pop(req_id, None)

    def _record_saved(self, req_id: str, token_len: int) -> None:
        # Guard on liveness so a concurrent finish/preempt pop isn't recreated.
        with self.done_task_lock:
            if req_id in self.stored_requests:
                self._saved_offset[req_id] = token_len

    def _should_skip_request(self, req_id: str) -> bool:
        with self.done_task_lock:
            return self._store_pressure_active and req_id in self._skip_store_requests

    def _mark_request_skipped_for_pressure(self, req_id: str) -> bool:
        with self.done_task_lock:
            already_skipped = req_id in self._skip_store_requests
            self._store_pressure_active = True
            self._skip_store_requests.add(req_id)
        return already_skipped

    def _clear_store_pressure(self) -> bool:
        with self.done_task_lock:
            if not self._store_pressure_active and not self._skip_store_requests:
                return False
            self._store_pressure_active = False
            self._skip_store_requests.clear()
        return True

    def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool:
        """Offload the request's sub-block partial tail (its last prompt hash
        boundary) so a later request can hit the sub-block prefix.

        Covers every block from the normal save's lcm floor to the boundary:
        the normal save floors to ``lcm_block_size``, so a smaller-block
        group's full blocks in that gap are never persisted elsewhere, and
        the consumer's lookup needs every group at every probed boundary.
        Full blocks are keyed by their block-end hash, the partial boundary
        block by the boundary sub-hash; the mamba "align" boundary block is
        the core-provided CoW block. All keys are deduped against the store.

        Returns:
            True when no put is needed or every put succeeds, False otherwise.
        """
        if not self.coord.enable_partial_hash_hits or not req_meta.block_hashes:
            return True
        partial_tail_offloads = req_meta.partial_tail_offloads
        if not partial_tail_offloads:
            return True
        hash_block_size = self.coord.hash_block_size
        boundaries = {boundary for _, _, boundary in partial_tail_offloads}
        if len(boundaries) != 1:
            raise ValueError(
                "Partial-tail offloads for one request must share a boundary"
            )
        boundary = boundaries.pop()
        if boundary == 0:
            return True
        if boundary // hash_block_size - 1 >= len(req_meta.block_hashes):
            return True
        mamba_offloads = {
            group_id: block_id for group_id, block_id, _ in partial_tail_offloads
        }

        keys: list[str] = []
        addrs: list[list[int]] = []
        sizes: list[list[int]] = []
        group_ids: list[str] | None = (
            [] if self.enable_group_semantics and self.supports_group_ids else None
        )
        saved = self._saved_offset.get(req_meta.req_id, 0)
        for g_idx, db in enumerate(self.token_databases):
            group_blocks = req_meta.block_ids[g_idx]
            # Distribute across ranks by the same rule as normal chunks.
            put_step = self.group_put_steps[g_idx]
            put_step_rank = (self.tp_rank + g_idx) % put_step
            # Always include the boundary block: its sub-hash key is written
            # only here, even if normal saves already advanced past it.
            last_block = cdiv(boundary, db.block_size) - 1
            for block_idx in range(
                min(saved // db.block_size, last_block), last_block + 1
            ):
                if block_idx % put_step != put_step_rank:
                    continue
                valid_end = min((block_idx + 1) * db.block_size, boundary)
                key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1]
                if (
                    g_idx in mamba_offloads
                    and valid_end == boundary
                    and boundary % db.block_size != 0
                ):
                    block_id = mamba_offloads[g_idx]
                else:
                    if block_idx >= len(group_blocks):
                        continue
                    block_id = group_blocks[block_idx]
                if block_id == NULL_BLOCK_ID:
                    logger.debug(
                        "Skipping unavailable partial-tail source block "
                        "(req=%s, group=%d, block=%d)",
                        req_meta.req_id,
                        g_idx,
                        block_idx,
                    )
                    continue
                addr, size = db.prepare_value_for_block(block_id)
                key = db.key_for(key_hash)
                keys.append(key)
                addrs.append(addr)
                sizes.append(size)
                if group_ids is not None:
                    group_ids.append(
                        _make_mooncake_group_id(
                            db.metadata,
                            key.rsplit("@", 1)[-1],
                        )
                    )

        if not keys:
            return True
        exists_start = time.perf_counter()
        try:
            exists = self.store.batch_is_exist(keys)
        except Exception as e:
            self._record_operation(
                "save_exists",
                exists_start,
                len(keys),
                status="error",
                num_failed_keys=len(keys),
            )
            logger.error(
                "Failed to check partial-tail keys for request %s: %s",
                req_meta.req_id,
                e,
            )
            return False
        self._record_operation("save_exists", exists_start, len(keys))
        missing = [i for i, e in enumerate(exists) if e != 1]
        if not missing:
            return True
        keys = [keys[i] for i in missing]
        addrs = [addrs[i] for i in missing]
        sizes = [sizes[i] for i in missing]
        if group_ids is not None:
            group_ids = [group_ids[i] for i in missing]
        if req_meta.current_event is not None:
            # Fence the CoW block copy enqueued earlier this step.
            req_meta.current_event.synchronize()
        if group_ids is not None:
            assert len(group_ids) == len(keys)
            self.replicate_config.group_ids = group_ids
        batch_bytes = _sum_batch_bytes(sizes)
        put_start = time.perf_counter()
        try:
            res = self.store.batch_put_from_multi_buffers(
                keys, addrs, sizes, self.replicate_config
            )
        except Exception as e:
            self._record_operation(
                "save_put",
                put_start,
                len(keys),
                num_bytes=batch_bytes,
                status="error",
                num_failed_keys=len(keys),
            )
            logger.error(
                "Failed to put partial-tail keys for request %s: %s",
                req_meta.req_id,
                e,
            )
            return False

        failed = [i for i, value in enumerate(res) if value < 0]
        self._record_operation(
            "save_put",
            put_start,
            len(keys),
            num_bytes=batch_bytes,
            status="partial_failure" if failed else "ok",
            num_failed_keys=len(failed),
        )
        if failed:
            failed_codes = {res[i] for i in failed}
            logger.warning(
                "Partial-tail put failed for request %s: %d/%d keys failed (codes=%s)",
                req_meta.req_id,
                len(failed),
                len(keys),
                failed_codes,
            )
            if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes:
                self._mark_request_skipped_for_pressure(req_meta.req_id)
            return False

        if self._clear_store_pressure():
            logger.info(
                "Mooncake CPU/disk offloading pressure cleared after a "
                "successful partial-tail batch"
            )
        return True

    def _handle_request(self, req_meta: ReqMeta):
        # Cache hits are always a multiple of ``lcm_block_size`` tokens, which
        # is also ``store_mask``'s precondition.
        lcm_block_size = self.coord.lcm_block_size
        token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size
        block_ids_per_group = req_meta.block_ids
        req_id = req_meta.req_id
        current_event = req_meta.current_event

        if req_id not in self.stored_requests:
            self.request_queue.task_done()
            return

        # Decrement the in-flight counter and signal task_done() in `finally`
        # so the scheduler can release the GPU blocks it pinned for this
        # request (via `delay_free_blocks`) even when the store path raises.
        try:
            if self._should_skip_request(req_id):
                logger.debug(
                    "Skipping Mooncake store for request %s while CPU/disk "
                    "offloading is under pressure",
                    req_id,
                )
                return

            # Offload the sub-block partial tail (independent of the normal
            # block-aligned save, which may be skipped this step).
            if req_meta.partial_tail_offloads is not None and not (
                self._maybe_offload_partial_tail(req_meta)
            ):
                return

            if token_len == 0:
                return

            # Resume from where this rank left off; only the new suffix is saved.
            save_start = self._saved_offset.get(req_id, 0)

            # Within each lcm region only per-spec relevant chunks are loaded
            # (e.g., SWA or linear attn), so mask out irrelevant chunks
            store_masks = self.coord.store_mask(
                token_len,
                save_start,
                num_prompt_tokens=req_meta.num_prompt_tokens,
            )

            starts: list[int] = []
            ends: list[int] = []
            keys: list[str] = []
            kv_event_block_hashes: list[BlockHash] = []
            group_indices: list[int] = []
            for g_idx, db in enumerate(self.token_databases):
                # Rotate the stride phase per group to balance load across ranks.
                put_step = self.group_put_steps[g_idx]
                put_step_rank = (self.tp_rank + g_idx) % put_step
                for start, end, block_hash in db.process_tokens(
                    token_len,
                    req_meta.block_hashes,
                    mask_num=save_start,
                    chunk_mask=store_masks[g_idx],
                    put_step=put_step,
                    put_step_rank=put_step_rank,
                ):
                    starts.append(start)
                    ends.append(end)
                    keys.append(db.key_for(block_hash))
                    if self.enable_kv_event:
                        kv_event_block_hashes.append(block_hash)
                    group_indices.append(g_idx)

            if not keys:
                self._record_saved(req_id, token_len)
                return

            # Check which blocks already exist (dedup)
            save_exists_start = time.perf_counter()
            try:
                exists_states = self.store.batch_is_exist(keys)
            except Exception:
                self._record_operation(
                    "save_exists",
                    save_exists_start,
                    len(keys),
                    status="error",
                    num_failed_keys=len(keys),
                )
                raise
            self._record_operation(
                "save_exists",
                save_exists_start,
                len(keys),
            )
            missing_indices = [
                i for i, exists in enumerate(exists_states) if exists != 1
            ]

            if not missing_indices:
                self._record_saved(req_id, token_len)
                return

            if len(missing_indices) != len(keys):
                starts = [starts[i] for i in missing_indices]
                ends = [ends[i] for i in missing_indices]
                keys = [keys[i] for i in missing_indices]
                if self.enable_kv_event:
                    kv_event_block_hashes = [
                        kv_event_block_hashes[i] for i in missing_indices
                    ]
                group_indices = [group_indices[i] for i in missing_indices]

            group_ids = (
                [
                    _make_mooncake_group_id(
                        self.token_databases[g_idx].metadata,
                        key.rsplit("@", 1)[-1],
                    )
                    for key, g_idx in zip(keys, group_indices, strict=True)
                ]
                if self.enable_group_semantics and self.supports_group_ids
                else None
            )

            logger.debug(
                "Storing KV cache for %d blocks (groups=%s) for request %s",
                len(keys),
                set(group_indices),
                req_id,
            )

            addrs: list[list[int]] = []
            sizes: list[list[int]] = []
            stored_events: list[BlockStored] = []
            chunks_per_group: list[list[tuple[int, int]]] = [
                [] for _ in self.token_databases
            ]
            for start, end, g_idx in zip(starts, ends, group_indices, strict=True):
                chunks_per_group[g_idx].append((start, end))
            for g_idx, chunks in enumerate(chunks_per_group):
                if not chunks:
                    continue
                db = self.token_databases[g_idx]
                group_addrs, group_sizes, _ = db.prepare_values(
                    chunks, block_ids_per_group[g_idx]
                )
                addrs.extend(group_addrs)
                sizes.extend(group_sizes)

            # parent_block_hash chains live within a group, not across.
            if self.enable_kv_event:
                prev_key_per_group: dict[int, Any] = {}
                new_block_hashes = [
                    maybe_convert_block_hash(bh) for bh in kv_event_block_hashes
                ]

            for idx, (s, e, g_idx) in enumerate(
                zip(starts, ends, group_indices, strict=True)
            ):
                db = self.token_databases[g_idx]
                if self.enable_kv_event:
                    token_ids = (
                        req_meta.token_ids[s:e]
                        if req_meta.token_ids is not None
                        else None
                    )
                    stored_event = BlockStored(
                        block_hashes=[new_block_hashes[idx]],
                        parent_block_hash=prev_key_per_group.get(g_idx),
                        token_ids=token_ids,
                        block_size=db.block_size,
                        lora_id=None,
                        medium="cpu",
                        lora_name=None,
                        group_idx=g_idx,
                    )
                    stored_events.append(stored_event)
                    prev_key_per_group[g_idx] = new_block_hashes[idx]

            if current_event is not None:
                current_event.synchronize()

            if group_ids is not None:
                assert len(group_ids) == len(keys)
                self.replicate_config.group_ids = group_ids

            batch_bytes = _sum_batch_bytes(sizes)
            put_start = time.perf_counter()
            try:
                res = self.store.batch_put_from_multi_buffers(
                    keys,
                    addrs,
                    sizes,
                    self.replicate_config,
                )
                failed = [i for i, v in enumerate(res) if v < 0]
                self._record_operation(
                    "save_put",
                    put_start,
                    len(keys),
                    num_bytes=batch_bytes,
                    status="partial_failure" if failed else "ok",
                    num_failed_keys=len(failed),
                )
                if failed:
                    failed_codes = set(res[i] for i in failed)
                    logger.warning(
                        "batch_put failed: %d/%d keys failed "
                        "(codes=%s, batch_bytes=%d, num_keys=%d), "
                        "first_key=%s",
                        len(failed),
                        len(keys),
                        failed_codes,
                        batch_bytes,
                        len(keys),
                        keys[0] if keys else "N/A",
                    )
                    if (
                        MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes
                        and not self._mark_request_skipped_for_pressure(req_id)
                    ):
                        logger.warning(
                            "Detected Mooncake CPU/disk offloading pressure "
                            "(NO_AVAILABLE_HANDLE); skipping future store "
                            "batches for request %s until a later store "
                            "batch succeeds",
                            req_id,
                        )
                else:
                    self._record_saved(req_id, token_len)
                    if self._clear_store_pressure():
                        logger.info(
                            "Mooncake CPU/disk offloading pressure cleared "
                            "after a successful store batch"
                        )
            except Exception as e:
                self._record_operation(
                    "save_put",
                    put_start,
                    len(keys),
                    num_bytes=batch_bytes,
                    status="error",
                    num_failed_keys=len(keys),
                )
                logger.error("Failed to put key %s, error: %s", keys, e)

            if self.enable_kv_event and stored_events:
                self.update_kv_event(stored_events)
        finally:
            self.dec_stored_request(req_id)
            self.request_queue.task_done()