feat(vllm): route embedding worker through Rust preprocessor (opt-in via --use-vllm-tokenizer)

[tzulingk]

Overview:

Route the vLLM embedding worker through Dynamo's Rust OpenAIPreprocessor (tokenize off the engine process) instead of only handing raw text to vLLM — selectable at launch via the existing --use-vllm-tokenizer flag, exactly like chat/completions. Motivation: on text-input embeddings, vLLM tokenizes in-engine; moving tokenization to the Rust frontend is measurably faster (early A/B below).

Details:

EmbeddingWorkerHandler.generate now dispatches on use_vllm_tokenizer:

• Text mode (--use-vllm-tokenizer, baseline): unchanged — raw {input} in, OpenAI {data:[…base64…]} out.

• Tokens mode (default now): reads {token_ids} from the PreprocessedEmbeddingRequest, wraps each in TokensPrompt (so vLLM skips its own tokenizer), and returns the raw EmbeddingsEngineOutput shape so the Rust postprocessor does OpenAI formatting/base64.

• worker_factory.py — embedding registration picks ModelInput.Text if config.use_vllm_tokenizer else ModelInput.Tokens (was hardcoded Text); threads the flag into the handler and the health-check payload.

• handlers.py — extracted a shared _run_encode plus _parse_embedding_dimensions / _check_embedding_dimensions so both paths validate + encode identically (incl. the task="embed" pooling + Matryoshka dimensions handling from fix(vllm): pool embedding worker output via PoolingParams(task="embed") #10248).

• health_check.py — VllmEmbeddingHealthCheckPayload gains use_text_input; probes with {input:"probe"} for Text mode and {token_ids:[[1]]} for Tokens mode (otherwise the probe fails with a missing-field error).

• Tests — embedding registration-contract test (flag → ModelInput), tokens-mode handler tests, and updated health-check payload tests for the new default; existing text-path tests pinned to use_vllm_tokenizer=True.

Behavior change: the default is now the Tokens/Rust path, so any embedding model whose tokenizer the frontend can't resolve must be launched with --use-vllm-tokenizer.

Early A/B (single GB200, vLLM 0.24.0, vllm bench serve --backend openai-embeddings, text input, 4096 prompts, Qwen/Qwen3-Embedding-0.6B — embeddinggemma-300m is gated, run pending an HF token): candidate (Rust) vs baseline (--use-vllm-tokenizer) throughput = +27% / +46% / +13% at max-concurrency 128 / 512 / 1024. Tracking in DIS-2307.

Where should the reviewer start?

• components/src/dynamo/vllm/handlers.py — EmbeddingWorkerHandler.generate dispatch, the new _generate_tokens_mode, and the shared _run_encode.
• components/src/dynamo/vllm/worker_factory.py — the ModelInput.Text/Tokens selection driven by use_vllm_tokenizer.

🤖 Generated with Claude Code

[datadog-official]

⚠️ Warnings

🚦 7 Pipeline jobs failed

Docs link check | lychee     

PR | backend-status-check     

PR | vllm-runtime / Multi-GPU Test cuda13.0, amd64     

View all 7 failed jobs.

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 0b55132 | Docs | Give us feedback!}

[github-actions]

This PR is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.