Skip to content

[TRTLLM-4987][feat] Support context logits in TRTLLMSampler #4538

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
dcampora merged 7 commits into from
May 31, 2025
Merged

[TRTLLM-4987][feat] Support context logits in TRTLLMSampler #4538

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
04a1062
First versionl
dcampora May 21, 2025
65956a7
Remove comments.
dcampora May 21, 2025
36bdfb3
Partially support context logits.
dcampora May 21, 2025
2ac1cbf
Revert quickstart advanced.
dcampora May 21, 2025
6b8c066
Fix return context logits.
dcampora May 28, 2025
a4e937e
Remove unsqueeze.
dcampora May 29, 2025
3d75010
Add gather_context_logits.
dcampora May 30, 2025
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
138 changes: 16 additions & 122 deletions cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,121 +102,25 @@ void initBindings(pybind11::module_& m)
.def("get_tokens", py::overload_cast<>(&GenLlmReq::getTokens, py::const_))
.def("get_last_tokens", py::overload_cast<GenLlmReq::SizeType32>(&GenLlmReq::getLastTokens), py::arg("beam"))
.def("get_last_tokens", py::overload_cast<>(&GenLlmReq::getLastTokens))
.def("get_beam_width_by_iter", &GenLlmReq::getBeamWidthByIter, py::arg("for_next_iteration") = false)
.def_property_readonly("max_num_generated_tokens", &GenLlmReq::getMaxNumGeneratedTokens)
.def("add_new_token", &GenLlmReq::addNewToken, py::arg("token"), py::arg("beam"))
.def("add_new_tokens", &GenLlmReq::addNewTokens, py::arg("beam_tokens"))
.def_property_readonly("num_draft_tokens", &GenLlmReq::getNumDraftTokens)
.def("set_generated_tokens", &GenLlmReq::setGeneratedTokens, py::arg("generated_beam_tokens"))
.def("pause", &GenLlmReq::pause, py::arg("max_input_len"))
.def_property("max_sent_token_len", &GenLlmReq::getMaxSentTokenLen, &GenLlmReq::setMaxSentTokenLen)
.def("prompt_embedding_table",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getPromptEmbeddingTable();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def("multimodal_embedding",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getMultimodalEmbedding();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def("get_mrope_rotary_cos_sin",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getMropeRotaryCosSin();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def("bad_words_list",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getBadWordsList();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def_property(
"draft_logits",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getDraftLogits();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
},
[](GenLlmReq& self, at::Tensor& logits)
{ self.setDraftLogits(std::make_optional<GenLlmReq::TensorPtr>(tr::TorchView::of(logits))); })
.def("embedding_bias",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getEmbeddingBias();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def_property(
"lora_config",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getLoraConfig();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
},
[](GenLlmReq& self, at::Tensor& loraConfig)
{ self.setLoraConfig(static_cast<GenLlmReq::TensorPtr>(tr::TorchView::of(loraConfig))); })
.def_property(
"lora_weights",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getLoraWeights();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
},
[](GenLlmReq& self, at::Tensor& loraWeights)
{ self.setLoraWeights(static_cast<GenLlmReq::TensorPtr>(tr::TorchView::of(loraWeights))); })
.def("stop_words_list",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
auto tensor = self.getStopWordsList();
if (tensor)
{
value = tr::Torch::tensor(*tensor);
}
return value;
})
.def_property_readonly("prompt_embedding_table", &GenLlmReq::getPromptEmbeddingTable)
.def_property_readonly("multimodal_embedding", &GenLlmReq::getMultimodalEmbedding)
.def_property_readonly("mrope_rotary_cos_sin", &GenLlmReq::getMropeRotaryCosSin)
.def_property_readonly("bad_words_list", &GenLlmReq::getBadWordsList)
.def_property("draft_logits", &GenLlmReq::getDraftLogits, &GenLlmReq::setDraftLogits)
.def_property_readonly("embedding_bias", &GenLlmReq::getEmbeddingBias)
.def_property("lora_config", &GenLlmReq::getLoraConfig, &GenLlmReq::setLoraConfig)
.def_property("lora_weights", &GenLlmReq::getLoraWeights, &GenLlmReq::setLoraWeights)
.def_property_readonly("stop_words_list", &GenLlmReq::getStopWordsList)
.def_property_readonly("context_logits", &GenLlmReq::getContextLogitsHost)
.def_property_readonly("generation_logits", &GenLlmReq::getGenerationLogitsHost)
.def_property_readonly("prompt_vocab_size", &GenLlmReq::getPromptVocabSize)
.def_property_readonly("mrope_position_deltas", &GenLlmReq::getMropePositionDeltas)
.def_property_readonly("lora_task_id", &GenLlmReq::getLoraTaskId)
Expand Down Expand Up @@ -253,6 +157,8 @@ void initBindings(pybind11::module_& m)
.def("is_last_context_chunk", py::overload_cast<>(&GenLlmReq::isLastContextChunk, py::const_))
.def("is_first_context_chunk", py::overload_cast<>(&GenLlmReq::isFirstContextChunk, py::const_))
.def("get_context_remaining_length", py::overload_cast<>(&GenLlmReq::getContextRemainingLength, py::const_))
.def_property_readonly("context_logits", &GenLlmReq::getContextLogitsHost)
.def_property_readonly("num_draft_tokens", &GenLlmReq::getNumDraftTokens)
.def("set_finished_reason", &GenLlmReq::setFinishedReason, py::arg("finish_reason"), py::arg("beam"))
.def_property_readonly("is_finished", &GenLlmReq::isFinished)
.def_property_readonly("is_finished_due_to_length", &GenLlmReq::isFinishedDueToLength)
Expand Down Expand Up @@ -280,6 +186,7 @@ void initBindings(pybind11::module_& m)
.def_property_readonly("avg_decoded_tokens_per_iter", &GenLlmReq::getAvgDecodedTokensPerIter)
.def_property_readonly("alloc_total_blocks", &GenLlmReq::getAllocTotalBlocksPerRequest)
.def_property_readonly("alloc_new_blocks", &GenLlmReq::getAllocNewBlocksPerRequest)
.def("alloc_context_logits", &GenLlmReq::allocContextLogitsHost, py::arg("vocab_size"), py::arg("logit_dtype"))
.def_property_readonly("reused_blocks", &GenLlmReq::getReusedBlocksPerRequest)
.def_property_readonly("missed_blocks", &GenLlmReq::getMissedBlocksPerRequest)
.def_property_readonly("kv_cache_hit_rate", &GenLlmReq::getKVCacheHitRatePerRequest)
Expand Down Expand Up @@ -311,20 +218,7 @@ void initBindings(pybind11::module_& m)
{
self.setDraftTokens(std::make_shared<GenLlmReq::VecTokens>(draftTokens.value()));
}
})
.def_property(
"context_logits",
[](GenLlmReq& self)
{
std::optional<at::Tensor> value{std::nullopt};
GenLlmReq::TensorPtr const& tensor = self.getContextLogitsHost();
if (tensor)
{
value = tr::Torch::tensor(tensor);
}
return value;
},
[](GenLlmReq& self, at::Tensor& logits) { self.setContextLogitsHost(tr::TorchView::of(logits)); });
});

