EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix: Complete int32 to int64 conversion (#4465)

Browse Source
This commit is contained in:
Zhiqiang Xie
2025-03-16 18:14:27 -07:00
committed by GitHub
parent 5cbd709ea1
commit f5bbf6037d
3 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -305,7 +305,7 @@ class HiRadixCache(RadixCache):
if value: if value:
value = torch.cat(value) value = torch.cat(value)
else: else:
value = torch.tensor([], dtype=torch.int32) value = torch.tensor([], dtype=torch.int64)
last_node_global = last_node last_node_global = last_node
while last_node.evicted: while last_node.evicted:

5
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -622,11 +622,10 @@ class HostKVCache(abc.ABC):
self.mem_state = torch.zeros( self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device (self.size,), dtype=torch.uint8, device=self.device
) )
self.free_slots = torch.arange(self.size, dtype=torch.int32)
self.can_use_mem_size = self.size
# A lock for synchronized operations on memory allocation and state transitions. # A lock for synchronized operations on memory allocation and state transitions.
self.lock = threading.RLock() self.lock = threading.RLock()
self.clear()
@abc.abstractmethod @abc.abstractmethod
def get_size_per_token(self): def get_size_per_token(self):
@@ -656,7 +655,7 @@ class HostKVCache(abc.ABC):
def clear(self): def clear(self):
self.mem_state.fill_(0) self.mem_state.fill_(0)
self.can_use_mem_size = self.size self.can_use_mem_size = self.size
self.free_slots = torch.arange(self.size, dtype=torch.int32) self.free_slots = torch.arange(self.size, dtype=torch.int64)
@synchronized @synchronized
def get_state(self, indices: torch.Tensor) -> MemoryStateInt: def get_state(self, indices: torch.Tensor) -> MemoryStateInt:

4
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -140,7 +140,7 @@ class RadixCache(BasePrefixCache):
return ( return (
torch.empty( torch.empty(
(0,), (0,),
dtype=torch.int32, dtype=torch.int64,
device=self.device, device=self.device,
), ),
self.root_node, self.root_node,
@@ -154,7 +154,7 @@ class RadixCache(BasePrefixCache):
if value: if value:
value = torch.cat(value) value = torch.cat(value)
else: else:
value = torch.empty((0,), dtype=torch.int32, device=self.device) value = torch.empty((0,), dtype=torch.int64, device=self.device)
return value, last_node return value, last_node
def insert(self, key: List, value=None): def insert(self, key: List, value=None):