P/D handoff eagerly allocates KV state on the decode worker and can abort on insufficient NPU memory

[Veilwalker]

Summary

In P/D disaggregation mode, in the non-prefix-caching handoff path, when a request finishes prefill, it is transferred to a decode scheduler through:

Router.transfer_prefill_request() -> Scheduler.add_decode()

The decode side then immediately allocates the request's prefilled KV state on the selected decode worker's NPU memory. If the selected decode worker does not have enough free NPU
memory, MemoryModel.allocate(..., Device.NPU) raises a RuntimeError, causing the simulator to abort.

This seems to happen before the handoff path verifies whether the selected decode worker can actually admit the request.

Observed behavior

The relevant path appears to be:

• serving/__main__.py

  • finished prefill requests are passed to router.transfer_prefill_request(finished_reqs)

• serving/core/router.py

  • transfer_prefill_request() selects a decode scheduler and calls sched.add_decode(req)

• serving/core/scheduler.py

  • add_decode() appends the request to the decode queue
  • then allocates the request's full KV on the decode worker

• serving/core/memory_model.py

  • allocate(..., Device.NPU) raises if the NPU memory is insufficient

In the non-prefix-caching path, the code appears to be:

def add_decode(self, req):
    req.instance_id = self.instance_id
    self.request.append(req)
    ...
    kv_size = self.memory.get_total_kv(req)
    self.memory.allocate(kv_size, Device.NPU)

MemoryModel.allocate() then raises when the requested size exceeds available NPU memory:

if self.npu_used + size > self.npu_mem:
    raise RuntimeError(...)

As a result, the current behavior is:

1. a request finishes prefill
2. a decode worker is selected
3. the full prefilled KV is allocated immediately on that decode worker
4. if the selected decode worker does not have enough free NPU memory, the simulator aborts

Reproduction conditions

This behavior should be reproducible under a P/D configuration where:

• the prefill completion rate exceeds the decode side's ability to admit or drain requests
• decode workers are configured with relatively limited available NPU memory
• the workload includes sufficiently long prompts, resulting in large prefilled KV states during handoff

Under these conditions, once a prefilled request is transferred to a decode worker without enough free NPU memory, the simulator aborts in MemoryModel.allocate().

Question

I would like to clarify whether this behavior is expected in the current simulator design.

Specifically, should decode-side memory exhaustion during P/D handoff be treated as a fatal simulation error, or should it be considered a missing admission-control / backpressure mechanism in the P/D handoff path?

I am asking because decode-side memory pressure seems like a natural condition in memory-constrained P/D disaggregation experiments. If the simulator terminates immediately in this case, it becomes difficult to study handoff queuing, decode-side congestion, or memory-aware decode scheduling.

If this is a known limitation rather than intended behavior, would future support for a recoverable P/D handoff path be within the project scope, such as decode-side admission checks, retrying another decode worker, or delaying the handoff until decode memory becomes available?

[JaehongCS20]

Ah, I see. Thanks for pointing this out — I think I had not fully considered this case.

I agree that if the selected decode worker does not have enough memory, the request should wait in the decode scheduler queue instead of aborting the simulation.

One question I have is where the prefilled KV cache should reside while the request is waiting. Since the prefill side produces the KV cache and the decode side needs to receive it before decoding, the KV cache cannot be allocated on the decode worker if there is not enough decode-side memory yet.

Then, should the KV cache remain on the prefill worker until the decode worker can admit the request?

Also, if a retry policy is introduced, would that mean retrying admission or placement on another decode worker using the already-prefilled KV cache, or would the request need to restart from the prefill stage?

I think this detail matters because it affects memory accounting and backpressure. For example, if the KV cache remains on the prefill worker, decode-side congestion would also keep prefill-side memory occupied.

[Veilwalker]

Thanks for the explanation. I think one possible direction is to model the handoff more like vLLM’s disaggregated prefill path.

In vLLM, the prefill worker does not have to eagerly materialize the KV cache on the decode worker at the moment prefill finishes. Instead, the prefill side can keep the generated KV blocks pinned and return metadata describing where those blocks are. The decode side only allocates local KV space when its scheduler can actually admit the request. Until then, the request can stay in a handoff/decode waiting queue.

So for this simulator, maybe the handoff could be split into two steps:

1. Prefill finishes:
   The KV remains accounted on the prefill worker, and the request enters a pending handoff state.

2. Decode admission succeeds:
   Once the selected decode worker has enough KV memory, it reserves local KV space and loads the corresponding KV cache from the prefill worker. After the KV is loaded, the request becomes ready for decode, and the prefill-side KV cache can be released.

This would answer the memory-location question as: yes, by default the KV should remain on the prefill worker while the decode worker cannot admit the request. That also naturally models backpressure, because decode-side congestion keeps prefill-side KV memory occupied.

For retry policy, I think retrying decode admission or trying another decode worker should reuse the already-prefilled KV. Restarting from prefill should probably only be a fallback if the prefill-side KV is evicted, the lease expires, or no decode worker can fit the request at all.

A minimal version could keep the current decode placement policy and simply defer decode-side KV allocation until admission succeeds. A more advanced version could add memory-aware decode placement or retry another decode worker before waiting.