[wdykas]

This MR completes end-to-end prefill/decode disaggregation for Megatron dynamic inference. Prefill and decode run as dedicated, independently parallelized shards. A coordinator routes requests and manages capacity and ownership, while KV and SSM tensors move directly between GPU workers through NCCL or NIXL. The coordinator carries request metadata and lifecycle messages only; it is not in the tensor data path.

Native coordinator request flow

sequenceDiagram
    autonumber
    participant U as Client / RL rollout
    participant C as Coordinator
    participant P as Prefill shard
    participant D as Decode shard

    P->>C: REGISTER_ROLE(prefill, capacity, transfer metadata)
    D->>C: REGISTER_ROLE(decode, capacity, transfer metadata)
    C-->>P: REGISTER_ROLE_ACK
    C-->>D: REGISTER_ROLE_ACK

    U->>C: Submit prompt and sampling parameters
    C->>P: Route using prefix affinity and available capacity

    P->>P: Execute prompt
    P->>P: Produce first token and optional MTP proposals
    P->>P: Pin prompt KV and exact final live SSM state
    P-->>C: Prefill result and handoff metadata

    C->>D: SUBMIT_REQUEST_WITH_KV
    D->>D: Match contiguous locally cached KV prefix
    D->>D: Reserve missing KV, continuation blocks, and a live SSM request slot

    alt NIXL
        P-->>D: One-sided read initiated by decode
    else NCCL
        D-->>C: KV_TRANSFER_READY
        C-->>P: SEND_KV
        P->>D: Send uncached KV suffix and SSM state
    end

    D->>D: Aggregate completion across model-parallel ranks
    D->>D: Broadcast KV_HANDOFF_COMPLETE within the MP group
    D->>D: Register imported complete-block hashes
    D->>D: Admit directly into decode without prompt replay

    D-->>C: KV_READ_DONE
    C-->>P: RELEASE_KV
    P->>P: Release pinned KV and SSM storage

    D-->>C: Generated tokens / final result
    C-->>U: Generated tokens / final result

Loading

Request lifecycle

1. Role registration

   Each engine registers as either a prefill or decode instance. Registration includes its request capacity, transport backend, KV layout, and hybrid SSM layout when present. The coordinator validates that paired engines have compatible model geometry.

2. Prefill routing

   The coordinator selects a prefill instance using prefix affinity, available capacity, and current load. The prefill engine executes the prompt and captures everything required to continue decoding.

3. Source-state retention

   When prefill finishes, its prompt KV blocks are temporarily retained so the allocator cannot reuse them during transfer. For hybrid models, the exact final live SSM slot is detached from the completed request and retained in the same way.

4. Decode-side cache reuse

   The coordinator selects a decode instance using prefix affinity, available capacity, and current load. Decode recomputes the prompt's complete-block hash chain and finds the longest contiguous prefix already present in its local cache. Only the uncached KV suffix is allocated and transferred. Complete blocks are hash-indexed for reuse; a partial final prompt block is transferred but remains request-owned rather than being registered as a reusable prefix. Concurrent imports reserve their hashes so two requests cannot independently import and register the same prefix.

5. State transfer

   NIXL performs decode-initiated one-sided reads. NCCL uses matching receives and sends: decode first commits its destination storage and reports KV_TRANSFER_READY, after which the coordinator tells prefill to post its sends. Different prefill and decode TP, PP, EP, and expert-TP layouts are supported by the KV and SSM re-sharding planners.

6. Decode admission

   Decode waits until every model-parallel rank has completed the transfer, then admits the request directly into the active decode batch. It does not replay any prompt tokens. The handoff includes the sampled first token and any MTP proposals so decoding starts at the exact continuation point. These continuation tokens have not passed through the base model yet, so decode produces their KV locally and allocates continuation blocks only when they cross a block boundary.

