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

[mem_cache][2/N] refactor: move SWATokenToKVPoolAllocator to allocator/swa.py (#26676)

Browse Source
This commit is contained in:
shuwenn
2026-06-04 15:10:37 +08:00
committed by GitHub
parent a10bd785be
commit 1af53f67b2
18 changed files with 391 additions and 387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/kv_canary/api.py
+1 -1
View File
@@ -11,7 +11,7 @@ from sglang.srt.kv_canary.perturb.config import PerturbConfig
from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
if TYPE_CHECKING:
python/sglang/srt/kv_canary/runner/canary_manager.py
+1 -1
View File
@@ -31,9 +31,9 @@ from sglang.srt.kv_canary.state import CanaryDeviceState
from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
python/sglang/srt/kv_canary/runner/swa_divergence.py
+1 -1
View File
@@ -13,8 +13,8 @@ from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
python/sglang/srt/layers/attention/flashinfer_backend.py
+1 -1
View File
@@ -24,7 +24,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import (
python/sglang/srt/managers/schedule_policy.py
+1 -1
View File
@@ -38,6 +38,7 @@ from sglang.srt.dllm.config import DllmConfig
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
InitLoadBackParams,
@@ -49,7 +50,6 @@ from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.server_args import ServerArgs, get_global_server_args
if TYPE_CHECKING:
python/sglang/srt/mem_cache/allocator/swa.py
+369
View File
@@ -0,0 +1,369 @@
import torch
from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.utils import is_npu
from sglang.srt.utils.common import get_num_new_pages
_is_npu = is_npu()
if _is_npu:
from sglang.srt.hardware_backend.npu.allocator_npu import (
NPUPagedTokenToKVPoolAllocator,
)
class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
"""Allocator for SWA hybrid KV cache."""
def __init__(
self,
size: int,
size_swa: int,
page_size: int,
dtype: torch.dtype,
device: str,
kvcache: BaseSWAKVPool,
need_sort: bool,
):
assert isinstance(kvcache, BaseSWAKVPool)
self._size_full = size
self._size_swa = size_swa
self.dtype = dtype
self.device = device
self.page_size = page_size
full_kv_pool = getattr(kvcache, "full_kv_pool", None)
swa_kv_pool = getattr(kvcache, "swa_kv_pool", None)
if page_size == 1:
self.full_attn_allocator = TokenToKVPoolAllocator(
size,
dtype,
device,
full_kv_pool,
need_sort,
)
self.swa_attn_allocator = TokenToKVPoolAllocator(
size_swa,
dtype,
device,
swa_kv_pool,
need_sort,
)
else:
if _is_npu:
PagedTokenToKVPoolAllocatorClass = NPUPagedTokenToKVPoolAllocator
else:
PagedTokenToKVPoolAllocatorClass = PagedTokenToKVPoolAllocator
self.full_attn_allocator = PagedTokenToKVPoolAllocatorClass(
size,
page_size,
dtype,
device,
full_kv_pool,
need_sort,
)
self.swa_attn_allocator = PagedTokenToKVPoolAllocatorClass(
size_swa,
page_size,
dtype,
device,
swa_kv_pool,
need_sort,
)
# Note: append one more item of value -1 in the end so -1 maps to -1.
# It is needed for the last_loc in alloc_extend, where the first full_last_loc
# is -1, and we need to map it to swa_last_loc -1 as well.
self.full_to_swa_index_mapping = torch.cat(
[
torch.zeros(
size + self.page_size,
dtype=torch.int64,
device=device,
),
torch.tensor([-1], dtype=torch.int64, device=device),
]
)
self.need_sort = need_sort
self.free_pages = None
self.release_pages = None
self.is_not_in_free_group = True
self.free_group = []
self._kvcache = kvcache
self.clear()
self._kvcache.register_mapping(self.full_to_swa_index_mapping)
def available_size(self):
return min(
self.full_attn_allocator.available_size(),
self.swa_attn_allocator.available_size(),
)
def full_available_size(self):
return self.full_attn_allocator.available_size()
def swa_available_size(self):
return self.swa_attn_allocator.available_size()
@property
def size(self):
return min(self._size_full, self._size_swa)
@property
def size_swa(self):
return self._size_swa
@property
def size_full(self):
return self._size_full
def debug_print(self) -> str:
msg = ""
msg += f"#swa-available-size: {self.swa_attn_allocator.available_size()}, "
msg += (
f"#full-attn-available-size: {self.full_attn_allocator.available_size()}, "
)
return msg
def get_kvcache(self):
return self._kvcache
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
assert self._kvcache.full_to_swa_index_mapping is not None
return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
def alloc(self, need_size: int):
assert self.page_size == 1
if need_size > self.full_attn_allocator.available_size():
return None
if need_size > self.swa_attn_allocator.available_size():
return None
alloc_full_indices = self.full_attn_allocator.alloc(need_size)
alloc_swa_indices = self.swa_attn_allocator.alloc(need_size)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_extend(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
):
assert self.page_size > 1
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
if num_new_pages > self.full_attn_allocator.available_size() // self.page_size:
return None
if num_new_pages > self.swa_attn_allocator.available_size() // self.page_size:
return None
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
swa_last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_extend_swa_tail(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
swa_tail_len: int,
):
"""Allocate full KV for the whole extend and SWA KV only for the tail.
This is used by disaggregated decode preallocation: decode receives full
prompt KV for full-attention layers, but only the sliding-window state is
transferred for SWA layers.
"""
assert self.page_size > 1
assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1"
assert len(prefix_lens_cpu) == 1
assert 0 <= swa_tail_len <= extend_num_tokens
num_full_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size
if num_full_pages > self.full_attn_allocator.available_size() // self.page_size:
return None
if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size:
return None
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
)
assert alloc_full_indices is not None
if swa_tail_len == 0:
return alloc_full_indices
device = self.device
swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device)
swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64)
swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device)
swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64)
swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
swa_prefix_lens,
swa_prefix_lens_cpu,
swa_seq_lens,
swa_seq_lens_cpu,
swa_last_loc,
swa_tail_len,
)
assert alloc_swa_indices is not None
self.full_to_swa_index_mapping[alloc_full_indices[-swa_tail_len:]] = (
alloc_swa_indices
)
if swa_tail_len < extend_num_tokens:
self.full_to_swa_index_mapping[alloc_full_indices[:-swa_tail_len]] = 0
return alloc_full_indices
def alloc_decode(
self,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
):
assert self.page_size > 1
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, last_loc
)
alloc_swa_indices = self.swa_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, swa_last_loc
)
if alloc_full_indices is None or alloc_swa_indices is None:
return None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def free(self, free_index: torch.Tensor):
if free_index.numel() == 0:
return
# NOTE: the API is not idempotent.
if self.is_not_in_free_group:
self.full_attn_allocator.free(free_index)
self.free_swa(free_index)
else:
self.free_group.append(free_index)
assert (
self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
)
assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
def set_full_to_swa_mapping(
self, full_indices: torch.Tensor, swa_indices: torch.Tensor
) -> None:
"""Write full_to_swa_index_mapping[full_indices[i]] = swa_indices[i].
Used by HiCache load-back path to rebuild the mapping after FULL and SWA device alloc.
"""
if full_indices.numel() == 0:
return
assert full_indices.numel() == swa_indices.numel()
if _is_npu:
self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[full_indices] = swa_indices
def free_swa(self, free_index: torch.Tensor):
swa_indices = self.full_to_swa_index_mapping[free_index]
swa_indices = swa_indices[swa_indices > 0]
self.swa_attn_allocator.free(swa_indices)
self.full_to_swa_index_mapping[free_index] = 0
def backup_state(self):
return [
self.full_attn_allocator.backup_state(),
self.swa_attn_allocator.backup_state(),
]
def restore_state(self, state):
assert len(state) == 2
self.full_attn_allocator.restore_state(state[0])
self.swa_attn_allocator.restore_state(state[1])
def clear(self):
self.swa_attn_allocator.clear()
self.full_attn_allocator.clear()
# Note: the last item is -1, we don't clear it, see the comment in __init__
self.full_to_swa_index_mapping[:-1].fill_(0)
self.is_not_in_free_group = True
self.free_group = []
def get_cpu_copy(self, indices, mamba_indices=None):
return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
return self._kvcache.load_cpu_copy(
kv_cache_cpu, indices, mamba_indices=mamba_indices
)
python/sglang/srt/mem_cache/chunk_cache.py
+1 -1
View File
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional
import torch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
@@ -22,7 +23,6 @@ from sglang.srt.mem_cache.base_prefix_cache import (
from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
python/sglang/srt/mem_cache/common.py
+1 -1
View File
@@ -6,9 +6,9 @@ from typing import TYPE_CHECKING
import numpy as np
import torch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.triton_ops.common import (
_get_last_loc_safe_kernel as _get_last_loc_safe_kernel,
)
python/sglang/srt/mem_cache/swa_memory_pool.py
-368
View File
@@ -4,23 +4,9 @@ from typing import Dict, List, Optional, Tuple
import torch
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.allocator import (
BaseTokenToKVPoolAllocator,
PagedTokenToKVPoolAllocator,
TokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool
from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool
from sglang.srt.utils import is_npu
from sglang.srt.utils.common import get_num_new_pages
_is_npu = is_npu()
if _is_npu:
from sglang.srt.hardware_backend.npu.allocator_npu import (
NPUPagedTokenToKVPoolAllocator,
)
logger = logging.getLogger(__name__)
GB = 1024 * 1024 * 1024
@@ -280,357 +266,3 @@ class SWAKVPool(BaseSWAKVPool):
swa_kv_cpu = self._filter_swa_cpu_copy(swa_kv_cpu, row_mask)
self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices)
class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
"""Allocator for SWA hybrid KV cache."""
def __init__(
self,
size: int,
size_swa: int,
page_size: int,
dtype: torch.dtype,
device: str,
kvcache: BaseSWAKVPool,
need_sort: bool,
):
assert isinstance(kvcache, BaseSWAKVPool)
self._size_full = size
self._size_swa = size_swa
self.dtype = dtype
self.device = device
self.page_size = page_size
full_kv_pool = getattr(kvcache, "full_kv_pool", None)
swa_kv_pool = getattr(kvcache, "swa_kv_pool", None)
if page_size == 1:
self.full_attn_allocator = TokenToKVPoolAllocator(
size,
dtype,
device,
full_kv_pool,
need_sort,
)
self.swa_attn_allocator = TokenToKVPoolAllocator(
size_swa,
dtype,
device,
swa_kv_pool,
need_sort,
)
else:
if _is_npu:
PagedTokenToKVPoolAllocatorClass = NPUPagedTokenToKVPoolAllocator
else:
PagedTokenToKVPoolAllocatorClass = PagedTokenToKVPoolAllocator
self.full_attn_allocator = PagedTokenToKVPoolAllocatorClass(
size,
page_size,
dtype,
device,
full_kv_pool,
need_sort,
)
self.swa_attn_allocator = PagedTokenToKVPoolAllocatorClass(
size_swa,
page_size,
dtype,
device,
swa_kv_pool,
need_sort,
)
# Note: append one more item of value -1 in the end so -1 maps to -1.
# It is needed for the last_loc in alloc_extend, where the first full_last_loc
# is -1, and we need to map it to swa_last_loc -1 as well.
self.full_to_swa_index_mapping = torch.cat(
[
torch.zeros(
size + self.page_size,
dtype=torch.int64,
device=device,
),
torch.tensor([-1], dtype=torch.int64, device=device),
]
)
self.need_sort = need_sort
self.free_pages = None
self.release_pages = None
self.is_not_in_free_group = True
self.free_group = []
self._kvcache = kvcache
self.clear()
self._kvcache.register_mapping(self.full_to_swa_index_mapping)
def available_size(self):
return min(
self.full_attn_allocator.available_size(),
self.swa_attn_allocator.available_size(),
)
def full_available_size(self):
return self.full_attn_allocator.available_size()
def swa_available_size(self):
return self.swa_attn_allocator.available_size()
@property
def size(self):
return min(self._size_full, self._size_swa)
@property
def size_swa(self):
return self._size_swa
@property
def size_full(self):
return self._size_full
def debug_print(self) -> str:
msg = ""
msg += f"#swa-available-size: {self.swa_attn_allocator.available_size()}, "
msg += (
f"#full-attn-available-size: {self.full_attn_allocator.available_size()}, "
)
return msg
def get_kvcache(self):
return self._kvcache
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
assert self._kvcache.full_to_swa_index_mapping is not None
return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
def alloc(self, need_size: int):
assert self.page_size == 1
if need_size > self.full_attn_allocator.available_size():
return None
if need_size > self.swa_attn_allocator.available_size():
return None
alloc_full_indices = self.full_attn_allocator.alloc(need_size)
alloc_swa_indices = self.swa_attn_allocator.alloc(need_size)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_extend(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
):
assert self.page_size > 1
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
if num_new_pages > self.full_attn_allocator.available_size() // self.page_size:
return None
if num_new_pages > self.swa_attn_allocator.available_size() // self.page_size:
return None
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
swa_last_loc,
extend_num_tokens,
num_new_pages=num_new_pages,
)
assert alloc_full_indices is not None
assert alloc_swa_indices is not None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def alloc_extend_swa_tail(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
swa_tail_len: int,
):
"""Allocate full KV for the whole extend and SWA KV only for the tail.
This is used by disaggregated decode preallocation: decode receives full
prompt KV for full-attention layers, but only the sliding-window state is
transferred for SWA layers.
"""
assert self.page_size > 1
assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1"
assert len(prefix_lens_cpu) == 1
assert 0 <= swa_tail_len <= extend_num_tokens
num_full_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size
if num_full_pages > self.full_attn_allocator.available_size() // self.page_size:
return None
if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size:
return None
alloc_full_indices = self.full_attn_allocator.alloc_extend(
prefix_lens,
prefix_lens_cpu,
seq_lens,
seq_lens_cpu,
last_loc,
extend_num_tokens,
)
assert alloc_full_indices is not None
if swa_tail_len == 0:
return alloc_full_indices
device = self.device
swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device)
swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64)
swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device)
swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64)
swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device)
alloc_swa_indices = self.swa_attn_allocator.alloc_extend(
swa_prefix_lens,
swa_prefix_lens_cpu,
swa_seq_lens,
swa_seq_lens_cpu,
swa_last_loc,
swa_tail_len,
)
assert alloc_swa_indices is not None
self.full_to_swa_index_mapping[alloc_full_indices[-swa_tail_len:]] = (
alloc_swa_indices
)
if swa_tail_len < extend_num_tokens:
self.full_to_swa_index_mapping[alloc_full_indices[:-swa_tail_len]] = 0
return alloc_full_indices
def alloc_decode(
self,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
):
assert self.page_size > 1
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
alloc_full_indices = self.full_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, last_loc
)
alloc_swa_indices = self.swa_attn_allocator.alloc_decode(
seq_lens, seq_lens_cpu, swa_last_loc
)
if alloc_full_indices is None or alloc_swa_indices is None:
return None
if _is_npu:
self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
alloc_swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
return alloc_full_indices
def free(self, free_index: torch.Tensor):
if free_index.numel() == 0:
return
# NOTE: the API is not idempotent.
if self.is_not_in_free_group:
self.full_attn_allocator.free(free_index)
self.free_swa(free_index)
else:
self.free_group.append(free_index)
assert (
self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
)
assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
def set_full_to_swa_mapping(
self, full_indices: torch.Tensor, swa_indices: torch.Tensor
) -> None:
"""Write full_to_swa_index_mapping[full_indices[i]] = swa_indices[i].
Used by HiCache load-back path to rebuild the mapping after FULL and SWA device alloc.
"""
if full_indices.numel() == 0:
return
assert full_indices.numel() == swa_indices.numel()
if _is_npu:
self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[full_indices] = swa_indices
def free_swa(self, free_index: torch.Tensor):
swa_indices = self.full_to_swa_index_mapping[free_index]
swa_indices = swa_indices[swa_indices > 0]
self.swa_attn_allocator.free(swa_indices)
self.full_to_swa_index_mapping[free_index] = 0
def backup_state(self):
return [
self.full_attn_allocator.backup_state(),
self.swa_attn_allocator.backup_state(),
]
def restore_state(self, state):
assert len(state) == 2
self.full_attn_allocator.restore_state(state[0])
self.swa_attn_allocator.restore_state(state[1])
def clear(self):
self.swa_attn_allocator.clear()
self.full_attn_allocator.clear()
# Note: the last item is -1, we don't clear it, see the comment in __init__
self.full_to_swa_index_mapping[:-1].fill_(0)
self.is_not_in_free_group = True
self.free_group = []
def get_cpu_copy(self, indices, mamba_indices=None):
return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
return self._kvcache.load_cpu_copy(
kv_cache_cpu, indices, mamba_indices=mamba_indices
)
python/sglang/srt/mem_cache/swa_radix_cache.py
+1 -1
View File
@@ -28,6 +28,7 @@ import torch
from numpy import float64
from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
@@ -43,7 +44,6 @@ from sglang.srt.mem_cache.base_prefix_cache import (
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.events import KVCacheEventMixin
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.utils import split_node_hash_value
if TYPE_CHECKING:
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+1 -1
View File
@@ -49,7 +49,7 @@ class SWAComponent(TreeComponent):
"""
def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams):
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
assert isinstance(
cache.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+2 -1
View File
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.allocator import (
PagedTokenToKVPoolAllocator,
TokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
@@ -34,7 +35,7 @@ from sglang.srt.mem_cache.memory_pool import (
NoOpMHATokenToKVPool,
ReqToTokenPool,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.utils.common import (
get_available_gpu_memory,
is_float4_e2m1fn_x2,
test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py
+1 -1
View File
@@ -11,7 +11,7 @@ from unittest.mock import MagicMock
import torch
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+2 -1
View File
@@ -19,9 +19,10 @@ from types import SimpleNamespace
import torch
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py
+2 -1
View File
@@ -16,6 +16,7 @@ from array import array
import torch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
@@ -25,7 +26,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_cuda_ci
test/registered/unit/mem_cache/test_swa_unittest.py
+2 -1
View File
@@ -5,6 +5,7 @@ import torch
from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
@@ -16,7 +17,7 @@ from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+2 -4
View File
@@ -178,10 +178,8 @@ def create_bench_cache(
# --- KV pool + allocator ---
if has_swa:
from sglang.srt.mem_cache.swa_memory_pool import (
SWAKVPool,
SWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
pool = SWAKVPool(
size=kv_size,
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+2 -1
View File
@@ -22,6 +22,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
@@ -41,7 +42,7 @@ from sglang.srt.mem_cache.memory_pool import (
ReqToTokenPool,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
CacheTransferPhase,
ComponentType,