Inference: Fixes for large scale and concurrency by sidsingh-nvidia · Pull Request #6497 · NVIDIA/Megatron-LM

sidsingh-nvidia · 2026-08-12T21:22:16Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Coordinator wire-format change for faster routing

Splits the coordinator's ZMQ wire format into multipart frames: a small, constant-size metadata frame followed by opaque body frames. The coordinator unpacks and repacks only the metadata frame and forwards the bodies untouched. The metadata is of a small constant size, independent of the sequence length.
Todo - apply the same fix for multimodal. Move large multimodal data to the body from the metadata.

Hash prompts on the client, not the coordinator

Hashing is embarassingly parallel and can be handled on frontend replicas - each of which spawns a client.

Apply a weight versioned salt to prefix hashes in the engine

This prevents prompts in RL from matching with stale prefix caches.

Change routing cost function

A linear combination of the prefix match score and the relative load on each engine. Works well in practice and is able to balance load and affinity.

Load is measured against the fleet mean rather than against max_requests. Against a configured ceiling a sizeable imbalance stays a small fraction of it, so affinity keeps winning however lopsided the fleet gets; against the mean the term vanishes while ranks are even and grows only as they diverge — the drain at the end of a batch, when work should spread. It subtracts rather than blending, so a full cache hit cannot cancel the load term and strand work on a saturated rank.

Prefix depth is now graded per rank and counted forward from the first block. The previous match credited only the ranks holding the deepest cached block, so a rank holding 8 of 10 blocks scored the same as one holding none, leaving the load term nothing to trade against. Counting forward also stops a rank being credited for a deep block whose prefix has been evicted — unusable KV, since the hashes chain.

Expire coordinator hash entries after a TTL

The coordinator's hash table is a guess about what each engine still holds: it sees blocks being routed, never blocks being evicted, so the guess only gets more optimistic and the table grows without bound. Entries untouched for prefix_cache_ttl_seconds (default 300) are dropped. The queue is insertion-ordered, so the sweep stops at the first live entry and costs what it evicts rather than the size of the table.

Behaviour changes worth flagging

prefix_caching_routing_alpha changes meaning. It was a blend weight in [0, 1]; it is now a dimensionless, uncapped coefficient on the load term. A value tuned for the old scoring does not carry over — and at 1.0 a single request of imbalance across two ranks exactly cancels a full cache hit, so the default stays 0.5.

detokenize_generations is a new SamplingParams field, defaulting to True. The coordinator used to detokenize every reply unconditionally. Now that reply bodies are opaque frames, it detokenizes only when asked — defaulting on so no existing client changes behaviour. Clients that detokenize for themselves (the HTTP frontend does) can opt out and let the coordinator forward the body without decoding and re-encoding it.

New knobs: --inference-dynamic-batching-prefix-cache-ttl-seconds.

Issue tracking

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Existing coordinator, client, streaming, engine, and prefix-caching-coordinator unit tests were updated to the new framing, and roughly 25 tests added across the pieces:

Client-side hashing — the client hashes and salts multimodal with the media key, string prompts are left unhashed, LOAD_BALANCED skips hashing, prefix routing reads those hashes without decoding the prompt (using a deliberately undecodable prompt frame, so a decode attempt would raise), an empty hash list is not treated as a fallback and a null frame is.
Weight salt — generations are disjoint, epoch 0 is byte-identical to unsalted, media identity stays distinct within a generation, a text request admitted after a resume cannot match blocks from before it, and caller-supplied hashes are adopted rather than re-salted (a disaggregated handoff depends on this).
Cost function — one test per property: alpha=0 is pure affinity, mild imbalance does not overturn a hit, a saturated rank loses to an idle one, a balanced fleet ignores load entirely, deeper prefixes outscore shallower ones, a deep block without its prefix earns nothing, and ties go to the least loaded rank.
TTL — expiry, refresh-on-reuse, an expired entry ceasing to win routing, and survival across an engine removal (which renumbers rank indices, so the expiry queue is deliberately keyed on the hash rather than the rank).
Framing robustness — a malformed submission is dropped rather than raising out of the event loop, and handle_engine_reply is covered in both the detokenizing and forward-untouched cases.

Verified against main at a3344a6f: across all of tests/unit_tests/inference/, including engines/, the set of failing tests is identical with and without this branch (218 either side), and this branch passes 25 more tests than main.

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-08-12T21:22:20Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-08-12T21:22:29Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

…is metadata of constant size. This is the only thing that the coordinator needs to unpack/read and pack Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

wdykas

Looks good from disag side

santhnm2

LGTM, but as a suggestion for a future PR: the design here is passing around lists of frames where frames are arbitrary packed bytes, which requires indexing of the frame list at certain indices and doesn't assign any semantic meaning to each value in the list. Would it make sense to have a wrapper class so we could refer to the header frame more explicitly?

sidsingh-nvidia · 2026-08-19T22:57:05Z

@santhnm2 that's a fair critique. I can impose the following structure on all messages

each message can only consist of two frames - [metadata, payload]
metadata and payload will be dicts with bespoke keys for easy readabiity - this can remove seemingly arbitrary indexing.

santhnm2 · 2026-08-19T22:58:50Z

@santhnm2 that's a fair critique. I can impose the following structure on all messages

each message can only consist of two frames - [metadata, payload]

metadata and payload will be dicts with bespoke keys for easy readabiity - this can remove seemingly arbitrary indexing.

I think this makes sense. An even tighter contract might be a custom dataclass with metadata and payload fields, but that's a small design decision I'll defer to you

wdykas · 2026-08-19T23:00:30Z

Since we wanna pack metadata and payload separately, we will have to separate them out prior to the send operation.

do you know what we are allowed to serialize with msg pack? I thought it was real stingy

Addresses review feedback on NVIDIA#6497: the frame lists were passed around as arbitrary packed bytes, indexed positionally at each site, with no semantic meaning attached to any element. A message's layout previously existed only as an unwritten agreement between a packb site and an unpack site in another module -- `[header, request_id, sampling_params]` matched by `_, request_id, sampling_params = metadata` three files away. Nothing tied the two together, which is how a 5-field metadata frame came to be read by a handler expecting 6. Each message type is now declared once as a MessageSpec naming its metadata fields and payload frames in wire order. pack and parse are both derived from that declaration, so no frame index is written by hand and the two directions cannot drift. Adding a field means editing one tuple. The split between the two tuples is the contract that matters, and it is now structural rather than prose: metadata_fields are decoded by the coordinator and must be constant-size, payload_frames are forwarded untouched and may grow with the prompt. Putting src_block_ids in one rather than the other is now a visible review-time decision. Batched engine replies keep a dedicated BatchedMessageSpec. They carry one payload frame per finished request because the batch fans out to different clients: collapsing it into a single frame would force the coordinator to decode every reply just to find the boundaries, and a finished request echoes the prompt back. Measured on a 50-reply batch with 2000-token prompts, that split costs the coordinator 0.001 ms against 3.7 ms for a decode. pack now asserts that a batch's entry count matches its frame count, so a mismatch fails at the send rather than silently truncating the batch. No wire-format change: every layout was verified to produce byte-identical frames to the previous code. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

Lowering media slots to model tokens was left on the event loop when the text path moved off it. Only the chat-template render was awaited, via asyncio.to_thread, and rendering is the first of 3N+1 tokenizer calls for N media slots: each slot encodes the text preceding it, then the model token's prefix and suffix, and a final call encodes the tail. Every one of those ran inline on the loop, stalling delivery of responses for every other request the replica owns -- exactly what the offload exists to prevent, on the requests with the longest prompts. The whole lowering is now one synchronous function dispatched to the tokenize executor in a single hop. That also fixes a thread-safety hazard: the render went to the default executor, which is shared and many-threaded, holding the tokenizer every other request also holds. HF tokenizers are not thread-safe. Running on the dedicated single-worker executor means one thread touches it, and it is that executor's private copy rather than the shared one. Fetching media stays on asyncio.to_thread. That work is network I/O rather than tokenizer work, so it wants the shared pool's concurrency and holds nothing that a second thread could corrupt. This also makes an existing comment true: `chat_tok` is now only read for the capability checks, having previously also been handed to this path. Verified against nvidia/main across all of tests/unit_tests/inference including engines/: the set of failing tests is identical, 218 on both sides. The existing slot-lowering test moves to the synchronous function, and a new one pins the shape -- that it is not a coroutine, and that the endpoint reaches it only through the executor. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

