[Hybrid] Decoupling block_size from allocation block size

[s3woz]

Purpose

Note: This is a draft PR and needs testing & verification.

Hybrid models often comprise attention layers with relatively small KVcache states and other layer types (e.g. Mamba, GDN) with relatively large cache states. Current vLLM implementation sets the block_size in a way to keep the KVcache allocation unit of block the same for both attention and hybrid cache. This typically involves considerably increasing the block_size of attention, resulting in lower cache hit chances.

This PR separates the block_size setting from the size of hybrid cache entries by allowing different sizes of block allocations for attention and SSM layers. This allows to set any block_size for hybrid models and by default to keep the block_size of the attention, increasing the cache hit ratio.

Note: Currently enabled only for align mode, to avoid too fast memory filling of all mode.
When running Mamba-based (non GDN) models switch from all to align with: --mamba-cache-mode align

In a prefill-oriented synthetic test with potential full cache hits (see test plan below), current PR substantially increases the cache hits (avoids the "steps" in the figure below), and speeds up the prefill (2.7x times) for Qwen/Qwen3.6-35B-A3B:

Technical details

High-level design idea:

1. Get back to the logic of flexible setting of block size also for hybrid models
2. Handle different KV cache block sizes by introducing size information in class SingleTypeKVCacheManager(ABC): that is set by each layer type accordingly (current PR just uses a flag _is_large_block: bool and sets to True for SSM layers, but we could make it to keep the exact size).
3. Use the size information when getting new blocks in kv cache manager:

       allocated_blocks = self.block_pool.get_new_blocks(
            cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks),
            # NEW size information -> current flag could be replaced with size specifier:
            large_block=self._is_large_block,
            # NEW handle fragmentation - see below:
            last_hit_block_id=_last_non_null_block_id(req_blocks),
        )

Primary challenge with variable-sized allocations is the KV cache blocks memory fragmentation. To handle this, we pass information about previous allocation for each request to the block_pool, so that it can avoid fragmentation. Note: currently computed with _last_non_null_block_id(req_blocks), but I would rather just pass req_blocks there and let the block_pool figure out how to defragment.

At this point we have:

1. [+] Minimal changes to KV cache interface API
2. [+] Variable-sized KV cache blocks allocation
3. [+] Ability to freely tune the block size
4. [-] Potential fragmentation issues

Fragmentation avoidance policy:
All of the complexity gets hidden into block_pool as follows:

1. Different block_pool implementation can be used, based on model configuration and defragmentation policy.
2. Current implementation uses standard code paths for non-hybrid models. For hybrid-models it uses the following solution:
   1. It creates various views of the same memory: large blocks (size of SSM state) and small blocks (size of attention state). The standard block list keeps the pointers to large blocks. Each large block keeps a list of pointers to small blocks inside it. All of this is just list and pointer arithmetic. Number of python objects/list entries is comparable to the baseline case with small block set for attention models.
   2. Assume e.g. 1 large block = 100 small blocks. These pointers overlap and have aligned memory addresses, i.e. GPU tensor address for large block number 2 is the same as small block number 200. This way all the current pointer arithmetic in gpu_model_runner can be reused with no changes, and num_blocks for tensor shape can be determined based on state_page_bytes for each layer type. All kernels have their own views of the shared memory.
   3. When SSM allocates, it gets a large block form the list and marks as allocated. Entire large block is assigned a single hash value and can be matched by APC.
   4. When attention allocates, it requests small blocks as follows: it marks large block as allocated but then returns the requested number of small blocks from that large block. Each small block gets assigned a hash and can be matched by APC.
   5. Fragmentation is an issue only for attention requests. There the block_pool checks what was the last_hit_block_id for that request and checks if there are small blocks available within that large block. It returns them if they are available, if not it takes small blocks from the next free large block. This bundles together memory blocks from the same request or requests that have a shared prefix, which seems like a viable heuristic. (It's much better than allowing random requests to allocate, as one frequently used small block could lock the entire large block for a long time.)

This is one possible implementation. By having a clear API decoupling between the KV cache manager allocations and the underlying block_pool implementation, the community could contribute other alternative implementations. (current implementation has a limitation that it handles only two allocation size: large and small, but it opens a path to new alternative implementations with more options)

@tdoublep

Test Plan

if __name__ == '__main__':
    SETUP = {'model': 'Qwen/Qwen3.6-35B-A3B', 'max_model_len': '75000', 'max_num_seqs': 50}
    import os
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
    from vllm import LLM, SamplingParams
    from vllm.distributed import cleanup_dist_env_and_memory
    import time, string
    sampling_params = SamplingParams(temperature=0.0, max_tokens=1)
    prefix = ( # examples/offline_inference/prefix_caching.py
        "You are an expert school principal, skilled in effectively managing "
        "faculty and staff. Draft 10-15 questions for a potential first grade "
        "Head Teacher for my K-12, all-girls', independent school that emphasizes "
        "community, joyful discovery, and life-long learning. The candidate is "
        "coming in for a first-round panel interview for a 8th grade Math "
        "teaching role. They have 5 years of previous teaching experience "
        "as an assistant teacher at a co-ed, public school with experience "
        "in middle school math teaching.")
    NUM_PROMPTS = 200
    engine = LLM(enable_prefix_caching=True,
        disable_log_stats=False, **SETUP)
    so_far_cached = 0
    for MULTIPLE in range(5, 35, 2):
        prompt = str(MULTIPLE) + MULTIPLE * prefix + ("What is the capital of France?")
        # Initial prompt
        outputs = engine.generate(prompt, sampling_params, use_tqdm=False)
        start_time = time.time()
        for i in range(NUM_PROMPTS):
            outputs = engine.generate(prompt, sampling_params, use_tqdm=False)
            #print(f"Generated text: {outputs[0].outputs[0].text!r}")
        total_time = time.time() - start_time
        for m in engine.llm_engine.get_metrics():
            if 'vllm:prompt_tokens_cached' in m.name:
                print("Multiple", MULTIPLE,
                    "Hits/prompt", (m.value - so_far_cached) // NUM_PROMPTS,
                    "Took:", total_time, "seconds")
                so_far_cached = m.value

Test Result

• Main:

Multiple 5 Hits/prompt 0 Took: 12.20611047744751 seconds
Multiple 7 Hits/prompt 0 Took: 12.233927726745605 seconds
Multiple 9 Hits/prompt 0 Took: 12.590906381607056 seconds
Multiple 11 Hits/prompt 1056 Took: 5.146499872207642 seconds
Multiple 13 Hits/prompt 1056 Took: 12.73249340057373 seconds
Multiple 15 Hits/prompt 1056 Took: 12.883115530014038 seconds
Multiple 17 Hits/prompt 1056 Took: 12.990457773208618 seconds
Multiple 19 Hits/prompt 1056 Took: 13.023258686065674 seconds
Multiple 21 Hits/prompt 2112 Took: 5.588909864425659 seconds
Multiple 23 Hits/prompt 2112 Took: 13.212487697601318 seconds
Multiple 25 Hits/prompt 2112 Took: 13.240469217300415 seconds
Multiple 27 Hits/prompt 2112 Took: 13.315996170043945 seconds
Multiple 29 Hits/prompt 2112 Took: 13.515663623809814 seconds
Multiple 31 Hits/prompt 3168 Took: 5.614190578460693 seconds
Multiple 33 Hits/prompt 3168 Took: 14.520177841186523 seconds

• This PR:

Multiple 5 Hits/prompt 512 Took: 3.692441463470459 seconds
Multiple 7 Hits/prompt 720 Took: 3.752166748046875 seconds
Multiple 9 Hits/prompt 928 Took: 3.8311171531677246 seconds
Multiple 11 Hits/prompt 1136 Took: 3.9413886070251465 seconds
Multiple 13 Hits/prompt 1344 Took: 3.9924542903900146 seconds
Multiple 15 Hits/prompt 1552 Took: 4.049445152282715 seconds
Multiple 17 Hits/prompt 1735 Took: 4.199482679367065 seconds
Multiple 19 Hits/prompt 1952 Took: 4.246189594268799 seconds
Multiple 21 Hits/prompt 2160 Took: 4.319143056869507 seconds
Multiple 23 Hits/prompt 2368 Took: 4.434829950332642 seconds
Multiple 25 Hits/prompt 2576 Took: 4.527528524398804 seconds
Multiple 27 Hits/prompt 2784 Took: 4.6004273891448975 seconds
Multiple 29 Hits/prompt 2992 Took: 4.739585876464844 seconds
Multiple 31 Hits/prompt 3200 Took: 4.832530736923218 seconds
Multiple 33 Hits/prompt 3375 Took: 4.971061706542969 seconds

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @s3woz.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[eugr]

@mgoin - can anyone look into this? This PR results in 2x TTFT reduction for follow up requests on Qwen3.6-35B-A3B which is a significant improvement. Current mamba cache performance is pretty poor, especially compared to llama.cpp to the point that despite having 3x slower initial prefill, llama.cpp still wins over vLLM by being able to utilize the cache more efficiently. That's on DGX Spark.

[s3woz]

This is a quickly drafted PoC. I've just updated it with a more extensive description of the idea above in Technical details. If there is interest, I could refactor and clean up the code.
FYI: @tdoublep

[danielrmay]

Hi @s3woz - apologies for the tag noise above, I've been iterating on a continuation of this draft in a fork as an experiment and didn't realize each cleanup push would ref the thread here, lesson learned! 🙏

The headline is that your design here roughly triples usable capacity (74k to ~245k tokens) with prefix caching turned on for the latest Nemotron 75B model at W4A16 in VRAM-constrained scenarios like 2x24GB consumer GPUs, with fine-grained cache hits intact. I used AI agents heavily in the implementation and as such have been subjecting it to heavy review, and intended to come back to this thread once I had it into a feasible state I was ready to defend. I'm still assessing what's left to get it fit for a PR, it's not there yet, but wanted to provide an early explanation and share the positive results I'd seen.

[Dao007forever]

#45702 is working on improving the cache hit rate for hybrid model.