py::classh<tb::LlmRequest, GenLlmReq>(m, "LlmRequest", pybind11::dynamic_attr())
.def(py::init(
Expand Down
21 changes: 10 additions & 11 deletions cpp/tensorrt_llm/pybind/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,8 +114,14 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m)
.def("get_device", &tr::CudaStream::getDevice);

// Create submodule for executor bindings.
py::module_ executor_submodule = m.def_submodule("executor", "Executor bindings");
tensorrt_llm::pybind::executor::initBindings(executor_submodule);
auto mExecutor = m.def_submodule("executor", "Executor bindings");
auto mInternal = m.def_submodule("internal", "Internal submodule of TRTLLM runtime");
auto mInternalRuntime = mInternal.def_submodule("runtime", "Runtime internal bindings");
auto mInternalTesting = mInternal.def_submodule("testing", "Testing internal bindings");
auto mInternalBatchManager = mInternal.def_submodule("batch_manager", "Batch manager internal bindings");

tensorrt_llm::pybind::executor::initBindings(mExecutor);
tensorrt_llm::pybind::runtime::initBindingsEarly(mInternalRuntime);

auto buildInfo = m.def_submodule("BuildInfo");
buildInfo.attr("ENABLE_MULTI_DEVICE") = py::int_(ENABLE_MULTI_DEVICE);
Expand Down Expand Up @@ -329,6 +335,7 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m)
.def_property_readonly("hidden_size", &tr::ModelConfig::getHiddenSize)
.def_property_readonly("size_per_head", &tr::ModelConfig::getSizePerHead)
.def_property_readonly("data_type", &tr::ModelConfig::getDataType)
.def_property_readonly("speculative_decoding_mode", &tr::ModelConfig::getSpeculativeDecodingMode)
.def_property("head_size", &tr::ModelConfig::getSizePerHead, &tr::ModelConfig::setSizePerHead)
.def_property(
"num_kv_heads_per_layer", &tr::ModelConfig::getNumKvHeadsPerLayer, &tr::ModelConfig::setNumKvHeadsPerLayer)
Expand Down Expand Up @@ -456,11 +463,10 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m)
.def_readwrite("num_return_sequences", &tr::SamplingConfig::numReturnSequences)
.def_readwrite("min_p", &tr::SamplingConfig::minP)
.def_readwrite("beam_width_array", &tr::SamplingConfig::beamWidthArray)
.def_readwrite("normalize_log_probs", &tr::SamplingConfig::normalizeLogProbs)
.def(py::pickle(SamplingConfigGetState, SamplingConfigSetState))
.def("__eq__", &tr::SamplingConfig::operator==);

py::bind_vector<std::vector<tr::SamplingConfig>>(m, "VectorSamplingConfig");

m.def("make_sampling_config", &makeSamplingConfig, py::arg("configs"));

py::class_<tr::GptJsonConfig>(m, "GptJsonConfig")
Expand Down Expand Up @@ -548,15 +554,8 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m)
.def_property_readonly("pinned", &tr::MemoryCounters::getPinned)
.def_property_readonly("uvm", &tr::MemoryCounters::getUVM);

auto mInternal = m.def_submodule("internal", "Internal submodule of TRTLLM runtime");

auto mInternalRuntime = mInternal.def_submodule("runtime", "Runtime internal bindings");
tensorrt_llm::pybind::runtime::initBindings(mInternalRuntime);

auto mInternalTesting = mInternal.def_submodule("testing", "Testing internal bindings");
tensorrt_llm::pybind::testing::initBindings(mInternalTesting);

auto mInternalBatchManager = mInternal.def_submodule("batch_manager", "Batch manager internal bindings");
tpb::initBindings(mInternalBatchManager);
tb::kv_cache_manager::KVCacheManagerBindings::initBindings(mInternalBatchManager);
tb::BasePeftCacheManagerBindings::initBindings(mInternalBatchManager);
Expand Down
65 changes: 63 additions & 2 deletions cpp/tensorrt_llm/pybind/common/customCasters.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,6 @@
// Opaque bindings
PYBIND11_MAKE_OPAQUE(tensorrt_llm::batch_manager::ReqIdsSet)
PYBIND11_MAKE_OPAQUE(std::vector<tensorrt_llm::batch_manager::SlotDecoderBuffers>)
PYBIND11_MAKE_OPAQUE(std::vector<tensorrt_llm::runtime::decoder_batch::Request>)
PYBIND11_MAKE_OPAQUE(std::vector<tensorrt_llm::runtime::SamplingConfig>)

// Custom casters
namespace PYBIND11_NAMESPACE
Expand Down Expand Up @@ -204,5 +202,68 @@ struct type_caster<tensorrt_llm::executor::Tensor>
}
};

