Add trait for BasePrefixCache (#12436)
This commit is contained in:
@@ -1,12 +1,31 @@
|
|||||||
|
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):
|
|||||||
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
|
||||||
|
|||||||
Reference in New Issue
Block a user