tensorrt_llm/_torch/models/modeling_speculative.py (1)

171-175: Register non-trainable d2t as a buffer instead of a Parameter
d2t is static; registering it as a buffer excludes it from optimizer and state-dict overhead—existing .data and tensor indexing still work.
Apply:

-            self.d2t = nn.Parameter(torch.empty((config.draft_vocab_size, ),
-                                                dtype=torch.int32),
-                                    requires_grad=False)
+            self.register_buffer(
+                "d2t",
+                torch.empty((config.draft_vocab_size,), dtype=torch.int32),
+                persistent=False,
+            )

cpp/tensorrt_llm/kernels/logitsBitmask.h (1)

32-35: Document invokeContiguousLogitsBitmask API contract

• logits: contiguous row-major [batchSize, vocabSizePadded]
• bitmask: [batchSize, bitmaskSize] where bitmaskSize = ceilDiv(vocabSizePadded, 32) uint32 elements
• tokenMask: optional (nullptr allowed); when non-null, 0 at index batchIdx skips that row
• d2t: optional (nullptr allowed); when non-null, adds per-token index delta
• batchSize, vocabSizePadded, bitmaskSize: int32 dimensions

Replace #pragma once in logitsBitmask.h with include guards [nitpick]

tensorrt_llm/_torch/attention_backend/interface.py (1)

338-342: Guard restore against empty state and add idempotence

Avoid no-op loops if nothing was saved.

-    def restore_from_spec_dec(self) -> None:
-        for f, v in self._saved_tensors.items():
-            setattr(self, f, v)
-        self._saved_tensors.clear()
+    def restore_from_spec_dec(self) -> None:
+        if not self._saved_tensors:
+            return
+        for f, v in self._saved_tensors.items():
+            setattr(self, f, v)
+        self._saved_tensors.clear()

tensorrt_llm/_torch/speculative/drafting_loops.py (1)

133-136: Avoid .data; use the tensor directly

.data is discouraged; it bypasses safety checks. This tensor is non-trainable, so direct usage is fine.

-        if hasattr(self.draft_model.model, "d2t"):
-            d2t = self.draft_model.model.d2t.data
-            return tokens + d2t[tokens]
+        if hasattr(self.draft_model.model, "d2t"):
+            d2t = self.draft_model.model.d2t
+            return tokens + d2t[tokens]

cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp (2)

27-32: Add validation for empty batch size edge case.

Consider adding an assertion or warning for empty batch to ensure the early return is intentional and documented.

 int32_t const batchSize = logits.size(0);
 if (batchSize == 0)
 {
+    // Early return for empty batch - no work to do
     return;
 }

---

33-44: Consider extracting common validation logic.

The repeated validation pattern could be refactored into a helper function for better maintainability.

+namespace
+{
+void validateTensor2D(torch::Tensor const& tensor, char const* name, int32_t expectedBatchSize)
+{
+    TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor.");
+    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous.");
+    TORCH_CHECK(tensor.dim() == 2, name, " must be a 2D tensor.");
+    TORCH_CHECK(tensor.size(0) == expectedBatchSize, name, " must have the same batch size as logits.");
+}
+} // namespace
+
 TORCH_CHECK(bitmask.size(0) == batchSize, "bitmask must have the same batch size as logits.");
 
 int32_t vocabSizePadded = logits.size(1);
 int32_t bitmaskSize = bitmask.size(1);
-TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor.");
-TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous.");
-TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor.");
-TORCH_CHECK(bitmask.is_cuda(), "bitmask must be a CUDA tensor.");
-TORCH_CHECK(bitmask.is_contiguous(), "bitmask must be contiguous.");
-TORCH_CHECK(bitmask.dim() == 2, "bitmask must be a 2D tensor.");
+validateTensor2D(logits, "logits", batchSize);
+validateTensor2D(bitmask, "bitmask", batchSize);

tests/unittest/_torch/thop/parallel/test_logits_bitmask_op.py (1)

12-37: Good coverage of token_mask path; small cleanup possible.

The ref uses slicing [::stride] to mirror token_mask semantics and validates in-place behavior. Consider adding one randomized token_mask case (not only stride-based) to catch irregular masks; also you can drop the redundant re-cast since (% stride == 0) already yields bool:

-    if stride > 1:
-        token_mask = torch.arange(batch_size, dtype=torch.int32,
-                                  device="cuda") % stride == 0
-        token_mask = token_mask.to(torch.int32)
+    if stride > 1:
+        token_mask = (torch.arange(batch_size, device="cuda") % stride == 0).to(torch.int32)

tensorrt_llm/_torch/speculative/eagle3.py (1)

299-299: Guard prepare_for_spec_dec for optional CUDA fields
prepare_for_spec_dec asserts each field is a Tensor; since _seq_lens_cuda defaults to None until seq_lens is set, calling prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") can trigger an AssertionError. Replace with a guarded call:

-        attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+        fields = ["_seq_lens"]
+        if isinstance(attn_metadata._seq_lens_cuda, torch.Tensor):
+            fields.append("_seq_lens_cuda")
+        attn_metadata.prepare_for_spec_dec(*fields)

cpp/tensorrt_llm/kernels/logitsBitmask.cu (2)

257-292: Static SM count caching may be wrong across multi-GPU contexts.

static int const smCount = getMultiProcessorCount(); is process-static. If device changes between calls, grid sizing can be off. Prefer querying per-call or caching per-device id.

-    static int const smCount = tensorrt_llm::common::getMultiProcessorCount();
+    int const smCount = tensorrt_llm::common::getMultiProcessorCount();

If you want caching, key it by cudaGetDevice() id.

---

51-63: Packed reinterpret_cast relies on alignment; consider safer packing.

reinterpret_cast<PackedT*> of stack arrays can violate alignment rules and hamper SASS generation on some archs. Use a pack helper that writes through a PackedT lvalue:

-template <typename T, typename PackedT>
-__device__ PackedT packedNegativeInfinity()
+template <typename T, typename PackedT>
+__device__ PackedT packedNegativeInfinity()
 {
-    int constexpr kAlignment = sizeof(PackedT) / sizeof(T);
-    T packed[kAlignment];
-#pragma unroll
-    for (int i = 0; i < kAlignment; i++)
-    {
-        packed[i] = negativeInfinity<T>();
-    }
-    return *reinterpret_cast<PackedT*>(packed);
+    int constexpr kAlignment = sizeof(PackedT) / sizeof(T);
+    PackedT v;
+    T* elems = reinterpret_cast<T*>(&v);
+#pragma unroll
+    for (int i = 0; i < kAlignment; i++) { elems[i] = negativeInfinity<T>(); }
+    return v;
 }

Similarly, consider using the same pattern where you load/store via PackedT to avoid potential misalignment penalties.

Also applies to: 100-114

tensorrt_llm/_torch/pyexecutor/guided_decoder.py (5)

140-141: Confirm token_mask dtype contract (int32) with the custom op.

If the kernel accepts narrower or bool masks, consider switching to torch.bool/torch.int8 to cut memory bandwidth. Otherwise, add a brief comment stating int32 is required by torch.ops.trtllm.logits_bitmask.

---

273-282: Add fast-path and sanity checks in _copy_bitmask.

Short-circuit when there’s nothing to copy and assert basic shape/dtype invariants to catch integration errors early.

 def _copy_bitmask(self,
                   requests: GuidedRequests,
                   num_bitmask_tokens: Optional[int] = None) -> None:
     if num_bitmask_tokens is None:
         num_bitmask_tokens = requests.num_bitmask_tokens
+    if num_bitmask_tokens == 0:
+        return
+    # Basic invariants (cheap and helpful in CI)
+    assert self.bitmask.dtype == self.bitmask_dtype
+    assert self.token_mask.dtype == self.token_mask_dtype
+    assert self.bitmask_host.size(0) >= num_bitmask_tokens
+    assert self.token_mask_host.size(0) >= num_bitmask_tokens
     self.bitmask[:num_bitmask_tokens].copy_(
         self.bitmask_host[:num_bitmask_tokens], non_blocking=True)
     self.token_mask[:num_bitmask_tokens].copy_(
         self.token_mask_host[:num_bitmask_tokens], non_blocking=True)

---

287-299: Gracefully handle op availability and empty work; validate shapes.

Add early return for zero rows, verify the custom op exists, and assert batch-row alignment to fail fast on mismatches.

 @torch.inference_mode()
 def _apply_bitmask(self,
                    requests: GuidedRequests,
                    logits: torch.Tensor,
                    d2t: Optional[torch.Tensor] = None,
                    num_bitmask_tokens: Optional[int] = None) -> None:
     if num_bitmask_tokens is None:
         num_bitmask_tokens = requests.num_bitmask_tokens
+    if num_bitmask_tokens == 0:
+        return
+    # Ensure op is registered
+    if not hasattr(torch.ops, "trtllm") or not hasattr(torch.ops.trtllm, "logits_bitmask"):
+        raise RuntimeError("trtllm.logits_bitmask custom op is not available. Ensure extensions are built/loaded.")
+    # Basic shape checks
+    assert logits.size(0) >= num_bitmask_tokens, (
+        f"logits rows {logits.size(0)} < num_bitmask_tokens {num_bitmask_tokens}")
+    assert self.bitmask.size(0) >= num_bitmask_tokens
+    assert self.token_mask.size(0) >= num_bitmask_tokens
     torch.ops.trtllm.logits_bitmask(
         logits[:num_bitmask_tokens],
         self.bitmask[:num_bitmask_tokens],
         token_mask=self.token_mask[:num_bitmask_tokens],
         d2t=d2t)

If you want, I can push a follow-up commit with these guards.

---

311-313: Public API: briefly document num_bitmask_tokens override.

Add a one-line docstring explaining when/why callers should pass num_bitmask_tokens (e.g., CUDA stream timing during drafting).

---

317-323: Public API: include d2t/num_bitmask_tokens in docstring.

The method’s behavior changed; reflect the new parameters in the docstring for discoverability.

Configuration used: Path: .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.

Learnt from: pcastonguay
PR: NVIDIA/TensorRT-LLM#7455
File: tensorrt_llm/_torch/pyexecutor/py_executor.py:1852-1860
Timestamp: 2025-09-02T13:42:44.885Z
Learning: In MPI communication within TensorRT-LLM pipeline parallelism, different communication types (tokens, logits, termination sync) must use disjoint tag namespaces to avoid message routing collisions when using the same source/destination patterns.