[Router][Bugfix] prefixaware: accept token-id prompts

[clateral912]

Summary

--routing-logic prefixaware returns a 500 for any /v1/completions request that sends token ids instead of text.

The OpenAI completions API accepts prompt as str, list[str], list[int] or list[list[int]]. PrefixAwareRouter.route_request passes the raw value to HashTrie, which slices it by characters and hashes each slice:

# src/vllm_router/prefix/hashtrie.py
for i in range(0, len(request), self.chunk_size):
    yield xxhash.xxh64(request[i : i + self.chunk_size]).intdigest()

For a list, request[i:i+chunk_size] is a sublist, and xxhash raises TypeError: unicode/bytes-like object expected. The exception propagates out of routing and the request fails.

Reproduce on main

vllm-router --routing-logic prefixaware --service-discovery static \
            --static-backends http://127.0.0.1:8000 --static-models <model>

curl -s http://127.0.0.1:8000/v1/completions \
     -H 'Content-Type: application/json' \
     -d '{"model":"<model>","prompt":[1,2,3,4],"max_tokens":1}'
# -> 500; router log shows TypeError at hashtrie.py:57

The same request succeeds with --routing-logic roundrobin, so it is specific to the prefix-aware path.

Fix

Normalize the prompt to a string before it reaches the trie (_prompt_to_trie_key). The routing algorithm is unchanged — prompts sharing a token prefix still share a character prefix, so prefix affinity behaves as before, only the alphabet differs. list[str] is joined rather than str()-ified so existing text batches keep matching exactly as they do today.

There is precedent for this normalization: SGLang's model gateway renders token-id input to a space-joined decimal string for its own routing decisions on the native /generate endpoint (openai-protocol, generate.rs, extract_text_for_routing). Hashing a rendering of the ids is not as good as hashing the ids themselves, but it preserves prefix structure, which is all the character-chunked trie needs.

Test plan

pytest src/tests/test_prefixaware_router.py -q

Six new cases. Five of them fail on main with the TypeError above:

case	on main	with fix
token-ids ([1,2,3])	TypeError	pass
token-id-batch ([[1,2],[3]])	TypeError	pass
string-list (["a","b"])	TypeError	pass
empty-list ([])	TypeError	pass
string (unchanged path)	pass	pass
prefix affinity over real HashTrie	TypeError	pass

The last one is not a mock: it drives the real HashTrie with a token-id prompt and then a continuation of it, asserting the continuation routes back to the endpoint that served the prefix. That is what guards the "affinity still works after normalization" claim rather than just "it no longer crashes".

Test result

$ pytest src/tests/test_prefixaware_router.py -q
..........                    [100%]
10 passed

$ pytest src/tests/ -q --ignore=src/tests/perftest --ignore=src/tests/external_providers
153 passed

black, isort, ruff and codespell clean on both changed files.

[gemini-code-assist]

Code Review

This pull request introduces prompt normalization in the prefix-aware router to handle non-string prompt types (such as token IDs and batches) and prevent type errors during trie lookup. The review feedback highlights two key improvements: first, adding namespace prefixes to normalized keys to prevent collisions between different prompt types (e.g., list of token IDs vs. raw strings) which would otherwise cause KV cache misses; and second, using safe dictionary access (.get()) when retrieving the prompt to avoid potential KeyError exceptions when the prompt is omitted.

[gemini-code-assist]

Without namespaces or prefixes, different prompt types can collide and produce identical trie keys. For example, a token ID prompt like [1, 2, 3] and a string prompt like "1,2,3" would both normalize to "1,2,3". Similarly, a list of strings ["a", "b"] and a single string "a b" would both normalize to "a b".

When these collisions occur, the router will incorrectly assume they share a prefix and route them to the same backend engine to reuse the KV cache. However, because the actual tokens are completely different, this will result in a KV cache miss on the backend, defeating the purpose of prefix-aware routing.

To prevent this, we should prefix non-string keys with a namespace indicator (e.g., token_ids:, token_ids_batch:, str_batch:) while keeping plain string prompts as-is to maintain backward compatibility.

def _prompt_to_trie_key(prompt) -> str:
    """Normalize a `/v1/completions` prompt into a string for `HashTrie`.

    The OpenAI completions API accepts token ids as well as text: `prompt` may
    be `str`, `list[str]`, `list[int]`, or `list[list[int]]`. `HashTrie` slices
    the value by characters and feeds each slice to `xxhash`, so anything that
    is not a string raises `TypeError` and fails the request.

    Rendering token ids to a stable decimal string keeps the trie's algorithm
    unchanged -- prompts sharing a token prefix still share a character prefix,
    so prefix affinity works as intended, only the alphabet differs.
    """
    if isinstance(prompt, str):
        return prompt
    if isinstance(prompt, list):
        if not prompt:
            return ""
        # list[list[int]]: a batch of token-id sequences.
        if isinstance(prompt[0], list):
            return "token_ids_batch:" + " ".join(",".join(str(t) for t in seq) for seq in prompt)
        # list[str]: already text.
        if isinstance(prompt[0], str):
            return "str_batch:" + " ".join(prompt)
        # list[int]: token ids.
        return "token_ids:" + ",".join(str(t) for t in prompt)
    return str(prompt)

[gemini-code-assist]

In the OpenAI completions API, the prompt field is optional and can be omitted. Accessing request_json["prompt"] directly will raise a KeyError if the prompt is missing, leading to a 500 Internal Server Error. Using .get("prompt", "") provides a safe fallback and aligns with how prompts are extracted elsewhere in this file.

        else:
            # Handle regular completions
            prompt = _prompt_to_trie_key(request_json.get("prompt", ""))

[gemini-code-assist]

Update the test expectations to match the new prefixed trie keys to ensure namespaces are correctly validated.

Suggested change

	@pytest.mark.parametrize(
	"prompt, expected_key",
	[
	pytest.param([1, 2, 3], "1,2,3", id="token-ids"),
	pytest.param([[1, 2], [3]], "1,2 3", id="token-id-batch"),
	pytest.param(["a", "b"], "a b", id="string-list"),
	pytest.param("plain text", "plain text", id="string"),
	pytest.param([], "", id="empty-list"),
	],
	)
	@pytest.mark.parametrize(
	"prompt, expected_key",
	[
	pytest.param([1, 2, 3], "token_ids:1,2,3", id="token-ids"),
	pytest.param([[1, 2], [3]], "token_ids_batch:1,2 3", id="token-id-batch"),
	pytest.param(["a", "b"], "str_batch:a b", id="string-list"),
	pytest.param("plain text", "plain text", id="string"),
	pytest.param([], "", id="empty-list"),
	],
	)