cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp (1)

124-165: Add sanity checks and minor MPI argument fix to harden the receive path.

Null and dtype checks protect against misuse; validating dims positivity avoids invalid shapes; passing dims instead of &dims avoids pointer-to-array confusion.

Apply this diff:

-void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
-    executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype)
+void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
+    executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype)
 {
 #if ENABLE_MULTI_DEVICE
+    TLLM_CHECK(draftLogitsHost && "draftLogitsHost must be non-null");
+    TLLM_CHECK(draftLogitsHost->getDataType() == logitsDtype && "draftLogitsHost dtype must match logitsDtype");
     auto const& worldComm = tensorrt_llm::mpi::MpiComm::world();
@@
-    int64_t dims[2];
-    MPICHECK(MPI_Mrecv(&dims, count, MPI_INT64_T, &msg, &status));
+    int64_t dims[2];
+    MPICHECK(MPI_Mrecv(dims, count, MPI_INT64_T, &msg, &status));
+    TLLM_CHECK(dims[0] > 0 && dims[1] > 0);
@@
-    MPICHECK(MPI_Mrecv(draftLogitsHost->data(), count, MPI_UINT8_T, &msg, &status));
+    MPICHECK(MPI_Mrecv(draftLogitsHost->data(), count, MPI_UINT8_T, &msg, &status));
 #endif // ENABLE_MULTI_DEVICE
 }

cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp (1)

25-36: Add missing standard headers used in this file

This file uses std::transform (Line 122) and std::memcpy (Line 343) but does not include or . Relying on transitive includes is brittle.

Apply this diff to include the required headers:

 #include "tensorrt_llm/common/nvtxUtils.h"
 #include "tensorrt_llm/runtime/common.h"
 #include "tensorrt_llm/runtime/decoderState.h"
 #include "tensorrt_llm/runtime/decodingInput.h"
 #include "tensorrt_llm/runtime/decodingOutput.h"
 #include "tensorrt_llm/runtime/iBuffer.h"
 #include "tensorrt_llm/runtime/runtimeKernels.h"
 #include "tensorrt_llm/runtime/speculativeDecodingMode.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h"

 #include <NvInferRuntimeBase.h>
+#include <algorithm>
+#include <cstring>

cpp/include/tensorrt_llm/runtime/decodingInput.h (2)

105-105: Add brief docs for new host-side tensor member (shape, location expectations).

Please document the expected memory type (pinned host) and shape conventions of draftLogitsHost to mirror draftLogits. This reduces ambiguity for future maintainers and call sites.

Apply this diff to add a short Doxygen note:

-        TensorPtr draftLogitsHost;
+        //! Host-pinned buffer mirroring draftLogits. Expected shape:
+        //! [batchSize, maxTokensPerEngineStep, vocabSizePadded] on pinned host.
+        TensorPtr draftLogitsHost;

---

111-111: Add brief docs for new host-side token-ids tensor (shape, location expectations).

Similarly, document draftTokenIdsHost to mirror draftTokenIds and clarify that it lives in pinned host memory and matches device shape.

Apply this diff to add a short Doxygen note:

-        TensorPtr draftTokenIdsHost;
+        //! Host-pinned buffer mirroring draftTokenIds. Expected shape:
+        //! [batchSize, maxTokensPerEngineStep] on pinned host.
+        TensorPtr draftTokenIdsHost;

cpp/tensorrt_llm/batch_manager/utils/logitsThread.h (1)

53-55: API change looks good; add preconditions and usage doc for clarity.

The switch to an in/out buffer is clean and removes optional churn. Please add a short comment describing:

• draftLogitsHost must be a valid ITensor (host-pinned) view/slice sized 2D by the caller, since the callee reshapes it to [numRows, vocab].
• logitsDtype must match draftLogitsHost->getDataType().

This helps avoid misuse at call sites.

Apply this diff to document expectations:

-void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
-    executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype);
+// Receives fast logits over MPI into a caller-provided host-pinned tensor.
+// Preconditions:
+// - draftLogitsHost is non-null and points to a host-pinned tensor/view; dtype must equal logitsDtype.
+// - The provided tensor may be a 2D slice/view; this function will reshape it to [numRows, vocab].
+void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
+    executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype);

