From 3017c57825f37c4bfd1d3a8ac6593c3acac41208 Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <itabrizian@nvidia.com>
Date: Mon, 8 Jun 2026 20:23:42 -0700
Subject: [PATCH] [None][perf] disagg: serialize Request input_token_ids as
 int32 bytes

The nanobind Request __getstate__ returned input_token_ids as a Python
list[int] (and __setstate__ cast it back to a vector), so every Request
pickle/unpickle paid a PyLong per token. On the disaggregated generation
path the Request is pickled on hot, GIL-held paths -- the per-iteration
request broadcast across DP ranks and the RpcWorker submit IPC -- so for
long (ISL-sized) prompts this dominated host-side time and starved the
executor loop.

Serialize input_token_ids as a raw little-endian int32 byte buffer in
__getstate__/__setstate__ (vector<->bytes is a memcpy, no per-element
Python objects). The pickle format is symmetric and only consumed within
a single build (broadcast + IPC), so there is no cross-version concern.

Measured (DSv4-Pro disagg, GB300, ISL ~40k):
- Request pickle+unpickle: 1.21 ms -> 0.06 ms (~19x)
- broadcast_requests (GEN executor loop): 54.6 ms -> 6.6 ms avg (~8x)

Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>
---
 .../nanobind/executor/request.cpp             | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp
index 97b77e2214a2..510bd2fd343b 100644
--- a/cpp/tensorrt_llm/nanobind/executor/request.cpp
+++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp
@@ -36,6 +36,7 @@
 #include <nanobind/stl/vector.h>
 #include <sstream>
 
+#include <cstring>
 #include <optional>
 #include <vector>
 
@@ -591,7 +592,15 @@ void initRequestBindings(nb::module_& m)
 
     auto requestGetstate = [](tle::Request const& self)
     {
-        return nb::make_tuple(self.getInputTokenIds(), self.getMaxTokens(), self.getStreaming(),
+        // Serialize input_token_ids as a raw int32 byte buffer instead of a Python
+        // list[int]: nanobind casts VecTokens element-by-element (a PyLong storm,
+        // ISL-proportional) on every Request pickle -- the request broadcast and the
+        // RPC IPC submit. A bytes blob is a memcpy, ~order-of-magnitude cheaper.
+        // Paired with requestSetstate.
+        auto const& inputTokenIds = self.getInputTokenIds();
+        auto inputTokenIdsBytes = nb::bytes(
+            reinterpret_cast<char const*>(inputTokenIds.data()), inputTokenIds.size() * sizeof(VecTokens::value_type));
+        return nb::make_tuple(std::move(inputTokenIdsBytes), self.getMaxTokens(), self.getStreaming(),
             self.getSamplingConfig(), self.getOutputConfig(), self.getEndId(), self.getPadId(), self.getPositionIds(),
             self.getBadWords(), self.getStopWords(), self.getEmbeddingBias(), self.getExternalDraftTokensConfig(),
             self.getPromptTuningConfig(), self.getMultimodalInput(), self.getMultimodalEmbedding(),
@@ -608,8 +617,12 @@ void initRequestBindings(nb::module_& m)
         {
             throw std::runtime_error("Invalid Request state!");
         }
-        new (&self) tle::Request(nb::cast<VecTokens>(state[0]), nb::cast<SizeType32>(state[1]),
-            nb::cast<bool>(state[2]), nb::cast<tle::SamplingConfig>(state[3]), nb::cast<tle::OutputConfig>(state[4]),
+        // input_token_ids is a raw int32 byte buffer (see requestGetstate).
+        auto const inputTokenIdsBytes = nb::cast<nb::bytes>(state[0]);
+        VecTokens inputTokenIds(inputTokenIdsBytes.size() / sizeof(VecTokens::value_type));
+        std::memcpy(inputTokenIds.data(), inputTokenIdsBytes.c_str(), inputTokenIdsBytes.size());
+        new (&self) tle::Request(std::move(inputTokenIds), nb::cast<SizeType32>(state[1]), nb::cast<bool>(state[2]),
+            nb::cast<tle::SamplingConfig>(state[3]), nb::cast<tle::OutputConfig>(state[4]),
             nb::cast<std::optional<SizeType32>>(state[5]), nb::cast<std::optional<SizeType32>>(state[6]),
             nb::cast<std::optional<std::vector<SizeType32>>>(state[7]),
             nb::cast<std::optional<std::list<VecTokens>>>(state[8]),