feat(sglang): OpenAI embeddings dimensions (DIS-2093) by tzulingk · Pull Request #9722 · ai-dynamo/dynamo

tzulingk · 2026-05-19T01:07:17Z

Overview:

Closes one of the two OpenAI /v1/embeddings spec gaps tracked in DIS-2093: the dimensions field (Matryoshka truncation).

Scope: This PR only changes the SGLang embedding handler. The vLLM backend doesn't have an embedding worker yet (that's PR #9713), so there's nothing on the vLLM side to update here. This PR can be reviewed and merged on its own — it doesn't depend on #9713, and #9713 doesn't depend on it. Once both land, the vLLM handler should pick up the same _transform_response logic (and at that point a shared helper may be worth lifting out).

The other field in DIS-2093 — encoding_format=base64 — was originally in scope but verified on dynamo-aks-dev to not work end-to-end because the Rust frontend's response type rejects a base64 string. It's been split out to a separate follow-up; see the "Deferred" section below.

Details:

dimensions (Matryoshka truncation): The field was declared on EmbeddingRequest and accepted by the Rust frontend, but the SGLang handler ignored it. Now _transform_response slices each embedding to the first N floats. Validates 1 <= dimensions <= model_hidden_dim; otherwise raises ValueError which the frontend surfaces as HTTP 400.

Note on Matryoshka correctness: truncation is semantically meaningful only for models trained with Matryoshka representation learning (Qwen3-Embedding, BGE-M3, OpenAI text-embedding-3-*). On non-Matryoshka models the leading slice is still a valid vector but loses quality. Per the OpenAI spec we honor the request regardless; downstream caveat for users.

Test plan:

One new payload added to tests/serve/test_sglang.py's embedding_agg config:

dimensions=128 → assert vector length is exactly 128 (passes on Qwen3-Embedding-4B which has hidden dim 2560).

The existing default + single-string + 3-string-list payloads also continue to pass.

embedding_payload and embedding_payload_default gained an extra_body kwarg in tests/utils/payload_builder.py so future callers can attach arbitrary OpenAI fields without introducing new helpers per field.

Deferred — `encoding_format=base64`

Implementation attempt was reverted from this PR. Verified on dynamo-aks-dev with an A100:

The Python handler can pack a (possibly-truncated) float list as little-endian float32 + base64-encode it.
But the Rust frontend's NvCreateEmbeddingResponse::from_annotated_stream aggregator deserializes data[].embedding as Vec<f32> (inherited from the upstream async_openai::types::embeddings::* re-export). A base64 string is rejected with:

invalid type: string "AAACOQAA...", expected a sequence

End-to-end encoding_format=base64 support requires owning the embedding response type in lib/protocols/ (per the rubric in lib/protocols/CLAUDE.md) and updating the aggregator. That's a sizable Rust-side change orthogonal to the SGLang handler change. Tracked as DIS-2099 — the Python implementation reverted from this PR is preserved in its git history so the follow-up doesn't start from scratch.

Where should the reviewer start?

components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py — _transform_response() now accepts an optional dimensions arg and slices each embedding accordingly.
components/src/dynamo/sglang/protocol.py — the field on EmbeddingRequest was already declared in main; only the handler change is functional.
tests/utils/payload_builder.py — extra_body parameter is a small reusable surface for OpenAI-field-passthrough tests.
tests/serve/test_sglang.py — the new dimensions=128 payload.

Closes the dimensions part of DIS-2093.

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

New Features
- Embedding responses can now be truncated to specified dimensions, enabling flexible vector size control.
Tests
- Added test coverage for dimension truncation in embedding requests.

coderabbitai · 2026-05-19T16:59:34Z

Walkthrough

This PR adds Matryoshka-style dimension truncation to embedding responses. The handler now accepts an optional dimensions parameter, validates it, and slices embeddings accordingly. Test infrastructure is updated to support passing custom request fields, with an integration test validating the new truncation behavior.

Changes

Embedding Dimension Truncation

