[Hybrid] Marconi-style admission policy for hybrid cache

[s3woz]

Purpose

This PR implements Marconi-style admission policy for hybrid cache used e.g. for Qwen models. In Marconi paper an effective cache admission policy is proposed that caches in two cases:

1. Last state - important for chat workloads
2. Shared prefix (cache after the 2nd hit: observe prompts, if a shared prefix is detected, then cache it) - important for system prompts, instructions, few-shot examples, self-consistency, long-document Q&A, RL workloads

Logic before this PR:

• Only the last block-aligned block is cached, which conceptually corresponds to Last state caching

Logic with this PR:

• Enable additionally shared prefix caching
• The implementation in this PR differs from Marconi (that uses radix tree).
  • Here, existing regular attention KVCache is used for detecting presence of a shared prefix (i.e. a non-cached shared prefix exists if standard attention KVCache has hits, but SSM attention doesn't, i.e. it lags behind the KVCache)
  • If a shared prefix is detected, it is cached at the last block-aligned position (to avoid the need for any kernel adjustments and minimize the impact to vLLM logic)

A synthetic test below demonstrates:

• Cache hits for shared prefix when using this PR, and no cache hits for shared prefix in Main.
• With APC enabled, the Main is worse by up to 40% (Qwen/Qwen3.5-0.8B) or 66% (Qwen/Qwen3.5-35B-A3B) in terms of latency over the proposed PR.
• With APC disabled, no noticeable overhead.

@tdoublep @bohnstingl

Test Plan

if __name__ == "__main__":
    from vllm import LLM, SamplingParams
    from vllm.distributed import cleanup_dist_env_and_memory
    import time, string
    MODEL = "Qwen/Qwen3.5-0.8B"
    sampling_params = SamplingParams(temperature=0.0)
    prefix1 = ( # 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. ")
    prefix2 = ("Based on these information, fulfill "
                "the following paragraph: ")
    # Long context
    SHARED_LONG_PREFIX_MULTIPLE = 150
    prefix = SHARED_LONG_PREFIX_MULTIPLE * prefix1
    # How many times to diverge after Long context above
    DIVERGENT_PROMPTS = 20    
    # Moderately-sized prompt that spans more than 1 cache block
    MULTIPLE = 15
    prompt = MULTIPLE * prefix1 + prefix2 + ("Hello, my name is")

    for APC in [True, False]:
        engine = LLM(model=MODEL, enable_prefix_caching=APC, 
            gpu_memory_utilization=0.4, disable_log_stats=False)
        # Initial prompt
        outputs = engine.generate(prefix + prompt, sampling_params)
        # Measure diverging prompts:
        start_time = time.time()
        for i in range(DIVERGENT_PROMPTS):
            divergence = (' ' + string.ascii_letters[i] + ' ')
            outputs = engine.generate(prefix + divergence + prompt, sampling_params)
            # print(f"Generated text: {outputs[0].outputs[0].text!r}")
        total_time = time.time() - start_time
        # Summary
        print('Execution with APC:', APC, "took --- %s seconds ---" % total_time)
        for m in engine.llm_engine.get_metrics():
            if 'vllm:prompt_tokens_cached' in m.name:
                print(m.name, m.value)
        del engine
        cleanup_dist_env_and_memory()

Test Result

Main (Qwen/Qwen3.5-0.8B):

Execution with APC: True took --- 3.2461838722229004 seconds ---
vllm:prompt_tokens_cached 0
Execution with APC: False took --- 3.215266704559326 seconds ---
vllm:prompt_tokens_cached 0

This PR (Qwen/Qwen3.5-0.8B):

Execution with APC: True took --- 2.3055942058563232 seconds ---
vllm:prompt_tokens_cached 289408
Execution with APC: False took --- 3.2060539722442627 seconds ---
vllm:prompt_tokens_cached 0

This PR (for Qwen/Qwen3.5-35B-A3B):

Execution with APC: True took --- 5.365646839141846 seconds ---
vllm:prompt_tokens_cached 280896
Execution with APC: False took --- 8.88869309425354 seconds ---
vllm:prompt_tokens_cached 0

[gemini-code-assist]

Code Review

This pull request introduces a Marconi-style admission policy for the hybrid cache, which is a valuable performance optimization. The implementation correctly identifies shared prefixes by comparing cache hits between different attention mechanisms and forces caching at divergence points. However, I've identified a few areas where the code could be made more robust and maintainable. Specifically, there's a potential for a critical UnboundLocalError, some confusing variable reuse that hinders readability, and a piece of code with an uncertain assertion that could lead to subtle bugs. Addressing these points will improve the overall quality and stability of this new feature.

[tdoublep]

cc @peakcrosser7 @heheda12345

This is how we can catch system prompts using align mode

[mergify]

Hi @s3woz, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[tdoublep]

LGTM but would like @heheda12345 to also review

[s3woz]

Re: @QilaiZhang :

@s3woz @heheda12345 Hi! Is there any progress on this PR?

I've just fixed a merge conflict. Otherwise, waiting for feedback / reviews. FYI: @tdoublep

[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

[heheda12345]

LGTM! Thanks for this improvement. Please simplify the comments to only include necessary ones.

[mergify]

Hi @s3woz, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.