[Feature][MP] Negotiate chunk_size from vLLM KV geometry at first registration

Motivation

The MP server currently starts with a fixed chunk_size (default: 256) before it knows the serving model's KV-cache geometry. Later, the vLLM connector derives per-engine-group tokens_per_block from vLLM's KVCacheConfig / KVCacheGroupSpec and checks that the server chunk size is compatible.

This is correct for safety, but it is hard to configure for models whose cacheable groups do not use common powers-of-two boundaries.

For example, in some hybrid / Mamba / sub-paged-attention layouts, an attention backend may still use 32-token physical kernel pages, while vLLM's manager/logical block for that group covers a larger aligned page, e.g. 17 * 32 = 544 tokens. In that case:

tokens_per_block = 544
chunk_size = 256   -> invalid
chunk_size = 512   -> invalid
chunk_size = 1024  -> invalid
chunk_size = 544   -> valid for that group
chunk_size = 1088  -> valid for that group

Concrete examples from real vLLM + LMCache runs

This is not only a theoretical 544-token example. Recent PR #5042 validation hit the same class of issue with real hybrid layouts.

Example 1: GLM-5.3-Flash

When serving zai-org/GLM-5.3-Flash through vLLM with the LMCache MP connector, vLLM reports this model-specific geometry during connector initialization:

group_tokens_per_block = [640, 0, 640, 640, 640, 640]
scheduler_block_size = 640

The 0 entry is a non-cacheable / scratch group, so it should not constrain LMCache chunking. The positive cacheable groups require the server chunk size to be a multiple of 640.

That means a server started with common fixed values such as 256, 512, or 1024 would be incompatible with this model geometry. The operator does not naturally know the magic number 640 when starting the LMCache server; it is only known after vLLM initializes the model and exports the KV-cache group metadata.

With a compatible server chunk_size=640, the H200 validation run behaved as expected:

cold request: num_lmcache_cached_tokens = 0
warm request: num_lmcache_cached_tokens = 1280

The warm hit count is exactly two 640-token LMCache chunks, which confirms that this geometry is not just metadata bookkeeping; it directly determines what can be stored and reused.

Example 2: Qwen3.8-Flash-Next-FP8

When serving Qwen3.8-Flash-Next-FP8 with the same connector path, the compatible run used chunk_size=400. The registered cache context contained cacheable groups with tokens_per_block=400; one group also had slots_per_block=100, showing that logical token grouping and physical slot grouping can differ for these hybrid layouts.

With the compatible 400-token chunk size, the H200 validation run showed:

cold request: num_lmcache_cached_tokens = 0
warm request: num_lmcache_cached_tokens = 800

Again, the warm hit count is two model-derived chunks. A fixed default 256 is therefore not a generally safe launch-time guess for this kind of model.

If the model has several cacheable groups, the real requirement is that chunk_size be a common multiple of all positive group tokens_per_block values. Non-prefix-cacheable / scratch groups (tokens_per_block = 0) should not participate in this constraint.

Current behavior

Today the server-side chunk size is configured up front:

• MPCacheServerContext.__init__(..., chunk_size=256, ...) stores the process-wide chunk size.
• TokenHasher(chunk_size=chunk_size, ...) and SessionManager are initialized from that value.
• GET_CHUNK_SIZE returns this fixed value to clients.

The vLLM connector then detects the model geometry and fails fast if the server value is incompatible:

• LMCacheMPConnector computes group_tokens_per_block and math.lcm(*cached_spans) for hit alignment.
• The scheduler role checks lmcache_tokens_per_chunk % tokens_per_block == 0 for each cacheable group.
• LMCacheMPWorkerAdapter.register_kv_caches() checks the same condition on the final EngineGroupInfo list.
• KVLayerGroupsManager also validates that lmcache_tokens_per_chunk is a multiple of each kernel group's tokens_per_block.

That failure is much better than silently slicing block IDs incorrectly, but the user often cannot know the correct chunk size when launching the MP server, because the required geometry is only available after vLLM initializes the model.

Proposed direction

Support an opt-in auto / late-bound chunk size mode for the MP server.

Conceptually:

LMCache server starts with chunk_size = auto / unset

First REGISTER_KV_CACHE request carries:
  required_alignment = lcm(all positive group tokens_per_block)
  preferred_chunk_size = user_desired_or_default rounded up to a legal multiple

Server receives first registration:
  if chunk_size is unset:
      bind chunk_size = preferred_chunk_size
      initialize or finalize TokenHasher / SessionManager / chunk-dependent runtime state
  else:
      validate existing chunk_size against required_alignment

This should be a one-shot binding, not arbitrary runtime mutation. Once the server has accepted any lookup/store/retrieve/register state using a chunk size, changing it would invalidate token hashing, object boundaries, stored keys, sessions, layout descriptors, and transfer planning.

Suggested API / protocol shape

One possible design:

1. Add an optional chunk_size_hint / recommended_chunk_size and required_chunk_alignment to the REGISTER_KV_CACHE payload, or add a separate pre-registration negotiation RPC.
2. Let the vLLM connector compute:

cacheable_spans = [x for x in group_tokens_per_block if x > 0]
required_alignment = math.lcm(*cacheable_spans)
recommended_chunk_size = round_up(user_base_chunk_size, required_alignment)

For example:

group_tokens_per_block = [544, 64]
required_alignment = lcm(544, 64) = 1088
base/default preference = 1024
recommended_chunk_size = 1088

1. On the server, allow the hint only when the context is still unbound / empty.
2. Reject conflicting later registrations unless the already-bound chunk size is compatible with the new model's required_alignment.
3. Keep explicit fixed chunk_size behavior unchanged for existing deployments.

Important constraints

• This should not allow the second worker/model to silently override the first selected chunk size.
• Multi-worker registration needs a lock / compare-and-set boundary so only one registration binds the server chunk size.
• Multi-server deployments must converge on the same chunk size, otherwise token hashes and object keys diverge across servers.
• Any already-created token hasher, session manager, layout descriptor registry entries, and storage objects must be treated as chunk-size-bound state.
• Existing chunk_size: 256 configs should remain fixed and fail fast on incompatible models, preserving current safety.

Why this matters

Without this, users have to guess a server chunk_size before vLLM exposes the model-specific KV geometry. Common values like 256, 512, or 1024 are not guaranteed to work for hybrid / MLA / Mamba layouts. The current fail-fast checks are correct, but a registration-time negotiation path would make the system self-adaptive while preserving correctness.

Related context: PR #5042 made more hybrid/MLA/Mamba KV layouts registerable, which makes this configuration issue easier to hit in real deployments.