From 663137da921d2a2c3b0fc29aa18503603ede8571 Mon Sep 17 00:00:00 2001 From: Eran Geva Date: Tue, 17 Feb 2026 05:44:30 -0800 Subject: [PATCH 1/5] improved host time Signed-off-by: Eran Geva --- .../custom_ops/attention_interface.py | 87 ++++++++++++++----- 1 file changed, 66 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 415b434fb16f..ae5326cfd831 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/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 @@ -181,17 +182,36 @@ def get_current_length(self, name: str) -> int: """Get the current stored length for the specified tensor.""" return self._current_lengths[name] + # Mapping from torch dtype to numpy dtype for fast list-to-pinned-memory writes + _TORCH_TO_NUMPY_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 store( self, name: str, - data: List[Number], + data: "Union[List[Number], torch.Tensor]", fill_value: Optional[Number] = None, ) -> int: """Store data into the host buffer. + Accepts either a Python list or a torch.Tensor. When a Tensor is provided, it is + copied directly into pinned memory (fast path, avoids torch.tensor() from list). + When a list is provided, numpy is used to write directly into pinned memory, + which is faster than torch.tensor(list) for large lists. + Args: name: Name of the tensor to store to. - data: List of values to store. + data: List of values or a 1-D torch.Tensor to store. fill_value: Optional value to fill the entire tensor with before storing. Returns: @@ -203,11 +223,18 @@ def store( if fill_value is not None: host_view.fill_(fill_value) - length = len(data) - 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) + if isinstance(data, torch.Tensor): + length = data.numel() + assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" + host_view[:length].copy_(data if data.dtype == dtype else data.to(dtype)) + else: + length = len(data) + assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" + np_dtype = self._TORCH_TO_NUMPY_DTYPE.get(dtype) + if np_dtype is not None: + host_view[:length].numpy()[:] = np.array(data, dtype=np_dtype) + else: + host_view[:length].copy_(torch.tensor(data, dtype=dtype)) self._current_lengths[name] = length return length @@ -661,8 +688,12 @@ 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)) + current = self._args_list[tensor_name] + if isinstance(current, list): + old_size = len(current) + current.extend([0] * (estimated_capacity - old_size)) + else: + self._args_list[tensor_name] = [0] * estimated_capacity @staticmethod def _get_page_assignments( @@ -824,7 +855,7 @@ def _flatten(nested_seqs: Sequence[Sequence[int]]) -> List[int]: def _store_arg( self, name: str, - tnsr_like: List[Number], + tnsr_like: "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. @@ -832,14 +863,22 @@ def _store_arg( 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(). + Accepts either a Python list or a 1-D torch.Tensor. Tensor inputs use a fast path + that avoids the expensive torch.tensor(list) conversion in InputBuffer.store(). + Args: name: Name of the argument to store. - tnsr_like: List of values to store. + tnsr_like: 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() + # Store for Python access. Tensor data is stored as-is to avoid expensive + # .tolist() conversion on the hot path. Properties that need lists (seq_len, + # input_pos, cache_loc, pages_per_seq) always receive list inputs. + if isinstance(tnsr_like, torch.Tensor): + self._args_list[name] = tnsr_like + else: + self._args_list[name] = tnsr_like.copy() # 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 @@ -981,16 +1020,22 @@ 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. + # Auto-compute page_seq_indices and page_in_seq from pages_per_seq using + # vectorized torch ops. This avoids the slow Python-list-to-tensor conversion + # that dominates host time for large page counts (e.g., 50k ISL models). # 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) + pages_per_seq_t = torch.tensor(pages_per_seq, dtype=torch.int) + seq_indices = torch.arange(len(pages_per_seq), dtype=torch.int) + page_seq_indices_t = torch.repeat_interleave(seq_indices, pages_per_seq_t) + cu_pages = torch.zeros(len(pages_per_seq) + 1, dtype=torch.int) + cu_pages[1:] = pages_per_seq_t.cumsum(0) + total_pages = cu_pages[-1].item() + page_in_seq_t = torch.arange(total_pages, dtype=torch.int) - torch.repeat_interleave( + cu_pages[:-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 cumulative number of pages if cu_num_pages is None: From 7953ba4c6090d0a88079d5332f4878f0a5d43520 Mon Sep 17 00:00:00 2001 From: Eran Geva Date: Tue, 17 Feb 2026 06:24:41 -0800 Subject: [PATCH 2/5] cleaned the host optimization to work only with tensors Signed-off-by: Eran Geva --- .../custom_ops/attention_interface.py | 112 +++++++++--------- 1 file changed, 53 insertions(+), 59 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index ae5326cfd831..a8a9ddb909c4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -41,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: ... @@ -182,37 +205,18 @@ def get_current_length(self, name: str) -> int: """Get the current stored length for the specified tensor.""" return self._current_lengths[name] - # Mapping from torch dtype to numpy dtype for fast list-to-pinned-memory writes - _TORCH_TO_NUMPY_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 store( self, name: str, - data: "Union[List[Number], torch.Tensor]", + data: torch.Tensor, fill_value: Optional[Number] = None, ) -> int: - """Store data into the host buffer. - - Accepts either a Python list or a torch.Tensor. When a Tensor is provided, it is - copied directly into pinned memory (fast path, avoids torch.tensor() from list). - When a list is provided, numpy is used to write directly into pinned memory, - which is faster than torch.tensor(list) for large lists. + """Store a tensor into the pinned host buffer. Args: name: Name of the tensor to store to. - data: List of values or a 1-D torch.Tensor 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. @@ -223,18 +227,12 @@ def store( if fill_value is not None: host_view.fill_(fill_value) - if isinstance(data, torch.Tensor): - length = data.numel() - assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" - host_view[:length].copy_(data if data.dtype == dtype else data.to(dtype)) - else: - length = len(data) - assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" - np_dtype = self._TORCH_TO_NUMPY_DTYPE.get(dtype) - if np_dtype is not None: - host_view[:length].numpy()[:] = np.array(data, dtype=np_dtype) - else: - host_view[:length].copy_(torch.tensor(data, dtype=dtype)) + length = data.numel() + assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" + # 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 @@ -508,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") @@ -604,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: @@ -688,12 +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) - current = self._args_list[tensor_name] - if isinstance(current, list): - old_size = len(current) - current.extend([0] * (estimated_capacity - old_size)) - else: - self._args_list[tensor_name] = [0] * estimated_capacity + self._args_list[tensor_name] = torch.zeros(estimated_capacity, dtype=torch.int) @staticmethod def _get_page_assignments( @@ -855,7 +850,7 @@ def _flatten(nested_seqs: Sequence[Sequence[int]]) -> List[int]: def _store_arg( self, name: str, - tnsr_like: "Union[List[Number], torch.Tensor]", + 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. @@ -863,22 +858,21 @@ def _store_arg( 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(). - Accepts either a Python list or a 1-D torch.Tensor. Tensor inputs use a fast path - that avoids the expensive torch.tensor(list) conversion in InputBuffer.store(). + 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 or a 1-D torch.Tensor 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}"): - # Store for Python access. Tensor data is stored as-is to avoid expensive - # .tolist() conversion on the hot path. Properties that need lists (seq_len, - # input_pos, cache_loc, pages_per_seq) always receive list inputs. - if isinstance(tnsr_like, torch.Tensor): - self._args_list[name] = tnsr_like - else: - self._args_list[name] = tnsr_like.copy() + # Convert to tensor at the boundary (numpy is ~2-3x faster than torch.tensor for large lists) + 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 @@ -887,7 +881,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]]] From e9edfcc8f9754a1ea52d2eb3b236cada694bdcf6 Mon Sep 17 00:00:00 2001 From: Eran Geva Date: Tue, 24 Feb 2026 06:11:28 -0800 Subject: [PATCH 3/5] fixed CR comments Signed-off-by: Eran Geva --- .../_torch/auto_deploy/custom_ops/attention_interface.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index a8a9ddb909c4..ae9cad23c9c6 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -1019,7 +1019,7 @@ def nest_sequences( # that dominates host time for large page counts (e.g., 50k ISL models). # page_seq_indices[j] = which sequence page j belongs to # page_in_seq[j] = which page within that sequence (0-indexed) - pages_per_seq_t = torch.tensor(pages_per_seq, dtype=torch.int) + pages_per_seq_t = self._args_list["pages_per_seq"] seq_indices = torch.arange(len(pages_per_seq), dtype=torch.int) page_seq_indices_t = torch.repeat_interleave(seq_indices, pages_per_seq_t) cu_pages = torch.zeros(len(pages_per_seq) + 1, dtype=torch.int) @@ -1031,6 +1031,9 @@ def nest_sequences( self._store_arg("page_seq_indices", page_seq_indices_t) self._store_arg("page_in_seq", page_in_seq_t) + if cu_num_pages is None: + cu_num_pages = cu_pages.tolist() + # update cumulative number of pages if cu_num_pages is None: pages_per_seq = self.pages_per_seq From cd8488816f8b0e67ec30e14155eaffa50d2534e1 Mon Sep 17 00:00:00 2001 From: Eran Geva <19514940+MrGeva@users.noreply.github.com> Date: Wed, 25 Feb 2026 00:52:55 -0800 Subject: [PATCH 4/5] fixed CR Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com> --- .../auto_deploy/custom_ops/attention_interface.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index ae9cad23c9c6..1b9c2dc39249 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -868,6 +868,7 @@ def _store_arg( """ with nvtx_range(f"ad_store_on_host_seq_info_arg_{name}"): # 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) @@ -1031,14 +1032,10 @@ def nest_sequences( self._store_arg("page_seq_indices", page_seq_indices_t) self._store_arg("page_in_seq", page_in_seq_t) - if cu_num_pages is None: - cu_num_pages = cu_pages.tolist() - - # 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) + 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) cu_num_pages = cu_num_pages.tolist() self._store_arg("cu_num_pages", cu_num_pages) From 433645bb131959d9e76960cd69352d82bb89c0d3 Mon Sep 17 00:00:00 2001 From: Eran Geva <19514940+MrGeva@users.noreply.github.com> Date: Wed, 25 Feb 2026 03:03:59 -0800 Subject: [PATCH 5/5] optimized cu_num_pages calculation Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com> --- .../custom_ops/attention_interface.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 1b9c2dc39249..bb2d3a3594d3 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -1015,30 +1015,28 @@ 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 using - # vectorized torch ops. This avoids the slow Python-list-to-tensor conversion - # that dominates host time for large page counts (e.g., 50k ISL models). + # 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) 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) - cu_pages = torch.zeros(len(pages_per_seq) + 1, dtype=torch.int) - cu_pages[1:] = pages_per_seq_t.cumsum(0) - total_pages = cu_pages[-1].item() + total_pages = cu_pages_t[-1].item() page_in_seq_t = torch.arange(total_pages, dtype=torch.int) - torch.repeat_interleave( - cu_pages[:-1], pages_per_seq_t + 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) - 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) - cu_num_pages = cu_num_pages.tolist() - self._store_arg("cu_num_pages", cu_num_pages) - # update sequence length with cache if seq_len_with_cache is None: seq_len_with_cache = [i_p + s_l for i_p, s_l in zip(self.input_pos, self.seq_len)]