diff --git a/docker/Dockerfile-NPU b/docker/Dockerfile-NPU
index 5fb3da5e7..1579f37dd 100644
--- a/docker/Dockerfile-NPU
+++ b/docker/Dockerfile-NPU
@@ -10,7 +10,7 @@ COPY . /workspace/unified-cache-management
 
 RUN pip config set global.index-url ${PIP_INDEX_URL}
 
-RUN export PLATFORM="ascend" && \
+RUN export PLATFORM="ascend" ENABLE_SPARSE=true && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     pip install -v -e /workspace/unified-cache-management --no-build-isolation
 
diff --git a/ucm/sparse/kvcomp/CMakeLists.txt b/ucm/sparse/kvcomp/CMakeLists.txt
index ee5fa0c59..957ef2278 100644
--- a/ucm/sparse/kvcomp/CMakeLists.txt
+++ b/ucm/sparse/kvcomp/CMakeLists.txt
@@ -48,7 +48,7 @@ else()
 endif()
 
 string(TOLOWER "$ENV{PLATFORM}" PLATFORM_ENV)
-if(PLATFORM_ENV STREQUAL "cuda")
+if(RUNTIME_ENVIRONMENT STREQUAL "cuda")
     message(STATUS "Building kvcomp for CUDA...")
     add_subdirectory(hash_retrieval)
     add_subdirectory(ham_dist)
diff --git a/ucm/sparse/kvcomp/ham_dist/CMakeLists.txt b/ucm/sparse/kvcomp/ham_dist/CMakeLists.txt
index 6dbd8e22a..87a4a26de 100644
--- a/ucm/sparse/kvcomp/ham_dist/CMakeLists.txt
+++ b/ucm/sparse/kvcomp/ham_dist/CMakeLists.txt
@@ -6,6 +6,10 @@ set(CMAKE_CUDA_COMPILER ${CUDA_ROOT}/bin/nvcc)
 set(CMAKE_CUDA_ARCHITECTURES 75 80 86 89 90)
 enable_language(CUDA)
 
+# ---- CUDA Toolkit ----
+find_package(CUDAToolkit REQUIRED)
+message(STATUS "Found CUDAToolkit: ${CUDAToolkit_INCLUDE_DIRS}")
+
 # ---- Python ----
 find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
 
@@ -52,6 +56,7 @@ set(INCLUDE_DIRS
     ${PYTORCH_PATH}/include
     ${PYTORCH_PATH}/include/torch/csrc/api/include
     ${CMAKE_CURRENT_SOURCE_DIR}
+    ${CUDAToolkit_INCLUDE_DIRS}
 )
 
 set(LIBRARY_DIRS
@@ -66,6 +71,7 @@ set(LIBRARIES
     torch_cpu
     torch_python
     pthread
+    CUDA::cudart
 )
 
 # ---- Build: hamming python module (.so) ----