refactor(ops): route Llama paths through InfiniOps

src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc

@@ -40,7 +40,8 @@ bool canUseInfiniOps(const Tensor &out,
                      const Tensor &block_table,
                      const std::optional<Tensor> &alibi_slopes) {
     const auto dtype = q->dtype();
-    if (out->device().getType() != Device::Type::NVIDIA
+    const auto device_type = out->device().getType();
+    if ((device_type != Device::Type::NVIDIA && device_type != Device::Type::METAX)
         || q->ndim() != 4
         || out->ndim() != 4
         || k_cache->ndim() != 4

src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc

@@ -42,7 +42,8 @@ bool canUseInfiniOps(const Tensor &out,
                      const std::optional<Tensor> &alibi_slopes) {
     const bool paged = block_table.has_value();
     const auto dtype = q->dtype();
-    if (out->device().getType() != Device::Type::NVIDIA
+    const auto device_type = out->device().getType();
+    if ((device_type != Device::Type::NVIDIA && device_type != Device::Type::METAX)
         || q->ndim() != 3
         || out->ndim() != 3
         || ((paged && (k->ndim() != 4 || v->ndim() != 4))

src/infinicore/ops/select_last_token_hidden/select_last_token_hidden_infiniops.cc

@@ -0,0 +1,85 @@
+#include "infinicore/ops/select_last_token_hidden.hpp"
+
+#ifdef ENABLE_INFINIOPS_API
+#include "../infiniops_impl.hpp"
+
+#include "base/add.h"
+#include "base/index_select.h"
+
+namespace infinicore::op::select_last_token_hidden_impl::infiniops {
+namespace {
+using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
+
+struct PlannedMeta {
+    TensorMeta output, hidden_states, input_offsets, one, indices;
+    graph::GraphTensor output_tensor, hidden_states_tensor, input_offsets_tensor, one_tensor, indices_tensor;
+};
+} // namespace
+
+void *plan(Tensor output, const Tensor &hidden_states, const Tensor &input_offsets) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, hidden_states, input_offsets);
+
+    const auto hidden_size = hidden_states->size(2);
+    auto hidden_states_view = hidden_states->view({hidden_states->numel() / hidden_size, hidden_size});
+    auto output_view = output->view({output->numel() / hidden_size, hidden_size});
+    const auto num_requests = input_offsets->numel() - 1;
+    auto input_offsets_view = input_offsets->narrow({{0, 1, num_requests}});
+    auto indices = Tensor::empty({num_requests}, DataType::I32, input_offsets->device());
+    auto one = Tensor::empty({1}, DataType::I32, input_offsets->device());
+    const int32_t one_value = 1;
+    context::memcpyH2D(one->data(), &one_value, sizeof(one_value), false);
+
+    return new PlannedMeta{
+        TensorMeta(output_view),
+        TensorMeta(hidden_states_view),
+        TensorMeta(input_offsets_view),
+        TensorMeta(one),
+        TensorMeta(indices),
+        graph::GraphTensor(output_view),
+        graph::GraphTensor(hidden_states_view),
+        graph::GraphTensor(input_offsets_view),
+        graph::GraphTensor(one),
+        graph::GraphTensor(indices)};
+}
+
+void run(void *planned_meta) {
+    auto *planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+    infini::ops::Handle handle;
+    handle.set_stream(context::getStream());
+    infini::ops::Config add_config;
+
+    infini::ops::Add::Call(
+        handle,
+        add_config,
+        planned->input_offsets.tensor(planned->input_offsets_tensor),
+        planned->one.tensor(planned->one_tensor),
+        -1.0,
+        planned->indices.tensor(planned->indices_tensor));
+    infini::ops::Config index_select_config;
+    index_select_config.set_implementation_index(8);
+    infini::ops::IndexSelect::Call(
+        handle,
+        index_select_config,
+        planned->hidden_states.tensor(planned->hidden_states_tensor),
+        planned->indices.tensor(planned->indices_tensor),
+        int64_t{0},
+        planned->output.tensor(planned->output_tensor));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+static bool registered = []() {
+    SelectLastTokenHidden::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
+    SelectLastTokenHidden::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
+    SelectLastTokenHidden::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
+    SelectLastTokenHidden::plan_dispatcher().registerDevice(Device::Type::METAX, &plan);
+    SelectLastTokenHidden::run_dispatcher().registerDevice(Device::Type::METAX, &run);
+    SelectLastTokenHidden::cleanup_dispatcher().registerDevice(Device::Type::METAX, &cleanup);
+    return true;
+}();
+
+} // namespace infinicore::op::select_last_token_hidden_impl::infiniops
+#endif

xmake.lua

@@ -357,26 +357,37 @@ local infiniops_external_built = false
 
 local function configure_infiniops_ops(infiniops_ops)
     if not infiniops_ops or #infiniops_ops == 0 then
-        return infiniops_ops, false, false
+        return nil, false, false
     end
 
-    local selected = {}
+    local selection = {}
     local with_linked_flash_attn_with_kvcache = false
     local with_linked_flash_attn_varlen_func = false
     for _, op in ipairs(infiniops_ops:split("[,;]")) do
         op = op:trim()
         if #op > 0 then
-            table.insert(selected, op)
-            if has_config("nv-gpu") and op == "flash_attn_with_kvcache" then
+            local use_linked_implementation = false
+            if (has_config("nv-gpu") or has_config("metax-gpu")) and op == "flash_attn_with_kvcache" then
                 with_linked_flash_attn_with_kvcache = true
+                use_linked_implementation = true
             end
-            if has_config("nv-gpu") and op == "flash_attn_varlen_func" then
+            if (has_config("nv-gpu") or has_config("metax-gpu")) and op == "flash_attn_varlen_func" then
                 with_linked_flash_attn_varlen_func = true
+                use_linked_implementation = true
             end
+            local implementations = "all"
+            if use_linked_implementation then
+                implementations = {16}
+            elseif op == "argmax" or op == "index_select" then
+                implementations = {8}
+            elseif op == "rms_norm" or op == "silu_and_mul" or op == "topk_softmax" then
+                implementations = {0}
+            end
+            selection[op] = {implementations = implementations}
         end
     end
 
-    return table.concat(selected, ","), with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func
+    return selection, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func
 end
 
 local function get_infiniops_backend_cmake_arg()
@@ -399,7 +410,7 @@ local function get_infiniops_backend_cmake_arg()
     return enabled[1]
 end
 
-local function build_infiniops_external(xmake_os)
+local function build_infiniops_external(xmake_os, json)
     if not has_config("infiniops") or infiniops_external_built then
         return
     end
@@ -420,7 +431,11 @@ local function build_infiniops_external(xmake_os)
         or has_config("metax-gpu")
         or (has_config("iluvatar-gpu") and has_config("aten")) then
         table.insert(cmake_config_args, "-DWITH_TORCH=ON")
-        table.insert(cmake_config_args, "-DINFINI_OPS_TORCH_OPS=argmax")
+        local torch_ops = "argmax"
+        if has_config("nv-gpu") or has_config("metax-gpu") then
+            torch_ops = torch_ops .. ",index_select"
+        end
+        table.insert(cmake_config_args, "-DINFINI_OPS_TORCH_OPS=" .. torch_ops)
     end
     if has_config("iluvatar-gpu") and has_config("aten") then
         table.insert(cmake_config_args, "-DTORCH_CXX11_ABI=0")
@@ -432,12 +447,15 @@ local function build_infiniops_external(xmake_os)
             table.insert(cmake_config_args, "-DILUVATAR_ARCH=" .. iluvatar_arch)
         end
     end
-    local infiniops_ops, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
+    local infiniops_ops_config, with_linked_flash_attn_with_kvcache, with_linked_flash_attn_varlen_func = configure_infiniops_ops(os.getenv("INFINI_OPS_OPS"))
     if with_linked_flash_attn_with_kvcache or with_linked_flash_attn_varlen_func then
         table.insert(cmake_config_args, "-DWITH_LINKED=ON")
     end
-    if infiniops_ops and #infiniops_ops > 0 then
-        table.insert(cmake_config_args, "-DINFINI_OPS_OPS=" .. infiniops_ops)
+    if infiniops_ops_config then
+        xmake_os.mkdir(infiniops_builddir)
+        local infiniops_ops_config_path = path.join(infiniops_builddir, "ops.json")
+        json.savefile(infiniops_ops_config_path, infiniops_ops_config)
+        table.insert(cmake_config_args, "-DINFINI_OPS_OPS=" .. infiniops_ops_config_path)
     end
     local cmake_cuda_architectures = get_infiniops_cuda_architectures()
     if cmake_cuda_architectures and cmake_cuda_architectures ~= "" then
@@ -689,7 +707,7 @@ target("infiniops_external")
     set_default(false)
 
     on_build(function (target)
-        build_infiniops_external(os)
+        build_infiniops_external(os, import("core.base.json"))
     end)
 target_end()
 
@@ -753,7 +771,7 @@ target("infinicore_cpp_api")
         add_links("infiniops")
         add_rpathdirs(INFINI_ROOT .. "/lib")
         on_load(function (target)
-            build_infiniops_external(os)
+            build_infiniops_external(os, import("core.base.json"))
         end)
         after_install(function (target)
             local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")