[https://nvbugs/6104831][fix] Disaggregated KV transfer: lifecycle, cancellation, and quiescence hardening

[chienchunhung]

@coderabbitai summary

tl;dr

This PR fixes a permanent disaggregated-serving wedge triggered by long-prompt, cancellation-heavy traffic in the PyTorch backend. The root cause is a stack of cleanup/lifetime invariant gaps across C++ KV transfer, Python request termination, NIXL cancellation, and block reuse. The final solution makes disaggregated KV cleanup transport-aware and block-reuse-safe: cancelled requests stop scheduling immediately, but KV resources are only released after C++ transfer status becomes terminal, and deferred termination is retried exactly once when it becomes safe.

Background

Disaggregated serving splits a request across:

• a context worker that runs prefill and sends KV cache;
• a generation worker that receives KV cache and decodes tokens;
• a disaggregated frontend that routes OpenAI-compatible requests.

Under long-prompt bursts with client-side cancellations, the old cleanup path could leave the deployment alive but permanently non-responsive: workers stayed up, health endpoints could still return 200, but every post-burst request timed out.

The investigation found that this was not one bug. It was a stack of independent lifecycle gaps. Fixing one exposed the next.

Failure signatures observed

Signature	Symptom	Root cause
#1	Sender-side Broken promise	cancel-after-ready path erased a promise without fulfilling it
#4	Generation event loop blocked forever	checkGenTransferStatus(atLeastNum=1) could call future.get() on an unready future
#5	Receiver-side Broken promise	queued cancel erased receiver promise without fulfilling it
#6	Receiver buffer pool wedge	requestSync() early-return path leaked a recv-buffer slot
#7	First-request SIGSEGV / mutex wedge variants	request lifetime and eval-order bugs after moving from raw pointer to shared_ptr<LlmRequest>
rc13 regression	Block reuse enabled caused no recovery even after the above fixes	Python deferred termination while context transfer was in progress, but later skipped cleanup because it assumed early termination had already succeeded

Correctness model

The central invariant is:

A request can be cancelled at the API level immediately,
but its KV resources cannot be released/reused until:
  1. C++ KV transfer is terminal, or the transfer/buffer state is fail-closed;
  2. transfer buffer ownership has been released or poisoned;
  3. block-reuse/prefix-cache references are accounted for;
  4. Python resource cleanup runs exactly once.

Python request state alone is not sufficient. With block reuse enabled, a cancelled request may still have KV blocks referenced by the reuse tree or pinned by an in-flight context transfer. Freeing those blocks too early risks memory corruption; forgetting to free them later causes a permanent wedge.

What this PR changes

1. C++ transfer request lifetime

The C++ transceiver now uses std::shared_ptr<LlmRequest> across async send/receive state so Python-side termination cannot destroy a request object while C++ futures/workers still dereference it.

This closes the UAF class that previously manifested as corrupted request IDs and first-request crashes.

2. Sender/receiver cancellation bookkeeping

Cancellation paths now fulfill promises before dropping queued/ready work. This prevents std::future_error: Broken promise from escaping on both sender and receiver sides.

3. Non-blocking generation transfer polling

checkGenTransferStatus(atLeastNum=1) no longer blocks indefinitely on an unready future. It skips unready entries during non-blocking polling and lets timeout handling drive cancellation/error state.

4. NIXL transfer cancellation and release

The NIXL path now releases transfer handles on cancellation (TransferStatus::release() → nixlAgent::releaseXferReq()), so cancelled requests do not leave backend transfer state stranded.

5. Eval-order fix after shared_ptr

After moving Response::mRequest to shared_ptr, this pattern became unsafe:

sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp));

C++ does not specify function-argument evaluation order, so std::move(resp) could run before resp.mRequest->mRequestId. This PR materializes reqId before the move.

6. Python idempotency guards

Generation-side disagg initialization and KV receive setup are now guarded by py_request_id, preventing repeated side effects such as duplicate KVCacheManager::addSequence() or duplicate request_and_receive_async() for the same request.

