[None][fix] Use prefer_pinned() in pard.py

[mikeiovine]

Description

Use prefer_pinned() in pard.py to unblock pre-commit pipeline.

Test Coverage

Existing tests.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

• Chores
  • Refined memory pinning behavior in speculative decoding operations to improve runtime efficiency and resource utilization.

[brb-nv]

LGTM.

[mikeiovine]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

Replaces hard-coded pin_memory=True with a dynamic pin_memory=prefer_pinned() call for batch indices in PARD speculative decoding. Adds corresponding import of prefer_pinned utility function.

Changes

Cohort / File(s)	Summary
Memory pinning configuration tensorrt_llm/_torch/speculative/pard.py	Replaced hard-coded pin_memory=True with dynamic prefer_pinned() function call in batch_indices tensor creation. Added import of prefer_pinned from tensorrt_llm._utils to enable runtime-determined memory pinning behavior.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: replacing hard-coded pin_memory behavior with prefer_pinned() in pard.py to resolve a pre-commit pipeline issue.
Description check	✅ Passed	The description briefly explains the issue (unblock pre-commit pipeline) and solution (use prefer_pinned()), though it lacks specific technical details about what the pre-commit check was and why this fix resolves it.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

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

7-7: Use module import to preserve namespace in Python imports.

Line 7 imports a function directly; project guidelines require importing the module and calling through its namespace.

Proposed change

-from tensorrt_llm._utils import prefer_pinned
+from tensorrt_llm import _utils
@@
-            num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned()
+            num_seqs, dtype=torch.int, device="cpu", pin_memory=_utils.prefer_pinned()

As per coding guidelines: "When importing in Python, always maintain the namespace. Import the module, not individual classes or functions."

Also applies to: 39-39

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/speculative/pard.py` at line 7, Replace the direct import
of the function with a module import to preserve namespace: change "from
tensorrt_llm._utils import prefer_pinned" to import the module (e.g., "import
tensorrt_llm._utils as _utils") and update all calls to prefer_pinned to use the
module namespace (e.g., _utils.prefer_pinned); apply the same change for the
other direct import referenced at line 39 so every usage references
tensorrt_llm._utils.<function_name> instead of the bare function name.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tensorrt_llm/_torch/speculative/pard.py`:
- Line 7: Replace the direct import of the function with a module import to
preserve namespace: change "from tensorrt_llm._utils import prefer_pinned" to
import the module (e.g., "import tensorrt_llm._utils as _utils") and update all
calls to prefer_pinned to use the module namespace (e.g., _utils.prefer_pinned);
apply the same change for the other direct import referenced at line 39 so every
usage references tensorrt_llm._utils.<function_name> instead of the bare
function name.

---

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 80aa8ca and 31c492f.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/speculative/pard.py

[mikeiovine]

/bot skip --comment "Pre-commit pipeline is fixed"

[tensorrt-cicd]

PR_Github #36975 [ skip ] triggered by Bot. Commit: 31c492f Link to invocation

[tensorrt-cicd]

PR_Github #36975 [ skip ] completed with state SUCCESS. Commit: 31c492f
Skipping testing for commit 31c492f

Link to invocation