init

vllm_vacc/vllm/v1/core/block_pool.py

@@ -0,0 +1,78 @@
+
+
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import Callable, Optional
+
+from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
+                                        BlockStored, KVCacheEvent)
+from vllm.logger import init_logger
+
+from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
+                                         FreeKVCacheBlockQueue, KVCacheBlock,
+                                         generate_block_hash_extra_keys,
+                                         hash_block_tokens)
+from vllm.v1.request import Request
+from vllm.v1.core.block_pool import BlockHashToBlockMap
+
+logger = init_logger(__name__)
+
+class BlockPool:
+    """BlockPool that manages KVCacheBlocks.
+    It provides methods to allocate, free and cache the kv cache blocks. The
+    free_block_queue stores the free blocks in eviction order to enable
+    allocation, free, and cache eviction. The cached_block_hash_to_block
+    maps between block hash and cached block to support finding cached blocks
+    by their block hash.
+
+    Args:
+        num_gpu_blocks: The number of blocks in the pool.
+        enable_caching: Whether to enable prefix caching.
+        enable_kv_cache_events: Whether to enable kv cache events.
+    """
+
+    def __init__(
+        self,
+        num_gpu_blocks: int,
+        enable_caching: bool,
+        enable_kv_cache_events: bool = False,
+    ):
+        assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
+        self.num_gpu_blocks = num_gpu_blocks
+        self.enable_caching = enable_caching
+        # All kv-cache blocks.
+        self.blocks: list[KVCacheBlock] = [
+            KVCacheBlock(idx) for idx in range(num_gpu_blocks)
+        ]
+        # Free block queue that constructs and manipulates a doubly linked
+        # list of free blocks (including eviction candidates when caching is
+        # enabled).
+        self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
+
+        # Cache for block lookup
+        self.cached_block_hash_to_block: BlockHashToBlockMap = \
+            BlockHashToBlockMap()
+
+        # To represent a placeholder block with block_id=0.
+        # The ref_cnt of null_block is not maintained, needs special care to
+        # avoid freeing it.
+        # self.null_block = self.free_block_queue.popleft()
+        # self.null_block.is_null = True
+        self.null_block = self.free_block_queue.fake_free_list_head.next_free_block #self.free_block_queue.popleft()
+        self.null_block.is_null = False
+
+        self.enable_kv_cache_events = enable_kv_cache_events
+        self.kv_event_queue: list[KVCacheEvent] = []
+    
+    def get_usage(self) -> float:
+        """Get the KV cache usage.
+
+        Returns:
+            The KV cache usage (between 0.0 and 1.0).
+        """
+
+        total_gpu_blocks = self.num_gpu_blocks
+        if not total_gpu_blocks:
+            return 0
+        
+        return 1.0 - (self.get_num_free_blocks() / total_gpu_blocks)

vllm_vacc/vllm/v1/core/kv_cache_utils.py

@@ -0,0 +1,156 @@
+
+from vllm.config import VllmConfig
+from vllm.utils import GiB_bytes
+from vllm.v1.core.kv_cache_utils import logger
+from vllm_vacc.vllm.model_executor.models.vars import BLOCK_GROUP_SIZE as env_blk_grp_size
+from vllm.v1.kv_cache_interface import KVCacheSpec
+
+def query_device_avaliable_memory(vllm_config):
+
+    import os    
+    import torch
+    import torch_vacc
+
+    available_kv_cache_memory= int(os.getenv("VLLM_VACC_KVCACHE_SPACE", "16")) * GiB_bytes
+    if available_kv_cache_memory ==0:
+        torch.vacc.empty_cache()
+        torch.vacc.reset_peak_memory_stats()
+        total_memory = torch.vacc.mem_get_info()[1]
+        torch.vacc.synchronize()
+        peak_memory = torch.vacc.max_memory_allocated()
+        torch.vacc.empty_cache()
+        torch_allocated_bytes = torch.vacc.memory_stats(
+        )["allocated_bytes.all.current"]
+        total_allocated_bytes = torch.vacc.mem_get_info(
+        )[1] - torch.vacc.mem_get_info()[0]
+        non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
+        if non_torch_allocations > 0:
+            peak_memory += non_torch_allocations
+        available_kv_cache_memory=total_memory * vllm_config.cache_config.gpu_memory_utilization - peak_memory
+
+    return available_kv_cache_memory
+
+
+def get_num_blocks(vllm_config: VllmConfig, num_layers: int,
+                   available_memory: int, page_size: int) -> int:
+    """
+    Get the number of kv cache blocks.
+
+    Args:
+        vllm_config: The global VllmConfig
+        num_layers: The number of layers
+        available_memory: Memory available for KV cache in bytes.
+        page_size: The page size of the KV cache.
+    """
+    num_blocks = int(available_memory // page_size // num_layers)
+    num_blocks = max(num_blocks, 0)
+    if vllm_config.cache_config.num_gpu_blocks_override is not None:
+        num_gpu_blocks_override = \
+            vllm_config.cache_config.num_gpu_blocks_override
+        logger.info(
+            "Overriding num_gpu_blocks=%d with "
+            "num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)
+        num_blocks = num_gpu_blocks_override
+    
+    block_num_per_group = env_blk_grp_size // 16
+    num_blocks = num_blocks // block_num_per_group * block_num_per_group
+
+    return num_blocks
+
+def estimate_max_model_len(vllm_config: VllmConfig,
+                           kv_cache_spec: dict[str, KVCacheSpec],
+                           available_memory: int) -> int:
+    """
+    Estimates the maximum model length that can fit in the available memory
+    using binary search.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+        available_memory: Memory available for KV cache in bytes.
+
+    Returns:
+        The estimated maximum model length that can fit in the available memory.
+    """
+    from vllm.v1.core.kv_cache_utils import max_memory_usage_bytes
+    available_kv_cache_memory = query_device_avaliable_memory(vllm_config)
+
+    # Define a function to check if a given model length fits in memory
+    def fits_in_memory(model_len: int) -> bool:
+        # Modify the max_model_len for this calculation
+        vllm_config.model_config.max_model_len = model_len
+        # Calculate memory needed for the given model length
+        memory_needed = max_memory_usage_bytes(vllm_config,
+                                               kv_cache_spec.values())
+        # 增加vacc buffer 判断,与原始设计不一致
+        # 如果所需memory 小于总的预分配空间/剩余物理空间，支持其继续分配，
+        # 无需受到score模型max_model_len个预分配空间限制
+        return memory_needed <= max(available_memory, available_kv_cache_memory)
+
+    # Binary search for the maximum model length
+    current_max = vllm_config.model_config.max_model_len
+    left, right = 1, current_max
+
+    # If even the smallest model length doesn't fit, return 0
+    if not fits_in_memory(left):
+        return 0
+
+    # Binary search for the maximum model length that fits
+    result = 1
+    while left <= right:
+        mid = (left + right) // 2
+        if fits_in_memory(mid):
+            result = mid
+            left = mid + 1
+        else:
+            right = mid - 1
+    return result
+
+def check_enough_kv_cache_memory(vllm_config: VllmConfig,
+                                 kv_cache_spec: dict[str, KVCacheSpec],
+                                 available_memory: int):
+    """
+    Checks whether `available_memory` is enough for the KV cache to hold at
+    least one request with the model's max_model_len.
+
+    Args:
+        vllm_config: The global VllmConfig
+        kv_cache_spec: The kv cache spec of each attention layer in the model
+        available_memory: Memory available for KV cache in bytes.
+
+    Raises:
+        ValueError: If there is not enough memory available for the KV cache.
+    """
+    from vllm.v1.core.kv_cache_utils import max_memory_usage_bytes
+
+    if available_memory <= 0:
+        raise ValueError("No available memory for the cache blocks. "
+                         "Try increasing `gpu_memory_utilization` when "
+                         "initializing the engine.")
+
+    max_model_len = vllm_config.model_config.max_model_len
+    needed_memory = max_memory_usage_bytes(vllm_config, kv_cache_spec.values())
+    available_kv_cache_memory = query_device_avaliable_memory(vllm_config)
+
+    # 增加vacc buffer 判断,与原始设计不一致
+    # 如果所需memory 小于总的预分配空间/剩余物理空间，支持其继续分配，
+    # 无需受到score模型max_model_len个预分配空间限制
+    if needed_memory > max(available_memory, available_kv_cache_memory):
+        # Estimate the maximum model length that can fit in the available memory
+        estimated_max_len = estimate_max_model_len(vllm_config, kv_cache_spec,
+                                                   available_memory)
+        print("exec max model len is:", estimated_max_len)
+        estimated_msg = ""
+        if estimated_max_len > 0:
+            estimated_msg = (
+                "Based on the available memory, "
+                f"the estimated maximum model length is {estimated_max_len}.")
+
+        raise ValueError(
+            f"To serve at least one request with the models's max seq len "
+            f"({max_model_len}), ({needed_memory/GiB_bytes:.2f} GiB KV "
+            f"cache is needed, which is larger than the available KV cache "
+            f"memory ({available_memory/GiB_bytes:.2f} GiB). "
+            f"{estimated_msg} "
+            f"Try increasing `gpu_memory_utilization` or decreasing "
+            f"`max_model_len` when initializing the engine.")

vllm_vacc/vllm/v1/core/sched/output.py

@@ -0,0 +1,89 @@
+
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from vllm._bc_linter import bc_linter_include
+
+if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
+    import torch
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+        KVConnectorMetadata)
+    from vllm.lora.request import LoRARequest
+    from vllm.multimodal.inputs import MultiModalFeatureSpec
+    from vllm.pooling_params import PoolingParams
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.request import Request
+    
+    
+
+@bc_linter_include
+@dataclass
+class NewRequestData:
+
+    req_id: str
+    prompt_token_ids: Optional[list[int]]
+    mm_features: list[MultiModalFeatureSpec]
+    sampling_params: Optional[SamplingParams]
+    pooling_params: Optional[PoolingParams]
+    block_ids: tuple[list[int], ...]
+    num_computed_tokens: int
+    lora_request: Optional[LoRARequest]
+    prompt_embeds: Optional[torch.Tensor] = None
+    deepstack_input_embeds: Optional[torch.Tensor] = None #patch
+
+    @classmethod
+    def from_request(
+        cls,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> NewRequestData:
+        return cls(
+            req_id=request.request_id,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            block_ids=block_ids,
+            num_computed_tokens=request.num_computed_tokens,
+            lora_request=request.lora_request,
+            prompt_embeds=request.prompt_embeds,
+            deepstack_input_embeds=request.deepstack_input_embeds,
+        )
+
+    def __repr__(self) -> str:
+        prompt_embeds_shape = (self.prompt_embeds.shape
+                               if self.prompt_embeds else None)
+        return (f"NewRequestData("
+                f"req_id={self.req_id},"
+                f"prompt_token_ids={self.prompt_token_ids},"
+                f"mm_features={self.mm_features},"
+                f"sampling_params={self.sampling_params},"
+                f"block_ids={self.block_ids},"
+                f"num_computed_tokens={self.num_computed_tokens},"
+                f"lora_request={self.lora_request},"
+                f"prompt_embeds_shape={prompt_embeds_shape}"
+                ")")
+
+    # Version of __repr__ with the prompt data obfuscated
+    def anon_repr(self) -> str:
+        prompt_token_ids_len = len(
+            self.prompt_token_ids
+        ) if self.prompt_token_ids is not None else None
+        prompt_embeds_shape = (self.prompt_embeds.shape
+                               if self.prompt_embeds else None)
+        return (f"NewRequestData("
+                f"req_id={self.req_id},"
+                f"prompt_token_ids_len={prompt_token_ids_len},"
+                f"mm_features={self.mm_features},"
+                f"sampling_params={self.sampling_params},"
+                f"block_ids={self.block_ids},"
+                f"num_computed_tokens={self.num_computed_tokens},"
+                f"lora_request={self.lora_request},"
+                f"prompt_embeds_shape={prompt_embeds_shape}"
+                ")")

vllm_vacc/vllm/v1/core/sched/scheduler.py

@@ -0,0 +1,930 @@
+from __future__ import annotations
+
+import itertools
+import time
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
+from vllm.distributed.kv_transfer.kv_connector.factory import (
+    KVConnectorFactory)
+from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
+                                                          KVConnectorRole)
+from vllm.logger import init_logger
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
+from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
+                                                compute_encoder_budget)
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+from vllm.v1.core.sched.interface import SchedulerInterface
+from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
+                                       SchedulerOutput)
+from vllm.v1.core.sched.request_queue import (SchedulingPolicy,
+                                              create_request_queue)
+from vllm.v1.core.sched.utils import check_stop
+from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,
+                            EngineCoreOutputs)
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.metrics.stats import SchedulerStats
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
+from vllm.v1.structured_output import StructuredOutputManager
+
+from vllm_vacc.vllm.model_executor.models.vars import BLOCK_GROUP_SIZE as env_blk_grp_size
+from vllm_vacc.vllm.model_executor.models.vars import LLM_MAX_PREFILL_SEQ_LEN
+
+logger = init_logger(__name__)
+
+prefill_first = True
+
+def _is_in_prefill(req: Request) -> bool:
+    # 还没把 prompt 全算完
+    return req.num_computed_tokens < req.num_prompt_tokens
+
+def align_up(value: int, alignment: int) -> int:
+    if alignment <= 0:
+        return value
+    return ((value + alignment - 1) // alignment) * alignment
+
+class Scheduler(SchedulerInterface):
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_config: KVCacheConfig,
+        structured_output_manager: StructuredOutputManager,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+        include_finished_set: bool = False,
+        log_stats: bool = False,
+    ) -> None:
+        self.vllm_config = vllm_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.cache_config = vllm_config.cache_config
+        self.lora_config = vllm_config.lora_config
+        self.kv_cache_config = kv_cache_config
+        self.kv_events_config = vllm_config.kv_events_config
+        self.parallel_config = vllm_config.parallel_config
+        self.log_stats = log_stats
+        self.structured_output_manager = structured_output_manager
+        self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder
+
+        # include_finished_set controls whether a separate set of finished
+        # request ids should be included in the EngineCoreOutputs returned
+        # by update_from_outputs(). This is currently used in the multi-engine
+        # case to track request lifetimes efficiently.
+        self.finished_req_ids_dict: Optional[dict[int, set[str]]] = (
+            defaultdict(set) if include_finished_set else None)
+
+        # Scheduling constraints.
+        self.max_num_running_reqs = self.scheduler_config.max_num_seqs
+        self.max_num_scheduled_tokens = \
+            self.scheduler_config.max_num_batched_tokens
+        self.max_model_len = self.scheduler_config.max_model_len
+        self.enable_kv_cache_events = (
+            self.kv_events_config is not None
+            and self.kv_events_config.enable_kv_cache_events)
+
+        # Create KVConnector for the Scheduler. Note that each Worker
+        # will have a corresponding KVConnector with Role=WORKER.
+        # KV Connector pushes/pull of remote KVs for P/D and offloading.
+        self.connector = None
+        if self.vllm_config.kv_transfer_config is not None:
+            assert len(self.kv_cache_config.kv_cache_groups) == 1, (
+                "Multiple KV cache groups are not currently supported "
+                "with KV connectors")
+            assert not self.is_encoder_decoder, (
+                "Encoder-decoder models are not currently supported "
+                "with KV connectors")
+            self.connector = KVConnectorFactory.create_connector(
+                config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
+
+        self.kv_event_publisher = EventPublisherFactory.create(
+            self.kv_events_config,
+            self.parallel_config.data_parallel_rank,
+        )
+
+        num_gpu_blocks = self.cache_config.num_gpu_blocks
+        assert num_gpu_blocks is not None and num_gpu_blocks > 0
+
+        self.block_size = self.cache_config.block_size
+        self.max_total_kv_tokens: Optional[int] = None
+        if self.block_size is not None:
+            self.max_total_kv_tokens = (
+                self.kv_cache_config.num_blocks * self.block_size)
+        self.dcp_world_size = \
+            vllm_config.parallel_config.decode_context_parallel_size
+        
+        # Note(hc): The scheduler’s block_size must be multiplied
+        # by dcp_world_size, since block hashes are computed on the
+        # original full token sequence at a granularity of
+        # original_block_size × dcp_world_size.
+        if self.dcp_world_size > 1:
+            self.block_size *= self.dcp_world_size
+
+        # req_id -> Request
+        self.requests: dict[str, Request] = {}
+        # Scheduling policy
+        if self.scheduler_config.policy == "priority":
+            self.policy = SchedulingPolicy.PRIORITY
+        elif self.scheduler_config.policy == "fcfs":
+            self.policy = SchedulingPolicy.FCFS
+        else:
+            raise ValueError(
+                f"Unknown scheduling policy: {self.scheduler_config.policy}")
+        # Priority queues for requests.
+        self.waiting = create_request_queue(self.policy)
+        self.running: list[Request] = []
+        # Pending barrier groups: gid -> list[Request]
+        self._barrier_groups: dict[str, list[Request]] = {}
+        # Expected group sizes: gid -> expected size
+        self._barrier_expected: dict[str, int] = {}
+        # Force decode-only scheduling in the next step when prefills block.
+        self.force_decode_next_step = False
+
+        # The request IDs that are finished in between the previous and the
+        # current steps. This is used to notify the workers about the finished
+        # requests so that they can free the cached states for those requests.
+        # This is flushed at the end of each scheduling step.
+        self.finished_req_ids: set[str] = set()
+
+        # KV Connector: requests in process of async KV loading or recving
+        self.finished_recving_kv_req_ids: set[str] = set()
+
+        # Encoder-related.
+        # Calculate encoder cache size if applicable
+        # NOTE: For now we use the same budget for both compute and space.
+        # This can be changed when we make encoder cache for embedding caching
+        # across requests.
+        encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
+            model_config=vllm_config.model_config,
+            scheduler_config=vllm_config.scheduler_config,
+            mm_registry=mm_registry,
+        )
+
+        # NOTE(woosuk): Here, "encoder" includes the vision encoder (and
+        # projector if needed) for MM models as well as encoder-decoder
+        # transformers.
+        self.max_num_encoder_input_tokens = encoder_compute_budget
+        # NOTE: For the models without encoder (e.g., text-only models),
+        # the encoder cache will not be initialized because cache size is 0
+        # for these models.
+        self.encoder_cache_manager = EncoderCacheManager(
+            cache_size=encoder_cache_size)
+
+        speculative_config = vllm_config.speculative_config
+        self.use_eagle = False
+        self.num_spec_tokens = self.num_lookahead_tokens = 0
+        if speculative_config:
+            self.num_spec_tokens = speculative_config.num_speculative_tokens
+            if speculative_config.use_eagle():
+                self.use_eagle = True
+                self.num_lookahead_tokens = self.num_spec_tokens
+
+        # Create the KV cache manager.
+        self.kv_cache_manager = KVCacheManager(
+            kv_cache_config=kv_cache_config,
+            max_model_len=self.max_model_len,
+            enable_caching=self.cache_config.enable_prefix_caching,
+            use_eagle=self.use_eagle,
+            log_stats=self.log_stats,
+            enable_kv_cache_events=self.enable_kv_cache_events,
+            dcp_world_size=self.dcp_world_size,
+        )
+        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
+
+    def add_request(self, request: Request) -> None:
+        # self.waiting.add_request(request)
+        # Always track in global dict first
+        self.requests[request.request_id] = request
+        gid = getattr(request, "barrier_group_id", None)
+        gsz = getattr(request, "barrier_group_size", 0) or 0
+
+        if gid and gsz > 1:
+            # Stage into pending barrier group
+            group_list = self._barrier_groups.get(gid)
+            if group_list is None:
+                self._barrier_groups[gid] = group_list = []
+                self._barrier_expected[gid] = gsz
+            else:
+                # Reconcile inconsistent sizes by taking the max
+                self._barrier_expected[gid] = max(self._barrier_expected[gid],
+                                                  gsz)
+            group_list.append(request)
+            if self.log_stats:
+                request.record_event(EngineCoreEventType.QUEUED)
+
+            # Flush group when complete
+            if len(group_list) >= self._barrier_expected[gid]:
+                # Preserve arrival order
+                for r in sorted(group_list, key=lambda x: x.arrival_time):
+                    self.waiting.add_request(r)
+                # Cleanup
+                del self._barrier_groups[gid]
+                del self._barrier_expected[gid]
+            return
+
+        # No barrier group; enqueue directly
+        self.waiting.add_request(request)
+        if self.log_stats:
+            request.record_event(EngineCoreEventType.QUEUED)
+            
+    def _schedule_running_requests_for_mode(
+        self,
+        prefill_mode: bool,
+        token_budget: int,
+        encoder_compute_budget: int,
+        scheduled_timestamp: float,
+        scheduled_running_reqs: list[Request],
+        preempted_reqs: list[Request],
+        req_to_new_blocks: dict[str, KVCacheBlocks],
+        num_scheduled_tokens: dict[str, int],
+        scheduled_spec_decode_tokens: dict[str, list[int]],
+        scheduled_encoder_inputs: dict[str, list[int]],
+        scan_index: int,
+        next_req_index: int,
+        skip_request_ids: set[str],
+    ) -> tuple[int, int, int, int]:
+        """Schedule running requests under the given mode.
+
+        Args:
+            scan_index: The position in ``self.running`` to start scanning from.
+            next_req_index: The next logical "running index" used when we
+                populate ``structured_output_request_ids``.
+        """
+        while scan_index < len(self.running) and token_budget > 0:
+            request = self.running[scan_index]
+
+            if request.request_id in skip_request_ids:
+                scan_index += 1
+                continue
+
+            if prefill_mode and not _is_in_prefill(request):
+                scan_index += 1
+                continue
+
+            if (not prefill_mode) and _is_in_prefill(request):
+                scan_index += 1
+                continue
+
+            num_new_tokens = max(
+                request.num_tokens_with_spec - request.num_computed_tokens, 0)
+
+            if (0 < self.scheduler_config.long_prefill_token_threshold <
+                    num_new_tokens):
+                num_new_tokens = (
+                    self.scheduler_config.long_prefill_token_threshold)
+
+            # chunked prefill has to be enabled explicitly to allow pooling
+            # requests to be chunked
+            if (not self.scheduler_config.chunked_prefill_enabled
+                    and num_new_tokens > token_budget):
+                break
+
+            num_new_tokens = min(num_new_tokens, token_budget)
+
+            # Make sure the input position does not exceed the max model len.
+            # This is necessary when using spec decoding.
+            num_new_tokens = min(
+                num_new_tokens,
+                self.max_model_len - 1 - request.num_computed_tokens)
+
+            # Schedule encoder inputs.
+            encoder_inputs_to_schedule = None
+            # new_encoder_budget = encoder_budget
+            # if request.has_encoder_inputs:
+            #     (encoder_inputs_to_schedule, num_new_tokens,
+            #      new_encoder_budget) = self._try_schedule_encoder_inputs(
+            #          request, request.num_computed_tokens, num_new_tokens,
+            #          encoder_budget)
+            new_encoder_compute_budget = encoder_compute_budget
+            if request.has_encoder_inputs:
+                (encoder_inputs_to_schedule, num_new_tokens,
+                 new_encoder_compute_budget
+                 ) = self._try_schedule_encoder_inputs(
+                     request, request.num_computed_tokens, num_new_tokens,
+                     encoder_compute_budget)
+
+            if num_new_tokens == 0:
+                scan_index += 1
+                continue
+
+            # num_draft_tokens = max(
+            #     num_new_tokens + request.num_computed_tokens -
+            #     request.num_tokens, 0)
+
+            while True:
+                new_blocks = self.kv_cache_manager.allocate_slots(
+                    request,
+                    num_new_tokens,
+                    # num_draft_tokens=num_draft_tokens,
+                    num_lookahead_tokens=self.num_lookahead_tokens)
+                if new_blocks is None:
+                    if self.policy == SchedulingPolicy.PRIORITY:
+                        preempted_req = max(
+                            self.running,
+                            key=lambda r: (r.priority, r.arrival_time),
+                        )
+                        self.running.remove(preempted_req)
+                        if preempted_req in scheduled_running_reqs:
+                            scheduled_running_reqs.remove(preempted_req)
+     
+                    else:
+                        preempted_req = self.running.pop()
+
+                    self.kv_cache_manager.free(preempted_req)
+                    self.encoder_cache_manager.free(preempted_req)
+                    preempted_req.status = RequestStatus.PREEMPTED
+                    if _is_in_prefill(preempted_req):
+                        preempted_req.num_computed_tokens = 0
+                    if self.log_stats:
+                        preempted_req.record_event(
+                            EngineCoreEventType.PREEMPTED, scheduled_timestamp)
+
+                    self.waiting.prepend_request(preempted_req)
+                    preempted_reqs.append(preempted_req)
+                    if preempted_req == request:
+                        can_schedule = False
+                        break
+                else:
+                    can_schedule = True
+                    break
+            if not can_schedule:
+                break
+            assert new_blocks is not None
+
+            scheduled_running_reqs.append(request)
+            req_to_new_blocks[request.request_id] = new_blocks
+            num_scheduled_tokens[request.request_id] = num_new_tokens
+            token_budget -= num_new_tokens
+            next_req_index += 1
+
+            if request.spec_token_ids:
+                num_scheduled_spec_tokens = (num_new_tokens +
+                                             request.num_computed_tokens -
+                                             request.num_tokens)
+                if num_scheduled_spec_tokens > 0:
+                    del request.spec_token_ids[num_scheduled_spec_tokens:]
+                    scheduled_spec_decode_tokens[request.request_id] = (
+                        request.spec_token_ids)
+
+            if encoder_inputs_to_schedule:
+                scheduled_encoder_inputs[request.request_id] = (
+                    encoder_inputs_to_schedule)
+                for i in encoder_inputs_to_schedule:
+                    self.encoder_cache_manager.allocate(request, i)
+                # encoder_budget = new_encoder_budget
+                encoder_compute_budget = new_encoder_compute_budget
+
+            scan_index += 1
+
+        return token_budget, encoder_compute_budget, scan_index, next_req_index
+
+    def finish_requests(
+        self,
+        request_ids: Union[str, Iterable[str]],
+        finished_status: RequestStatus,
+    ) -> None:
+        """Handles the finish signal from outside the scheduler.
+
+        For example, the API server can abort a request when the client
+        disconnects.
+        """
+        assert RequestStatus.is_finished(finished_status)
+        if isinstance(request_ids, str):
+            request_ids = (request_ids, )
+        else:
+            request_ids = set(request_ids)
+
+        running_requests_to_remove = []
+        waiting_requests_to_remove = []
+        valid_requests = []
+
+        # First pass: collect requests to remove from queues
+        for req_id in request_ids:
+            request = self.requests.get(req_id)
+            if request is None:
+                # Invalid request ID.
+                continue
+
+            valid_requests.append(request)
+            if request.status == RequestStatus.RUNNING:
+                running_requests_to_remove.append(request)
+            else:
+                waiting_requests_to_remove.append(request)
+
+        # Remove all requests from queues at once for better efficiency
+        for request in running_requests_to_remove:
+            self.running.remove(request)
+        if waiting_requests_to_remove:
+            self.waiting.remove_requests(waiting_requests_to_remove)
+
+        # Handle requests that are pending in barrier groups:
+        # if any item from a pending group is finished/aborted, flush the rest.
+        for request in valid_requests:
+            gid = getattr(request, "barrier_group_id", None)
+            if not gid:
+                continue
+            group_list = self._barrier_groups.get(gid)
+            if not group_list:
+                continue
+            # Remove this request from pending group if present
+            try:
+                if request in group_list:
+                    group_list.remove(request)
+            except ValueError:
+                pass
+            # Flush remaining members (if any) to waiting, then cleanup.
+            if group_list:
+                for r in sorted(group_list, key=lambda x: x.arrival_time):
+                    self.waiting.add_request(r)
+            self._barrier_groups.pop(gid, None)
+            self._barrier_expected.pop(gid, None)
+
+
+
+        # Second pass: set status and free requests
+        for request in valid_requests:
+            request.status = finished_status
+            self._free_request(request)
+
+    def _estimate_future_kv_tokens(self, request: Request) -> int:
+        """Estimate the KV tokens a request may need for its full lifetime."""
+        alignment = min(env_blk_grp_size, self.max_model_len)
+        alignment = max(alignment, 1)
+
+        prompt_tokens = request.num_prompt_tokens  # prefill length
+        prompt_reserved = align_up(prompt_tokens, alignment)
+
+        remaining_room = max(self.max_model_len - prompt_tokens, 0)
+        max_decode_tokens = request.max_tokens # output length
+        if max_decode_tokens < 0:
+            max_decode_tokens = remaining_room
+        else:
+            max_decode_tokens = min(max_decode_tokens, remaining_room)
+
+        decode_budget = max_decode_tokens
+        if self.num_lookahead_tokens > 0:
+            decode_budget = min(decode_budget + self.num_lookahead_tokens,
+                                remaining_room)
+
+        remaining_capacity_in_prompt = max(prompt_reserved - prompt_tokens, 0)
+        decode_after_prompt = max(decode_budget - remaining_capacity_in_prompt,
+                                  0)
+        if decode_after_prompt == 0:
+            total_reserved = prompt_reserved
+        else:
+            total_reserved = prompt_reserved + \
+                align_up(decode_after_prompt, alignment)
+
+        return min(total_reserved, self.max_model_len)
+
+    def _compute_total_future_kv_tokens(
+        self, requests: Iterable[Request]
+    ) -> int:
+        return sum(self._estimate_future_kv_tokens(req) for req in requests)
+
+    def make_stats(
+        self,
+        spec_decoding_stats: Optional[SpecDecodingStats] = None,
+        kv_connector_stats: Optional["KVConnectorStats"] = None,
+    ) -> Optional[SchedulerStats]:
+        if not self.log_stats:
+            return None
+        prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats()
+        assert prefix_cache_stats is not None
+        return SchedulerStats(num_running_reqs=len(self.running),
+                              num_waiting_reqs=len(self.waiting),
+                              running_seqlens = [len(running_i._all_token_ids) for running_i in self.running],
+                              kv_cache_usage=self.kv_cache_manager.usage,
+                              prefix_cache_stats=prefix_cache_stats,
+                              spec_decoding_stats=spec_decoding_stats,
+                              num_corrupted_reqs=sum(req.is_output_corrupted
+                                                     for req in self.running),
+                              kv_connector_stats=kv_connector_stats.data
+                              if kv_connector_stats else None)
+
+    def schedule(self) -> SchedulerOutput:
+        # NOTE(woosuk) on the scheduling algorithm:
+        # There's no "decoding phase" nor "prefill phase" in the scheduler.
+        # Each request just has the num_computed_tokens and
+        # num_tokens_with_spec. num_tokens_with_spec =
+        # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids).
+        # At each step, the scheduler tries to assign tokens to the requests
+        # so that each request's num_computed_tokens can catch up its
+        # num_tokens_with_spec. This is general enough to cover
+        # chunked prefills, prefix caching, speculative decoding,
+        # and the "jump decoding" optimization in the future.
+
+        scheduled_new_reqs: list[Request] = []
+        scheduled_resumed_reqs: list[Request] = []
+        scheduled_running_reqs: list[Request] = []
+        preempted_reqs: list[Request] = []
+
+        # NOTE: structured_output_request_ids maps
+        # a request's (request that uses structured output)
+        # request_id to the running request index.
+        # This will helps us determine to slice the grammar bitmask
+        # and only applies valid mask for requests that
+        # uses structured decoding.
+
+        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
+        num_scheduled_tokens: dict[str, int] = {}
+        token_budget = self.max_num_scheduled_tokens
+        # Encoder-related.
+        scheduled_encoder_inputs: dict[str, list[int]] = {}
+        encoder_compute_budget = self.max_num_encoder_input_tokens
+        # Spec decode-related.
+        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+
+        # For logging.
+        scheduled_timestamp = time.monotonic()
+        
+        # Determine whether we should prioritize prefills in this step.
+        has_running_prefill = any(_is_in_prefill(r) for r in self.running)
+        has_waiting_prefill = False
+        if self.waiting and prefill_first:
+            for waiting_req in self.waiting:
+                if _is_in_prefill(waiting_req):
+                    has_waiting_prefill = True
+                    break
+        prefill_mode = prefill_first and (has_running_prefill or has_waiting_prefill or not self.running) and len(self.running) < self.max_num_running_reqs
+        if self.force_decode_next_step:
+            prefill_mode = False
+             
+        # First, schedule the RUNNING requests.
+        req_index = 0
+        token_budget, encoder_budget, _, req_index = (
+            self._schedule_running_requests_for_mode(
+                prefill_mode,
+                token_budget,
+                encoder_compute_budget,
+                scheduled_timestamp,
+                scheduled_running_reqs,
+                preempted_reqs,
+                req_to_new_blocks,
+                num_scheduled_tokens,
+                scheduled_spec_decode_tokens,
+                scheduled_encoder_inputs,
+                0,
+                req_index,
+                set(),
+            ))
+        # Record the LoRAs in scheduled_running_reqs
+        scheduled_loras: set[int] = set()
+        if self.lora_config:
+            scheduled_loras = set(
+                req.lora_request.lora_int_id for req in scheduled_running_reqs
+                if req.lora_request and req.lora_request.lora_int_id > 0)
+            assert len(scheduled_loras) <= self.lora_config.max_loras
+
+        # Use a temporary RequestQueue to collect requests that need to be
+        # skipped and put back at the head of the waiting queue later
+        skipped_waiting_requests = create_request_queue(self.policy)
+        # prefill_allocation_failed = False
+
+        # Next, schedule the WAITING requests.
+        if not preempted_reqs:
+            future_kv_token_usage = 0
+            prefill_all_len = 0
+            if self.max_total_kv_tokens is not None:
+                future_kv_token_usage = (
+                    self._compute_total_future_kv_tokens(self.running))
+
+            while self.waiting and token_budget > 0:
+                if len(self.running) == self.max_num_running_reqs:
+                    break
+
+                request = self.waiting.peek_request()
+                request_future_tokens: Optional[int] = None
+                if self.max_total_kv_tokens is not None:
+                    request_future_tokens = self._estimate_future_kv_tokens(
+                        request)
+                    if prefill_mode:
+                        prefill_all_len += request.num_prompt_tokens
+                    if prefill_mode and ((future_kv_token_usage +
+                                         request_future_tokens
+                                         > self.max_total_kv_tokens) or prefill_all_len > LLM_MAX_PREFILL_SEQ_LEN):
+                        self.waiting.pop_request()
+                        skipped_waiting_requests.prepend_request(request)
+                        # prefill_allocation_failed = True
+                        self.force_decode_next_step = True
+                        continue
+                in_prefill = _is_in_prefill(request)
+                if prefill_mode and not in_prefill:
+                    self.waiting.pop_request()
+                    skipped_waiting_requests.prepend_request(request)
+                    continue
+                if (not prefill_mode) and in_prefill:
+                    self.waiting.pop_request()
+                    skipped_waiting_requests.prepend_request(request)
+                    continue
+
+                # KVTransfer: skip request if still waiting for remote kvs.
+                if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
+                    is_ready = self._update_waiting_for_remote_kv(request)
+                    if is_ready:
+                        request.status = RequestStatus.WAITING
+                    else:
+                        logger.debug(
+                            "%s is still in WAITING_FOR_REMOTE_KVS state.",
+                            request.request_id)
+                        self.waiting.pop_request()
+                        skipped_waiting_requests.prepend_request(request)
+                        continue
+
+                # Skip request if the structured output request is still waiting
+                # for FSM compilation.
+                if request.status == RequestStatus.WAITING_FOR_FSM:
+                    structured_output_req = request.structured_output_request
+                    if structured_output_req and structured_output_req.grammar:
+                        request.status = RequestStatus.WAITING
+                    else:
+                        self.waiting.pop_request()
+                        skipped_waiting_requests.prepend_request(request)
+                        continue
+
+                # Check that adding the request still respects the max_loras
+                # constraint.
+                if (self.lora_config and request.lora_request and
+                    (len(scheduled_loras) == self.lora_config.max_loras and
+                     request.lora_request.lora_int_id not in scheduled_loras)):
+                    # Scheduling would exceed max_loras, skip.
+                    self.waiting.pop_request()
+                    skipped_waiting_requests.prepend_request(request)
+                    continue
+
+                num_external_computed_tokens = 0
+                load_kv_async = False
+
+                # Get already-cached tokens.
+                if request.num_computed_tokens == 0:
+                    # Get locally-cached tokens.
+                    new_computed_blocks, num_new_local_computed_tokens = \
+                        self.kv_cache_manager.get_computed_blocks(
+                            request)
+
+                    # Get externally-cached tokens if using a KVConnector.
+                    if self.connector is not None:
+                        num_external_computed_tokens, load_kv_async = (
+                            self.connector.get_num_new_matched_tokens(
+                                request, num_new_local_computed_tokens))
+                        if num_external_computed_tokens is None:
+                            # The request cannot be scheduled because
+                            # the KVConnector couldn't determine
+                            # the number of matched tokens.
+                            self.waiting.pop_request()
+                            skipped_waiting_requests.prepend_request(request)
+                            continue
+
+                    # Total computed tokens (local + external).
+                    num_computed_tokens = (num_new_local_computed_tokens +
+                                           num_external_computed_tokens)
+                # KVTransfer: WAITING reqs have num_computed_tokens > 0
+                # after async KV recvs are completed.
+                else:
+                    new_computed_blocks = (
+                        self.kv_cache_manager.create_empty_block_list())
+                    num_new_local_computed_tokens = 0
+                    num_computed_tokens = request.num_computed_tokens
+
+                encoder_inputs_to_schedule = None
+                # new_encoder_budget = encoder_budget
+                new_encoder_compute_budget = encoder_compute_budget
+
+                # KVTransfer: loading remote KV, do not allocate for new work.
+                if load_kv_async:
+                    assert num_external_computed_tokens > 0
+                    num_new_tokens = 0
+                # Number of tokens to be scheduled.
+                else:
+                    # We use `request.num_tokens` instead of
+                    # `request.num_prompt_tokens` to consider the resumed
+                    # requests, which have output tokens.
+                    num_new_tokens = request.num_tokens - num_computed_tokens
+                    if (0 < self.scheduler_config.long_prefill_token_threshold
+                            < num_new_tokens):
+                        num_new_tokens = (
+                            self.scheduler_config.long_prefill_token_threshold)
+
+                    # chunked prefill has to be enabled explicitly to allow
+                    # pooling requests to be chunked
+                    if not self.scheduler_config.chunked_prefill_enabled and \
+                        num_new_tokens > token_budget:
+                        self.waiting.pop_request()
+                        skipped_waiting_requests.prepend_request(request)
+                        continue
+
+                    num_new_tokens = min(num_new_tokens, token_budget)
+                    assert num_new_tokens > 0
+
+                    # Schedule encoder inputs.
+                    if request.has_encoder_inputs:
+                        (encoder_inputs_to_schedule, num_new_tokens,
+                         new_encoder_compute_budget
+                         ) = self._try_schedule_encoder_inputs(
+                             request, num_computed_tokens, num_new_tokens,
+                             encoder_compute_budget)
+                        if num_new_tokens == 0:
+                            # The request cannot be scheduled.
+                            break
+
+                # Handles an edge case when P/D Disaggregation
+                # is used with Spec Decoding where an
+                # extra block gets allocated which
+                # creates a mismatch between the number
+                # of local and remote blocks.
+                effective_lookahead_tokens = (0 if request.num_computed_tokens
+                                              == 0 else
+                                              self.num_lookahead_tokens)
+
+                # Determine if we need to allocate cross-attention blocks.
+                if self.is_encoder_decoder and request.has_encoder_inputs:
+                    # TODO(russellb): For Whisper, we know that the input is
+                    # always padded to the maximum length. If we support other
+                    # encoder-decoder models, this will need to be updated if we
+                    # want to only allocate what is needed.
+                    num_encoder_tokens =\
+                        self.scheduler_config.max_num_encoder_input_tokens
+                else:
+                    num_encoder_tokens = 0
+
+                new_blocks = self.kv_cache_manager.allocate_slots(
+                    request,
+                    num_new_tokens + num_external_computed_tokens,
+                    num_new_local_computed_tokens,
+                    new_computed_blocks,
+                    num_lookahead_tokens=effective_lookahead_tokens,
+                    delay_cache_blocks=load_kv_async,
+                    num_encoder_tokens=num_encoder_tokens,
+                )
+                if new_blocks is None:
+                    # The request cannot be scheduled in this step.
+                    if prefill_mode:
+                        # prefill_allocation_failed = True
+                        self.force_decode_next_step = True
+                    break
+
+                # KVTransfer: the connector uses this info to determine
+                # if a load is needed. Note that
+                # This information is used to determine if a load is
+                # needed for this request.
+                if self.connector is not None:
+                    self.connector.update_state_after_alloc(
+                        request,
+                        new_computed_blocks + new_blocks,
+                        num_external_computed_tokens,
+                    )
+
+                # Request was already popped from self.waiting
+                # unless it was re-added above due to new_blocks being None.
+                request = self.waiting.pop_request()
+                if load_kv_async:
+                    # If loading async, allocate memory and put request
+                    # into the WAITING_FOR_REMOTE_KV state.
+                    skipped_waiting_requests.prepend_request(request)
+                    request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
+                    continue
+
+                req_index += 1
+                self.running.append(request)
+                if self.max_total_kv_tokens is not None:
+                    if request_future_tokens is None:
+                        request_future_tokens = (
+                            self._estimate_future_kv_tokens(request))
+                    future_kv_token_usage += request_future_tokens
+                if self.log_stats:
+                    request.record_event(EngineCoreEventType.SCHEDULED,
+                                         scheduled_timestamp)
+                if request.status == RequestStatus.WAITING:
+                    scheduled_new_reqs.append(request)
+                elif request.status == RequestStatus.PREEMPTED:
+                    scheduled_resumed_reqs.append(request)
+                else:
+                    raise RuntimeError(
+                        f"Invalid request status: {request.status}")
+
+                if self.lora_config and request.lora_request:
+                    scheduled_loras.add(request.lora_request.lora_int_id)
+                # req_to_new_block_ids[request.request_id] = (
+                #     self.kv_cache_manager.get_block_ids(request.request_id))
+                req_to_new_blocks[request.request_id] = new_blocks
+
+                num_scheduled_tokens[request.request_id] = num_new_tokens
+                token_budget -= num_new_tokens
+                request.status = RequestStatus.RUNNING
+                request.num_computed_tokens = num_computed_tokens
+                # Count the number of prefix cached tokens.
+                if request.num_cached_tokens < 0:
+                    request.num_cached_tokens = num_computed_tokens
+                # Encoder-related.
+                if encoder_inputs_to_schedule:
+                    scheduled_encoder_inputs[request.request_id] = (
+                        encoder_inputs_to_schedule)
+                    # Allocate the encoder cache.
+                    for i in encoder_inputs_to_schedule:
+                        self.encoder_cache_manager.allocate(request, i)
+                    # encoder_budget = new_encoder_budget
+                    encoder_compute_budget = new_encoder_compute_budget
+
+        # Put back any skipped requests at the head of the waiting queue
+        if skipped_waiting_requests:
+            self.waiting.prepend_requests(skipped_waiting_requests)
+
+        if not prefill_mode:
+            self.force_decode_next_step = False
+
+        # Check if the scheduling constraints are satisfied.
+        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+        assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
+        assert token_budget >= 0
+        assert len(self.running) <= self.max_num_running_reqs
+        # Since some requests in the RUNNING queue may not be scheduled in
+        # this step, the total number of scheduled requests can be smaller than
+        # len(self.running).
+        assert (len(scheduled_new_reqs) + len(scheduled_resumed_reqs) +
+                len(scheduled_running_reqs) <= len(self.running))
+
+        # Get the longest common prefix among all requests in the running queue.
+        # This can be potentially used for cascade attention.
+        num_common_prefix_blocks = [0] * len(
+            self.kv_cache_config.kv_cache_groups)
+        if self.running:
+            any_request = self.running[0]
+            num_common_prefix_blocks = (
+                self.kv_cache_manager.get_num_common_prefix_blocks(
+                    any_request, len(self.running)))
+
+        # Construct the scheduler output.
+        # new_reqs_data = [
+        #     NewRequestData.from_request(req,
+        #                                 req_to_new_block_ids[req.request_id])
+        #     for req in scheduled_new_reqs
+        # ]
+        new_reqs_data = [
+            NewRequestData.from_request(
+                req, req_to_new_blocks[req.request_id].get_block_ids())
+            for req in scheduled_new_reqs
+        ]
+        cached_reqs_data = self._make_cached_request_data(
+            scheduled_running_reqs,
+            scheduled_resumed_reqs,
+            num_scheduled_tokens,
+            scheduled_spec_decode_tokens,
+            # req_to_new_block_ids,
+            req_to_new_blocks,
+        )
+        scheduled_requests = (scheduled_new_reqs + scheduled_running_reqs +
+                              scheduled_resumed_reqs)
+        structured_output_request_ids, grammar_bitmask = (
+            self.get_grammar_bitmask(scheduled_requests,
+                                     scheduled_spec_decode_tokens))
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=new_reqs_data,
+            scheduled_cached_reqs=cached_reqs_data,
+            num_scheduled_tokens=num_scheduled_tokens,
+            total_num_scheduled_tokens=total_num_scheduled_tokens,
+            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
+            scheduled_encoder_inputs=scheduled_encoder_inputs,
+            num_common_prefix_blocks=num_common_prefix_blocks,
+            # finished_req_ids is an existing state in the scheduler,
+            # instead of being newly scheduled in this step.
+            # It contains the request IDs that are finished in between
+            # the previous and the current steps.
+            finished_req_ids=self.finished_req_ids,
+            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+            structured_output_request_ids=structured_output_request_ids,
+            grammar_bitmask=grammar_bitmask,
+        )
+
+        # NOTE(Kuntai): this function is designed for multiple purposes:
+        # 1. Plan the KV cache store
+        # 2. Wrap up all the KV cache load / save ops into an opaque object
+        # 3. Clear the internal states of the connector
+        if self.connector is not None:
+            meta = self.connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_connector_metadata = meta
+
+        events = self.kv_cache_manager.take_events()
+        
+        # collect KV cache events from connector
+        if self.connector is not None:
+            print('if self.connector is not None:')
+            connector_events = self.connector.take_events()
+            if connector_events:
+                if events is None:
+                    events = list(connector_events)
+                else:
+                    events.extend(connector_events)
+
+        # publish collected KV cache events
+        
+        if events:
+            batch = KVEventBatch(ts=time.time(), events=events)
+            self.kv_event_publisher.publish(batch)
+
+        self._update_after_schedule(scheduler_output)
+        return scheduler_output
+
+    

vllm_vacc/vllm/v1/core/single_type_kv_cache_manager.py

@@ -0,0 +1,61 @@
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Callable
+
+from vllm.utils import cdiv
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
+from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,
+                                        MambaSpec, SlidingWindowSpec)
+from vllm.v1.request import Request
+
+import os
+from vllm_vacc.vllm.model_executor.models.vars import BLOCK_GROUP_SIZE as env_blk_grp_size
+
+
+class SingleTypeKVCacheManager(ABC):
+
+    def free(self, request_id: str) -> None:
+        """
+        Free the blocks for the request.
+
+        Args:
+            request_id: The request ID.
+        """
+        # Default to [] in case a request is freed (aborted) before alloc.
+        req_blocks = self.req_to_blocks.pop(request_id, [])
+
+        # Free blocks in reverse order so that the tail blocks are
+        # freed first.
+        # ordered_blocks = reversed(req_blocks)
+        self.block_pool.free_blocks(req_blocks)
+        self.num_cached_block.pop(request_id, None)
+    
+    
+    def allocate_new_blocks(self, request_id: str,
+                            num_tokens: int) -> list[KVCacheBlock]:
+        """
+        Allocate new blocks for the request to give it at least `num_tokens` 
+        token slots.
+
+        Args:
+            request_id: The request ID.
+            num_tokens: The total number of tokens that need a slot (including 
+                tokens that are already allocated).
+
+        Returns:
+            The new allocated blocks.
+        """
+        req_blocks = self.req_to_blocks[request_id]
+        num_required_blocks = cdiv(num_tokens, self.block_size)
+        
+        block_size_number_per_group = env_blk_grp_size // self.block_size #512
+        num_required_blocks = (num_required_blocks + block_size_number_per_group - 1) // block_size_number_per_group * block_size_number_per_group
+        num_new_blocks = num_required_blocks - len(req_blocks)
+        if num_new_blocks <= 0:
+            return []
+        else:
+            new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
+            req_blocks.extend(new_blocks)
+            return new_blocks