[https://nvbugs/5911143][fix] add async worker to MTP/Eagle3 sampler,… (NVIDIA#11573)

Signed-off-by: Dan Hansen <1+dhansen-nvidia@users.noreply.github.com>
Signed-off-by: dhansen-nvidia <218031328+dhansen-nvidia@users.noreply.github.com>
Co-authored-by: Dan Hansen <1+dhansen-nvidia@users.noreply.github.com>
Co-authored-by: Venky <23023424+venkywonka@users.noreply.github.com>

Author: 3 people

.pre-commit-config.yaml

@@ -1473,3 +1473,8 @@ repos:
         entry: ./scripts/dco_check.py
         language: script
         stages: [commit-msg]
+    -   id: pinned memory policy
+        name: Disallow raw pinned-memory APIs in runtime code
+        entry: ./scripts/check_pinned_memory_usage.py
+        language: script
+        files: ^(tensorrt_llm|triton_backend)/.*\.py$

scripts/check_pinned_memory_usage.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+import pathlib
+import sys
+
+
+class PinnedMemoryUsageChecker(ast.NodeVisitor):
+    def __init__(self, *, allow_direct_pin_memory: bool) -> None:
+        self.allow_direct_pin_memory = allow_direct_pin_memory
+        self.violations: list[tuple[int, str]] = []
+
+    def visit_Call(self, node: ast.Call) -> None:
+        if isinstance(node.func, ast.Attribute) and node.func.attr == "pin_memory":
+            if not self.allow_direct_pin_memory:
+                self.violations.append(
+                    (
+                        node.lineno,
+                        "Use `maybe_pin_memory(tensor)` instead of direct `.pin_memory()`.",
+                    )
+                )
+
+        for keyword in node.keywords:
+            if (
+                keyword.arg == "pin_memory"
+                and isinstance(keyword.value, ast.Constant)
+                and keyword.value.value is True
+            ):
+                self.violations.append(
+                    (
+                        node.lineno,
+                        "Use `pin_memory=prefer_pinned()` instead of `pin_memory=True`.",
+                    )
+                )
+
+        self.generic_visit(node)
+
+
+def _check_file(path: pathlib.Path) -> list[tuple[int, str]]:
+    try:
+        source = path.read_text(encoding="utf-8")
+    except OSError as exc:
+        return [(0, f"Failed to read file: {exc}")]
+
+    try:
+        tree = ast.parse(source, filename=str(path))
+    except SyntaxError as exc:
+        return [(exc.lineno or 0, f"Failed to parse file: {exc.msg}")]
+
+    allow_direct_pin_memory = path.as_posix().endswith("tensorrt_llm/_utils.py")
+    checker = PinnedMemoryUsageChecker(allow_direct_pin_memory=allow_direct_pin_memory)
+    checker.visit(tree)
+    return checker.violations
+
+
+def main(argv: list[str]) -> int:
+    if len(argv) <= 1:
+        return 0
+
+    has_violations = False
+    for file_arg in argv[1:]:
+        path = pathlib.Path(file_arg)
+        violations = _check_file(path)
+        for lineno, message in violations:
+            has_violations = True
+            print(f"{path}:{lineno}: {message}")
+
+    if has_violations:
+        print("\nPinned-memory policy check failed.")
+        print("Use `tensorrt_llm._utils.maybe_pin_memory()` for direct pinning.")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv))

tensorrt_llm/_torch/attention_backend/interface.py

@@ -14,6 +14,7 @@
     from ..speculative.interface import SpecMetadata
     from ..speculative.spec_tree_manager import SpecTreeManager
 
+from tensorrt_llm._utils import maybe_pin_memory
 from tensorrt_llm.functional import (PositionEmbeddingType, RopeEmbeddingUtils,
                                      RotaryScalingType)
 from tensorrt_llm.mapping import Mapping
@@ -199,7 +200,7 @@ def seq_lens(self, value: Optional[torch.Tensor]):
 
         # The model executor sets seq_lens to None initially.
         if self._seq_lens is not None:
-            self._seq_lens = self._seq_lens.pin_memory()
+            self._seq_lens = maybe_pin_memory(self._seq_lens)
 
             if self.is_cuda_graph and self._seq_lens_cuda is not None:
                 # Very important: do not reallocate if we are using CUDA graphs.
@@ -249,7 +250,7 @@ def seq_lens_kv(self, value: Optional[torch.Tensor]):
         self.on_update()
         # The model executor sets seqlens to None initially.
         if self._seq_lens_kv is not None:
-            self._seq_lens_kv = self._seq_lens_kv.pin_memory()
+            self._seq_lens_kv = maybe_pin_memory(self._seq_lens_kv)
             self._seq_lens_kv_cuda = self._seq_lens_kv.cuda(non_blocking=True)
 
     @property

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -18,7 +18,7 @@
 from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm._torch.utils import maybe_compile, maybe_compiled_cat
-from tensorrt_llm._utils import get_size_in_bytes, get_sm_version
+from tensorrt_llm._utils import get_size_in_bytes, get_sm_version, prefer_pinned
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.bindings.internal.batch_manager import \
@@ -339,7 +339,7 @@ def __post_init__(self):
         self.host_indexer_k_cache_block_offsets = torch.zeros_like(
             self.indexer_k_cache_block_offsets,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
 
         if not self.enable_context_mla_with_cached_kv:
@@ -353,7 +353,7 @@ def __post_init__(self):
             self.host_ctx_cached_token_indptr = torch.zeros_like(
                 self.ctx_cached_token_indptr,
                 device='cpu',
-                pin_memory=True,
+                pin_memory=prefer_pinned(),
             )
             self.ctx_kv_indptr = self.get_empty(
                 self.cuda_graph_buffers,
@@ -365,7 +365,7 @@ def __post_init__(self):
             self.host_ctx_kv_indptr = torch.zeros_like(
                 self.ctx_kv_indptr,
                 device='cpu',
-                pin_memory=True,
+                pin_memory=prefer_pinned(),
             )
 
         # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation.
@@ -385,7 +385,7 @@ def __post_init__(self):
         self.host_gen_cached_token_indptr = torch.zeros_like(
             self.gen_cached_token_indptr,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         self.gen_kv_indptr = self.get_empty(
             self.cuda_graph_buffers,
@@ -397,7 +397,7 @@ def __post_init__(self):
         self.host_gen_kv_indptr = torch.zeros_like(
             self.gen_kv_indptr,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         # Indexer metadata
         # Separate slot mappings for non-interleaved layout (flat byte indices)
@@ -411,7 +411,7 @@ def __post_init__(self):
         self.host_slot_mapping_fp8 = torch.zeros_like(
             self.slot_mapping_fp8,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         self.slot_mapping_scale = self.get_empty(
             self.cuda_graph_buffers,
@@ -423,7 +423,7 @@ def __post_init__(self):
         self.host_slot_mapping_scale = torch.zeros_like(
             self.slot_mapping_scale,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         # Per-token request index buffer for topk_indices conversion
         self.req_idx_per_token = self.get_empty(
@@ -474,7 +474,7 @@ def __post_init__(self):
             self.host_topk_indices_buffer = torch.zeros_like(
                 self.topk_indices_buffer,
                 device='cpu',
-                pin_memory=True,
+                pin_memory=prefer_pinned(),
             )
         # Create expanded buffers for MTP support
         self.create_expanded_buffers(capture_graph=capture_graph)
@@ -491,7 +491,7 @@ def create_expanded_buffers(self, capture_graph=False):
         self.kv_lens_expanded_host = torch.zeros_like(
             self.kv_lens_expanded_cuda,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         self.block_table_expanded = self.get_empty(
             self.cuda_graph_buffers,
@@ -506,7 +506,7 @@ def create_expanded_buffers(self, capture_graph=False):
         self.host_block_table_expanded = torch.zeros_like(
             self.block_table_expanded,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
         self.scheduler_metadata_buffer_expanded = self.get_empty(
             self.cuda_graph_buffers,
@@ -1171,12 +1171,10 @@ def prepare(metadata: DSAtrtllmAttentionMetadata):
             total_kv_per_request = seq_lens[:
                                             num_contexts] + start_positions[:
                                                                             num_contexts]
-            host_slot_mapping_fp8_fullkv = torch.empty(total_kv_len,
-                                                       dtype=torch.int64,
-                                                       pin_memory=True)
-            host_slot_mapping_scale_fullkv = torch.empty(total_kv_len,
-                                                         dtype=torch.int64,
-                                                         pin_memory=True)
+            host_slot_mapping_fp8_fullkv = torch.empty(
+                total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned())
+            host_slot_mapping_scale_fullkv = torch.empty(
+                total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned())
 
             req_indices = torch.repeat_interleave(
                 torch.arange(num_contexts, dtype=torch.int64, device='cpu'),

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -15,7 +15,7 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
 from tensorrt_llm._torch.pyexecutor.resource_manager import (BlockManager,
                                                              KVCacheManager)
-from tensorrt_llm._utils import get_size_in_bytes
+from tensorrt_llm._utils import get_size_in_bytes, prefer_pinned
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.bindings.internal.batch_manager import \
@@ -143,7 +143,7 @@ def __post_init__(self):
         self.host_kt_cache_block_offsets = torch.zeros_like(
             self.kt_cache_block_offsets,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
 
         # Number of KT tokens for each sequence
@@ -594,7 +594,7 @@ def __post_init__(self):
         self.host_kt_cache_block_offsets = torch.zeros_like(
             self.kt_cache_block_offsets,
             device='cpu',
-            pin_memory=True,
+            pin_memory=prefer_pinned(),
         )
 
     def prepare(self) -> None: