[bugfix] kvcomp config (#584)

Author: wangwenxin0312

examples/offline_inference_kvcomphbm.py

@@ -77,11 +77,7 @@ def build_llm_with_uc(module_path: str, name: str, model: str):
                     },
                 }
             ],
-            "ucm_sparse_config": {
-                "GSA": {
-                    "kvcompOnDevice_config_path": "/workspace/unified-cache-management/ucm/sparse/kvcomp/configs/kvcomp_deepseek_v2_lite_config.json"
-                }
-            },
+            "ucm_sparse_config": {"GSA": {}},
         },
     )
 

setup.py

@@ -148,4 +148,5 @@ def build_cmake(self, ext: CMakeExtension):
     cmdclass={"build_ext": CMakeBuild},
     zip_safe=False,
     include_package_data=False,
+    package_data={"ucm": ["sparse/kvcomp/configs/**/*.json"]},
 )

ucm/sparse/kvcomp/kvcomp_hbm.py

@@ -1,3 +1,5 @@
+from importlib import resources
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
 import torch
@@ -24,6 +26,49 @@
 ReqType = Union[str, int]
 
 
+def kvcomp_config_path_for_model(vllm_config) -> str:
+    model = vllm_config.model_config.model.lower()
+    logger.info("[KvComp] model name: %s", model)
+
+    if "deepseek" in model and "r1" in model:
+        rel = "ucm/sparse/kvcomp/configs/kvcomp_deepseek_r1_awq_config.json"
+    elif "qwen3" in model and "32b" in model:
+        rel = "ucm/sparse/kvcomp/configs/kvcomp_qwen3_32B_config.json"
+    else:
+        raise ValueError(f"[KvCompOnDevice] Unsupported model for kvcomp: {model}")
+
+    logger.info("[KvComp] target relative path: %s", rel)
+
+    cur = Path(__file__).resolve()
+    repo = cur
+    for depth in range(30):
+        if (
+            (repo / "pyproject.toml").is_file()
+            or (repo / "setup.cfg").is_file()
+            or (repo / ".git").exists()
+        ):
+
+            p = repo / rel
+            logger.info("[KvComp] repo root detected at depth=%d: %s", depth, repo)
+            if p.is_file():
+                logger.info("[KvComp] config loaded from SOURCE tree: %s", p)
+                return str(p)
+            logger.warning("[KvComp] repo root found but config missing: %s", p)
+            break
+        if repo.parent == repo:
+            logger.debug("[KvComp] reached filesystem root, stop searching")
+            break
+
+        repo = repo.parent
+
+    sub = rel[len("ucm/") :] if rel.startswith("ucm/") else rel
+    res = resources.files("ucm").joinpath(*sub.split("/"))
+
+    with resources.as_file(res) as p:
+        logger.info("[KvComp] config loaded from PACKAGE resource (wheel): %s", p)
+        return str(p)
+
+
 class KvCompOnDevice(UcmSparseBase):
     # handle batch
     def __init__(self, vllm_config: VllmConfig, role: UcmSparseRole):
@@ -39,16 +84,7 @@ def __init__(self, vllm_config: VllmConfig, role: UcmSparseRole):
         )
         self.block_size = vllm_config.cache_config.block_size
 
-        self.kvcompOnDevice_cfg = (
-            Config(vllm_config.kv_transfer_config)
-            .get_config()
-            .get("ucm_sparse_config")
-            .get("GSA")
-        )
-
-        kvcompOnDevice_config_path = self.kvcompOnDevice_cfg[
-            "kvcompOnDevice_config_path"
-        ]
+        kvcompOnDevice_config_path = kvcomp_config_path_for_model(vllm_config)
         self.kvcompOnDevice_config = KvCompConfig.from_json(kvcompOnDevice_config_path)
         logger.info(f"read kvcomp config file : {kvcompOnDevice_config_path} ")
         self.hash_topk_tokens = self.kvcompOnDevice_config.vllm_hash_attention_topk

ucm/sparse/state.py

@@ -41,7 +41,7 @@ def ensure_ucm_sparse_initialized(
 
     # Check if UCM sparse is enabled
     ucm_config = Config(vllm_config.kv_transfer_config)
-    ucm_sparse_config = ucm_config.get_config().get("ucm_sparse_config")
+    ucm_sparse_config = ucm_config.get_config().get("ucm_sparse_method")
     if not ucm_sparse_config:
         return