Tiny adjust hybrid swa handling (#16292)

python/sglang/srt/managers/schedule_policy.py

@@ -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):