# NIXL push-mode KV transfer

The default NIXL connector is **pull-based**: the decode (D) instance
reads KV blocks from the prefill (P) instance via `NIXL READ` after
prefill completes. `NixlPushConnector` adds a **push-based** alternative
in which P writes the KV blocks directly into D's pre-allocated memory
via `NIXL WRITE`.

This document describes the threading, queues, and scheduling
interactions specific to the push design. The pull-mode design is
unchanged; the push connector reuses the same handshake, NIXL agent
setup, and metadata path wherever possible.

## High-level flow

```mermaid
sequenceDiagram
    autonumber
    participant Client
    participant Proxy
    participant DSched as D Scheduler
    participant DWorker as D Worker (main)
    participant DWriter as D Writer
    participant PWriter as P Writer
    participant PWorker as P Worker (main)
    participant PSched as P Scheduler

    Client->>Proxy: POST /v1/completions
    Proxy->>PSched: prefill leg (do_remote_decode=True, max_tokens=1)
    Proxy->>DSched: decode leg (do_remote_prefill=True, P coordinates)

    note over DSched,DWriter: D side - register blocks with P
    DSched->>DSched: update_state_after_alloc, stash registration, arm watchdog
    DSched->>DWorker: build_connector_meta -> meta.push_registrations
    DWorker->>DWriter: enqueue (req_id, reg_data) on _reg_send_inbox
    DWriter->>PWriter: NIXL send_notif PUSH_REG msgpack

    note over PSched,PWriter: P side - prefill, stage finished blocks
    PSched->>PSched: request_finished, stash blocks
    PSched->>PWorker: build_connector_meta -> meta.push_finished_blocks
    PWorker->>PWriter: enqueue (req_id, blocks) on _finished_blocks_inbox

    note over PWriter: P writer matches and WRITEs
    PWriter->>PWriter: get_new_notifs returns PUSH_REG, route via _handle_push_reg_notif
    alt PUSH_REG and finished blocks both present
        PWriter->>PWriter: pop matching pair, fire WRITE
    else only one side present
        PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched
    end
    PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE)
    PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake
    PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif

    note over DWorker,DWriter: D side - completion accounting
    DWriter-->>DWorker: forward HB and completion notifs via _pending_completion_notifs
    DWorker->>DWorker: _get_new_notifs drains, HB extends lease, completion marks recv done
    DWorker->>DSched: update_connector_output(finished_recving)
    DSched->>DSched: clear watchdog deadline

    note over PWorker,PWriter: P side - reclaim
    PWorker->>PWorker: get_finished, drain _sending_transfers, queue eviction
    PWriter->>PWriter: drain _evict_finished_inbox, drop stale state
    PWorker->>PSched: update_connector_output(finished_sending)
    PSched->>PSched: free lease

    DWorker-->>Proxy: stream decode tokens
    Proxy-->>Client: response
```

## Threads

``NixlPushConnectorWorker`` introduces a single dedicated background
thread per worker (i.e. per TP rank), named ``nixl-push-writer``.
Each owns the new push-specific NIXL operations on its rank:

* ``nixl_wrapper.get_new_notifs()`` — receive notifications.
* ``nixl_wrapper.send_notif(...)`` for the ``PUSH_REG:<msgpack>`` (D
  side) and for the per-WRITE completion notif (P side).
* ``nixl_wrapper.make_prepped_xfer(...) / transfer(...)`` — submit the
  WRITE itself.

Heartbeats continue to go out from the engine main thread via the
existing base-worker ``_send_heartbeats`` plumbing inside
``start_load_kv``.

### Wake model

The writer thread blocks on ``_push_writer_wake`` (a
``threading.Event``) when it has no work. Three callers set the
event:

1. **``start_load_kv``** (worker main thread, called once per engine
   step with the scheduler's metadata) — sets the wake only when the
   step actually hands the writer new work, i.e. when
   ``meta.push_registrations`` or ``meta.push_finished_blocks`` is
   non-empty. This is the wake for new transfers.
2. **``get_finished``** (worker main thread, called once per engine
   step to report completions) — always sets the wake. The writer is
   the sole consumer of ``nixl_wrapper.get_new_notifs()`` for push,
   so this gives it a chance to drain inbound notifs (heartbeats from
   D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
   even when there is no new metadata to act on.
3. **Handshake-completion callback** (background handshake executor
   thread) — both handshakes run on the executor and never block the
   writer; their done-callbacks re-enqueue the deferred op and set the
   wake, since neither ``send_notif`` nor the NIXL WRITE may run off the
   writer thread:
    * the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the
      registration onto ``_reg_send_inbox``;
    * the **P→D** handshake (before a WRITE) re-enqueues the matched
      ``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``.

   On this second pass ``_ensure_handshake`` returns ``None`` (the agent
   is now connected), so the writer sends the ``PUSH_REG`` / issues the
   WRITE directly. If a handshake *failed*, the callback fails or drops
   the request instead of re-enqueuing, so there is no retry loop (see
   Failure handling).

In addition to event-driven wakes, the writer self-polls at
``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
finished blocks waiting for an unmatched ``PUSH_REG``.

When a request completes on P (lease expires or the WRITE finishes),
``get_finished`` enqueues the request id onto ``_evict_finished_inbox``,
which the writer drains to drop stale ``_push_finished_blocks`` /
``_pending_d_registrations`` and stop self-polling.

## Writer-local matching tables

| Table                          | Owner            | Holds                                                                  |
|--------------------------------|------------------|------------------------------------------------------------------------|
| `_pending_d_registrations`     | writer           | D registrations received from a remote D, waiting for P's blocks       |
| `_push_finished_blocks`        | writer           | P blocks staged by the scheduler, waiting for a remote D registration  |

Either side can arrive first. The writer matches in both directions:
when a ``PUSH_REG`` arrives we look up ``_push_finished_blocks``, and
when finished blocks arrive we look up ``_pending_d_registrations``.
Both lookups try an exact ``request_id`` match first, then fall back
to comparing the ids after stripping the trailing per-engine random
suffix (via ``get_base_request_id``). The fallback exists because the
proxy hands the same ``X-Request-Id`` to both legs, so P and D wrap it
into the same ``cmpl-<uuid>-<index>`` form and differ only by the
8-hex randomization suffix that ``input_processor.assign_request_id``
appends per engine. Stripping just that suffix normalizes both sides
to the same id while preserving the completion index (so multi-prompt
sub-requests stay distinct). It also works whether or not
``VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`` is set, which matters since
that env var is slated for removal upstream.

## Wire format

A push registration is sent as a NIXL notification:

```text
PUSH_REG:<msgpack-encoded dict>
```

Fields in the dict:

| Field                | Set by | Meaning                                                                |
|----------------------|--------|------------------------------------------------------------------------|
| ``request_id``       | D      | D's own vLLM request id; P's match key, echoed in the completion notif |
| ``decode_engine_id`` | D      | D's engine id (P uses this for the reverse handshake)                  |
| ``decode_host``      | D      | D's NIXL side-channel host                                             |
| ``decode_port``      | D      | D's NIXL side-channel port                                             |
| ``decode_tp_size``   | D      | D's tensor-parallel size                                               |
| ``local_block_ids``  | D      | per-group lists of D's *logical* block ids (preallocated)              |
| ``remote_engine_id`` | D      | P's engine id (for the existing P-side handshake)                      |
| ``remote_host``      | D      | P's NIXL side-channel host                                             |
| ``remote_port``      | D      | P's NIXL side-channel port                                             |
| ``remote_tp_size``   | D      | P's tensor-parallel size                                               |

D ships **logical** block ids; P expands them to physical block ids at
WRITE-submission time using the ratio learned during the NIXL
handshake (`remote_physical_blocks_per_logical`). This matches the
pull-mode contract — schedulers ship logical ids, workers expand to
physical at submission.

The completion notif sent from P to D after a WRITE is the existing
`<request_id>:<tp_size>` format used in pull mode (here ``request_id``
is D's own request id, taken from the registration), so the D-side
accounting code is unchanged.

## Scheduler-side responsibilities

`NixlPushConnectorScheduler` extends the base scheduler with:

* **D side** — `update_state_after_alloc` stashes registration data in
  `_push_pending_registrations` and arms a soft watchdog
  (`_push_registration_deadlines`). `build_connector_meta` drains the
  stash into `meta.push_registrations` and any expired entries are
  dropped with a warning.
* **P side** — `request_finished` stashes block IDs in
  `_finished_request_blocks` (for the lease and for
  `has_pending_push_work`) and `_newly_finished_push_blocks` (for the
  next worker step via `meta.push_finished_blocks`).
* **Both sides** — `has_pending_push_work` keeps the engine main loop
  stepping while there is in-flight push state, so the writer always
  gets at least one wake per step.

`update_connector_output`:

* `finished_sending` (P side) clears the lease entry.
* `finished_recving` (D side) clears the watchdog deadline.

## Timeouts and watchdogs

Two per-request timers are armed on the scheduler:

* **D-side registration watchdog** — ``_push_registration_deadlines``.
  If a registered request does not see a push completion within
  ``push_registration_timeout`` seconds (defaults to
  ``decoder_kv_blocks_ttl``), ``build_connector_meta`` drops the stale
  registration and the pending entry, logs a warning, and stops trying
  to resend the registration. The corresponding request remains tracked
  in ``_reqs_need_recv``; it is the engine's request-level abort path
  (or the user / proxy timing out the HTTP call) that ultimately fails
  the request.
* **P-side block lease** — same ``_kv_lease_duration`` used by pull
  mode. ``request_finished`` sets the expiration in ``_reqs_need_send``
  and ``update_connector_output(finished_sending=...)`` clears it on
  successful WRITE. Stale leases are reaped by ``get_finished`` in the
  base worker, which then enqueues the eviction onto
  ``_evict_finished_inbox`` so the writer also stops self-polling.

## Failure handling

* **D-side handshake failure (P→D handshake before sending PUSH_REG)** —
  the future's done-callback calls ``_handle_failed_transfer(rid, None)``,
  which marks D's pre-allocated blocks invalid and enqueues onto
  ``_failed_recv_reqs`` so the next ``get_finished`` reports the
  request as a failed recv. Same recv-side accounting as pull mode.
* **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
  identical handling: ``_handle_failed_transfer`` marks the recv as
  failed.
* **P-side handshake failure (P→D handshake before a WRITE)** — the
  future's done-callback logs ``push_handshake_failed`` and drops the
  request without re-queuing. It deliberately does *not* call
  ``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to
  invalidate on the producer side, same reasoning as the WRITE-submission
  failure below). P's blocks are reclaimed by the ``_kv_lease_duration``
  lease and D's stale registration by its watchdog.
* **P-side WRITE submission failure** — the WRITE handle (if any) is
  released and ``xfer_stats.record_failed_transfer()`` bumps the
  failure counter. We deliberately do not call
  ``_handle_failed_transfer`` here: ``req_id`` on the P side has no
  entry in ``_recving_metadata`` (P is not the receiver), so the
  helper would put a P-local request id into ``_failed_recv_reqs``
  and trip the assertion in the base worker's ``get_finished``. The
  outbound WRITE is dropped on the floor; D's lease watchdog handles
  the missing completion.

## Summary

The push design is a small, well-contained extension on top of the
existing NIXL connector:

* one new connector class, one new scheduler class, one new worker
  class — all subclasses of the existing base classes;
* one dedicated background thread per worker;
* a few cross-thread queues, each with a single consumer (the writer);
  most have one producer, except the two replay queues fed by both the
  engine main thread and a handshake-completion callback:
  ``_reg_send_inbox`` (registrations replayed after their D→P handshake)
  and ``_deferred_push_inbox`` (matched pushes replayed after their P→D
  handshake);
* one new notification type (`PUSH_REG:<msgpack>`).

Behavior on the engine main thread is otherwise unchanged. The writer
thread is event-driven and idle when there is no push work.
