minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(swa): eliminate spurious translate_loc_from_full_to_swa warning in BCG and CG paths (#26152)

Browse Source 
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
Cheng Wan
2026-05-23 00:01:12 -07:00
committed by GitHub
co-authored by Cheng Wan
parent 75427c9ca4
commit 5964d30233
2 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/triton_backend.py
+11 -8
View File
@@ -329,7 +329,7 @@ class TritonAttnBackend(AttentionBackend):
seq_lens, seq_lens,
req_pool_indices, req_pool_indices,
bs, bs,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, token_to_kv_pool=self.token_to_kv_pool,
window_kv_indices=self.cuda_graph_window_kv_indices, window_kv_indices=self.cuda_graph_window_kv_indices,
) )
return kv_indptr, window_kv_indptr, window_kv_lens return kv_indptr, window_kv_indptr, window_kv_lens
@@ -374,7 +374,7 @@ class TritonAttnBackend(AttentionBackend):
seq_lens[:bs], seq_lens[:bs],
req_pool_indices, req_pool_indices,
bs, bs,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, token_to_kv_pool=self.token_to_kv_pool,
window_kv_indices=window_kv_indices, window_kv_indices=window_kv_indices,
) )
) )
@@ -461,7 +461,7 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
bs, bs,
self.device, self.device,
self.token_to_kv_pool_allocator, self.token_to_kv_pool,
) )
) )
window_num_kv_splits = torch.empty( window_num_kv_splits = torch.empty(
@@ -532,7 +532,7 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
bs, bs,
self.device, self.device,
self.token_to_kv_pool_allocator, self.token_to_kv_pool,
) )
custom_mask = spec_info.custom_mask custom_mask = spec_info.custom_mask
@@ -592,7 +592,7 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
bs, bs,
self.device, self.device,
self.token_to_kv_pool_allocator, self.token_to_kv_pool,
) )
qo_indptr = self.qo_indptr qo_indptr = self.qo_indptr
@@ -1445,7 +1445,7 @@ def update_sliding_window_buffer(
req_pool_indices, req_pool_indices,
bs, bs,
device=None, device=None,
token_to_kv_pool_allocator=None, token_to_kv_pool=None,
window_kv_indices=None, window_kv_indices=None,
): ):
"""Fill window KV buffers for sliding-window attention. """Fill window KV buffers for sliding-window attention.
@@ -1474,11 +1474,14 @@ def update_sliding_window_buffer(
window_kv_indices, window_kv_indices,
req_to_token.stride(0), req_to_token.stride(0),
) )
if hasattr(token_to_kv_pool_allocator, "translate_loc_from_full_to_swa"): if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"):
kv_last_index = window_kv_indptr[-1] kv_last_index = window_kv_indptr[-1]
# Flush before+after: window_kv_indices is a different tensor than out_cache_loc.
token_to_kv_pool.invalidate_loc_cache()
window_kv_indices[:kv_last_index] = ( window_kv_indices[:kv_last_index] = (
token_to_kv_pool_allocator.translate_loc_from_full_to_swa( token_to_kv_pool.translate_loc_from_full_to_swa(
window_kv_indices[:kv_last_index] window_kv_indices[:kv_last_index]
) )
) )
token_to_kv_pool.invalidate_loc_cache()
return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx
python/sglang/srt/model_executor/model_runner.py
+3 -3
View File
@@ -3266,6 +3266,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.hisparse_coordinator.wait_for_pending_backup() self.hisparse_coordinator.wait_for_pending_backup()
self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size) self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
if self.is_hybrid_swa:
self.token_to_kv_pool.invalidate_loc_cache()
# Replay cuda graph if applicable # Replay cuda graph if applicable
if can_run_graph: if can_run_graph:
ret = self.graph_runner.replay( ret = self.graph_runner.replay(
@@ -3292,9 +3295,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
server_args=self.server_args, server_args=self.server_args,
) )
if self.is_hybrid_swa:
self.token_to_kv_pool.invalidate_loc_cache()
# Hisparse coordinator — backends now read it from self.model_runner. # Hisparse coordinator — backends now read it from self.model_runner.
if self.hisparse_coordinator is not None: if self.hisparse_coordinator is not None:
self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size) self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)