From 079d41a1c6a0395d3fb170e92792c5a89a5cc01e Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 16 Oct 2026 22:53:32 +0000 Subject: [PATCH] perf(attention): copy triton metadata without blocking the host --- python/freetoken/attention/triton.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 8bf9792d55..30bedf4ffa 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -230,17 +230,16 @@ def prepare_metadata(self, batch: Batch) -> None: cached_lens = [req.cached_len for req in reqs] num_query_tokens = sum(seqlens_q) is_decode = max(seqlens_q) == 1 - prefix_lens = torch.tensor(cached_lens, dtype=torch.int32, device=device) + CPU_KWARGS = {"device": "cpu", "dtype": torch.int32, "pin_memory": device.type == "cuda"} + prefix_lens = torch.tensor(cached_lens, **CPU_KWARGS).to(device, non_blocking=True) - indptr = torch.tensor([0] + seqlens_k, dtype=torch.int32, device=device).cumsum_(0) + indptr = torch.tensor([0] + seqlens_k, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True) if is_decode: cu_seqlens_q_gpu = torch.arange(0, padded_size + 1, device=device, dtype=torch.int32) elif all(l == 0 for l in cached_lens): cu_seqlens_q_gpu = indptr else: - cu_seqlens_q_gpu = torch.tensor( - [0] + seqlens_q, dtype=torch.int32, device=device - ).cumsum_(0) + cu_seqlens_q_gpu = torch.tensor([0] + seqlens_q, **CPU_KWARGS).cumsum_(0).to(device, non_blocking=True) indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) swa_indices = None if getattr(self.kvcache, "swa_paged", False):