Multimodal requests were the last thing making the coordinator decode a prompt. The endpoint hashed text but skipped media, because salting those hashes needs the media key and the endpoint does not have it -- serialize_multimodal_data derives it a few frames later, inside the client, and deriving it again would digest the media a second time, hundreds of MB for video. Moving hashing into the client removes the problem rather than working around it. The client is the one place holding both inputs at once: the tokens, and the key it has just derived. Text and multimodal now take the same path, the media key salts the hashes at no extra cost, and the coordinator no longer decodes a prompt for any request the frontend submits. An unconfigured client reports None, not an empty list. Callers that build an InferenceClient directly -- MegatronAsyncLLM, megatron.rl, the coordinator example -- pass neither the block size nor the policy. None means "I did not hash" and the coordinator hashes on their behalf, exactly as today; an empty list would mean "hashed, nothing matched", so it would skip its own hashing and turn prefix-affinity routing into load balancing with nothing raising to say so. That distinction is the whole reason the hashes are their own frame rather than a field in the metadata, and it is now pinned by a test. Documents the frame contract on both sides while here, since the three frames have three different rules: metadata is decoded and repacked every request so it must stay constant-size, the prompt is forwarded verbatim and never decoded, and the hashes are decoded, consumed for routing, and not forwarded -- one int per block rather than per token, which is why they do not belong in metadata. Verified against nvidia/main across all of tests/unit_tests/inference including engines/: the set of failing tests is identical, 218 on both sides. Five tests added, covering the unconfigured-client regression, text hashing, multimodal salting, the LOAD_BALANCED skip, and string prompts. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia · 2026-09-01T19:30:54Z

/ok to test d611c58

The coordinator open-coded the "does anyone read the hashes" test as `policy != LOAD_BALANCED` while the client asked `routes_on_prefix`. Two spellings of one predicate, which is what routes_on_prefix was written to prevent -- its docstring even claims a new prefix-aware policy only has to be added in one place. They agree today only because there are three members: add a fourth and the client would stop hashing for a policy this handler still expects hashes from, or hash for one that discards them. Also drops the functools import from chat_completions, unused since lowering media slots stopped being a coroutine and no longer needed a partial. CI's pylint caught it; the container used for local runs has no pylint, and the version installed there rejects the repo's .pylintrc, so this did not surface until the job ran. Verified against nvidia/main across all of tests/unit_tests/inference including engines/: the set of failing tests is identical, 218 on both sides. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia · 2026-09-01T19:43:33Z

/ok to test d8ed7e5

The metadata frame is decoded and repacked for every request on the one serial loop shared by every rank, so everything in it has to be bounded. The serialized multimodal dict was not: it carries the raw image or video bytes, up to hundreds of megabytes, which is three orders of magnitude more than the prompt decode an earlier commit removed. Split it. The bounded descriptor -- the 64-character content key plus the modality and token-expansion flags -- stays in metadata, which is all the coordinator ever reads: the key salts the routing hashes when nobody hashed upstream, and drives vision-embedding-cache affinity. The payload moves to a fourth frame the coordinator forwards without decoding, and the engine reassembles the two. split_multimodal_data / merge_multimodal_data live next to serialize_multimodal_data so no frame-shape knowledge leaks into the client or the engine. Both serialization branches go through them: the preprocessed tensor dict is as unbounded as raw bytes. The frame count is fixed at four rather than varying with media, so a malformed submission is still caught by an arity check; a text-only request pays one byte for a None media frame. The TP broadcast needs no change -- it carries a per-message frame-count manifest precisely so grouping survives without payloads being copied or decoded. Also folds the duplicated frame construction in add_request and add_request_streaming into one _pack_submit_frames, and corrects a handler comment that claimed multimodal never reaches the coordinator's hashing fallback. It does, for an unconfigured client, which is why the media key is still passed there as the salt. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia · 2026-09-01T20:42:32Z

/ok to test 698d7d4

sidsingh-nvidia · 2026-09-01T20:44:11Z

/ok to test a8125e4

sidsingh-nvidia · 2026-09-02T00:15:36Z

/ok to test fbb4d54

Brings in NVIDIA#6916, which splits add_request into a thin wrapper over a new add_request_with_id returning (id, future) so an HTTP handler can name a request to abort_request when its client disconnects mid-generation. Both conflicts were the two branches editing adjacent lines rather than disagreeing: - inference_client.py: keep upstream's add_request / add_request_with_id split, but build the submission with _pack_submit_frames so a request still travels as four frames. add_request_with_id has to return the tuple that its callers in completions.py and chat_completions.py unpack; returning the bare future silently breaks both. - test_inference_client.py: both sides appended tests at the same point, so keep both. Upstream's new test read the submission off socket.send, but submissions are send_multipart on this branch, so it now reads frame 0. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia · 2026-09-02T18:32:12Z

