refactor(v1/core): re-port BlockPool + Request block-hashes to e24d1b24 API (was classic); fix free_blocks LRU split

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

Author: mudler

include/vllm/v1/core/block_pool.h

@@ -8,39 +8,42 @@
 // (no CUDA, no model): the physical KV memory itself is not touched here, only
 // the block bookkeeping.
 //
-// VERSION NOTE (which upstream is mirrored): the landed kv_cache_utils.{h,cpp}
-// mirror the CLASSIC vLLM V1 block-hashing API — free functions
-// hash_block_tokens / hash_request_tokens that take a pluggable hash_function,
-// BlockHash as raw hash bytes, and make_block_hash_with_group_id packing a
-// 4-byte group id. BlockPool is ported to match THOSE utils, i.e. the classic
-// BlockPool shape:
-//   - cached_block_hash_to_block is a plain
-//     {BlockHashWithGroupId -> {block_id -> KVCacheBlock*}} nesting (the classic
-//     defaultdict(dict)); the e24d1b24 BlockHashToBlockMap single-vs-dict union
-//     optimization is intentionally NOT ported — it is a GC micro-optimization
-//     with identical observable semantics.
-//   - cache_full_blocks carries the classic (block_hashes, hash_fn) parameters:
-//     the caller passes the precomputed per-block hashes (from
-//     hash_request_tokens) and a hash function used to fill in any trailing
-//     blocks whose hashes were not precomputed. The e24d1b24 signature instead
-//     reads request.block_hashes (a Request field not present in the T0 Request
-//     port) and takes hash_block_size on the pool; that path is not portable
-//     against the landed utils, so the classic signature is used.
-// The @ e24d1b24 pin marks the upstream tree tracked for the surrounding V1
-// core; the BlockPool logic mirrored is upstream's classic pre-BlockHashToBlockMap
-// form, consistent with the already-landed classic-hashing utils.
+// API FIDELITY: this mirrors the e24d1b24 BlockPool API 1:1 for the COMMON
+// prefix-caching path. cache_full_blocks takes the pinned signature
+// (request, blocks, num_cached_blocks, num_full_blocks, block_size,
+// kv_cache_group_id, block_mask) and reads request.block_hashes — the per-block
+// hashes the Request computes incrementally via its _block_hasher
+// (get_request_block_hasher). The pool carries hash_block_size (per the pin).
 //
-// DEFERRED, recorded (kept as stubs / no-ops so signatures stay 1:1 and the
-// later units fill them in without a call-site change):
+// DEFERRED behind 1:1 stubs (documented; the gate models — text-only GDN/MoE —
+// never exercise these, and later units fill them in without a call-site change):
+//   - The align / multiple-block-size path in cache_full_blocks
+//     (block_size != hash_block_size, upstream's BlockHashListWithBlockSize):
+//     guarded with a throw. Only the block_size == hash_block_size common case
+//     is ported.
+//   - The partial->full promotion branch inside cache_full_blocks (a "new full
+//     block" that already carries a hash): guarded with a throw. Requires the
+//     deferred partial primitives.
+//   - cache_partial_block / _get_partial_block_hash /
+//     _get_partial_block_parent_hash_and_start: throw-if-called stubs (partial
+//     prefix entries, not needed by the gate models).
+//   - evict_blocks (KV-connector-driven cache eviction): throw-if-called stub.
 //   - KV-cache events: enable_kv_cache_events, kv_event_queue, take_events(),
-//     and the BlockStored / BlockRemoved emission inside cache_full_blocks /
-//     _maybe_evict_cached_block. KVCacheEvent is a placeholder struct; the
-//     emission branches are omitted (marked in the .cpp). This mirrors upstream
-//     vllm/distributed/kv_events.py, deferred with the rest of the V1 event bus.
+//     and the BlockStored / BlockRemoved / AllBlocksCleared emission inside
+//     cache_full_blocks / _remove_cached_block_hashes / reset_prefix_cache.
+//     KVCacheEvent is a placeholder struct; the emission branches are omitted
+//     (marked in the .cpp). Mirrors upstream vllm/distributed/kv_events.py.
 //   - metrics_collector (KVCacheMetricsCollector): the on_block_allocated /
-//     _accessed / _evicted hooks are omitted. Not part of the correctness core.
-//   - cache_partial_block / evict_blocks: not in the classic shape and not
-//     needed by the gate models; omitted.
+//     _accessed / _evicted / reset hooks are omitted. Not part of the
+//     correctness core.
+//   - block_mask: the parameter is present and honored (masked blocks skipped),
+//     but only ever passed None by the ported (non-SWA/non-mamba) call sites.
+//
+// DEVIATION, recorded: the BlockHashToBlockMap single-block-vs-dict union
+// (upstream's GC micro-optimization) is NOT ported. cached_block_hash_to_block
+// is the plain {BlockHashWithGroupId -> {block_id -> KVCacheBlock*}} nesting.
+// Observable semantics are identical (get_one_block / contain / insert / pop
+// map onto the nested-map operations); only the inner GC cost differs.
 //
 // DEVIATIONS, recorded:
 //   - The null block (block_id 0) is made un-allocatable exactly as upstream:
@@ -64,6 +67,7 @@
 #include <cstdint>
 #include <map>
 #include <optional>
+#include <set>
 #include <unordered_map>
 #include <vector>
 
@@ -88,8 +92,11 @@ class BlockPool {
   // Args:
   //   num_gpu_blocks: The number of blocks in the pool (must be > 0).
   //   enable_caching: Whether to enable prefix caching.
+  //   hash_block_size: The block size at which block hashes are computed. The
+  //     actual block size usually equals it; with differently-sized KV cache
+  //     groups it can be a multiple (the align path, DEFERRED).
   //   enable_kv_cache_events: Whether to enable kv cache events (DEFERRED).
-  BlockPool(int64_t num_gpu_blocks, bool enable_caching,
+  BlockPool(int64_t num_gpu_blocks, bool enable_caching, int hash_block_size,
             bool enable_kv_cache_events = false);
 
   // Non-copyable / non-movable (see DEVIATIONS in the file header).
@@ -104,25 +111,39 @@ class BlockPool {
   std::optional<std::vector<KVCacheBlock*>> get_cached_block(
       const BlockHash& block_hash, const std::vector<int>& kv_cache_group_ids);
 
-  // Cache a list of full blocks for prefix caching. Updates each newly-full
-  // block's hash metadata and inserts it into cached_block_hash_to_block.
+  // Cache a list of full blocks for prefix caching. Reads request.block_hashes
+  // (computed incrementally by the Request's _block_hasher), updates each
+  // newly-full block's hash metadata and inserts it into
+  // cached_block_hash_to_block.
   //
-  // block_hashes: the request's per-block hashes; block_hashes[k] is the hash
-  //   of the k-th full block. May be extended in place for trailing blocks
-  //   whose hash was not precomputed (the else branch computes it via hash_fn
-  //   and appends it — matching upstream, which caches it on the request for a
-  //   possible future preemption).
+  // request: the request whose block_hashes drive caching.
+  // blocks: all blocks in the request (the range
+  //   [num_cached_blocks, num_full_blocks) is cached).
   // num_cached_blocks: number of blocks already cached.
   // num_full_blocks: number of blocks that are full and should be cached now.
   // block_size: number of tokens per block.
   // kv_cache_group_id: id of the KV cache group.
-  // hash_fn: the pluggable block-hash function (see kv_cache_utils.h).
-  void cache_full_blocks(const Request& request,
-                         const std::vector<KVCacheBlock*>& blocks,
-                         std::vector<BlockHash>& block_hashes,
-                         int num_cached_blocks, int num_full_blocks,
-                         int block_size, int kv_cache_group_id,
-                         const HashFn& hash_fn);
+  // block_mask: optional mask aligned with blocks[num_cached_blocks:
+  //   num_full_blocks]; masked-off (false) blocks are skipped like null blocks.
+  //   Only ever None from the ported call sites (see header DEFERRED note).
+  //
+  // Common path only (block_size == hash_block_size); the align path and the
+  // partial->full promotion branch are DEFERRED (throw). See the header.
+  void cache_full_blocks(
+      const Request& request, const std::vector<KVCacheBlock*>& blocks,
+      int num_cached_blocks, int num_full_blocks, int block_size,
+      int kv_cache_group_id,
+      const std::optional<std::vector<bool>>& block_mask = std::nullopt);
+
+  // Register a partial prefix-cache entry for an existing block. DEFERRED 1:1
+  // stub (partial primitives, not needed by the gate models): throws if called.
+  std::optional<BlockHashWithGroupId> cache_partial_block(
+      const Request& request, KVCacheBlock* block, int num_tokens,
+      int kv_cache_group_id, int block_size);
+
+  // Evict blocks from the prefix cache by their block IDs (KV-connector-driven).
+  // DEFERRED 1:1 stub: throws if called.
+  void evict_blocks(const std::set<int>& block_ids);
 
   // Get num_blocks new blocks from the free block pool. Does NOT check the
   // prefix cache. Throws std::runtime_error (upstream ValueError) if there are
@@ -133,6 +154,20 @@ class BlockPool {
   // and evict it from the cache. Returns true iff the block was evicted.
   bool _maybe_evict_cached_block(KVCacheBlock* block);
 
+  // Remove every hash key that points to `block` (its primary hash plus any
+  // partial-alias hashes in cached_block_hashes_by_block), reset the block's
+  // hash, and return the keys actually removed from the map. Mirrors upstream
+  // _remove_cached_block_hashes.
+  std::vector<BlockHashWithGroupId> _remove_cached_block_hashes(
+      KVCacheBlock* block);
+
+  // Insert a hash key for `block` into cached_block_hash_to_block. If the block
+  // has no primary hash yet, the key becomes its primary hash (with num_tokens);
+  // otherwise the key is tracked as a partial alias in
+  // cached_block_hashes_by_block. Mirrors upstream _insert_block_hash.
+  void _insert_block_hash(const BlockHashWithGroupId& block_hash_with_group_id,
+                          KVCacheBlock* block, std::optional<int> num_tokens);
+
   // Touch a block: increase its reference count by 1, and remove it from the
   // free queue if it was an eviction candidate (ref_cnt == 0). Used when a
   // block is hit by another request with the same prefix.
@@ -162,6 +197,7 @@ class BlockPool {
 
   int64_t num_gpu_blocks;
   bool enable_caching;
+  int hash_block_size;
 
   // All kv-cache blocks, indexed by block_id (blocks[0] is the null block).
   std::vector<KVCacheBlock> blocks;
@@ -176,6 +212,14 @@ class BlockPool {
   std::unordered_map<BlockHashWithGroupId, std::map<int, KVCacheBlock*>>
       cached_block_hash_to_block;
 
+  // {block_id -> set of partial-alias hash keys}. Tracks the extra hash keys
+  // (beyond a block's primary block_hash) that point to the block, so eviction /
+  // reset / promotion can remove every key. Only populated by the deferred
+  // partial primitives; empty on the common path. Mirrors upstream
+  // cached_block_hashes_by_block.
+  std::unordered_map<int, std::set<BlockHashWithGroupId>>
+      cached_block_hashes_by_block;
+
   // The null placeholder block (block_id 0). Its ref_cnt is NOT maintained;
   // it is popped out of the free queue so it can never be allocated.
   KVCacheBlock* null_block;

include/vllm/v1/core/kv_cache_utils.h

@@ -268,10 +268,18 @@ class FreeKVCacheBlockQueue {
   KVCacheBlock fake_free_list_tail{-1};
 };
 
-// Forward declaration for generate_block_hash_extra_keys (defined in
-// vllm/v1/request.h; only a const-ref is needed here).
+// Forward declaration for generate_block_hash_extra_keys / the request block
+// hasher (Request is defined in vllm/v1/request.h; only a const-ref is needed
+// here).
 struct Request;
 
+// The request block hasher. Mirrors upstream's
+// `Callable[[Request], list[BlockHash]]` (Request._block_hasher): given a
+// request, returns the block hashes for its newly-complete full blocks (the
+// ones not yet in request.block_hashes). A null BlockHasher mirrors upstream
+// `block_hasher=None` => prefix caching off (update_block_hashes is a no-op).
+using BlockHasher = std::function<std::vector<BlockHash>(const Request&)>;
+
 // generate_block_hash_extra_keys — DEFERRED derivation.
 //
 // Upstream derives a block's extra hash keys from a Request's multi-modal
@@ -310,6 +318,17 @@ std::vector<BlockHash> hash_request_tokens(
     const std::vector<int32_t>& token_ids,
     const std::vector<ExtraKeys>& per_block_extra_keys = {});
 
+// Build the incremental request block hasher. Mirrors upstream
+// get_request_block_hasher: the returned closure computes ONLY the not-yet-hashed
+// full blocks of a request (starting at
+// request.block_hashes.size() * hash_block_size), chaining each block's hash
+// onto the previous one (request.block_hashes.back() as the parent). It hashes
+// full blocks only; a partial trailing block is left unhashed. Request stores
+// this closure as _block_hasher and calls it from update_block_hashes() at
+// construction and after each append.
+BlockHasher get_request_block_hasher(int hash_block_size,
+                                     const HashFn& caching_hash_fn);
+
 }  // namespace vllm::v1
 
 #endif  // VLLM_V1_CORE_KV_CACHE_UTILS_H_

include/vllm/v1/request.h

@@ -14,8 +14,7 @@
 // slot these in without reshaping the struct:
 //   - prompt_embeds / prompt_is_token_ids / _prompt_embeds_per_block_hashes,
 //     mm_features (multimodal), pooling_params, structured_output_request,
-//     lora_request, cache_salt (prefix caching salt), block_hashes /
-//     _block_hasher / update_block_hashes (M1.2 BlockPool), events /
+//     lora_request, cache_salt (prefix caching salt), events /
 //     stop_reason / kv_transfer_params, spec_token_ids, priority /
 //     client_index / __lt__ (priority scheduling), streaming / resumable
 //     state, prefill_stats, async-scheduling counters
@@ -48,6 +47,7 @@
 #include <vector>
 
 #include "vllm/sampling_params.h"
+#include "vllm/v1/core/kv_cache_utils.h"  // BlockHash, BlockHasher
 
 namespace vllm::v1 {
 
@@ -106,16 +106,22 @@ std::optional<FinishReason> GetFinishedReason(RequestStatus status);
 // A generation request tracked by the V1 engine (T0 field subset). The
 // scheduler / model runner mutate this in place exactly as upstream does.
 struct Request {
+  // block_hasher mirrors upstream's `block_hasher=None` default: null => prefix
+  // caching off (update_block_hashes() is a no-op, block_hashes stays empty).
   Request(std::string request_id, std::vector<int32_t> prompt_token_ids,
-          SamplingParams sampling_params, double arrival_time);
+          SamplingParams sampling_params, double arrival_time,
+          BlockHasher block_hasher = nullptr);
 
   // from_engine_core_request: build a Request from the frontend->core message.
   // Mirrors upstream Request.from_engine_core_request for the T0 fields
   // (request_id, prompt_token_ids, sampling_params, arrival_time); status
   // starts kWaiting, num_computed_tokens 0, output empty. The params arrived
   // already PostInit'd / validated by the frontend, so this does NOT
-  // re-validate (upstream's factory doesn't either).
-  static Request FromEngineCoreRequest(const EngineCoreRequest& request);
+  // re-validate (upstream's factory doesn't either). block_hasher is injected
+  // by the engine exactly as upstream from_engine_core_request(request,
+  // block_hasher).
+  static Request FromEngineCoreRequest(const EngineCoreRequest& request,
+                                       BlockHasher block_hasher = nullptr);
 
   std::string request_id;
   std::vector<int32_t> prompt_token_ids;
@@ -138,12 +144,28 @@ struct Request {
   // num_output_tokens: len(_output_token_ids).
   int NumOutputTokens() const;
 
+  // Per-block hashes of this request's full blocks, computed at
+  // hash_block_size granularity and chained over the full prefix. Populated by
+  // update_block_hashes() via _block_hasher (empty when caching is off).
+  std::vector<BlockHash> block_hashes;
+
+  // The block hasher (upstream _block_hasher). Stored without binding a
+  // back-reference to this Request (upstream avoids the Request->partial->Request
+  // reference cycle); here it is a plain std::function taking the request by
+  // const ref. Null => prefix caching off.
+  BlockHasher block_hasher_;
+
   // append_output_token_ids(int | list[int]): append the sampled token(s) to
-  // output_token_ids. (Upstream also mirrors into _all_token_ids and updates
-  // block hashes; NumTokens recomputes here and block hashing is deferred.)
+  // output_token_ids, then update_block_hashes() (upstream mirrors into
+  // _all_token_ids too; NumTokens recomputes prompt+output here).
   void AppendOutputToken(int32_t token_id);
   void AppendOutputToken(const std::vector<int32_t>& token_ids);
 
+  // update_block_hashes: compute block hashes for any newly-complete full
+  // blocks and append them to block_hashes. No-op when _block_hasher is null.
+  // Mirrors upstream Request.update_block_hashes.
+  void update_block_hashes();
+
   // is_finished / get_finished_reason: delegate to RequestStatus.
   bool IsFinished() const;
   std::optional<FinishReason> GetFinishedReason() const;