From 04a10629e1e4ab4b72a0eef9424e48972e9248d2 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 21 May 2025 11:22:19 +0000 Subject: [PATCH 1/7] First versionl Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- .../pybind/batch_manager/bindings.cpp | 138 ++---------------- cpp/tensorrt_llm/pybind/bindings.cpp | 21 ++- .../pybind/common/customCasters.h | 65 ++++++++- cpp/tensorrt_llm/pybind/runtime/bindings.cpp | 39 ++--- cpp/tensorrt_llm/pybind/runtime/bindings.h | 1 + examples/pytorch/quickstart_advanced.py | 9 +- .../pyexecutor/handle_context_logits.py | 94 ++++++++++++ .../_torch/pyexecutor/model_engine.py | 6 +- tensorrt_llm/_torch/pyexecutor/sampler.py | 12 +- tests/unittest/_torch/test_return_logits.py | 20 ++- tests/unittest/bindings/test_bindings_ut.py | 8 +- 11 files changed, 240 insertions(+), 173 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/handle_context_logits.py diff --git a/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp index 186e4f9b6b57..1403da1285c9 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp @@ -102,6 +102,7 @@ void initBindings(pybind11::module_& m) .def("get_tokens", py::overload_cast<>(&GenLlmReq::getTokens, py::const_)) .def("get_last_tokens", py::overload_cast(&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")) @@ -109,114 +110,17 @@ void initBindings(pybind11::module_& m) .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 value{std::nullopt}; - auto tensor = self.getPromptEmbeddingTable(); - if (tensor) - { - value = tr::Torch::tensor(*tensor); - } - return value; - }) - .def("multimodal_embedding", - [](GenLlmReq& self) - { - std::optional 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 value{std::nullopt}; - auto tensor = self.getMropeRotaryCosSin(); - if (tensor) - { - value = tr::Torch::tensor(*tensor); - } - return value; - }) - .def("bad_words_list", - [](GenLlmReq& self) - { - std::optional value{std::nullopt}; - auto tensor = self.getBadWordsList(); - if (tensor) - { - value = tr::Torch::tensor(*tensor); - } - return value; - }) - .def_property( - "draft_logits", - [](GenLlmReq& self) - { - std::optional 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(tr::TorchView::of(logits))); }) - .def("embedding_bias", - [](GenLlmReq& self) - { - std::optional value{std::nullopt}; - auto tensor = self.getEmbeddingBias(); - if (tensor) - { - value = tr::Torch::tensor(*tensor); - } - return value; - }) - .def_property( - "lora_config", - [](GenLlmReq& self) - { - std::optional 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(tr::TorchView::of(loraConfig))); }) - .def_property( - "lora_weights", - [](GenLlmReq& self) - { - std::optional 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(tr::TorchView::of(loraWeights))); }) - .def("stop_words_list", - [](GenLlmReq& self) - { - std::optional 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) @@ -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) @@ -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) @@ -311,20 +218,7 @@ void initBindings(pybind11::module_& m) { self.setDraftTokens(std::make_shared(draftTokens.value())); } - }) - .def_property( - "context_logits", - [](GenLlmReq& self) - { - std::optional 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(m, "LlmRequest", pybind11::dynamic_attr()) .def(py::init( diff --git a/cpp/tensorrt_llm/pybind/bindings.cpp b/cpp/tensorrt_llm/pybind/bindings.cpp index c6f040bfa628..f54414162c09 100644 --- a/cpp/tensorrt_llm/pybind/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/bindings.cpp @@ -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); @@ -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) @@ -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>(m, "VectorSamplingConfig"); - m.def("make_sampling_config", &makeSamplingConfig, py::arg("configs")); py::class_(m, "GptJsonConfig") @@ -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); diff --git a/cpp/tensorrt_llm/pybind/common/customCasters.h b/cpp/tensorrt_llm/pybind/common/customCasters.h index 45cc876b2fd6..3d1eea7e3f34 100644 --- a/cpp/tensorrt_llm/pybind/common/customCasters.h +++ b/cpp/tensorrt_llm/pybind/common/customCasters.h @@ -43,8 +43,6 @@ // Opaque bindings PYBIND11_MAKE_OPAQUE(tensorrt_llm::batch_manager::ReqIdsSet) PYBIND11_MAKE_OPAQUE(std::vector) -PYBIND11_MAKE_OPAQUE(std::vector) -PYBIND11_MAKE_OPAQUE(std::vector) // Custom casters namespace PYBIND11_NAMESPACE @@ -204,5 +202,68 @@ struct type_caster } }; +template <> +struct type_caster +{ +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 +{ +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(src))); + } +}; + } // namespace detail } // namespace PYBIND11_NAMESPACE diff --git a/cpp/tensorrt_llm/pybind/runtime/bindings.cpp b/cpp/tensorrt_llm/pybind/runtime/bindings.cpp index 6fff7e6bc54a..9498e8b3d983 100644 --- a/cpp/tensorrt_llm/pybind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/runtime/bindings.cpp @@ -216,23 +216,6 @@ void initBindings(pybind11::module_& m) .def(py::init(), py::arg("stream"), py::arg("trim_pool") = false) .def_property_readonly("stream", &tr::BufferManager::getStream); - py::class_(m, "SpeculativeDecodingMode") - .def(py::init(), 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(m, "TllmRuntime") .def(py::init( [](std::filesystem::path engine_path, float gpu_weights_percent = 1.0f, bool use_shape_inference = true) @@ -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>(m, "VectorRequest"); py::class_(m, "DecoderBatchInput") .def(py::init>, tr::SizeType32>(), py::arg("logits"), @@ -431,4 +413,25 @@ void initBindings(pybind11::module_& m) initMoeBindings(m); } +void initBindingsEarly(py::module_& m) +{ + py::class_(m, "SpeculativeDecodingMode") + .def(py::init(), 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 diff --git a/cpp/tensorrt_llm/pybind/runtime/bindings.h b/cpp/tensorrt_llm/pybind/runtime/bindings.h index 95c9720f35ea..b8e1ab66574f 100644 --- a/cpp/tensorrt_llm/pybind/runtime/bindings.h +++ b/cpp/tensorrt_llm/pybind/runtime/bindings.h @@ -26,5 +26,6 @@ namespace tensorrt_llm::pybind::runtime { void initBindings(py::module_& m); +void initBindingsEarly(py::module_& m); } // namespace tensorrt_llm::pybind::runtime diff --git a/examples/pytorch/quickstart_advanced.py b/examples/pytorch/quickstart_advanced.py index 4092e1f0d41b..d00d790cba81 100644 --- a/examples/pytorch/quickstart_advanced.py +++ b/examples/pytorch/quickstart_advanced.py @@ -7,9 +7,9 @@ example_prompts = [ "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", + # "The president of the United States is", + # "The capital of France is", + # "The future of AI is", ] @@ -182,6 +182,7 @@ def setup_llm(args): temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, + return_context_logits=True, ) return llm, sampling_params @@ -196,7 +197,9 @@ def main(): for i, output in enumerate(outputs): prompt = output.prompt generated_text = output.outputs[0].text + context_logits = output.context_logits print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}") + print(f"[{i}] Context logits: {context_logits!r}") if __name__ == '__main__': diff --git a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py new file mode 100644 index 000000000000..e6f4902c1152 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py @@ -0,0 +1,94 @@ +import torch +from typing import List, Optional + +from tensorrt_llm.bindings.internal.batch_manager import DecoderBuffers +from tensorrt_llm.bindings import ModelConfig +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + + +# TODO: Implement this once return generation logits is supported +# def copy_last_context_logits(context_logits: torch.Tensor, llm_req: LlmRequest): +# """Copy logits from context phase to beginning of generation logits. + +# Usually, this concerns logits of 1 token. In speculative decoding this concerns draftLen + 1 tokens. +# """ +# num_logits = context_logits.shape[0] +# for beam in range(llm_req.get_beam_width_by_iter()): +# # [beamWidth, mMaxNewTokens, vocabSizePadded] -> [numLogits, vocabSizePadded] +# beam_host_tensor = llm_req.get_generation_logits_host()[beam, :num_logits] +# beam_host_tensor.copy_(context_logits, non_blocking=True) + + +class HandleContextLogits: + + def __call__(self, + context_requests: List[LlmRequest], + num_context_logits_vec: List[int], + logits: torch.Tensor, + decoder_buffers: DecoderBuffers) -> int: + """Handle context logits for a batch of requests. + + Args: + context_requests: List of context requests to process + num_context_logits_vec: Number of context logits for each request + logits: Input logits tensor + decoder_buffers: Decoder buffers for storing intermediate results + + Returns: + int: Index into logits tensor after processing all requests + """ + logits_index = 0 + + # Copy logits into decoderBuffers.logits + decoder_buffer_logits = [torch.empty(0)] * len(decoder_buffers.logits) + for batch_index, llm_req in enumerate(context_requests): + num_context_logits = num_context_logits_vec[batch_index] + draft_length = llm_req.num_draft_tokens if llm_req.is_last_context_chunk() else 0 + + print("Is last context chunk: ", llm_req.is_last_context_chunk(), "py_return_context_logits: ", llm_req.py_return_context_logits) + + if llm_req.py_return_context_logits: + print("py_return_context_logits: ", llm_req.py_return_context_logits) + if llm_req.prepopulated_prompt_len > 0: + print(f"Warning: Because of KV cache reuse, not all context logits could be produced for request {llm_req.request_id}.") + + print("logits shape: ", logits.shape) + + context_logits_device_view = logits[logits_index:logits_index + num_context_logits] + llm_req.py_result.append_context_logits(context_logits_device_view) + + logits_index += num_context_logits + draft_length + + # Get the logits from the last context token and draft tokens + num_decoder_logits = 1 + draft_length + seq_slot = llm_req.seq_slot + logits_view = logits[logits_index - num_decoder_logits:logits_index] + logits_view_shape = logits_view.shape + + # Create a view of logits_view with shape (logits_view_shape[0], 1, logits_view_shape[1]) + # This creates a new tensor that shares the same underlying data + decoder_buffer_logits[seq_slot] = logits_view[:logits_view_shape[0], :1, :logits_view_shape[1]] + + # TODO: Implement this once return generation logits is supported + # # Save the last token logits of context into generation logits or + # # save the accepted token logits from target model + # if llm_req.get_return_generation_logits(): + # copy_last_context_logits(logits_view, llm_req) + + # TODO: Implement this once we have beam width support + # Scatter the output logits to the decoderLogits + # req_beam_width = llm_req.get_beam_width_by_iter() + # if req_beam_width > 1: + # # Tile logits of context requests + # logits_shape = logits_view.shape + # logits_type = logits_view.dtype + # # decoder_logits = buffer_manager.gpu((req_beam_width, logits_shape[1]), logits_type) + # # tensorrt_llm.runtime.kernels.tile_tensor(decoder_logits, logits_view, req_beam_width, stream) + # decoder_logits = decoder_logits.unsqueeze(0) + # else: + # decoder_buffer_logits[seq_slot] = logits_view[:logits_view_shape[0], :1, :logits_view_shape[1]] + + # Needs to be done in bulk for the copy to work + decoder_buffers.logits = decoder_buffer_logits + + return logits_index diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 176f5a51c3fb..d7388140c09a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1062,11 +1062,11 @@ def _prepare_tp_inputs( prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = request.context_current_position num_cached_tokens_per_seq.append(past_seen_token_num) - multimodal_embedding = request.multimodal_embedding() + multimodal_embedding = request.multimodal_embedding if multimodal_embedding is not None: multi_modal_data.append(multimodal_embedding) - mrope_rotary_cos_sin = request.get_mrope_rotary_cos_sin() + mrope_rotary_cos_sin = request.mrope_rotary_cos_sin if mrope_rotary_cos_sin is not None: mrope_config['mrope_rotary_cos_sin'].append( mrope_rotary_cos_sin) @@ -1380,7 +1380,7 @@ def _prepare_tp_inputs_no_cache( gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) draft_lens.append(0) - multimodal_embedding = request.multimodal_embedding() + multimodal_embedding = request.multimodal_embedding if multimodal_embedding is not None: multi_modal_data.append(multimodal_embedding) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 9ed4c5897345..5dc323d01ec9 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -4,13 +4,15 @@ import torch +from tensorrt_llm._torch.pyexecutor.handle_context_logits import \ + HandleContextLogits from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings import (CudaStream, DataType, ModelConfig, WorldConfig, make_sampling_config) from tensorrt_llm.bindings.executor import (DecodingConfig, DecodingMode, ExecutorConfig, FinishReason) from tensorrt_llm.bindings.internal.algorithms import ( - CreateNewDecoderRequests, HandleContextLogits, HandleGenerationLogits, + CreateNewDecoderRequests, HandleGenerationLogits, MakeDecodingBatchInputOutput) from tensorrt_llm.bindings.internal.batch_manager import (DecoderBuffers, DecoderInputBuffers) @@ -601,10 +603,9 @@ def sample_async(self, scheduled_requests: ScheduledRequests, # numContextLogits.at(batchIdx) = modelConfig.computeContextLogits() ? contextChunkSize : 1; # Revisit this when we support chunked context. num_context_logits = [1] * batch_size - logits_index = self.algs.handle_context_logits( - scheduled_requests.context_requests, num_context_logits, logits, - self.store["decoder_buffers"], self.model_config, - self.store["buffer_manager"], self.store["cuda_stream"]) + logits_index = self.algs.handle_context_logits(scheduled_requests.context_requests, + num_context_logits, logits, + self.store["decoder_buffers"]) self.algs.handle_generation_logits( logits_index, scheduled_requests.generation_requests, @@ -682,7 +683,6 @@ def update_requests(self, state: SampleStateTRTLLM): seq_slot = request.seq_slot num_generated_tokens = request.num_draft_tokens + 1 current_num_of_tokens = request.max_beam_num_tokens - num_new_tokens = [0] * beam_width for beam in range(beam_width): diff --git a/tests/unittest/_torch/test_return_logits.py b/tests/unittest/_torch/test_return_logits.py index 38aed5d40d5b..91b56efbd897 100644 --- a/tests/unittest/_torch/test_return_logits.py +++ b/tests/unittest/_torch/test_return_logits.py @@ -9,6 +9,7 @@ from tensorrt_llm import SamplingParams from tensorrt_llm._torch import LLM from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse, PyResult +from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig from tensorrt_llm.bindings.executor import Response, Result from tensorrt_llm.executor.result import Logprob from tensorrt_llm.llmapi.llm_utils import BuildConfig, KvCacheConfig @@ -54,7 +55,7 @@ def test_LlmResponse_pickle(): assert pickle_result.log_probs == logprobs -@force_ampere # Save H100 resource +# @force_ampere # Save H100 resource @pytest.mark.parametrize("gather_context_logits", [False, True]) @pytest.mark.parametrize("gather_generation_logits", [False, True]) @pytest.mark.parametrize("return_log_probs", [False, True]) @@ -64,22 +65,30 @@ def test_generate_with_return_logits(gather_context_logits: bool, if not (gather_context_logits or gather_generation_logits or return_log_probs): # prune space pytest.skip("Nothing to test") - - if gather_context_logits: - pytest.skip("gather_context_logits unimplemented yet") + + enable_trtllm_sampler = gather_context_logits + if enable_trtllm_sampler and (gather_generation_logits or return_log_probs): + pytest.skip("TRTLLMSampler does not support gather_generation_logits or return_log_probs") build_config = BuildConfig() build_config.gather_context_logits = gather_context_logits + pytorch_config = PyTorchConfig( + enable_trtllm_sampler=enable_trtllm_sampler) + llm = LLM( model=os.path.join(llm_models_root(), "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0"), kv_cache_config=global_kvcache_config, build_config=build_config, + gather_context_logits=gather_context_logits, gather_generation_logits=gather_generation_logits, max_batch_size= 128, # reduce buffer sizes, specially for generation logits + pytorch_backend_config=pytorch_config, ) + + print("gather_context_logits: ", gather_context_logits) sampling_params = SamplingParams( max_tokens=8, return_context_logits=gather_context_logits, @@ -89,6 +98,9 @@ def test_generate_with_return_logits(gather_context_logits: bool, for output in llm.generate(prompts, sampling_params=sampling_params): if gather_context_logits: assert output.context_logits is not None + print(f"context_logits: {output.context_logits.shape}") + print(f"prompts: {prompts}") + print(f"len(prompts[0].split()): {len(prompts[0].split())}") assert len(prompts[0].split()) + \ 1 == output.context_logits.shape[0] else: diff --git a/tests/unittest/bindings/test_bindings_ut.py b/tests/unittest/bindings/test_bindings_ut.py index 59e19bab1cbc..8a727de65a64 100644 --- a/tests/unittest/bindings/test_bindings_ut.py +++ b/tests/unittest/bindings/test_bindings_ut.py @@ -338,12 +338,12 @@ def test_llm_request(): assert llm_request.pad_id == 99 assert llm_request.end_id == 100 assert llm_request.seq_slot is None - assert torch.equal(llm_request.prompt_embedding_table(), + assert torch.equal(llm_request.prompt_embedding_table, kwargs["prompt_embedding_table"]) assert llm_request.prompt_vocab_size == 2 - assert torch.equal(llm_request.embedding_bias(), kwargs["embedding_bias"]) - assert torch.equal(llm_request.stop_words_list(), kwargs["stop_words_list"]) - assert torch.equal(llm_request.bad_words_list(), kwargs["bad_words_list"]) + assert torch.equal(llm_request.embedding_bias, kwargs["embedding_bias"]) + assert torch.equal(llm_request.stop_words_list, kwargs["stop_words_list"]) + assert torch.equal(llm_request.bad_words_list, kwargs["bad_words_list"]) assert llm_request.get_num_tokens(0) == 3 assert llm_request.max_beam_num_tokens == 3 From 65956a7371a18c71ebbf8d11e2653474fcd2e7ed Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 21 May 2025 12:28:53 +0000 Subject: [PATCH 2/7] Remove comments. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/handle_context_logits.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py index e6f4902c1152..dfcc7923aec4 100644 --- a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py +++ b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py @@ -45,14 +45,9 @@ def __call__(self, num_context_logits = num_context_logits_vec[batch_index] draft_length = llm_req.num_draft_tokens if llm_req.is_last_context_chunk() else 0 - print("Is last context chunk: ", llm_req.is_last_context_chunk(), "py_return_context_logits: ", llm_req.py_return_context_logits) - if llm_req.py_return_context_logits: - print("py_return_context_logits: ", llm_req.py_return_context_logits) if llm_req.prepopulated_prompt_len > 0: print(f"Warning: Because of KV cache reuse, not all context logits could be produced for request {llm_req.request_id}.") - - print("logits shape: ", logits.shape) context_logits_device_view = logits[logits_index:logits_index + num_context_logits] llm_req.py_result.append_context_logits(context_logits_device_view) From 36bdfb3515f84a57d31b596c3c5e7e566c320eb8 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 21 May 2025 12:43:03 +0000 Subject: [PATCH 3/7] Partially support context logits. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- .../pyexecutor/handle_context_logits.py | 46 +++++++++-------- tensorrt_llm/_torch/pyexecutor/sampler.py | 6 +-- tests/unittest/_torch/test_return_logits.py | 50 ++++++++++++------- 3 files changed, 59 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py index dfcc7923aec4..5f908fa29c17 100644 --- a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py +++ b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py @@ -1,15 +1,15 @@ +from typing import List + import torch -from typing import List, Optional -from tensorrt_llm.bindings.internal.batch_manager import DecoderBuffers -from tensorrt_llm.bindings import ModelConfig from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest - +from tensorrt_llm.bindings.internal.batch_manager import DecoderBuffers +from tensorrt_llm.logger import logger # TODO: Implement this once return generation logits is supported # def copy_last_context_logits(context_logits: torch.Tensor, llm_req: LlmRequest): # """Copy logits from context phase to beginning of generation logits. - + # Usually, this concerns logits of 1 token. In speculative decoding this concerns draftLen + 1 tokens. # """ # num_logits = context_logits.shape[0] @@ -21,36 +21,39 @@ class HandleContextLogits: - def __call__(self, - context_requests: List[LlmRequest], - num_context_logits_vec: List[int], - logits: torch.Tensor, + def __call__(self, context_requests: List[LlmRequest], + num_context_logits_vec: List[int], logits: torch.Tensor, decoder_buffers: DecoderBuffers) -> int: """Handle context logits for a batch of requests. - + Args: context_requests: List of context requests to process num_context_logits_vec: Number of context logits for each request logits: Input logits tensor decoder_buffers: Decoder buffers for storing intermediate results - + Returns: int: Index into logits tensor after processing all requests """ logits_index = 0 - + # Copy logits into decoderBuffers.logits decoder_buffer_logits = [torch.empty(0)] * len(decoder_buffers.logits) for batch_index, llm_req in enumerate(context_requests): num_context_logits = num_context_logits_vec[batch_index] - draft_length = llm_req.num_draft_tokens if llm_req.is_last_context_chunk() else 0 + draft_length = llm_req.num_draft_tokens if llm_req.is_last_context_chunk( + ) else 0 if llm_req.py_return_context_logits: if llm_req.prepopulated_prompt_len > 0: - print(f"Warning: Because of KV cache reuse, not all context logits could be produced for request {llm_req.request_id}.") + logger.warning( + f"Because of KV cache reuse, not all context logits could be produced for request {llm_req.request_id}." + ) - context_logits_device_view = logits[logits_index:logits_index + num_context_logits] - llm_req.py_result.append_context_logits(context_logits_device_view) + context_logits_device_view = logits[logits_index:logits_index + + num_context_logits] + llm_req.py_result.append_context_logits( + context_logits_device_view) logits_index += num_context_logits + draft_length @@ -58,11 +61,12 @@ def __call__(self, num_decoder_logits = 1 + draft_length seq_slot = llm_req.seq_slot logits_view = logits[logits_index - num_decoder_logits:logits_index] - logits_view_shape = logits_view.shape - # Create a view of logits_view with shape (logits_view_shape[0], 1, logits_view_shape[1]) + # Create a view of logits_view with shape (logits_view.shape[0], 1, logits_view.shape[1]) # This creates a new tensor that shares the same underlying data - decoder_buffer_logits[seq_slot] = logits_view[:logits_view_shape[0], :1, :logits_view_shape[1]] + decoder_buffer_logits[ + seq_slot] = logits_view[:logits_view.shape[0], :1, :logits_view. + shape[1]] # TODO: Implement this once return generation logits is supported # # Save the last token logits of context into generation logits or @@ -81,8 +85,8 @@ def __call__(self, # # tensorrt_llm.runtime.kernels.tile_tensor(decoder_logits, logits_view, req_beam_width, stream) # decoder_logits = decoder_logits.unsqueeze(0) # else: - # decoder_buffer_logits[seq_slot] = logits_view[:logits_view_shape[0], :1, :logits_view_shape[1]] - + # decoder_buffer_logits[seq_slot] = logits_view[:logits_view.shape[0], :1, :logits_view.shape[1]] + # Needs to be done in bulk for the copy to work decoder_buffers.logits = decoder_buffer_logits diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 5dc323d01ec9..8d9e0e10fe16 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -603,9 +603,9 @@ def sample_async(self, scheduled_requests: ScheduledRequests, # numContextLogits.at(batchIdx) = modelConfig.computeContextLogits() ? contextChunkSize : 1; # Revisit this when we support chunked context. num_context_logits = [1] * batch_size - logits_index = self.algs.handle_context_logits(scheduled_requests.context_requests, - num_context_logits, logits, - self.store["decoder_buffers"]) + logits_index = self.algs.handle_context_logits( + scheduled_requests.context_requests, num_context_logits, logits, + self.store["decoder_buffers"]) self.algs.handle_generation_logits( logits_index, scheduled_requests.generation_requests, diff --git a/tests/unittest/_torch/test_return_logits.py b/tests/unittest/_torch/test_return_logits.py index 91b56efbd897..43a6f0e7299a 100644 --- a/tests/unittest/_torch/test_return_logits.py +++ b/tests/unittest/_torch/test_return_logits.py @@ -8,8 +8,8 @@ from tensorrt_llm import SamplingParams from tensorrt_llm._torch import LLM -from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse, PyResult from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig +from tensorrt_llm._torch.pyexecutor.llm_request import LlmResponse, PyResult from tensorrt_llm.bindings.executor import Response, Result from tensorrt_llm.executor.result import Logprob from tensorrt_llm.llmapi.llm_utils import BuildConfig, KvCacheConfig @@ -55,26 +55,30 @@ def test_LlmResponse_pickle(): assert pickle_result.log_probs == logprobs -# @force_ampere # Save H100 resource +@force_ampere # Save H100 resource +@pytest.mark.parametrize("enable_trtllm_sampler", [False, True]) @pytest.mark.parametrize("gather_context_logits", [False, True]) @pytest.mark.parametrize("gather_generation_logits", [False, True]) @pytest.mark.parametrize("return_log_probs", [False, True]) -def test_generate_with_return_logits(gather_context_logits: bool, +def test_generate_with_return_logits(enable_trtllm_sampler: bool, + gather_context_logits: bool, gather_generation_logits: bool, return_log_probs: bool): if not (gather_context_logits or gather_generation_logits or return_log_probs): # prune space pytest.skip("Nothing to test") - - enable_trtllm_sampler = gather_context_logits + if enable_trtllm_sampler and (gather_generation_logits or return_log_probs): - pytest.skip("TRTLLMSampler does not support gather_generation_logits or return_log_probs") + pytest.skip( + "TRTLLMSampler does not support gather_generation_logits or return_log_probs" + ) + elif not enable_trtllm_sampler and gather_context_logits: + pytest.skip("TorchSampler does not support gather_context_logits") build_config = BuildConfig() build_config.gather_context_logits = gather_context_logits - pytorch_config = PyTorchConfig( - enable_trtllm_sampler=enable_trtllm_sampler) + pytorch_config = PyTorchConfig(enable_trtllm_sampler=enable_trtllm_sampler) llm = LLM( model=os.path.join(llm_models_root(), "llama-models-v2", @@ -88,7 +92,6 @@ def test_generate_with_return_logits(gather_context_logits: bool, pytorch_backend_config=pytorch_config, ) - print("gather_context_logits: ", gather_context_logits) sampling_params = SamplingParams( max_tokens=8, return_context_logits=gather_context_logits, @@ -98,11 +101,11 @@ def test_generate_with_return_logits(gather_context_logits: bool, for output in llm.generate(prompts, sampling_params=sampling_params): if gather_context_logits: assert output.context_logits is not None - print(f"context_logits: {output.context_logits.shape}") - print(f"prompts: {prompts}") - print(f"len(prompts[0].split()): {len(prompts[0].split())}") - assert len(prompts[0].split()) + \ - 1 == output.context_logits.shape[0] + # TODO: The intended behaviour is to return all context logits. + # However, the logits received in the sampler do not contain all context logits. + # For now, only the last context token logits are returned. + # When it is fixed, it should be len(prompts[0].split()) + 1. + assert 1 == output.context_logits.shape[0] else: assert output.context_logits is None @@ -121,22 +124,30 @@ def test_generate_with_return_logits(gather_context_logits: bool, @force_ampere # Save H100 resource +@pytest.mark.parametrize("enable_trtllm_sampler", [False, True]) @pytest.mark.parametrize("gather_context_logits", [False, True]) @pytest.mark.parametrize("gather_generation_logits", [False, True]) @pytest.mark.parametrize("return_log_probs", [False, True]) -def test_generate_async_with_return_logits(gather_context_logits: bool, +def test_generate_async_with_return_logits(enable_trtllm_sampler: bool, + gather_context_logits: bool, gather_generation_logits: bool, return_log_probs: bool): if not (gather_context_logits or gather_generation_logits or return_log_probs): # prune space pytest.skip("Nothing to test") - if gather_context_logits: - pytest.skip("gather_context_logits unimplemented yet") + if enable_trtllm_sampler and (gather_generation_logits or return_log_probs): + pytest.skip( + "TRTLLMSampler does not support gather_generation_logits or return_log_probs" + ) + elif not enable_trtllm_sampler and gather_context_logits: + pytest.skip("TorchSampler does not support gather_context_logits") build_config = BuildConfig() build_config.gather_context_logits = gather_context_logits + pytorch_config = PyTorchConfig(enable_trtllm_sampler=enable_trtllm_sampler) + llm = LLM( model=os.path.join(llm_models_root(), "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0"), @@ -145,6 +156,7 @@ def test_generate_async_with_return_logits(gather_context_logits: bool, gather_generation_logits=gather_generation_logits, max_batch_size= 128, # reduce buffer sizes, specially for generation logits + pytorch_backend_config=pytorch_config, ) sampling_params = SamplingParams( max_tokens=8, @@ -158,8 +170,8 @@ def test_generate_async_with_return_logits(gather_context_logits: bool, streaming=True)): if gather_context_logits: assert output.context_logits is not None - assert len(prompts[0].split()) + \ - 1 == output.context_logits.shape[0] + # TODO: The intended behaviour is to return all context logits. See above. + assert 1 == output.context_logits.shape[0] else: assert output.context_logits is None From 2ac1cbfeaf89e36c6ec42a9032018ea714c8c3da Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 21 May 2025 12:51:58 +0000 Subject: [PATCH 4/7] Revert quickstart advanced. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- examples/pytorch/quickstart_advanced.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/pytorch/quickstart_advanced.py b/examples/pytorch/quickstart_advanced.py index d00d790cba81..4092e1f0d41b 100644 --- a/examples/pytorch/quickstart_advanced.py +++ b/examples/pytorch/quickstart_advanced.py @@ -7,9 +7,9 @@ example_prompts = [ "Hello, my name is", - # "The president of the United States is", - # "The capital of France is", - # "The future of AI is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", ] @@ -182,7 +182,6 @@ def setup_llm(args): temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, - return_context_logits=True, ) return llm, sampling_params @@ -197,9 +196,7 @@ def main(): for i, output in enumerate(outputs): prompt = output.prompt generated_text = output.outputs[0].text - context_logits = output.context_logits print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}") - print(f"[{i}] Context logits: {context_logits!r}") if __name__ == '__main__': From 6b8c066abe824e17659fc91ac7db8ebfa5adbd2a Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 28 May 2025 10:34:06 +0000 Subject: [PATCH 5/7] Fix return context logits. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 27 ++++++++++++------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 +++++++---- tensorrt_llm/_torch/pyexecutor/sampler.py | 11 +++++--- tests/unittest/_torch/test_return_logits.py | 12 +++------ 4 files changed, 40 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d7388140c09a..64f91d87cf19 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -66,10 +66,12 @@ def get_max_num_sequences(self) -> int: raise NotImplementedError @abstractmethod - def forward(self, scheduled_requests: ScheduledRequests, + def forward(self, + scheduled_requests: ScheduledRequests, resource_manager: ResourceManager, new_tensors_device: Optional[SampleStateTensors], - extra_model_inputs: Optional[Dict[str, Any]]): + extra_model_inputs: Optional[Dict[str, Any]], + gather_context_logits: bool = False): raise NotImplementedError def warmup(self, resource_manager: ResourceManager) -> None: @@ -1844,7 +1846,8 @@ def forward(self, scheduled_requests: ScheduledRequests, resource_manager: ResourceManager, new_tensors_device: Optional[SampleStateTensors] = None, - extra_model_inputs: Optional[Dict[str, Any]] = None): + extra_model_inputs: Optional[Dict[str, Any]] = None, + gather_context_logits: bool = False): kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) @@ -1866,7 +1869,7 @@ def forward(self, inputs.update(extra_model_inputs) self.last_spec_metadata = spec_metadata - return self._forward_step(inputs, gather_ids) + return self._forward_step(inputs, gather_ids, gather_context_logits) with self._maybe_pad_batch(scheduled_requests, kv_cache_manager) as scheduled_requests: @@ -1893,12 +1896,15 @@ def forward(self, self.iter_counter += 1 if maybe_graph is None: - outputs = self._forward_step(inputs, gather_ids) + outputs = self._forward_step(inputs, gather_ids, + gather_context_logits) else: if maybe_graph.needs_capture(): pool = maybe_graph.capture( lambda inputs: self._forward_step( - inputs, gather_ids=gather_ids), + inputs, + gather_ids=gather_ids, + gather_context_logits=gather_context_logits), self._cuda_graph_mem_pool, extra_model_inputs, ) @@ -1934,8 +1940,10 @@ def model_forward(self, **kwargs): return self.model.forward(**kwargs) @nvtx_range("_forward_step") - def _forward_step(self, inputs: Dict[str, Any], - gather_ids: Optional[torch.Tensor]) -> Dict[str, Any]: + def _forward_step(self, + inputs: Dict[str, Any], + gather_ids: Optional[torch.Tensor], + gather_context_logits: bool = False) -> Dict[str, Any]: inputs = self._preprocess_inputs(inputs) if self.without_logits: outputs = self.model_forward(**inputs) @@ -1945,7 +1953,8 @@ def _forward_step(self, inputs: Dict[str, Any], # from speculative decoding. logits = self.model_forward( **inputs, - return_context_logits=gather_ids is not None, + return_context_logits=gather_ids is not None + or gather_context_logits, ) if gather_ids is not None: return {'logits': logits[gather_ids]} diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 57c04cc1c20f..20eba1146d8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1588,14 +1588,20 @@ def _forward_step(self, @nvtx_range( f"[Executor] _forward_step: {len(scheduled_requests.context_requests)} ctx reqs, {len(scheduled_requests.generation_requests)} gen reqs" ) - def forward(scheduled_requests, resource_manager, new_tensors_device): - return self.model_engine.forward(scheduled_requests, - resource_manager, - new_tensors_device) + def forward(scheduled_requests, resource_manager, new_tensors_device, + gather_context_logits): + return self.model_engine.forward( + scheduled_requests, + resource_manager, + new_tensors_device, + gather_context_logits=gather_context_logits) try: + gather_context_logits = any( + a.py_return_context_logits + for a in scheduled_requests.context_requests) outputs = forward(scheduled_requests, self.resource_manager, - new_tensors_device) + new_tensors_device, gather_context_logits) return outputs except Exception as e: traceback.print_exc() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 8d9e0e10fe16..af3500f8b71b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -595,14 +595,17 @@ def sample_async(self, scheduled_requests: ScheduledRequests, batch_size = scheduled_requests.batch_size beam_width = self.beam_width(scheduled_requests.all_requests) - logits = model_outputs["logits"].reshape((batch_size, beam_width, -1)) + # TODO: Remove this unsqueezing once we get beam width support. + logits = model_outputs["logits"].unsqueeze(1) self.setup_sampler_step(scheduled_requests.context_requests) - # Note: In runtimeBuffers.cpp, num_context_logits is set to: - # numContextLogits.at(batchIdx) = modelConfig.computeContextLogits() ? contextChunkSize : 1; - # Revisit this when we support chunked context. num_context_logits = [1] * batch_size + for batch_index, request in enumerate( + scheduled_requests.context_requests): + num_context_logits[ + batch_index] = request.context_chunk_size if request.py_return_context_logits else 1 + logits_index = self.algs.handle_context_logits( scheduled_requests.context_requests, num_context_logits, logits, self.store["decoder_buffers"]) diff --git a/tests/unittest/_torch/test_return_logits.py b/tests/unittest/_torch/test_return_logits.py index 43a6f0e7299a..f09db8d35bb5 100644 --- a/tests/unittest/_torch/test_return_logits.py +++ b/tests/unittest/_torch/test_return_logits.py @@ -15,7 +15,7 @@ from tensorrt_llm.llmapi.llm_utils import BuildConfig, KvCacheConfig prompts = ["A B C"] -global_kvcache_config = KvCacheConfig(max_tokens=256) +global_kvcache_config = KvCacheConfig(max_tokens=2048) def test_LlmResponse_pickle(): @@ -101,11 +101,7 @@ def test_generate_with_return_logits(enable_trtllm_sampler: bool, for output in llm.generate(prompts, sampling_params=sampling_params): if gather_context_logits: assert output.context_logits is not None - # TODO: The intended behaviour is to return all context logits. - # However, the logits received in the sampler do not contain all context logits. - # For now, only the last context token logits are returned. - # When it is fixed, it should be len(prompts[0].split()) + 1. - assert 1 == output.context_logits.shape[0] + assert len(prompts[0].split()) == output.context_logits.shape[0] else: assert output.context_logits is None @@ -153,6 +149,7 @@ def test_generate_async_with_return_logits(enable_trtllm_sampler: bool, "TinyLlama-1.1B-Chat-v1.0"), kv_cache_config=global_kvcache_config, build_config=build_config, + gather_context_logits=gather_context_logits, gather_generation_logits=gather_generation_logits, max_batch_size= 128, # reduce buffer sizes, specially for generation logits @@ -170,8 +167,7 @@ def test_generate_async_with_return_logits(enable_trtllm_sampler: bool, streaming=True)): if gather_context_logits: assert output.context_logits is not None - # TODO: The intended behaviour is to return all context logits. See above. - assert 1 == output.context_logits.shape[0] + assert len(prompts[0].split()) == output.context_logits.shape[0] else: assert output.context_logits is None From a4e937ecb4608e4f90cf567acecbadac226821c9 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Thu, 29 May 2025 09:10:11 +0000 Subject: [PATCH 6/7] Remove unsqueeze. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- .../_torch/pyexecutor/handle_context_logits.py | 5 ++--- tensorrt_llm/_torch/pyexecutor/sampler.py | 11 ++++------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py index 5f908fa29c17..ac9a6a3d13db 100644 --- a/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py +++ b/tensorrt_llm/_torch/pyexecutor/handle_context_logits.py @@ -64,9 +64,8 @@ def __call__(self, context_requests: List[LlmRequest], # Create a view of logits_view with shape (logits_view.shape[0], 1, logits_view.shape[1]) # This creates a new tensor that shares the same underlying data - decoder_buffer_logits[ - seq_slot] = logits_view[:logits_view.shape[0], :1, :logits_view. - shape[1]] + decoder_buffer_logits[seq_slot] = logits_view.reshape( + logits_view.shape[0], 1, logits_view.shape[1]) # TODO: Implement this once return generation logits is supported # # Save the last token logits of context into generation logits or diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index af3500f8b71b..69f7dfd64517 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -595,9 +595,6 @@ def sample_async(self, scheduled_requests: ScheduledRequests, batch_size = scheduled_requests.batch_size beam_width = self.beam_width(scheduled_requests.all_requests) - # TODO: Remove this unsqueezing once we get beam width support. - logits = model_outputs["logits"].unsqueeze(1) - self.setup_sampler_step(scheduled_requests.context_requests) num_context_logits = [1] * batch_size @@ -607,13 +604,13 @@ def sample_async(self, scheduled_requests: ScheduledRequests, batch_index] = request.context_chunk_size if request.py_return_context_logits else 1 logits_index = self.algs.handle_context_logits( - scheduled_requests.context_requests, num_context_logits, logits, - self.store["decoder_buffers"]) + scheduled_requests.context_requests, num_context_logits, + model_outputs["logits"], self.store["decoder_buffers"]) self.algs.handle_generation_logits( logits_index, scheduled_requests.generation_requests, self.store["decoder_buffers"], self.model_config, - self.store["buffer_manager"], logits) + self.store["buffer_manager"], model_outputs["logits"]) decoding_input, self.decoding_output = self.algs.make_decoding_batch_input_output( scheduled_requests.context_requests, @@ -658,7 +655,7 @@ def sample_async(self, scheduled_requests: ScheduledRequests, sampler_event.record() return SampleStateTRTLLM(scheduled_requests=scheduled_requests, - logits=logits, + logits=model_outputs["logits"], device=device, host=host, sampler_event=sampler_event) From 3d75010f3fdb55a6cd54036c43aac62d175f67d5 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Fri, 30 May 2025 07:19:49 +0000 Subject: [PATCH 7/7] Add gather_context_logits. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b51f3ba3b858..65bd8fd2f81d 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -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