[#11529][perf] AD host time attention MD optimization for large context

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -28,6 +28,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, Literal, Optional, Protocol, Sequence, Set, Tuple, Type, Union
 
+import numpy as np
 import torch
 from torch._ops import OpOverloadPacket
 from torch.fx import Node
@@ -40,6 +41,29 @@
 
 Constant = Union[int, float, str, None]
 
+# Torch dtype → numpy dtype for fast list-to-tensor conversion.
+# numpy's list→array conversion is ~2-3x faster than torch.tensor(list) for large lists.
+_TORCH_TO_NUMPY_DTYPE: Dict[torch.dtype, np.dtype] = {
+    torch.int: np.int32,
+    torch.int32: np.int32,
+    torch.int64: np.int64,
+    torch.long: np.int64,
+    torch.float: np.float32,
+    torch.float32: np.float32,
+    torch.float64: np.float64,
+    torch.double: np.float64,
+    torch.float16: np.float16,
+    torch.bool: np.bool_,
+}
+
+
+def _list_to_tensor(data: list, dtype: torch.dtype) -> torch.Tensor:
+    """Convert a Python list to a tensor, using numpy for speed."""
+    np_dtype = _TORCH_TO_NUMPY_DTYPE.get(dtype)
+    if np_dtype is not None:
+        return torch.from_numpy(np.array(data, dtype=np_dtype))
+    return torch.tensor(data, dtype=dtype)
+
 
 class PrepareMetadataHostCallable(Protocol):
     def __call__(self, **sequence_info_args: torch.Tensor) -> None: ...
@@ -184,15 +208,15 @@ def get_current_length(self, name: str) -> int:
     def store(
         self,
         name: str,
-        data: List[Number],
+        data: torch.Tensor,
         fill_value: Optional[Number] = None,
     ) -> int:
-        """Store data into the host buffer.
+        """Store a tensor into the pinned host buffer.
 
         Args:
             name: Name of the tensor to store to.
-            data: List of values to store.
-            fill_value: Optional value to fill the entire tensor with before storing.
+            data: 1-D torch.Tensor to store.
+            fill_value: Optional value to fill the entire buffer with before storing.
 
         Returns:
             Number of elements stored.
@@ -203,11 +227,12 @@ def store(
         if fill_value is not None:
             host_view.fill_(fill_value)
 
-        length = len(data)
+        length = data.numel()
         assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}"
-
-        temp_tensor = torch.tensor(data, dtype=dtype)
-        host_view[:length].copy_(temp_tensor)
+        # Use numpy for the memcpy into pinned memory — avoids torch dispatcher overhead
+        dst = host_view[:length].numpy()
+        src = (data if data.dtype == dtype else data.to(dtype)).numpy()
+        np.copyto(dst, src)
 
         self._current_lengths[name] = length
         return length
@@ -481,8 +506,10 @@ def __init__(
             f"{name}_host" for name in self._input_buffer.tensor_names
         }
 
-        # Initialize args_list from tensor specs
-        self._args_list: Dict[str, List[int]] = {spec[0]: [0] * spec[1] for spec in tensor_specs}
+        # Initialize args_list from tensor specs (all entries are tensors)
+        self._args_list: Dict[str, torch.Tensor] = {
+            spec[0]: torch.zeros(spec[1], dtype=spec[2]) for spec in tensor_specs
+        }
 
         self._active_args = ("input_ids", "position_ids")
         self._shapeable_args = ("input_ids", "position_ids", "input_ids_host", "position_ids_host")
@@ -577,19 +604,19 @@ def args(self) -> Tuple[torch.Tensor, ...]:
 
     @property
     def seq_len(self) -> List[int]:
-        return self._args_list["seq_len"].copy()
+        return self._args_list["seq_len"].tolist()
 
     @property
     def input_pos(self) -> List[int]:
-        return self._args_list["input_pos"].copy()
+        return self._args_list["input_pos"].tolist()
 
     @property
     def cache_loc(self) -> List[int]:
-        return self._args_list["cache_loc"].copy()
+        return self._args_list["cache_loc"].tolist()
 
     @property
     def pages_per_seq(self) -> List[int]:
-        return self._args_list["pages_per_seq"].copy()
+        return self._args_list["pages_per_seq"].tolist()
 
     @property
     def num_sequences(self) -> int:
@@ -661,8 +688,7 @@ def update_cache_information(self, num_blocks: int, block_offset_multiplier: int
             # Resize all truncatable page tensors together (they share the same max size)
             for tensor_name in ("cache_loc", "page_seq_indices", "page_in_seq"):
                 self._input_buffer.resize(tensor_name, estimated_capacity)
-                old_size = len(self._args_list[tensor_name])
-                self._args_list[tensor_name].extend([0] * (estimated_capacity - old_size))
+                self._args_list[tensor_name] = torch.zeros(estimated_capacity, dtype=torch.int)
 
     @staticmethod
     def _get_page_assignments(
@@ -824,22 +850,30 @@ def _flatten(nested_seqs: Sequence[Sequence[int]]) -> List[int]:
     def _store_arg(
         self,
         name: str,
-        tnsr_like: List[Number],
+        data: "Union[List[Number], torch.Tensor]",
         reset_val: Optional[Number] = None,
     ) -> None:
         """Store the argument into the pinned host buffer for later batch transfer to device.
 
         The data is stored in the host-side pinned memory buffer managed by InputBuffer.
         The actual H2D transfer happens in a single batch at the end of nest_sequences().
 
+        Lists are converted to tensors at the boundary so the rest of the pipeline is
+        tensor-only.
+
         Args:
             name: Name of the argument to store.
-            tnsr_like: List of values to store.
+            data: List of values or a 1-D torch.Tensor to store.
             reset_val: Value to reset/fill the tensor with before writing data.
         """
         with nvtx_range(f"ad_store_on_host_seq_info_arg_{name}"):
-            # Always store list object for Python access
-            self._args_list[name] = tnsr_like.copy()
+            # Convert to tensor at the boundary (numpy is ~2-3x faster than torch.tensor for large lists)
+            # TODO: move this to self._input_buffer.store() when _args_list get deprecated
+            if not isinstance(data, torch.Tensor):
+                _, dtype = self._input_buffer._tensor_specs[name]
+                data = _list_to_tensor(data, dtype)
+
+            self._args_list[name] = data
 
             # Only store to buffer when the argument is active or requires copy
             is_active = name in self._active_args or f"{name}_host" in self._active_args
@@ -848,7 +882,7 @@ def _store_arg(
                 return
 
             # Store to the InputBuffer's pinned host memory
-            self._input_buffer.store(name, tnsr_like, fill_value=reset_val)
+            self._input_buffer.store(name, data, fill_value=reset_val)
 
     def _store_extra_arg(
         self, name: str, tnsr_like: Optional[Union[torch.Tensor, Sequence[torch.Tensor]]]
@@ -981,24 +1015,27 @@ def nest_sequences(
             self._store_arg("cache_loc", cache_loc)
             self._store_arg("pages_per_seq", pages_per_seq)
 
-            # Auto-compute page_seq_indices and page_in_seq from pages_per_seq.
+            # Resolve cu_num_pages: use caller-provided value or derive from pages_per_seq
+            if cu_num_pages is None:
+                pps_t = self._args_list["pages_per_seq"]
+                cu_num_pages = torch.zeros(len(pps_t) + 1, dtype=torch.int)
+                cu_num_pages[1:] = pps_t.cumsum(0)
+            self._store_arg("cu_num_pages", cu_num_pages)
+
+            # Compute page_seq_indices and page_in_seq using vectorized torch ops,
+            # reusing the stored cu_num_pages tensor instead of recomputing the cumsum.
             # page_seq_indices[j] = which sequence page j belongs to
             # page_in_seq[j] = which page within that sequence (0-indexed)
-            page_seq_indices = []
-            page_in_seq_vals = []
-            for i, n_pages in enumerate(pages_per_seq):
-                page_seq_indices.extend([i] * n_pages)
-                page_in_seq_vals.extend(range(n_pages))
-            self._store_arg("page_seq_indices", page_seq_indices)
-            self._store_arg("page_in_seq", page_in_seq_vals)
-
-        # update cumulative number of pages
-        if cu_num_pages is None:
-            pages_per_seq = self.pages_per_seq
-            cu_num_pages = torch.zeros(len(pages_per_seq) + 1, dtype=torch.int)
-            cu_num_pages[1:] = torch.cumsum(torch.tensor(pages_per_seq), dim=0)
-            cu_num_pages = cu_num_pages.tolist()
-        self._store_arg("cu_num_pages", cu_num_pages)
+            pages_per_seq_t = self._args_list["pages_per_seq"]
+            cu_pages_t = self._args_list["cu_num_pages"]
+            seq_indices = torch.arange(len(pages_per_seq), dtype=torch.int)
+            page_seq_indices_t = torch.repeat_interleave(seq_indices, pages_per_seq_t)
+            total_pages = cu_pages_t[-1].item()
+            page_in_seq_t = torch.arange(total_pages, dtype=torch.int) - torch.repeat_interleave(
+                cu_pages_t[:-1], pages_per_seq_t
+            )
+            self._store_arg("page_seq_indices", page_seq_indices_t)
+            self._store_arg("page_in_seq", page_in_seq_t)
 
         # update sequence length with cache
         if seq_len_with_cache is None: