Skip to content

[None][feat] Dis-agg transceiver mass integration from the DSV4 branch #15222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Shixiaowei02 merged 9 commits into from
Jun 25, 2026
Merged

[None][feat] Dis-agg transceiver mass integration from the DSV4 branch #15222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
72f11eb
[TRTLLM-12403][fix] Fix deepseekv4 stall
chuangz0 Apr 30, 2026
f143f3b
[None][fix] Avoid TP collective deadlock in ADP disagg error path (#1…
Shixiaowei02 May 9, 2026
7331953
[None][perf] Skip transceiver tp_allgather when no sessions ever open…
Shixiaowei02 May 13, 2026
1d5fffe
[None][test] Add dsv4 dis-agg module-level unit tests (#13936)
Shixiaowei02 May 13, 2026
669b2b0
Fix disagg test regressions: SWA clamp, router session, URL assertion
Tabrizian May 13, 2026
8993276
[None][fix] NIXL agent transfer and shutdown race (#14137)
Shixiaowei02 May 18, 2026
0d0aed9
[None][fix] Return 3-tuple from check_gen_transfer_status fast path (…
qiaoxj07 May 19, 2026
209e879
[None][fix] fix token_range_end add extra_kv_num_tokens (#14258)
chuangz0 May 19, 2026
4a08c69
[None][fix] Post-rebase fixes + coverage tests for DeepSeek-V4 disagg…
Shixiaowei02 Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 28 additions & 14 deletions cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,39 +214,53 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
.def_rw("backend_params", &kvc::BaseAgentConfig::backendParams);

// BaseTransferAgent class (abstract base)
// All transfer-engine operations release the GIL: they may block on NIXL /
// UCX / network / GPU memory pinning, and holding the GIL across them
// starves Python listener / progress threads in the same process.
nb::class_<kvc::BaseTransferAgent>(m, "BaseTransferAgent")
.def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs"))
.def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs"))
.def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs"),
nb::call_guard<nb::gil_scoped_release>())
.def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs"),
nb::call_guard<nb::gil_scoped_release>())
.def("load_remote_agent",
nb::overload_cast<std::string const&, kvc::AgentDesc const&>(&kvc::BaseTransferAgent::loadRemoteAgent),
nb::arg("name"), nb::arg("agent_desc"))
nb::arg("name"), nb::arg("agent_desc"), nb::call_guard<nb::gil_scoped_release>())
.def("load_remote_agent_by_connection",
nb::overload_cast<std::string const&, kvc::ConnectionInfoType const&>(
&kvc::BaseTransferAgent::loadRemoteAgent),
nb::arg("name"), nb::arg("connection_info"))
.def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc)
.def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name"))
nb::arg("name"), nb::arg("connection_info"), nb::call_guard<nb::gil_scoped_release>())
.def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc,
nb::call_guard<nb::gil_scoped_release>())
.def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name"),
nb::call_guard<nb::gil_scoped_release>())
.def(
"submit_transfer_requests",
[](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request)
{ return self.submitTransferRequests(request).release(); },
nb::arg("request"), nb::rv_policy::take_ownership, nb::keep_alive<0, 1>())
.def(
"notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message"))
.def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages)
.def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo)
.def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"));
nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard<nb::gil_scoped_release>(),
nb::keep_alive<0, 1>())
.def("notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"),
nb::arg("sync_message"), nb::call_guard<nb::gil_scoped_release>())
.def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages,
nb::call_guard<nb::gil_scoped_release>())
.def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo,
nb::call_guard<nb::gil_scoped_release>())
.def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"),
nb::call_guard<nb::gil_scoped_release>());

#ifdef ENABLE_NIXL
// NixlTransferStatus class - release GIL for blocking operations
nb::class_<kvc::NixlTransferStatus, kvc::TransferStatus>(m, "NixlTransferStatus")
.def("is_completed", &kvc::NixlTransferStatus::isCompleted, nb::call_guard<nb::gil_scoped_release>())
.def("wait", &kvc::NixlTransferStatus::wait, nb::arg("timeout_ms") = -1,
nb::call_guard<nb::gil_scoped_release>());
nb::call_guard<nb::gil_scoped_release>())
.def("get_last_status", &kvc::NixlTransferStatus::getLastStatus)
.def("get_last_status_str", &kvc::NixlTransferStatus::getLastStatusStr);

// NixlTransferAgent class
nb::class_<kvc::NixlTransferAgent, kvc::BaseTransferAgent>(m, "NixlTransferAgent")
.def(nb::init<kvc::BaseAgentConfig const&>(), nb::arg("config"))
.def(nb::init<kvc::BaseAgentConfig const&>(), nb::arg("config"), nb::call_guard<nb::gil_scoped_release>())
.def("shutdown", &kvc::NixlTransferAgent::shutdown, nb::call_guard<nb::gil_scoped_release>())
.def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs"))
.def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs"))
.def("load_remote_agent",
Expand Down
Loading
Loading