diff --git a/src/infinicore/ops/paged_caching/paged_caching.cc b/src/infinicore/ops/paged_caching/paged_caching.cc index afc8bf0c6..7eecc10f1 100644 --- a/src/infinicore/ops/paged_caching/paged_caching.cc +++ b/src/infinicore/ops/paged_caching/paged_caching.cc @@ -15,7 +15,17 @@ void PagedCaching::execute(Tensor k_cache, Tensor v_cache, const Tensor &k, cons } void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - PagedCaching::execute(k_cache, v_cache, k, v, slot_mapping); + constexpr Size MAX_TOKENS_PER_LAUNCH = 32768; + const Size num_tokens = k->size(0); + + for (Size start = 0; start < num_tokens; start += MAX_TOKENS_PER_LAUNCH) { + const Size remaining = num_tokens - start; + const Size chunk_size = remaining < MAX_TOKENS_PER_LAUNCH ? remaining : MAX_TOKENS_PER_LAUNCH; + PagedCaching::execute(k_cache, v_cache, + k->narrow({{0, start, chunk_size}}), + v->narrow({{0, start, chunk_size}}), + slot_mapping->narrow({{0, start, chunk_size}})); + } } } // namespace infinicore::op diff --git a/src/infinicore/ops/swiglu/swiglu.cc b/src/infinicore/ops/swiglu/swiglu.cc index 8ee0682ad..dcdf9344c 100644 --- a/src/infinicore/ops/swiglu/swiglu.cc +++ b/src/infinicore/ops/swiglu/swiglu.cc @@ -20,7 +20,32 @@ Tensor swiglu(const Tensor &a, const Tensor &b) { } void swiglu_(Tensor c, const Tensor &a, const Tensor &b) { - SwiGLU::execute(c, a, b); + constexpr Size MAX_ELEMENTS_PER_LAUNCH = Size{1} << 30; + INFINICORE_ASSERT(c->shape() == a->shape()); + INFINICORE_ASSERT(a->shape() == b->shape()); + INFINICORE_ASSERT(c->dtype() == a->dtype()); + INFINICORE_ASSERT(a->dtype() == b->dtype()); + + if (c->numel() <= MAX_ELEMENTS_PER_LAUNCH) { + SwiGLU::execute(c, a, b); + return; + } + + const Size row_width = c->size(c->ndim() - 1); + INFINICORE_ASSERT(row_width > 0 && row_width <= MAX_ELEMENTS_PER_LAUNCH); + const Size num_rows = c->numel() / row_width; + const Size max_rows = MAX_ELEMENTS_PER_LAUNCH / row_width; + auto c_rows = c->view({num_rows, row_width}); + auto a_rows = a->view({num_rows, row_width}); + auto b_rows = b->view({num_rows, row_width}); + + for (Size start = 0; start < num_rows; start += max_rows) { + const Size remaining = num_rows - start; + const Size rows = remaining < max_rows ? remaining : max_rows; + SwiGLU::execute(c_rows->narrow({{0, start, rows}}), + a_rows->narrow({{0, start, rows}}), + b_rows->narrow({{0, start, rows}})); + } } } // namespace infinicore::op diff --git a/src/infinicore/ops/swiglu/swiglu_infiniop.cc b/src/infinicore/ops/swiglu/swiglu_infiniop.cc index fbb76b570..e874318ed 100644 --- a/src/infinicore/ops/swiglu/swiglu_infiniop.cc +++ b/src/infinicore/ops/swiglu/swiglu_infiniop.cc @@ -1,3 +1,4 @@ +#include "swiglu_infiniop.hpp" #include "infinicore/ops/swiglu.hpp" #include "../infiniop_impl.hpp" diff --git a/src/infinicore/ops/swiglu/swiglu_infiniop.hpp b/src/infinicore/ops/swiglu/swiglu_infiniop.hpp new file mode 100644 index 000000000..d887c5c7b --- /dev/null +++ b/src/infinicore/ops/swiglu/swiglu_infiniop.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "infinicore/tensor.hpp" + +namespace infinicore::op::swiglu_impl::infiniop { + +void *plan(Tensor c, const Tensor &a, const Tensor &b); +void run(void *planned_meta); +void cleanup(void **planned_meta_ptr); + +} // namespace infinicore::op::swiglu_impl::infiniop diff --git a/src/infinicore/ops/swiglu/swiglu_infiniops.cc b/src/infinicore/ops/swiglu/swiglu_infiniops.cc index feacfd3fc..2ff99d96c 100644 --- a/src/infinicore/ops/swiglu/swiglu_infiniops.cc +++ b/src/infinicore/ops/swiglu/swiglu_infiniops.cc @@ -2,88 +2,15 @@ #ifdef ENABLE_INFINIOPS_API #include "../infiniops_impl.hpp" - -#include "base/copy.h" -#include "base/silu_and_mul.h" +#include "swiglu_infiniop.hpp" namespace infinicore::op::swiglu_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta c, a, b, packed, gate, up; - graph::GraphTensor c_tensor, a_tensor, b_tensor, packed_tensor, gate_tensor, up_tensor; -}; - -} // namespace - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - INFINICORE_ASSERT(c->shape() == a->shape()); - INFINICORE_ASSERT(a->shape() == b->shape()); - INFINICORE_ASSERT(c->dtype() == a->dtype()); - INFINICORE_ASSERT(a->dtype() == b->dtype()); - INFINICORE_ASSERT(!a->shape().empty()); - - auto packed_shape = a->shape(); - packed_shape.back() *= 2; - auto packed = Tensor::empty(packed_shape, a->dtype(), a->device()); - auto hidden_size = a->size(a->ndim() - 1); - auto gate = packed->narrow({{packed->ndim() - 1, 0, hidden_size}}); - auto up = packed->narrow({{packed->ndim() - 1, hidden_size, hidden_size}}); - - return new PlannedMeta{ - TensorMeta(c), - TensorMeta(a), - TensorMeta(b), - TensorMeta(packed), - TensorMeta(gate), - TensorMeta(up), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(packed), - graph::GraphTensor(gate), - graph::GraphTensor(up)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::Copy::Call( - handle, - config, - planned->b.tensor(planned->b_tensor), - false, - planned->gate.tensor(planned->gate_tensor)); - infini::ops::Copy::Call( - handle, - config, - planned->a.tensor(planned->a_tensor), - false, - planned->up.tensor(planned->up_tensor)); - infini::ops::SiluAndMul::Call( - handle, - config, - planned->packed.tensor(planned->packed_tensor), - planned->c.tensor(planned->c_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::cleanup_dispatcher(), &cleanup); + // Reuse the shared InfiniOp elementwise implementation in InfiniOps-enabled builds. + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::plan_dispatcher(), &infiniop::plan); + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::run_dispatcher(), &infiniop::run); + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::cleanup_dispatcher(), &infiniop::cleanup); return true; }();