feat(llama-infini): 添加 ascend 选项

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Dec 31, 2024 · a3ba3da · a3ba3da
1 parent d79b107
commit a3ba3da
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/docs/user-guide/doc.md b/docs/user-guide/doc.md
@@ -162,7 +162,7 @@ cargo test --release --package `model` --lib -- `test` --exact --nocapture
 | `model`        | `test`                      | `DEVICES` | 说明
 |:--------------:|:---------------------------:|:----------|:-
 | `llama-cpu`    | `infer::test_infer`         | 默认值“1”。任意间隔的正整数数组，表示每个线程分布模型的份数，数组的项数必须是 2 的幂 | 纯 cpu 后端，不需要任何额外依赖
-| `llama-infini` | `infer::test_infer`         | 默认值“cpu;0”。格式“硬件类型; 卡号” | 九源统一软件栈后端
+| `llama-infini` | `infer::test_infer`         | 默认值“cpu;0”。格式“硬件类型; 卡号”，硬件类型目前支持 `cpu`、`nv`、`ascend` | 九源统一软件栈后端
 | `llama-cl`     | `infer::test_infer`         | TODO | OpenCL 后端
 | `llama-nv`     | `infer::test_infer`         | 默认值“0”。单个非负整数，推理使用的卡号 | 原生 CUDA Toolkit 后端
 | `llama-nv`     | `nccl_parallel::test_infer` | 默认值“0”。任意间隔的非负整数集合，参与分布式推理的卡号 | 原生 CUDA Toolkit 后端，同时依赖 NCCL 实现分布式

diff --git a/models/llama/infini/src/infer.rs b/models/llama/infini/src/infer.rs
@@ -72,6 +72,10 @@ fn test_infer() {
             infini_rt::init(infini_rt::DEVICE_NVIDIA);
             WorkerSeed::new(InfiniNode::nv_gpu(&indices))
         }
+        "ascend" => {
+            infini_rt::init(infini_rt::DEVICE_ASCEND);
+            WorkerSeed::new(InfiniNode::ascend_npu(&indices))
+        }
         _ => todo!(),
     };
     thread::scope(|s| {