perf(triton): avoid per-step D2H .item() sync in cuda-graph loc translate (#29921)
Co-authored-by: lch1475369 <lch1475369@gmail.com>
This commit is contained in:
@@ -636,15 +636,31 @@ class TritonAttnBackend(AttentionBackend):
|
|||||||
"""
|
"""
|
||||||
if self._translate_kv_loc is None:
|
if self._translate_kv_loc is None:
|
||||||
return None
|
return None
|
||||||
# Full-attention read path.
|
# seq_lens_sum is the reliable "mirror present" signal: it is
|
||||||
n_kv = int(self.kv_indptr[bs].item())
|
# 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:
|
if n_kv > 0:
|
||||||
self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc(
|
self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc(
|
||||||
self.cuda_graph_kv_indices[:n_kv]
|
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:
|
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:
|
if n_win > 0:
|
||||||
self.cuda_graph_window_kv_indices[:n_win] = (
|
self.cuda_graph_window_kv_indices[:n_win] = (
|
||||||
self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||||
|
|||||||
Reference in New Issue
Block a user