[AMD] Fix nightly-8-gpu-mi35x-deepseek-v4-flash-rocm720 OOM issue (#28941)
This commit is contained in:
@@ -578,29 +578,29 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
global_page_size=swa_page_size,
|
||||
)
|
||||
|
||||
c4_kv_pool_type = DeepSeekV4SingleKVPool
|
||||
if enable_hisparse:
|
||||
c4_kv_pool_type = HiSparseC4DevicePool
|
||||
self.c4_kv_pool = self._make_kv_pool(
|
||||
size=c4_size,
|
||||
page_size=c4_page_size,
|
||||
dtype=dtype,
|
||||
layer_num=c4_layer_num,
|
||||
device=device,
|
||||
enable_memory_saver=enable_memory_saver,
|
||||
global_page_size=page_size,
|
||||
cls=c4_kv_pool_type,
|
||||
)
|
||||
c4_kv_pool_type = DeepSeekV4SingleKVPool
|
||||
if enable_hisparse:
|
||||
c4_kv_pool_type = HiSparseC4DevicePool
|
||||
self.c4_kv_pool = self._make_kv_pool(
|
||||
size=c4_size,
|
||||
page_size=c4_page_size,
|
||||
dtype=dtype,
|
||||
layer_num=c4_layer_num,
|
||||
device=device,
|
||||
enable_memory_saver=enable_memory_saver,
|
||||
global_page_size=page_size,
|
||||
cls=c4_kv_pool_type,
|
||||
)
|
||||
|
||||
self.c128_kv_pool = self._make_kv_pool(
|
||||
size=c128_size,
|
||||
page_size=c128_page_size,
|
||||
dtype=dtype,
|
||||
layer_num=c128_layer_num,
|
||||
device=device,
|
||||
enable_memory_saver=enable_memory_saver,
|
||||
global_page_size=page_size,
|
||||
)
|
||||
self.c128_kv_pool = self._make_kv_pool(
|
||||
size=c128_size,
|
||||
page_size=c128_page_size,
|
||||
dtype=dtype,
|
||||
layer_num=c128_layer_num,
|
||||
device=device,
|
||||
enable_memory_saver=enable_memory_saver,
|
||||
global_page_size=page_size,
|
||||
)
|
||||
|
||||
indexer_size = (
|
||||
self.c4_logical_size
|
||||
|
||||
Reference in New Issue
Block a user