From 9dcb55be49c177ecb964d070818b6ff2ad4cfa2d Mon Sep 17 00:00:00 2001 From: William Zhang Date: Tue, 11 Aug 2026 13:01:02 -0700 Subject: [PATCH] [https://nvbugs/6248648][fix] Fix H2D copy * Why? Under confidential compute, we usually intentionally do not use pageable host memory for H2D transfers for performance reasons. Prior to this commit, the KV cache manager V2 would launch copy kernels that access unmapped pageable host memory. * What? This commit fixes that issue. Signed-off-by: William Zhang --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2c567824a8cb..616540003c1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -3628,6 +3628,31 @@ def copy_batch_block_offsets( ) return + # The CUDA copy kernel directly dereferences these host tensors. Under + # confidential compute, prefer_pinned() intentionally returns False, + # so launching that kernel would access unmapped pageable host memory. + host_inputs = ( + self.host_kv_cache_block_offsets, + copy_idx, + self.index_scales, + self.kv_offset, + ) + if not all(tensor.is_pinned() for tensor in host_inputs): + source = self.host_kv_cache_block_offsets[ + :, + copy_idx.to(dtype=torch.long), + 0, + :, + ] + valid = source != BAD_PAGE_INDEX + scales = self.index_scales.view(-1, 1, 1) + offsets = self.kv_offset.view(-1, 1, 1) + keys = torch.where(valid, source * scales, 0) + values = torch.where(valid, keys + offsets, 0) + staged_block_offsets = torch.stack((keys, values), dim=2) + dst_tensor[:, :num_seqs].copy_(staged_block_offsets, non_blocking=False) + return + copy_batch_block_offsets_to_device( self.host_kv_cache_block_offsets, dst_tensor,