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

Fuse more ops & Simplify token mapping (#1758)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-10-22 23:20:43 -07:00
committed by GitHub
parent 17536e7e3d
commit ad4125d1a9
9 changed files with 99 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -51,7 +51,7 @@ class ReqToTokenPool:
self.write = self.write_without_records
def write(self, indices, values):
# Keep the signature for type checking, will be initialized during runtime
# Keep the signature for type checking. It will be assigned during runtime.
raise NotImplementedError()
def available_size(self):
@@ -221,16 +221,21 @@ class MHATokenToKVPool(BaseTokenToKVPool):
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
if cache_v.dtype != self.dtype:
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
self.k_buffer[layer_id][loc] = cache_k.view(self.store_dtype)
self.v_buffer[layer_id][loc] = cache_v.view(self.store_dtype)
else:
self.k_buffer[layer_id][loc] = cache_k
self.v_buffer[layer_id][loc] = cache_v
copy_two_array(
loc,
self.k_buffer[layer_id],
cache_k,
self.v_buffer[layer_id],
cache_v,
self.dtype,
self.store_dtype,
)
@torch.compile(dynamic=True)
def copy_two_array(loc, dst_1, src_1, dst_2, src_2, dtype, store_dtype):
dst_1[loc] = src_1.to(dtype).view(store_dtype)
dst_2[loc] = src_2.to(dtype).view(store_dtype)
class MLATokenToKVPool(BaseTokenToKVPool):