template <>
struct type_caster<tensorrt_llm::runtime::ITensor::SharedPtr>
{
public:
PYBIND11_TYPE_CASTER(tensorrt_llm::runtime::ITensor::SharedPtr, _("torch.Tensor"));

// Convert PyObject(torch.Tensor) -> tensorrt_llm::runtime::ITensor::SharedPtr
bool load(handle src, bool)
{
PyObject* obj = src.ptr();
if (THPVariable_Check(obj))
{
at::Tensor const& t = THPVariable_Unpack(obj);
value = std::move(tensorrt_llm::runtime::TorchView::of(t));
return true;
}
return false;
}

// Convert tensorrt_llm::runtime::ITensor::SharedPtr -> PyObject(torch.Tensor)
static handle cast(
tensorrt_llm::runtime::ITensor::SharedPtr const& src, return_value_policy /* policy */, handle /* parent */)
{
if (src == nullptr)
{
return none().release();
}
return THPVariable_Wrap(tensorrt_llm::runtime::Torch::tensor(src));
}
};

template <>
struct type_caster<tensorrt_llm::runtime::ITensor::SharedConstPtr>
{
public:
PYBIND11_TYPE_CASTER(tensorrt_llm::runtime::ITensor::SharedConstPtr, _("torch.Tensor"));

// Convert PyObject(torch.Tensor) -> tensorrt_llm::runtime::ITensor::SharedConstPtr
bool load(handle src, bool)
{
PyObject* obj = src.ptr();
if (THPVariable_Check(obj))
{
at::Tensor const& t = THPVariable_Unpack(obj);
value = std::move(tensorrt_llm::runtime::TorchView::of(t));
return true;
}
return false;
}

// Convert tensorrt_llm::runtime::ITensor::SharedConstPtr -> PyObject(torch.Tensor)
static handle cast(tensorrt_llm::runtime::ITensor::SharedConstPtr const& src, return_value_policy /* policy */,
handle /* parent */)
{
if (src == nullptr)
{
return none().release();
}
return THPVariable_Wrap(tensorrt_llm::runtime::Torch::tensor(
reinterpret_cast<tensorrt_llm::runtime::ITensor::SharedPtr const&>(src)));
}
};

} // namespace detail
} // namespace PYBIND11_NAMESPACE
39 changes: 21 additions & 18 deletions cpp/tensorrt_llm/pybind/runtime/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,23 +216,6 @@ void initBindings(pybind11::module_& m)
.def(py::init<tr::BufferManager::CudaStreamPtr, bool>(), py::arg("stream"), py::arg("trim_pool") = false)
.def_property_readonly("stream", &tr::BufferManager::getStream);

py::class_<tr::SpeculativeDecodingMode>(m, "SpeculativeDecodingMode")
.def(py::init<tr::SpeculativeDecodingMode::UnderlyingType>(), py::arg("state"))
.def_static("NoneType", &tr::SpeculativeDecodingMode::None)
.def_static("DraftTokensExternal", &tr::SpeculativeDecodingMode::DraftTokensExternal)
.def_static("Medusa", &tr::SpeculativeDecodingMode::Medusa)
.def_static("LookaheadDecoding", &tr::SpeculativeDecodingMode::LookaheadDecoding)
.def_static("ExplicitDraftTokens", &tr::SpeculativeDecodingMode::ExplicitDraftTokens)
.def_property_readonly("is_none", &tr::SpeculativeDecodingMode::isNone)
.def_property_readonly("is_draft_tokens_external", &tr::SpeculativeDecodingMode::isDraftTokensExternal)
.def_property_readonly("is_medusa", &tr::SpeculativeDecodingMode::isMedusa)
.def_property_readonly("is_lookahead_decoding", &tr::SpeculativeDecodingMode::isLookaheadDecoding)
.def_property_readonly("is_explicit_draft_tokens", &tr::SpeculativeDecodingMode::isExplicitDraftTokens)
.def_property_readonly("needs_kv_cache_rewind", &tr::SpeculativeDecodingMode::needsKVCacheRewind)
.def_property_readonly("needs_decoder_prologue", &tr::SpeculativeDecodingMode::needsDecoderPrologue)
.def_property_readonly("predicts_draft_tokens", &tr::SpeculativeDecodingMode::predictsDraftTokens)
.def_property_readonly("needs_kv_cache_rewind", &tr::SpeculativeDecodingMode::needsKVCacheRewind);

