[Feature] dyc8 support prefixcache (#5125)
Some checks failed
CE Compile Job / ce_job_pre_check (push) Has been cancelled
CE Compile Job / print_ce_job_pre_check_outputs (push) Has been cancelled
CE Compile Job / FD-Clone-Linux (push) Has been cancelled
CE Compile Job / Show Code Archive Output (push) Has been cancelled
CE Compile Job / BUILD_SM8090 (push) Has been cancelled
CE Compile Job / BUILD_SM8689 (push) Has been cancelled
CE Compile Job / CE_UPLOAD (push) Has been cancelled
Deploy GitHub Pages / deploy (push) Has been cancelled
Publish Job / publish_pre_check (push) Has been cancelled
Publish Job / print_publish_pre_check_outputs (push) Has been cancelled
Publish Job / FD-Clone-Linux (push) Has been cancelled
Publish Job / Show Code Archive Output (push) Has been cancelled
Publish Job / BUILD_SM8090 (push) Has been cancelled
Publish Job / BUILD_SM8689 (push) Has been cancelled
Publish Job / PADDLE_PYPI_UPLOAD_8090 (push) Has been cancelled
Publish Job / PADDLE_PYPI_UPLOAD_8689 (push) Has been cancelled
Publish Job / Run FD Image Build (push) Has been cancelled
Publish Job / Run FastDeploy Unit Tests and Coverage (push) Has been cancelled
Publish Job / Run FastDeploy LogProb Tests (push) Has been cancelled
Publish Job / Extracted partial CE model tasks to run in CI. (push) Has been cancelled
Publish Job / Run Base Tests (push) Has been cancelled
Publish Job / Run Accuracy Tests (push) Has been cancelled
Publish Job / Run Stable Tests (push) Has been cancelled
CI Images Build / FD-Clone-Linux (push) Has been cancelled
CI Images Build / Show Code Archive Output (push) Has been cancelled
CI Images Build / CI Images Build (push) Has been cancelled
CI Images Build / BUILD_SM8090 (push) Has been cancelled
CI Images Build / Run FastDeploy Unit Tests and Coverage (push) Has been cancelled
CI Images Build / Run FastDeploy LogProb Tests (push) Has been cancelled
CI Images Build / Extracted partial CE model tasks to run in CI. (push) Has been cancelled
CI Images Build / Run Base Tests (push) Has been cancelled
CI Images Build / Publish Docker Images Pre Check (push) Has been cancelled

* dyc8 support prefixcache

* fix cache_trans test case

* update code
This commit is contained in:
kevin
2025-11-21 19:46:26 +08:00
committed by GitHub
parent ab3a2e45ff
commit c068a4f642
5 changed files with 272 additions and 121 deletions

View File

@@ -1450,8 +1450,10 @@ class GPUModelRunner(ModelRunnerBase):
for i in range(self.model_config.num_hidden_layers):
# init key cache
key_cache_name = f"key_caches_{i}_rank{local_rank}.device{self.device_id}"
key_cache_scales_name = f"key_cache_scales_{i}_rank{local_rank}.device{self.device}"
if value_cache_shape:
val_cache_name = f"value_caches_{i}_rank{local_rank}.device{self.device_id}"
value_cache_scales_name = f"value_cache_scales_{i}_rank{local_rank}.device{self.device}"
if create_cache_tensor:
logger.info(f"..creating kv cache for layer {i}: key:{key_cache_shape}, value:{value_cache_shape}")
key_cache = paddle.full(shape=key_cache_shape, fill_value=0, dtype=cache_type)
@@ -1477,12 +1479,25 @@ class GPUModelRunner(ModelRunnerBase):
logger.info(f"..attaching kv cache for layer {i}: key:{key_cache_shape}, value:{value_cache_shape}")
key_cache = paddle.empty(shape=[], dtype=cache_type)
key_cache = share_external_data(key_cache, key_cache_name, key_cache_shape)
if kv_cache_quant_type == "block_wise_fp8":
key_cache_scales = paddle.empty(shape=[], dtype=paddle.get_default_dtype())
key_cache_scales = share_external_data(
key_cache_scales, key_cache_scales_name, kv_cache_scale_shape
)
if value_cache_shape:
val_cache = paddle.empty(shape=[], dtype=cache_type)
val_cache = share_external_data(val_cache, val_cache_name, value_cache_shape)
cache_kvs_list.extend([key_cache, val_cache])
if kv_cache_quant_type == "block_wise_fp8":
val_cache_scales = paddle.empty(shape=[], dtype=paddle.get_default_dtype())
val_cache_scales = share_external_data(
val_cache_scales, value_cache_scales_name, kv_cache_scale_shape
)
cache_kvs_list.extend([key_cache_scales, val_cache_scales])
else:
cache_kvs_list.extend([key_cache])
if kv_cache_quant_type == "block_wise_fp8":
cache_kvs_list.extend([key_cache_scales])
self.share_inputs["caches"] = cache_kvs_list