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

Tiny adjust hybrid swa handling (#16292)

Browse Source
This commit is contained in:
Ke Bao
2026-01-02 23:11:23 +08:00
committed by GitHub
parent f26f6c2c99
commit c483a5f45f
2 changed files with 8 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/managers/schedule_batch.py
View File

@@ -63,7 +63,6 @@ from sglang.srt.mem_cache.allocator import (
SWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
from sglang.srt.mem_cache.common import (
alloc_for_decode,
alloc_for_extend,
@@ -73,7 +72,6 @@ from sglang.srt.mem_cache.common import (
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.metrics.collector import (
DPCooperationInfo,
SchedulerMetricsCollector,
@@ -1299,11 +1297,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
is_hybrid_swa = False
if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator):
assert (
tree_cache is None
or isinstance(tree_cache, SWARadixCache)
or isinstance(tree_cache, SWAChunkCache)
), "SWARadixCache or SWAChunkCache is required for SWATokenToKVPoolAllocator"
is_hybrid_swa = True
return cls(

19
python/sglang/srt/managers/schedule_policy.py
View File

@@ -30,6 +30,7 @@ from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.server_args import ServerArgs
if TYPE_CHECKING:
@@ -472,17 +473,13 @@ class PrefillAdder:
@contextmanager
def _lock_node(self, last_node: TreeNode):
if self.is_hybrid_swa:
try:
swa_uuid_for_lock = self.tree_cache.inc_lock_ref(last_node)
yield None
finally:
self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock)
else:
try:
self.tree_cache.inc_lock_ref(last_node)
yield None
finally:
try:
ret = self.tree_cache.inc_lock_ref(last_node)
yield None
finally:
if isinstance(self.tree_cache, SWARadixCache):
self.tree_cache.dec_lock_ref(last_node, ret)
else:
self.tree_cache.dec_lock_ref(last_node)
def add_one_req_ignore_eos(self, req: Req):