Merge pull request #232 from InfiniTensor/issue/231

issue/231 - remove ninetoothed dependency by default

README.md

@@ -63,6 +63,12 @@ python scripts/test_ppl.py --model-path MODEL_PATH [--ndev NDEV] [--max-batch MA
     ```
 
 
+  - 选择是否使用九齿计算路径，默认为false，即不依赖九齿算子
+    ```bash
+      xmake f --ninetoothed= [true | false] -cv
+    ```
+
+
   - 安装 InfiniLM Python 包
     ```bash
       pip install -e .

csrc/cache/kv_cache.cpp

@@ -93,16 +93,14 @@ StaticKVCache::update(size_t layer_idx,
 
     auto device = k_cache_layer->device();
 
-    if (device.getType() == infinicore::Device::Type::NVIDIA
-        || device.getType() == infinicore::Device::Type::ILUVATAR
-        || device.getType() == infinicore::Device::Type::METAX) {
+#ifdef ENABLE_NINETOOTHED
     infinicore::op::kv_caching_(
         k_cache_layer,
         v_cache_layer,
         k,
         v,
         past_sequence_lengths);
-    } else {
+#else
     size_t cache_pos = reinterpret_cast<int64_t *>(past_sequence_lengths->to(infinicore::Device::cpu())->data())[0];
     auto result_len = cache_pos + update_len;
     ASSERT(result_len <= cache_len_);
@@ -112,7 +110,7 @@ StaticKVCache::update(size_t layer_idx,
 
     k_cache_update->copy_from(k);
     v_cache_update->copy_from(v);
-    }
+#endif
 
     return {k_cache_layer, v_cache_layer};
 }

xmake.lua

@@ -8,6 +8,16 @@ set_toolchains("gcc")
 add_includedirs("third_party/spdlog/include")
 add_includedirs("third_party/json/single_include/")
 
+option("ninetoothed")
+    set_default(false)
+    set_showmenu(true)
+    set_description("Whether to complie NineToothed specifc path")
+option_end()
+
+if has_config("ninetoothed") then
+    add_defines("ENABLE_NINETOOTHED")
+end
+
 target("infinicore_infer")
     set_kind("shared")