[Bugfix][Pooling] Preserve cache_salt for chunked embeddings by Sunt-ing · Pull Request #47696 · vllm-project/vllm

Sunt-ing · 2026-07-06T06:20:28Z

Purpose

Chunked embedding preprocessing dropped request-level cache_salt when it rebuilt long inputs into per-chunk token prompts. With prefix caching enabled, a second /v1/embeddings request using a different salt could still reuse the first request's cached chunks.

This preserves the original salt when creating chunked embedding engine inputs. Unsalted requests keep the existing behavior.

Test Plan

Run a real /v1/embeddings server with prefix caching and long-text chunked processing, then send the same 8102-token embedding input with cache_salt=tenant-a, cache_salt=tenant-b, and cache_salt=tenant-b again. The second request must not report prefix-cache hits; the third request should report hits as the same-salt positive control.

Run scoped diff and ruff checks for the changed file.

Test Result

On upstream main, the second request used a different salt but still reported 7856 prefix-cache hits. With this patch, that second request reports 0 hits and the same-salt repeat still reports 7856 hits. Reverting the patch brings the 7856-hit failure back.

server repro command and client script

The output below was captured on upstream main 394edc81082c04cb38e3455fc0c03d62ba403b72 with one RTX 4090 24 GB GPU. I used a local snapshot of Qwen/Qwen3-Embedding-0.6B; the command below uses the public model id.

Server command used from a clean checkout:

VLLM_USE_FLASHINFER_SAMPLER=0 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
python -m vllm.entrypoints.openai.api_server \
  --host 127.0.0.1 \
  --port 8025 \
  --model Qwen/Qwen3-Embedding-0.6B \
  --served-model-name qwen3-embed \
  --runner pooling \
  --dtype bfloat16 \
  --enforce-eager \
  --enable-prefix-caching \
  --max-model-len 512 \
  --gpu-memory-utilization 0.45 \
  --pooler-config '{"pooling_type":"LAST","use_activation":true,"enable_chunked_processing":true,"max_embed_len":10000}'

Client script:

import json
import re
import sys
import time
import urllib.error
import urllib.request

BASE_URL = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:8025"
MODEL = sys.argv[2] if len(sys.argv) > 2 else "qwen3-embed"


def request_json(method: str, path: str, payload: dict | None = None):
    data = None
    headers = {}
    if payload is not None:
        data = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(
        f"{BASE_URL}{path}", data=data, headers=headers, method=method
    )
    with urllib.request.urlopen(req, timeout=90) as resp:
        body = resp.read()
        if not body:
            return None
        return json.loads(body.decode("utf-8"))


def request_text(path: str) -> str:
    with urllib.request.urlopen(f"{BASE_URL}{path}", timeout=30) as resp:
        return resp.read().decode("utf-8")


def wait_ready() -> None:
    deadline = time.time() + 420
    last_error = None
    while time.time() < deadline:
        try:
            request_json("GET", "/v1/models")
            return
        except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
            last_error = exc
            time.sleep(2)
    raise SystemExit(f"server did not become ready: {last_error}")


def metric_value(metrics: str, metric: str) -> float:
    pattern = re.compile(
        rf"^{re.escape(metric)}(?:_total)?(?:\{{[^}}]*\}})?\s+([0-9.]+)$"
    )
    total = 0.0
    for line in metrics.splitlines():
        match = pattern.match(line)
        if match:
            total += float(match.group(1))
    return total


def read_counters() -> dict[str, float]:
    metrics = request_text("/metrics")
    return {
        "queries": metric_value(metrics, "vllm:prefix_cache_queries"),
        "hits": metric_value(metrics, "vllm:prefix_cache_hits"),
        "cached_tokens": metric_value(metrics, "vllm:prompt_tokens_cached"),
    }


def send_embedding(cache_salt: str) -> None:
    token_ids = list(range(1000, 9102))
    response = request_json(
        "POST",
        "/v1/embeddings",
        {
            "model": MODEL,
            "input": token_ids,
            "encoding_format": "float",
            "cache_salt": cache_salt,
        },
    )
    print(
        "response",
        cache_salt,
        "status=ok",
        "items",
        len(response.get("data", [])),
        "usage",
        response.get("usage", {}),
        flush=True,
    )


def main() -> None:
    wait_ready()
    labels = ["baseline"]
    counters = [read_counters()]

    for salt in ("tenant-a", "tenant-b", "tenant-b"):
        send_embedding(salt)
        time.sleep(2)
        labels.append(salt)
        counters.append(read_counters())

    for label, counter in zip(labels, counters):
        print("metrics", label, json.dumps(counter, sort_keys=True), flush=True)

    deltas = []
    for before, after in zip(counters, counters[1:]):
        deltas.append(
            {
                key: after[key] - before[key]
                for key in ("queries", "hits", "cached_tokens")
            }
        )
    for label, delta in zip(("tenant-a", "tenant-b", "tenant-b-repeat"), deltas):
        print("delta", label, json.dumps(delta, sort_keys=True), flush=True)

    second_delta = deltas[1]
    if second_delta["hits"] > 0 or second_delta["cached_tokens"] > 0:
        raise SystemExit(
            "BUG: second request used a different cache_salt but reported "
            f"prefix-cache reuse: {json.dumps(second_delta, sort_keys=True)}"
        )

    repeat_delta = deltas[2]
    if repeat_delta["hits"] <= 0 and repeat_delta["cached_tokens"] <= 0:
        raise SystemExit(
            "inconclusive: repeated same-salt request did not show prefix-cache "
            f"reuse either: {json.dumps(repeat_delta, sort_keys=True)}"
        )

    print("OK: different cache_salt avoided prefix-cache reuse", flush=True)


if __name__ == "__main__":
    main()

before, after, and revert output

Upstream main:

response tenant-a status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
metrics baseline {"cached_tokens": 0.0, "hits": 0.0, "queries": 0.0}
metrics tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
metrics tenant-b {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 16204.0}
metrics tenant-b {"cached_tokens": 15712.0, "hits": 15712.0, "queries": 24306.0}
delta tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
delta tenant-b {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}
delta tenant-b-repeat {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}
BUG: second request used a different cache_salt but reported prefix-cache reuse: {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}

With this patch:

response tenant-a status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
metrics baseline {"cached_tokens": 0.0, "hits": 0.0, "queries": 0.0}
metrics tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
metrics tenant-b {"cached_tokens": 0.0, "hits": 0.0, "queries": 16204.0}
metrics tenant-b {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 24306.0}
delta tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
delta tenant-b {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
delta tenant-b-repeat {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}
OK: different cache_salt avoided prefix-cache reuse

After reverting this patch:

response tenant-a status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
response tenant-b status=ok items 1 usage {'prompt_tokens': 8102, 'total_tokens': 8102, 'completion_tokens': 0, 'prompt_tokens_details': None}
BUG: second request used a different cache_salt but reported prefix-cache reuse: {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}
metrics baseline {"cached_tokens": 0.0, "hits": 0.0, "queries": 0.0}
metrics tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
metrics tenant-b {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 16204.0}
metrics tenant-b {"cached_tokens": 15712.0, "hits": 15712.0, "queries": 24306.0}
delta tenant-a {"cached_tokens": 0.0, "hits": 0.0, "queries": 8102.0}
delta tenant-b {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}
delta tenant-b-repeat {"cached_tokens": 7856.0, "hits": 7856.0, "queries": 8102.0}

Scoped checks:

git diff --check

python -m ruff check vllm/entrypoints/pooling/embed/io_processor.py
All checks passed!

python -m ruff format --check vllm/entrypoints/pooling/embed/io_processor.py
1 file already formatted

AI assistance was used to prepare this PR.

cc @noooop

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mergify · 2026-07-22T13:05:51Z

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

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

Signed-off-by: Ting Sun <suntcrick@gmail.com>

Sunt-ing · 2026-10-08T06:03:13Z

@noooop @DarkLight1337 PTAL?

Signed-off-by: Ting Sun <suntcrick@gmail.com>

noooop

thanks!

noooop · 2026-10-08T06:10:49Z

/ci run

github-actions · 2026-10-08T06:11:09Z

✅ Triggered Buildkite CI #93526 for commit 43fcd8a6e1f8.

vllm-agent · 2026-10-08T06:35:31Z

CI selector (shadow): 5 test steps (8 jobs) instead of 39 (55 jobs)

Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works.

Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.

steps (jobs)	Today's rules	Selector	Would skip	Would add
NVIDIA, CPU and others	39 (55)	5 (8)	36 (49)	2 (2)
AMD mirrors	38 (49)	5 (8)	35 (43)	2 (2)

Selector would run (5)

entrypoints-integration-api-server ×4
entrypoints-integration-pooling
entrypoints-unit-tests
pipeline-context-parallelism-4-gpus
pytorch-fullgraph-test

Would skip (today's rules run them) (36)

basic-correctness ×2
basic-correctness-cpu-offload
basic-correctness-cumem
basic-correctness-prefetch-offload
basic-correctness-sleep-mode
basic-models-test-other-cpu
basic-models-tests-initialization
basic-models-tests-other
benchmarks-cli-test
cpu-language-generation-and-pooling-model-tests ×3
entrypoints-integration-api-server-generate
entrypoints-integration-api-server-openai-chat_completion
entrypoints-integration-api-server-openai-completion
entrypoints-integration-llm
entrypoints-integration-multimodal
entrypoints-integration-responses-api
entrypoints-integration-speech_to_text
examples
kernels-fla-ops-test-b200
kernels-mhc-test-b200
kernels-root-misc-test-b200
language-models-tests-granite-l4-compatibility
language-models-tests-hybrid ×2
language-models-tests-standard
multi-modal-models-standard-1-qwen2
multi-modal-models-standard-2-qwen3-gemma
multi-modal-models-standard-3-llava-qwen2-vl
multi-modal-models-standard-4-other-whisper
multi-modal-processor ×4
multi-modal-processor-cpu ×4
rl-entrypoints-tests
v1-core
v1-executor-worker
v1-kv-connectors ×4
v1-kv-offload
v1-metrics-lmeval

Would add (today's rules do not run them) (2)

pipeline-context-parallelism-4-gpus (Python record)
pytorch-fullgraph-test (Python record)

AMD mirrors: would skip (35)

basic-correctness ×2
basic-correctness-cpu-offload
basic-correctness-cumem
basic-correctness-prefetch-offload
basic-correctness-sleep-mode
basic-models-tests-initialization
basic-models-tests-other
benchmarks-cli-test
entrypoints-integration-api-server-generate
entrypoints-integration-api-server-openai-chat_completion
entrypoints-integration-api-server-openai-completion
entrypoints-integration-llm
entrypoints-integration-multimodal
entrypoints-integration-responses-api
entrypoints-integration-speech_to_text
examples
kernels-fla-ops-test-b200
kernels-mhc-test-b200
kernels-root-misc-test-b200
language-models-tests-granite-l4-compatibility
language-models-tests-hybrid ×2
language-models-tests-standard
multi-modal-models-standard-1-qwen2
multi-modal-models-standard-2-qwen3-gemma
multi-modal-models-standard-3-llava-qwen2-vl
multi-modal-models-standard-4-other-whisper
multi-modal-processor ×4
platform-tests
pytorch-compilation-passes-unit-tests
rl-entrypoints-tests
v1-core
v1-executor-worker
v1-kv-connectors ×4
v1-kv-offload
v1-metrics-lmeval

AMD mirrors: would add (2)

pipeline-context-parallelism-4-gpus (Python record)
pytorch-fullgraph-test (Python record)

_{1 changed files · base 95307ac6d4 · head 43fcd8a6e1 · Python record: build 93039 at 1e5d0ea888 · kernel record: table 43b4aae (build 93244), map 43b4aae · not counted: 10 build steps, 5 A100 steps the generator no longer emits, 2 optional steps the selector would also run}

Sunt-ing requested a review from noooop as a code owner July 6, 2026 06:20

claude Bot reviewed Jul 6, 2026

View reviewed changes

mergify Bot added frontend bug Something isn't working labels Jul 6, 2026

mergify Bot added the needs-rebase label Jul 22, 2026

Tejasvichand mentioned this pull request Aug 21, 2026

[RFC]: A conformance suite for KV-cache key partitioning #53194

Open

1 task done

Tejasvichand mentioned this pull request Sep 8, 2026

[Test] Conformance suite for cache_salt across OpenAI entrypoints #55886

Open

Sunt-ing added 2 commits October 3, 2026 06:31

[Bugfix][Pooling] Preserve cache_salt for chunked embeddings

3d9f3ed

Signed-off-by: Ting Sun <suntcrick@gmail.com>

Test cache salt propagation through chunked embedding inputs

16dce1e

Signed-off-by: Ting Sun <suntcrick@gmail.com>

Sunt-ing force-pushed the 15 branch from fbd835f to 16dce1e Compare October 3, 2026 18:15

mergify Bot added pooling and removed needs-rebase labels Oct 3, 2026

Keep cache salt fix focused on input propagation

6321acd

Signed-off-by: Ting Sun <suntcrick@gmail.com>

noooop approved these changes Oct 8, 2026

View reviewed changes

noooop enabled auto-merge (squash) October 8, 2026 06:08

github-actions Bot added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 8, 2026

Merge branch 'main' into 15

43fcd8a

noooop merged commit 0f112d1 into vllm-project:main Oct 8, 2026
105 checks passed

Sunt-ing deleted the 15 branch October 10, 2026 20:45

Navigation Menu

Repository navigation

[Bugfix][Pooling] Preserve cache_salt for chunked embeddings - #47696

[Bugfix][Pooling] Preserve cache_salt for chunked embeddings#47696
noooop merged 4 commits into
vllm-project:mainfrom
Sunt-ing:15

Sunt-ing commented Jul 6, 2026

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Jul 22, 2026

Uh oh!

Sunt-ing commented Oct 8, 2026

Uh oh!

noooop left a comment

Uh oh!

noooop commented Oct 8, 2026

Uh oh!

github-actions Bot commented Oct 8, 2026

Uh oh!

vllm-agent commented Oct 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Sunt-ing commented Jul 6, 2026

Purpose

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Jul 22, 2026

Uh oh!

Sunt-ing commented Oct 8, 2026

Uh oh!

noooop left a comment

Choose a reason for hiding this comment

Uh oh!

noooop commented Oct 8, 2026

Uh oh!

github-actions Bot commented Oct 8, 2026

Uh oh!

vllm-agent commented Oct 8, 2026

CI selector (shadow): 5 test steps (8 jobs) instead of 39 (55 jobs)

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants