Skip to content

Add phase-disaggregated serving with streamed prefill-to-decode state handoff #1427

Description

@i386

Goal

Add a phase-disaggregated Mesh serving mode in which prefill and decode may run on different machines or topologies:

  • a high-throughput prefill worker/topology processes the prompt;
  • exact continuation state is streamed to the decode worker while prefill is still running;
  • the decoder imports and atomically commits that state; and
  • generation continues locally on the decode worker, including native MTP and suffix N-gram speculation.

Initial lab target:

This is phase disaggregation, not ordinary layer splitting.

Motivation

Current layer splitting distributes every decode token across the stage pipeline. That retains a network dependency on every generated token.

Phase disaggregation instead pays one prompt-state transfer:

Prompt ──► M3 prefill ──► streamed exact state ──► M1 decode + MTP/suffix
                  │
                  └── begins transferring before prefill completes

Once the handoff commits, decode is local to the M1. This can:

  • use the faster/larger M3 for prompt ingestion;
  • remove per-token decode traffic between machines;
  • overlap prefill for a later request with decode of an earlier request;
  • retain the M1 as the decode/MTP owner; and
  • improve aggregate goodput when prompt and generation phases overlap.

It will not necessarily improve a single short request. The topology must only be selected when measured state-transfer economics predict a benefit.

Key constraint: this is hybrid state, not just KV

Nemotron 3 Super contains:

  • 8 attention layers;
  • 40 Mamba recurrent layers; and
  • 40 MoE layers.

An exact handoff must include every state component required for continuation:

  • attention KV pages;
  • recurrent/SSM state;
  • convolution state;
  • logical token position and sequence metadata;
  • target sampler/session metadata where required;
  • native-MTP context or sufficient state to reprime it exactly;
  • native-MTP session bookkeeping and pending hidden row;
  • committed token history required by the request-local suffix proposer; and
  • an immutable identity tying the payload to the model, package, tokenizer, context configuration, and state schema.

A KV-only handoff is not correct for Nemotron-H.

State-transfer economics

For Nemotron Super with F16 KV:

8 attention layers
× 2 (K and V)
× 2 KV heads
× 128 dimensions
× 2 bytes
≈ 8 KiB per prompt token

Approximate attention-KV payloads:

Prompt length Attention KV
8K 64 MiB
32K 256 MiB
64K 512 MiB
128K 1 GiB
262K 2 GiB

The fixed recurrent/SSM state is roughly another 160 MiB before convolution state, metadata, and MTP-specific state. Exact sizes must be measured from runtime export rather than assumed from this estimate.

At ideal line rate, a roughly 420 MiB 32K-context state takes approximately:

  • 3.5 seconds over 1 Gb/s;
  • 350 ms over 10 Gb/s;
  • 140 ms over 25 Gb/s; or
  • 90 ms over 40 Gb/s.

Streaming append-only KV pages while later prompt chunks compute can hide most of the bulk transfer. The final recurrent snapshot, MTP state/bookkeeping, import, and commit remain on the uncovered handoff tail.

The current binary protocol is not sufficient as the production transport:

  • StateImport and StateExport message kinds still exist, but skippy-server rejects them;
  • MAX_STAGE_STATE_IMPORT_BYTES is 512 MiB;
  • useful Nemotron contexts can exceed that limit; and
  • a monolithic Vec<u8> export/import adds allocation and critical-path latency.

Do not simply re-enable the old messages.

Architecture

Roles

Introduce additive phase-role capabilities:

  • prefill;
  • decode; and
  • prefill+decode.

A node may advertise more than one role. Older nodes ignore these optional capabilities and are never selected for a disaggregated topology.

Supported topology shapes

1. Full-prefill to full-decode

Initial correctness implementation:

M3 full target prefill ──► M1 full target decode

Both machines load compatible package revisions. The M1 additionally owns the active MTP/suffix generation path.

2. Split-prefill to collapsed decode

Generalized form:

Prefill stage 0 ─┐
Prefill stage 1 ─┼──► state ranges streamed to decode owner
Prefill stage N ─┘

Each prefill stage exports only the state for its owned layer range. The decode worker may own the complete model and import those ranges into one decode session.

This allows an existing Skippy layer pipeline to compute prefill while avoiding that pipeline during subsequent decode.

The first implementation should prove full-prefill/full-decode before adding split-prefill/collapsed-decode.

Streaming handoff

Attention KV

Attention KV is append-only during prefill and should stream by token page:

  1. Prefill a bounded prompt chunk.
  2. Export completed KV pages for that token range.
  3. Begin sending those pages immediately.
  4. Continue computing the next prompt chunk while transfer proceeds.
  5. Import pages into an uncommitted decoder session as they arrive.

Each chunk is identified by:

  • handoff id;
  • package/manifest identity;
  • state schema version;
  • layer range;
  • token range;
  • KV codec and dtype;
  • sequence position configuration;
  • payload length;
  • checksum; and
  • final/non-final marker.

Use dedicated backpressured QUIC streams. Do not require a network round trip for every KV page.

Recurrent state

Recurrent/SSM state is overwritten as the prompt advances. Sending a full snapshot after every prompt chunk would waste bandwidth.

For the first implementation:

  • stream attention KV incrementally;
  • export and transfer one final recurrent/SSM snapshot after the last prompt chunk; and
  • import it before handoff commit.

Investigate recurrent deltas only if the final snapshot remains a material uncovered tail.

MTP state

Native MTP consumes target hidden rows and maintains its own context/session bookkeeping. The design must inventory whether Nemotron's MTP context:

  • shares all required memory with the target context;
  • has independent attention KV;
  • requires a separate state export/import; or
  • can be deterministically reprised from a bounded final hidden-state payload.

Do not enable MTP immediately after handoff until exact equivalence has been proven.

Possible implementations:

  1. Prefiller loads and primes MTPv2, then transfers target and MTP state.
  2. Decoder receives the required hidden-state stream and primes its MTP context concurrently.
  3. Decoder initially runs target-only while a bounded MTP reprime completes.

Option 3 is acceptable only as an explicitly measured fallback; it must not silently report MTP as active.

Suffix N-gram

Dan's suffix N-gram proposer requires no model-state transfer. Send the committed request token history with the handoff and construct its request-local exact-suffix index on the decoder.

History must remain request-local and contain only prompt and target-committed tokens.

Atomic ownership transfer

Use a two-phase handoff:

  1. Coordinator creates a handoff_id.
  2. Decoder reserves an uncommitted session and returns transfer credit.
  3. Prefiller streams state components.
  4. Prefiller sends a final manifest containing component hashes and logical state.
  5. Decoder validates and imports every required component.
  6. Decoder performs an exact continuation probe where required.
  7. Decoder returns CommitAck.
  8. Coordinator changes request affinity to the decoder.
  9. Prefiller may release transient state.

The decoder must never generate from a partially imported session.

The prefiller retains state until commit acknowledgement so that connection loss can retry or safely fall back.

Prefix-cache interaction

The prefill worker should perform exact-prefix lookup before computing missing prompt chunks.

On a prefix-cache hit:

  • restore the cached target state;
  • compute only the uncached suffix;
  • stream restored/new pages using the same handoff contract; and
  • report cached versus newly computed state separately.

For hybrid families, prefix reuse must include attention KV plus recurrent state. KV-only reuse remains invalid.

The decoder may retain imported state for request affinity, but the design must not introduce a standalone KV server. State remains owned by Skippy runtimes and Mesh routing.

Placement and eligibility

Disaggregation must be cost-based, not always enabled.

The planner should estimate:

handoff_tail =
    uncovered_KV_transfer
  + final_recurrent_transfer
  + MTP_state_transfer_or_reprime
  + import
  + commit_round_trip

benefit =
    decoder_prefill_time_avoided
  + overlap/goodput benefit
  - prefiller contention

Select a disaggregated topology only when:

  • both peers advertise the required state schema and role capabilities;
  • package and tokenizer identities match;
  • the link has recent bandwidth and RTT measurements;
  • estimated state fits decoder memory;
  • estimated handoff beats recompute or the best non-disaggregated topology; and
  • the relevant family has exact state-handoff certification.

Bandwidth should dominate bulk-transfer eligibility. RTT still matters for reservation, final commit, cancellation, and first-token latency.

When the link degrades, stop admitting new disaggregated requests and fall back to a supported local or layer-split topology. Do not migrate an active decoder session without another exact handoff.

Scheduler relationship

This issue owns:

  • cross-role placement;
  • exact state schema;
  • streaming handoff transport;
  • atomic ownership transfer;
  • cost-based routing; and
  • disaggregation correctness/economics.

#1416 owns iteration-level scheduling and concurrency within each worker.

A single-request correctness prototype can precede #1416. Production multi-request cutover requires #1416 or an equivalent scheduler so that:

  • the prefiller can queue/chunk multiple prompts;
  • the decoder can continuously batch active sessions;
  • prefill and decode overlap without one runtime lock serializing both; and
  • KV pressure is handled through admission rather than allocation failure.

Protocol compatibility

Mesh protocol additions must be additive:

  • optional phase-role capability advertisement;
  • optional exact-state schema/codec capabilities;
  • new dedicated handoff stream type;
  • versioned state-component headers; and
  • explicit unsupported-capability fallback.

Never route a disaggregated request through an older peer lacking the capability. Existing mixed-version layer serving must remain functional.

The Skippy native ABI should add only the state streaming/chunking primitives required for exact export/import. Bump the Skippy ABI version and Rust mirror together for any new symbols.

Security and privacy

Transferred state can encode private prompt content.

Requirements:

  • transfer only across authenticated mesh peers;
  • retain QUIC transport encryption;
  • bind payload identity to the request and destination;
  • enforce request/tenant isolation;
  • do not log tokens, state contents, or content-derived hashes;
  • do not persist transferred state by default;
  • bound retained source state and erase it after commit/abort; and
  • reject replayed, cross-model, cross-session, or cross-tenant handoffs.

Telemetry

Emit per handoff:

  • source and destination role;
  • prompt tokens;
  • cache-hit and newly computed tokens;
  • bytes by state component;
  • export, queue, wire, import, and commit time;
  • amount of transfer overlapped with prefill;
  • uncovered handoff-tail time;
  • measured bandwidth and RTT;
  • retry/fallback reason;
  • target and MTP state readiness;
  • first-token delay after prefill; and
  • end-to-end TTFT/TPOT/goodput.

Do not include prompt content or stable content fingerprints.

Correctness gates

Use and extend skippy-correctness state-handoff.

For Nemotron-H, compare local continuation with remote-import continuation across:

  • prompt lengths: 1K, 8K, 32K, 64K;
  • cold and prefix-cache-assisted prefill;
  • target-only decode;
  • MTP-only decode;
  • MTP plus suffix N-gram;
  • cancellation during transfer;
  • connection loss before commit;
  • duplicate/replayed chunks; and
  • mismatched model, tokenizer, state schema, dtype, and layer range.

Deterministic next-token/logit and generated-token output must match the non-disaggregated target.

Add representative gates for dense KV, ISWA, and recurrent/hybrid families before making the mode generic.

Performance matrix

Compare:

  1. M3 all-in-one.
  2. M1 all-in-one.
  3. Existing M1/M3 layer split.
  4. M3 full prefill to M1 full decode.
  5. Split-prefill to M1 collapsed decode.
  6. Disaggregation with exact-prefix cache.
  7. Disaggregation with MTP.
  8. Disaggregation with MTP plus suffix N-gram.

Sweep:

  • prompt lengths;
  • output lengths;
  • concurrency;
  • prefix hit rates;
  • link bandwidth: 1/10/25/40/100 Gb/s;
  • RTT and packet loss;
  • prefill chunk/page sizes; and
  • MTP/suffix proposal settings.

Primary metric: goodput at fixed TTFT and TPOT SLOs.

Also report single-request latency. Disaggregation must not be advertised as universally faster if its benefit exists only with long prompts or concurrent phase overlap.

Cutover gate

Enable automatic disaggregated placement only when:

  1. exact continuation parity passes;
  2. no partially imported state can generate;
  3. overload queues or falls back rather than producing an allocation failure;
  4. the cost model rejects unsuitable links and short prompts;
  5. an eligible high-bandwidth topology beats the best non-disaggregated baseline on goodput at the target SLO;
  6. state-transfer memory remains bounded under concurrency;
  7. MTP and suffix attribution remains correct after handoff; and
  8. link failure produces a safe retry, recompute, or fail-closed response.

Implementation tasks

  • Inventory exact Nemotron target, recurrent, convolution, MTP, and logical session state.
  • Measure payload bytes at representative context lengths.
  • Define versioned component and handoff manifests.
  • Add chunked/streaming native export/import APIs where current page APIs are insufficient.
  • Add the dedicated backpressured QUIC handoff stream.
  • Implement decoder reservation and two-phase commit.
  • Implement full-prefill to full-decode correctness prototype.
  • Stream attention KV pages concurrently with prefill chunks.
  • Transfer the final recurrent snapshot.
  • Prove or implement MTP context handoff/reprime.
  • Rebuild request-local suffix history on the decoder.
  • Integrate prefiller exact-prefix cache.
  • Add bandwidth/RTT measurement and cost-based placement.
  • Add failure, cancellation, retry, and replay handling.
  • Add full correctness and performance matrices.
  • Implement split-prefill to collapsed-decode as the generalized second phase.
  • Add Design: from-scratch iteration-level scheduler for concurrent staged serving #1416 multi-request scheduling integration.

Non-goals

  • A standalone KV server.
  • Unauthenticated state transfer.
  • Assuming KV-only state is exact for hybrid families.
  • Requiring all meshes to use disaggregation.
  • Selecting phase disaggregation over a low-bandwidth or high-latency link.
  • Replacing ordinary layer splitting where it remains faster.
  • Combining the Design: from-scratch iteration-level scheduler for concurrent staged serving #1416 scheduler implementation into this issue.
  • Persisting prompt-derived state by default.

Dependencies and references

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions