fix(kv-router): fold SGLang wide image pad token ids in KV events by KrishnanPrash · Pull Request #17046 · ai-dynamo/dynamo

KrishnanPrash · 2026-10-21T21:30:17Z

Summary

SGLang main writes four image pad tokens of about 62 bits for each image (sgl-project/sglang#43043). Dynamo reads KV-event token ids as u32, so one wide id fails the whole event batch, and the router also loses the text blocks in that batch. This change reads event token ids as u64 and folds each id above u32 into the existing image pad range, 1_000_000 + (id & (2^30 - 1)). Ids that fit in u32 pass through unchanged, so block hashes for vLLM, TensorRT-LLM and the mocker keep their exact bytes. Repeated images still miss until the request side writes the same pads, which the next PR in this stack does.

Fold, not widen: Token is u32 in every block hash, in KVBM and in the mocker, so a wider Token would change the hashes of every engine to fix one.
One choke point: every ZMQ KV-event path decodes through normalize_token_ids, so the fold sits there and no code after it changes.
Shared helper: fold_token_id sits next to pad_value_for_mm_hash, so the request side can fold with the same function.

Validation

Two SGLang workers (Qwen3-VL-2B) behind the KV router on SGLang main 5e79333c95, same client script before and after:

main:  12 BlockStored events, 7 decoded, 5 "Failed to decode ZMQ KV batch" warnings
PR:    12 BlockStored events, 12 applied ok, 0 warnings
both:  text repeat kv_hit_rate 0.968, image repeat kv_hit_rate 0.0

The new wire tests decode an SGLang-shaped batch with wide pads, ids on both sides of the u32 boundary, a text event and a bigram event. On main they fail with data did not match any variant of untagged enum KvTokenIds. kv-router: 1107 passed. mocker: 272 passed. dynamo-llm kv_router: 441 passed.

Related Issues

Summary by CodeRabbit

Bug Fixes
- Improved token ID handling for values larger than the supported 32-bit range.
- Preserved standard token IDs while safely normalizing oversized IDs.
- Updated single-token and bigram processing to maintain consistent token data across batch events.
Tests
- Added coverage for oversized, standard, and nested token ID inputs.

Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>

coderabbitai · 2026-10-21T21:32:36Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: ai-dynamo/dynamo/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0271830b-6858-404a-8190-892a5477861b

📥 Commits

Reviewing files that changed from the base of the PR and between 5004f3f and dfae347.

📒 Files selected for processing: 4

lib/kv-router/src/protocols.rs
lib/kv-router/src/zmq_wire/deserialize.rs
lib/kv-router/src/zmq_wire/tests.rs
lib/kv-router/src/zmq_wire/types.rs

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The change accepts 64-bit token IDs in the ZMQ wire types and folds IDs larger than u32::MAX during normalization. Added tests cover wide and ordinary IDs, including nested token-ID arrays.

Changes

Token ID normalization

Layer / File(s)	Summary
Accept and normalize wide token IDs `lib/kv-router/src/zmq_wire/types.rs`, `lib/kv-router/src/protocols.rs`, `lib/kv-router/src/zmq_wire/deserialize.rs`, `lib/kv-router/src/zmq_wire/tests.rs`	`KvTokenIds` stores 64-bit IDs. Normalization folds each ID to the 32-bit token space. Tests cover oversized and ordinary IDs, including nested arrays.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to dfae3

The wide-token handling has no identified merge blocker. Repeated-image cache hits remain deferred work, not a regression established for this change.

Pre-merge checks: 5

Check name	Status	Explanation
Description check		The description explains the problem, implementation, validation results, and related issue. It does not include a separate “Where should the reviewer start?” section, and DIS-3202 is not written as a…
Title check		The title clearly and concisely describes the main change: folding wide SGLang image pad token IDs in KV events.
Docstring Coverage		Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 4 files.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description explains the problem, implementation, validation results, and related issue. It does not include a separate “Where should the reviewer start?” section, and DIS-3202 is not written as an explicit issue reference, but the description is otherwise complete and relevant.

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

dynamo-review-agent · 2026-10-21T21:47:20Z

 fn normalize_token_ids(token_ids: KvTokenIds) -> (Vec<u32>, bool) {
    match token_ids {
-        KvTokenIds::Single(tids) => (tids, false),
+        KvTokenIds::Single(tids) => (tids.into_iter().map(fold_token_id).collect(), false),


This removes the zero-copy path for every ordinary single-form KV event: deserialization now first allocates a Vec<u64>, then this collect() allocates and copies a second Vec<u32>. All live and replayed ZMQ BlockStored batches pass through here, so vLLM, TensorRT-LLM, mocker, and text-only SGLang traffic pay the extra per-event allocation even when every ID fits in u32. Keep a direct u32 decoding path and use folding only for inputs that actually contain wide IDs (for example with a custom visitor or separate untagged wide variants).

🤖 AI Fix

Deserialize fitting token IDs directly into the final Vec<u32> and allocate a folded replacement only for token sequences that contain values above u32::MAX.

dynamo-review-agent · 2026-10-21T21:47:20Z

+    );
+}
+
+#[test]


This test duplicates pass-through coverage already retained by test_decode_batch_folds_sglang_wide_pad_token_ids: that test exercises the same batch decoder with ordinary single-form IDs, u32::MAX, and an ordinary bigram element while also detecting the wide-ID regression. The existing test_deserialize_bigram_block_stored_sequence additionally preserves an all-u32 bigram. The extra values here traverse no distinct normalization boundary; the supported behavior remains fully protected without this case.

🤖 AI Fix

Remove test_decode_batch_keeps_u32_token_ids and retain the wide-pad and existing bigram tests.

devin-ai-integration · 2026-10-21T22:27:14Z

✅ Dynamo PR CI passed — run 39679335695 (attempt 1) on `dfae347f8c`

Gate checks: ✅ backend-status-check · ✅ deploy-status-check · ✅ dynamo-status-check

_{Posted automatically by Devin for run 39679335695. Updated on every full-CI run of this PR.}

KrishnanPrash · 2026-10-28T00:36:48Z

Closing in favor of the DEP #17044 work in #17126, #17082 and sgl-project/sglang#45737.

fix(kv-router): fold SGLang wide image pad token ids in KV events

dfae347

Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>

KrishnanPrash requested a review from a team as a code owner October 21, 2026 21:30

pull-request-size Bot added the size/M label Oct 21, 2026

github-actions Bot added the fix label Oct 21, 2026

This was referenced Oct 21, 2026

fix(router): write SGLang wide image pads in MM routing tokens #17048

Closed

[BUG]: SGLang after sglang#43043 breaks multimodal routing for Kimi #16975

Open

dynamo-review-agent Bot reviewed Oct 21, 2026

View reviewed changes

KrishnanPrash closed this Oct 28, 2026

Navigation Menu

Repository navigation

fix(kv-router): fold SGLang wide image pad token ids in KV events - #17046

fix(kv-router): fold SGLang wide image pad token ids in KV events#17046
KrishnanPrash wants to merge 1 commit into
mainfrom
kprashanth/dis-3202-sglang-wide-pads

KrishnanPrash commented Oct 21, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Oct 21, 2026

Uh oh!

dynamo-review-agent Bot Oct 21, 2026

Uh oh!

dynamo-review-agent Bot Oct 21, 2026

Uh oh!

devin-ai-integration Bot commented Oct 21, 2026

Uh oh!

KrishnanPrash commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

+                  );
+              }
+              #[test]

Conversation

KrishnanPrash commented Oct 21, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Related Issues

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Oct 21, 2026

Walkthrough

Changes

Uh oh!

dynamo-review-agent Bot Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

dynamo-review-agent Bot Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot commented Oct 21, 2026

✅ Dynamo PR CI passed — run 39679335695 (attempt 1) on dfae347f8c

Uh oh!

KrishnanPrash commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

KrishnanPrash commented Oct 21, 2026 •

edited by coderabbitai Bot

Loading

✅ Dynamo PR CI passed — run 39679335695 (attempt 1) on `dfae347f8c`