From bc25abb7864bd0ac0e9446fe5656dde875abe4b2 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:20:00 -0700 Subject: [PATCH] perf(triton): avoid per-step D2H .item() sync in cuda-graph loc translate (#29921) Co-authored-by: lch1475369 --- .../srt/layers/attention/triton_backend.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 4e2ff1cb4..2bbed5bd4 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -636,15 +636,31 @@ class TritonAttnBackend(AttentionBackend): """ if self._translate_kv_loc is None: return None - # Full-attention read path. - n_kv = int(self.kv_indptr[bs].item()) + # seq_lens_sum is the reliable "mirror present" signal: it is + # None-preserving into the replay view, unlike seq_lens_cpu (always a + # non-None but stale slice for gpu_only batches). None -> fall back to a + # per-step D2H `.item()` on the indptr. + have_cpu_mirror = forward_batch.seq_lens_sum is not None + # Full-attention read path. kv_indptr[bs] == seq_lens_sum. + n_kv = ( + forward_batch.seq_lens_sum + if have_cpu_mirror + else int(self.kv_indptr[bs].item()) + ) if n_kv > 0: self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc( self.cuda_graph_kv_indices[:n_kv] ) - # SWA window read path (hybrid-SWA unified pools only). + # SWA window read path. window_kv_indptr[bs] == sum(min(seq_len, window)). if self.sliding_window_size is not None and self.sliding_window_size > 0: - n_win = int(self.window_kv_indptr[bs].item()) + if have_cpu_mirror: + n_win = int( + forward_batch.seq_lens_cpu[:bs] + .clamp(max=self.sliding_window_size) + .sum() + ) + else: + n_win = int(self.window_kv_indptr[bs].item()) if n_win > 0: self.cuda_graph_window_kv_indices[:n_win] = ( self.token_to_kv_pool.translate_loc_from_full_to_swa(