refactor(qwen3.5): register stateful GDN runtime ops

mllm/backends/cpu/CPUBackend.cpp

@@ -25,6 +25,7 @@
 #include "mllm/backends/cpu/ops/FlashAttn2WithSinkAndSwaOp.hpp"
 #include "mllm/backends/cpu/ops/GELUOp.hpp"
 #include "mllm/backends/cpu/ops/GatherOp.hpp"
+#include "mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp"
 #include "mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp"
 #include "mllm/backends/cpu/ops/InterpolateOp.hpp"
 #include "mllm/backends/cpu/ops/LayerNorm2DOp.hpp"
@@ -86,7 +87,7 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) {
                CPUConv2DOpFactory, CPULayerNorm2DOpFactory, CPUInterpolateOpFactory, CPUPadOpFactory, CPUMaskedScatterOpFactory,
                CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory,
                CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory,
-               CPUGatherOpFactory, CPUCausalDepthwiseConv1DOpFactory,
+               CPUGatherOpFactory, CPUCausalDepthwiseConv1DOpFactory, CPUGatedDeltaRuleOpFactory,
                CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory>();
 }
 

mllm/backends/cpu/ops/GatedDeltaRuleOp.cpp

@@ -0,0 +1,28 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#include "mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp"
+
+#include <cstring>
+
+#include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp"
+#include "mllm/utils/Common.hpp"
+
+namespace mllm::cpu {
+
+CPUGatedDeltaRuleOp::CPUGatedDeltaRuleOp(const aops::GatedDeltaRuleOpOptions& options) : aops::GatedDeltaRuleOp(options) {}
+
+void CPUGatedDeltaRuleOp::forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  for (const auto& input : inputs) { MLLM_RT_ASSERT(input.isContiguous()); }
+  const auto& q = inputs[0];
+  const auto& v = inputs[2];
+  auto& output = outputs[0];
+  auto& updated_state = outputs[1];
+  if (!options_.state_inplace) { std::memcpy(updated_state.ptr<float>(), inputs[7].ptr<float>(), inputs[7].bytes()); }
+  gdn::gatedDeltaRuleF32(inputs[0].ptr<float>(), inputs[1].ptr<float>(), inputs[2].ptr<float>(), inputs[3].ptr<float>(),
+                         inputs[4].ptr<float>(), inputs[5].ptr<float>(), inputs[6].ptr<float>(), updated_state.ptr<float>(),
+                         output.ptr<float>(), q.shape()[0], q.shape()[1], q.shape()[2], v.shape()[2], q.shape()[3],
+                         v.shape()[3], options_.getThreads());
+}
+
+}  // namespace mllm::cpu

mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp

@@ -0,0 +1,23 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "mllm/core/aops/GatedDeltaRuleOp.hpp"
+
+namespace mllm::cpu {
+
+class CPUGatedDeltaRuleOp final : public aops::GatedDeltaRuleOp {
+ public:
+  explicit CPUGatedDeltaRuleOp(const aops::GatedDeltaRuleOpOptions& options);
+  void forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+};
+
+class CPUGatedDeltaRuleOpFactory : public TypedOpFactory<OpTypes::kGatedDeltaRule, aops::GatedDeltaRuleOpOptions> {
+ protected:
+  std::shared_ptr<BaseOp> createOpImpl(const aops::GatedDeltaRuleOpOptions& options) override {
+    return std::make_shared<CPUGatedDeltaRuleOp>(options);
+  }
+};
+
+}  // namespace mllm::cpu

mllm/compile/ir/GeneratedRTTIKind.hpp

@@ -44,6 +44,7 @@ enum NodeKind : uint32_t {
   RK_Op_LinalgIROp_CausalDepthwiseConv1DOp,
   RK_Op_LinalgIROp_GroupedQueryAttentionOp,
   RK_Op_LinalgIROp_ParallelLinearOp,
+  RK_Op_LinalgIROp_GatedDeltaRuleOp,
   RK_Op_LinalgIROp_RepeatOp,
   RK_Op_LinalgIROp_PermuteOp,
   RK_Op_LinalgIROp_Conv1DOp,

mllm/compile/ir/NodeRTTIClassOfImpl.hpp

@@ -104,6 +104,9 @@ struct NodeRTTIClassOfImpl {
 #define RTTI_RK_OP_LINALGIROP_PARALLELLINEAROP_IMPL(v) \
   return (v)->getKind() >= RK_Op_LinalgIROp_ParallelLinearOp && (v)->getKind() <= RK_Op_LinalgIROp_ParallelLinearOp
 
+#define RTTI_RK_OP_LINALGIROP_GATEDDELTARULEOP_IMPL(v) \
+  return (v)->getKind() >= RK_Op_LinalgIROp_GatedDeltaRuleOp && (v)->getKind() <= RK_Op_LinalgIROp_GatedDeltaRuleOp
+
 #define RTTI_RK_OP_LINALGIROP_REPEATOP_IMPL(v) \
   return (v)->getKind() >= RK_Op_LinalgIROp_RepeatOp && (v)->getKind() <= RK_Op_LinalgIROp_RepeatOp
 

mllm/compile/ir/linalg/Op.cpp

@@ -69,6 +69,7 @@ LINALG_AOPS_DECL(OpTypes::kFlashAttention2, FlashAttention2Op);
 LINALG_AOPS_DECL(OpTypes::kCausalDepthwiseConv1D, CausalDepthwiseConv1DOp);
 LINALG_AOPS_DECL(OpTypes::kGroupedQueryAttention, GroupedQueryAttentionOp);
 LINALG_AOPS_DECL(OpTypes::kParallelLinear, ParallelLinearOp);
+LINALG_AOPS_DECL(OpTypes::kGatedDeltaRule, GatedDeltaRuleOp);
 LINALG_AOPS_DECL(OpTypes::kRepeat, RepeatOp);
 LINALG_AOPS_DECL(OpTypes::kPermute, PermuteOp);
 

mllm/compile/ir/linalg/Op.hpp

@@ -38,6 +38,7 @@ class FlashAttention2Op;
 class CausalDepthwiseConv1DOp;
 class GroupedQueryAttentionOp;
 class ParallelLinearOp;
+class GatedDeltaRuleOp;
 class RepeatOp;
 class PermuteOp;
 class Conv1DOp;
@@ -203,6 +204,7 @@ LINALG_AOPS_DEFINE(FlashAttention2Op, FLASHATTENTION2OP);
 LINALG_AOPS_DEFINE(CausalDepthwiseConv1DOp, CAUSALDEPTHWISECONV1DOP);
 LINALG_AOPS_DEFINE(GroupedQueryAttentionOp, GROUPEDQUERYATTENTIONOP);
 LINALG_AOPS_DEFINE(ParallelLinearOp, PARALLELLINEAROP);
+LINALG_AOPS_DEFINE(GatedDeltaRuleOp, GATEDDELTARULEOP);
 LINALG_AOPS_DEFINE(RepeatOp, REPEATOP);
 LINALG_AOPS_DEFINE(PermuteOp, PERMUTEOP);
 

mllm/compile/ir/rtti_kind_gen.py

@@ -250,6 +250,7 @@ def define_lianlg_ir(ir: dict):
     op.derive(Cls("CausalDepthwiseConv1DOp"))
     op.derive(Cls("GroupedQueryAttentionOp"))
     op.derive(Cls("ParallelLinearOp"))
+    op.derive(Cls("GatedDeltaRuleOp"))
     op.derive(Cls("RepeatOp"))
     op.derive(Cls("PermuteOp"))
     op.derive(Cls("Conv1DOp"))

mllm/compile/jit/binary/LinalgIRSerialization.cpp

@@ -14,6 +14,7 @@
 #include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp"
 #include "mllm/core/aops/GroupedQueryAttentionOp.hpp"
 #include "mllm/core/aops/ParallelLinearOp.hpp"
+#include "mllm/core/aops/GatedDeltaRuleOp.hpp"
 #include "mllm/core/aops/KVCacheOp.hpp"
 #include "mllm/core/aops/MultimodalRoPEOp.hpp"
 #include "mllm/core/aops/VisionRoPEOp.hpp"
@@ -75,6 +76,7 @@ nlohmann::json dumpLinalgIROptions(const ir::linalg::LinalgIROp::ptr_t& op) {
     CASE(CausalDepthwiseConv1D)
     CASE(GroupedQueryAttention)
     CASE(ParallelLinear)
+    CASE(GatedDeltaRule)
     CASE(Repeat)
     CASE(Permute)
     CASE(Conv1D)
@@ -155,6 +157,11 @@ nlohmann::json dumpCausalDepthwiseConv1DOpIROptions(const ir::linalg::LinalgIROp
           {"accumulation_order", aops::causalDepthwiseConv1DAccumulationOrder2Str(options.accumulation_order)}};
 }
 
+nlohmann::json dumpGatedDeltaRuleOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) {
+  const auto options = static_cast<aops::GatedDeltaRuleOp*>(op->getAOp())->options();
+  return {{"state_inplace", options.state_inplace}};
+}
+
 nlohmann::json dumpGroupedQueryAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) {
   const auto options = static_cast<aops::GroupedQueryAttentionOp*>(op->getAOp())->options();
   return {{"implementation", aops::groupedQueryAttentionImplementation2Str(options.implementation)}};

mllm/compile/jit/binary/LinalgIRSerialization.hpp

@@ -40,6 +40,7 @@ nlohmann::json dumpFlashAttention2OpIROptions(const ir::linalg::LinalgIROp::ptr_
 nlohmann::json dumpCausalDepthwiseConv1DOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
 nlohmann::json dumpGroupedQueryAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
 nlohmann::json dumpParallelLinearOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
+nlohmann::json dumpGatedDeltaRuleOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
 nlohmann::json dumpRepeatOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
 nlohmann::json dumpPermuteOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);
 nlohmann::json dumpConv1DOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op);

mllm/compile/jit/interpreter/AopsFromJson.cpp

@@ -25,6 +25,7 @@
 #include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp"
 #include "mllm/core/aops/GroupedQueryAttentionOp.hpp"
 #include "mllm/core/aops/ParallelLinearOp.hpp"
+#include "mllm/core/aops/GatedDeltaRuleOp.hpp"
 #include "mllm/core/aops/RepeatOp.hpp"
 #include "mllm/core/aops/PermuteOp.hpp"
 #include "mllm/core/aops/GELUOp.hpp"
@@ -114,6 +115,8 @@ BaseOp::ptr_t aopsFromJson(const nlohmann::json& json) {
       return __groupedQueryAttentionFromJson(json);
     } else if (op_type == "ParallelLinear") {
       return __parallelLinearFromJson(json);
+    } else if (op_type == "GatedDeltaRule") {
+      return __gatedDeltaRuleFromJson(json);
     } else if (op_type == "Repeat") {
       return __repeatFromJson(json);
     } else if (op_type == "Permute") {
@@ -677,6 +680,16 @@ BaseOp::ptr_t __parallelLinearFromJson(const nlohmann::json& json) {
   return Context::instance().getBackend(backend)->createOp(OpTypes::kParallelLinear, options);
 }
 
+BaseOp::ptr_t __gatedDeltaRuleFromJson(const nlohmann::json& json) {
+  aops::GatedDeltaRuleOpOptions options;
+  if (json.contains("op_options") && json["op_options"].contains("state_inplace")) {
+    options.state_inplace = json["op_options"]["state_inplace"];
+  }
+  DeviceTypes backend = DeviceTypes::kCPU;
+  if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); }
+  return Context::instance().getBackend(backend)->createOp(OpTypes::kGatedDeltaRule, options);
+}
+
 BaseOp::ptr_t __repeatFromJson(const nlohmann::json& json) {
   aops::RepeatOpOptions options;
 

mllm/compile/jit/interpreter/AopsFromJson.hpp

@@ -38,6 +38,7 @@ BaseOp::ptr_t __groupedQueryAttentionDecodeFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __causalDepthwiseConv1DFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __groupedQueryAttentionFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __parallelLinearFromJson(const nlohmann::json& json);
+BaseOp::ptr_t __gatedDeltaRuleFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __repeatFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __permuteFromJson(const nlohmann::json& json);
 BaseOp::ptr_t __conv2dFromJson(const nlohmann::json& json);

mllm/core/OpTypes.hpp

@@ -105,6 +105,7 @@ enum class OpTypes : int32_t {
   kCausalDepthwiseConv1D = 77,
   kGroupedQueryAttention = 78,
   kParallelLinear = 79,
+  kGatedDeltaRule = 80,
 
   // Dynamic Op Start for user to register there own ops.
   kDynamicOp_Start = 4096,
@@ -193,6 +194,7 @@ inline std::string optype2Str(OpTypes type) {
     case OpTypes::kCausalDepthwiseConv1D: return "CausalDepthwiseConv1D";
     case OpTypes::kGroupedQueryAttention: return "GroupedQueryAttention";
     case OpTypes::kParallelLinear: return "ParallelLinear";
+    case OpTypes::kGatedDeltaRule: return "GatedDeltaRule";
     case OpTypes::kDynamicOp_Start: return "DynamicOp_Start";
     case OpTypes::kOpType_End: return "OpType_End";
     default: return "Unknown";

mllm/core/aops/GatedDeltaRuleOp.cpp

@@ -0,0 +1,74 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#include "mllm/core/aops/GatedDeltaRuleOp.hpp"
+
+#include "mllm/compile/ir/linalg/Op.hpp"
+#include "mllm/core/Tensor.hpp"
+#include "mllm/utils/Common.hpp"
+
+namespace mllm::aops {
+
+GatedDeltaRuleOp::GatedDeltaRuleOp(const GatedDeltaRuleOpOptions& options)
+    : BaseOp(OpTypes::kGatedDeltaRule), options_(options) {}
+
+void GatedDeltaRuleOp::load(const ParameterFile::ptr_t& ploader) { MLLM_EMPTY_SCOPE; }
+
+void GatedDeltaRuleOp::trace(void* trace_context, const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  auto* ir_ctx = static_cast<ir::IRContext*>(trace_context);
+  const auto input_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, inputs);
+  const auto output_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, outputs);
+  ir_ctx->create<ir::linalg::GatedDeltaRuleOp>(shared_from_this(), input_irs, output_irs);
+}
+
+void GatedDeltaRuleOp::forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  NYI("GatedDeltaRuleOp::forward not implemented in aops base.");
+}
+
+void GatedDeltaRuleOp::reshape(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  MLLM_RT_ASSERT_EQ(inputs.size(), 8);
+  const auto& q = inputs[0];
+  const auto& k = inputs[1];
+  const auto& v = inputs[2];
+  const auto& a = inputs[3];
+  const auto& b = inputs[4];
+  const auto& a_log = inputs[5];
+  const auto& dt_bias = inputs[6];
+  const auto& state = inputs[7];
+
+  MLLM_RT_ASSERT_EQ(q.rank(), 4);
+  MLLM_RT_ASSERT_EQ(k.rank(), 4);
+  MLLM_RT_ASSERT_EQ(v.rank(), 4);
+  const int32_t batch = q.shape()[0];
+  const int32_t sequence = q.shape()[1];
+  const int32_t key_heads = q.shape()[2];
+  const int32_t key_dim = q.shape()[3];
+  const int32_t value_heads = v.shape()[2];
+  const int32_t value_dim = v.shape()[3];
+  MLLM_RT_ASSERT_EQ(k.shape(), q.shape());
+  MLLM_RT_ASSERT_EQ(v.shape()[0], batch);
+  MLLM_RT_ASSERT_EQ(v.shape()[1], sequence);
+  MLLM_RT_ASSERT(key_heads > 0 && value_heads > 0 && value_heads % key_heads == 0);
+  MLLM_RT_ASSERT_EQ(a.shape(), (Tensor::shape_t{batch, sequence, value_heads}));
+  MLLM_RT_ASSERT_EQ(b.shape(), (Tensor::shape_t{batch, sequence, value_heads}));
+  MLLM_RT_ASSERT_EQ(a_log.numel(), static_cast<std::size_t>(value_heads));
+  MLLM_RT_ASSERT_EQ(dt_bias.numel(), static_cast<std::size_t>(value_heads));
+  MLLM_RT_ASSERT_EQ(state.shape(), (Tensor::shape_t{batch, value_heads, value_dim, key_dim}));
+  for (const auto& tensor : inputs) {
+    MLLM_RT_ASSERT_EQ(tensor.dtype(), kFloat32);
+    MLLM_RT_ASSERT_EQ(tensor.device(), q.device());
+  }
+
+  outputs.emplace_back(Tensor::empty({batch, sequence, value_heads, value_dim}, q.dtype(), q.device()));
+  outputs.emplace_back(options_.state_inplace ? state : Tensor::empty(state.shape(), state.dtype(), state.device()));
+}
+
+void GatedDeltaRuleOp::setup(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  if (options_.state_inplace) {
+    outputs[0].alloc();
+  } else {
+    BaseOp::setup(inputs, outputs);
+  }
+}
+
+}  // namespace mllm::aops

mllm/core/aops/GatedDeltaRuleOp.hpp

@@ -0,0 +1,35 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "mllm/core/BaseOp.hpp"
+#include "mllm/core/ParameterFile.hpp"
+
+namespace mllm::aops {
+
+struct GatedDeltaRuleOpOptions : public BaseOpOptions<GatedDeltaRuleOpOptions> {
+  bool state_inplace = false;
+};
+
+// Stateful grouped-head gated delta recurrence.
+// Inputs: q/k [B, S, Hk, Dk], v [B, S, Hv, Dv], a/b [B, S, Hv],
+// A_log/dt_bias [Hv], state [B, Hv, Dv, Dk].
+// Outputs: output [B, S, Hv, Dv], updated_state [B, Hv, Dv, Dk].
+class GatedDeltaRuleOp : public BaseOp {
+ public:
+  explicit GatedDeltaRuleOp(const GatedDeltaRuleOpOptions& options);
+
+  void load(const ParameterFile::ptr_t& ploader) override;
+  void trace(void* trace_context, const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+  void forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+  void reshape(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+  void setup(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+
+  [[nodiscard]] const GatedDeltaRuleOpOptions& options() const { return options_; }
+
+ protected:
+  GatedDeltaRuleOpOptions options_;
+};
+
+}  // namespace mllm::aops