diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index cdefb904b..fae0340fe 100644 --- a/python/infinicore/__init__.py +++ b/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 diff --git a/src/infinicore/ops/linear/linear.cc b/src/infinicore/ops/linear/linear.cc index 976bf0bf9..dd529b790 100644 --- a/src/infinicore/ops/linear/linear.cc +++ b/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 diff --git a/src/infiniop/ops/rope/kunlun/rope_kunlun.xpu b/src/infiniop/ops/rope/kunlun/rope_kunlun.xpu index 5e7683d21..0a13f04f6 100644 --- a/src/infiniop/ops/rope/kunlun/rope_kunlun.xpu +++ b/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: