feat(infiniop): add lightning_attention op (CPU + NVIDIA) for MiniMax

include/infinicore/ops.hpp

@@ -45,6 +45,7 @@
 #include "ops/hardswish.hpp"
 #include "ops/hardtanh.hpp"
 #include "ops/kimi_delta_attention.hpp"
+#include "ops/lightning_attention.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/layer_norm.hpp"
 #include "ops/linear.hpp"
@@ -114,3 +115,4 @@
 #include "ops/w4a8_group_gemm.hpp"
 #include "ops/w8a8_group_gemm.hpp"
 #endif
+

include/infinicore/ops/lightning_attention.hpp

@@ -0,0 +1,41 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(LightningAttention,
+                          Tensor,
+                          Tensor,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &);
+
+// Indexed-pool lightning attention (MiniMax-01 style).
+// Returns out [B, T, H, D] and updates `initial_state` in place at
+// `final_state_indices` rows.
+__export Tensor lightning_attention(const Tensor &q,
+                                    const Tensor &k,
+                                    const Tensor &v,
+                                    const Tensor &slope,
+                                    Tensor initial_state,
+                                    const Tensor &initial_state_indices,
+                                    const Tensor &final_state_indices);
+
+__export void lightning_attention_(Tensor out,
+                                   Tensor initial_state,
+                                   const Tensor &q,
+                                   const Tensor &k,
+                                   const Tensor &v,
+                                   const Tensor &slope,
+                                   const Tensor &initial_state_indices,
+                                   const Tensor &final_state_indices);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -85,6 +85,7 @@
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/ldexp.h"
 #include "infiniop/ops/lerp.h"
+#include "infiniop/ops/lightning_attention.h"
 #include "infiniop/ops/linear_mxfp4.h"
 #include "infiniop/ops/log10.h"
 #include "infiniop/ops/log1p.h"
@@ -169,3 +170,5 @@
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
 #endif // __INFINIOP_API_H__
+
+

include/infiniop/ops/lightning_attention.h

@@ -0,0 +1,60 @@
+#ifndef __INFINIOP_LIGHTNING_ATTENTION_API_H__
+#define __INFINIOP_LIGHTNING_ATTENTION_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLightningAttentionDescriptor_t;
+
+// Lightning attention (MiniMax-01 style, ALiBi-style per-head decay) with an
+// indexed recurrent-state pool.
+//
+// Recurrence (per head h, per token t):
+//   S  = ratio[h] * S + k_t^T v_t          (ratio[h] = exp(-slope[h]))
+//   o_t = q_t @ S
+// i.e. the state is updated *before* the output is read, so each output token
+// attends to itself with weight 1 (no decay within the same position).
+//
+// Tensor layouts:
+//   out                  [B, T, H, D]      (last dim contiguous)
+//   initial_state (pool) [pool_size, H, D, D]
+//   q/k/v                [B, T, H, D]      (last dim contiguous)
+//   slope                [H]               (fp32)
+//   initial/final_state_indices [B]        (int32 or int64)
+//
+// Indexed-pool mode only: for each request b the op reads the state row
+// `initial_state[initial_state_indices[b]]` and writes the final state in place
+// to `initial_state[final_state_indices[b]]`.
+__INFINI_C __export infiniStatus_t infiniopCreateLightningAttentionDescriptor(
+    infiniopHandle_t handle,
+    infiniopLightningAttentionDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t initial_state_desc,
+    infiniopTensorDescriptor_t q_desc,
+    infiniopTensorDescriptor_t k_desc,
+    infiniopTensorDescriptor_t v_desc,
+    infiniopTensorDescriptor_t slope_desc,
+    infiniopTensorDescriptor_t initial_state_indices_desc,
+    infiniopTensorDescriptor_t final_state_indices_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetLightningAttentionWorkspaceSize(
+    infiniopLightningAttentionDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopLightningAttention(
+    infiniopLightningAttentionDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    void *initial_state,
+    const void *q,
+    const void *k,
+    const void *v,
+    const void *slope,
+    const void *initial_state_indices,
+    const void *final_state_indices,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyLightningAttentionDescriptor(
+    infiniopLightningAttentionDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -121,6 +121,7 @@
 from infinicore.ops.kv_caching import kv_caching
 from infinicore.ops.ldexp import ldexp
 from infinicore.ops.lerp import lerp
+from infinicore.ops.lightning_attention import lightning_attention
 from infinicore.ops.logaddexp import logaddexp
 from infinicore.ops.logaddexp2 import logaddexp2
 from infinicore.ops.logcumsumexp import logcumsumexp
@@ -388,3 +389,5 @@
     "w4a8_group_gemm_",
     "w8a8_group_gemm_",
 ]
+
+

python/infinicore/ops/lightning_attention.py

@@ -0,0 +1,21 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def lightning_attention(q, k, v, slope, initial_state, initial_state_indices, final_state_indices):
+    """Indexed-pool lightning attention (MiniMax-01 style).
+
+    Returns out [B, T, H, D]; `initial_state` is updated in place at the
+    `final_state_indices` rows.
+    """
+    return Tensor(
+        _infinicore.lightning_attention(
+            q._underlying,
+            k._underlying,
+            v._underlying,
+            slope._underlying,
+            initial_state._underlying,
+            initial_state_indices._underlying,
+            final_state_indices._underlying,
+        )
+    )

src/infinicore/ops/lightning_attention/lightning_attention.cc

@@ -0,0 +1,90 @@
+#include "infinicore/ops/lightning_attention.hpp"
+#include "../../utils.hpp"
+
+#include <stdexcept>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LightningAttention);
+
+LightningAttention::LightningAttention(Tensor out,
+                                       Tensor initial_state,
+                                       const Tensor &q,
+                                       const Tensor &k,
+                                       const Tensor &v,
+                                       const Tensor &slope,
+                                       const Tensor &initial_state_indices,
+                                       const Tensor &final_state_indices) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, slope);
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices);
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices);
+    INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(),
+                                 out,
+                                 initial_state,
+                                 q,
+                                 k,
+                                 v,
+                                 slope,
+                                 initial_state_indices,
+                                 final_state_indices);
+}
+
+void LightningAttention::execute(Tensor out,
+                                 Tensor initial_state,
+                                 const Tensor &q,
+                                 const Tensor &k,
+                                 const Tensor &v,
+                                 const Tensor &slope,
+                                 const Tensor &initial_state_indices,
+                                 const Tensor &final_state_indices) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(LightningAttention,
+                                      out,
+                                      initial_state,
+                                      q,
+                                      k,
+                                      v,
+                                      slope,
+                                      initial_state_indices,
+                                      final_state_indices);
+}
+
+static Tensor ensure_4d_sequence_tensor(const Tensor &x, const char *name) {
+    if (x->shape().size() == 4) {
+        return x;
+    }
+    if (x->shape().size() == 3) {
+        return x->unsqueeze(1);
+    }
+    throw std::runtime_error(std::string("lightning_attention expects ") + name + " with shape [B, T, H, D] or [B, H, D]");
+}
+
+Tensor lightning_attention(const Tensor &q,
+                           const Tensor &k,
+                           const Tensor &v,
+                           const Tensor &slope,
+                           Tensor initial_state,
+                           const Tensor &initial_state_indices,
+                           const Tensor &final_state_indices) {
+    Tensor q4 = ensure_4d_sequence_tensor(q, "q");
+    Tensor k4 = ensure_4d_sequence_tensor(k, "k");
+    Tensor v4 = ensure_4d_sequence_tensor(v, "v");
+    auto out = Tensor::empty(v4->shape(), v4->dtype(), v4->device());
+    lightning_attention_(out, initial_state, q4, k4, v4, slope, initial_state_indices, final_state_indices);
+    return out;
+}
+
+void lightning_attention_(Tensor out,
+                          Tensor initial_state,
+                          const Tensor &q,
+                          const Tensor &k,
+                          const Tensor &v,
+                          const Tensor &slope,
+                          const Tensor &initial_state_indices,
+                          const Tensor &final_state_indices) {
+    Tensor q4 = ensure_4d_sequence_tensor(q, "q");
+    Tensor k4 = ensure_4d_sequence_tensor(k, "k");
+    Tensor v4 = ensure_4d_sequence_tensor(v, "v");
+    LightningAttention::execute(out, initial_state, q4, k4, v4, slope, initial_state_indices, final_state_indices);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/lightning_attention/lightning_attention_infiniop.cc

@@ -0,0 +1,85 @@
+#include "infinicore/ops/lightning_attention.hpp"
+
+#include "../infiniop_impl.hpp"
+
+namespace infinicore::op::lightning_attention_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, LightningAttention, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace, out, initial_state, q, k, v, slope;
+    graph::GraphTensor initial_state_indices;
+    graph::GraphTensor final_state_indices;
+};
+
+void *plan(Tensor out,
+           Tensor initial_state,
+           const Tensor &q,
+           const Tensor &k,
+           const Tensor &v,
+           const Tensor &slope,
+           const Tensor &initial_state_indices,
+           const Tensor &final_state_indices) {
+    size_t seed = hash_combine(out,
+                               initial_state,
+                               q,
+                               k,
+                               v,
+                               slope,
+                               initial_state_indices,
+                               final_state_indices);
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, LightningAttention,
+        seed,
+        out->desc(),
+        initial_state->desc(),
+        q->desc(),
+        k->desc(),
+        v->desc(),
+        slope->desc(),
+        initial_state_indices->desc(),
+        final_state_indices->desc());
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, LightningAttention, descriptor);
+
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(out),
+        graph::GraphTensor(initial_state),
+        graph::GraphTensor(q),
+        graph::GraphTensor(k),
+        graph::GraphTensor(v),
+        graph::GraphTensor(slope),
+        graph::GraphTensor(initial_state_indices),
+        graph::GraphTensor(final_state_indices)};
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+
+    INFINICORE_CHECK_ERROR(infiniopLightningAttention(
+        planned->descriptor->desc,
+        planned->workspace->data(),
+        planned->workspace->numel(),
+        planned->out->data(),
+        planned->initial_state->data(),
+        planned->q->data(),
+        planned->k->data(),
+        planned->v->data(),
+        planned->slope->data(),
+        planned->initial_state_indices->data(),
+        planned->final_state_indices->data(),
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(LightningAttention, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::lightning_attention_impl::infiniop

src/infinicore/pybind11/ops.hpp

@@ -77,6 +77,7 @@
 #include "ops/ldexp.hpp"
 #include "ops/lerp.hpp"
 #include "ops/linear.hpp"
+#include "ops/lightning_attention.hpp"
 #include "ops/linear_mxfp4.hpp"
 #include "ops/linear_w8a8i8.hpp"
 #include "ops/log_softmax.hpp"
@@ -229,6 +230,7 @@ inline void bind(py::module &m) {
     bind_hinge_embedding_loss(m);
     bind_kv_caching(m);
     bind_kimi_delta_attention(m);
+    bind_lightning_attention(m);
     bind_fmod(m);
     bind_fp8_indexer_logits(m);
     bind_fp8_indexer_quant(m);
@@ -345,3 +347,5 @@ inline void bind(py::module &m) {
 }
 
 } // namespace infinicore::ops
+
+

src/infinicore/pybind11/ops/lightning_attention.hpp

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/lightning_attention.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_lightning_attention(py::module &m) {
+    m.def("lightning_attention",
+          &op::lightning_attention,
+          py::arg("q"),
+          py::arg("k"),
+          py::arg("v"),
+          py::arg("slope"),
+          py::arg("initial_state"),
+          py::arg("initial_state_indices"),
+          py::arg("final_state_indices"),
+          R"doc(Indexed-pool lightning attention (MiniMax-01 style).
+Returns out [B, T, H, D]; updates `initial_state` in place at final_state_indices rows.)doc");
+}
+
+} // namespace infinicore::ops