[Core] Proactively free KV cache blocks when aborting finished requests

[jianzs]

This is an enhancement to PR #25067 which ignored aborts on finished requests and relied on timeout-based cleanup. Instead of waiting for the connector timeout to free blocks, immediately free them when receiving FINISHED_ABORTED for an already-finished request.

This enables earlier KV cache memory reclamation, which is especially important under heavy load in multi-node scenarios where memory pressure is high.

[gemini-code-assist]

Code Review

This pull request enhances memory management by proactively freeing KV cache blocks for requests that are aborted after they have already finished. The changes in vllm/v1/core/sched/scheduler.py correctly implement this by adding logic to finish_requests to immediately free blocks for finished requests upon receiving a FINISHED_ABORTED status. The related change to handle a potential race condition in _update_from_kv_xfer_finished is also correct and makes the system more robust. The test case in tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py has been updated appropriately to validate the new behavior. Overall, the changes are well-implemented and achieve the goal of earlier memory reclamation.

[orozery]

This changes basically allows the scheduler to unilateraly break the contract "async-save" contract between the scheduler and the connector.
The contract is:
If the connector returns delay_free_blocks=True when the request finishes, the request blocks will not free until the connector returns the request id in finished_sending.
With this unilateral break of contract, you replace an assert with a warning.
And now it is up to the connector to figure things out on this contract break.

I think a better alternative is to allow the scheduler to report the "abort after request finished" case to the connector, to allow the connector to clean up and make sure both sides agree on cleanup.

[jianzs]

This changes basically allows the scheduler to unilateraly break the contract "async-save" contract between the scheduler and the connector. The contract is: If the connector returns delay_free_blocks=True when the request finishes, the request blocks will not free until the connector returns the request id in finished_sending. With this unilateral break of contract, you replace an assert with a warning. And now it is up to the connector to figure things out on this contract break.

I think a better alternative is to allow the scheduler to report the "abort after request finished" case to the connector, to allow the connector to clean up and make sure both sides agree on cleanup.

@orozery Thanks for the feedback—do you have any thoughts on how to send this message to the connector?

[orozery]

@orozery Thanks for the feedback—do you have any thoughts on how to send this message to the connector?

This will require some changes in the connector API, which we are trying to minimize for now while we are re-iterating its design towards a new connector v2 API.
This is one point we should consider for the v2 API.

For a short-term solution, I think we can handle it this way:

1. At Scheduler.finish_requests, if request.is_finished() set request.status = finished_status so that the Request object will reflect the fact that the request was aborted.
2. At NixlConnectorScheduler.request_finished save the request at a new self._reqs_being_sent: dict[req_id, Request].
3. At NixlConnectorScheduler.build_connector_meta check if any of the requests from self._reqs_being_sent was aborted (by checking their status). For any such aborted request, pass this information to the workers via a new NixlConnectorMetadata.reqs_to_abort: set[ReqId].
4. At NixlConnectorWorker.get_finished add the request IDs of all aborted requests to done_sending, while also performing any necessary cleanups.

cc @NickLucche @njhill

[jianzs]

@orozery Thanks for the detailed feedback! I've implemented your suggested approach:

Implementation Summary:

1. Scheduler.finish_requests: When request.is_finished(), set request.status = FINISHED_ABORTED so the connector can detect this case.

2. NixlConnectorScheduler.request_finished: Check if request.status == FINISHED_ABORTED, then:

   • Add to self._reqs_abort_done: set[ReqId]
   • Clean up internal state (_reqs_not_processed, _reqs_need_send)
   • Return delay_free_blocks=False to not delay further

3. NixlConnectorMetadata: Added reqs_abort_done: set[ReqId] field to pass abort info to workers.

4. NixlConnectorWorker:

   • start_load_kv: Extract reqs_abort_done from metadata and store in self._aborted_reqs
   • get_finished: Add aborted requests to done_sending so the scheduler can free blocks via the normal contract

Key Design Decision:

• A finished request can only exist in self.requests when the connector delays block freeing (P/D scenario)
• Therefore, we use assert self.connector is not None to express this invariant
• Blocks are freed through the normal finished_sending mechanism, maintaining the scheduler-connector contract

Testing:

• Unit tests pass (test_abort_during_kv_transfer validates the flow)
• E2E basic inference tests pass

Please let me know if this aligns with your expectations or if any adjustments are needed.

cc @NickLucche @njhill

[mergify]

Hi @jianzs, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[orozery]

I think we need to remove this.
Connectors may assume that request_finished is called only once.
See #33377 for example.

Instead, NixlConnector will have to poll each of its "being sent" requests to see if their status was changed.