py::classh<tr::TllmRuntime>(m, "TllmRuntime")
.def(py::init(
[](std::filesystem::path engine_path, float gpu_weights_percent = 1.0f, bool use_shape_inference = true)
Expand Down Expand Up @@ -282,7 +265,6 @@ void initBindings(pybind11::module_& m)
.def_readwrite("medusa_paths", &tr::decoder_batch::Request::medusaPaths)
.def_readwrite("medusa_tree_ids", &tr::decoder_batch::Request::medusaTreeIds)
.def_readwrite("lookahead_runtime_config", &tr::decoder_batch::Request::lookaheadRuntimeConfig);
py::bind_vector<std::vector<tr::decoder_batch::Request>>(m, "VectorRequest");

py::class_<tr::decoder_batch::Input>(m, "DecoderBatchInput")
.def(py::init<std::vector<std::vector<tr::ITensor::SharedConstPtr>>, tr::SizeType32>(), py::arg("logits"),
Expand Down Expand Up @@ -431,4 +413,25 @@ void initBindings(pybind11::module_& m)
initMoeBindings(m);
}

void initBindingsEarly(py::module_& m)
{
py::class_<tr::SpeculativeDecodingMode>(m, "SpeculativeDecodingMode")
.def(py::init<tr::SpeculativeDecodingMode::UnderlyingType>(), py::arg("state"))
.def_static("NoneType", &tr::SpeculativeDecodingMode::None)
.def_static("DraftTokensExternal", &tr::SpeculativeDecodingMode::DraftTokensExternal)
.def_static("Medusa", &tr::SpeculativeDecodingMode::Medusa)
.def_static("Eagle", &tr::SpeculativeDecodingMode::Eagle)
.def_static("LookaheadDecoding", &tr::SpeculativeDecodingMode::LookaheadDecoding)
.def_static("ExplicitDraftTokens", &tr::SpeculativeDecodingMode::ExplicitDraftTokens)
.def_property_readonly("is_none", &tr::SpeculativeDecodingMode::isNone)
.def_property_readonly("is_draft_tokens_external", &tr::SpeculativeDecodingMode::isDraftTokensExternal)
.def_property_readonly("is_medusa", &tr::SpeculativeDecodingMode::isMedusa)
.def_property_readonly("is_eagle", &tr::SpeculativeDecodingMode::isEagle)
.def_property_readonly("is_lookahead_decoding", &tr::SpeculativeDecodingMode::isLookaheadDecoding)
.def_property_readonly("is_explicit_draft_tokens", &tr::SpeculativeDecodingMode::isExplicitDraftTokens)
.def_property_readonly("needs_kv_cache_rewind", &tr::SpeculativeDecodingMode::needsKVCacheRewind)
.def_property_readonly("needs_decoder_prologue", &tr::SpeculativeDecodingMode::needsDecoderPrologue)
.def_property_readonly("predicts_draft_tokens", &tr::SpeculativeDecodingMode::predictsDraftTokens)
.def_property_readonly("needs_kv_cache_rewind", &tr::SpeculativeDecodingMode::needsKVCacheRewind);
}
} // namespace tensorrt_llm::pybind::runtime
1 change: 1 addition & 0 deletions cpp/tensorrt_llm/pybind/runtime/bindings.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,5 +26,6 @@ namespace tensorrt_llm::pybind::runtime
{

void initBindings(py::module_& m);
void initBindingsEarly(py::module_& m);

} // namespace tensorrt_llm::pybind::runtime
1 change: 1 addition & 0 deletions tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def forward(
scheduled_requests: ScheduledRequests,
resource_manager: ResourceManager,
new_tokens_device: Optional[torch.Tensor] = None,
gather_context_logits: bool = False,
):
"""Run forward from scheduled requests; main entrypoint that gets called by the executor."""
# convert requests and store in sequence info object
Expand Down
Loading