cpp/tensorrt_llm/runtime/decoderState.cpp (1)

182-194: Host buffers allocation aligns with device counterparts; consider adding a brief assertion of pinned pool usage where consumed.

Allocating draftLogitsHost/draftTokenIdsHost/numDraftTokensHost with MemoryType::kPINNEDPOOL is correct and consistent. Ensure downstream consumers (e.g., MPI receive) assume host-pinned buffers to enable direct writes.

Optionally add a comment near allocations indicating these are intended for MPI/direct host writes:

-        externalDraftTokensInputs.draftLogitsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, dtype);
+        // Host-pinned buffer for receiving fast logits via MPI/direct writes.
+        externalDraftTokensInputs.draftLogitsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, dtype);
...
-        externalDraftTokensInputs.draftTokenIdsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvTokenIdType);
+        // Host-pinned buffer for incoming draft token IDs.
+        externalDraftTokensInputs.draftTokenIdsHost = bufferManager.emptyTensor(MemoryType::kPINNEDPOOL, nvTokenIdType);

cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp (2)

410-420: Ensure draft logits presence and shape match numDraftTokens when provided

If useDraftLogits is true, we should assert that the provided logits buffer is consistent with the number of draft tokens, and clarify expected dtype (matches modelConfig.getLogitsDtype()).

Consider adding checks similar to:

• reqDraftLogits is non-null
• reqDraftLogits->getShape().nbDims == 2 and leading dimension equals numDraftTokens
• reqDraftLogits dtype matches modelConfig.getLogitsDtype() when speculativeDecodingFastLogits is false

Do you want me to generate a patch adding these validations?

---

569-621: Zeroing speculative outputs per-request may be unnecessary work

The unconditional zeroing of nextDraftTokens (and nextDraftTokensLen when variable length) on every new speculative request adds memory traffic per request. If the downstream kernels overwrite these buffers every step, this can be skipped.

• Verify whether the first-step kernels always write these buffers entirely; if yes, avoid zeroing for performance.
• If partial writes occur, limit zeroing to the active span or use a header field indicating valid length.

I can help benchmark the impact if you share typical batch sizes and token counts.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp (1)

124-166: Harden MPI receive: validate buffer, fix dims receive pointer, and guard against invalid shapes.

Issues:

• No null/memory-type check for draftLogitsHost; receiving into GPU memory would be invalid.
• dims receive uses &dims; pass dims to be conventional and avoid type pitfalls.
• Missing non-negative shape validation; negative dims lead to invalid reshape.
• expectedSize multiplies mixed signed/unsigned; cast both dims to uint64_t.

Apply this diff:

 #include "logitsThread.h"
 
+#include "tensorrt_llm/runtime/common.h"
@@
-void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
+void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost,
     executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype)
 {
 #if ENABLE_MULTI_DEVICE
     auto const& worldComm = tensorrt_llm::mpi::MpiComm::world();
 
+    TLLM_CHECK_WITH_INFO(draftLogitsHost, "draftLogitsHost must be a valid host tensor.");
+    TLLM_CHECK_WITH_INFO(
+        draftLogitsHost->getMemoryType() != runtime::MemoryType::kGPU,
+        "draftLogitsHost must reside in host/pinned memory.");
@@
-    int64_t dims[2];
-    MPICHECK(MPI_Mrecv(&dims, count, MPI_INT64_T, &msg, &status));
+    int64_t dims[2];
+    MPICHECK(MPI_Mrecv(dims, count, MPI_INT64_T, &msg, &status));
+
+    TLLM_CHECK_WITH_INFO(dims[0] >= 0 && dims[1] >= 0, "Received invalid logits shape: [%ld, %ld].", dims[0], dims[1]);
@@
-    uint64_t const expectedSize = static_cast<uint64_t>(dims[0]) * dims[1] * tc::getDTypeSize(logitsDtype);
+    uint64_t const expectedSize = static_cast<uint64_t>(dims[0]) * static_cast<uint64_t>(dims[1])
+        * tc::getDTypeSize(logitsDtype);
     TLLM_CHECK((uint64_t) count == expectedSize);
 
     MPICHECK(MPI_Mrecv(draftLogitsHost->data(), count, MPI_UINT8_T, &msg, &status));
 #endif // ENABLE_MULTI_DEVICE
 }

cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp (3)

327-369: Guard MPI bcast on non-leader path and validate metadata before memcpy (still unresolved).

Two correctness issues remain (previously flagged):

• Potential deadlock: Non-leader unconditionally calls COMM_SESSION.bcast. If tensor parallel is disabled, leader doesn’t broadcast and non-leader hangs.
• Safety: memcpy from reqDraftLogits->data() lacks size and memory-type validation for SpeculativeDecodingFastLogitsInfo.

Fix by validating metadata size and wrapping the non-leader bcast in a worldConfig.isTensorParallel() check; also remove trailing semicolon after the function for consistency.

 void retrieveDraftLogits(TensorPtr& draftLogitsHost, std::shared_ptr<runtime::ITensor> const& reqDraftLogits,
     ModelConfig const& modelConfig, WorldConfig const& worldConfig, bool speculativeDecodingFastLogits,
     bool isLeaderInOrchMode, BufferManager const& bufferManager)
 {
     TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__);
 
     if (!speculativeDecodingFastLogits)
     {
         TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__);
         bufferManager.copy(*reqDraftLogits, *draftLogitsHost);
         return;
     }
 
     if (isLeaderInOrchMode)
     {
-        te::SpeculativeDecodingFastLogitsInfo fastLogitsInfo;
-        std::memcpy(&fastLogitsInfo, reqDraftLogits->data(), sizeof(fastLogitsInfo));
+        // Validate metadata buffer
+        TLLM_CHECK_WITH_INFO(
+            reqDraftLogits->getSizeInBytes() >= sizeof(te::SpeculativeDecodingFastLogitsInfo),
+            "Draft logits metadata buffer too small for SpeculativeDecodingFastLogitsInfo.");
+        TLLM_CHECK_WITH_INFO(
+            reqDraftLogits->getMemoryType() != MemoryType::kGPU,
+            "Fast logits metadata must be in host/pinned memory.");
+        te::SpeculativeDecodingFastLogitsInfo fastLogitsInfo{};
+        std::memcpy(&fastLogitsInfo, reqDraftLogits->data(), sizeof(te::SpeculativeDecodingFastLogitsInfo));
         utils::targetModelReceiveLogits(draftLogitsHost, fastLogitsInfo, modelConfig.getLogitsDtype());
 
         // Broadcast to other ranks if needed
         if (worldConfig.isTensorParallel())
         {
             auto const& commSession = COMM_SESSION;
             auto shape = draftLogitsHost->getShape();
             commSession.bcastValue(shape.d[0], 0);
             commSession.bcastValue(shape.d[1], 0);
             commSession.bcast(draftLogitsHost->data(), draftLogitsHost->getSizeInBytes(), mpi::MpiType::kUINT8, 0);
         }
     }
     else
     {
-        // Get logits from leader rank
-        auto const& commSession = COMM_SESSION;
-        int64_t dims[2];
-        commSession.bcastValue(dims[0], 0);
-        commSession.bcastValue(dims[1], 0);
-        draftLogitsHost->reshape(ITensor::makeShape({dims[0], dims[1]}));
-        commSession.bcast(draftLogitsHost->data(), draftLogitsHost->getSizeInBytes(), mpi::MpiType::kUINT8, 0);
+        // Receive logits from leader rank only when TP is enabled
+        if (worldConfig.isTensorParallel())
+        {
+            auto const& commSession = COMM_SESSION;
+            int64_t dims[2];
+            commSession.bcastValue(dims[0], 0);
+            commSession.bcastValue(dims[1], 0);
+            draftLogitsHost->reshape(ITensor::makeShape({dims[0], dims[1]}));
+            commSession.bcast(draftLogitsHost->data(), draftLogitsHost->getSizeInBytes(), mpi::MpiType::kUINT8, 0);
+        }
+        else
+        {
+            TLLM_THROW("Fast logits path requires tensor-parallel broadcast for non-leader ranks.");
+        }
     }
 
     TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__);
-};
+}

---

672-680: Sanity-check external draft-token availability before deriving numDecodingEngineTokens.

If getNumDraftTokens() > 0 but the external tensor is missing, downstream copies will OOB or null-deref. Fail fast here.

-        SizeType32 numDecodingEngineTokens{1};
+        SizeType32 numDecodingEngineTokens{1};
         if (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal())
         {
-            numDecodingEngineTokens = llmReq->getNumDraftTokens() + 1;
+            auto const draftLen = llmReq->getNumDraftTokens();
+            TLLM_CHECK_WITH_INFO(draftLen == 0 || llmReq->getDraftTokens(),
+                "External speculative decoding requested with %d draft tokens, but no draft-token tensor provided.",
+                draftLen);
+            numDecodingEngineTokens = draftLen + 1;
         }
         else if (!modelConfig.getSpeculativeDecodingMode().isNone())
         {
             numDecodingEngineTokens = modelConfig.getMaxDecodingTokens();
         }

---

370-421: Validate presence and shape of external draft tokens before copying.

When numDraftTokens > 0, the code dereferences llmReq.getDraftTokens() without validating presence or shape. This risks OOB copies.

 void newRequestDraftTokensExternal(DecodingInput& jointDecodingInput, SizeType32 batchIdx, LlmRequest const& llmReq,
     SizeType32 numDecodingEngineTokens, runtime::ModelConfig const& modelConfig, WorldConfig const& worldConfig,
     bool speculativeDecodingFastLogits, bool isLeaderInOrchMode, CudaStream const& decoderStream)
 {
@@
     auto const& draftTokens = llmReq.getDraftTokens();
     auto const numDraftTokens = numDecodingEngineTokens - 1;
@@
-    if (numDraftTokens > 0)
+    if (numDraftTokens > 0)
     {
+        TLLM_CHECK_WITH_INFO(draftTokens, "External draft tokens mode expects non-empty draft tokens.");
+        auto const dtShape = draftTokens->getShape();
+        TLLM_CHECK_WITH_INFO(
+            dtShape.nbDims == 2 && dtShape.d[0] == 1 && dtShape.d[1] >= numDraftTokens,
+            tc::fmtstr("Unexpected draft tokens shape. Expected [1, >= %d], got [%d, %d].",
+                numDraftTokens, (dtShape.nbDims > 0 ? (int) dtShape.d[0] : -1),
+                (dtShape.nbDims > 1 ? (int) dtShape.d[1] : -1)));
         TensorPtr draftTokenIdsHostSlice
             = ITensor::slice(externalDraftTokensInputs->draftTokenIdsHost, {batchIdx, 0}, numDraftTokens);
         // Copy to pinned host memory (don't care about stream of bufferManager)
         decoderBufferManager.copy(draftTokens->data(), *draftTokenIdsHostSlice);
@@
     if (useDraftLogits)
     {
         TensorPtr draftLogitsHostSlice
             = ITensor::slice(externalDraftTokensInputs->draftLogitsHost, {batchIdx, 0}, numDraftTokens);
         retrieveDraftLogits(draftLogitsHostSlice, draftLogits.value(), modelConfig, worldConfig,
             speculativeDecodingFastLogits, isLeaderInOrchMode, decoderBufferManager);

cpp/include/tensorrt_llm/runtime/decodingInput.h (1)

104-118: Document memory-type and shape invariants for host draft buffers.

Please add brief Doxygen comments for draftLogitsHost and draftTokenIdsHost clarifying:

• Expected memory type (pinned host via MemoryType::kPINNEDPOOL).
• Shape invariants mirroring device tensors (same 2D shapes as draftLogits and draftTokenIds per step).
• Dtype consistency with logits/token IDs.

This will reduce future misuse and ease maintenance.

 class ExternalDraftTokensInputs
 {
 public:
     TensorPtr draftLogits;
+    //! Host-pinned counterpart of draftLogits. Shape and dtype must mirror draftLogits.
+    //! Allocated from pinned pool (MemoryType::kPINNEDPOOL).
     TensorPtr draftLogitsHost;
@@
     TensorPtr numDraftTokensHost;
     TensorPtr draftTokenIds;
+    //! Host-pinned counterpart of draftTokenIds. Shape must mirror draftTokenIds.
+    //! Allocated from pinned pool (MemoryType::kPINNEDPOOL).
     TensorPtr draftTokenIdsHost;

cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp (2)

505-567: Eagle setup: host-side control tensors and static-tree init are reasonable.

Initialization of ctx/gen request types and lengths, and static-tree path generation with validation against module depth, look sound. Consider adding a brief comment clarifying that dynamic-tree handling is deferred when eagleConfig->useDynamicTree() is true.

---

569-621: Avoid zeroing large nextDraftTokens buffers unless required.

Zeroing nextDraftTokens (and length) eagerly can be costly at scale. If kernels overwrite these fully, skip zeroing or restrict it to variableDraftLength cases where length buffers are read before write.

-        TensorPtr nextDraftTokens = ITensor::slice(speculativeDecodingOutputs.nextDraftTokens, batchIdx, 1);
-        // FIXME: can we skip this?
-        manager.setZero(*nextDraftTokens);
+        TensorPtr nextDraftTokens = ITensor::slice(speculativeDecodingOutputs.nextDraftTokens, batchIdx, 1);
+        // Zero only if downstream reads before write; otherwise omit to save time.
+        // manager.setZero(*nextDraftTokens);
         if (speculativeDecodingMode.variableDraftLength())
         {
             TensorPtr nextDraftTokensLen = ITensor::slice(speculativeDecodingOutputs.nextDraftTokensLen, batchIdx, 1);
             manager.setZero(*nextDraftTokensLen);
         }

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

cpp/include/tensorrt_llm/batch_manager/llmRequest.h (7)

1111-1114: Null-safe draft-token count is correct; add explicit cast and noexcept (minor).

Good defensive change. To avoid implicit narrowing warnings and match other trivial getters, add an explicit cast and noexcept.

-    [[nodiscard]] SizeType32 getNumDraftTokens() const
+    [[nodiscard]] SizeType32 getNumDraftTokens() const noexcept
     {
-        return hasDraftTokens() ? mDraftTokens->size() : 0;
+        return hasDraftTokens() ? static_cast<SizeType32>(mDraftTokens->size()) : 0;
     }

---

1067-1071: Remove redundant check.

The second TLLM_CHECK duplicates the preceding TLLM_CHECK_WITH_INFO on the same condition and adds no value.

         TLLM_CHECK_WITH_INFO(prepopulatedPromptLen < promptLen,
             "Invalid state: prepopulatedPromptLen (%d) >= promptLen (%d) for request %lu", prepopulatedPromptLen,
             promptLen, mRequestId);
-        TLLM_CHECK(prepopulatedPromptLen < promptLen);

---

1384-1388: Make hasAdditionalOutputs() a const member.

Pure query; should be const.

-    bool hasAdditionalOutputs()
+    bool hasAdditionalOutputs() const
     {
         return !mAdditionalContextOutputTensors.empty() || !mAdditionalGenerationOutputTensors.empty();
     }

---

1374-1377: Const-correctness and explicit cast for getGenerationLogitsFragmentsSize().

This simple accessor can be const; add explicit cast to SizeType32.

-    SizeType32 getGenerationLogitsFragmentsSize()
+    SizeType32 getGenerationLogitsFragmentsSize() const
     {
-        return mGenerationLogitsFragments.size();
+        return static_cast<SizeType32>(mGenerationLogitsFragments.size());
     }

---

2156-2165: Prefer std::copy over std::memcpy for type safety.

Avoid C-library functions when possible (guideline). Using std::copy improves type-safety and readability.

-        auto* data = runtime::bufferCast<int32_t>(*tensor);
-        std::memcpy(data, words.data(), numWords * sizeof(int32_t));
-        std::memcpy(data + numWords, offsets.data(), numWords * sizeof(int32_t));
+        auto* data = runtime::bufferCast<int32_t>(*tensor);
+        std::copy(words.begin(), words.end(), data);
+        std::copy(offsets.begin(), offsets.end(), data + numWords);

---

1-15: Nit: update copyright year.

Guideline asks for current year; consider updating to 2025.

- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.

---

17-19: Header guards per guideline (non-blocking).

Project guideline prefers include guards (TRTLLM__H). You currently use #pragma once. Consider adding guards for portability.

+#ifndef TRTLLM_LLMREQUEST_H
+#define TRTLLM_LLMREQUEST_H
 #pragma once
 ...
-} // namespace tensorrt_llm::batch_manager
+} // namespace tensorrt_llm::batch_manager
+
+#endif // TRTLLM_LLMREQUEST_H

Also applies to: 2377-2378

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.