7. Source release

   Once decode no longer reads prefill storage, it reports KV_READ_DONE. The coordinator sends RELEASE_KV, allowing prefill to release its temporary KV references and retained SSM slot. NCCL sources remain pinned until their local send handles are also complete.

Control plane and data plane

The coordinator carries request and control metadata:

• Prompt and sampling parameters
• Source block IDs
• Buffer addresses and shard layouts
• First sampled token and MTP proposals
• Readiness, completion, abort, and release messages

The GPU workers transfer model state directly:

• Uncached KV suffix
• Exact convolution state for hybrid models
• Exact recurrent SSM state for hybrid models

Large KV and SSM tensors are never serialized through the coordinator.

Key design decisions

• Decode engines never execute prompt prefill.
• Prefix caching is required on both roles because decode admission reuses the prefix-cache allocator.
• Hybrid decode receives the exact live SSM state directly into a live request slot; it does not maintain a durable SSM prefix-snapshot cache.
• For hybrid models with expert parallelism, each disaggregated engine reserves one additional recurrent-state slot outside max_requests for EP dummy forwards. Real requests and retained handoffs cannot consume it, so an idle rank can still participate in expert collectives when every request-owned slot is occupied. The slot is reused between dummy steps, and its memory is included when sizing the recurrent-state and KV buffers.
• Transfer completion is aggregated over a CPU control path rather than adding a per-request collective to the compute critical path.
• Ranks in an inference shard receive the same request stream and use deterministic allocators, so corresponding requests receive the same physical block IDs. Rank-specific transfer layouts are exchanged during setup; unexpected allocator divergence is treated as an engine failure rather than adding a collective to every request admission.
• Requests are queued FIFO when request, KV, or live SSM capacity is unavailable.
• Cancellation does not immediately free storage involved in an active transfer. Source state remains pinned and decode destinations may be quarantined until the backend reports that reuse is safe. Reset also refuses to release storage while a transfer can still access it, preferring temporary capacity loss over memory corruption.
• Disaggregated behavior is isolated behind role-tagged inference shards; ordinary inference continues through the existing engine.
• The native coordinator, Dynamo integration, and Megatron RL use the same underlying handoff and transfer mechanisms. Each control plane supplies its own routing and orchestration.

RL checkpoint loading and refit

Each role-tagged RL inference shard is an independent checkpoint and refit destination. Checkpoint loading therefore uses the shard's process group rather than the global world group. Refit performs one collective pass for each disjoint destination pool: source ranks participate in every pass, while only ranks belonging to that pool materialize the destination weights.

CUDA graph-aware prefill scheduling

When --inference-cuda-graph-all-prefills is enabled, the scheduler treats each captured (token count, prefill count, decode count) shape as capacity rather than requiring the real batch to land exactly on its token boundary. For example, a 128-token graph can execute a 100-token chunk with padding, so the scheduler uses the full 100-token budget instead of shrinking the chunk to the next smaller captured size. This reduces the number of prompt chunks and avoids unnecessary eager execution.

The same graph-coverage check applies to both initial and continuation chunks. If the current resident batch prevents a compatible shape, the queue-head request waits while that finite batch drains and changes shape. Hybrid models use strict matching because captured prefill and decode SSM slots are not interchangeable; attention-only models can reuse unused request capacity more flexibly. Graph admission is a local CPU scan and does not add a model-parallel collective.

If an idle engine cannot cover a new request, it reports a configuration error instead of silently ignoring all-prefill graph mode. The only eager fallback is a lone continuation that has already started and cannot match any captured shape; without other resident work, deferring it could never change the batch and would deadlock progress.

Disaggregation protocol headers

These messages form the control plane. KV and SSM tensors are transferred directly through NCCL or NIXL and are never serialized through the coordinator.

Disaggregation-specific headers