/ok to test c7bb112

nemo-automation-bot · 2026-09-03T09:00:58Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33736490650

sidsingh-nvidia requested review from a team as code owners August 12, 2026 21:22

svcnvidia-nemo-ci marked this pull request as draft August 12, 2026 21:22

sidsingh-nvidia force-pushed the siddharth/co-ordinator-fixes branch from 8082e93 to 570a137 Compare August 12, 2026 21:29

sidsingh-nvidia mentioned this pull request Aug 19, 2026

feat(inference): prefix-affinity routing and a configurable cost policy in the DP coordinator #6668

Closed

sidsingh-nvidia requested a review from nvcsathe August 19, 2026 18:45

sidsingh-nvidia marked this pull request as ready for review August 19, 2026 20:58

sidsingh-nvidia force-pushed the siddharth/co-ordinator-fixes branch from 570a137 to 1e4999a Compare August 19, 2026 21:10

sidsingh-nvidia requested a review from wdykas August 19, 2026 21:15

sidsingh-nvidia force-pushed the siddharth/co-ordinator-fixes branch 5 times, most recently from 2d70fa8 to 2849a1c Compare August 19, 2026 21:57

modify coordinator wire format - split into two chunks - first chunk …

2de90c0

…is metadata of constant size. This is the only thing that the coordinator needs to unpack/read and pack Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia force-pushed the siddharth/co-ordinator-fixes branch from 2849a1c to 815842c Compare August 19, 2026 22:06

github-actions Bot added the complexity: medium label Aug 19, 2026

wdykas approved these changes Aug 19, 2026

View reviewed changes

santhnm2 approved these changes Aug 19, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Aug 19, 2026

kvareddy approved these changes Aug 20, 2026

View reviewed changes

nvcsathe approved these changes Aug 20, 2026

View reviewed changes

lhb8125 mentioned this pull request Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

sidsingh-nvidia added 2 commits September 1, 2026 11:52

copy-pr-bot Bot temporarily deployed to public September 1, 2026 19:34 Inactive

copy-pr-bot Bot had a problem deploying to test September 1, 2026 19:44 Error

copy-pr-bot Bot temporarily deployed to public September 1, 2026 19:47 Inactive

sidsingh-nvidia removed Run tests Run functional tests labels Sep 1, 2026

Merge branch 'main' into siddharth/co-ordinator-fixes

a8125e4

sidsingh-nvidia added Run tests Run functional tests labels Sep 1, 2026

copy-pr-bot Bot had a problem deploying to test September 1, 2026 20:45 Error

copy-pr-bot Bot temporarily deployed to public September 1, 2026 20:48 Inactive

Merge branch 'main' into siddharth/co-ordinator-fixes

fbb4d54

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

sidsingh-nvidia mentioned this pull request Oct 6, 2026

Disaggregated Inference: End2End #6675

Merged

Conversation

sidsingh-nvidia commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Coordinator wire-format change for faster routing

Hash prompts on the client, not the coordinator

Apply a weight versioned salt to prefix hashes in the engine

Change routing cost function

Expire coordinator hash entries after a TTL

Behaviour changes worth flagging

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Aug 12, 2026

Uh oh!

github-actions Bot commented Aug 12, 2026

Uh oh!

wdykas left a comment

Choose a reason for hiding this comment

Uh oh!

santhnm2 left a comment

Choose a reason for hiding this comment

Uh oh!

sidsingh-nvidia commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

santhnm2 commented Aug 19, 2026

Uh oh!

wdykas commented Aug 19, 2026

Uh oh!

sidsingh-nvidia commented Sep 1, 2026

Uh oh!

sidsingh-nvidia commented Sep 1, 2026

Uh oh!

sidsingh-nvidia commented Sep 1, 2026

Uh oh!

sidsingh-nvidia commented Sep 1, 2026

Uh oh!

sidsingh-nvidia commented Sep 2, 2026

Uh oh!

sidsingh-nvidia commented Sep 2, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 3, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

sidsingh-nvidia commented Aug 12, 2026 •

edited

Loading

sidsingh-nvidia commented Aug 19, 2026 •

edited

Loading