7. Fail-closed behavior for unknown transfer quiescence

If TRT-LLM cannot prove transport quiescence, it fails closed instead of returning buffers to the pool. This includes:

• BufferIndexHolder::poison()
• tri-state ready-signal handling (ready, not-ready, cancelled/unknown)
• poison-on-NIXL-throw in cacheFormatter.cpp
• equivalent MLA formatter hardening in mlaCacheFormatter.cpp
• Python-side deferral/fail-safe handling for unquiesced transfer states

This protects against silent buffer-pool corruption under cancellation races.

8. Deferred cleanup for timed-out context/generation transfer

A successful cancel_request() is not a quiescence proof. The PR therefore defers Python-side resource cleanup after timeout cancellation until C++ transfer status reports the request terminal.

This preserves memory safety: resources are not freed while C++/NIXL may still reference transfer buffers or KV blocks.

9. Block-reuse-safe deferred termination

rc13 enabled block reuse with the overlap scheduler, which exposed a new liveness hole.

Before the final fix:

_handle_responses() decides request is done
→ attempts _terminate_request()
→ _terminate_request() refuses because context KV transfer is still in progress
→ later transfer completes
→ AsyncTransferManager.end_transfer() assumes early termination already happened
→ cleanup obligation is lost
→ request/resources remain pinned
→ frontend times out forever

This PR records whether termination was requested and whether resources were actually freed:

RequestTransferMetadata:
  termination_requested
  resources_freed

AsyncTransferManager.end_transfer() now returns:

EndTransferResult(
    completed: bool,
    needs_termination: bool,
)

When transfer becomes terminal, _end_transfer_and_maybe_terminate() retries termination if an earlier attempt was deferred. This preserves both invariants:

• do not free/reuse KV resources while context transfer is still in progress;
• do not drop the cleanup obligation after deferring it.

Why deferring cleanup is intentional

This PR intentionally defers cleanup in some cancellation paths.

That is the correct behavior because cancel_request() only requests cancellation. It does not prove that C++/NIXL has stopped touching transfer buffers or KV blocks.

The final behavior is:

request cancelled
→ stop scheduling it
→ request C++ cancellation once
→ keep transfer/KV resources alive
→ wait for terminal C++ transfer status
→ release/poison transfer buffers
→ unpin block-reuse refs
→ free Python resources exactly once

So the PR does not mean "never clean up." It means "do not clean up until it is safe, and then guarantee cleanup happens."

Tests

New/updated unit coverage includes:

• sender cancel-after-ready promise fulfillment;
• receiver queued-cancel promise fulfillment;
• non-blocking checkGenTransferStatus(atLeastNum=1);
• deferred termination surviving until async transfer completion;
• partial-reuse path still skipping termination when no termination was requested;
• successful early free preventing duplicate termination;
• direct EndTransferResult.needs_termination coverage.

Local test runs:

python -m py_compile tensorrt_llm/_torch/pyexecutor/py_executor.py
pytest -v tests/unittest/_torch/executor/test_async_transfer_manager.py
  9 passed
pytest -v tests/unittest/others/test_kv_cache_transceiver.py -k \
  "test_cancel_request_in_transmission_does_not_break_sender_future or \
   test_check_gen_transfer_status_at_least_one_does_not_block_on_unready_future or \
   test_cancel_queued_gen_request_fulfills_receiver_future"
  6 passed

Notes / limitations

• The direct UCX path has a separate throughput/backpressure boundary under high concurrency. This PR primarily fixes the NIXL customer path and the shared disaggregated cleanup invariants.
• MLA send-path hardening is included for parity with the non-MLA formatter path, but the stress validation above uses Qwen3-0.6B and therefore primarily exercises the non-MLA formatter.
• The local CuTe DSL import guard used during development is not part of this PR.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[chienchunhung]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #47132 [ run ] triggered by Bot. Commit: 484655e Link to invocation