Layer / File(s)	Summary
Handler dimension truncation implementation `components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py`	Import `List` type; `generate` extracts `embedding_request.dimensions` and passes it to `_transform_response`; `_transform_response` validates `dimensions >= 1`, ensures it does not exceed embedding length, slices embeddings to requested dimensions, and returns OpenAI-format response.
Test payload builders and integration test `tests/utils/payload_builder.py`, `tests/serve/test_sglang.py`	`embedding_payload_default` and `embedding_payload` now accept optional `extra_body` parameter and merge it into the request body; integration test adds a case that sets `extra_body={"dimensions": 128}` and validates the response reports 128 dimensions.

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main feature added (OpenAI embeddings dimensions) and references the tracking issue (DIS-2093), matching the core implementation work in the PR.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description follows the template structure with Overview, Details, Test plan, and Where to review sections. All required information is present and well-documented.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Mirrors what PR #9722 did for SGLang's embedding handler: read the optional `dimensions` field from the request and slice each embedding to the first N floats before responding. Validates `1 <= dimensions <= model_hidden_dim`; otherwise raises `ValueError` which the Rust frontend surfaces as HTTP 400. Matches the OpenAI hosted API's behavior of rejecting out-of-range `dimensions` values rather than silently zero-padding or clamping. Test plan additions in `tests/serve/test_vllm.py::embedding_agg`: - New payload with `dimensions=128` against Qwen3-Embedding-0.6B (hidden dim 1024) → response must report dimension exactly 128. `EmbeddingPayload` is constructed inline here because the `extra_body` kwarg on `embedding_payload()` lives in PR #9722 and isn't in this branch's base. Once both PRs land we can simplify this test to use `embedding_payload(..., extra_body={"dimensions": 128})`. Stacks on: #9713 (vLLM embedding worker MVP — adds the handler being modified here). Refs DIS-2093, parallel to #9722 for SGLang. Note: `encoding_format=base64` is NOT included. Same reason as the SGLang side: the Rust frontend's response type is `Vec<f32>` via the upstream `async_openai::types::embeddings` re-export and rejects base64 strings. Tracked in DIS-2099 (requires owning the embedding response type in lib/protocols/). Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

Move the lower-bound `dimensions >= 1` check out of `_transform_response` into `generate`, before `engine.async_encode` is called. An obviously invalid request (e.g. `dimensions=0`) now fails as HTTP 400 without spending a full pooling forward pass on the GPU. The upper-bound check (`dimensions <= len(embedding)`) stays in `_transform_response` because it needs to compare against the actual embedding length returned by SGLang. Addresses dynamo-ops review on PR #9722. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

Mirrors what PR #9722 did for SGLang's embedding handler: read the optional `dimensions` field from the request and slice each embedding to the first N floats before responding. Validates `1 <= dimensions <= model_hidden_dim`; otherwise raises `ValueError` which the Rust frontend surfaces as HTTP 400. Matches the OpenAI hosted API's behavior of rejecting out-of-range `dimensions` values rather than silently zero-padding or clamping. Test plan additions in `tests/serve/test_vllm.py::embedding_agg`: - New payload with `dimensions=128` against Qwen3-Embedding-0.6B (hidden dim 1024) → response must report dimension exactly 128. `EmbeddingPayload` is constructed inline here because the `extra_body` kwarg on `embedding_payload()` lives in PR #9722 and isn't in this branch's base. Once both PRs land we can simplify this test to use `embedding_payload(..., extra_body={"dimensions": 128})`. Stacks on: #9713 (vLLM embedding worker MVP — adds the handler being modified here). Refs DIS-2093, parallel to #9722 for SGLang. Note: `encoding_format=base64` is NOT included. Same reason as the SGLang side: the Rust frontend's response type is `Vec<f32>` via the upstream `async_openai::types::embeddings` re-export and rejects base64 strings. Tracked in DIS-2099 (requires owning the embedding response type in lib/protocols/). Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

Adds OpenAI `/v1/embeddings` `dimensions` (Matryoshka truncation) support to the vLLM embedding worker. Mirrors PR #9722 which did the same for SGLang. In `EmbeddingWorkerHandler.generate`, the handler now reads `request.get("dimensions")` once before the per-input loop, validates it (must be a positive int), then per embedding asserts `dimensions <= len(embedding)` and slices to `embedding[:dimensions]`. Out-of-range values raise `ValueError` which the Rust frontend surfaces as HTTP 400. `encoding_format=base64` is not yet supported end-to-end because the Rust frontend's response type is `Vec<f32>` and rejects base64 strings; that requires owning the embedding response type in `lib/protocols/` and is tracked as a separate follow-up. Test: `embedding_agg` config in `tests/serve/test_vllm.py` gains a fourth payload — `dimensions=128` against Qwen3-Embedding-0.6B (hidden dim 1024). Built inline using `EmbeddingPayload(...)` directly because the `embedding_payload(..., extra_body=...)` helper added in PR #9722 isn't on this branch's base. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