Header	Direction	Purpose
REGISTER_ROLE	Prefill or decode -> coordinator	Registers the engine's role, transport backend, capacity, and transfer metadata.
REGISTER_ROLE_ACK	Coordinator -> prefill or decode	Confirms successful role registration.
SUBMIT_REQUEST_WITH_KV	Coordinator -> decode	Starts the decode hop with the request and source transfer metadata.
KV_TRANSFER_READY	Decode -> coordinator	NCCL only. Indicates that decode has reserved its destinations and prepared the receives.
SEND_KV	Coordinator -> prefill	NCCL only. Tells prefill to post sends for the uncached suffix.
KV_READ_DONE	Decode -> coordinator	Indicates that all decode ranks completed the handoff and the source state can be released.
RELEASE_KV	Coordinator -> prefill	Releases the KV blocks and live SSM slot retained by prefill for the handoff.
KV_HANDOFF_COMPLETE	Decode MP coordinator -> decode MP ranks	Distributes the aggregated transfer result so all ranks make the same admission or failure decision.

NIXL does not use KV_TRANSFER_READY or SEND_KV; decode initiates direct reads after allocating its destination storage.

Existing headers reused by disaggregation

Header	Direction	Purpose
CONNECT / CONNECT_ACK	Client <-> coordinator	Establishes the client connection.
SUBMIT_REQUEST	Client -> coordinator -> prefill	Starts the prefill hop.
ENGINE_REPLY	Prefill or decode -> coordinator	A prefill reply starts the decode handoff; a decode reply contains the final result.
ENGINE_REPLY_PARTIAL	Decode -> coordinator -> client	Streams incremental decode output.
ABORT_REQUEST	Client -> coordinator -> current owner	Cancels the request on whichever engine currently owns it.
REQUEST_ABORTED	Engine -> coordinator -> client	Confirms that cancellation and associated storage cleanup are safe.
REQUEST_ERROR	Engine -> coordinator -> client	Reports a terminal failure and whether the source storage is safe to release.
DISCONNECT	Prefill or decode -> coordinator	Removes the engine from its role and routing pool.
SHUTDOWN	Client -> coordinator	Stops the coordinator.
TP_BROADCAST	MP coordinator -> other MP ranks	Frames the existing internal ZMQ broadcast. Disaggregation uses it to carry KV_HANDOFF_COMPLETE, but it is not specific to disaggregation.
PAUSE, UNPAUSE, SUSPEND, RESUME, SET_GENERATION_EPOCH, STOP	Client -> coordinator -> engines	Existing engine lifecycle controls.
START_CUDA_PROFILER, STOP_CUDA_PROFILER	Client -> coordinator -> engines	Existing profiler controls.

Current limitations

• Native disaggregation requires prefix caching on both prefill and decode engines.
• All participating prefill and decode engines must use the same transfer backend and compatible model-state geometry.
• The native coordinator path currently supports text-only requests.
• Prompt log probabilities and top-N token log probabilities are not currently transferred across the handoff.
• Context parallelism is not currently supported for disaggregated RL inference.

Configuration

A native deployment is described using role-tagged inference shards, for example:

--inference-shards "tp=1,dp=2,role=prefill+tp=1,dp=2,role=decode"
--disagg-kv-transport-backend nccl

Each role can use its own supported TP, PP, DP, EP, and expert-TP configuration.

[copy-pr-bot]

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.

[wdykas]

/claude strict-review

[wdykas]

@nvcsathe I am going to try to get it all cleaned up in here. If you could see what the heck is going on for dynamo I will work on the other clean up. Not sure why it added like 2k extra lines in the rebase...

[wdykas]

/claude strict-review

[wdykas]

/ok to test a2e9402

[wdykas]

/ok to test 098fafb

[wdykas]

thanks for the feedback @tdene !

[wdykas]

/ok to test f192b94

[wdykas]

/ok to test 5458fb3

[wdykas]

/ok to test 5458fb3

[nemo-automation-bot]

🔄 Merge queue validation started!

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

[wdykas]

/ok to test e962afe

[wdykas]

/ok to test e962afe

[nemo-automation-bot]

🔄 Merge queue validation started!

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