Repository navigation
Conversation
1 task done
Contributor
📝 WalkthroughWalkthroughAdds a KV Transfer Only Mode
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_kv_transfer_endpoint.py (1)
44-66: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd background-transfer failure coverage.
Coverage is insufficient for
_do_kv_transfer: the stubbed task never exercises remote/local non-2xx responses, faileddisaggregated_paramsextraction,_kvt_seencleanup, or_kvt_inflightdecrement. Add focused tests in this file, e.g.test_do_kv_transfer_counts_failed_on_non_2xxandtest_failed_transfer_allows_retry. As per path instructions, tests undertests/**should be reviewed for coverage sufficiency.Also applies to: 141-158
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/disaggregated/test_kv_transfer_endpoint.py` around lines 44 - 66, Coverage around _do_kv_transfer is missing for failure paths because the current _make_server setup only stubs success and never exercises non-2xx remote/local responses or disaggregated_params extraction failures. Add focused tests in test_kv_transfer_endpoint.py, such as test_do_kv_transfer_counts_failed_on_non_2xx and test_failed_transfer_allows_retry, that drive the real _do_kv_transfer flow and assert _kvt_stats["failed"], _kvt_seen cleanup, and _kvt_inflight decrement on error. Use the existing _make_server helper and related state fields (_kvt_seen, _kvt_inflight, _kvt_stats) to verify retries are possible after a failed transfer.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 4731-4735: The kv_transfer_only request is being marked finished
in py_executor.py but still remains in scheduled_batch.generation_requests, so
the batch can continue into resource prep and _forward_step. Update the batching
flow around the request-handling logic in py_executor.py to remove these
transfer-only requests from the scheduled batch immediately after finish_by,
then recompute or short-circuit can_queue if nothing remains. In the PP path,
also clear any inflight id that was assigned before this method so the finished
request is fully unscheduled.
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 219-221: The KV transfer source allowlist handling in
openai_server.py currently treats a missing KV_TRANSFER_SOURCE_ALLOW as
unrestricted access, which leaves /v1/kv_transfer open to arbitrary source
hosts. Update the allowlist initialization and the request path in the relevant
server logic so the endpoint fails closed when no allowlist is configured, or is
disabled until an explicit allowlist is provided; make sure the validation uses
the existing _kvt_allow setting and related /v1/kv_transfer handling
consistently.
- Around line 670-759: The KV transfer flow in _do_kv_transfer and the request
de-duplication path in the handler currently treat failures as dedupes and can
mark failed POSTs as successful transfers. Update the logic so _kvt_seen[h] is
cleared or not retained when the transfer fails, verify both the source and
local POST responses before incrementing "transferred", and only record success
after those checks pass. Also narrow the broad exception handling in
_do_kv_transfer to the expected JSON, HTTP, and timeout cases so genuine bugs
are not swallowed.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_kv_transfer_endpoint.py`:
- Around line 44-66: Coverage around _do_kv_transfer is missing for failure
paths because the current _make_server setup only stubs success and never
exercises non-2xx remote/local responses or disaggregated_params extraction
failures. Add focused tests in test_kv_transfer_endpoint.py, such as
test_do_kv_transfer_counts_failed_on_non_2xx and
test_failed_transfer_allows_retry, that drive the real _do_kv_transfer flow and
assert _kvt_stats["failed"], _kvt_seen cleanup, and _kvt_inflight decrement on
error. Use the existing _make_server helper and related state fields (_kvt_seen,
_kvt_inflight, _kvt_stats) to verify retries are possible after a failed
transfer.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5c9c2f76-ed1c-451f-92de-24009edb0af1
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/disaggregated_params.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_server.pytests/unittest/disaggregated/test_kv_transfer_endpoint.py
added 2 commits
June 29, 2026 22:46
…stances Adds a general `kv_transfer_only` flag on generation_only disaggregated requests: the destination instance pulls the prompt's KV from a source instance over the disagg transceiver, commits it for reuse, and finishes WITHOUT a generation forward pass. General mechanism (cache migration/ rebalancing, replication). The two instances just need the cache transceiver enabled (not limited to a dedicated context/generation pair, need not be co-located). Plus a fire-and-forget `POST /v1/kv_transfer` endpoint (+`GET /v1/kv_transfer/stats`; dedup, bounded concurrency, optional source allowlist). Engine changes (all additive): - disaggregated_params / openai_protocol: kv_transfer_only field + passthrough. - py_executor _prepare_disagg_gen_transmission_complete: when kv_transfer_only, finish via finish_by(LENGTH, 0) after normal sampler setup -> no forward pass, KV committed for reuse. - py_executor _handle_first_token_response: skip already-finished requests so a kv_transfer_only req emits a single is_final result via the terminal path (mirrors the is_finished skip in _emit_first_token_responses). No-op for stock requests. kv_transfer_only requests compose with the existing idle-only transfer-progress blocking, so they pull and commit without stalling live traffic. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
Covers request handling for the kv_transfer endpoint: bad JSON / missing fields -> 400, source allowlist -> 403, accept + counters, dedup within TTL, in-flight cap -> skipped_busy, and the /stats counters. The background transfer is stubbed, so no engine/GPU is required. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
nafis271
force-pushed
the
kv-transfer-api
branch
from
June 29, 2026 22:46
13b14eb to
2fc17f8
Compare
…allowlist, robust transfer failures - py_executor: drop finished kv_transfer_only requests from the forward batch when other work remains, so they don't run a wasted decode; keep them when they are the whole batch (empirically their own forward is needed to finalize the reuse commit, and removing them would leave an empty batch the overlap loop mishandles). - openai_server.kv_transfer: fail closed -- require KV_TRANSFER_SOURCE_ALLOW and 403 when unset, since the endpoint makes the server POST to `source` (SSRF guard). - openai_server._do_kv_transfer: check both POST responses before counting a transfer as successful; clear the dedup entry on failure so a transient error can retry before the TTL; narrow the except to HTTP/JSON/timeout; add docstrings. - tests: add a fail-closed (no allowlist -> 403) case. Signed-off-by: Md Nafis Ul Haque Shifat <nafis@deepinfra.com>
nafis271
force-pushed
the
kv-transfer-api
branch
from
June 30, 2026 00:48
65a02d3 to
3490f71
Compare
Member
JunyiXu-nv
requested changes
Oct 26, 2026
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
Collaborator
There was a problem hiding this comment.
Requesting changes: the branch no longer merges (the auto-merged py_executor.py doesn't even parse), and the main capability looks covered by #13055 already. I'll hold the detailed comments until that's settled.
Comment on lines
+4732
to
+4737
| # kv_transfer_only: transfer done + first token adopted; KV already | ||
| # committed for reuse above, so finish now and record for the trim below. | ||
| if (req.py_disaggregated_params is not None and getattr( | ||
| req.py_disaggregated_params, "kv_transfer_only", None)): | ||
| req.finish_by(FinishReason.LENGTH, 0) | ||
| kv_transfer_only_done.append(req.py_request_id) |
Collaborator
There was a problem hiding this comment.
This block no longer applies on main: _prepare_disagg_gen_transmission_complete was restructured there (the loop body is one level shallower and goes through self.disagg.try_finish_gen_receive(req)). Git still merges py_executor.py without a conflict, but the result does not parse: this if ends up one level deeper than the loop body (IndentationError, checked with git merge-tree against current main). openai_protocol.py and openai_server.py conflict too. Please rebase and re-validate the executor change on main.
| return self.generator._check_health() | ||
| return True | ||
|
|
||
| async def kv_transfer(self, request: Request): |
Collaborator
There was a problem hiding this comment.
#13055 merged on 07-24, after this was opened: a server can now pull a prompt's KV straight from another server's reuse tree with a generation_only request built from GET /v1/data_transceiver_state, without a context_only prefill on the source, and test_arbitrary_kv_cache_transfer checks the reuse hit using a plain generation_only with max_tokens=1. Tabrizian asked about this on 06-30. What do this endpoint and kv_transfer_only add on top of that (e.g. prompts not yet cached on the source)? If it is mostly orchestration, could it be rebuilt on the merged path?
1 task done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Implements #15735.
Builds on the disagg transceiver to add a general, cross-instance KV-transfer mechanism. A
generation_onlyrequest with a newkv_transfer_onlyflag pulls the prompt's KV from asource instance over the transceiver, commits it to the local reuse cache, and finishes
immediately, producing no output tokens. The two instances just need the cache transceiver enabled — they
are not limited to a dedicated context/generation pair, and need not be co-located (UCX does
cross-node over IB too).
New API
POST /v1/kv_transfer— accepts the request, starts the transfer in the background, and returns202immediately (does not wait for the transfer to finish).source(required) — base URL of the instance to pull the prompt's KV from.messagesorprompt(one required) — the prompt, in chat or completions form.model(optional) — defaults to the server's served model.202with{status, request_hash}, wherestatusis one ofaccepted,deduped(an identical prompt was seen within the dedup TTL), or
skipped_busy(the in-flight transfer capwas reached). Errors:
400(missingsourceor prompt),403(no allowlist configured, orsourcenot on it).excess as
skipped_busy; and requires asourceallowlist — the endpoint fails closed (403)until
KV_TRANSFER_SOURCE_ALLOWis set, to avoid an SSRF primitive. Tunables:KV_TRANSFER_MAX_INFLIGHT,KV_TRANSFER_DEDUP_TTL,KV_TRANSFER_SOURCE_ALLOW.GET /v1/kv_transfer/stats— observability counters:fired,transferred,failed,deduped,skipped_busy, plus liveinflight,seen(dedup-table size), andmax_inflight.How it works
It composes the existing disaggregated-serving interface rather than adding new transport. The
endpoint runs the two-phase disagg flow internally: the
context_onlyleg makessourceprefill theprompt and return its disagg ticket (the opaque transceiver handle for that prompt's KV), and
replaying that ticket back to this instance as a
generation_onlyrequest (withkv_transfer_only)makes the engine pull the KV over the transceiver and commit it for reuse. The internal
generation_onlyleg is issued over loopback HTTP so it goes through the normal serving path(validation,
disaggregated_paramshandling, executor submission).Changes (all additive; regular disagg/serving paths unchanged)
disaggregated_params.py/serve/openai_protocol.py— add thekv_transfer_onlyfield andpass it through the request converter.
_torch/pyexecutor/py_executor.py_prepare_disagg_gen_transmission_complete: whenkv_transfer_only, commit the KV andfinish_by(LENGTH, 0)after the normal sampler setup (skipping that setup corrupts the batchedsampler under concurrency), then drop the request from the forward batch to skip a wasted
decode — but only when other work remains; a lone transfer keeps its (output-discarded) forward,
which is what finalizes the reuse commit (and avoids an empty batch).
_handle_first_token_response: skip already-finished requests. Akv_transfer_onlyreq isfinished at transmission-complete, so without this it would emit a first-token response and the
terminal response — two
is_finalresults for oneclient_id. The skip mirrors the existingis_finishedguard in_emit_first_token_responses, and is a no-op for stock requests (whichare never finished at first-token).
serve/openai_server.py— thePOST /v1/kv_transferandGET /v1/kv_transfer/statshandlers.The endpoint fails closed without an allowlist (SSRF guard), and
_do_kv_transferchecks both POSTresponses before counting a transfer, clears the dedup entry on failure (so transient errors can
retry), and narrows its
exceptto HTTP/JSON/timeout.kv_transfer_onlyrequests compose with the existing idle-only transfer-progress blocking, so theypull and commit without stalling live traffic.
Test Coverage
Adds
tests/unittest/disaggregated/test_kv_transfer_endpoint.py— unit tests for the endpoint'srequest handling (bad JSON / missing fields → 400, allowlist + fail-closed → 403, accept + counters,
dedup within TTL, in-flight cap →
skipped_busy, and the/statscounters), with the backgroundtransfer stubbed so no engine/GPU is required (10 tests).
The engine-level transfer + reuse path was validated manually on the rc19 release image with
gpt-oss-120b across two instances (one source, one destination): reuse correctness (the transferred KV
is reused on a subsequent request with correct output), zero cross-contamination under sustained mixed
load (concurrent unrelated requests + mid-stream cancellations), and no fatal engine errors. Happy to
add an automated 2-GPU disagg integration test following the
tests/unittest/disaggregatedpatternsif maintainers prefer.
(Note vs rc19: rc19 still has an always-on transfer-progress wait, so there
kv_transfer_onlyalsoneeds excluding from that wait;
mainblocks only when idle, so this PR does not touch the blockingpath.)
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added — this PR is api-compatible (the flag and endpoints are additive and off by default).
Any new dependencies have been scanned for license and vulnerabilities — no new dependencies.
CODEOWNERS updated if ownership changes — n/a.
Documentation updated as needed.
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.@coderabbitai summary