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

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

* test(cpu): align GDN kernels with unified suite

* test(cpu): expose GDN kernel case matrices

* fix(test): avoid cross-running Android test binaries

Author: Aharrypotter

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);