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

Reduce CP shared-KV request-boundary stalls

Browse Source 
CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.

Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.

Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.

Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.

Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.

Tested: Remote container py_compile for modified runtime files; git diff --cached --check.

Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.
This commit is contained in:
laoyao0822
2026-05-05 05:18:35 +08:00
parent a638d71d53
commit 49eaf9ffde
11 changed files with 551 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
python/sglang/srt/disaggregation/decode.py
View File

@@ -43,7 +43,6 @@ from sglang.srt.disaggregation.utils import (
TransferBackend,
get_kv_class,
is_mla_backend,
kv_to_page_indices,
poll_and_all_reduce,
prepare_abort,
)
@@ -78,6 +77,29 @@ if TYPE_CHECKING:
CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get()
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
*,
num_tokens: Optional[int] = None,
):
"""Return compact int32 page indices for a token-location tensor.
PD bootstrap only needs page-level destination indices. Avoid materializing
a full token-level CPU copy before converting to pages.
"""
if num_tokens is not None:
kv_locs = kv_locs[:num_tokens]
if page_size == 1:
page_locs = kv_locs
else:
page_locs = kv_locs[::page_size] // page_size
return page_locs.to(dtype=torch.int32).cpu().numpy()
def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
return req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
req.bootstrap_host is None
@@ -409,7 +431,11 @@ class DecodePreallocQueue:
)
self.queue.append(
DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
DecodeRequest(
req=req,
kv_receiver=kv_receiver,
waiting_for_input=False,
)
)
def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool:
@@ -675,16 +701,17 @@ class DecodePreallocQueue:
continue
allocatable_tokens -= required_tokens_for_request
self._pre_alloc(decode_req.req)
kv_indices = (
self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][
: len(decode_req.req.origin_input_ids)
]
.cpu()
.numpy()
kv_loc = self._pre_alloc(
decode_req.req,
write_req_to_token=False,
)
page_size = self.token_to_kv_pool_allocator.page_size
page_indices = _kv_locs_to_page_indices_cpu(
kv_loc,
page_size,
num_tokens=origin_input_len,
)
# Prepare extra pool indices for hybrid models
if isinstance(self.token_to_kv_pool, HybridLinearKVPool):
@@ -703,9 +730,7 @@ class DecodePreallocQueue:
window_start = max(0, seq_len - window_size)
window_start = (window_start // page_size) * page_size
window_kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, window_start:seq_len
]
window_kv_indices_full = kv_loc[window_start:seq_len]
# Translate to SWA pool indices
window_kv_indices_swa = (
@@ -713,23 +738,23 @@ class DecodePreallocQueue:
window_kv_indices_full
)
)
state_indices = window_kv_indices_swa.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
window_kv_indices_swa,
page_size,
)
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
seq_len = len(decode_req.req.origin_input_ids)
kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, :seq_len
]
state_indices = kv_indices_full.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = page_indices
else:
state_indices = None
self.req_to_token_pool.write(
(decode_req.req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
decode_req.metadata_buffer_index = (
self.req_to_metadata_buffer_idx_allocator.alloc()
)
assert decode_req.metadata_buffer_index is not None
page_indices = kv_to_page_indices(kv_indices, page_size)
decode_req.kv_receiver.init(
page_indices, decode_req.metadata_buffer_index, state_indices
)
@@ -840,7 +865,12 @@ class DecodePreallocQueue:
last_loc,
)
def _pre_alloc(self, req: Req) -> torch.Tensor:
def _pre_alloc(
self,
req: Req,
*,
write_req_to_token: bool = True,
) -> torch.Tensor:
"""Pre-allocate the memory for req_to_token and token_kv_pool"""
req_pool_indices = self.req_to_token_pool.alloc([req])
@@ -875,7 +905,10 @@ class DecodePreallocQueue:
"KV cache is full! There is a bug in memory estimation."
)
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
if write_req_to_token:
self.req_to_token_pool.write(
(req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
# populate metadata
req.fill_ids = req.origin_input_ids + req.output_ids

38
python/sglang/srt/disaggregation/prefill.py
View File

@@ -37,7 +37,6 @@ from sglang.srt.disaggregation.utils import (
TransferBackend,
get_kv_class,
is_mla_backend,
kv_to_page_indices,
kv_to_page_num,
poll_and_all_reduce_attn_cp_tp_group,
prepare_abort,
@@ -62,6 +61,20 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _kv_locs_to_page_indices_cpu(
kv_locs: torch.Tensor,
page_size: int,
):
"""Return int32 page indices without materializing token-level CPU indices."""
if page_size == 1:
page_locs = kv_locs
else:
page_locs = kv_locs[::page_size] // page_size
return page_locs.to(dtype=torch.int32).cpu().numpy()
def release_req_to_metadata_buffer(
req: Req, allocator: ReqToMetadataIdxAllocator
) -> None:
@@ -723,10 +736,9 @@ class SchedulerDisaggregationPrefillMixin:
# if not the last chunk and the last page is partial, delay the last partial page to the next send
end_idx = end_idx - end_idx % page_size
kv_indices = (
self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
.cpu()
.numpy()
page_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx],
page_size,
)
req.start_send_idx = end_idx
state_indices = None
@@ -762,19 +774,19 @@ class SchedulerDisaggregationPrefillMixin:
window_kv_indices_full
)
)
state_indices = window_kv_indices_swa.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
window_kv_indices_swa,
page_size,
)
elif isinstance(
self.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
):
seq_len = len(req.fill_ids)
kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len
]
state_indices = kv_indices_full.cpu().numpy()
state_indices = kv_to_page_indices(state_indices, page_size)
state_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, :seq_len],
page_size,
)
page_indices = kv_to_page_indices(kv_indices, page_size)
if len(page_indices) == 0:
logger.info(
f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"

2
python/sglang/srt/environ.py
View File

@@ -203,6 +203,8 @@ class Envs:
SGLANG_FORCE_SHUTDOWN = EnvBool(False)
SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)
SGLANG_DEBUG_SORT_NVTX = EnvBool(False)
SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)
SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)

82
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -15,6 +15,7 @@ logger = logging.getLogger(__name__)
_DEBUG_LOG_COUNTS: dict[str, int] = {}
_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_MLA_PREFETCH_LOG_PROBE_LAYER = 2
_SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
def cp_shared_kv_debug_enabled() -> bool:
@@ -25,6 +26,10 @@ def cp_shared_kv_current_reuse_enabled() -> bool:
return envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.get()
def cp_shared_kv_sort_nvtx_enabled() -> bool:
return _SORT_NVTX_ENABLED
def cp_shared_kv_tai_materialize_enabled() -> bool:
return envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get()
@@ -725,6 +730,9 @@ def remap_logical_locs_to_dense_locs(
def build_current_loc_remap(
query_locs: torch.Tensor,
current_locs: torch.Tensor,
*,
page_size: int | None = None,
logical_page_capacity: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Map logical locs into rows of the gathered current chunk tensor.
@@ -732,6 +740,11 @@ def build_current_loc_remap(
`current_locs` is `forward_batch.out_cache_loc`; its row order is the row
order of the already CP-all-gathered current KV/index tensor.
When `page_size` and `logical_page_capacity` are provided, this uses a
page-inverse fast path for the single-request current-only case. That avoids
the PyTorch `torch.sort + searchsorted` path, which dispatches ATen/CCCL
radix-sort kernels in the attention hot path.
Returns:
- is_current_mask: true where query_locs is a non-negative current loc.
- compact_row_ids: row id into the current compact tensor where valid,
@@ -745,7 +758,74 @@ def build_current_loc_remap(
query_flat_long = query_locs.reshape(-1).to(torch.long)
current_flat_long = current_locs.reshape(-1).to(torch.long)
sorted_current_locs, sorted_to_current_rows = torch.sort(current_flat_long)
if page_size is not None and logical_page_capacity is not None:
if cp_shared_kv_sort_nvtx_enabled():
torch.cuda.nvtx.range_push(
f"CP_SHARED_KV:current_loc_remap:page_inverse num_current={current_flat_long.numel()} num_query={query_flat_long.numel()} page_size={page_size}"
)
try:
current_pages = torch.div(
current_flat_long[::page_size], page_size, rounding_mode="floor"
)
page_inverse = torch.full(
(logical_page_capacity,),
-1,
device=current_flat_long.device,
dtype=torch.long,
)
row_page_ids = torch.arange(
current_pages.numel(), device=current_flat_long.device, dtype=torch.long
)
valid_current_pages = (current_pages >= 0) & (
current_pages < logical_page_capacity
)
safe_current_pages = torch.where(
valid_current_pages, current_pages, torch.zeros_like(current_pages)
)
safe_row_page_ids = torch.where(
valid_current_pages, row_page_ids, torch.zeros_like(row_page_ids)
)
page_inverse.scatter_(0, safe_current_pages, safe_row_page_ids)
valid_query = query_flat_long >= 0
safe_query_locs = torch.where(
valid_query, query_flat_long, torch.zeros_like(query_flat_long)
)
query_pages = torch.div(
safe_query_locs, page_size, rounding_mode="floor"
)
query_offsets = torch.remainder(safe_query_locs, page_size)
query_pages_in_range = query_pages < logical_page_capacity
safe_query_pages = torch.clamp(query_pages, max=logical_page_capacity - 1)
row_pages = page_inverse[safe_query_pages]
row_values = row_pages * page_size + query_offsets
matched = (
valid_query
& query_pages_in_range
& (row_pages >= 0)
& (row_values < current_flat_long.numel())
)
compact_flat = torch.where(
matched,
row_values.to(compact_row_ids.dtype),
torch.full_like(compact_row_ids.reshape(-1), -1),
)
return matched.reshape(query_locs.shape), compact_flat.reshape(query_locs.shape)
finally:
if cp_shared_kv_sort_nvtx_enabled():
torch.cuda.nvtx.range_pop()
if cp_shared_kv_sort_nvtx_enabled():
torch.cuda.nvtx.range_push(
f"CP_SHARED_KV:current_loc_remap:torch.sort_fallback num_current={current_flat_long.numel()} num_query={query_flat_long.numel()}"
)
try:
sorted_current_locs, sorted_to_current_rows = torch.sort(current_flat_long)
finally:
torch.cuda.nvtx.range_pop()
else:
sorted_current_locs, sorted_to_current_rows = torch.sort(current_flat_long)
insert_positions = torch.searchsorted(sorted_current_locs, query_flat_long)
safe_positions = torch.clamp(insert_positions, max=current_flat_long.numel() - 1)

16
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -1667,9 +1667,25 @@ class NativeSparseAttnBackend(
)
if can_reuse_current_kv:
logical_page_table_1 = page_table_1
current_remap_page_size = None
current_remap_logical_page_capacity = None
if len(getattr(forward_batch, "extend_seq_lens_cpu", []) or []) == 1:
current_remap_page_size = forward_batch.token_to_kv_pool.page_size
current_remap_logical_page_capacity = (
max(
forward_batch.token_to_kv_pool.size
// current_remap_page_size
- 1,
0,
)
* forward_batch.cp_shared_kv_layout.cp_size
+ 1
)
current_mask, page_table_1 = build_current_loc_remap(
logical_page_table_1,
forward_batch.out_cache_loc,
page_size=current_remap_page_size,
logical_page_capacity=current_remap_logical_page_capacity,
)
if cp_shared_kv_debug_enabled():
missing_current = (logical_page_table_1 >= 0) & (~current_mask)

38
python/sglang/srt/layers/moe/ep_moe/kernels.py
View File

@@ -3,10 +3,13 @@ import logging
import torch
import triton
from sglang.srt.environ import envs
from sglang.srt.utils import ceil_div, is_cuda
logger = logging.getLogger(__name__)
_MOE_SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_MOE_SORT_NVTX.get()
_is_cuda = is_cuda()
if _is_cuda:
from sglang.srt.layers.quantization.fp8_kernel import (
@@ -16,6 +19,23 @@ if _is_cuda:
import triton.language as tl
def _debug_sort_nvtx_enabled() -> bool:
return _MOE_SORT_NVTX_ENABLED
def _sort_with_optional_nvtx(
tensor: torch.Tensor, *, stable: bool, nvtx_name: str
):
if not _debug_sort_nvtx_enabled():
return torch.sort(tensor, stable=stable)
torch.cuda.nvtx.range_push(f"{nvtx_name} numel={tensor.numel()} dtype={tensor.dtype}")
try:
return torch.sort(tensor, stable=stable)
finally:
torch.cuda.nvtx.range_pop()
def _get_launch_config_1d(device, numel):
MAX_THREADS_PER_BLOCK = 1024
MIN_THREADS_PER_BLOCK = 512
@@ -158,7 +178,11 @@ def deepep_compute_src2dst_triton_kernel(
def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int):
reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
reorder_topk_ids, reorder_ids = _sort_with_optional_nvtx(
topk_ids.view(-1),
stable=True,
nvtx_name="MOE_EP:deepep_preprocess:topk_ids:torch.sort",
)
seg_indptr = torch.empty(num_experts + 1, device=topk_ids.device, dtype=torch.int64)
src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int64)
@@ -197,7 +221,11 @@ def compute_seg_indptr_triton_kernel(reorder_topk_ids, seg_indptr, num_toks):
def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor):
_, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
_, reorder_ids = _sort_with_optional_nvtx(
topk_ids.view(-1),
stable=True,
nvtx_name="MOE_EP:cutlass_w4_preprocess:topk_ids:torch.sort",
)
BLOCK_SIZE = 512
grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),)
@@ -1046,7 +1074,11 @@ def moe_ep_deepgemm_preprocess(
block_shape,
output_dtype: torch.dtype = torch.float8_e4m3fn,
):
reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
reorder_topk_ids, reorder_ids = _sort_with_optional_nvtx(
topk_ids.view(-1),
stable=True,
nvtx_name="MOE_EP:deepgemm_preprocess:topk_ids:torch.sort",
)
seg_indptr = torch.zeros(
num_local_experts + 1, device=topk_ids.device, dtype=torch.int64
)

103
python/sglang/srt/mem_cache/allocator.py
View File

@@ -24,14 +24,22 @@ from typing import TYPE_CHECKING, List, Optional
import torch
import triton
from sglang.srt.environ import envs
import triton.language as tl
from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
_SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import KVCache
def _debug_sort_nvtx_enabled() -> bool:
return _SORT_NVTX_ENABLED
class BaseTokenToKVPoolAllocator(abc.ABC):
@abc.abstractmethod
def __init__(
@@ -81,8 +89,19 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
def merge_and_sort_free(self):
if len(self.release_pages) > 0:
num_free_pages = len(self.free_pages)
num_release_pages = len(self.release_pages)
self.free_pages = torch.cat((self.free_pages, self.release_pages))
self.free_pages, _ = torch.sort(self.free_pages)
if _debug_sort_nvtx_enabled():
torch.cuda.nvtx.range_push(
f"KV_ALLOCATOR:merge_and_sort_free:torch.sort free_pages={num_free_pages} release_pages={num_release_pages}"
)
try:
self.free_pages, _ = torch.sort(self.free_pages)
finally:
torch.cuda.nvtx.range_pop()
else:
self.free_pages, _ = torch.sort(self.free_pages)
self.release_pages = torch.empty(
(0,), dtype=self.release_pages.dtype, device=self.device
)
@@ -588,30 +607,69 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
def _select_compute_owner_pages(
self,
page_compute_owners: List[int],
) -> Optional[torch.Tensor]:
selected_pages = []
lane_offsets = [0 for _ in range(self.cp_size)]
lane_pages = [
self.free_pages[
torch.remainder(self.free_pages - 1, self.cp_size) == owner
]
for owner in range(self.cp_size)
]
) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
if not page_compute_owners:
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
torch.zeros_like(self.free_pages, dtype=torch.bool),
torch.zeros_like(self.release_pages, dtype=torch.bool),
)
required_by_owner = [0 for _ in range(self.cp_size)]
for owner in page_compute_owners:
if owner < 0 or owner >= self.cp_size:
raise ValueError(
f"compute owner must be in [0, {self.cp_size}), got {owner}"
)
required_by_owner[owner] += 1
lane_pages = [None for _ in range(self.cp_size)]
selected_free_mask = torch.zeros_like(self.free_pages, dtype=torch.bool)
selected_release_mask = torch.zeros_like(self.release_pages, dtype=torch.bool)
for owner, required_count in enumerate(required_by_owner):
if required_count == 0:
continue
owner_mask = torch.remainder(self.free_pages - 1, self.cp_size) == owner
selected_owner_free_mask = owner_mask & (
torch.cumsum(owner_mask.to(torch.int64), dim=0) <= required_count
)
selected_owner_pages = self.free_pages[selected_owner_free_mask]
remaining_count = required_count - selected_owner_pages.numel()
if remaining_count > 0:
release_owner_mask = (
torch.remainder(self.release_pages - 1, self.cp_size) == owner
)
selected_owner_release_mask = release_owner_mask & (
torch.cumsum(release_owner_mask.to(torch.int64), dim=0)
<= remaining_count
)
selected_owner_release_pages = self.release_pages[
selected_owner_release_mask
]
if remaining_count > selected_owner_release_pages.numel():
return None
selected_owner_pages = torch.cat(
(selected_owner_pages, selected_owner_release_pages)
)
selected_release_mask |= selected_owner_release_mask
lane_pages[owner] = selected_owner_pages
selected_free_mask |= selected_owner_free_mask
selected_pages = []
lane_offsets = [0 for _ in range(self.cp_size)]
for owner in page_compute_owners:
lane_offset = lane_offsets[owner]
if lane_offset >= lane_pages[owner].numel():
return None
selected_pages.append(lane_pages[owner][lane_offset])
lane_offsets[owner] = lane_offset + 1
if not selected_pages:
return torch.empty((0,), dtype=torch.int64, device=self.device)
return torch.stack(selected_pages).to(torch.int64)
return (
torch.stack(selected_pages).to(torch.int64),
selected_free_mask,
selected_release_mask,
)
def alloc_extend_compute_owner(
self,
@@ -645,15 +703,10 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
f"{num_new_pages=} page_compute_owners={len(page_compute_owners)}"
)
if self.need_sort and num_new_pages > len(self.free_pages):
self.merge_and_sort_free()
selected_pages = self._select_compute_owner_pages(page_compute_owners)
if selected_pages is None and self.need_sort and len(self.release_pages) > 0:
self.merge_and_sort_free()
selected_pages = self._select_compute_owner_pages(page_compute_owners)
if selected_pages is None:
selected = self._select_compute_owner_pages(page_compute_owners)
if selected is None:
return None
selected_pages, selected_free_mask, selected_release_mask = selected
out_indices = torch.empty(
(extend_num_tokens,), dtype=torch.int64, device=self.device
@@ -668,8 +721,8 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
self.device,
)
selected_mask = torch.isin(self.free_pages, selected_pages)
self.free_pages = self.free_pages[~selected_mask]
self.free_pages = self.free_pages[~selected_free_mask]
self.release_pages = self.release_pages[~selected_release_mask]
if self.debug_mode:
assert len(torch.unique(out_indices)) == len(out_indices)

24
python/sglang/srt/mem_cache/common.py
View File

@@ -364,13 +364,6 @@ def alloc_paged_token_slots_extend(
"multi_batch" if len(prefix_lens_cpu) != 1 else "unknown"
)
if page_compute_owners is not None:
_evict_for_compute_owner_lanes(
tree_cache=tree_cache,
allocator=allocator,
page_compute_owners=page_compute_owners,
)
state = None
if backup_state:
state = allocator.backup_state()
@@ -385,6 +378,23 @@ def alloc_paged_token_slots_extend(
extend_num_tokens,
page_compute_owners,
)
if out_cache_loc is None:
_evict_for_compute_owner_lanes(
tree_cache=tree_cache,
allocator=allocator,
page_compute_owners=page_compute_owners,
)
if backup_state:
state = allocator.backup_state()
out_cache_loc = alloc_extend_compute_owner(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
page_compute_owners,
)
if out_cache_loc is None:
required = available = deficits = None
compute_owner_lane_stats = getattr(