fix: llama long context precision

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

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

src/infinicore/ops/swiglu/swiglu_infiniop.cc

@@ -1,3 +1,4 @@
+#include "swiglu_infiniop.hpp"
 #include "infinicore/ops/swiglu.hpp"
 
 #include "../infiniop_impl.hpp"

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

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<PlannedMeta *>(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<PlannedMeta **>(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;
 }();