Unverified Commit 5cc0d25a authored by Liangsheng Yin's avatar Liangsheng Yin Committed by GitHub
Browse files

Add trait for `BasePrefixCache` (#12436)

parent a076ec1a
from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple from typing import (
TYPE_CHECKING,
Any,
NamedTuple,
Optional,
Protocol,
Tuple,
runtime_checkable,
)
import torch import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import Req
else:
Req = Any # Placeholder for Req type when not type checking
@runtime_checkable
class PrefixCacheTrait(Protocol):
req_to_token_pool: ReqToTokenPool
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
page_size: int
disable: bool
class MatchResult(NamedTuple): class MatchResult(NamedTuple):
...@@ -28,7 +47,7 @@ class MatchResult(NamedTuple): ...@@ -28,7 +47,7 @@ class MatchResult(NamedTuple):
host_hit_length: int = 0 host_hit_length: int = 0
class BasePrefixCache(ABC): class BasePrefixCache(ABC, PrefixCacheTrait):
"""Cache can be indexed by either rid or key.""" """Cache can be indexed by either rid or key."""
@abstractmethod @abstractmethod
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment