Inference: Prevent the coordinator from choking at large concurrency - #6497
Inference: Prevent the coordinator from choking at large concurrency#6497sidsingh-nvidia wants to merge 2 commits into
Conversation
|
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:
See the contribution guide for more details. |
8082e93 to
570a137
Compare
570a137 to
1e4999a
Compare
2d70fa8 to
2849a1c
Compare
…is metadata of constant size. This is the only thing that the coordinator needs to unpack/read and pack
2849a1c to
815842c
Compare
wdykas
left a comment
wdykas
left a comment
There was a problem hiding this comment.
Looks good from disag side
santhnm2
left a comment
santhnm2
left a comment
There was a problem hiding this comment.
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?
|
@santhnm2 that's a fair critique. I can impose the following structure on all messages
|
I think this makes sense. An even tighter contract might be a custom dataclass with |
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>
5 participants
What does this PR do?
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.
Problem
The DP inference coordinator is a single serial event loop shared by every data-parallel rank and every frontend replica. Previously each message was one
msgpackblob, so to do its job — read a header, read a request id, pick a rank — the coordinator had to unpack and repack the entire payload.That payload is dominated by data the coordinator never looks at:
The result is per-request coordinator CPU that scales with prompt length rather than with request count. At high concurrency with long prompts the serial loop saturates and becomes the bottleneck for every rank behind it.
Fix
Move everything the coordinator needs into frame 0 and leave the rest as opaque frames:
SUBMIT_REQUEST[header, request_id, sampling_params][prompt]ENGINE_REPLY[header, [[request_id, detokenize], ...]]ENGINE_REPLY_PARTIAL[header, [request_id, ...]]Issue tracking
Linked issue:
Contribution process
Pre-checks
Existing coordinator, client, streaming, engine, and prefix-caching-coordinator unit tests were updated to the new framing (5 test files).
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"
.github/CODEOWNERS.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, theFinal Reviewlabel 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
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.