[codex] Enable Kunlun Qwen2.5 inference support

python/infinicore/__init__.py

@@ -108,10 +108,21 @@
 from infinicore.ops.mha import mha
 from infinicore.ops.mha_kvcache import mha_kvcache
 from infinicore.ops.mha_varlen import mha_varlen
-from infinicore.ops.moore_mate_flash_attn import (
-    moore_mate_flash_attn_decode,
-    moore_mate_flash_attn_prefill,
-)
+_optional_ops = []
+
+try:
+    from infinicore.ops.moore_mate_flash_attn import (
+        moore_mate_flash_attn_decode,
+        moore_mate_flash_attn_prefill,
+    )
+
+    _optional_ops.extend(
+        ["moore_mate_flash_attn_prefill", "moore_mate_flash_attn_decode"]
+    )
+except Exception:
+    # Optional Moore Threads flash-attn bridge. Keep other backends importable
+    # when its Python-side dependencies are absent or incompatible.
+    pass
 from infinicore.ops.mrope import mrope
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
@@ -288,14 +299,12 @@
     "zeros",
     "sum",
     "var_mean",
-    "moore_mate_flash_attn_prefill",
-    "moore_mate_flash_attn_decode",
     "var",
     "topk",
     "all",
     "set_printoptions",
     "printoptions",
-]
+] + _optional_ops
 
 use_ntops = False
 

src/infinicore/ops/linear/linear.cc

@@ -1,6 +1,6 @@
 #include "infinicore/ops/linear.hpp"
+#include "infinicore/ops/add.hpp"
 #include "infinicore/ops/gemm.hpp"
-#include "infinicore/ops/rearrange.hpp"
 
 namespace infinicore::op {
 
@@ -44,17 +44,16 @@ void linear_(Tensor out,
 
     // linear transformation
     Tensor out_view = out->view({N, out_features});
-    // Add bias
-    float beta = 0.0f;
-    if (bias.has_value()) {
-        rearrange_(out_view,
-                   bias.value()->as_strided({N, out_features}, {0, 1}));
-        beta = 1.0f;
-    }
-
+    auto weight_t = weight->permute({1, 0})->contiguous();
     gemm_(out_view,
           input->view({N, in_features}),
-          weight->permute({1, 0}), alpha, beta);
+          weight_t, alpha, 0.0f);
+
+    if (bias.has_value()) {
+        add_(out_view,
+             out_view,
+             bias.value()->as_strided({N, out_features}, {0, 1}));
+    }
 }
 
 } // namespace infinicore::op

src/infiniop/ops/rope/kunlun/rope_kunlun.xpu

@@ -206,6 +206,20 @@ infiniStatus_t Descriptor::calculate(
         default:
             return INFINI_STATUS_BAD_TENSOR_DTYPE;
         }
+    } else if (_info.pos_type == INFINI_DTYPE_I64) {
+        switch (_info.data_type) {
+        case INFINI_DTYPE_F32:
+            LAUNCH_KERNEL(float, int64_t);
+            break;
+        case INFINI_DTYPE_F16:
+            LAUNCH_KERNEL(half, int64_t);
+            break;
+        case INFINI_DTYPE_BF16:
+            LAUNCH_KERNEL(bfloat16_t, int64_t);
+            break;
+        default:
+            return INFINI_STATUS_BAD_TENSOR_DTYPE;
+        }
     } else if (_info.pos_type == INFINI_DTYPE_U32) {
         switch (_info.data_type) {
         case INFINI_DTYPE_F32: