fix(tokenizer): preserve render request key order - fixes #2641 by Lucas-Fernandes-Martins · Pull Request #2644 · llm-d/llm-d-router

Lucas-Fernandes-Martins · 2026-09-01T12:23:57Z

What type of PR is this?

What this PR does / why we need it:

Fixes #2641, which seemingly breaks precise KV cache as the /render code path marshals and unmarshals the raw body, potentially leading to llm-d index seeing different token ID's than those received via ZMQ messages.

Tentative solutuion through which we check if EPP has not mutated the body, and if so, we use the rawBody directly to call /render

Which issue(s) this PR fixes:

#2641
Fixes #

Release note (write NONE if no user-facing change):

NONE

Signed-off-by: Lucas-Fernandes-Martins <lucasfmartins16@gmail.com>

Copilot

Pull request overview

This PR addresses issue #2641 by ensuring the vLLM /render tokenization path can use the original request JSON bytes (when the request body is forwarded unchanged), preventing Go map re-marshaling from reordering object keys and causing token mismatches that break precise KV-cache routing.

Changes:

Propagate the original request RawBody bytes into scheduling.InferenceRequest so downstream dataproducers can access the exact incoming JSON.
Extend the tokenizer backend to preferentially call vLLM chat render using the raw JSON payload when the body has not been mutated (preserving key order).
Add tests verifying raw payload preservation, correct fallback behavior when the payload is mutated, and correct fallback when the model differs.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
pkg/epp/requestcontrol/director.go	Plumbs request `RawBody` into the scheduling request context.
pkg/epp/requestcontrol/director_test.go	Asserts `SchedulingRequest.RawBody` carries the original request bytes.
pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http.go	Adds a raw-payload POST path to avoid re-marshaling for `/render`.
pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/vllm_http_test.go	Adds coverage ensuring key order is preserved via raw payload and that fallbacks behave correctly.
pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go	Allows backends to tokenize using the full `InferenceRequest` when needed.
pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/backend.go	Implements raw-body tokenization for vLLM chat renders when the request is unmutated.
pkg/epp/framework/interface/scheduling/types.go	Adds `RawBody []byte` to `InferenceRequest` to carry original bytes through scheduling/plugins.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

nilig

The fix goes in the right direction: the only reliable way to keep the router and vLLM in agreement is to send both of them the exact same bytes. My comments are about the cases the fix doesn't cover, the fact that you can't tell when it's active, and test coverage. None of them block the approach. Details are inline.

Two points on the PR itself:

The release note says NONE, but this fixes routing behavior operators can see, so it should get a line. Otherwise the fix won't appear in the release notes at all.
nit: "solutuion" typo in the description; the description says the body is "marshaled and unmarshaled", but what actually happens is the parsed map gets re-serialized with sorted keys while the engine gets the original bytes; and /kind bug is still inside the HTML comment, so the PR never got its kind label.

nilig · 2026-09-01T17:53:46Z

+	if !ok || body.ChatCompletions == nil || body.Mutated || len(request.RawBody) == 0 ||
+		request.TargetModel != renderer.modelName || body.Model != renderer.modelName || body.Payload == nil {


The raw path only turns on when the model name in the request exactly matches the tokenizer's configured modelName. If clients call the model by a different accepted name (a vLLM --served-model-name alias, or a LoRA adapter name), the request is not considered mutated, so the engine still gets the original bytes while /render gets the reordered copy. For those deployments the #2641 bug stays. That case was broken before #2318 too, so I don't think it needs to hold up this PR. Suggest we file a follow-up issue for it (a likely shape: replace just the model value inside the raw bytes without reordering anything else, the way #2511 keys purely off Mutated on the Anthropic path), and in this PR just note the limitation in a comment on produceRequest and in the README's modelName row so the difference between the two paths is on record.

nilig · 2026-09-01T17:53:46Z

+	renderer, ok := b.tk.(*vllmHTTPRenderer)
+	if !ok || body.ChatCompletions == nil || body.Mutated || len(request.RawBody) == 0 ||
+		request.TargetModel != renderer.modelName || body.Model != renderer.modelName || body.Payload == nil {
+		return b.produce(ctx, body)


There is no way to see which path a request took. No log line, no metric, and both paths produce the same trace span. So if a deployment fails one of the conditions, it looks exactly like the original bug: routing quietly stops matching and nothing tells you why. The original issue took days to notice for the same reason. A DEBUG log on the fallback saying which condition failed (and printing both model names when they differ) would fix that.

nilig · 2026-09-01T17:53:46Z

+// also the body that will be forwarded downstream, preserving object key order.
+func (b renderBackend) produceRequest(ctx context.Context, request *scheduling.InferenceRequest) (*fwkrh.TokenizedRequest, error) {
+	body := request.Body
+	renderer, ok := b.tk.(*vllmHTTPRenderer)


nit: this is a concrete-type check; everywhere else this package uses small interfaces (timeoutAware, anthropicRenderer). Same behavior either way today, so style only.

nilig · 2026-09-01T17:53:46Z

+	// RawBody contains the request body exactly as received from the downstream proxy.
+	RawBody []byte


The whole fix assumes nothing modifies the request body after the tokenizer runs. That's true today, but only by accident: MutatePayloadMap is public, and admission and PreRequest plugins run after the tokenizer, so a third-party plugin that adds any field to the body would quietly reintroduce this exact bug, and no test would catch it. Let's write the rule down: the RawBody doc should say the director forwards these bytes unchanged only while Mutated is false, and produceRequest should get one sentence explaining why it needs the same bytes.

nit: RequestSizeBytes and RawBody now store the same information separately. Setting one from the other would keep them from disagreeing.

nilig · 2026-09-01T17:53:46Z

+	assert.Equal(t, rawBody, captured.chat)
+}
+
+func TestProduce_ChatCompletionsVLLMHTTPUsesMutatedPayload(t *testing.T) {


Each of the fallback conditions should have its own test. Right now this test doesn't set Body.Model, so it falls back because of the model check, not the mutated check; removing the Mutated condition wouldn't fail any test. The TargetModel check, empty RawBody, nil payload, and protobuf payload have no tests at all. A small table test that flips one condition at a time would cover all of them.

nilig · 2026-09-01T17:53:46Z

+
+	p := newTestPlugin(newHTTPRenderer(t, srv))
+	require.NoError(t, p.Produce(context.Background(), req, nil))
+	assert.Equal(t, rawBody, captured.chat)


nit: this test checks the bytes that were sent but never checks the tokens that came back, so the response handling of postRawChatRender is untested. Copying the token assertions from the neighboring tests would close that.

nilig · 2026-09-01T17:53:46Z

 	return resp.TokenIDs, toKVCacheMM(resp.Features), nil
 }

+func (r *vllmHTTPRenderer) postRawChatRender(ctx context.Context, body []byte, timeout time.Duration) ([]uint32, *tokenization.MultiModalFeatures, error) {


nit: this duplicates postChatRender minus the marshal; postChatRender could just marshal and call postRawChatRender. Same for produceRequest copying the body of produce's chat branch.

nilig · 2026-09-01T17:53:46Z

 		RequestID:        reqCtx.Request.Headers[reqcommon.RequestIDHeaderKey],
 		TargetModel:      reqCtx.TargetModelName,
 		Body:             inferenceRequestBody,
+		RawBody:          reqCtx.Request.RawBody,


nit: RawBody is attached even for mutated requests, where nothing ever reads it, and it keeps the original body in memory until the stream ends. Wonder if we could set it only when the body wasn't mutated.

vMaroon · 2026-09-02T13:36:26Z

+		return b.produce(ctx, body)
+	}
+
+	payload, ok := body.Payload.AsMap()


nit: neither the body.Payload == nil guard nor the payload.AsMap() failure fallback (lines 124-127) is covered. A case with a nil or non-map payload asserting the fallback to produce would lock that in.

vMaroon · 2026-09-02T13:36:50Z

+	if err != nil {
+		return nil, fmt.Errorf("tokenization failed: %w", err)
+	}
+	return &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{


nit: this TokenizedRequest construction is now the third verbatim copy of the same {Prompts: []PromptTokens{{TokenIDs, MultiModalFeatures: convertMMFeaturesToUpstream(...)}}} shape (also at lines 150-153 and 159-162 in produce). Wonder if we could extract a helper, e.g. newRenderResult(tokenIDs []uint32, mm *tokenization.MultiModalFeatures) *fwkrh.TokenizedRequest, and reuse it in all three.

fix(tokenizer): preserve render request key order

16578cc

Signed-off-by: Lucas-Fernandes-Martins <lucasfmartins16@gmail.com>

Copilot AI lite review requested due to automatic review settings September 1, 2026 12:23

Lucas-Fernandes-Martins requested review from a team, liu-cong, sagearc, vMaroon and zetxqx as code owners September 1, 2026 12:23

github-actions Bot added size/L Denotes a PR that changes 100-499 lines, ignoring generated files. area/epp area/scheduling labels Sep 1, 2026

Copilot started reviewing on behalf of Lucas-Fernandes-Martins September 1, 2026 12:24 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

nilig reviewed Sep 1, 2026

View reviewed changes

vMaroon reviewed Sep 2, 2026

View reviewed changes

Navigation Menu

Uh oh!

fix(tokenizer): preserve render request key order - fixes #2641 - #2644

fix(tokenizer): preserve render request key order - fixes #2641#2644
Lucas-Fernandes-Martins wants to merge 1 commit into
llm-d:mainfrom
mistralai:lfm/fix-render-payload-parity

Lucas-Fernandes-Martins commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

nilig left a comment

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

nilig Sep 1, 2026

Uh oh!

vMaroon Sep 2, 2026

Uh oh!

vMaroon Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if !ok \|\| body.ChatCompletions == nil \|\| body.Mutated \|\| len(request.RawBody) == 0 \|\|
		request.TargetModel != renderer.modelName \|\| body.Model != renderer.modelName \|\| body.Payload == nil {

		// RawBody contains the request body exactly as received from the downstream proxy.
		RawBody []byte

Uh oh!

Conversation

Lucas-Fernandes-Martins commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

nilig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants