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

misc: correct the int data type for token ids and indices (#969)

Browse Source
This commit is contained in:
Zhiqiang Xie
2024-08-07 13:40:07 -07:00
committed by GitHub
parent 4d929107ae
commit 6db27f7b3b
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -780,7 +780,7 @@ def top_k_top_p_sampling_from_probs_torch(
sampled_index = torch.multinomial(probs_sort, num_samples=1) sampled_index = torch.multinomial(probs_sort, num_samples=1)
except RuntimeError: except RuntimeError:
batch_next_token_ids = torch.zeros( batch_next_token_ids = torch.zeros(
(probs_sort.shape[0],), dtype=torch.int64, device=probs.device (probs_sort.shape[0],), dtype=torch.int32, device=probs.device
) )
success = torch.zeros(probs.shape[0], dtype=torch.bool, device=probs.device) success = torch.zeros(probs.shape[0], dtype=torch.bool, device=probs.device)
return batch_next_token_ids, success return batch_next_token_ids, success

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

@@ -74,7 +74,7 @@ class RadixCache(BasePrefixCache):
if value: if value:
value = torch.concat(value) value = torch.concat(value)
else: else:
value = torch.tensor([], dtype=torch.int64) value = torch.tensor([], dtype=torch.int32)
return value, last_node[0] return value, last_node[0]
def insert(self, key, value=None): def insert(self, key, value=None):
@@ -102,7 +102,7 @@ class RadixCache(BasePrefixCache):
if del_in_memory_pool: if del_in_memory_pool:
self.token_to_kv_pool.free(indices) self.token_to_kv_pool.free(indices)
else: else:
return torch.tensor([], dtype=torch.int64), self.root_node return torch.tensor([], dtype=torch.int32), self.root_node
# Radix Cache takes one ref in memory pool # Radix Cache takes one ref in memory pool
self.token_to_kv_pool.free(indices[last_uncached_pos:new_prefix_len]) self.token_to_kv_pool.free(indices[last_uncached_pos:new_prefix_len])