diff --git a/qwen3_6_scripts/patch_corex_copy_blocks.py b/qwen3_6_scripts/patch_corex_copy_blocks.py new file mode 100644 index 00000000..60d48cee --- /dev/null +++ b/qwen3_6_scripts/patch_corex_copy_blocks.py @@ -0,0 +1,33 @@ +from patch_utils import package_root, replace_once + + +CUSTOM_OPS = package_root("vllm") / "_custom_ops.py" + +CLEAN_BLOCK = """\ +def copy_blocks(key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor) -> None: + ixf_F.copy_blocks(key_caches, value_caches, block_mapping) +""" + +COMPATIBLE_BLOCK = """\ +def copy_blocks(key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor) -> None: + # BI100 CoreX 3.2.3 exposes vllm_copy_blocks, not copy_blocks. + _fn = getattr(ixf_F, "copy_blocks", None) or getattr( + ixf_F, "vllm_copy_blocks", None) + if _fn is None: + raise RuntimeError( + "ixformer exposes neither copy_blocks nor vllm_copy_blocks") + _fn(key_caches, value_caches, block_mapping) +""" + + +replace_once( + CUSTOM_OPS, + CLEAN_BLOCK, + COMPATIBLE_BLOCK, + required=True, + already_contains="BI100 CoreX 3.2.3 exposes vllm_copy_blocks", +) \ No newline at end of file diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 6e5f5125..d288a31e 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -149,8 +149,9 @@ cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py" cp ./block_major_kv_cache.py "${VLLM_ROOT}/block_major_kv_cache.py" cp ./gdn_prefix.py "${VLLM_ROOT}/gdn_prefix.py" -build_stage "installing CoreX paged-KV swap compatibility" +build_stage "installing CoreX paged-KV swap/copy compatibility" python3 ./patch_corex_swap_blocks.py +python3 ./patch_corex_copy_blocks.py python3 ./patch_block_major_cache_engine.py python3 ./patch_worker_cache_transfer_order.py @@ -458,7 +459,4 @@ python3 ./verify_dlopen_chain.py --vllm-root "${VLLM_ROOT}" || { echo "[WARN] dlopen chain verification found issues (non-fatal)" } -build_stage "patch script completed" - - - +build_stage "patch script completed" \ No newline at end of file