[AMD] DSv4: bring HIP compress-state pool into the memory_saver KV_CACHE region (#31747)
Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
This commit is contained in:
co-authored by
Zhiyao Jiang
parent
c32c4ef79c
commit
f5bcd00e16
@@ -153,30 +153,25 @@ class CompressStatePool:
|
||||
allocation boilerplate. Requires ``self._size`` and ``self.last_dim``
|
||||
to be set already.
|
||||
"""
|
||||
if _is_hip:
|
||||
self.kv_score_buffer = KVAndScore(
|
||||
torch.empty((self._size, self.last_dim), dtype=dtype, device=device)
|
||||
)
|
||||
else:
|
||||
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
|
||||
enable=enable_memory_saver
|
||||
)
|
||||
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
|
||||
maybe_init_custom_mem_pool(device=device)
|
||||
)
|
||||
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
|
||||
with (
|
||||
torch.cuda.use_mem_pool(self.custom_mem_pool)
|
||||
if self.custom_mem_pool
|
||||
else nullcontext()
|
||||
):
|
||||
self.kv_score_buffer = KVAndScore(
|
||||
torch.empty(
|
||||
(self._size, self.last_dim),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
|
||||
enable=enable_memory_saver
|
||||
)
|
||||
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
|
||||
maybe_init_custom_mem_pool(device=device)
|
||||
)
|
||||
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
|
||||
with (
|
||||
torch.cuda.use_mem_pool(self.custom_mem_pool)
|
||||
if self.custom_mem_pool
|
||||
else nullcontext()
|
||||
):
|
||||
self.kv_score_buffer = KVAndScore(
|
||||
torch.empty(
|
||||
(self._size, self.last_dim),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
|
||||
@property
|
||||
def state_cache_3d(self) -> torch.Tensor:
|
||||
|
||||
@@ -646,10 +646,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
|
||||
self._init_compressed_layer_mapping()
|
||||
|
||||
if _is_hip:
|
||||
self._init_paged_compress_states(False)
|
||||
else:
|
||||
self._init_paged_compress_states(enable_memory_saver)
|
||||
self._init_paged_compress_states(enable_memory_saver)
|
||||
|
||||
def get_unified_kv(self, layer_id: int) -> torch.Tensor:
|
||||
# Under HiCache the compressed region is loaded H->D per layer; wait for this
|
||||
|
||||
Reference in New Issue
Block a user