biswapanda

lgtm

dynamo-ops

Previous review comments have been addressed. Approving.

Closes the `dimensions` field gap in the SGLang embedding handler. The field was declared on EmbeddingRequest and accepted by the Rust frontend, but the handler ignored it. The handler now: - Validates `dimensions >= 1` upfront in `generate`, before `engine.async_encode` is called, so an invalid request fails as HTTP 400 without spending a pooling forward pass on the GPU. - Per-embedding asserts `dimensions <= len(embedding)` in `_transform_response`, then slices each embedding to the first N floats. Out-of-range values raise `ValueError` which the frontend surfaces as HTTP 400. Matryoshka caveat: truncation is semantically meaningful only for models trained with Matryoshka representation learning (Qwen3-Embedding, BGE-M3, OpenAI text-embedding-3-*). On non-Matryoshka models the leading slice is still a valid vector but loses quality. Per the OpenAI spec we honor the request regardless. Not included — `encoding_format=base64` ======================================== The Rust frontend's `NvCreateEmbeddingResponse::from_annotated_stream` aggregator deserializes `data[].embedding` as `Vec<f32>` (inherited from the upstream `async_openai::types::embeddings::*` re-export). A base64 string is rejected with: invalid type: string "AAACOQAA...", expected a sequence End-to-end `encoding_format=base64` support requires owning the embedding response type in `lib/protocols/` (per the rubric in `lib/protocols/CLAUDE.md`) and updating the aggregator. Tracked as a separate follow-up. Tests ===== One new payload added to tests/serve/test_sglang.py's `embedding_agg` config: `dimensions=128` against Qwen3-Embedding-4B (hidden dim 2560) — asserts the returned vector length is exactly 128. `embedding_payload` and `embedding_payload_default` gained an `extra_body` kwarg in tests/utils/payload_builder.py so future callers can attach arbitrary OpenAI fields without introducing new helpers. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

Adds OpenAI `/v1/embeddings` `dimensions` (Matryoshka truncation) support to the vLLM embedding worker. Mirrors PR #9722 which did the same for SGLang. In `EmbeddingWorkerHandler.generate`, the handler now reads `request.get("dimensions")` once before the per-input loop, validates it (must be a positive int), then per embedding asserts `dimensions <= len(embedding)` and slices to `embedding[:dimensions]`. Out-of-range values raise `ValueError` which the Rust frontend surfaces as HTTP 400. `encoding_format=base64` is not yet supported end-to-end because the Rust frontend's response type is `Vec<f32>` and rejects base64 strings; that requires owning the embedding response type in `lib/protocols/` and is tracked as a separate follow-up. Test: `embedding_agg` config in `tests/serve/test_vllm.py` gains a fourth payload — `dimensions=128` against Qwen3-Embedding-0.6B (hidden dim 1024). Built inline using `EmbeddingPayload(...)` directly because the `embedding_payload(..., extra_body=...)` helper added in PR #9722 isn't on this branch's base. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

dynamo-ops

Previous review comments have been addressed. Approving.

Adds OpenAI `/v1/embeddings` `dimensions` (Matryoshka truncation) support to the vLLM embedding worker. Mirrors PR #9722 which did the same for SGLang. In `EmbeddingWorkerHandler.generate`, the handler now reads `request.get("dimensions")` once before the per-input loop, validates it (must be a positive int), then per embedding asserts `dimensions <= len(embedding)` and slices to `embedding[:dimensions]`. Out-of-range values raise `ValueError` which the Rust frontend surfaces as HTTP 400. `encoding_format=base64` is not yet supported end-to-end because the Rust frontend's response type is `Vec<f32>` and rejects base64 strings; that requires owning the embedding response type in `lib/protocols/` and is tracked as a separate follow-up. Test: `embedding_agg` config in `tests/serve/test_vllm.py` gains a fourth payload — `dimensions=128` against Qwen3-Embedding-0.6B (hidden dim 1024). Built inline using `EmbeddingPayload(...)` directly because the `embedding_payload(..., extra_body=...)` helper added in PR #9722 isn't on this branch's base. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com> Signed-off-by: Kirill Emelianov <kirill.emelianov@gcore.com>

pull-request-size Bot added the size/L label May 19, 2026

github-actions Bot added feat backend::sglang Relates to the sglang backend labels May 19, 2026

pull-request-size Bot added size/M and removed size/L labels May 19, 2026

copy-pr-bot Bot temporarily deployed to GITLAB May 19, 2026 04:17 Inactive

copy-pr-bot Bot temporarily deployed to GITLAB May 19, 2026 07:16 Inactive

tzulingk marked this pull request as ready for review May 19, 2026 16:56

tzulingk requested review from a team as code owners May 19, 2026 16:56

tzulingk requested review from biswapanda and nnshah1 May 19, 2026 16:58

tzulingk changed the title ~~feat(sglang): OpenAI embeddings dimensions + encoding_format=base64 (DIS-2093)~~ feat(sglang): OpenAI embeddings dimensions (DIS-2093) May 19, 2026

tzulingk mentioned this pull request May 19, 2026

feat(vllm): OpenAI embeddings dimensions truncation (DIS-2093 vLLM side) #9751

Merged

dynamo-ops reviewed May 19, 2026

View reviewed changes

Comment thread components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py Outdated

tzulingk force-pushed the feat/embedding-spec-gaps branch from d5c9747 to ffbc7c9 Compare May 19, 2026 19:45

copy-pr-bot Bot temporarily deployed to GITLAB May 19, 2026 19:45 Inactive

tzulingk force-pushed the feat/embedding-spec-gaps branch from ffbc7c9 to 0ccfc69 Compare May 19, 2026 20:55

copy-pr-bot Bot temporarily deployed to GITLAB May 19, 2026 20:55 Inactive

copy-pr-bot Bot temporarily deployed to GITLAB May 19, 2026 22:47 Inactive

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 15:05 Inactive

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 15:11 Inactive

tzulingk force-pushed the feat/embedding-spec-gaps branch from 86f78a7 to deb74ad Compare May 20, 2026 15:25

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 15:25 Inactive

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 15:59 Inactive

biswapanda approved these changes May 20, 2026

View reviewed changes

dynamo-ops approved these changes May 20, 2026

View reviewed changes

tzulingk force-pushed the feat/embedding-spec-gaps branch from deb74ad to 10c0f9f Compare May 20, 2026 20:50

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 20:50 Inactive

dynamo-ops approved these changes May 20, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to GITLAB May 20, 2026 22:00 Inactive

tzulingk merged commit 92f5b3b into main May 21, 2026
94 checks passed

tzulingk deleted the feat/embedding-spec-gaps branch May 21, 2026 00:43

tzulingk mentioned this pull request May 22, 2026

feat(embeddings): honor OpenAI encoding_format=base64 end-to-end #9887

Merged

kirillemilio pushed a commit to kirillemilio/dynamo that referenced this pull request May 25, 2026

feat(sglang): OpenAI embeddings dimensions (DIS-2093) (ai-dynamo#9722)

769eed6

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com> Signed-off-by: Kirill Emelianov <kirill.emelianov@gcore.com>

Navigation Menu

Uh oh!

feat(sglang): OpenAI embeddings dimensions (DIS-2093) - #9722

feat(sglang): OpenAI embeddings dimensions (DIS-2093)#9722
tzulingk merged 1 commit into
mainfrom
feat/embedding-spec-gaps

tzulingk commented May 19, 2026 •

edited by itay

Loading

Uh oh!

coderabbitai Bot commented May 19, 2026 •

edited

Loading

❌ Failed checks (1 warning)

Uh oh!

Uh oh!

biswapanda left a comment

Uh oh!

dynamo-ops left a comment

Uh oh!

dynamo-ops left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

tzulingk commented May 19, 2026 • edited by itay Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview:

Details:

Test plan:

Deferred — encoding_format=base64

Where should the reviewer start?

Summary by CodeRabbit

Release Notes

Uh oh!

coderabbitai Bot commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

Uh oh!

biswapanda left a comment

Choose a reason for hiding this comment

Uh oh!

dynamo-ops left a comment

Choose a reason for hiding this comment

Uh oh!

dynamo-ops left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tzulingk commented May 19, 2026 •

edited by itay

Loading

Deferred — `encoding_format=base64`

coderabbitai Bot commented May 19, 2026 •

edited

Loading