fix(build): 消除COPY ./vllm_overrides — vendor_overrides预置到qwen3_6_scripts/

竞赛平台docker build失败，无日志。最大嫌疑：
  COPY ./vllm_overrides /workspace/vllm_overrides
26e6cb4(成功)只有3个COPY，HEAD多了这第4个COPY。

修复：把9个vllm_overrides文件直接放进qwen3_6_scripts/vendor_overrides/
Dockerfile回到3个COPY（和26e6cb4结构一致），去掉Step 4 staging。
patch_ops.sh不需要改——它已经从./vendor_overrides/读取。

COPY数量: 4→3 (匹配26e6cb4)
Dockerfile行数: 74→48 (更简洁)
Step数: 8→7 (去掉staging step)

Dockerfile

@@ -3,64 +3,37 @@ FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.1
 RUN mkdir -p /workspace
 WORKDIR /workspace/
 
-# Copy all sources — ex_engine for build-time .so compilation,
-# qwen3_6_scripts for patches+prebuilt, vllm_overrides for core fixes
+# Copy all sources (vendor_overrides pre-staged inside qwen3_6_scripts/)
 COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
 COPY ./computility-run.yaml /workspace/computility-run.yaml
 COPY ./ex_engine /workspace/ex_engine
-COPY ./vllm_overrides /workspace/vllm_overrides
 
-# Step 1: Build EX Engine .so libraries (tolerant of compile failures)
+# Step 1: Build EX Engine .so libraries
 RUN chmod +x /workspace/ex_engine/build.sh && \
     bash /workspace/ex_engine/build.sh --corex 2>&1 | tee /workspace/ex_build.log ; \
     echo "[Dockerfile] ex_engine build exit code: $?"
 
-# Step 2: Precompile MoE CUDA kernels (tolerant)
+# Step 2: Precompile MoE CUDA kernels
 RUN python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 | tee -a /workspace/ex_build.log ; \
     echo "[Dockerfile] moe_topk precompile exit code: $?"
 
-# Step 3: Precompile vllm v0.5.5 MoE kernels (tolerant)
+# Step 3: Precompile vllm v0.5.5 MoE kernels
 RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee -a /workspace/ex_build.log ; \
     echo "[Dockerfile] moe_v055 precompile exit code: $?"
 
-# Step 4: Stage vendor_overrides into qwen3_6_scripts/ so patch_ops.sh finds them
-RUN mkdir -p /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block \
-             /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers && \
-    cp /workspace/vllm_overrides/core/evictor_v2.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/core/block/cpu_kv_content_cache.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/core/block/cpu_gpu_block_allocator.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/core/block/prefix_caching_block.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/core/block/block_table.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/core/block_manager_v2.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/sampling_params.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/model_executor/sampling_metadata.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py 2>/dev/null || true && \
-    cp /workspace/vllm_overrides/model_executor/layers/sampler.py \
-       /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py 2>/dev/null || true ; \
-    echo "[Dockerfile] vendor_overrides staged"
-
-# Step 5: Deploy patches (serving + engine fixes + prebuilt .so)
+# Step 4: Deploy patches (serving + engine fixes + prebuilt .so)
 RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
     cd /workspace/qwen3_6_scripts && \
     bash ./patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
     echo "[Dockerfile] patch_ops exit code: $?"
 
-# Step 6: Build ix_unified_bridge.so (ixformer::infer symbols resolved at runtime via RTLD_GLOBAL preload)
-# Tolerant: bridge is optional — corex_moe.py has ixformer.functions fallback
+# Step 5: Build ix_unified_bridge.so (ixformer symbols resolved at runtime)
 RUN chmod +x /workspace/ex_engine/build_unified_bridge.sh && \
     (bash /workspace/ex_engine/build_unified_bridge.sh 2>&1 || echo "[Dockerfile] bridge build FAILED (non-fatal)") | tee -a /workspace/ex_build.log ; \
     echo "[Dockerfile] ix_unified_bridge build exit code: $?"
 
-# Step 7: Deploy ix_unified and ex_engine Python modules to vllm path
+# Step 6: Deploy ex_engine Python modules to vllm path
 RUN VLLM_ROOT=$(python3 -c "import vllm; print(vllm.__path__[0])" 2>/dev/null | tail -1 || echo "/usr/local/corex/lib64/python3/dist-packages/vllm") && \
-    echo "VLLM_ROOT=${VLLM_ROOT}" && \
     cp /workspace/ex_engine/python/ix_unified.py "${VLLM_ROOT}/ix_unified.py" 2>/dev/null || true && \
     cp /workspace/ex_engine/python/corex_so_loader.py "${VLLM_ROOT}/corex_so_loader.py" 2>/dev/null || true && \
     cp /workspace/ex_engine/python/moe_fused_dispatch.py "${VLLM_ROOT}/moe_fused_dispatch.py" 2>/dev/null || true && \
@@ -69,7 +42,7 @@ RUN VLLM_ROOT=$(python3 -c "import vllm; print(vllm.__path__[0])" 2>/dev/null |
     fi ; \
     echo "[Dockerfile] ex_engine Python modules deployed"
 
-# Step 8: Precompile GDN kernel (needs vllm in path, so after patch_ops)
+# Step 7: Precompile GDN kernel (needs vllm in path, so after patch_ops)
 RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
     /workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
     echo "[Dockerfile] gdn precompile exit code: $?"

qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py

@@ -0,0 +1,456 @@
+import math
+from typing import List, Optional
+
+from vllm.core.block.common import BlockList
+from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator
+from vllm.utils import Device, cdiv, chunk_list
+
+
+class BlockTable:
+    """A class to manage blocks for a specific sequence.
+
+    The BlockTable maps a sequence of tokens to a list of blocks, where each
+    block represents a contiguous memory allocation for a portion of the 
+    sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is
+    responsible for allocating and freeing memory for the blocks.
+
+    Args:
+        block_size (int): The maximum number of tokens that can be stored in a
+            single block.
+        block_allocator (DeviceAwareBlockAllocator): The block allocator used to
+            manage memory for the blocks.
+        _blocks (Optional[List[Block]], optional): An optional list of existing
+            blocks to initialize the BlockTable with. If not provided, an empty
+            BlockTable is created.
+        max_block_sliding_window (Optional[int], optional): The number of
+            blocks to keep around for each sequance. If None, all blocks
+            are kept (eg., when sliding window is not used).
+            It should at least fit the sliding window size of the model.
+
+    Attributes:
+        _block_size (int): The maximum number of tokens that can be stored in a
+            single block.
+        _allocator (DeviceAwareBlockAllocator): The block allocator used to
+            manage memory for the blocks.
+        _blocks (Optional[List[Block]]): The list of blocks managed by this
+            BlockTable.
+        _num_full_slots (int): The number of tokens currently stored in the
+            blocks.
+    """
+
+    def __init__(
+        self,
+        block_size: int,
+        block_allocator: DeviceAwareBlockAllocator,
+        _blocks: Optional[List[Block]] = None,
+        max_block_sliding_window: Optional[int] = None,
+        cache_namespace: Optional[bytes] = None,
+    ):
+        self._block_size = block_size
+        self._allocator = block_allocator
+        self._cache_namespace = cache_namespace
+        if _blocks is None:
+            _blocks = []
+        self._blocks: BlockList = BlockList(_blocks)
+
+        self._max_block_sliding_window = max_block_sliding_window
+        self._num_full_slots = self._get_num_token_ids()
+
+    @staticmethod
+    def get_num_required_blocks(token_ids: List[int],
+                                block_size: int,
+                                num_lookahead_slots: int = 0) -> int:
+        """Calculates the minimum number of blocks required to store a given
+        sequence of token IDs along with any look-ahead slots that may be
+        required (like in multi-step + chunked-prefill).
+
+        This assumes worst-case scenario, where every block requires a new
+        allocation (e.g. ignoring prefix caching).
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be stored.
+            block_size (int): The maximum number of tokens that can be stored in
+                a single block.
+            num_lookahead_slots (int): look-ahead slots that the sequence may
+                require.
+
+        Returns:
+            int: The minimum number of blocks required to store the given
+                sequence of token IDs along with any required look-ahead slots.
+        """
+        return cdiv(len(token_ids) + num_lookahead_slots, block_size)
+
+    def allocate(self,
+                 token_ids: List[int],
+                 device: Device = Device.GPU) -> None:
+        """Allocates memory blocks for storing the given sequence of token IDs.
+
+        This method allocates the required number of blocks to store the given
+        sequence of token IDs.
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be stored.
+            device (Device, optional): The device on which the blocks should be
+                allocated. Defaults to Device.GPU.
+        """
+        assert not self._is_allocated
+        assert token_ids
+        blocks = self._allocate_blocks_for_token_ids(prev_block=None,
+                                                     token_ids=token_ids,
+                                                     device=device)
+        self.update(blocks)
+        self._num_full_slots = len(token_ids)
+
+    def update(self, blocks: List[Block]) -> None:
+        """Resets the table to the newly provided blocks
+        (with their corresponding block ids)
+        """
+        self._blocks.update(blocks)
+
+    def get_content_hashes(self) -> List[bytes]:
+        """Returns block-level content hashes for full blocks in order."""
+        content_hashes: List[bytes] = []
+        for block in self._blocks:
+            block_hash = block.content_hash
+            if block_hash is not None:
+                content_hashes.append(block_hash)
+        return content_hashes
+
+    def append_token_ids(self,
+                         token_ids: List[int],
+                         num_lookahead_slots: int = 0,
+                         num_computed_slots: Optional[int] = None) -> None:
+        """Appends a sequence of token IDs to the existing blocks in the
+        BlockTable.
+
+        This method appends the given sequence of token IDs to the existing
+        blocks in the BlockTable. If there is not enough space in the existing
+        blocks, new blocks are allocated using the `ensure_num_empty_slots`
+        method to accommodate the additional tokens.
+
+        The token IDs are divided into chunks of size `block_size` (except for
+        the first chunk, which may be smaller), and each chunk is appended to a
+        separate block.
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be appended.
+            num_computed_slots (Optional[int]): The number of KV cache slots
+                that are already filled (computed).
+                When sliding window is enabled, this is used to compute how many
+                blocks to drop at the front of the sequence.
+                Without sliding window, None can be passed.
+                Without chunked prefill, it should be the same as
+                _num_full_slots.
+        """
+        assert self._is_allocated, "no blocks have been allocated"
+        assert len(self._blocks) > 0
+
+        # Drop blocks that are no longer needed due to sliding window
+        if self._max_block_sliding_window is not None:
+            null_block = self._allocator.allocate_or_get_null_block()
+            assert num_computed_slots is not None
+            end_block_idx = (num_computed_slots //
+                             self._block_size) - self._max_block_sliding_window
+            for idx in range(0, end_block_idx):
+                b = self._blocks[idx]
+                if b is not null_block:
+                    self._allocator.free(b)
+                    self._blocks[idx] = null_block
+
+        # Ensure there are enough empty slots for the new tokens plus
+        # lookahead slots
+        self.ensure_num_empty_slots(num_empty_slots=len(token_ids) +
+                                    num_lookahead_slots)
+
+        # Update the blocks with the new tokens
+        first_block_idx = self._num_full_slots // self._block_size
+        token_blocks = self._chunk_token_blocks_for_append(token_ids)
+
+        for i, token_block in enumerate(token_blocks):
+            self._blocks.append_token_ids(first_block_idx + i, token_block)
+
+        self._num_full_slots += len(token_ids)
+
+    def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
+        """Ensures that the BlockTable has at least the specified number of
+        empty slots available.
+
+        This method checks if the BlockTable has enough empty slots (i.e.,
+        available space) to accommodate the requested number of tokens. If not,
+        it allocates additional blocks on the GPU to ensure that the required
+        number of empty slots is available.
+
+        Args:
+            num_empty_slots (int): The minimum number of empty slots required.
+        """
+        # Currently the block table only supports
+        # appending tokens to GPU blocks.
+        device = Device.GPU
+        assert self._is_allocated
+
+        if self._num_empty_slots >= num_empty_slots:
+            return
+
+        slots_to_allocate = num_empty_slots - self._num_empty_slots
+        blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)
+
+        for _ in range(blocks_to_allocate):
+            assert len(self._blocks) > 0
+            self._blocks.append(
+                self._allocator.allocate_mutable_block(
+                    prev_block=self._blocks[-1], device=device))
+
+    def fork(self) -> "BlockTable":
+        """Creates a new BlockTable instance with a copy of the blocks from the
+        current instance.
+
+        This method creates a new BlockTable instance with the same block size,
+        block allocator, and a copy of the blocks from the current instance. The
+        new BlockTable has its own independent set of blocks, but shares the
+        same underlying memory allocation with the original BlockTable.
+
+        Returns:
+            BlockTable: A new BlockTable instance with a copy of the blocks from
+                the current instance.
+        """
+        assert self._is_allocated
+        assert len(self._blocks) > 0
+        forked_blocks = self._allocator.fork(self._blocks[-1])
+        return BlockTable(
+            block_size=self._block_size,
+            block_allocator=self._allocator,
+            _blocks=forked_blocks,
+            max_block_sliding_window=self._max_block_sliding_window,
+            cache_namespace=self._cache_namespace,
+        )
+
+    def free(self) -> None:
+        """Frees the memory occupied by the blocks in the BlockTable.
+
+        This method iterates over all the blocks in the `_blocks` list and calls
+        the `free` method of the `_allocator` object to release the memory
+        occupied by each block. After freeing all the blocks, the `_blocks` list
+        is set to `None`.
+        """
+        for block in self.blocks:
+            self._allocator.free(block)
+        self._blocks.reset()
+
+    @property
+    def physical_block_ids(self) -> List[int]:
+        """Returns a list of physical block indices for the blocks in the
+        BlockTable.
+
+        This property returns a list of integers, where each integer represents
+        the physical block index of a corresponding block in the `_blocks` list.
+        The physical block index is a unique identifier for the memory location
+        occupied by the block.
+
+        Returns:
+            List[int]: A list of physical block indices for the blocks in the
+                BlockTable.
+        """
+        return self._blocks.ids()
+
+    def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:
+        """Get the number of "unseen" tokens in the sequence.
+
+        Unseen tokens are tokens in the sequence corresponding to this block
+        table, but are not yet appended to this block table.
+
+        Args:
+            sequence_token_ids (List[int]): The list of token ids in the
+                sequence.
+
+        Returns:
+            List[int]: The postfix of sequence_token_ids that has not yet been
+                appended to the block table.
+        """
+
+        # Since the block table is append-only, the unseen token ids are the
+        # ones after the appended ones.
+        return sequence_token_ids[self.num_full_slots:]
+
+    def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
+                                       token_ids: List[int],
+                                       device: Device) -> List[Block]:
+        blocks: List[Block] = []
+
+        block_token_ids = []
+        tail_token_ids = []
+        for cur_token_ids in chunk_list(token_ids, self._block_size):
+            if len(cur_token_ids) == self._block_size:
+                block_token_ids.append(cur_token_ids)
+            else:
+                tail_token_ids.append(cur_token_ids)
+
+        if block_token_ids:
+            blocks.extend(self._allocate_immutable_blocks(
+                prev_block=prev_block,
+                block_token_ids=block_token_ids,
+                device=device))
+            prev_block = blocks[-1]
+
+        if tail_token_ids:
+            assert len(tail_token_ids) == 1
+            cur_token_ids = tail_token_ids[0]
+
+            block = self._allocate_mutable_block(prev_block=prev_block,
+                                                 device=device)
+            block.append_token_ids(cur_token_ids)
+
+            blocks.append(block)
+
+        return blocks
+
+    def _allocate_mutable_block(self, prev_block: Optional[Block],
+                                device: Device) -> Block:
+        if self._cache_namespace is None:
+            return self._allocator.allocate_mutable_block(
+                prev_block=prev_block, device=device)
+
+        with_cache_namespace = getattr(
+            self._allocator, "allocate_mutable_block_with_cache_namespace",
+            None)
+        if callable(with_cache_namespace):
+            return with_cache_namespace(
+                prev_block=prev_block,
+                cache_namespace=self._cache_namespace,
+                device=device)
+
+        backend_allocators = getattr(self._allocator, "_allocators", None)
+        if isinstance(backend_allocators, dict):
+            device_allocator = backend_allocators.get(device)
+            if device_allocator is not None:
+                with_cache_namespace = getattr(
+                    device_allocator,
+                    "allocate_mutable_block_with_cache_namespace", None)
+                if callable(with_cache_namespace):
+                    return with_cache_namespace(
+                        prev_block=prev_block,
+                        cache_namespace=self._cache_namespace)
+
+        return self._allocator.allocate_mutable_block(
+            prev_block=prev_block, device=device)
+
+    def _allocate_immutable_blocks(self,
+                                  prev_block: Optional[Block],
+                                  block_token_ids: List[List[int]],
+                                  device: Device) -> List[Block]:
+        if self._cache_namespace is None:
+            return self._allocator.allocate_immutable_blocks(
+                prev_block,
+                block_token_ids=block_token_ids,
+                device=device)
+
+        with_cache_namespace = getattr(
+            self._allocator, "allocate_immutable_blocks_with_cache_namespace", None)
+        if callable(with_cache_namespace):
+            return with_cache_namespace(
+                prev_block=prev_block,
+                block_token_ids=block_token_ids,
+                cache_namespace=self._cache_namespace,
+                device=device)
+
+        backend_allocator = getattr(self._allocator, "_allocators", None)
+        if isinstance(backend_allocator, dict):
+            device_allocator = backend_allocator.get(device)
+            if device_allocator is not None:
+                with_cache_namespace = getattr(
+                    device_allocator,
+                    "allocate_immutable_blocks_with_cache_namespace",
+                    None)
+                if callable(with_cache_namespace):
+                    return with_cache_namespace(
+                        prev_block=prev_block,
+                        block_token_ids=block_token_ids,
+                        cache_namespace=self._cache_namespace)
+
+        # Fallback: keep behavior identical when no namespace-aware allocator
+        # is available.
+        return self._allocator.allocate_immutable_blocks(
+            prev_block,
+            block_token_ids=block_token_ids,
+            device=device)
+
+    def _get_all_token_ids(self) -> List[int]:
+        # NOTE: This function is O(seq_len); use sparingly.
+        token_ids: List[int] = []
+
+        if not self._is_allocated:
+            return token_ids
+
+        for block in self.blocks:
+            token_ids.extend(block.token_ids)
+
+        return token_ids
+
+    def _get_num_token_ids(self) -> int:
+        res = 0
+        for block in self.blocks:
+            res += len(block.token_ids)
+
+        return res
+
+    @property
+    def _is_allocated(self) -> bool:
+        return len(self._blocks) > 0
+
+    @property
+    def blocks(self) -> List[Block]:
+        return self._blocks.list()
+
+    @property
+    def _num_empty_slots(self) -> int:
+        assert self._is_allocated
+        return len(self._blocks) * self._block_size - self._num_full_slots
+
+    @property
+    def num_full_slots(self) -> int:
+        """Returns the total number of tokens currently stored in the
+        BlockTable.
+
+        Returns:
+            int: The total number of tokens currently stored in the BlockTable.
+        """
+        return self._num_full_slots
+
+    def get_num_blocks_touched_by_append_slots(
+            self, token_ids: List[int], num_lookahead_slots: int) -> int:
+        """Determine how many blocks will be "touched" by appending the token
+        ids.
+
+        This is required for the scheduler to determine whether a sequence can
+        continue generation, or if it must be preempted.
+        """
+        # Math below is equivalent to:
+        # all_token_ids = token_ids + [-1] * num_lookahead_slots
+        # token_blocks = self._chunk_token_blocks_for_append(all_token_ids)
+        # return len(token_blocks)
+
+        num_token_ids = len(token_ids) + num_lookahead_slots
+        first_chunk_size = self._block_size - (self._num_full_slots %
+                                               self._block_size)
+        num_token_blocks = (1 + math.ceil(
+            (num_token_ids - first_chunk_size) / self._block_size))
+        return num_token_blocks
+
+    def _chunk_token_blocks_for_append(
+            self, token_ids: List[int]) -> List[List[int]]:
+        """Split the token ids into block-sized chunks so they can be easily
+        appended to blocks. The first such "token block" may have less token ids
+        than the block size, since the last allocated block may be partially
+        full.
+
+        If no token ids are provided, then no chunks are returned.
+        """
+
+        if not token_ids:
+            return []
+
+        first_chunk_size = self._block_size - (self._num_full_slots %
+                                               self._block_size)
+        token_blocks = [token_ids[:first_chunk_size]]
+        token_blocks.extend(
+            chunk_list(token_ids[first_chunk_size:], self._block_size))
+        return token_blocks

qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py

@@ -0,0 +1,475 @@
+from typing import Dict, FrozenSet, List, Optional, Tuple
+
+from vllm.core.block.cpu_kv_content_cache import (CpuKvContentCache,
+                                                   cpu_kv_offload_enabled)
+from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,
+                                        DeviceAwareBlockAllocator)
+from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
+from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator
+from vllm.utils import Device
+
+
+class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):
+    """A block allocator that can allocate blocks on both CPU and GPU memory.
+
+    This class implements the `DeviceAwareBlockAllocator` interface and provides
+    functionality for allocating and managing blocks of memory on both CPU and
+    GPU devices.
+
+    The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU
+    blocks, and allows for allocation, deallocation, forking, and swapping of
+    blocks across these memory pools.
+    """
+
+    @staticmethod
+    def create(
+        allocator_type: str,
+        num_gpu_blocks: int,
+        num_cpu_blocks: int,
+        block_size: int,
+    ) -> DeviceAwareBlockAllocator:
+        """Creates a CpuGpuBlockAllocator instance with the specified
+        configuration.
+
+        This static method creates and returns a CpuGpuBlockAllocator instance
+        based on the provided parameters. It initializes the CPU and GPU block
+        allocators with the specified number of blocks, block size, and
+        allocator type.
+
+        Args:
+            allocator_type (str): The type of block allocator to use for CPU
+                and GPU blocks. Currently supported values are "naive" and
+                "prefix_caching".
+            num_gpu_blocks (int): The number of blocks to allocate for GPU
+                memory.
+            num_cpu_blocks (int): The number of blocks to allocate for CPU
+                memory.
+            block_size (int): The size of each block in number of tokens.
+
+        Returns:
+            DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the
+                specified configuration.
+
+        Notes:
+            - The block IDs are assigned contiguously, with GPU block IDs coming
+                before CPU block IDs.
+        """
+        content_offload = cpu_kv_offload_enabled()
+        if content_offload and allocator_type != "prefix_caching":
+            raise RuntimeError(
+                "BI100_CPU_KV_OFFLOAD=1 requires prefix caching")
+        if content_offload and num_cpu_blocks <= 0:
+            raise RuntimeError(
+                "BI100_CPU_KV_OFFLOAD=1 requires at least one CPU KV block")
+
+        block_ids = list(range(num_gpu_blocks + num_cpu_blocks))
+        gpu_block_ids = block_ids[:num_gpu_blocks]
+        cpu_block_ids = block_ids[num_gpu_blocks:]
+
+        if allocator_type == "naive":
+            gpu_allocator: BlockAllocator = NaiveBlockAllocator(
+                create_block=NaiveBlock,  # type: ignore
+                num_blocks=num_gpu_blocks,
+                block_size=block_size,
+                block_ids=gpu_block_ids,
+            )
+
+            cpu_allocator: BlockAllocator = NaiveBlockAllocator(
+                create_block=NaiveBlock,  # type: ignore
+                num_blocks=num_cpu_blocks,
+                block_size=block_size,
+                block_ids=cpu_block_ids,
+            )
+        elif allocator_type == "prefix_caching":
+            gpu_allocator = PrefixCachingBlockAllocator(
+                num_blocks=num_gpu_blocks,
+                block_size=block_size,
+                block_ids=gpu_block_ids,
+            )
+
+            cpu_allocator = PrefixCachingBlockAllocator(
+                num_blocks=num_cpu_blocks,
+                block_size=block_size,
+                block_ids=cpu_block_ids,
+            )
+        else:
+            raise ValueError(f"Unknown allocator type {allocator_type=}")
+
+        return CpuGpuBlockAllocator(
+            cpu_block_allocator=cpu_allocator,
+            gpu_block_allocator=gpu_allocator,
+            cpu_content_cache=(CpuKvContentCache(num_cpu_blocks)
+                               if content_offload else None),
+        )
+
+    def __init__(self, cpu_block_allocator: BlockAllocator,
+                 gpu_block_allocator: BlockAllocator,
+                 cpu_content_cache: Optional[CpuKvContentCache] = None):
+        assert not (
+            cpu_block_allocator.all_block_ids
+            & gpu_block_allocator.all_block_ids
+        ), "cpu and gpu block allocators can't have intersection of block ids"
+
+        self._allocators = {
+            Device.CPU: cpu_block_allocator,
+            Device.GPU: gpu_block_allocator,
+        }
+
+        self._swap_mapping: Dict[int, int] = {}
+        self._null_block: Optional[Block] = None
+        self._cpu_content_cache = cpu_content_cache
+
+        self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
+        for _, allocator in self._allocators.items():
+            for block_id in allocator.all_block_ids:
+                self._block_ids_to_allocator[block_id] = allocator
+
+        if self._cpu_content_cache is not None:
+            if not isinstance(gpu_block_allocator,
+                              PrefixCachingBlockAllocator):
+                raise RuntimeError(
+                    "CPU KV content tier requires PrefixCachingBlockAllocator")
+            if (self._cpu_content_cache.capacity !=
+                    cpu_block_allocator.get_num_total_blocks()):
+                raise RuntimeError(
+                    "CPU KV content capacity must cover the complete CPU cache")
+            gpu_block_allocator.set_external_cache_callbacks(
+                claim=self._claim_cpu_content,
+                load=self._stage_cpu_to_gpu,
+                cancel=self._cancel_cpu_claim,
+                store=self._stage_gpu_to_cpu,
+            )
+
+    @property
+    def content_offload_enabled(self) -> bool:
+        return self._cpu_content_cache is not None
+
+    def _claim_cpu_content(self, content_hash: bytes) -> Optional[int]:
+        assert self._cpu_content_cache is not None
+        return self._cpu_content_cache.claim_load(content_hash)
+
+    def _cancel_cpu_claim(self, content_hash: bytes, cpu_slot: int) -> None:
+        assert self._cpu_content_cache is not None
+        self._cpu_content_cache.cancel_load(content_hash, cpu_slot)
+
+    def _stage_cpu_to_gpu(self, content_hash: bytes, cpu_slot: int,
+                          gpu_block_id: BlockId) -> None:
+        assert self._cpu_content_cache is not None
+        gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
+        self._cpu_content_cache.stage_load(
+            content_hash, cpu_slot, gpu_slot)
+
+    def _stage_gpu_to_cpu(self, content_hash: bytes,
+                          gpu_block_id: BlockId) -> bool:
+        assert self._cpu_content_cache is not None
+        gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
+        return self._cpu_content_cache.stage_store(content_hash, gpu_slot)
+
+    def allocate_or_get_null_block(self) -> Block:
+        if self._null_block is None:
+            self._null_block = NullBlock(
+                self.allocate_mutable_block(None, Device.GPU))
+        return self._null_block
+
+    def allocate_mutable_block(self, prev_block: Optional[Block],
+                               device: Device) -> Block:
+        """Allocates a new mutable block on the specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block to in the sequence.
+                Used for prefix hashing.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            Block: The newly allocated mutable block.
+        """
+        return self._allocators[device].allocate_mutable_block(prev_block)
+
+    def allocate_immutable_blocks(self, prev_block: Optional[Block],
+                                  block_token_ids: List[List[int]],
+                                  device: Device) -> List[Block]:
+        """Allocates a new group of immutable blocks with the provided block 
+        token IDs on the specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block in the sequence.
+                Used for prefix hashing.
+            block_token_ids (List[int]): The list of block token IDs to be 
+                stored in the new blocks.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            List[Block]: The newly allocated list of immutable blocks 
+                containing the provided block token IDs.
+        """
+        return self._allocators[device].allocate_immutable_blocks(
+            prev_block, block_token_ids)
+
+    def allocate_immutable_block(self, prev_block: Optional[Block],
+                                 token_ids: List[int],
+                                 device: Device) -> Block:
+        """Allocates a new immutable block with the provided token IDs on the
+        specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block in the sequence.
+                Used for prefix hashing.
+            token_ids (List[int]): The list of token IDs to be stored in the new
+                block.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            Block: The newly allocated immutable block containing the provided
+                token IDs.
+        """
+        return self._allocators[device].allocate_immutable_block(
+            prev_block, token_ids)
+
+    def free(self, block: Block) -> None:
+        """Frees the memory occupied by the given block.
+
+        Args:
+            block (Block): The block to be freed.
+        """
+        # Null block should never be freed
+        if isinstance(block, NullBlock):
+            return
+        block_id = block.block_id
+        assert block_id is not None
+        allocator = self._block_ids_to_allocator[block_id]
+        allocator.free(block)
+
+    def fork(self, last_block: Block) -> List[Block]:
+        """Creates a new sequence of blocks that shares the same underlying
+            memory as the original sequence.
+
+        Args:
+            last_block (Block): The last block in the original sequence.
+
+        Returns:
+            List[Block]: A new list of blocks that shares the same memory as the
+                original sequence.
+        """
+        # do not attempt to fork the null block
+        assert not isinstance(last_block, NullBlock)
+        block_id = last_block.block_id
+        assert block_id is not None
+        allocator = self._block_ids_to_allocator[block_id]
+        return allocator.fork(last_block)
+
+    def get_num_free_blocks(self, device: Device) -> int:
+        """Returns the number of free blocks available on the specified device.
+
+        Args:
+            device (Device): The device for which to query the number of free
+                blocks. AssertionError is raised if None is passed.
+
+        Returns:
+            int: The number of free blocks available on the specified device.
+        """
+        return self._allocators[device].get_num_free_blocks()
+
+    def get_num_total_blocks(self, device: Device) -> int:
+        return self._allocators[device].get_num_total_blocks()
+
+    def get_physical_block_id(self, device: Device, absolute_id: int) -> int:
+        """Returns the zero-offset block id on certain device given the 
+        absolute block id.
+
+        Args:
+            device (Device): The device for which to query relative block id.
+                absolute_id (int): The absolute block id for the block in 
+                whole allocator.
+
+        Returns:
+            int: The zero-offset block id on certain device.
+        """
+        return self._allocators[device].get_physical_block_id(absolute_id)
+
+    def swap(self, blocks: List[Block], src_device: Device,
+             dst_device: Device) -> Dict[int, int]:
+        """Execute the swap for the given blocks from source_device
+        on to dest_device, save the current swap mapping and append 
+        them to the accumulated `self._swap_mapping` for each 
+        scheduling move.
+
+        Args:
+            blocks: List of blocks to be swapped.
+            src_device (Device): Device to swap the 'blocks' from.
+            dst_device (Device): Device to swap the 'blocks' to.
+        
+        Returns:
+            Dict[int, int]: Swap mapping from source_device
+                on to dest_device.
+        """
+        if self.content_offload_enabled:
+            raise RuntimeError(
+                "request-level preemption swap cannot share CPU slots with "
+                "BI100_CPU_KV_OFFLOAD")
+
+        src_block_ids = [block.block_id for block in blocks]
+        self._allocators[src_device].swap_out(blocks)
+        self._allocators[dst_device].swap_in(blocks)
+        dst_block_ids = [block.block_id for block in blocks]
+
+        current_swap_mapping: Dict[int, int] = {}
+        for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids):
+            if src_block_id is not None and dst_block_id is not None:
+                self._swap_mapping[src_block_id] = dst_block_id
+                current_swap_mapping[src_block_id] = dst_block_id
+        return current_swap_mapping
+
+    def get_num_full_blocks_touched(self, blocks: List[Block],
+                                    device: Device) -> int:
+        """Returns the number of full blocks that will be touched by
+        swapping in/out the given blocks on to the 'device'.
+
+        Args:
+            blocks: List of blocks to be swapped.
+            device (Device): Device to swap the 'blocks' on.
+
+        Returns:
+            int: the number of full blocks that will be touched by
+                swapping in/out the given blocks on to the 'device'.
+                Non full blocks are ignored when deciding the number
+                of blocks to touch.
+        """
+        return self._allocators[device].get_num_full_blocks_touched(blocks)
+
+    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
+        """Clears the copy-on-write (CoW) state and returns the mapping of
+            source to destination block IDs.
+
+        Returns:
+            List[Tuple[int, int]]: A list mapping source block IDs to 
+                destination block IDs.
+        """
+        # CoW only supported on GPU
+        device = Device.GPU
+        return self._allocators[device].clear_copy_on_writes()
+
+    def mark_blocks_as_accessed(self, block_ids: List[int],
+                                now: float) -> None:
+        """Mark blocks as accessed, only use for prefix caching."""
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].mark_blocks_as_accessed(block_ids, now)
+
+    def mark_blocks_as_computed(self, block_ids: List[int]) -> None:
+        """Mark blocks as accessed, only use for prefix caching."""
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].mark_blocks_as_computed(block_ids)
+
+    def get_computed_block_ids(self, prev_computed_block_ids: List[int],
+                               block_ids: List[int],
+                               skip_last_block_id: bool) -> List[int]:
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].get_computed_block_ids(
+            prev_computed_block_ids, block_ids, skip_last_block_id)
+
+    def get_common_computed_block_ids(
+            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].get_common_computed_block_ids(
+            computed_seq_block_ids)
+
+    @property
+    def all_block_ids(self) -> FrozenSet[int]:
+        return frozenset(self._block_ids_to_allocator.keys())
+
+    def get_prefix_cache_hit_rate(self, device: Device) -> float:
+        """Prefix cache hit rate. -1 means not supported or disabled."""
+        assert device in self._allocators
+        return self._allocators[device].get_prefix_cache_hit_rate()
+
+    def get_and_reset_swaps(self) -> List[Tuple[int, int]]:
+        """Returns and clears the mapping of source to destination block IDs.
+        Will be called after every swapping operations for now, and after every
+        schedule when BlockManagerV2 become default. Currently not useful.
+
+        Returns:
+            List[Tuple[int, int]]: A mapping of source to destination block IDs.
+        """
+        mapping = self._swap_mapping.copy()
+        self._swap_mapping.clear()
+        return list(mapping.items())
+
+    def get_and_reset_prefix_swaps(
+            self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
+        """Return scheduler-owned (CPU->GPU, GPU->CPU) content maps."""
+        if self._cpu_content_cache is None:
+            return [], []
+        return self._cpu_content_cache.drain_step()
+
+    def begin_prefix_cache_step(self) -> None:
+        if self._cpu_content_cache is not None:
+            self._cpu_content_cache.begin_step()
+
+
+class NullBlock(Block):
+    """
+    Null blocks are used as a placeholders for KV cache blocks that have
+    been dropped due to sliding window.
+    This implementation just wraps an ordinary block and prevents it from
+    being modified. It also allows for testing if a block is NullBlock
+    via isinstance().
+    """
+
+    def __init__(self, proxy: Block):
+        super().__init__()
+        self._proxy = proxy
+
+    def append_token_ids(self, token_ids: List[BlockId]):
+        raise ValueError("null block should not be modified")
+
+    @property
+    def block_id(self):
+        return self._proxy.block_id
+
+    @block_id.setter
+    def block_id(self, value: Optional[BlockId]):
+        raise ValueError("null block should not be modified")
+
+    @property
+    def token_ids(self) -> List[BlockId]:
+        return self._proxy.token_ids
+
+    @property
+    def num_tokens_total(self) -> int:
+        raise NotImplementedError(
+            "num_tokens_total is not used for null block")
+
+    @property
+    def num_empty_slots(self) -> BlockId:
+        return self._proxy.num_empty_slots
+
+    @property
+    def is_full(self):
+        return self._proxy.is_full
+
+    @property
+    def prev_block(self):
+        return self._proxy.prev_block
+
+    @property
+    def computed(self):
+        return self._proxy.computed
+
+    @computed.setter
+    def computed(self, value):
+        self._proxy.computed = value
+
+    @property
+    def last_accessed(self) -> float:
+        return self._proxy.last_accessed
+
+    @last_accessed.setter
+    def last_accessed(self, last_accessed_ts: float):
+        self._proxy.last_accessed = last_accessed_ts
+
+    @property
+    def content_hash(self):
+        return self._proxy.content_hash

qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py

@@ -0,0 +1,255 @@
+"""Scheduler-owned content index for an inclusive CPU KV cache tier."""
+
+from __future__ import annotations
+
+import heapq
+import os
+from collections import OrderedDict
+from typing import Dict, List, Mapping, Optional, Set, Tuple
+
+
+ContentHash = bytes
+SwapMapping = List[Tuple[int, int]]
+
+
+def cpu_kv_offload_enabled(
+    environ: Optional[Mapping[str, str]] = None,
+) -> bool:
+    """Read the experimental selector without accepting ambiguous values."""
+    source = os.environ if environ is None else environ
+    value = source.get("BI100_CPU_KV_OFFLOAD", "0")
+    if value == "0":
+        return False
+    if value == "1":
+        return True
+    raise RuntimeError(
+        "BI100_CPU_KV_OFFLOAD must be exactly '0' or '1', "
+        f"got {value!r}")
+
+
+class CpuKvContentCache:
+    """Track immutable KV blocks held in the worker's pinned CPU cache.
+
+    The scheduler owns this metadata and sends identical physical block maps
+    to every tensor-parallel worker. CPU copies are inclusive: loading a block
+    back to GPU does not remove its CPU entry. Slots touched by either transfer
+    direction are pinned for the whole scheduling step so a D2H destination
+    can never overwrite an H2D source before workers execute the maps.
+    """
+
+    def __init__(self, capacity: int) -> None:
+        if capacity <= 0:
+            raise ValueError("CPU KV content cache capacity must be positive")
+        self.capacity = capacity
+        self._hash_to_slot: Dict[ContentHash, int] = {}
+        self._slot_to_hash: Dict[int, ContentHash] = {}
+        self._ready_slots: Set[int] = set()
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self._free_slots = list(range(capacity))
+        heapq.heapify(self._free_slots)
+
+        self._step_slots_in_use: Set[int] = set()
+        self._step_load_slots: Set[int] = set()
+        self._step_h2d: Dict[int, int] = {}
+        self._step_d2h: Dict[int, int] = {}
+        self._deferred_d2h: Dict[int, ContentHash] = {}
+        self._deferred_hashes: Set[ContentHash] = set()
+        self._pending_ready_slots: Set[int] = set()
+
+        self.hits = 0
+        self.misses = 0
+        self.stores = 0
+        self.deduplicated_stores = 0
+        self.evictions = 0
+        self.skipped_stores = 0
+
+    @staticmethod
+    def _validate_hash(content_hash: ContentHash) -> None:
+        if not isinstance(content_hash, bytes) or len(content_hash) != 32:
+            raise ValueError("CPU KV cache key must be a 32-byte content hash")
+
+    @staticmethod
+    def _validate_block_id(name: str, block_id: int) -> None:
+        if not isinstance(block_id, int) or isinstance(block_id, bool):
+            raise TypeError(f"{name} must be an integer")
+        if block_id < 0:
+            raise ValueError(f"{name} must be non-negative")
+
+    def _touch(self, slot: int) -> None:
+        self._lru.pop(slot, None)
+        self._lru[slot] = None
+
+    def _select_store_slot(self) -> Optional[int]:
+        if self._free_slots:
+            return heapq.heappop(self._free_slots)
+        for slot in self._lru:
+            if slot not in self._step_slots_in_use:
+                return slot
+        return None
+
+    def _commit_store(self, content_hash: ContentHash,
+                      gpu_block: int, slot: int) -> None:
+        old_hash = self._slot_to_hash.get(slot)
+        if old_hash is not None:
+            if slot in self._step_slots_in_use:
+                raise RuntimeError("selected an in-use CPU KV slot for eviction")
+            del self._hash_to_slot[old_hash]
+            self._ready_slots.discard(slot)
+            self.evictions += 1
+
+        if slot in self._step_h2d:
+            raise RuntimeError(
+                "a CPU KV slot cannot be an H2D source and D2H destination "
+                "in one scheduler step")
+        if slot in self._step_d2h.values():
+            raise RuntimeError(f"duplicate D2H destination CPU slot {slot}")
+
+        self._hash_to_slot[content_hash] = slot
+        self._slot_to_hash[slot] = content_hash
+        self._ready_slots.discard(slot)
+        self._step_slots_in_use.add(slot)
+        self._step_d2h[gpu_block] = slot
+        self._touch(slot)
+        self.stores += 1
+
+    def begin_step(self) -> None:
+        """Publish D2H stores returned by the preceding synchronous step."""
+        if (self._step_slots_in_use or self._step_h2d or self._step_d2h
+                or self._deferred_d2h or self._deferred_hashes):
+            raise RuntimeError("cannot begin a CPU KV step before draining it")
+        self._ready_slots.update(self._pending_ready_slots)
+        self._pending_ready_slots.clear()
+
+    def _require_step_started(self) -> None:
+        if self._pending_ready_slots:
+            raise RuntimeError(
+                "CPU KV step must begin before content lookup or eviction")
+
+    def claim_load(self, content_hash: ContentHash) -> Optional[int]:
+        """Pin and return a ready CPU source for this scheduling step."""
+        self._validate_hash(content_hash)
+        self._require_step_started()
+        slot = self._hash_to_slot.get(content_hash)
+        if slot is None or slot not in self._ready_slots:
+            self.misses += 1
+            return None
+        if slot in self._step_slots_in_use:
+            raise RuntimeError(
+                f"CPU KV slot {slot} was claimed twice in one scheduler step")
+        self._step_slots_in_use.add(slot)
+        self._step_load_slots.add(slot)
+        self._touch(slot)
+        self.hits += 1
+        return slot
+
+    def cancel_load(self, content_hash: ContentHash, cpu_slot: int) -> None:
+        """Release a claim when GPU allocation fails before H2D is staged."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("cpu_slot", cpu_slot)
+        if self._hash_to_slot.get(content_hash) != cpu_slot:
+            raise RuntimeError("CPU KV load cancellation key/slot mismatch")
+        if cpu_slot in self._step_h2d:
+            raise RuntimeError("cannot cancel a CPU KV load after H2D staging")
+        if cpu_slot not in self._step_slots_in_use:
+            raise RuntimeError("cannot cancel an unclaimed CPU KV load")
+        self._step_slots_in_use.remove(cpu_slot)
+        self._step_load_slots.remove(cpu_slot)
+
+    def stage_load(self, content_hash: ContentHash, cpu_slot: int,
+                   gpu_block: int) -> None:
+        """Stage one CPU-to-GPU promotion after the GPU slot is reserved."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("cpu_slot", cpu_slot)
+        self._validate_block_id("gpu_block", gpu_block)
+        if self._hash_to_slot.get(content_hash) != cpu_slot:
+            raise RuntimeError("CPU KV load key/slot mismatch")
+        if cpu_slot not in self._ready_slots:
+            raise RuntimeError("CPU KV load source is not ready")
+        if cpu_slot not in self._step_slots_in_use:
+            raise RuntimeError("CPU KV load source was not claimed")
+        if cpu_slot in self._step_h2d:
+            raise RuntimeError(f"duplicate H2D source CPU slot {cpu_slot}")
+        if gpu_block in self._step_h2d.values():
+            raise RuntimeError(f"duplicate H2D destination GPU block {gpu_block}")
+        if cpu_slot in self._step_d2h.values():
+            raise RuntimeError(
+                "a CPU KV slot cannot be an H2D source and D2H destination "
+                "in one scheduler step")
+        self._step_h2d[cpu_slot] = gpu_block
+
+    def stage_store(self, content_hash: ContentHash,
+                    gpu_block: int) -> bool:
+        """Stage a lazy GPU-to-CPU copy for an evicted immutable block."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("gpu_block", gpu_block)
+        self._require_step_started()
+        if (content_hash in self._hash_to_slot
+                or content_hash in self._deferred_hashes):
+            slot = self._hash_to_slot.get(content_hash)
+            if slot is not None:
+                self._touch(slot)
+            self.deduplicated_stores += 1
+            return False
+        if gpu_block in self._step_d2h or gpu_block in self._deferred_d2h:
+            raise RuntimeError(f"duplicate D2H source GPU block {gpu_block}")
+
+        if self._free_slots:
+            self._commit_store(
+                content_hash, gpu_block, heapq.heappop(self._free_slots))
+            return True
+
+        # Do not replace resident content until every lookup in this scheduler
+        # step is known. A later H2D claim can refer to any current LRU entry.
+        self._deferred_d2h[gpu_block] = content_hash
+        self._deferred_hashes.add(content_hash)
+        return True
+
+    def _resolve_deferred_stores(self) -> None:
+        if self._step_load_slots:
+            self.skipped_stores += len(self._deferred_d2h)
+        else:
+            for gpu_block, content_hash in self._deferred_d2h.items():
+                slot = self._select_store_slot()
+                if slot is None:
+                    self.skipped_stores += 1
+                    continue
+                self._commit_store(content_hash, gpu_block, slot)
+        self._deferred_d2h.clear()
+        self._deferred_hashes.clear()
+
+    def drain_step(self) -> Tuple[SwapMapping, SwapMapping]:
+        """Finalize this synchronous step and return (H2D, D2H) maps."""
+        self._resolve_deferred_stores()
+        transfer_slots = (
+            set(self._step_h2d) | set(self._step_d2h.values()))
+        if transfer_slots != self._step_slots_in_use:
+            raise RuntimeError(
+                "CPU KV scheduler step contains an uncommitted slot claim")
+        if set(self._step_h2d) & set(self._step_d2h.values()):
+            raise RuntimeError(
+                "CPU KV scheduler step reuses a CPU slot across directions")
+        if set(self._step_h2d) != self._step_load_slots:
+            raise RuntimeError(
+                "CPU KV scheduler step contains an unstaged load claim")
+
+        swap_in = sorted(self._step_h2d.items())
+        swap_out = sorted(self._step_d2h.items())
+        self._pending_ready_slots.update(self._step_d2h.values())
+        self._step_h2d.clear()
+        self._step_d2h.clear()
+        self._step_slots_in_use.clear()
+        self._step_load_slots.clear()
+        return swap_in, swap_out
+
+    def resident_slot(self, content_hash: ContentHash) -> Optional[int]:
+        self._validate_hash(content_hash)
+        return self._hash_to_slot.get(content_hash)
+
+    def is_ready(self, content_hash: ContentHash) -> bool:
+        self._validate_hash(content_hash)
+        slot = self._hash_to_slot.get(content_hash)
+        return slot is not None and slot in self._ready_slots
+
+    @property
+    def resident_count(self) -> int:
+        return len(self._hash_to_slot)

qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py

@@ -0,0 +1,769 @@
+"""A block manager that manages token blocks."""
+import hashlib
+import os
+import struct
+from collections.abc import Mapping
+from typing import Any, Dict, List, Optional, Sequence as GenericSequence, Tuple
+
+try:
+    from PIL import Image
+except Exception:  # pragma: no cover - optional dependency in some envs
+    Image = None  # type: ignore
+try:
+    import torch
+except Exception:  # pragma: no cover - optional dependency in some envs
+    torch = None  # type: ignore
+
+from vllm.core.block.block_table import BlockTable
+from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
+from vllm.core.block.interfaces import Block
+from vllm.core.block.prefix_caching_block import (ComputedBlocksTracker,
+                                                  LastAccessBlocksTracker)
+from vllm.core.block.utils import check_no_caching_or_swa_for_blockmgr_encdec
+from vllm.core.interfaces import AllocStatus, BlockSpaceManager
+from vllm.logger import init_logger
+from vllm.sequence import Sequence, SequenceGroup, SequenceStatus
+from vllm.utils import Device
+
+SeqId = int
+EncoderSeqId = str
+
+logger = init_logger(__name__)
+
+
+class BlockSpaceManagerV2(BlockSpaceManager):
+    """BlockSpaceManager which manages the allocation of KV cache.
+
+    It owns responsibility for allocation, swapping, allocating memory for
+    autoregressively-generated tokens, and other advanced features such as
+    prefix caching, forking/copy-on-write, and sliding-window memory allocation.
+
+    This class implements the design described in
+    https://github.com/vllm-project/vllm/pull/3492.
+
+    Lookahead slots
+        The block manager has the notion of a "lookahead slot". These are slots
+        in the KV cache that are allocated for a sequence. Unlike the other
+        allocated slots, the content of these slots is undefined -- the worker
+        may use the memory allocations in any way.
+
+        In practice, a worker could use these lookahead slots to run multiple
+        forward passes for a single scheduler invocation. Each successive
+        forward pass would write KV activations to the corresponding lookahead
+        slot. This allows low inter-token latency use-cases, where the overhead
+        of continuous batching scheduling is amortized over >1 generated tokens.
+
+        Speculative decoding uses lookahead slots to store KV activations of
+        proposal tokens.
+
+        See https://github.com/vllm-project/vllm/pull/3250 for more information
+        on lookahead scheduling.
+
+    Args:
+        block_size (int): The size of each memory block.
+        num_gpu_blocks (int): The number of memory blocks allocated on GPU.
+        num_cpu_blocks (int): The number of memory blocks allocated on CPU.
+        watermark (float, optional): The threshold used for memory swapping.
+            Defaults to 0.01.
+        sliding_window (Optional[int], optional): The size of the sliding
+            window. Defaults to None.
+        enable_caching (bool, optional): Flag indicating whether caching is
+            enabled. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        block_size: int,
+        num_gpu_blocks: int,
+        num_cpu_blocks: int,
+        watermark: float = 0.01,
+        sliding_window: Optional[int] = None,
+        enable_caching: bool = False,
+    ) -> None:
+        self.block_size = block_size
+        self.num_total_gpu_blocks = num_gpu_blocks
+        self.num_total_cpu_blocks = num_cpu_blocks
+
+        self.sliding_window = sliding_window
+        # max_block_sliding_window is the max number of blocks that need to be
+        # allocated
+        self.max_block_sliding_window = None
+        if sliding_window is not None:
+            # +1 here because // rounds down
+            num_blocks = sliding_window // block_size + 1
+            # +1 here because the last block may not be full,
+            # and so the sequence stretches one more block at the beginning
+            # For example, if sliding_window is 3 and block_size is 4,
+            # we may need 2 blocks when the second block only holds 1 token.
+            self.max_block_sliding_window = num_blocks + 1
+
+        self.watermark = watermark
+        assert watermark >= 0.0
+
+        self.enable_caching = enable_caching
+
+        self.watermark_blocks = int(watermark * num_gpu_blocks)
+
+        self.block_allocator = CpuGpuBlockAllocator.create(
+            allocator_type="prefix_caching" if enable_caching else "naive",
+            num_gpu_blocks=num_gpu_blocks,
+            num_cpu_blocks=num_cpu_blocks,
+            block_size=block_size,
+        )
+
+        self.block_tables: Dict[SeqId, BlockTable] = {}
+        self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {}
+        self._warned_mm_namespace_requests = set[str]()
+        self._request_local_namespace: Dict[str, bytes] = {}
+        self._runtime_cache_namespace = self._build_runtime_cache_namespace()
+
+        self._computed_blocks_tracker = ComputedBlocksTracker(
+            self.block_allocator)
+        self._last_access_blocks_tracker = LastAccessBlocksTracker(
+            self.block_allocator)
+
+    def can_allocate(self,
+                     seq_group: SequenceGroup,
+                     num_lookahead_slots: int = 0) -> AllocStatus:
+        # FIXME(woosuk): Here we assume that all sequences in the group share
+        # the same prompt. This may not be true for preempted sequences.
+
+        check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
+
+        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
+        num_required_blocks = BlockTable.get_num_required_blocks(
+            seq.get_token_ids(),
+            block_size=self.block_size,
+            num_lookahead_slots=num_lookahead_slots,
+        )
+
+        if seq_group.is_encoder_decoder():
+            encoder_seq = seq_group.get_encoder_seq()
+            assert encoder_seq is not None
+            num_required_blocks += BlockTable.get_num_required_blocks(
+                encoder_seq.get_token_ids(),
+                block_size=self.block_size,
+            )
+
+        if self.max_block_sliding_window is not None:
+            num_required_blocks = min(num_required_blocks,
+                                      self.max_block_sliding_window)
+
+        num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
+            device=Device.GPU)
+
+        # Use watermark to avoid frequent cache eviction.
+        if (self.num_total_gpu_blocks - num_required_blocks <
+                self.watermark_blocks):
+            return AllocStatus.NEVER
+        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
+            return AllocStatus.OK
+        else:
+            return AllocStatus.LATER
+
+    def _allocate_sequence(
+        self,
+        seq: Sequence,
+        cache_namespace: Optional[bytes] = None,
+    ) -> BlockTable:
+        block_table = BlockTable(
+            block_size=self.block_size,
+            block_allocator=self.block_allocator,
+            max_block_sliding_window=self.max_block_sliding_window,
+            cache_namespace=cache_namespace,
+        )
+        if seq.get_token_ids():
+            # Add blocks to the block table only if the sequence is non empty.
+            block_table.allocate(seq.get_token_ids())
+
+        return block_table
+
+    def allocate(self, seq_group: SequenceGroup) -> None:
+
+        # Allocate self-attention block tables for decoder sequences
+        waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
+        assert not (set(seq.seq_id for seq in waiting_seqs)
+                    & self.block_tables.keys()), "block table already exists"
+
+        # NOTE: Here we assume that all sequences in the group have the same
+        # prompt.
+        seq = waiting_seqs[0]
+        request_id = seq_group.request_id
+        cache_namespace = self._get_cache_namespace(
+            seq,
+            request_id=request_id,
+            seq_group=seq_group,
+        )
+        block_table: BlockTable = self._allocate_sequence(
+            seq,
+            cache_namespace=cache_namespace,
+        )
+        self.block_tables[seq.seq_id] = block_table
+
+        # Track seq
+        self._computed_blocks_tracker.add_seq(seq.seq_id)
+        self._last_access_blocks_tracker.add_seq(seq.seq_id)
+
+        # Assign the block table for each sequence.
+        for seq in waiting_seqs[1:]:
+            self.block_tables[seq.seq_id] = block_table.fork()
+
+            # Track seq
+            self._computed_blocks_tracker.add_seq(seq.seq_id)
+            self._last_access_blocks_tracker.add_seq(seq.seq_id)
+
+        # Allocate cross-attention block table for encoder sequence
+        #
+        # NOTE: Here we assume that all sequences in the group have the same
+        # encoder prompt.
+        request_id = seq_group.request_id
+
+        assert (request_id
+                not in self.cross_block_tables), \
+            "block table already exists"
+
+        check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
+
+        if seq_group.is_encoder_decoder():
+            encoder_seq = seq_group.get_encoder_seq()
+            assert encoder_seq is not None
+            encoder_cache_namespace = self._get_cache_namespace(
+                encoder_seq,
+                request_id=request_id,
+                seq_group=seq_group)
+            block_table = self._allocate_sequence(
+                encoder_seq, cache_namespace=encoder_cache_namespace)
+            self.cross_block_tables[request_id] = block_table
+
+    @staticmethod
+    def _has_multi_modal_payload(multi_modal_data: Any) -> bool:
+        if multi_modal_data is None:
+            return False
+        if isinstance(multi_modal_data, Mapping):
+            try:
+                return len(multi_modal_data) > 0
+            except (TypeError, ValueError, RuntimeError, OSError,
+                    OverflowError, AttributeError, LookupError, struct.error):
+                # Treat an unusual mapping as payload and let normalization
+                # either identify it or select request-local isolation.
+                return True
+        return True
+
+    def _get_cache_namespace(self, seq: Sequence, request_id: str,
+                             seq_group: SequenceGroup) -> bytes:
+        digest = hashlib.sha256()
+        digest.update(b"bi100-request-prefix-namespace-v1|")
+        digest.update(self._runtime_cache_namespace)
+        digest.update(self._adapter_cache_namespace(seq_group))
+
+        multi_modal_data = seq.multi_modal_data
+        if self._has_multi_modal_payload(multi_modal_data):
+            try:
+                mm_namespace = self._hash_multi_modal_namespace(
+                    multi_modal_data)
+            except (TypeError, ValueError, RuntimeError, OSError,
+                    OverflowError, AttributeError, LookupError, struct.error):
+                if request_id not in self._warned_mm_namespace_requests:
+                    logger.warning(
+                        "Request %s has multimodal input that cannot be "
+                        "normalized for cache namespace hashing. Falling "
+                        "back to "
+                        "request-local namespace isolation.",
+                        request_id,
+                    )
+                    self._warned_mm_namespace_requests.add(request_id)
+                mm_namespace = self._request_local_fallback_cache_namespace(
+                    request_id=request_id)
+            digest.update(b"mm|")
+            digest.update(mm_namespace)
+        else:
+            digest.update(b"text|")
+        return digest.digest()
+
+    def can_append_slots(self, seq_group: SequenceGroup,
+                         num_lookahead_slots: int) -> bool:
+        """Determine if there is enough space in the GPU KV cache to continue
+        generation of the specified sequence group.
+
+        We use a worst-case heuristic: assume each touched block will require a
+        new allocation (either via CoW or new block). We can append slots if the
+        number of touched blocks is less than the number of free blocks.
+
+        "Lookahead slots" are slots that are allocated in addition to the slots
+        for known tokens. The contents of the lookahead slots are not defined.
+        This is used by speculative decoding when speculating future tokens.
+        """
+
+        num_touched_blocks = 0
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            block_table = self.block_tables[seq.seq_id]
+
+            num_touched_blocks += (
+                block_table.get_num_blocks_touched_by_append_slots(
+                    token_ids=block_table.get_unseen_token_ids(
+                        seq.get_token_ids()),
+                    num_lookahead_slots=num_lookahead_slots,
+                ))
+
+        num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
+            Device.GPU)
+        return num_touched_blocks <= num_free_gpu_blocks
+
+    def append_slots(
+        self,
+        seq: Sequence,
+        num_lookahead_slots: int,
+    ) -> List[Tuple[int, int]]:
+
+        block_table = self.block_tables[seq.seq_id]
+
+        block_table.append_token_ids(
+            token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()),
+            num_lookahead_slots=num_lookahead_slots,
+            num_computed_slots=seq.data.get_num_computed_tokens(),
+        )
+        # Return any new copy-on-writes.
+        new_cows = self.block_allocator.clear_copy_on_writes()
+        return new_cows
+
+    def free(self, seq: Sequence) -> None:
+        seq_id = seq.seq_id
+
+        if seq_id not in self.block_tables:
+            # Already freed or haven't been scheduled yet.
+            return
+
+        # Update seq block ids with the latest access time
+        self._last_access_blocks_tracker.update_seq_blocks_last_access(
+            seq_id, self.block_tables[seq.seq_id].physical_block_ids)
+
+        # Untrack seq
+        self._last_access_blocks_tracker.remove_seq(seq_id)
+        self._computed_blocks_tracker.remove_seq(seq_id)
+
+        # Free table/blocks
+        self.block_tables[seq_id].free()
+        del self.block_tables[seq_id]
+
+    def free_cross(self, seq_group: SequenceGroup) -> None:
+        request_id = seq_group.request_id
+        if request_id not in self.cross_block_tables:
+            # Already freed or hasn't been scheduled yet.
+            return
+        self.cross_block_tables[request_id].free()
+        del self.cross_block_tables[request_id]
+
+    def get_block_table(self, seq: Sequence) -> List[int]:
+        block_ids = self.block_tables[seq.seq_id].physical_block_ids
+        return block_ids  # type: ignore
+
+    def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]:
+        request_id = seq_group.request_id
+        assert request_id in self.cross_block_tables
+        block_ids = self.cross_block_tables[request_id].physical_block_ids
+        assert all(b is not None for b in block_ids)
+        return block_ids  # type: ignore
+
+    def access_all_blocks_in_seq(self, seq: Sequence, now: float):
+        if self.enable_caching:
+            # Record the latest access time for the sequence. The actual update
+            # of the block ids is deferred to the sequence free(..) call, since
+            # only during freeing of block ids, the blocks are actually added to
+            # the evictor (which is when the most updated time is required)
+            # (This avoids expensive calls to mark_blocks_as_accessed(..))
+            self._last_access_blocks_tracker.update_last_access(
+                seq.seq_id, now)
+
+    def mark_blocks_as_computed(self, seq_group: SequenceGroup,
+                                token_chunk_size: int):
+        # If prefix caching is enabled, mark immutable blocks as computed
+        # right after they have been scheduled (for prefill). This assumes
+        # the scheduler is synchronous so blocks are actually computed when
+        # scheduling the next batch.
+        self.block_allocator.mark_blocks_as_computed([])
+
+    def get_common_computed_block_ids(
+            self, seqs: List[Sequence]) -> GenericSequence[int]:
+        """Determine which blocks for which we skip prefill.
+
+        With prefix caching we can skip prefill for previously-generated blocks.
+        Currently, the attention implementation only supports skipping cached
+        blocks if they are a contiguous prefix of cached blocks.
+
+        This method determines which blocks can be safely skipped for all
+        sequences in the sequence group.
+        """
+        computed_seq_block_ids = []
+        for seq in seqs:
+            computed_seq_block_ids.append(
+                self._computed_blocks_tracker.
+                get_cached_computed_blocks_and_update(
+                    seq.seq_id,
+                    self.block_tables[seq.seq_id].physical_block_ids))
+
+        # NOTE(sang): This assumes seq_block_ids doesn't contain any None.
+        return self.block_allocator.get_common_computed_block_ids(
+            computed_seq_block_ids)  # type: ignore
+
+    def get_content_hashes(self, seq: Sequence) -> List[bytes]:
+        return self.block_tables[seq.seq_id].get_content_hashes()
+
+    def get_and_reset_prefix_swaps(
+            self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
+        """Return scheduler-owned (CPU->GPU, GPU->CPU) content transfers."""
+        return self.block_allocator.get_and_reset_prefix_swaps()
+
+    def begin_prefix_cache_step(self) -> None:
+        self.block_allocator.begin_prefix_cache_step()
+
+    def _build_runtime_cache_namespace(self) -> bytes:
+        """Bind first-block hashes to the fixed model runtime identity."""
+        model = os.getenv("BI100_PREFIX_MODEL_FINGERPRINT",
+                          "Qwen3.6-35B-A3B")
+        dtype = os.getenv("BI100_PREFIX_DTYPE", "float16")
+        tp_raw = os.getenv("BI100_PREFIX_TP_SIZE", "4")
+        try:
+            tp_size = int(tp_raw)
+        except ValueError as exc:
+            raise RuntimeError(
+                "BI100_PREFIX_TP_SIZE must be a positive integer") from exc
+        if tp_size <= 0:
+            raise RuntimeError(
+                "BI100_PREFIX_TP_SIZE must be a positive integer")
+
+        digest = hashlib.sha256()
+        digest.update(b"bi100-runtime-prefix-identity-v1|")
+        for label, value in (
+                (b"model", model),
+                (b"dtype", dtype),
+                (b"tp", str(tp_size)),
+                (b"block_size", str(self.block_size))):
+            encoded = value.encode("utf-8")
+            digest.update(label)
+            digest.update(struct.pack("!Q", len(encoded)))
+            digest.update(encoded)
+        return digest.digest()
+
+    @staticmethod
+    def _adapter_cache_namespace(seq_group: SequenceGroup) -> bytes:
+        digest = hashlib.sha256()
+        digest.update(b"bi100-adapter-prefix-identity-v1|")
+        lora = getattr(seq_group, "lora_request", None)
+        prompt_adapter = getattr(seq_group, "prompt_adapter_request", None)
+        identities = (
+            ("lora", lora, ("lora_name", "lora_int_id", "lora_path",
+                             "base_model_name")),
+            ("prompt", prompt_adapter,
+             ("prompt_adapter_name", "prompt_adapter_id",
+              "prompt_adapter_local_path",
+              "prompt_adapter_num_virtual_tokens")),
+        )
+        for kind, adapter, fields in identities:
+            digest.update(kind.encode("ascii"))
+            if adapter is None:
+                digest.update(b"none|")
+                continue
+            for field in fields:
+                value = str(getattr(adapter, field, ""))
+                encoded = value.encode("utf-8")
+                digest.update(field.encode("ascii"))
+                digest.update(struct.pack("!Q", len(encoded)))
+                digest.update(encoded)
+        return digest.digest()
+
+    def _request_local_fallback_cache_namespace(self,
+                                               request_id: str) -> bytes:
+        namespace = self._request_local_namespace.get(request_id)
+        if namespace is None:
+            digest = hashlib.sha256()
+            digest.update(b"multimodal-unsupported-request-local-v1|")
+            digest.update(self._runtime_cache_namespace)
+            digest.update(os.urandom(32))
+            digest.update(request_id.encode("utf-8"))
+            namespace = digest.digest()
+            self._request_local_namespace[request_id] = namespace
+        return namespace
+
+    def release_request_cache_namespace(self, request_id: str) -> None:
+        """Release request-local isolation state after request completion."""
+        self._request_local_namespace.pop(request_id, None)
+        self._warned_mm_namespace_requests.discard(request_id)
+
+    def _hash_multi_modal_namespace(self, mm_data: Any) -> bytes:
+        digest = hashlib.sha256()
+        self._hash_multi_modal_obj(digest, mm_data)
+        return digest.digest()
+
+    @staticmethod
+    def _sort_map_keys(mm_map: Mapping[Any, Any]) -> List[Any]:
+        return sorted(mm_map.keys(), key=lambda key: repr(key))
+
+    @classmethod
+    def _hash_multi_modal_obj(cls, digest: Any, value: Any) -> None:
+        if value is None:
+            digest.update(b"none|")
+            return
+        if isinstance(value, Mapping):
+            digest.update(b"map|")
+            digest.update(struct.pack("!Q", len(value)))
+            for key in cls._sort_map_keys(value):
+                digest.update(b"k|")
+                cls._hash_multi_modal_obj(digest, key)
+                digest.update(b"v|")
+                cls._hash_multi_modal_obj(digest, value[key])
+            return
+        if isinstance(value, list):
+            digest.update(b"list|")
+            digest.update(struct.pack("!Q", len(value)))
+            for item in value:
+                cls._hash_multi_modal_obj(digest, item)
+            return
+        if isinstance(value, tuple):
+            digest.update(b"tuple|")
+            digest.update(struct.pack("!Q", len(value)))
+            for item in value:
+                cls._hash_multi_modal_obj(digest, item)
+            return
+        if isinstance(value, str):
+            encoded = value.encode()
+            digest.update(b"str|")
+            digest.update(struct.pack("!Q", len(encoded)))
+            digest.update(encoded)
+            return
+        if isinstance(value, bytes):
+            digest.update(b"bytes|")
+            digest.update(struct.pack("!Q", len(value)))
+            digest.update(value)
+            return
+        if isinstance(value, bytearray):
+            cls._hash_multi_modal_obj(digest, bytes(value))
+            return
+        if isinstance(value, bool):
+            digest.update(b"bool|")
+            digest.update(b"1" if value else b"0")
+            return
+        if isinstance(value, int):
+            digest.update(b"int|")
+            digest.update(str(value).encode())
+            return
+        if isinstance(value, float):
+            digest.update(b"float|")
+            digest.update(struct.pack("!d", value))
+            return
+        if torch is not None and isinstance(value, torch.Tensor):
+            digest.update(b"tensor|")
+            tensor = value.detach().cpu().contiguous()
+            digest.update(struct.pack("!Q", len(tensor.shape)))
+            for dim in tensor.shape:
+                digest.update(struct.pack("!Q", int(dim)))
+            digest.update(str(tensor.dtype).encode())
+            # Byte views work for bfloat16 and other dtypes that NumPy cannot
+            # materialize directly.
+            tensor_bytes = tensor.view(torch.uint8).numpy().tobytes()
+            digest.update(struct.pack("!Q", len(tensor_bytes)))
+            digest.update(tensor_bytes)
+            return
+        if Image is not None and isinstance(value, Image.Image):
+            digest.update(b"image|")
+            digest.update(value.mode.encode())
+            digest.update(struct.pack("!II", value.width, value.height))
+            image_bytes = value.tobytes()
+            digest.update(struct.pack("!Q", len(image_bytes)))
+            digest.update(image_bytes)
+            palette = value.getpalette()
+            digest.update(b"palette-mode|")
+            cls._hash_multi_modal_obj(
+                digest, getattr(getattr(value, "palette", None), "mode", None))
+            digest.update(b"palette|")
+            cls._hash_multi_modal_obj(digest, palette)
+            digest.update(b"transparency|")
+            cls._hash_multi_modal_obj(
+                digest, value.info.get("transparency"))
+            return
+
+        raise TypeError(f"Unsupported multimodal namespace value type {type(value)}")
+
+
+    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
+        if parent_seq.seq_id not in self.block_tables:
+            # Parent sequence has either been freed or never existed.
+            return
+        src_block_table = self.block_tables[parent_seq.seq_id]
+        self.block_tables[child_seq.seq_id] = src_block_table.fork()
+
+        # Track child seq
+        self._computed_blocks_tracker.add_seq(child_seq.seq_id)
+        self._last_access_blocks_tracker.add_seq(child_seq.seq_id)
+
+    def can_swap_in(self, seq_group: SequenceGroup,
+                    num_lookahead_slots: int) -> AllocStatus:
+        """Returns the AllocStatus for the given sequence_group 
+        with num_lookahead_slots.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            AllocStatus: The AllocStatus for the given sequence group.
+        """
+        if self.block_allocator.content_offload_enabled:
+            return AllocStatus.NEVER
+        return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED,
+                              num_lookahead_slots)
+
+    def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
+        """Returns the block id mapping (from CPU to GPU) generated by
+        swapping in the given seq_group with num_lookahead_slots.
+
+        Args:
+            seq_group (SequenceGroup): The sequence group to swap in.
+
+        Returns:
+            List[Tuple[int, int]]: The mapping of swapping block from CPU 
+                to GPU.
+        """
+        physical_block_id_mapping = []
+        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
+            blocks = self.block_tables[seq.seq_id].blocks
+            if len(blocks) == 0:
+                continue
+
+            seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
+                                                         src_device=Device.CPU,
+                                                         dst_device=Device.GPU)
+
+            # Refresh the block ids of the table (post-swap)
+            self.block_tables[seq.seq_id].update(blocks)
+
+            seq_physical_block_id_mapping = {
+                self.block_allocator.get_physical_block_id(
+                    Device.CPU, cpu_block_id):
+                self.block_allocator.get_physical_block_id(
+                    Device.GPU, gpu_block_id)
+                for cpu_block_id, gpu_block_id in seq_swap_mapping.items()
+            }
+
+            physical_block_id_mapping.extend(
+                list(seq_physical_block_id_mapping.items()))
+
+        return physical_block_id_mapping
+
+    def can_swap_out(self, seq_group: SequenceGroup) -> bool:
+        """Returns whether we can swap out the given sequence_group 
+        with num_lookahead_slots.
+
+        Args:
+            seq_group (SequenceGroup): The sequence group to swap in.
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            bool: Whether it's possible to swap out current sequence group.
+        """
+        if self.block_allocator.content_offload_enabled:
+            return False
+        alloc_status = self._can_swap(seq_group, Device.CPU,
+                                      SequenceStatus.RUNNING)
+        return alloc_status == AllocStatus.OK
+
+    def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
+        """Returns the block id mapping (from GPU to CPU) generated by
+        swapping out the given sequence_group with num_lookahead_slots.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+
+        Returns:
+            List[Tuple[int, int]]: The mapping of swapping block from 
+                GPU to CPU.
+        """
+        physical_block_id_mapping = []
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            blocks = self.block_tables[seq.seq_id].blocks
+            if len(blocks) == 0:
+                continue
+
+            seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
+                                                         src_device=Device.GPU,
+                                                         dst_device=Device.CPU)
+
+            # Refresh the block ids of the table (post-swap)
+            self.block_tables[seq.seq_id].update(blocks)
+
+            seq_physical_block_id_mapping = {
+                self.block_allocator.get_physical_block_id(
+                    Device.GPU, gpu_block_id):
+                self.block_allocator.get_physical_block_id(
+                    Device.CPU, cpu_block_id)
+                for gpu_block_id, cpu_block_id in seq_swap_mapping.items()
+            }
+
+            physical_block_id_mapping.extend(
+                list(seq_physical_block_id_mapping.items()))
+
+        return physical_block_id_mapping
+
+    def get_num_free_gpu_blocks(self) -> int:
+        return self.block_allocator.get_num_free_blocks(Device.GPU)
+
+    def get_num_free_cpu_blocks(self) -> int:
+        return self.block_allocator.get_num_free_blocks(Device.CPU)
+
+    def get_prefix_cache_hit_rate(self, device: Device) -> float:
+        return self.block_allocator.get_prefix_cache_hit_rate(device)
+
+    def _can_swap(self,
+                  seq_group: SequenceGroup,
+                  device: Device,
+                  status: SequenceStatus,
+                  num_lookahead_slots: int = 0) -> AllocStatus:
+        """Returns the AllocStatus for swapping in/out the given sequence_group 
+        on to the 'device'.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+            device (Device): device to swap the 'seq_group' on.
+            status (SequenceStatus): The status of sequence which is needed
+                for action. RUNNING for swap out and SWAPPED for swap in
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            AllocStatus: The AllocStatus for swapping in/out the given 
+                sequence_group on to the 'device'.
+        """
+        # First determine the number of blocks that will be touched by this
+        # swap. Then verify if there are available blocks in the device
+        # to perform the swap.
+        num_blocks_touched = 0
+        blocks: List[Block] = []
+        for seq in seq_group.get_seqs(status=status):
+            block_table = self.block_tables[seq.seq_id]
+            if block_table.blocks is not None:
+                # Compute the number blocks to touch for the tokens to be
+                # appended. This does NOT include the full blocks that need
+                # to be touched for the swap.
+                num_blocks_touched += \
+                    block_table.get_num_blocks_touched_by_append_slots(
+                        block_table.get_unseen_token_ids(seq.get_token_ids()),
+                        num_lookahead_slots=num_lookahead_slots)
+                blocks.extend(block_table.blocks)
+        # Compute the number of full blocks to touch and add it to the
+        # existing count of blocks to touch.
+        num_blocks_touched += self.block_allocator.get_num_full_blocks_touched(
+            blocks, device=device)
+
+        watermark_blocks = 0
+        if device == Device.GPU:
+            watermark_blocks = self.watermark_blocks
+
+        if self.block_allocator.get_num_total_blocks(
+                device) < num_blocks_touched:
+            return AllocStatus.NEVER
+        elif self.block_allocator.get_num_free_blocks(
+                device) - num_blocks_touched >= watermark_blocks:
+            return AllocStatus.OK
+        else:
+            return AllocStatus.LATER

qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py

@@ -0,0 +1,272 @@
+import enum
+import heapq
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Dict, List, OrderedDict, Tuple
+
+
+ContentHash = bytes
+
+
+class EvictionPolicy(enum.Enum):
+    """Enum for eviction policy used by make_evictor to instantiate the correct
+       Evictor subclass.
+    """
+    LRU = enum.auto()
+    FREQUENCY_AWARE = enum.auto()
+
+
+class Evictor(ABC):
+    """The Evictor subclasses should be used by the BlockAllocator class to
+    handle eviction of freed PhysicalTokenBlocks.
+    """
+
+    @abstractmethod
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def __contains__(self, block_id: int) -> bool:
+        pass
+
+    @abstractmethod
+    def evict(self) -> Tuple[int, ContentHash]:
+        """Runs the eviction algorithm and returns the evicted block's
+        content hash along with physical block id along with physical block id
+        """
+        pass
+
+    @abstractmethod
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int,
+            last_accessed: float):
+        """Adds block to the evictor, making it a candidate for eviction"""
+        pass
+
+    @abstractmethod
+    def update(self, block_id: int, last_accessed: float):
+        """Update corresponding block's access time in metadata"""
+        pass
+
+    @abstractmethod
+    def remove(self, block_id: int):
+        """Remove a given block id from the cache."""
+        pass
+
+    @property
+    @abstractmethod
+    def num_blocks(self) -> int:
+        pass
+
+
+class BlockMetaData():
+    """Data structure for storing key data describe cached block, so that
+    evitor could use to make its decision which one to choose for eviction
+
+    Here we use physical block id as the dict key, as there maybe several
+    blocks with the same content hash, but their physical id is unique.
+    """
+
+    def __init__(self, content_hash: ContentHash, num_hashed_tokens: int,
+                 last_accessed: float):
+        self.content_hash = content_hash
+        self.num_hashed_tokens = num_hashed_tokens
+        self.last_accessed = last_accessed
+
+
+class LRUEvictor(Evictor):
+    """Evicts in a least-recently-used order using the last_accessed timestamp
+    that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
+    the same last_accessed time, then the one with the largest num_hashed_tokens
+    will be evicted. If two blocks each have the lowest last_accessed time and
+    highest num_hashed_tokens value, then one will be chose arbitrarily
+    """
+
+    def __init__(self):
+        self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()
+
+    def __contains__(self, block_id: int) -> bool:
+        return block_id in self.free_table
+
+    def evict(self) -> Tuple[int, ContentHash]:
+        if len(self.free_table) == 0:
+            raise ValueError("No usable cache memory left")
+
+        evicted_block, evicted_block_id = None, None
+        # The blocks with the lowest timestamps should be placed consecutively
+        # at the start of OrderedDict. Loop through all these blocks to
+        # find the one with maximum number of hashed tokens.
+        for _id, block in self.free_table.items():
+            if evicted_block is None:
+                evicted_block, evicted_block_id = block, _id
+                continue
+            if evicted_block.last_accessed < block.last_accessed:
+                break
+            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
+                evicted_block, evicted_block_id = block, _id
+
+        assert evicted_block is not None
+        assert evicted_block_id is not None
+        self.free_table.pop(evicted_block_id)
+
+        return evicted_block_id, evicted_block.content_hash
+
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int,
+            last_accessed: float):
+        self.free_table[block_id] = BlockMetaData(content_hash,
+                                                  num_hashed_tokens,
+                                                  last_accessed)
+
+    def update(self, block_id: int, last_accessed: float):
+        self.free_table[block_id].last_accessed = last_accessed
+
+    def remove(self, block_id: int):
+        if block_id not in self.free_table:
+            raise ValueError(
+                "Attempting to remove block that's not in the evictor")
+        self.free_table.pop(block_id)
+
+    @property
+    def num_blocks(self) -> int:
+        return len(self.free_table)
+
+
+class FrequencyAwareEvictor(Evictor):
+    """Evict the least frequently reused logical prefix content first.
+
+    Content frequency survives physical block reuse. Heap entries carry a
+    generation and are lazily invalidated so eviction remains O(log N) without
+    allowing stale entries to grow without bound.
+    """
+
+    _COMPACTION_FACTOR = 2
+    _COMPACTION_SLACK = 1
+
+    def __init__(self):
+        self.free_table: Dict[int, BlockMetaData] = {}
+        self.frequency_by_hash: Dict[ContentHash, int] = {}
+        self._heap: List[Tuple[int, float, int, int, int]] = []
+        self._generations: Dict[int, int] = {}
+        self._next_generation = 0
+
+    @staticmethod
+    def _validate_content_hash(content_hash: ContentHash) -> None:
+        if not isinstance(content_hash, bytes) or len(content_hash) != 32:
+            raise ValueError(
+                "frequency-aware eviction requires a 32-byte content hash")
+
+    def __contains__(self, block_id: int) -> bool:
+        return block_id in self.free_table
+
+    def _heap_key(self, block_id: int, block: BlockMetaData,
+                  generation: int) -> Tuple[int, float, int, int, int]:
+        return (
+            self.frequency_by_hash[block.content_hash],
+            block.last_accessed,
+            -block.num_hashed_tokens,
+            block_id,
+            generation,
+        )
+
+    def _push(self, block_id: int) -> None:
+        self._next_generation += 1
+        generation = self._next_generation
+        self._generations[block_id] = generation
+        heapq.heappush(
+            self._heap,
+            self._heap_key(
+                block_id, self.free_table[block_id], generation),
+        )
+
+    def _compact_if_needed(self) -> None:
+        limit = (
+            self._COMPACTION_FACTOR * len(self.free_table)
+            + self._COMPACTION_SLACK
+        )
+        if len(self._heap) <= limit:
+            return
+        self._heap = [
+            self._heap_key(block_id, block, self._generations[block_id])
+            for block_id, block in self.free_table.items()
+        ]
+        heapq.heapify(self._heap)
+
+    def evict(self) -> Tuple[int, ContentHash]:
+        if not self.free_table:
+            raise ValueError("No usable cache memory left")
+
+        while self._heap:
+            entry = heapq.heappop(self._heap)
+            frequency, _, _, block_id, generation = entry
+            block = self.free_table.get(block_id)
+            if (
+                block is None
+                or self._generations.get(block_id) != generation
+            ):
+                continue
+            if self.frequency_by_hash[block.content_hash] != frequency:
+                heapq.heappush(
+                    self._heap,
+                    self._heap_key(block_id, block, generation),
+                )
+                continue
+
+            block = self.free_table.pop(block_id)
+            self._generations.pop(block_id)
+            self._compact_if_needed()
+            return block_id, block.content_hash
+
+        raise RuntimeError("Evictor heap has no usable entry")
+
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int, last_accessed: float):
+        self._validate_content_hash(content_hash)
+        self.frequency_by_hash[content_hash] = (
+            self.frequency_by_hash.get(content_hash, 0) + 1)
+        self.free_table[block_id] = BlockMetaData(
+            content_hash, num_hashed_tokens, last_accessed)
+        self._push(block_id)
+        self._compact_if_needed()
+
+    def update(self, block_id: int, last_accessed: float):
+        self.free_table[block_id].last_accessed = last_accessed
+        self._push(block_id)
+        self._compact_if_needed()
+
+    def remove(self, block_id: int):
+        if block_id not in self.free_table:
+            raise ValueError(
+                "Attempting to remove block that's not in the evictor")
+        self.free_table.pop(block_id)
+        self._generations.pop(block_id)
+        self._compact_if_needed()
+
+    @property
+    def num_blocks(self) -> int:
+        return len(self.free_table)
+
+
+def eviction_policy_from_env(
+    environ: Mapping[str, str] | None = None,
+) -> EvictionPolicy:
+    source = os.environ if environ is None else environ
+    value = source.get("BI100_KV_EVICTION_POLICY", "lru").strip().lower()
+    policies = {
+        "lru": EvictionPolicy.LRU,
+        "frequency": EvictionPolicy.FREQUENCY_AWARE,
+    }
+    if value not in policies:
+        raise ValueError(
+            "BI100_KV_EVICTION_POLICY must be one of: frequency, lru")
+    return policies[value]
+
+
+def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
+    if eviction_policy == EvictionPolicy.LRU:
+        return LRUEvictor()
+    elif eviction_policy == EvictionPolicy.FREQUENCY_AWARE:
+        return FrequencyAwareEvictor()
+    else:
+        raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")

qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py

@@ -0,0 +1,644 @@
+from array import array
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from vllm.sampling_params import SamplingParams, SamplingType
+from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData,
+                           SequenceGroupMetadata)
+from vllm.utils import (PyObjectCache, async_tensor_h2d,
+                        is_pin_memory_available, make_tensor_with_pad)
+
+_SAMPLING_EPS = 1e-5
+
+
+@dataclass
+class SequenceGroupToSample:
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ----------------------|
+    #                                   |-- query_len ---|
+
+    # Sequence ids for the sequence group in a previous step.
+    seq_ids: List[int]
+    sampling_params: SamplingParams
+    # seq_id -> sequence data.
+    seq_data: Dict[int, SequenceData]
+    # The length of the sequence (all tokens seen in the past + new token to
+    # compute attention) of the sequence group. None if it is in a decode
+    # stage.
+    seq_len: Optional[int]
+    # The length of new query tokens to compute in the current step. None if it
+    # is in a decode stage. The length of query_len <= seq_len if chunked
+    # prefill is enabled.
+    query_len: Optional[int]
+    # A random number generator for sampling.
+    generator: Optional[torch.Generator]
+    # True if the sequence group is in prefill stage. False if it is in a
+    # decode stage.
+    is_prompt: bool
+    # Query token indices from logits. to compute prompt logprob. Empty if
+    # prompt logprob is not required.
+    prompt_logprob_indices: List[int]
+    # Output offsets within this prefill chunk. Sparse diagnostic requests use
+    # this to retain the standard full-length prompt-logprob response shape.
+    prompt_logprob_output_indices: List[int]
+    # Sample token indices from logits. Empty if sampling is not required.
+    sample_indices: List[int]
+
+    @property
+    def do_sample(self):
+        return len(self.sample_indices) > 0
+
+    def __post_init__(self):
+        if len(self.prompt_logprob_indices) > 0:
+            assert self.sampling_params.prompt_logprobs is not None
+        assert (len(self.prompt_logprob_indices)
+                == len(self.prompt_logprob_output_indices))
+        assert self.prompt_logprob_output_indices == sorted(
+            set(self.prompt_logprob_output_indices))
+        if self.is_prompt:
+            assert self.seq_len is not None
+            assert self.query_len is not None
+            assert all(
+                0 <= index < self.query_len
+                for index in self.prompt_logprob_output_indices)
+
+
+def gen_seq_group_to_sample_builder(num_seqs: int):
+    return lambda: SequenceGroupToSample(
+        seq_ids=[0] * num_seqs,
+        sampling_params=None,
+        seq_data=None,  # type: ignore
+        seq_len=0,
+        query_len=0,
+        generator=None,
+        is_prompt=True,
+        prompt_logprob_indices=[],
+        prompt_logprob_output_indices=[],
+        sample_indices=[],
+    )
+
+
+class SamplingMetadataCache:
+    """Used to cache SamplingMetadata objects between scheduler iterations"""
+
+    def __init__(self):
+        self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {}
+
+    def get_cached_seq_group_to_sample(self, num_seqs):
+        if num_seqs not in self._seq_group_to_sample_cache:
+            self._seq_group_to_sample_cache[num_seqs] = PyObjectCache(
+                gen_seq_group_to_sample_builder(num_seqs))
+
+        obj = self._seq_group_to_sample_cache[num_seqs].get_object()
+        return obj
+
+    def reset(self):
+        for cache in self._seq_group_to_sample_cache.values():
+            cache.reset()
+
+
+class SamplingMetadata:
+    """Metadata for input sequences. Used in sampler.
+
+    The usage is as follow;
+    ```
+    hidden_states = execute_model(...)
+    logits = hidden_states[sampling_metadata.selected_token_indices]
+    sample(logits)
+
+    def sample(logits):
+        # Use categorized_sample_indices for sampling....
+    ```
+
+    Args:
+        seq_groups: List of batched sequence groups.
+        selected_token_indices: (num_query_tokens_to_logprob). Indices to find
+            logits from the initial model output hidden states.
+        categorized_sample_indices: SamplingType -> token indices to sample.
+            Each token indices is 2D tensor of (num_indices, num_indices) where
+            the first item means the sample index within the returned logit
+            (before pruning padding), and the second item means the sample
+            index after pruning using selected_token_indices.
+            For example, if the returned logit is [1, 2, 3], and we select
+            [1, 2] for sampling, the pruned logit will be [2, 3]. In this case,
+            The first tuple is [1, 2] (sampled index within original logit),
+            and the second tuple is [0, 1] (sampled index within pruned logit).
+        num_prompts: Number of prompt sequence groups in seq_groups.
+        skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU
+            serialization of token outputs.
+        reuse_sampling_tensors: Indicates if we want to reuse sampling
+            tensors that are part of the sampler forward pass. Currently,
+            it is mainly used for multi-step decode.
+
+    """
+
+    def __init__(
+        self,
+        seq_groups: List[SequenceGroupToSample],
+        selected_token_indices: torch.Tensor,
+        categorized_sample_indices: Dict[SamplingType, torch.Tensor],
+        num_prompts: int,
+        skip_sampler_cpu_output: bool = False,
+        reuse_sampling_tensors: bool = False,
+    ) -> None:
+        self.seq_groups = seq_groups
+        self.selected_token_indices = selected_token_indices
+        self.categorized_sample_indices = categorized_sample_indices
+        self.num_prompts = num_prompts
+        self.skip_sampler_cpu_output = skip_sampler_cpu_output
+        self.reuse_sampling_tensors = reuse_sampling_tensors
+
+    @staticmethod
+    def prepare(
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+        seq_lens: List[int],
+        query_lens: List[int],
+        device: str,
+        pin_memory: bool,
+        generators: Optional[Dict[str, torch.Generator]] = None,
+        cache: Optional[SamplingMetadataCache] = None,
+    ) -> "SamplingMetadata":
+        (
+            seq_groups,
+            selected_token_indices,
+            categorized_sample_indices,
+            num_prompts,
+        ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,
+                                device, generators, cache)
+        selected_token_indices = async_tensor_h2d(
+            selected_token_indices,
+            dtype=torch.long,
+            target_device=device,
+            pin_memory=pin_memory,
+        )
+        categorized_sample_indices = {
+            t: async_tensor_h2d(
+                seq_ids,
+                dtype=torch.int,
+                target_device=device,
+                pin_memory=pin_memory,
+            )
+            for t, seq_ids in categorized_sample_indices.items()
+        }
+
+        sampling_metadata = SamplingMetadata(
+            seq_groups=seq_groups,
+            selected_token_indices=selected_token_indices,
+            categorized_sample_indices=categorized_sample_indices,
+            num_prompts=num_prompts,
+        )
+        return sampling_metadata
+
+    def __repr__(self) -> str:
+        return (
+            "SamplingMetadata("
+            f"seq_groups={self.seq_groups}, "
+            f"selected_token_indices={self.selected_token_indices}, "
+            f"categorized_sample_indices={self.categorized_sample_indices}), ")
+
+
+def _get_prompt_logprob_output_indices(
+    sampling_params: SamplingParams,
+    seq_data: SequenceData,
+    prompt_logprob_len: int,
+) -> List[int]:
+    if sampling_params.prompt_logprobs is None or prompt_logprob_len <= 0:
+        return []
+    positions = sampling_params.prompt_logprob_positions
+    computed_len = seq_data.get_num_computed_tokens()
+    available_next_tokens = max(
+        0,
+        len(seq_data.prompt_token_ids) - computed_len - 1,
+    )
+    materialized_len = min(prompt_logprob_len, available_next_tokens)
+    if positions is None:
+        return list(range(materialized_len))
+
+    output_indices = [
+        position - computed_len - 1
+        for position in positions
+        if computed_len < position
+        <= computed_len + materialized_len
+    ]
+    assert output_indices == sorted(set(output_indices))
+    assert all(0 <= index < materialized_len for index in output_indices)
+    return output_indices
+
+
+def _prepare_seq_groups(
+    seq_group_metadata_list: List[SequenceGroupMetadata],
+    seq_lens: List[int],
+    query_lens: List[int],
+    device: str,
+    generators: Optional[Dict[str, torch.Generator]] = None,
+    cache: Optional[SamplingMetadataCache] = None,
+) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType,
+                                                        List[int]], int, ]:
+    """Prepare sequence groups and indices for sampling.
+
+    Args:
+        seq_group_metadata_list: A list of sequence group to batch.
+        seq_lens: A list of sequence lens per sequence group.
+            Index of prompt len should match with seq_group_metadata_list.
+        query_lens: A list of query lengths. Prompt lens include the length
+            of entire prompt tokens, and it could be shorter.
+        device: A device to use for random number generators,
+            `SequenceGroupToSample.generator`.
+        generators: A store of per-request random number generators used
+            for seeded requests.
+
+    Returns:
+        seq_groups: A list of sequence group to sample.
+        selected_token_indices: See the definition from `SamplingMetadata`.
+        categorized_sample_indices: See the definition from `SamplingMetadata`.
+        num_prompts: Total number of prompts from `seq_group_metadata_list`.
+    """
+    # Batched sequence groups for the current model forward stsep.
+    seq_groups: List[SequenceGroupToSample] = []
+    # A list of token indices to sample/compute logprob. It is used to
+    # prune the outcome logits from the model for the performance.
+    selected_token_indices: List[int] = []
+    # Used for selected_token_indices.
+    model_output_idx = 0
+
+    # Sampling type -> (
+    # indices to sample/prompt logprob within pruned output logits,
+    # indices to sample within pruned logits)
+    categorized_sample_indices: Dict[SamplingType, List[int]] = {
+        t: []
+        for t in SamplingType
+    }
+    # Index of logits to compute logprob. Logits include both prompt logprob
+    # and sample logprob indices.
+    logit_idx = 0
+    # Total number of prompts from given sequence groups.
+    num_prompts = 0
+
+    for i, seq_group_metadata in enumerate(seq_group_metadata_list):
+        seq_ids = seq_group_metadata.seq_data.keys()
+
+        if cache is not None:
+            sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids))
+
+            for j, seq_id in enumerate(seq_ids):
+                sample_obj.seq_ids[j] = seq_id
+
+            sample_obj.prompt_logprob_indices.clear()
+            sample_obj.prompt_logprob_output_indices.clear()
+            sample_obj.sample_indices.clear()
+
+        sampling_params = seq_group_metadata.sampling_params
+        is_prompt = seq_group_metadata.is_prompt
+        generator: Optional[torch.Generator] = None
+        # If the current seq group is in decode stage, it is None.
+        seq_len: Optional[int] = None
+        query_len: Optional[int] = None
+        prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices
+                                             if cache is not None else [])
+        prompt_logprob_output_indices: List[int] = (
+            sample_obj.prompt_logprob_output_indices
+            if cache is not None else [])
+        sample_indices: List[int] = (sample_obj.sample_indices
+                                     if cache is not None else [])
+        do_sample = seq_group_metadata.do_sample
+
+        if seq_group_metadata.is_prompt:
+            if sampling_params.seed is not None:
+                generator = torch.Generator(device=device).manual_seed(
+                    sampling_params.seed)
+                if generators is not None:
+                    generators[seq_group_metadata.request_id] = generator
+
+            num_prompts += 1
+            num_prefill_sample = len(seq_ids)
+            assert num_prefill_sample == 1
+            assert query_lens is not None and seq_lens is not None
+            query_len, seq_len = query_lens[i], seq_lens[i]
+            # If we need sampling, exclude num_prefill_sample tokens from
+            # prompt logprob.
+            prompt_logprob_len = (query_len - num_prefill_sample
+                                  if do_sample else query_len)
+            sample_len = num_prefill_sample if do_sample else 0
+        else:
+            # Decode
+            prompt_logprob_len = 0
+            query_len = query_lens[i] if query_lens is not None else 1
+            sample_len = len(seq_ids) * query_len if do_sample else 0
+
+            if sampling_params.seed is not None and generators is not None:
+                generator = generators.get(seq_group_metadata.request_id)
+
+        seq_data = next(iter(seq_group_metadata.seq_data.values()))
+        prompt_logprob_output_indices.extend(
+            _get_prompt_logprob_output_indices(
+                sampling_params,
+                seq_data,
+                prompt_logprob_len,
+            ))
+
+        # Update indices to select from the model output.
+        """
+        This blocks computes selected_token_indices which is used in the
+        following way.
+
+        hidden_states = model(...)
+        logits = hidden_states[selected_token_indices]
+        """
+
+        if sampling_params.prompt_logprobs is not None:
+            selected_token_indices.extend(
+                model_output_idx + output_index
+                for output_index in prompt_logprob_output_indices)
+        model_output_idx += prompt_logprob_len
+        if do_sample:
+            selected_token_indices.extend(
+                range(model_output_idx, model_output_idx + sample_len))
+        model_output_idx += sample_len
+
+        # We now find indices for logprob computation and sampling.
+        """
+        This block computes categorized_sample_indices which is used in the
+        following way.
+
+        hidden_states = model(...)
+        logits = hidden_states[selected_token_indices]
+        def sample(logits):
+           # Use categorized_sample_indices for sampling.
+           # prompt_logprob_indices to find prompt logprob indices.
+           # sample_indices to find sample indices.
+        """
+
+        if sampling_params.prompt_logprobs is not None:
+            prompt_logprob_indices.extend(
+                range(logit_idx,
+                      logit_idx + len(prompt_logprob_output_indices)))
+            logit_idx += len(prompt_logprob_output_indices)
+        if do_sample:
+            sample_indices.extend(range(logit_idx, logit_idx + sample_len))
+            categorized_sample_indices[sampling_params.sampling_type].extend(
+                list(range(logit_idx, logit_idx + sample_len)))
+            logit_idx += sample_len
+
+        if cache is not None:
+            sample_obj.sampling_params = sampling_params
+            sample_obj.seq_data = seq_group_metadata.seq_data
+            sample_obj.seq_len = seq_len
+            sample_obj.query_len = query_len
+            sample_obj.generator = generator
+            sample_obj.is_prompt = is_prompt
+        else:
+            sample_obj = SequenceGroupToSample(
+                seq_ids=list(seq_ids),
+                sampling_params=sampling_params,
+                seq_data=seq_group_metadata.seq_data,
+                seq_len=seq_len,
+                query_len=query_len,
+                generator=generator,
+                is_prompt=is_prompt,
+                prompt_logprob_indices=list(prompt_logprob_indices),
+                prompt_logprob_output_indices=list(
+                    prompt_logprob_output_indices),
+                sample_indices=list(sample_indices),
+            )
+
+        assert (len(sample_obj.prompt_logprob_indices)
+                == len(sample_obj.prompt_logprob_output_indices))
+        seq_groups.append(sample_obj)
+
+    if cache is not None:
+        cache.reset()
+
+    return (seq_groups, selected_token_indices, categorized_sample_indices,
+            num_prompts)
+
+
+@dataclass
+class SamplingTensors:
+    """Tensors for sampling."""
+
+    temperatures: torch.Tensor
+    top_ps: torch.Tensor
+    top_ks: torch.Tensor
+    min_ps: torch.Tensor
+    presence_penalties: torch.Tensor
+    frequency_penalties: torch.Tensor
+    repetition_penalties: torch.Tensor
+    prompt_tokens: torch.Tensor
+    output_tokens: torch.Tensor
+
+    @classmethod
+    def from_sampling_metadata(
+        cls,
+        sampling_metadata: "SamplingMetadata",
+        vocab_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tuple["SamplingTensors", bool, bool, bool]:
+        prompt_tokens: List[array] = []
+        output_tokens: List[array] = []
+        top_ks: List[int] = []
+        temperatures: List[float] = []
+        top_ps: List[float] = []
+        min_ps: List[float] = []
+        presence_penalties: List[float] = []
+        frequency_penalties: List[float] = []
+        repetition_penalties: List[float] = []
+        do_penalties = False
+        do_top_p_top_k = False
+        do_min_p = False
+
+        assert sampling_metadata.seq_groups is not None
+        for seq_group in sampling_metadata.seq_groups:
+            seq_ids = seq_group.seq_ids
+            sampling_params = seq_group.sampling_params
+            temperature = sampling_params.temperature
+            p = sampling_params.presence_penalty
+            f = sampling_params.frequency_penalty
+            r = sampling_params.repetition_penalty
+            top_p = sampling_params.top_p
+            min_p = sampling_params.min_p
+
+            # k should not be greater than the vocab size.
+            top_k = min(sampling_params.top_k, vocab_size)
+            top_k = vocab_size if top_k == -1 else top_k
+            if temperature < _SAMPLING_EPS:
+                # NOTE: Zero temperature means deterministic sampling
+                # (i.e., greedy sampling or beam search).
+                # Set the temperature to 1 to avoid division by zero.
+                temperature = 1.0
+            if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
+                                       or top_k != vocab_size):
+                do_top_p_top_k = True
+            if not do_min_p and min_p > _SAMPLING_EPS:
+                do_min_p = True
+            if not do_penalties and (abs(p) >= _SAMPLING_EPS
+                                     or abs(f) >= _SAMPLING_EPS
+                                     or abs(r - 1.0) >= _SAMPLING_EPS):
+                do_penalties = True
+
+            is_prompt = seq_group.is_prompt
+            if is_prompt and sampling_params.prompt_logprobs is not None:
+                # For tokens in the prompt that we only need to get
+                # their logprobs
+                query_len = seq_group.query_len
+                assert query_len is not None
+                prefill_len = len(seq_group.prompt_logprob_indices)
+                temperatures += [temperature] * prefill_len
+                top_ps += [top_p] * prefill_len
+                top_ks += [top_k] * prefill_len
+                min_ps += [min_p] * prefill_len
+                presence_penalties += [0] * prefill_len
+                frequency_penalties += [0] * prefill_len
+                repetition_penalties += [1] * prefill_len
+
+            if seq_group.do_sample:
+                sample_lens = len(seq_group.sample_indices)
+                assert sample_lens >= len(seq_ids)
+                temperatures += [temperature] * sample_lens
+                top_ps += [top_p] * sample_lens
+                top_ks += [top_k] * sample_lens
+                min_ps += [min_p] * sample_lens
+                presence_penalties += [p] * sample_lens
+                frequency_penalties += [f] * sample_lens
+                repetition_penalties += [r] * sample_lens
+
+        if do_penalties:
+            for seq_group in sampling_metadata.seq_groups:
+                seq_ids = seq_group.seq_ids
+                if (seq_group.is_prompt
+                        and sampling_params.prompt_logprobs is not None):
+                    prefill_len = len(seq_group.prompt_logprob_indices)
+                    prompt_tokens.extend(
+                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
+                        for _ in range(prefill_len))
+                    output_tokens.extend(
+                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
+                        for _ in range(prefill_len))
+                if seq_group.do_sample:
+                    for seq_id in seq_ids:
+                        seq_data = seq_group.seq_data[seq_id]
+                        prompt_tokens.append(seq_data.prompt_token_ids_array)
+                        output_tokens.append(seq_data.output_token_ids_array)
+
+        sampling_tensors = SamplingTensors.from_lists(
+            temperatures,
+            top_ps,
+            top_ks,
+            min_ps,
+            presence_penalties,
+            frequency_penalties,
+            repetition_penalties,
+            prompt_tokens,
+            output_tokens,
+            vocab_size,
+            device,
+            dtype,
+        )
+        return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)
+
+    @classmethod
+    def from_lists(
+        cls,
+        temperatures: List[float],
+        top_ps: List[float],
+        top_ks: List[int],
+        min_ps: List[float],
+        presence_penalties: List[float],
+        frequency_penalties: List[float],
+        repetition_penalties: List[float],
+        prompt_tokens: List[array],
+        output_tokens: List[array],
+        vocab_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> "SamplingTensors":
+        # Note that the performance will be very bad without
+        # pinned memory.
+        pin_memory = is_pin_memory_available()
+
+        do_penalties = prompt_tokens or output_tokens
+
+        if do_penalties:
+            prompt_t = make_tensor_with_pad(
+                prompt_tokens,
+                vocab_size,
+                device="cpu",
+                dtype=torch.int64,
+                pin_memory=pin_memory,
+            )
+            output_t = make_tensor_with_pad(
+                output_tokens,
+                vocab_size,
+                device="cpu",
+                dtype=torch.int64,
+                pin_memory=pin_memory,
+            )
+        else:
+            empty_tensor = torch.empty(0, device=device, dtype=torch.long)
+            prompt_t = empty_tensor
+            output_t = empty_tensor
+
+        temperatures_t = torch.tensor(
+            temperatures,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        top_ps_t = torch.tensor(
+            top_ps,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        min_ps_t = torch.tensor(
+            min_ps,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        presence_penalties_t = torch.tensor(
+            presence_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        frequency_penalties_t = torch.tensor(
+            frequency_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        repetition_penalties_t = torch.tensor(
+            repetition_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        top_ks_t = torch.tensor(
+            top_ks,
+            device="cpu",
+            dtype=torch.int,
+            pin_memory=pin_memory,
+        )
+        # Because the memory is pinned, we can do non-blocking
+        # transfer to device.
+
+        return cls(
+            temperatures=temperatures_t.to(device=device, non_blocking=True),
+            top_ps=top_ps_t.to(device=device, non_blocking=True),
+            top_ks=top_ks_t.to(device=device, non_blocking=True),
+            min_ps=min_ps_t.to(device=device, non_blocking=True),
+            presence_penalties=presence_penalties_t.to(device=device,
+                                                       non_blocking=True),
+            frequency_penalties=frequency_penalties_t.to(device=device,
+                                                         non_blocking=True),
+            repetition_penalties=repetition_penalties_t.to(device=device,
+                                                           non_blocking=True),
+            prompt_tokens=prompt_t.to(device=device, non_blocking=True),
+            output_tokens=output_t.to(device=device, non_blocking=True),
+        )

qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py

@@ -0,0 +1,520 @@
+"""Sampling parameters for text generation."""
+import copy
+from dataclasses import dataclass
+from enum import Enum, IntEnum
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import msgspec
+import torch
+from pydantic import BaseModel
+from typing_extensions import Annotated
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_SAMPLING_EPS = 1e-5
+_MAX_TEMP = 1e-2
+
+
+class SamplingType(IntEnum):
+    GREEDY = 0
+    RANDOM = 1
+    RANDOM_SEED = 2
+
+
+LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor],
+                        Callable[[List[int], List[int], torch.Tensor],
+                                 torch.Tensor]]
+"""LogitsProcessor is a function that takes a list
+of previously generated tokens, the logits tensor
+for the next token and, optionally, prompt tokens as a
+first argument, and returns a modified tensor of logits
+to sample from."""
+
+
+# maybe make msgspec?
+@dataclass
+class GuidedDecodingParams:
+    """One of these fields will be used to build a logit processor."""
+    json: Optional[Union[str, Dict]] = None
+    regex: Optional[str] = None
+    choice: Optional[List[str]] = None
+    grammar: Optional[str] = None
+    json_object: Optional[bool] = None
+    """These are other options that can be set"""
+    backend: Optional[str] = None
+    whitespace_pattern: Optional[str] = None
+
+    @staticmethod
+    def from_optional(
+        json: Optional[Union[Dict, BaseModel, str]],
+        regex: Optional[str] = None,
+        choice: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        json_object: Optional[bool] = None,
+        backend: Optional[str] = None,
+        whitespace_pattern: Optional[str] = None,
+    ) -> "GuidedDecodingParams":
+        # Extract json schemas from pydantic models
+        if isinstance(json, (BaseModel, type(BaseModel))):
+            json = json.model_json_schema()
+        return GuidedDecodingParams(
+            json=json,
+            regex=regex,
+            choice=choice,
+            grammar=grammar,
+            json_object=json_object,
+            backend=backend,
+            whitespace_pattern=whitespace_pattern,
+        )
+
+    def __post_init__(self):
+        """Validate that some fields are mutually exclusive."""
+        guide_count = sum([
+            self.json is not None, self.regex is not None, self.choice
+            is not None, self.grammar is not None, self.json_object is not None
+        ])
+        if guide_count > 1:
+            raise ValueError(
+                "You can only use one kind of guided decoding but multiple are "
+                f"specified: {self.__dict__}")
+
+
+class RequestOutputKind(Enum):
+    # Return entire output so far in every RequestOutput
+    CUMULATIVE = 0
+    # Return only deltas in each RequestOutput
+    DELTA = 1
+    # Do not return intermediate RequestOuputs
+    FINAL_ONLY = 2
+
+
+class SamplingParams(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        # required for @cached_property.
+        dict=True):  # type: ignore[call-arg]
+    """Sampling parameters for text generation.
+
+    Overall, we follow the sampling parameters from the OpenAI text completion
+    API (https://platform.openai.com/docs/api-reference/completions/create).
+    In addition, we support beam search, which is not supported by OpenAI.
+
+    Args:
+        n: Number of output sequences to return for the given prompt.
+        best_of: Number of output sequences that are generated from the prompt.
+            From these `best_of` sequences, the top `n` sequences are returned.
+            `best_of` must be greater than or equal to `n`. By default,
+            `best_of` is set to `n`.
+        presence_penalty: Float that penalizes new tokens based on whether they
+            appear in the generated text so far. Values > 0 encourage the model
+            to use new tokens, while values < 0 encourage the model to repeat
+            tokens.
+        frequency_penalty: Float that penalizes new tokens based on their
+            frequency in the generated text so far. Values > 0 encourage the
+            model to use new tokens, while values < 0 encourage the model to
+            repeat tokens.
+        repetition_penalty: Float that penalizes new tokens based on whether
+            they appear in the prompt and the generated text so far. Values > 1
+            encourage the model to use new tokens, while values < 1 encourage
+            the model to repeat tokens.
+        temperature: Float that controls the randomness of the sampling. Lower
+            values make the model more deterministic, while higher values make
+            the model more random. Zero means greedy sampling.
+        top_p: Float that controls the cumulative probability of the top tokens
+            to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
+        top_k: Integer that controls the number of top tokens to consider. Set
+            to -1 to consider all tokens.
+        min_p: Float that represents the minimum probability for a token to be
+            considered, relative to the probability of the most likely token.
+            Must be in [0, 1]. Set to 0 to disable this.
+        seed: Random seed to use for the generation.
+        stop: List of strings that stop the generation when they are generated.
+            The returned output will not contain the stop strings.
+        stop_token_ids: List of tokens that stop the generation when they are
+            generated. The returned output will contain the stop tokens unless
+            the stop tokens are special tokens.
+        include_stop_str_in_output: Whether to include the stop strings in
+            output text. Defaults to False.
+        ignore_eos: Whether to ignore the EOS token and continue generating
+            tokens after the EOS token is generated.
+        max_tokens: Maximum number of tokens to generate per output sequence.
+        min_tokens: Minimum number of tokens to generate per output sequence
+            before EOS or stop_token_ids can be generated
+        logprobs: Number of log probabilities to return per output token.
+            When set to None, no probability is returned. If set to a non-None
+            value, the result includes the log probabilities of the specified
+            number of most likely tokens, as well as the chosen tokens.
+            Note that the implementation follows the OpenAI API: The API will
+            always return the log probability of the sampled token, so there
+            may be up to `logprobs+1` elements in the response.
+        prompt_logprobs: Number of log probabilities to return per prompt token.
+        detokenize: Whether to detokenize the output. Defaults to True.
+        skip_special_tokens: Whether to skip special tokens in the output.
+        spaces_between_special_tokens: Whether to add spaces between special
+            tokens in the output.  Defaults to True.
+        logits_processors: List of functions that modify logits based on
+            previously generated tokens, and optionally prompt tokens as
+            a first argument.
+        truncate_prompt_tokens: If set to an integer k, will use only the last k
+            tokens from the prompt (i.e., left truncation). Defaults to None
+            (i.e., no truncation).
+        guided_decoding: If provided, the engine will construct a guided
+            decoding logits processor from these parameters. Defaults to None.
+        logit_bias: If provided, the engine will construct a logits processor
+            that applies these logit biases. Defaults to None.
+        allowed_token_ids: If provided, the engine will construct a logits
+            processor which only retains scores for the given token ids.
+            Defaults to None.
+        prompt_logprob_positions: Optional prompt-token positions whose logits
+            should be materialized. None preserves the standard all-position
+            prompt-logprob behavior.
+    """
+
+    n: int = 1
+    best_of: Optional[int] = None
+    _real_n: Optional[int] = None
+    presence_penalty: float = 0.0
+    frequency_penalty: float = 0.0
+    repetition_penalty: float = 1.0
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1
+    min_p: float = 0.0
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None
+    stop_token_ids: Optional[List[int]] = None
+    ignore_eos: bool = False
+    max_tokens: Optional[int] = 16
+    min_tokens: int = 0
+    logprobs: Optional[int] = None
+    prompt_logprobs: Optional[int] = None
+    # NOTE: This parameter is only exposed at the engine level for now.
+    # It is not exposed in the OpenAI API server, as the OpenAI API does
+    # not support returning only a list of token IDs.
+    detokenize: bool = True
+    skip_special_tokens: bool = True
+    spaces_between_special_tokens: bool = True
+    # Optional[List[LogitsProcessor]] type. We use Any here because
+    # Optional[List[LogitsProcessor]] type is not supported by msgspec.
+    logits_processors: Optional[Any] = None
+    include_stop_str_in_output: bool = False
+    truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None
+    output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE
+
+    # The below fields are not supposed to be used as an input.
+    # They are set in post_init.
+    output_text_buffer_length: int = 0
+    _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set)
+
+    # Fields used to construct logits processors
+    guided_decoding: Optional[GuidedDecodingParams] = None
+    logit_bias: Optional[Dict[int, float]] = None
+    allowed_token_ids: Optional[List[int]] = None
+    prompt_logprob_positions: Optional[List[int]] = None
+
+    @staticmethod
+    def from_optional(
+        n: Optional[int] = 1,
+        best_of: Optional[int] = None,
+        presence_penalty: Optional[float] = 0.0,
+        frequency_penalty: Optional[float] = 0.0,
+        repetition_penalty: Optional[float] = 1.0,
+        temperature: Optional[float] = 1.0,
+        top_p: Optional[float] = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        include_stop_str_in_output: bool = False,
+        ignore_eos: bool = False,
+        max_tokens: Optional[int] = 16,
+        min_tokens: int = 0,
+        logprobs: Optional[int] = None,
+        prompt_logprobs: Optional[int] = None,
+        detokenize: bool = True,
+        skip_special_tokens: bool = True,
+        spaces_between_special_tokens: bool = True,
+        logits_processors: Optional[List[LogitsProcessor]] = None,
+        truncate_prompt_tokens: Optional[Annotated[int,
+                                                   msgspec.Meta(ge=1)]] = None,
+        output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE,
+        guided_decoding: Optional[GuidedDecodingParams] = None,
+        logit_bias: Optional[Union[Dict[int, float], Dict[str, float]]] = None,
+        allowed_token_ids: Optional[List[int]] = None,
+        prompt_logprob_positions: Optional[List[int]] = None,
+    ) -> "SamplingParams":
+        if logit_bias is not None:
+            logit_bias = {
+                int(token): bias
+                for token, bias in logit_bias.items()
+            }
+
+        return SamplingParams(
+            n=1 if n is None else n,
+            best_of=best_of,
+            presence_penalty=0.0
+            if presence_penalty is None else presence_penalty,
+            frequency_penalty=0.0
+            if frequency_penalty is None else frequency_penalty,
+            repetition_penalty=1.0
+            if repetition_penalty is None else repetition_penalty,
+            temperature=1.0 if temperature is None else temperature,
+            top_p=1.0 if top_p is None else top_p,
+            top_k=top_k,
+            min_p=min_p,
+            seed=seed,
+            stop=stop,
+            stop_token_ids=stop_token_ids,
+            include_stop_str_in_output=include_stop_str_in_output,
+            ignore_eos=ignore_eos,
+            max_tokens=max_tokens,
+            min_tokens=min_tokens,
+            logprobs=logprobs,
+            prompt_logprobs=prompt_logprobs,
+            detokenize=detokenize,
+            skip_special_tokens=skip_special_tokens,
+            spaces_between_special_tokens=spaces_between_special_tokens,
+            logits_processors=logits_processors,
+            truncate_prompt_tokens=truncate_prompt_tokens,
+            output_kind=output_kind,
+            guided_decoding=guided_decoding,
+            logit_bias=logit_bias,
+            allowed_token_ids=allowed_token_ids,
+            prompt_logprob_positions=prompt_logprob_positions,
+        )
+
+    def __post_init__(self) -> None:
+        # how we deal with `best_of``:
+        # if `best_of`` is not set, we default to `n`;
+        # if `best_of`` is set, we set `n`` to `best_of`,
+        # and set `_real_n`` to the original `n`.
+        # when we return the result, we will check
+        # if we need to return `n` or `_real_n` results
+        if self.best_of:
+            if self.best_of < self.n:
+                raise ValueError(
+                    f"best_of must be greater than or equal to n, "
+                    f"got n={self.n} and best_of={self.best_of}.")
+            self._real_n = self.n
+            self.n = self.best_of
+        if 0 < self.temperature < _MAX_TEMP:
+            logger.warning(
+                "temperature %s is less than %s, which may cause numerical "
+                "errors nan or inf in tensors. We have maxed it out to %s.",
+                self.temperature, _MAX_TEMP, _MAX_TEMP)
+            self.temperature = max(self.temperature, _MAX_TEMP)
+        if self.seed == -1:
+            self.seed = None
+        else:
+            self.seed = self.seed
+        if self.stop is None:
+            self.stop = []
+        elif isinstance(self.stop, str):
+            self.stop = [self.stop]
+        else:
+            self.stop = list(self.stop)
+        if self.stop_token_ids is None:
+            self.stop_token_ids = []
+        else:
+            self.stop_token_ids = list(self.stop_token_ids)
+        self.logprobs = 1 if self.logprobs is True else self.logprobs
+        self.prompt_logprobs = (1 if self.prompt_logprobs is True else
+                                self.prompt_logprobs)
+        if self.prompt_logprob_positions is not None:
+            self.prompt_logprob_positions = list(
+                self.prompt_logprob_positions)
+
+        # Number of characters to hold back for stop string evaluation
+        # until sequence is finished.
+        if self.stop and not self.include_stop_str_in_output:
+            self.output_text_buffer_length = max(len(s) for s in self.stop) - 1
+
+        self._verify_args()
+
+        if self.temperature < _SAMPLING_EPS:
+            # Zero temperature means greedy sampling.
+            self.top_p = 1.0
+            self.top_k = -1
+            self.min_p = 0.0
+            self._verify_greedy_sampling()
+        # eos_token_id is added to this by the engine
+        self._all_stop_token_ids = set(self.stop_token_ids)
+
+    def _verify_args(self) -> None:
+        if not isinstance(self.n, int):
+            raise ValueError(f"n must be an int, but is of "
+                             f"type {type(self.n)}")
+        if self.n < 1:
+            raise ValueError(f"n must be at least 1, got {self.n}.")
+        if not -2.0 <= self.presence_penalty <= 2.0:
+            raise ValueError("presence_penalty must be in [-2, 2], got "
+                             f"{self.presence_penalty}.")
+        if not -2.0 <= self.frequency_penalty <= 2.0:
+            raise ValueError("frequency_penalty must be in [-2, 2], got "
+                             f"{self.frequency_penalty}.")
+        if not 0.0 < self.repetition_penalty <= 2.0:
+            raise ValueError("repetition_penalty must be in (0, 2], got "
+                             f"{self.repetition_penalty}.")
+        if self.temperature < 0.0:
+            raise ValueError(
+                f"temperature must be non-negative, got {self.temperature}.")
+        if not 0.0 < self.top_p <= 1.0:
+            raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
+        if self.top_k < -1 or self.top_k == 0:
+            raise ValueError(f"top_k must be -1 (disable), or at least 1, "
+                             f"got {self.top_k}.")
+        if not isinstance(self.top_k, int):
+            raise TypeError(
+                f"top_k must be an integer, got {type(self.top_k).__name__}")
+        if not 0.0 <= self.min_p <= 1.0:
+            raise ValueError("min_p must be in [0, 1], got "
+                             f"{self.min_p}.")
+        if self.max_tokens is not None and self.max_tokens < 1:
+            raise ValueError(
+                f"max_tokens must be at least 1, got {self.max_tokens}.")
+        if self.min_tokens < 0:
+            raise ValueError(f"min_tokens must be greater than or equal to 0, "
+                             f"got {self.min_tokens}.")
+        if self.max_tokens is not None and self.min_tokens > self.max_tokens:
+            raise ValueError(
+                f"min_tokens must be less than or equal to "
+                f"max_tokens={self.max_tokens}, got {self.min_tokens}.")
+        if self.logprobs is not None and self.logprobs < 0:
+            raise ValueError(
+                f"logprobs must be non-negative, got {self.logprobs}.")
+        if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
+            raise ValueError(f"prompt_logprobs must be non-negative, got "
+                             f"{self.prompt_logprobs}.")
+        if self.prompt_logprob_positions is not None:
+            if self.prompt_logprobs is None:
+                raise ValueError(
+                    "prompt_logprob_positions requires prompt_logprobs.")
+            if (
+                not self.prompt_logprob_positions
+                or any(
+                    not isinstance(position, int)
+                    or isinstance(position, bool)
+                    or position <= 0
+                    for position in self.prompt_logprob_positions
+                )
+                or self.prompt_logprob_positions
+                != sorted(set(self.prompt_logprob_positions))
+            ):
+                raise ValueError(
+                    "prompt_logprob_positions must be a sorted unique list "
+                    "of positive integers.")
+        if (self.truncate_prompt_tokens is not None
+                and self.truncate_prompt_tokens < 1):
+            raise ValueError(f"truncate_prompt_tokens must be >= 1, "
+                             f"got {self.truncate_prompt_tokens}")
+        assert isinstance(self.stop, list)
+        if any(not stop_str for stop_str in self.stop):
+            raise ValueError("stop cannot contain an empty string.")
+        if self.stop and not self.detokenize:
+            raise ValueError(
+                "stop strings are only supported when detokenize is True. "
+                "Set detokenize=True to use stop.")
+        if self.best_of != self._real_n and self.output_kind == (
+                RequestOutputKind.DELTA):
+            raise ValueError("best_of must equal n to use output_kind=DELTA")
+
+    def _verify_greedy_sampling(self) -> None:
+        if self.n > 1:
+            raise ValueError("n must be 1 when using greedy sampling, "
+                             f"got {self.n}.")
+
+    def update_from_generation_config(
+            self,
+            generation_config: Dict[str, Any],
+            model_eos_token_id: Optional[int] = None) -> None:
+        """Update if there are non-default values from generation_config"""
+
+        if model_eos_token_id is not None:
+            # Add the eos token id into the sampling_params to support
+            # min_tokens processing.
+            self._all_stop_token_ids.add(model_eos_token_id)
+
+        # Update eos_token_id for generation
+        if (eos_ids := generation_config.get("eos_token_id")) is not None:
+            # it can be either int or list of int
+            eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
+            if model_eos_token_id is not None:
+                # We don't need to include the primary eos_token_id in
+                # stop_token_ids since it's handled separately for stopping
+                # purposes.
+                eos_ids.discard(model_eos_token_id)
+            if eos_ids:
+                self._all_stop_token_ids.update(eos_ids)
+                if not self.ignore_eos:
+                    eos_ids.update(self.stop_token_ids)
+                    self.stop_token_ids = list(eos_ids)
+
+    @cached_property
+    def sampling_type(self) -> SamplingType:
+        if self.temperature < _SAMPLING_EPS:
+            return SamplingType.GREEDY
+        if self.seed is not None:
+            return SamplingType.RANDOM_SEED
+        return SamplingType.RANDOM
+
+    @property
+    def all_stop_token_ids(self) -> Set[int]:
+        return self._all_stop_token_ids
+
+    def clone(self) -> "SamplingParams":
+        """Deep copy excluding LogitsProcessor objects.
+
+        LogitsProcessor objects are excluded because they may contain an
+        arbitrary, nontrivial amount of data.
+        See https://github.com/vllm-project/vllm/issues/3087
+        """
+
+        logit_processor_refs = None if self.logits_processors is None else {
+            id(lp): lp
+            for lp in self.logits_processors
+        }
+        return copy.deepcopy(self, memo=logit_processor_refs)
+
+    def __repr__(self) -> str:
+        return (
+            f"SamplingParams(n={self.n}, "
+            f"presence_penalty={self.presence_penalty}, "
+            f"frequency_penalty={self.frequency_penalty}, "
+            f"repetition_penalty={self.repetition_penalty}, "
+            f"temperature={self.temperature}, "
+            f"top_p={self.top_p}, "
+            f"top_k={self.top_k}, "
+            f"min_p={self.min_p}, "
+            f"seed={self.seed}, "
+            f"stop={self.stop}, "
+            f"stop_token_ids={self.stop_token_ids}, "
+            f"include_stop_str_in_output={self.include_stop_str_in_output}, "
+            f"ignore_eos={self.ignore_eos}, "
+            f"max_tokens={self.max_tokens}, "
+            f"min_tokens={self.min_tokens}, "
+            f"logprobs={self.logprobs}, "
+            f"prompt_logprobs={self.prompt_logprobs}, "
+            "prompt_logprob_positions="
+            f"{self.prompt_logprob_positions}, "
+            f"skip_special_tokens={self.skip_special_tokens}, "
+            "spaces_between_special_tokens="
+            f"{self.spaces_between_special_tokens}, "
+            f"truncate_prompt_tokens={self.truncate_prompt_tokens}), "
+            f"guided_decoding={self.guided_decoding}")
+
+
+class BeamSearchParams(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        # required for @cached_property.
+        dict=True):  # type: ignore[call-arg]
+    """Beam search parameters for text generation."""
+    beam_width: int
+    max_tokens: int
+    ignore_eos: bool = False
+    temperature: float = 0.0
+    length_penalty: float = 1.0