From af36f848db8025d7d20f76f617e43d79401f5a9a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 19:55:36 +0800 Subject: [PATCH 1/6] refactor(ops): use canonical InfiniOps activation APIs --- src/infinicore/ops/gelu/gelu_infiniops.cc | 4 ++-- src/infinicore/ops/gelutanh/gelutanh_infiniops.cc | 7 +++++-- src/infinicore/ops/relu/relu_infiniops.cc | 4 ++-- src/infinicore/ops/sigmoid/sigmoid_infiniops.cc | 4 ++-- src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc | 4 ++-- submodules/InfiniOps | 2 +- 6 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/infinicore/ops/gelu/gelu_infiniops.cc b/src/infinicore/ops/gelu/gelu_infiniops.cc index de0b35fde..ca8b1a5f0 100644 --- a/src/infinicore/ops/gelu/gelu_infiniops.cc +++ b/src/infinicore/ops/gelu/gelu_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/gelu_infinilm.h" +#include "base/gelu.h" #include @@ -22,7 +22,7 @@ void calculate(Tensor output, Tensor input) { TensorMeta output_meta(output); TensorMeta input_meta(input); - infini::ops::GeluInfinilm::Call( + infini::ops::Gelu::Call( handle, config, input_meta.tensor(input), diff --git a/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc b/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc index e9faf4b35..cce7fc62e 100644 --- a/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc +++ b/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc @@ -3,7 +3,9 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/gelutanh_infinilm.h" +#include "base/gelu.h" + +#include namespace infinicore::op::gelutanh_impl::infiniops { namespace { @@ -20,10 +22,11 @@ void calculate(Tensor output, Tensor input) { TensorMeta output_meta(output); TensorMeta input_meta(input); - infini::ops::GelutanhInfinilm::Call( + infini::ops::Gelu::Call( handle, config, input_meta.tensor(input), + std::string{"tanh"}, output_meta.tensor(output)); } diff --git a/src/infinicore/ops/relu/relu_infiniops.cc b/src/infinicore/ops/relu/relu_infiniops.cc index 7890765fe..ba7c510ce 100644 --- a/src/infinicore/ops/relu/relu_infiniops.cc +++ b/src/infinicore/ops/relu/relu_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/relu_infinilm.h" +#include "base/relu.h" namespace infinicore::op::relu_impl::infiniops { namespace { @@ -20,7 +20,7 @@ void calculate(Tensor output, Tensor input) { TensorMeta output_meta(output); TensorMeta input_meta(input); - infini::ops::ReluInfinilm::Call( + infini::ops::Relu::Call( handle, config, input_meta.tensor(input), diff --git a/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc b/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc index 0327e6c7c..194fde6b3 100644 --- a/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc +++ b/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/sigmoid_infinilm.h" +#include "base/sigmoid.h" namespace infinicore::op::sigmoid_impl::infiniops { namespace { @@ -35,7 +35,7 @@ void run(void *planned_meta) { handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::SigmoidInfinilm::Call( + infini::ops::Sigmoid::Call( handle, config, planned->input.tensor(planned->input_tensor), diff --git a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc b/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc index aa1bf3413..4d4d2befc 100644 --- a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc +++ b/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/silu_and_mul_infinilm.h" +#include "base/silu_and_mul.h" namespace infinicore::op::silu_and_mul_impl::infiniops { namespace { @@ -25,7 +25,7 @@ void run(void *planned_meta) { infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::SiluAndMulInfinilm::Call( + infini::ops::SiluAndMul::Call( handle, config, planned->input.tensor(planned->input_tensor), planned->output.tensor(planned->output_tensor)); } diff --git a/submodules/InfiniOps b/submodules/InfiniOps index 1b9c37659..21b07ebcf 160000 --- a/submodules/InfiniOps +++ b/submodules/InfiniOps @@ -1 +1 @@ -Subproject commit 1b9c37659211eeaeedb48b9691f79e4c9e4412f9 +Subproject commit 21b07ebcfdb0f993d2f3b672a4e38788e489fb80 From 251c92133d79e6da5d9679e8c1ee148f3f0f2c45 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 6 Aug 2026 20:41:48 +0800 Subject: [PATCH 2/6] refactor(ops): migrate paged caching to canonical InfiniOps API --- .../paged_caching/paged_caching_infiniops.cc | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc index 664472904..0c5c6ece4 100644 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc +++ b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc @@ -1,23 +1,33 @@ #include "infinicore/ops/paged_caching.hpp" +#include + #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/paged_caching_infinilm.h" +#include "base/reshape_and_cache_flash.h" namespace infinicore::op::paged_caching_impl::infiniops { namespace { using TensorMeta = ::infinicore::op::infiniops::TensorMeta; struct PlannedMeta { - TensorMeta k_cache, v_cache, k, v, slot_mapping; - graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, slot_mapping_tensor; + TensorMeta k, v, slot_mapping, scale, k_cache, v_cache; + graph::GraphTensor k_tensor, v_tensor, slot_mapping_tensor, scale_tensor, k_cache_tensor, v_cache_tensor; }; } // namespace void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType())); INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); - return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping)}; + + // The "auto" cache path ignores scales, but the canonical API requires them. + auto scale = Tensor::empty({1}, DataType::F32, k_cache->device()); + auto k_cache_view = k_cache->permute({0, 2, 1, 3}); + auto v_cache_view = v_cache->permute({0, 2, 1, 3}); + + return new PlannedMeta{ + TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), TensorMeta(scale), TensorMeta(k_cache_view), TensorMeta(v_cache_view), + graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping), graph::GraphTensor(scale), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache)}; } void run(void *planned_meta) { @@ -25,12 +35,15 @@ void run(void *planned_meta) { infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::PagedCachingInfinilm::Call( + infini::ops::ReshapeAndCacheFlash::Call( handle, config, planned->k.tensor(planned->k_tensor), planned->v.tensor(planned->v_tensor), planned->slot_mapping.tensor(planned->slot_mapping_tensor), + planned->scale.tensor(planned->scale_tensor), + planned->scale.tensor(planned->scale_tensor), + std::string{"auto"}, planned->k_cache.tensor(planned->k_cache_tensor), planned->v_cache.tensor(planned->v_cache_tensor)); } From 4cb327293a403c49a929bcfdcf5e4d9f5da6b3b4 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 00:57:53 +0800 Subject: [PATCH 3/6] refactor(ops): migrate topksoftmax to canonical InfiniOps API --- .../ops/topksoftmax/topksoftmax_infiniops.cc | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc b/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc index 4a4679417..9cb7553ba 100644 --- a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc +++ b/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc @@ -3,15 +3,16 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/topksoftmax_infinilm.h" +#include "base/topk_softmax.h" + +#include namespace infinicore::op::topksoftmax_impl::infiniops { namespace { using TensorMeta = ::infinicore::op::infiniops::TensorMeta; struct PlannedMeta { - TensorMeta values, indices, x; - graph::GraphTensor values_tensor, indices_tensor, x_tensor; - size_t topk; + TensorMeta values, indices, token_expert_indices, x; + graph::GraphTensor values_tensor, indices_tensor, token_expert_indices_tensor, x_tensor; int norm; }; } // namespace @@ -19,7 +20,17 @@ struct PlannedMeta { void *plan(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(values->device().getType())); INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x); - return new PlannedMeta{TensorMeta(values), TensorMeta(indices), TensorMeta(x), graph::GraphTensor(values), graph::GraphTensor(indices), graph::GraphTensor(x), topk, norm}; + auto token_expert_indices = Tensor::empty({x->size(0), topk}, DataType::I32, indices->device()); + return new PlannedMeta{ + TensorMeta(values), + TensorMeta(indices), + TensorMeta(token_expert_indices), + TensorMeta(x), + graph::GraphTensor(values), + graph::GraphTensor(indices), + graph::GraphTensor(token_expert_indices), + graph::GraphTensor(x), + norm}; } void run(void *planned_meta) { @@ -27,14 +38,16 @@ void run(void *planned_meta) { infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::TopksoftmaxInfinilm::Call( + infini::ops::TopkSoftmax::Call( handle, config, planned->x.tensor(planned->x_tensor), - static_cast(planned->topk), + std::optional{}, + std::optional{}, planned->norm != 0, planned->values.tensor(planned->values_tensor), - planned->indices.tensor(planned->indices_tensor)); + planned->indices.tensor(planned->indices_tensor), + planned->token_expert_indices.tensor(planned->token_expert_indices_tensor)); } void cleanup(void **planned_meta_ptr) { From 25f46ab9de8564dd72191faa3794704f8ca67d8a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 19:56:42 +0800 Subject: [PATCH 4/6] refactor(ops): use canonical InfiniOps softmax API --- src/infinicore/ops/softmax/softmax_infiniops.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/infinicore/ops/softmax/softmax_infiniops.cc b/src/infinicore/ops/softmax/softmax_infiniops.cc index c0087e1f8..3765eeb86 100644 --- a/src/infinicore/ops/softmax/softmax_infiniops.cc +++ b/src/infinicore/ops/softmax/softmax_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/softmax_infinilm.h" +#include "base/softmax.h" #include @@ -22,7 +22,7 @@ void calculate(Tensor output, Tensor input, int axis) { TensorMeta output_meta(output); TensorMeta input_meta(input); - infini::ops::SoftmaxInfinilm::Call( + infini::ops::Softmax::Call( handle, config, input_meta.tensor(input), From 22334e188944490454e17e37c97dcd2ad9a2be4a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 20:02:57 +0800 Subject: [PATCH 5/6] refactor(ops): use InfiniOps Copy for rearrange --- src/infinicore/ops/rearrange/rearrange_infiniops.cc | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/infinicore/ops/rearrange/rearrange_infiniops.cc b/src/infinicore/ops/rearrange/rearrange_infiniops.cc index 24b275763..cd642c23e 100644 --- a/src/infinicore/ops/rearrange/rearrange_infiniops.cc +++ b/src/infinicore/ops/rearrange/rearrange_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/rearrange_infinilm.h" +#include "base/copy.h" namespace infinicore::op::rearrange_impl::infiniops { namespace { @@ -25,8 +25,8 @@ void run(void *planned_meta) { infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; - infini::ops::RearrangeInfinilm::Call( - handle, config, planned->x.tensor(planned->x_tensor), planned->y.tensor(planned->y_tensor)); + infini::ops::Copy::Call( + handle, config, planned->x.tensor(planned->x_tensor), false, planned->y.tensor(planned->y_tensor)); } void cleanup(void **planned_meta_ptr) { @@ -34,8 +34,5 @@ void cleanup(void **planned_meta_ptr) { *planned_meta_ptr = nullptr; } -// ReArrange is used by Tensor::copy_from during test result conversion. Keep -// the InfiniCore/InfiniOp implementation active until this adapter gets the same -// CUDA shim treatment as Gemm. } // namespace infinicore::op::rearrange_impl::infiniops #endif From 3a77f54a9dd58a975ba6ecb765f8461cd6d8e2c6 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 20:04:18 +0800 Subject: [PATCH 6/6] refactor(ops): use InfiniOps Convolution for conv2d --- src/infinicore/ops/conv2d/conv2d_infiniops.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/infinicore/ops/conv2d/conv2d_infiniops.cc b/src/infinicore/ops/conv2d/conv2d_infiniops.cc index 3cdf0c63e..f1a048569 100644 --- a/src/infinicore/ops/conv2d/conv2d_infiniops.cc +++ b/src/infinicore/ops/conv2d/conv2d_infiniops.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" -#include "base/conv_infinilm.h" +#include "base/convolution.h" #include #include @@ -48,15 +48,17 @@ void calculate(Tensor output, bias_meta.emplace(bias); } - infini::ops::ConvInfinilm::Call( + infini::ops::Convolution::Call( handle, config, input_meta.tensor(input), weight_meta.tensor(weight), bias_meta ? std::optional{bias_meta->tensor(bias)} : std::nullopt, - toInt64Vector(pads, n), toInt64Vector(strides, n), + toInt64Vector(pads, n), toInt64Vector(dilations, n), + false, + std::vector(n, 0), int64_t{1}, output_meta.tensor(output)); }