From 7bd764b84555ec820a3b23d3cc3072e620c991ea Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 2 Sep 2026 00:45:29 +0800 Subject: [PATCH 1/4] refactor(qwen3.5): register stateful GDN runtime ops --- mllm/backends/cpu/CPUBackend.cpp | 3 +- mllm/backends/cpu/ops/GatedDeltaRuleOp.cpp | 28 +++ mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp | 23 ++ mllm/compile/ir/GeneratedRTTIKind.hpp | 1 + mllm/compile/ir/NodeRTTIClassOfImpl.hpp | 3 + mllm/compile/ir/linalg/Op.cpp | 1 + mllm/compile/ir/linalg/Op.hpp | 2 + mllm/compile/ir/rtti_kind_gen.py | 1 + .../jit/binary/LinalgIRSerialization.cpp | 7 + .../jit/binary/LinalgIRSerialization.hpp | 1 + mllm/compile/jit/interpreter/AopsFromJson.cpp | 13 ++ mllm/compile/jit/interpreter/AopsFromJson.hpp | 1 + mllm/core/OpTypes.hpp | 2 + mllm/core/aops/GatedDeltaRuleOp.cpp | 74 ++++++ mllm/core/aops/GatedDeltaRuleOp.hpp | 35 +++ mllm/models/qwen3_5/modeling_qwen3_5.hpp | 40 ++-- mllm/nn/Functional.cpp | 9 + mllm/nn/Functional.hpp | 5 + mllm/nn/Nn.hpp | 2 + mllm/nn/layers/GatedDeltaRule.cpp | 15 ++ mllm/nn/layers/GatedDeltaRule.hpp | 20 ++ tests/CMakeLists.txt | 1 + tests/cpu/CMakeLists.txt | 35 +-- ...alDepthwiseConvCurrentFirstKernelTest.cpp} | 65 +++--- ...NTest.cpp => GatedDeltaRuleKernelTest.cpp} | 79 ++----- tests/models/CMakeLists.txt | 1 + tests/models/qwen3_5/CMakeLists.txt | 18 ++ .../qwen3_5}/Qwen35ConfigTest.cpp | 0 .../qwen3_5}/Qwen35MultimodalTest.cpp | 0 .../qwen3_5}/Qwen35TokenizerTest.cpp | 0 tests/nn/CMakeLists.txt | 5 + tests/nn/GatedDeltaRuleTest.cpp | 217 ++++++++++++++++++ 32 files changed, 562 insertions(+), 145 deletions(-) create mode 100644 mllm/backends/cpu/ops/GatedDeltaRuleOp.cpp create mode 100644 mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp create mode 100644 mllm/core/aops/GatedDeltaRuleOp.cpp create mode 100644 mllm/core/aops/GatedDeltaRuleOp.hpp create mode 100644 mllm/nn/layers/GatedDeltaRule.cpp create mode 100644 mllm/nn/layers/GatedDeltaRule.hpp rename tests/cpu/{Qwen35GDNConvTest.cpp => CausalDepthwiseConvCurrentFirstKernelTest.cpp} (82%) rename tests/cpu/{Qwen35GDNTest.cpp => GatedDeltaRuleKernelTest.cpp} (81%) create mode 100644 tests/models/CMakeLists.txt create mode 100644 tests/models/qwen3_5/CMakeLists.txt rename tests/{cpu => models/qwen3_5}/Qwen35ConfigTest.cpp (100%) rename tests/{cpu => models/qwen3_5}/Qwen35MultimodalTest.cpp (100%) rename tests/{cpu => models/qwen3_5}/Qwen35TokenizerTest.cpp (100%) create mode 100644 tests/nn/GatedDeltaRuleTest.cpp diff --git a/mllm/backends/cpu/CPUBackend.cpp b/mllm/backends/cpu/CPUBackend.cpp index 5b59e7c91..190f61109 100644 --- a/mllm/backends/cpu/CPUBackend.cpp +++ b/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>(); } diff --git a/mllm/backends/cpu/ops/GatedDeltaRuleOp.cpp b/mllm/backends/cpu/ops/GatedDeltaRuleOp.cpp new file mode 100644 index 000000000..e546f354c --- /dev/null +++ b/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 + +#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& inputs, std::vector& 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(), inputs[7].ptr(), inputs[7].bytes()); } + gdn::gatedDeltaRuleF32(inputs[0].ptr(), inputs[1].ptr(), inputs[2].ptr(), inputs[3].ptr(), + inputs[4].ptr(), inputs[5].ptr(), inputs[6].ptr(), updated_state.ptr(), + output.ptr(), q.shape()[0], q.shape()[1], q.shape()[2], v.shape()[2], q.shape()[3], + v.shape()[3], options_.getThreads()); +} + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp b/mllm/backends/cpu/ops/GatedDeltaRuleOp.hpp new file mode 100644 index 000000000..3ce437d7c --- /dev/null +++ b/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& inputs, std::vector& outputs) override; +}; + +class CPUGatedDeltaRuleOpFactory : public TypedOpFactory { + protected: + std::shared_ptr createOpImpl(const aops::GatedDeltaRuleOpOptions& options) override { + return std::make_shared(options); + } +}; + +} // namespace mllm::cpu diff --git a/mllm/compile/ir/GeneratedRTTIKind.hpp b/mllm/compile/ir/GeneratedRTTIKind.hpp index 726ec8657..7114ef549 100644 --- a/mllm/compile/ir/GeneratedRTTIKind.hpp +++ b/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, diff --git a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp index 0a08e255b..252c3fe4d 100644 --- a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp +++ b/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 diff --git a/mllm/compile/ir/linalg/Op.cpp b/mllm/compile/ir/linalg/Op.cpp index a33cd581b..3b315e6f2 100644 --- a/mllm/compile/ir/linalg/Op.cpp +++ b/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); diff --git a/mllm/compile/ir/linalg/Op.hpp b/mllm/compile/ir/linalg/Op.hpp index e79765cc7..f95911cc9 100644 --- a/mllm/compile/ir/linalg/Op.hpp +++ b/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); diff --git a/mllm/compile/ir/rtti_kind_gen.py b/mllm/compile/ir/rtti_kind_gen.py index 03029b23b..eda139259 100644 --- a/mllm/compile/ir/rtti_kind_gen.py +++ b/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")) diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.cpp b/mllm/compile/jit/binary/LinalgIRSerialization.cpp index 307556400..2a8138853 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.cpp +++ b/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(op->getAOp())->options(); + return {{"state_inplace", options.state_inplace}}; +} + nlohmann::json dumpGroupedQueryAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { const auto options = static_cast(op->getAOp())->options(); return {{"implementation", aops::groupedQueryAttentionImplementation2Str(options.implementation)}}; diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.hpp b/mllm/compile/jit/binary/LinalgIRSerialization.hpp index a3c2901d3..ca3c32d48 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.hpp +++ b/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); diff --git a/mllm/compile/jit/interpreter/AopsFromJson.cpp b/mllm/compile/jit/interpreter/AopsFromJson.cpp index f6c5c7a78..f3ecdcbb4 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.cpp +++ b/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; diff --git a/mllm/compile/jit/interpreter/AopsFromJson.hpp b/mllm/compile/jit/interpreter/AopsFromJson.hpp index d93e1a478..87d6f2e71 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.hpp +++ b/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); diff --git a/mllm/core/OpTypes.hpp b/mllm/core/OpTypes.hpp index 0a3263003..76777fe0c 100644 --- a/mllm/core/OpTypes.hpp +++ b/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"; diff --git a/mllm/core/aops/GatedDeltaRuleOp.cpp b/mllm/core/aops/GatedDeltaRuleOp.cpp new file mode 100644 index 000000000..6a40e99e3 --- /dev/null +++ b/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& inputs, std::vector& outputs) { + auto* ir_ctx = static_cast(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(shared_from_this(), input_irs, output_irs); +} + +void GatedDeltaRuleOp::forward(const std::vector& inputs, std::vector& outputs) { + NYI("GatedDeltaRuleOp::forward not implemented in aops base."); +} + +void GatedDeltaRuleOp::reshape(const std::vector& inputs, std::vector& 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(value_heads)); + MLLM_RT_ASSERT_EQ(dt_bias.numel(), static_cast(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& inputs, std::vector& outputs) { + if (options_.state_inplace) { + outputs[0].alloc(); + } else { + BaseOp::setup(inputs, outputs); + } +} + +} // namespace mllm::aops diff --git a/mllm/core/aops/GatedDeltaRuleOp.hpp b/mllm/core/aops/GatedDeltaRuleOp.hpp new file mode 100644 index 000000000..6a833d25b --- /dev/null +++ b/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 { + 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& inputs, std::vector& outputs) override; + void forward(const std::vector& inputs, std::vector& outputs) override; + void reshape(const std::vector& inputs, std::vector& outputs) override; + void setup(const std::vector& inputs, std::vector& outputs) override; + + [[nodiscard]] const GatedDeltaRuleOpOptions& options() const { return options_; } + + protected: + GatedDeltaRuleOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/models/qwen3_5/modeling_qwen3_5.hpp b/mllm/models/qwen3_5/modeling_qwen3_5.hpp index 4314965a0..59da4b224 100644 --- a/mllm/models/qwen3_5/modeling_qwen3_5.hpp +++ b/mllm/models/qwen3_5/modeling_qwen3_5.hpp @@ -20,7 +20,6 @@ #include "mllm/models/qwen3_5/configuration_qwen3_5.hpp" #include "mllm/models/qwen3_5/modeling_qwen3_5_vision.hpp" #include "mllm/models/qwen3_5/multimodal_qwen3_5.hpp" -#include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" #include "mllm/utils/Enumerate.hpp" #include "mllm/models/ARGeneration.hpp" @@ -263,8 +262,9 @@ class Qwen3_5GDNLayer final : public nn::Module { nn::Linear in_proj_a_; // hidden → num_v_heads (decay gate) nn::Linear in_proj_b_; // hidden → num_v_heads (beta gate) - // Causal Conv1D for sequence mixing - nn::Conv1D conv1d_; + // Stateful sequence primitives and their checkpoint-owned parameters. + nn::CausalDepthwiseConv1D causal_conv_; + nn::GatedDeltaRule gated_delta_rule_; // Learnable parameters for gating nn::Param A_log_; // [num_v_heads] @@ -310,15 +310,15 @@ class Qwen3_5GDNLayer final : public nn::Module { in_proj_a_ = reg("in_proj_a", hidden_size_, num_v_heads_, false, cfg.linear_impl_type); in_proj_b_ = reg("in_proj_b", hidden_size_, num_v_heads_, false, cfg.linear_impl_type); - // Causal Conv1D: groups = channels (depthwise), no bias, padding = kernel-1 for causal int conv_channels = key_dim_ * 2 + value_dim_; - conv1d_ = reg("conv1d", conv_channels, conv_channels, conv_kernel_size_, - /*stride=*/1, /*padding=*/conv_kernel_size_ - 1, /*dilation=*/1, - /*groups=*/conv_channels, /*bias=*/false); + causal_conv_ = reg( + "conv1d", conv_channels, conv_kernel_size_, /*bias=*/false, /*state_inplace=*/true, + aops::CausalDepthwiseConv1DAccumulationOrder::kCurrentFirst); + gated_delta_rule_ = reg("gated_delta_rule", /*state_inplace=*/true); // Learnable gating parameters (loaded from weight file) - A_log_ = reg("A_log", getModuleName() + ".A_log"); - dt_bias_ = reg("dt_bias", getModuleName() + ".dt_bias"); + A_log_ = reg("A_log", getModuleName() + ".A_log", Tensor::shape_t{num_v_heads_}); + dt_bias_ = reg("dt_bias", getModuleName() + ".dt_bias", Tensor::shape_t{num_v_heads_}); // Gated RMSNorm — standard (NOT GemmaRMSNorm, no add_unit_offset) norm_ = reg("norm", cfg.rms_norm_eps, /*add_unit_offset=*/false); @@ -342,7 +342,6 @@ class Qwen3_5GDNLayer final : public nn::Module { int B = x.shape()[0]; int S = x.shape()[1]; int conv_dim = key_dim_ * 2 + value_dim_; - int K = conv_kernel_size_; // Lazy init recurrent + conv state. A batch-size change starts a new // independent sequence set and therefore cannot reuse the old states. @@ -358,17 +357,15 @@ class Qwen3_5GDNLayer final : public nn::Module { // intentionally stay in float32. Mobile quantized Linear implementations // dequantize to float32 before this numerically sensitive recurrence. auto mixed_qkv_pre = mixed_qkv.contiguous(); - auto conv_weight = conv1d_.weight(); - auto a_log = A_log_.weight(); - auto dt_bias = dt_bias_.weight(); - if (mixed_qkv_pre.dtype() != kFloat32 || conv_weight.dtype() != kFloat32 || a_proj.dtype() != kFloat32 - || b_proj.dtype() != kFloat32 || a_log.dtype() != kFloat32 || dt_bias.dtype() != kFloat32) { + auto a_log = A_log_(); + auto dt_bias = dt_bias_(); + if (mixed_qkv_pre.dtype() != kFloat32 || a_proj.dtype() != kFloat32 || b_proj.dtype() != kFloat32 + || a_log.dtype() != kFloat32 || dt_bias.dtype() != kFloat32) { throw std::invalid_argument("Qwen3.5 GDN currently requires float32 activations and recurrent parameters"); } - auto conv_out = Tensor::empty({B, S, conv_dim}, kFloat32, kCPU).alloc(); - ::mllm::cpu::gdn::depthwiseCausalConvF32(mixed_qkv_pre.ptr(), conv_weight.ptr(), conv_state_.ptr(), - conv_out.ptr(), B, S, conv_dim, K); + auto [conv_out, updated_conv_state] = causal_conv_(mixed_qkv_pre, conv_state_); + conv_state_ = updated_conv_state; // Qwen3.5 applies SiLU to the depthwise-convolution result before // splitting it into q, k and v. @@ -383,11 +380,8 @@ class Qwen3_5GDNLayer final : public nn::Module { a_proj = a_proj.contiguous(); b_proj = b_proj.contiguous(); - auto output = Tensor::empty({B, S, num_v_heads_, head_v_dim_}, kFloat32, kCPU).alloc(); - ::mllm::cpu::gdn::gatedDeltaRuleF32(q.ptr(), k.ptr(), v.ptr(), a_proj.ptr(), - b_proj.ptr(), a_log.ptr(), dt_bias.ptr(), - recurrent_state_.ptr(), output.ptr(), B, S, num_k_heads_, num_v_heads_, - head_k_dim_, head_v_dim_); + auto [output, updated_recurrent_state] = gated_delta_rule_(q, k, v, a_proj, b_proj, a_log, dt_bias, recurrent_state_); + recurrent_state_ = updated_recurrent_state; // [B, S, num_v_heads, head_v_dim] -> [B, S, value_dim] output = output.view({B, S, value_dim_}); diff --git a/mllm/nn/Functional.cpp b/mllm/nn/Functional.cpp index 7ad399c67..5f78c13eb 100644 --- a/mllm/nn/Functional.cpp +++ b/mllm/nn/Functional.cpp @@ -26,6 +26,7 @@ #include "mllm/core/aops/RadixAttnDiffDimOp.hpp" #include "mllm/core/aops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" #include "mllm/core/aops/WhereOp.hpp" +#include "mllm/core/aops/GatedDeltaRuleOp.hpp" #include "mllm/engine/Context.hpp" namespace mllm::nn::functional { @@ -99,6 +100,14 @@ Tensor groupedQueryAttention(const Tensor& query, const Tensor& key, const Tenso {query, key, value})[0]; } +std::array gatedDeltaRule(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& a, const Tensor& b, + const Tensor& a_log, const Tensor& dt_bias, const Tensor& state, bool state_inplace) { + auto outputs = Context::instance().buildOpAndSubmitTask(OpTypes::kGatedDeltaRule, + aops::GatedDeltaRuleOpOptions{.state_inplace = state_inplace}, + {q, k, v, a, b, a_log, dt_bias, state}); + return {outputs[0], outputs[1]}; +} + Tensor softmax(const Tensor& x, int32_t dim) { return Context::instance().buildOpAndSubmitTask(OpTypes::kSoftmax, aops::SoftmaxOpOptions{.axis = dim}, {x})[0]; } diff --git a/mllm/nn/Functional.hpp b/mllm/nn/Functional.hpp index c3bccfad1..e7c1e135c 100644 --- a/mllm/nn/Functional.hpp +++ b/mllm/nn/Functional.hpp @@ -13,6 +13,7 @@ #include "mllm/core/aops/PadOp.hpp" #include "mllm/core/aops/InterpolateOp.hpp" #include "mllm/core/aops/GroupedQueryAttentionOp.hpp" +#include "mllm/core/aops/GatedDeltaRuleOp.hpp" #include "mllm/core/aops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" #include "mllm/engine/Context.hpp" @@ -117,6 +118,10 @@ Tensor groupedQueryAttention( const Tensor& query, const Tensor& key, const Tensor& value, aops::GroupedQueryAttentionImplementation implementation = aops::GroupedQueryAttentionImplementation::kDirectStrided); +std::array gatedDeltaRule(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& a, const Tensor& b, + const Tensor& a_log, const Tensor& dt_bias, const Tensor& state, + bool state_inplace = false); + Tensor softmax(const Tensor& x, int32_t dim); Tensor log(const Tensor& x); diff --git a/mllm/nn/Nn.hpp b/mllm/nn/Nn.hpp index a1492d9ae..3b0ab764a 100644 --- a/mllm/nn/Nn.hpp +++ b/mllm/nn/Nn.hpp @@ -34,3 +34,5 @@ #include "mllm/nn/layers/PagedAttn.hpp" // IWYU pragma: export #include "mllm/nn/layers/RadixAttn.hpp" // IWYU pragma: export #include "mllm/nn/layers/LayerNorm2D.hpp" // IWYU pragma: export +#include "mllm/nn/layers/CausalDepthwiseConv1D.hpp" // IWYU pragma: export +#include "mllm/nn/layers/GatedDeltaRule.hpp" // IWYU pragma: export diff --git a/mllm/nn/layers/GatedDeltaRule.cpp b/mllm/nn/layers/GatedDeltaRule.cpp new file mode 100644 index 000000000..9631144b4 --- /dev/null +++ b/mllm/nn/layers/GatedDeltaRule.cpp @@ -0,0 +1,15 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/GatedDeltaRule.hpp" + +namespace mllm::nn { + +GatedDeltaRule::GatedDeltaRule() : Layer(OpTypes::kGatedDeltaRule, aops::GatedDeltaRuleOpOptions{}) {} + +GatedDeltaRule::GatedDeltaRule(const aops::GatedDeltaRuleOpOptions& options) : Layer(OpTypes::kGatedDeltaRule, options) {} + +GatedDeltaRule::GatedDeltaRule(bool state_inplace) + : Layer(OpTypes::kGatedDeltaRule, aops::GatedDeltaRuleOpOptions{.state_inplace = state_inplace}) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/GatedDeltaRule.hpp b/mllm/nn/layers/GatedDeltaRule.hpp new file mode 100644 index 000000000..5be179cfe --- /dev/null +++ b/mllm/nn/layers/GatedDeltaRule.hpp @@ -0,0 +1,20 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/GatedDeltaRuleOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class GatedDeltaRule : public Layer { + public: + GatedDeltaRule(); + explicit GatedDeltaRule(const aops::GatedDeltaRuleOpOptions& options); + explicit GatedDeltaRule(bool state_inplace); + + MLLM_LAYER_ANY_INPUTS_2_OUTPUTS_FORWARD +}; + +} // namespace mllm::nn diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7e4e3642c..8be52f5ee 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(nn) +add_subdirectory(models) add_subdirectory(core) add_subdirectory(utils) add_subdirectory(plugin) diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 90ce8037b..c4e75d0e0 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -1,29 +1,14 @@ -add_executable(Mllm-Test-CPUKernel KernelTest.cpp) +add_executable(Mllm-Test-CPUKernel + KernelTest.cpp + CausalDepthwiseConvCurrentFirstKernelTest.cpp + GatedDeltaRuleKernelTest.cpp) target_link_libraries(Mllm-Test-CPUKernel PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUKernel PRIVATE ${MLLM_INCLUDE_DIR}) -add_executable(Mllm-Test-Qwen35-GDN Qwen35GDNTest.cpp) -target_link_libraries(Mllm-Test-Qwen35-GDN PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Qwen35-GDN PRIVATE ${MLLM_INCLUDE_DIR}) - -add_executable(Mllm-Test-Qwen35-GDN-Conv Qwen35GDNConvTest.cpp) -target_link_libraries(Mllm-Test-Qwen35-GDN-Conv PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Qwen35-GDN-Conv PRIVATE ${MLLM_INCLUDE_DIR}) - add_executable(Mllm-Test-KaiW4A32Pack KaiW4A32PackTest.cpp) target_link_libraries(Mllm-Test-KaiW4A32Pack PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-KaiW4A32Pack PRIVATE ${MLLM_INCLUDE_DIR}) -add_executable(Mllm-Test-Qwen35-Tokenizer Qwen35TokenizerTest.cpp) -target_link_libraries(Mllm-Test-Qwen35-Tokenizer PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Qwen35-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) - -add_executable(Mllm-Test-Qwen35-Config Qwen35ConfigTest.cpp) -target_link_libraries(Mllm-Test-Qwen35-Config PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Qwen35-Config PRIVATE ${MLLM_INCLUDE_DIR}) -target_compile_definitions(Mllm-Test-Qwen35-Config - PRIVATE QWEN35_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/qwen3_5") - add_executable(Mllm-Test-MiniCPM5-Config MiniCPM5ConfigTest.cpp) target_link_libraries(Mllm-Test-MiniCPM5-Config PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-MiniCPM5-Config PRIVATE ${MLLM_INCLUDE_DIR}) @@ -45,9 +30,9 @@ target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCP target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) include(GoogleTest) - -add_executable(Mllm-Test-Qwen35-Multimodal Qwen35MultimodalTest.cpp) -target_link_libraries(Mllm-Test-Qwen35-Multimodal PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Qwen35-Multimodal PRIVATE ${MLLM_INCLUDE_DIR}) -add_test(NAME Qwen35MultimodalFocused COMMAND Mllm-Test-Qwen35-Multimodal) -set_tests_properties(Qwen35MultimodalFocused PROPERTIES LABELS qwen35) +gtest_discover_tests( + Mllm-Test-CPUKernel + TEST_PREFIX "CPUKernelFocused." + TEST_FILTER + "CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*" + PROPERTIES LABELS cpu-kernel) diff --git a/tests/cpu/Qwen35GDNConvTest.cpp b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp similarity index 82% rename from tests/cpu/Qwen35GDNConvTest.cpp rename to tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp index 46bcaa079..232938700 100644 --- a/tests/cpu/Qwen35GDNConvTest.cpp +++ b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp @@ -1,7 +1,7 @@ // Copyright (c) MLLM Team. // Licensed under the MIT License. -// Focused oracle for the GDN depthwise causal convolution. +// Focused oracle for current-first depthwise causal convolution. // // The reference below is an independent scalar implementation of the frozen // contract. It is deliberately not routed through the production kernel, so a @@ -39,9 +39,9 @@ std::vector makeBuffer(std::size_t count, int salt) { // Independent scalar reference for the frozen contract: // input/output [B, S, C], weight [C, K], history [B, C, K - 1] updated in place. -void referenceDepthwiseCausalConv(const std::vector& input, const std::vector& weight, - std::vector& state, std::vector& output, int batch_size, - int sequence_length, int channels, int kernel_size) { +void referenceDepthwiseCausalConv(const std::vector& input, const std::vector& weight, std::vector& state, + std::vector& output, int batch_size, int sequence_length, int channels, + int kernel_size) { const int state_width = kernel_size - 1; for (int batch = 0; batch < batch_size; ++batch) { for (int token = 0; token < sequence_length; ++token) { @@ -77,10 +77,8 @@ struct ConvCase { // Runs one case through the production kernel and the reference, and requires // bitwise agreement on both the output and the final history. void expectBitwiseAgreement(const ConvCase& test_case) { - const auto element_count = - static_cast(test_case.batch) * test_case.sequence * test_case.channels; - const auto state_count = - static_cast(test_case.batch) * test_case.channels * (test_case.kernel - 1); + const auto element_count = static_cast(test_case.batch) * test_case.sequence * test_case.channels; + const auto state_count = static_cast(test_case.batch) * test_case.channels * (test_case.kernel - 1); const std::vector input = makeBuffer(element_count, test_case.channels + test_case.sequence); const std::vector weight = @@ -102,7 +100,7 @@ void expectBitwiseAgreement(const ConvCase& test_case) { ASSERT_EQ(kernel_state, reference_state) << "final history mismatch for " << test_case.describe(); } -TEST(Qwen35GDNConvTest, MatchesScalarReferenceAcrossFocusedMatrix) { +TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAcrossFocusedMatrix) { // Channel counts below, at, and above the natural four-channel vector width, // including several that leave a tail. const int channel_values[] = {1, 2, 3, 4, 5, 7, 130}; @@ -122,9 +120,9 @@ TEST(Qwen35GDNConvTest, MatchesScalarReferenceAcrossFocusedMatrix) { } } -TEST(Qwen35GDNConvTest, MatchesScalarReferenceAtProductionChannelWidths) { - // 6144 is the Qwen3.5-0.8B convolution width, 8192 the 4B width; both are - // multiples of four, so they never exercise a tail on their own. +TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAtProductionChannelWidths) { + // Representative production widths that are multiples of four and therefore + // do not exercise a vector tail on their own. for (int channels : {6144, 8192}) { for (int sequence : {1, 16, 69, 128, 517}) { ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, sequence, channels, 4, true})); @@ -132,15 +130,13 @@ TEST(Qwen35GDNConvTest, MatchesScalarReferenceAtProductionChannelWidths) { } } -TEST(Qwen35GDNConvTest, MatchesScalarReferenceWithChannelTailAtProductionScale) { +TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceWithChannelTailAtProductionScale) { // Production width minus one, two, and three channels: a full-width run plus // a tail of three, two, and one channel respectively. - for (int channels : {8189, 8190, 8191, 6141}) { - ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, 69, channels, 4, true})); - } + for (int channels : {8189, 8190, 8191, 6141}) { ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, 69, channels, 4, true})); } } -TEST(Qwen35GDNConvTest, ChunkedPartitionsMatchOneShot) { +TEST(CausalDepthwiseConvCurrentFirstKernelTest, ChunkedPartitionsMatchOneShot) { struct Partition { int channels; int kernel; @@ -178,8 +174,8 @@ TEST(Qwen35GDNConvTest, ChunkedPartitionsMatchOneShot) { int consumed = 0; for (int chunk : partition.chunks) { const auto offset = static_cast(consumed) * partition.channels; - depthwiseCausalConvF32(input.data() + offset, weight.data(), chunked_state.data(), chunked_output.data() + offset, - 1, chunk, partition.channels, partition.kernel); + depthwiseCausalConvF32(input.data() + offset, weight.data(), chunked_state.data(), chunked_output.data() + offset, 1, + chunk, partition.channels, partition.kernel); consumed += chunk; } @@ -190,7 +186,7 @@ TEST(Qwen35GDNConvTest, ChunkedPartitionsMatchOneShot) { } } -TEST(Qwen35GDNConvTest, ResetBetweenRequestsReproducesFirstRequest) { +TEST(CausalDepthwiseConvCurrentFirstKernelTest, ResetBetweenRequestsReproducesFirstRequest) { constexpr int kChannels = 8192; constexpr int kKernel = 4; constexpr int kSequence = 69; @@ -202,28 +198,25 @@ TEST(Qwen35GDNConvTest, ResetBetweenRequestsReproducesFirstRequest) { std::vector state(kStateCount, 0.0F); std::vector first_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), first_output.data(), 1, kSequence, kChannels, - kKernel); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), first_output.data(), 1, kSequence, kChannels, kKernel); const std::vector first_state = state; // A second request that continues the history must differ, proving the // history is really being carried. std::vector continued_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), continued_output.data(), 1, kSequence, kChannels, - kKernel); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), continued_output.data(), 1, kSequence, kChannels, kKernel); ASSERT_NE(continued_output, first_output); // Resetting the history reproduces the first request bit for bit. std::fill(state.begin(), state.end(), 0.0F); std::vector reset_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), reset_output.data(), 1, kSequence, kChannels, - kKernel); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), reset_output.data(), 1, kSequence, kChannels, kKernel); ASSERT_EQ(reset_output, first_output); ASSERT_EQ(state, first_state); } -TEST(Qwen35GDNConvTest, RejectsNullBuffersAndInvalidGeometry) { +TEST(CausalDepthwiseConvCurrentFirstKernelTest, RejectsNullBuffersAndInvalidGeometry) { constexpr int kBatch = 1; constexpr int kSequence = 2; constexpr int kChannels = 4; @@ -234,9 +227,9 @@ TEST(Qwen35GDNConvTest, RejectsNullBuffersAndInvalidGeometry) { std::vector state(static_cast(kChannels) * (kKernel - 1), 0.0F); std::vector output(input.size(), 0.0F); - EXPECT_THROW(depthwiseCausalConvF32(nullptr, weight.data(), state.data(), output.data(), kBatch, kSequence, kChannels, - kKernel), - std::invalid_argument); + EXPECT_THROW( + depthwiseCausalConvF32(nullptr, weight.data(), state.data(), output.data(), kBatch, kSequence, kChannels, kKernel), + std::invalid_argument); EXPECT_THROW( depthwiseCausalConvF32(input.data(), nullptr, state.data(), output.data(), kBatch, kSequence, kChannels, kKernel), std::invalid_argument); @@ -247,15 +240,13 @@ TEST(Qwen35GDNConvTest, RejectsNullBuffersAndInvalidGeometry) { depthwiseCausalConvF32(input.data(), weight.data(), state.data(), nullptr, kBatch, kSequence, kChannels, kKernel), std::invalid_argument); - EXPECT_THROW(depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), 0, kSequence, kChannels, - kKernel), - std::invalid_argument); EXPECT_THROW( - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), kBatch, 0, kChannels, kKernel), - std::invalid_argument); - EXPECT_THROW( - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), kBatch, kSequence, 0, kKernel), + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), 0, kSequence, kChannels, kKernel), std::invalid_argument); + EXPECT_THROW(depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), kBatch, 0, kChannels, kKernel), + std::invalid_argument); + EXPECT_THROW(depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), kBatch, kSequence, 0, kKernel), + std::invalid_argument); // kernel_size <= 1 leaves no history and is rejected by the frozen contract. EXPECT_THROW( depthwiseCausalConvF32(input.data(), weight.data(), state.data(), output.data(), kBatch, kSequence, kChannels, 1), diff --git a/tests/cpu/Qwen35GDNTest.cpp b/tests/cpu/GatedDeltaRuleKernelTest.cpp similarity index 81% rename from tests/cpu/Qwen35GDNTest.cpp rename to tests/cpu/GatedDeltaRuleKernelTest.cpp index 0a95704e2..153d7b642 100644 --- a/tests/cpu/Qwen35GDNTest.cpp +++ b/tests/cpu/GatedDeltaRuleKernelTest.cpp @@ -12,16 +12,11 @@ namespace { -using mllm::cpu::gdn::depthwiseCausalConvF32; using mllm::cpu::gdn::gatedDeltaRuleF32; class ScopedCpuOpThreads { public: explicit ScopedCpuOpThreads(int32_t thread_count) : original_thread_count_(mllm::Context::instance().getCpuOpThreads()) { - // initializeContext() registers the CPU backend; SymbolTable::reg aborts on - // a duplicate key, so call it exactly once (the tests have no fixture init). - static const bool kContextInitialized = [] { mllm::initializeContext(); return true; }(); - (void)kContextInitialized; mllm::Context::instance().setCpuOpThreads(thread_count); } @@ -31,38 +26,7 @@ class ScopedCpuOpThreads { int32_t original_thread_count_; }; -TEST(Qwen35GDNTest, CausalConvChunkingMatchesSinglePrefill) { - constexpr int kBatch = 1; - constexpr int kSequence = 4; - constexpr int kChannels = 2; - constexpr int kKernel = 3; - - const std::array input = { - 1.0F, 10.0F, 2.0F, 20.0F, 3.0F, 30.0F, 4.0F, 40.0F, - }; - const std::array weight = { - 0.25F, 0.5F, 1.0F, -0.5F, 0.25F, 2.0F, - }; - std::array full_state = {}; - std::array chunked_state = {}; - std::array full_output = {}; - std::array chunked_output = {}; - - depthwiseCausalConvF32(input.data(), weight.data(), full_state.data(), full_output.data(), kBatch, kSequence, kChannels, - kKernel); - depthwiseCausalConvF32(input.data(), weight.data(), chunked_state.data(), chunked_output.data(), kBatch, - /*sequence_length=*/1, kChannels, kKernel); - depthwiseCausalConvF32(input.data() + kChannels, weight.data(), chunked_state.data(), chunked_output.data() + kChannels, - kBatch, kSequence - 1, kChannels, kKernel); - - EXPECT_EQ(full_output, chunked_output); - EXPECT_EQ(full_state, chunked_state); - EXPECT_FLOAT_EQ(full_output[0], 1.0F); - EXPECT_FLOAT_EQ(full_output[2], 2.5F); - EXPECT_FLOAT_EQ(full_output[6], 6.0F); -} - -TEST(Qwen35GDNTest, DeltaRuleChunkingMatchesSinglePrefill) { +TEST(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { constexpr int kBatch = 1; constexpr int kSequence = 3; constexpr int kKeyHeads = 1; @@ -115,7 +79,7 @@ TEST(Qwen35GDNTest, DeltaRuleChunkingMatchesSinglePrefill) { } } -TEST(Qwen35GDNTest, RepeatedKeyHeadsMatchExplicitExpansion) { +TEST(GatedDeltaRuleKernelTest, GroupedKeyHeadsMatchExplicitExpansion) { constexpr int kBatch = 1; constexpr int kSequence = 3; constexpr int kKeyHeads = 2; @@ -182,7 +146,7 @@ TEST(Qwen35GDNTest, RepeatedKeyHeadsMatchExplicitExpansion) { } } -TEST(Qwen35GDNTest, RejectsIncompatibleHeadCounts) { +TEST(GatedDeltaRuleKernelTest, RejectsIncompatibleHeadCounts) { float scalar = 0.0F; EXPECT_THROW(gatedDeltaRuleF32(&scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, /*batch_size=*/1, /*sequence_length=*/1, @@ -191,7 +155,7 @@ TEST(Qwen35GDNTest, RejectsIncompatibleHeadCounts) { std::invalid_argument); } -TEST(Qwen35GDNTest, MatchesOfficialL2NormalizationEpsilonPlacement) { +TEST(GatedDeltaRuleKernelTest, MatchesFrozenL2NormalizationEpsilonPlacement) { const std::array q = {1.0e-4F, 0.0F}; const std::array k = {1.0e-4F, 0.0F}; const std::array v = {1.0F}; @@ -206,16 +170,16 @@ TEST(Qwen35GDNTest, MatchesOfficialL2NormalizationEpsilonPlacement) { /*batch_size=*/1, /*sequence_length=*/1, /*num_key_heads=*/1, /*num_value_heads=*/1, /*key_head_dim=*/2, /*value_head_dim=*/1); - // Qwen3.5/FLA normalizes with rsqrt(sum(x^2) + 1e-6), then applies + // The frozen operation normalizes with rsqrt(sum(x^2) + 1e-6), then applies // 1/sqrt(key_head_dim) to q. Placing epsilon outside sqrt changes this // tiny-vector result by two orders of magnitude. EXPECT_NEAR(output[0], 0.0035005286F, 1.0e-8F); } -TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) { +TEST(GatedDeltaRuleKernelTest, ParallelBatchValueHeadsMatchSerialBitwise) { constexpr int kBatch = 2; constexpr int kSequence = 4; - // Qwen3.5 4B/9B GDN geometry: each normalized key head is shared by + // Production grouped-head geometry: each normalized key head is shared by // two independently scheduled value-head recurrence tasks. constexpr int kKeyHeads = 16; constexpr int kValueHeads = 32; @@ -272,11 +236,9 @@ TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) { } } -// 4B real geometry (B=1, S=69, 16 key heads, 32 value heads, 128 dims) at the -// 8-lane cap — exercises the full task fan-out (32 tasks) that the small -// geometry above does not. Guards against the device crash observed on -// OnePlus with the 8-lane product build. -TEST(Qwen35GDNTest, FourBGeometry8LaneDoesNotCrash) { +// Full production geometry at the 8-lane cap exercises the 32-task fan-out +// that the small geometry above does not. +TEST(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStable) { constexpr int kBatch = 1; constexpr int kSequence = 69; constexpr int kKeyHeads = 16; @@ -302,7 +264,10 @@ TEST(Qwen35GDNTest, FourBGeometry8LaneDoesNotCrash) { a[i] = 0.001F * static_cast(static_cast(i % 3)); b[i] = 0.001F * static_cast(static_cast(i % 9)); } - for (int i = 0; i < kValueHeads; ++i) { a_log[i] = -1.0F; dt_bias[i] = 0.0F; } + for (int i = 0; i < kValueHeads; ++i) { + a_log[i] = -1.0F; + dt_bias[i] = 0.0F; + } std::vector state(kBatch * kValueHeads * kValueDim * kKeyDim, 0.0F); std::vector output(v.size()); @@ -314,18 +279,14 @@ TEST(Qwen35GDNTest, FourBGeometry8LaneDoesNotCrash) { ref_output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, /*thread_count=*/1); const ScopedCpuOpThreads scoped_threads(kThreadCount); - gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), state.data(), - output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount); + gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), state.data(), output.data(), + kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount); - for (std::size_t i = 0; i < output.size(); ++i) { - ASSERT_EQ(ref_output[i], output[i]) << "output index " << i; - } - for (std::size_t i = 0; i < state.size(); ++i) { - ASSERT_EQ(ref_state[i], state[i]) << "state index " << i; - } + for (std::size_t i = 0; i < output.size(); ++i) { ASSERT_EQ(ref_output[i], output[i]) << "output index " << i; } + for (std::size_t i = 0; i < state.size(); ++i) { ASSERT_EQ(ref_state[i], state[i]) << "state index " << i; } - // Repeat the full 4B GDN pass 24 times (one per layer) to mimic the real - // model's layer loop, which interleaves the recurrence with other parallel + // Repeat the full recurrence 24 times to mimic a deep model layer loop, + // which interleaves the recurrence with other parallel // ops on the shared thread pool. Context init is now once-only (see // ScopedCpuOpThreads), so this exercises multi-call thread-pool reuse. // Run the recurrence 24 times on a FRESH copy of the initial state each diff --git a/tests/models/CMakeLists.txt b/tests/models/CMakeLists.txt new file mode 100644 index 000000000..c427145b8 --- /dev/null +++ b/tests/models/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(qwen3_5) diff --git a/tests/models/qwen3_5/CMakeLists.txt b/tests/models/qwen3_5/CMakeLists.txt new file mode 100644 index 000000000..9880f3efc --- /dev/null +++ b/tests/models/qwen3_5/CMakeLists.txt @@ -0,0 +1,18 @@ +add_executable(Mllm-Test-Qwen35-Config Qwen35ConfigTest.cpp) +target_link_libraries(Mllm-Test-Qwen35-Config PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Qwen35-Config PRIVATE ${MLLM_INCLUDE_DIR}) +target_compile_definitions(Mllm-Test-Qwen35-Config + PRIVATE QWEN35_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../../examples/qwen3_5") + +add_executable(Mllm-Test-Qwen35-Tokenizer Qwen35TokenizerTest.cpp) +target_link_libraries(Mllm-Test-Qwen35-Tokenizer PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Qwen35-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) + +add_executable(Mllm-Test-Qwen35-Multimodal Qwen35MultimodalTest.cpp) +target_link_libraries(Mllm-Test-Qwen35-Multimodal PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Qwen35-Multimodal PRIVATE ${MLLM_INCLUDE_DIR}) + +include(GoogleTest) +gtest_discover_tests(Mllm-Test-Qwen35-Config TEST_PREFIX "Qwen35ConfigFocused." PROPERTIES LABELS qwen35) +gtest_discover_tests(Mllm-Test-Qwen35-Tokenizer TEST_PREFIX "Qwen35TokenizerFocused." PROPERTIES LABELS qwen35) +gtest_discover_tests(Mllm-Test-Qwen35-Multimodal TEST_PREFIX "Qwen35MultimodalFocused." PROPERTIES LABELS qwen35) diff --git a/tests/cpu/Qwen35ConfigTest.cpp b/tests/models/qwen3_5/Qwen35ConfigTest.cpp similarity index 100% rename from tests/cpu/Qwen35ConfigTest.cpp rename to tests/models/qwen3_5/Qwen35ConfigTest.cpp diff --git a/tests/cpu/Qwen35MultimodalTest.cpp b/tests/models/qwen3_5/Qwen35MultimodalTest.cpp similarity index 100% rename from tests/cpu/Qwen35MultimodalTest.cpp rename to tests/models/qwen3_5/Qwen35MultimodalTest.cpp diff --git a/tests/cpu/Qwen35TokenizerTest.cpp b/tests/models/qwen3_5/Qwen35TokenizerTest.cpp similarity index 100% rename from tests/cpu/Qwen35TokenizerTest.cpp rename to tests/models/qwen3_5/Qwen35TokenizerTest.cpp diff --git a/tests/nn/CMakeLists.txt b/tests/nn/CMakeLists.txt index f90afe27a..aaf642e20 100644 --- a/tests/nn/CMakeLists.txt +++ b/tests/nn/CMakeLists.txt @@ -22,4 +22,9 @@ add_executable(Mllm-Test-Nn-ParallelLinear ParallelLinearTest.cpp) target_link_libraries(Mllm-Test-Nn-ParallelLinear PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-Nn-ParallelLinear PRIVATE ${MLLM_INCLUDE_DIR}) +add_executable(Mllm-Test-Nn-GatedDeltaRule GatedDeltaRuleTest.cpp) +target_link_libraries(Mllm-Test-Nn-GatedDeltaRule PRIVATE gtest_main MllmRT MllmCPUBackend) +target_include_directories(Mllm-Test-Nn-GatedDeltaRule PRIVATE ${MLLM_INCLUDE_DIR}) + include(GoogleTest) +gtest_discover_tests(Mllm-Test-Nn-GatedDeltaRule TEST_PREFIX "GatedDeltaRuleFocused.") diff --git a/tests/nn/GatedDeltaRuleTest.cpp b/tests/nn/GatedDeltaRuleTest.cpp new file mode 100644 index 000000000..14f7124a7 --- /dev/null +++ b/tests/nn/GatedDeltaRuleTest.cpp @@ -0,0 +1,217 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include +#include + +#include +#include +#include +#include +#include + +#include "mllm/compile/ir/Trace.hpp" +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/compile/jit/binary/LinalgIRSerialization.hpp" +#include "mllm/compile/jit/interpreter/AopsFromJson.hpp" +#include "mllm/mllm.hpp" +#include "mllm/nn/Nn.hpp" + +namespace { + +using mllm::Tensor; + +class GatedDeltaRuleTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +class GatedDeltaRuleModule final : public mllm::nn::Module { + public: + GatedDeltaRuleModule(std::string name, bool state_inplace) : Module(std::move(name)) { + gated_delta_rule_ = reg("gated_delta_rule", state_inplace); + } + + std::vector forward(const std::vector& inputs, const std::vector&) override { + auto [output, state] = + gated_delta_rule_(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], inputs[5], inputs[6], inputs[7]); + return {output, state}; + } + + private: + mllm::nn::GatedDeltaRule gated_delta_rule_; +}; + +template +auto findOp(const mllm::ir::node_ptr_t& node) -> typename OpType::ptr_t { + if (node->isa_()) { return node->cast_(); } + if (!node->isa_()) { return nullptr; } + for (const auto& region : node->cast_()->regions()) { + for (const auto& op : region->ops()) { + if (auto found = findOp(op)) { return found; } + } + } + return nullptr; +} + +Tensor patterned(const Tensor::shape_t& shape, float scale, float offset = 0.0F) { + auto tensor = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < tensor.numel(); ++index) { + tensor.ptr()[index] = std::sin(static_cast(index + 1) * scale) + offset; + } + return tensor; +} + +float stableSoftplus(float value) { + if (value > 20.0F) { return value; } + if (value < -20.0F) { return std::exp(value); } + return std::log1p(std::exp(value)); +} + +float stableSigmoid(float value) { + if (value >= 0.0F) { + const float exp_value = std::exp(-value); + return 1.0F / (1.0F + exp_value); + } + const float exp_value = std::exp(value); + return exp_value / (1.0F + exp_value); +} + +std::pair, std::vector> referenceGatedDeltaRule( + const std::vector& q, const std::vector& k, const std::vector& v, const std::vector& a, + const std::vector& b, const std::vector& a_log, const std::vector& dt_bias, std::vector state, + int batch, int sequence, int key_heads, int value_heads, int key_dim, int value_dim) { + const int key_head_repeats = value_heads / key_heads; + const float query_dim_scale = 1.0F / std::sqrt(static_cast(key_dim)); + std::vector output(static_cast(batch) * sequence * value_heads * value_dim); + std::vector normalized_query(key_dim); + std::vector normalized_key(key_dim); + + for (int batch_index = 0; batch_index < batch; ++batch_index) { + for (int value_head = 0; value_head < value_heads; ++value_head) { + const int key_head = value_head / key_head_repeats; + const std::size_t state_base = (static_cast(batch_index) * value_heads + value_head) * value_dim * key_dim; + for (int token = 0; token < sequence; ++token) { + const std::size_t qk_base = + ((static_cast(batch_index) * sequence + token) * key_heads + key_head) * key_dim; + float query_norm_sq = 0.0F; + float key_norm_sq = 0.0F; + for (int dim = 0; dim < key_dim; ++dim) { + query_norm_sq += q[qk_base + dim] * q[qk_base + dim]; + key_norm_sq += k[qk_base + dim] * k[qk_base + dim]; + } + const float query_scale = query_dim_scale / std::sqrt(query_norm_sq + 1.0e-6F); + const float key_scale = 1.0F / std::sqrt(key_norm_sq + 1.0e-6F); + for (int dim = 0; dim < key_dim; ++dim) { + normalized_query[dim] = q[qk_base + dim] * query_scale; + normalized_key[dim] = k[qk_base + dim] * key_scale; + } + + const std::size_t gate_index = (static_cast(batch_index) * sequence + token) * value_heads + value_head; + const float gate = -std::exp(a_log[value_head]) * stableSoftplus(a[gate_index] + dt_bias[value_head]); + const float decay = std::exp(gate); + const float beta = stableSigmoid(b[gate_index]); + const std::size_t value_base = gate_index * value_dim; + + for (int value_index = 0; value_index < value_dim; ++value_index) { + const std::size_t state_row = state_base + static_cast(value_index) * key_dim; + float state_dot_key = 0.0F; + for (int dim = 0; dim < key_dim; ++dim) { + state[state_row + dim] *= decay; + state_dot_key += state[state_row + dim] * normalized_key[dim]; + } + const float delta = (v[value_base + value_index] - state_dot_key) * beta; + float result = 0.0F; + for (int dim = 0; dim < key_dim; ++dim) { + state[state_row + dim] += delta * normalized_key[dim]; + result += state[state_row + dim] * normalized_query[dim]; + } + output[value_base + value_index] = result; + } + } + } + } + return {output, state}; +} + +void expectNear(const Tensor& actual, const std::vector& expected, float tolerance = 1.0e-5F) { + ASSERT_EQ(actual.numel(), expected.size()); + for (int index = 0; index < actual.numel(); ++index) { + EXPECT_NEAR(actual.ptr()[index], expected[index], tolerance) << "index " << index; + } +} + +TEST_F(GatedDeltaRuleTest, EagerMatchesIndependentGroupedHeadReferenceAndPreservesInputState) { + constexpr int kBatch = 1; + constexpr int kSequence = 3; + constexpr int kKeyHeads = 2; + constexpr int kValueHeads = 4; + constexpr int kKeyDim = 4; + constexpr int kValueDim = 3; + auto q = patterned({kBatch, kSequence, kKeyHeads, kKeyDim}, 0.03F); + auto k = patterned(q.shape(), 0.05F); + auto v = patterned({kBatch, kSequence, kValueHeads, kValueDim}, 0.07F); + auto a = patterned({kBatch, kSequence, kValueHeads}, 0.09F, -0.2F); + auto b = patterned({kBatch, kSequence, kValueHeads}, 0.11F); + auto a_log = patterned({kValueHeads}, 0.13F, -0.4F); + auto dt_bias = patterned({kValueHeads}, 0.15F, -0.1F); + auto state = patterned({kBatch, kValueHeads, kValueDim, kKeyDim}, 0.017F); + const auto state_before = state.toVector(); + const auto [expected_output, expected_state] = + referenceGatedDeltaRule(q.toVector(), k.toVector(), v.toVector(), a.toVector(), + b.toVector(), a_log.toVector(), dt_bias.toVector(), state_before, kBatch, + kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim); + + GatedDeltaRuleModule module("gated_delta_rule_eager", false); + const auto outputs = module(q, k, v, a, b, a_log, dt_bias, state); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].shape(), v.shape()); + EXPECT_EQ(outputs[1].shape(), state.shape()); + EXPECT_NE(outputs[1].ptr(), state.ptr()); + EXPECT_EQ(state.toVector(), state_before); + expectNear(outputs[0], expected_output); + expectNear(outputs[1], expected_state); +} + +TEST_F(GatedDeltaRuleTest, InplaceStateOutputAliasesInput) { + auto q = patterned({1, 2, 2, 4}, 0.03F); + auto k = patterned(q.shape(), 0.05F); + auto v = patterned({1, 2, 4, 3}, 0.07F); + auto a = patterned({1, 2, 4}, 0.09F, -0.2F); + auto b = patterned({1, 2, 4}, 0.11F); + auto a_log = patterned({4}, 0.13F, -0.4F); + auto dt_bias = patterned({4}, 0.15F, -0.1F); + auto state = patterned({1, 4, 3, 4}, 0.017F); + const auto* state_storage = state.ptr(); + + GatedDeltaRuleModule module("gated_delta_rule_inplace", true); + const auto outputs = module(q, k, v, a, b, a_log, dt_bias, state); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[1].ptr(), state_storage); +} + +TEST_F(GatedDeltaRuleTest, TraceAndSerializationPreserveStateSemantics) { + GatedDeltaRuleModule module("gated_delta_rule_trace", true); + auto ir_context = mllm::ir::trace( + module, Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 4, 3}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU), Tensor::empty({4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({4}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 4, 3, 4}, mllm::kFloat32, mllm::kCPU)); + auto op = findOp(ir_context->topLevelOp()); + ASSERT_NE(op, nullptr); + EXPECT_EQ(op->getAOp()->getOpType(), mllm::OpTypes::kGatedDeltaRule); + EXPECT_EQ(op->inputs().size(), 8); + EXPECT_EQ(op->outputs().size(), 2); + + const auto options = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_TRUE(options.at("state_inplace").get()); + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "GatedDeltaRule"}, {"backend", "CPU"}, {"op_options", options}}); + ASSERT_NE(restored, nullptr); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGatedDeltaRule); + EXPECT_TRUE(std::static_pointer_cast(restored)->options().state_inplace); +} + +} // namespace From fef1d42cc73e4b94a6d4d6d13cf732a7fd859fe4 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 2 Sep 2026 10:27:32 +0800 Subject: [PATCH 2/4] test(cpu): align GDN kernels with unified suite --- tests/cpu/CMakeLists.txt | 5 +- ...alDepthwiseConvCurrentFirstKernelTest.hpp} | 21 +++++--- ...lTest.cpp => GatedDeltaRuleKernelTest.hpp} | 21 +++++--- tests/cpu/KernelTest.cpp | 53 +++++++++++++++++++ 4 files changed, 80 insertions(+), 20 deletions(-) rename tests/cpu/{CausalDepthwiseConvCurrentFirstKernelTest.cpp => CausalDepthwiseConvCurrentFirstKernelTest.hpp} (94%) rename tests/cpu/{GatedDeltaRuleKernelTest.cpp => GatedDeltaRuleKernelTest.hpp} (96%) diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index c4e75d0e0..1b0e4ba27 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -1,7 +1,4 @@ -add_executable(Mllm-Test-CPUKernel - KernelTest.cpp - CausalDepthwiseConvCurrentFirstKernelTest.cpp - GatedDeltaRuleKernelTest.cpp) +add_executable(Mllm-Test-CPUKernel KernelTest.cpp) target_link_libraries(Mllm-Test-CPUKernel PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUKernel PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp similarity index 94% rename from tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp rename to tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp index 232938700..013f34e66 100644 --- a/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.cpp +++ b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp @@ -1,6 +1,8 @@ // Copyright (c) MLLM Team. // Licensed under the MIT License. +#pragma once + // Focused oracle for current-first depthwise causal convolution. // // The reference below is an independent scalar implementation of the frozen @@ -19,8 +21,9 @@ #include #include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" +#include "KernelTestHelper.hpp" -namespace { +namespace causal_depthwise_conv_current_first_test { using mllm::cpu::gdn::depthwiseCausalConvF32; @@ -100,7 +103,7 @@ void expectBitwiseAgreement(const ConvCase& test_case) { ASSERT_EQ(kernel_state, reference_state) << "final history mismatch for " << test_case.describe(); } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAcrossFocusedMatrix) { +inline void testMatchesScalarReferenceAcrossFocusedMatrix() { // Channel counts below, at, and above the natural four-channel vector width, // including several that leave a tail. const int channel_values[] = {1, 2, 3, 4, 5, 7, 130}; @@ -120,7 +123,7 @@ TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAcrossFocu } } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAtProductionChannelWidths) { +inline void testMatchesScalarReferenceAtProductionChannelWidths() { // Representative production widths that are multiples of four and therefore // do not exercise a vector tail on their own. for (int channels : {6144, 8192}) { @@ -130,13 +133,13 @@ TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAtProducti } } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceWithChannelTailAtProductionScale) { +inline void testMatchesScalarReferenceWithChannelTailAtProductionScale() { // Production width minus one, two, and three channels: a full-width run plus // a tail of three, two, and one channel respectively. for (int channels : {8189, 8190, 8191, 6141}) { ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, 69, channels, 4, true})); } } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, ChunkedPartitionsMatchOneShot) { +inline void testChunkedPartitionsMatchOneShot() { struct Partition { int channels; int kernel; @@ -186,7 +189,7 @@ TEST(CausalDepthwiseConvCurrentFirstKernelTest, ChunkedPartitionsMatchOneShot) { } } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, ResetBetweenRequestsReproducesFirstRequest) { +inline void testResetBetweenRequestsReproducesFirstRequest() { constexpr int kChannels = 8192; constexpr int kKernel = 4; constexpr int kSequence = 69; @@ -216,7 +219,7 @@ TEST(CausalDepthwiseConvCurrentFirstKernelTest, ResetBetweenRequestsReproducesFi ASSERT_EQ(state, first_state); } -TEST(CausalDepthwiseConvCurrentFirstKernelTest, RejectsNullBuffersAndInvalidGeometry) { +inline void testRejectsNullBuffersAndInvalidGeometry() { constexpr int kBatch = 1; constexpr int kSequence = 2; constexpr int kChannels = 4; @@ -253,4 +256,6 @@ TEST(CausalDepthwiseConvCurrentFirstKernelTest, RejectsNullBuffersAndInvalidGeom std::invalid_argument); } -} // namespace +} // namespace causal_depthwise_conv_current_first_test + +class CausalDepthwiseConvCurrentFirstKernelTest : public KernelTest {}; diff --git a/tests/cpu/GatedDeltaRuleKernelTest.cpp b/tests/cpu/GatedDeltaRuleKernelTest.hpp similarity index 96% rename from tests/cpu/GatedDeltaRuleKernelTest.cpp rename to tests/cpu/GatedDeltaRuleKernelTest.hpp index 153d7b642..e5319aa99 100644 --- a/tests/cpu/GatedDeltaRuleKernelTest.cpp +++ b/tests/cpu/GatedDeltaRuleKernelTest.hpp @@ -1,6 +1,8 @@ // Copyright (c) MLLM Team. // Licensed under the MIT License. +#pragma once + #include #include @@ -9,8 +11,9 @@ #include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" #include "mllm/mllm.hpp" +#include "KernelTestHelper.hpp" -namespace { +namespace gated_delta_rule_kernel_test { using mllm::cpu::gdn::gatedDeltaRuleF32; @@ -26,7 +29,7 @@ class ScopedCpuOpThreads { int32_t original_thread_count_; }; -TEST(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { +inline void testChunkingMatchesSinglePrefill() { constexpr int kBatch = 1; constexpr int kSequence = 3; constexpr int kKeyHeads = 1; @@ -79,7 +82,7 @@ TEST(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { } } -TEST(GatedDeltaRuleKernelTest, GroupedKeyHeadsMatchExplicitExpansion) { +inline void testGroupedKeyHeadsMatchExplicitExpansion() { constexpr int kBatch = 1; constexpr int kSequence = 3; constexpr int kKeyHeads = 2; @@ -146,7 +149,7 @@ TEST(GatedDeltaRuleKernelTest, GroupedKeyHeadsMatchExplicitExpansion) { } } -TEST(GatedDeltaRuleKernelTest, RejectsIncompatibleHeadCounts) { +inline void testRejectsIncompatibleHeadCounts() { float scalar = 0.0F; EXPECT_THROW(gatedDeltaRuleF32(&scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, &scalar, /*batch_size=*/1, /*sequence_length=*/1, @@ -155,7 +158,7 @@ TEST(GatedDeltaRuleKernelTest, RejectsIncompatibleHeadCounts) { std::invalid_argument); } -TEST(GatedDeltaRuleKernelTest, MatchesFrozenL2NormalizationEpsilonPlacement) { +inline void testMatchesFrozenL2NormalizationEpsilonPlacement() { const std::array q = {1.0e-4F, 0.0F}; const std::array k = {1.0e-4F, 0.0F}; const std::array v = {1.0F}; @@ -176,7 +179,7 @@ TEST(GatedDeltaRuleKernelTest, MatchesFrozenL2NormalizationEpsilonPlacement) { EXPECT_NEAR(output[0], 0.0035005286F, 1.0e-8F); } -TEST(GatedDeltaRuleKernelTest, ParallelBatchValueHeadsMatchSerialBitwise) { +inline void testParallelBatchValueHeadsMatchSerialBitwise() { constexpr int kBatch = 2; constexpr int kSequence = 4; // Production grouped-head geometry: each normalized key head is shared by @@ -238,7 +241,7 @@ TEST(GatedDeltaRuleKernelTest, ParallelBatchValueHeadsMatchSerialBitwise) { // Full production geometry at the 8-lane cap exercises the 32-task fan-out // that the small geometry above does not. -TEST(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStable) { +inline void testProductionGroupedHeadGeometry8LaneIsBitwiseStable() { constexpr int kBatch = 1; constexpr int kSequence = 69; constexpr int kKeyHeads = 16; @@ -307,4 +310,6 @@ TEST(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStable } } -} // namespace +} // namespace gated_delta_rule_kernel_test + +class GatedDeltaRuleKernelTest : public KernelTest {}; diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index 277bb51d7..1449aa561 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -562,6 +562,59 @@ TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwis true); } +#include "CausalDepthwiseConvCurrentFirstKernelTest.hpp" +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAcrossFocusedMatrix) { + causal_depthwise_conv_current_first_test::testMatchesScalarReferenceAcrossFocusedMatrix(); +} + +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAtProductionChannelWidths) { + causal_depthwise_conv_current_first_test::testMatchesScalarReferenceAtProductionChannelWidths(); +} + +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceWithChannelTailAtProductionScale) { + causal_depthwise_conv_current_first_test::testMatchesScalarReferenceWithChannelTailAtProductionScale(); +} + +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, ChunkedPartitionsMatchOneShot) { + causal_depthwise_conv_current_first_test::testChunkedPartitionsMatchOneShot(); +} + +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, ResetBetweenRequestsReproducesFirstRequest) { + causal_depthwise_conv_current_first_test::testResetBetweenRequestsReproducesFirstRequest(); +} + +TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, RejectsNullBuffersAndInvalidGeometry) { + causal_depthwise_conv_current_first_test::testRejectsNullBuffersAndInvalidGeometry(); +} + +//===----------------------------------------------------------------------===// +// Gated delta rule +//===----------------------------------------------------------------------===// +#include "GatedDeltaRuleKernelTest.hpp" +TEST_F(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { + gated_delta_rule_kernel_test::testChunkingMatchesSinglePrefill(); +} + +TEST_F(GatedDeltaRuleKernelTest, GroupedKeyHeadsMatchExplicitExpansion) { + gated_delta_rule_kernel_test::testGroupedKeyHeadsMatchExplicitExpansion(); +} + +TEST_F(GatedDeltaRuleKernelTest, RejectsIncompatibleHeadCounts) { + gated_delta_rule_kernel_test::testRejectsIncompatibleHeadCounts(); +} + +TEST_F(GatedDeltaRuleKernelTest, MatchesFrozenL2NormalizationEpsilonPlacement) { + gated_delta_rule_kernel_test::testMatchesFrozenL2NormalizationEpsilonPlacement(); +} + +TEST_F(GatedDeltaRuleKernelTest, ParallelBatchValueHeadsMatchSerialBitwise) { + gated_delta_rule_kernel_test::testParallelBatchValueHeadsMatchSerialBitwise(); +} + +TEST_F(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStable) { + gated_delta_rule_kernel_test::testProductionGroupedHeadGeometry8LaneIsBitwiseStable(); +} + //===----------------------------------------------------------------------===// // Parallel linear //===----------------------------------------------------------------------===// From b65748819f144e05d04f223ea9a259f2b222da94 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 2 Sep 2026 10:38:54 +0800 Subject: [PATCH 3/4] test(cpu): expose GDN kernel case matrices --- ...salDepthwiseConvCurrentFirstKernelTest.hpp | 101 ++++++++---------- tests/cpu/GatedDeltaRuleKernelTest.hpp | 72 +++++++++---- tests/cpu/KernelTest.cpp | 43 +++++--- 3 files changed, 122 insertions(+), 94 deletions(-) diff --git a/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp index 013f34e66..2824cb6d3 100644 --- a/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp +++ b/tests/cpu/CausalDepthwiseConvCurrentFirstKernelTest.hpp @@ -103,18 +103,14 @@ void expectBitwiseAgreement(const ConvCase& test_case) { ASSERT_EQ(kernel_state, reference_state) << "final history mismatch for " << test_case.describe(); } -inline void testMatchesScalarReferenceAcrossFocusedMatrix() { - // Channel counts below, at, and above the natural four-channel vector width, - // including several that leave a tail. - const int channel_values[] = {1, 2, 3, 4, 5, 7, 130}; - const int sequence_values[] = {1, 2, 16, 69, 128, 517}; - const int kernel_values[] = {2, 3, 4, 5}; - - for (int batch : {1, 2}) { +inline void testMatchesScalarReferenceMatrix(const std::vector& batch_values, const std::vector& sequence_values, + const std::vector& channel_values, const std::vector& kernel_values, + const std::vector& non_zero_history_values) { + for (int batch : batch_values) { for (int sequence : sequence_values) { for (int channels : channel_values) { for (int kernel : kernel_values) { - for (bool non_zero_history : {false, true}) { + for (bool non_zero_history : non_zero_history_values) { ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({batch, sequence, channels, kernel, non_zero_history})); } } @@ -123,39 +119,13 @@ inline void testMatchesScalarReferenceAcrossFocusedMatrix() { } } -inline void testMatchesScalarReferenceAtProductionChannelWidths() { - // Representative production widths that are multiples of four and therefore - // do not exercise a vector tail on their own. - for (int channels : {6144, 8192}) { - for (int sequence : {1, 16, 69, 128, 517}) { - ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, sequence, channels, 4, true})); - } - } -} - -inline void testMatchesScalarReferenceWithChannelTailAtProductionScale() { - // Production width minus one, two, and three channels: a full-width run plus - // a tail of three, two, and one channel respectively. - for (int channels : {8189, 8190, 8191, 6141}) { ASSERT_NO_FATAL_FAILURE(expectBitwiseAgreement({1, 69, channels, 4, true})); } -} - -inline void testChunkedPartitionsMatchOneShot() { - struct Partition { - int channels; - int kernel; - std::vector chunks; - }; - - const std::vector partitions = { - {8192, 4, {517}}, // one-shot reference - {8192, 4, {1, 516}}, // prefill then continuation - {8192, 4, {128, 128, 128, 128, 5}}, // multi-chunk - {130, 4, {1, 15, 53}}, // tail channels across chunks - {7, 5, {1, 1, 14}}, // generic kernel size, odd channels - {6144, 4, {69}}, - {6144, 4, {16, 16, 16, 21}}, - }; +struct Partition { + int channels; + int kernel; + std::vector chunks; +}; +inline void testChunkedPartitionsMatchOneShot(const std::vector& partitions) { for (const auto& partition : partitions) { int total_sequence = 0; for (int chunk : partition.chunks) { total_sequence += chunk; } @@ -189,31 +159,28 @@ inline void testChunkedPartitionsMatchOneShot() { } } -inline void testResetBetweenRequestsReproducesFirstRequest() { - constexpr int kChannels = 8192; - constexpr int kKernel = 4; - constexpr int kSequence = 69; - constexpr auto kElements = static_cast(kSequence) * kChannels; - constexpr auto kStateCount = static_cast(kChannels) * (kKernel - 1); +inline void testResetBetweenRequestsReproducesFirstRequest(int channels, int kernel, int sequence) { + const auto element_count = static_cast(sequence) * channels; + const auto state_count = static_cast(channels) * (kernel - 1); - const std::vector input = makeBuffer(kElements, 5); - const std::vector weight = makeBuffer(static_cast(kChannels) * kKernel, kKernel); + const std::vector input = makeBuffer(element_count, 5); + const std::vector weight = makeBuffer(static_cast(channels) * kernel, kernel); - std::vector state(kStateCount, 0.0F); - std::vector first_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), first_output.data(), 1, kSequence, kChannels, kKernel); + std::vector state(state_count, 0.0F); + std::vector first_output(element_count, 0.0F); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), first_output.data(), 1, sequence, channels, kernel); const std::vector first_state = state; // A second request that continues the history must differ, proving the // history is really being carried. - std::vector continued_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), continued_output.data(), 1, kSequence, kChannels, kKernel); + std::vector continued_output(element_count, 0.0F); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), continued_output.data(), 1, sequence, channels, kernel); ASSERT_NE(continued_output, first_output); // Resetting the history reproduces the first request bit for bit. std::fill(state.begin(), state.end(), 0.0F); - std::vector reset_output(kElements, 0.0F); - depthwiseCausalConvF32(input.data(), weight.data(), state.data(), reset_output.data(), 1, kSequence, kChannels, kKernel); + std::vector reset_output(element_count, 0.0F); + depthwiseCausalConvF32(input.data(), weight.data(), state.data(), reset_output.data(), 1, sequence, channels, kernel); ASSERT_EQ(reset_output, first_output); ASSERT_EQ(state, first_state); @@ -258,4 +225,24 @@ inline void testRejectsNullBuffersAndInvalidGeometry() { } // namespace causal_depthwise_conv_current_first_test -class CausalDepthwiseConvCurrentFirstKernelTest : public KernelTest {}; +class CausalDepthwiseConvCurrentFirstKernelTest : public KernelTest { + public: + void testMatchesScalarReferenceMatrix(const std::vector& batch_values, const std::vector& sequence_values, + const std::vector& channel_values, const std::vector& kernel_values, + const std::vector& non_zero_history_values) { + causal_depthwise_conv_current_first_test::testMatchesScalarReferenceMatrix(batch_values, sequence_values, channel_values, + kernel_values, non_zero_history_values); + } + + void testChunkedPartitionsMatchOneShot(const std::vector& partitions) { + causal_depthwise_conv_current_first_test::testChunkedPartitionsMatchOneShot(partitions); + } + + void testResetBetweenRequestsReproducesFirstRequest(int channels, int kernel, int sequence) { + causal_depthwise_conv_current_first_test::testResetBetweenRequestsReproducesFirstRequest(channels, kernel, sequence); + } + + void testRejectsNullBuffersAndInvalidGeometry() { + causal_depthwise_conv_current_first_test::testRejectsNullBuffersAndInvalidGeometry(); + } +}; diff --git a/tests/cpu/GatedDeltaRuleKernelTest.hpp b/tests/cpu/GatedDeltaRuleKernelTest.hpp index e5319aa99..f9a28488f 100644 --- a/tests/cpu/GatedDeltaRuleKernelTest.hpp +++ b/tests/cpu/GatedDeltaRuleKernelTest.hpp @@ -29,6 +29,16 @@ class ScopedCpuOpThreads { int32_t original_thread_count_; }; +struct Geometry { + int batch; + int sequence; + int key_heads; + int value_heads; + int key_dim; + int value_dim; + int thread_count; +}; + inline void testChunkingMatchesSinglePrefill() { constexpr int kBatch = 1; constexpr int kSequence = 3; @@ -179,19 +189,14 @@ inline void testMatchesFrozenL2NormalizationEpsilonPlacement() { EXPECT_NEAR(output[0], 0.0035005286F, 1.0e-8F); } -inline void testParallelBatchValueHeadsMatchSerialBitwise() { - constexpr int kBatch = 2; - constexpr int kSequence = 4; - // Production grouped-head geometry: each normalized key head is shared by - // two independently scheduled value-head recurrence tasks. - constexpr int kKeyHeads = 16; - constexpr int kValueHeads = 32; - constexpr int kKeyDim = 128; - constexpr int kValueDim = 128; - // Exercises the parallel lane partition up to the 8-lane cap - // (kMaxParallelGDNLanes); tasks are disjoint so output must be bitwise - // identical regardless of how many lanes the scheduler picks. - constexpr int kThreadCount = 8; +inline void testParallelBatchValueHeadsMatchSerialBitwise(const Geometry& geometry) { + const int kBatch = geometry.batch; + const int kSequence = geometry.sequence; + const int kKeyHeads = geometry.key_heads; + const int kValueHeads = geometry.value_heads; + const int kKeyDim = geometry.key_dim; + const int kValueDim = geometry.value_dim; + const int kThreadCount = geometry.thread_count; std::vector q(kBatch * kSequence * kKeyHeads * kKeyDim); std::vector k(q.size()); @@ -241,14 +246,14 @@ inline void testParallelBatchValueHeadsMatchSerialBitwise() { // Full production geometry at the 8-lane cap exercises the 32-task fan-out // that the small geometry above does not. -inline void testProductionGroupedHeadGeometry8LaneIsBitwiseStable() { - constexpr int kBatch = 1; - constexpr int kSequence = 69; - constexpr int kKeyHeads = 16; - constexpr int kValueHeads = 32; - constexpr int kKeyDim = 128; - constexpr int kValueDim = 128; - constexpr int kThreadCount = 8; +inline void testProductionGroupedHeadGeometryIsBitwiseStable(const Geometry& geometry, int repeats) { + const int kBatch = geometry.batch; + const int kSequence = geometry.sequence; + const int kKeyHeads = geometry.key_heads; + const int kValueHeads = geometry.value_heads; + const int kKeyDim = geometry.key_dim; + const int kValueDim = geometry.value_dim; + const int kThreadCount = geometry.thread_count; std::vector q(kBatch * kSequence * kKeyHeads * kKeyDim); std::vector k(q.size()); @@ -297,7 +302,7 @@ inline void testProductionGroupedHeadGeometry8LaneIsBitwiseStable() { // and compare each run's output to the serial reference for that same input. // This exercises repeated thread-pool push/acquire/release cycles — the // multi-call reuse pattern that crashed on device. - for (int layer = 0; layer < 24; ++layer) { + for (int layer = 0; layer < repeats; ++layer) { std::vector layer_state(state.size(), 0.0F); std::vector layer_ref_state(state.size(), 0.0F); gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), layer_ref_state.data(), @@ -312,4 +317,25 @@ inline void testProductionGroupedHeadGeometry8LaneIsBitwiseStable() { } // namespace gated_delta_rule_kernel_test -class GatedDeltaRuleKernelTest : public KernelTest {}; +class GatedDeltaRuleKernelTest : public KernelTest { + public: + void testChunkingMatchesSinglePrefill() { gated_delta_rule_kernel_test::testChunkingMatchesSinglePrefill(); } + + void testGroupedKeyHeadsMatchExplicitExpansion() { + gated_delta_rule_kernel_test::testGroupedKeyHeadsMatchExplicitExpansion(); + } + + void testRejectsIncompatibleHeadCounts() { gated_delta_rule_kernel_test::testRejectsIncompatibleHeadCounts(); } + + void testMatchesFrozenL2NormalizationEpsilonPlacement() { + gated_delta_rule_kernel_test::testMatchesFrozenL2NormalizationEpsilonPlacement(); + } + + void testParallelBatchValueHeadsMatchSerialBitwise(const gated_delta_rule_kernel_test::Geometry& geometry) { + gated_delta_rule_kernel_test::testParallelBatchValueHeadsMatchSerialBitwise(geometry); + } + + void testProductionGroupedHeadGeometryIsBitwiseStable(const gated_delta_rule_kernel_test::Geometry& geometry, int repeats) { + gated_delta_rule_kernel_test::testProductionGroupedHeadGeometryIsBitwiseStable(geometry, repeats); + } +}; diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index 1449aa561..b3f3d4d99 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -564,55 +564,70 @@ TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwis #include "CausalDepthwiseConvCurrentFirstKernelTest.hpp" TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAcrossFocusedMatrix) { - causal_depthwise_conv_current_first_test::testMatchesScalarReferenceAcrossFocusedMatrix(); + EXPECT_NO_FATAL_FAILURE( + testMatchesScalarReferenceMatrix({1, 2}, {1, 2, 16, 69, 128, 517}, {1, 2, 3, 4, 5, 7, 130}, {2, 3, 4, 5}, {false, true})); } TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceAtProductionChannelWidths) { - causal_depthwise_conv_current_first_test::testMatchesScalarReferenceAtProductionChannelWidths(); + // Exact vector blocks at the production channel widths. + EXPECT_NO_FATAL_FAILURE(testMatchesScalarReferenceMatrix({1}, {1, 16, 69, 128, 517}, {6144, 8192}, {4}, {true})); } TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, MatchesScalarReferenceWithChannelTailAtProductionScale) { - causal_depthwise_conv_current_first_test::testMatchesScalarReferenceWithChannelTailAtProductionScale(); + // Production width minus one, two, and three channels exercises vector tails. + EXPECT_NO_FATAL_FAILURE(testMatchesScalarReferenceMatrix({1}, {69}, {6141, 8189, 8190, 8191}, {4}, {true})); } TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, ChunkedPartitionsMatchOneShot) { - causal_depthwise_conv_current_first_test::testChunkedPartitionsMatchOneShot(); + EXPECT_NO_FATAL_FAILURE(testChunkedPartitionsMatchOneShot({ + {8192, 4, {517}}, + {8192, 4, {1, 516}}, + {8192, 4, {128, 128, 128, 128, 5}}, + {130, 4, {1, 15, 53}}, + {7, 5, {1, 1, 14}}, + {6144, 4, {69}}, + {6144, 4, {16, 16, 16, 21}}, + })); } TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, ResetBetweenRequestsReproducesFirstRequest) { - causal_depthwise_conv_current_first_test::testResetBetweenRequestsReproducesFirstRequest(); + EXPECT_NO_FATAL_FAILURE(testResetBetweenRequestsReproducesFirstRequest(/*channels=*/8192, /*kernel=*/4, /*sequence=*/69)); } TEST_F(CausalDepthwiseConvCurrentFirstKernelTest, RejectsNullBuffersAndInvalidGeometry) { - causal_depthwise_conv_current_first_test::testRejectsNullBuffersAndInvalidGeometry(); + EXPECT_NO_FATAL_FAILURE(testRejectsNullBuffersAndInvalidGeometry()); } //===----------------------------------------------------------------------===// // Gated delta rule //===----------------------------------------------------------------------===// #include "GatedDeltaRuleKernelTest.hpp" -TEST_F(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { - gated_delta_rule_kernel_test::testChunkingMatchesSinglePrefill(); -} +TEST_F(GatedDeltaRuleKernelTest, ChunkingMatchesSinglePrefill) { EXPECT_NO_FATAL_FAILURE(testChunkingMatchesSinglePrefill()); } TEST_F(GatedDeltaRuleKernelTest, GroupedKeyHeadsMatchExplicitExpansion) { - gated_delta_rule_kernel_test::testGroupedKeyHeadsMatchExplicitExpansion(); + EXPECT_NO_FATAL_FAILURE(testGroupedKeyHeadsMatchExplicitExpansion()); } TEST_F(GatedDeltaRuleKernelTest, RejectsIncompatibleHeadCounts) { - gated_delta_rule_kernel_test::testRejectsIncompatibleHeadCounts(); + EXPECT_NO_FATAL_FAILURE(testRejectsIncompatibleHeadCounts()); } TEST_F(GatedDeltaRuleKernelTest, MatchesFrozenL2NormalizationEpsilonPlacement) { - gated_delta_rule_kernel_test::testMatchesFrozenL2NormalizationEpsilonPlacement(); + EXPECT_NO_FATAL_FAILURE(testMatchesFrozenL2NormalizationEpsilonPlacement()); } TEST_F(GatedDeltaRuleKernelTest, ParallelBatchValueHeadsMatchSerialBitwise) { - gated_delta_rule_kernel_test::testParallelBatchValueHeadsMatchSerialBitwise(); + // Each of 16 key heads is shared by two independently scheduled value heads. + EXPECT_NO_FATAL_FAILURE(testParallelBatchValueHeadsMatchSerialBitwise({/*batch=*/2, /*sequence=*/4, /*key_heads=*/16, + /*value_heads=*/32, /*key_dim=*/128, /*value_dim=*/128, + /*thread_count=*/8})); } TEST_F(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStable) { - gated_delta_rule_kernel_test::testProductionGroupedHeadGeometry8LaneIsBitwiseStable(); + EXPECT_NO_FATAL_FAILURE(testProductionGroupedHeadGeometryIsBitwiseStable( + {/*batch=*/1, /*sequence=*/69, /*key_heads=*/16, /*value_heads=*/32, /*key_dim=*/128, /*value_dim=*/128, + /*thread_count=*/8}, + /*repeats=*/24)); } //===----------------------------------------------------------------------===// From e0252d959a9d5f29912f8f12c85081d5ae1baaf4 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 2 Sep 2026 11:24:57 +0800 Subject: [PATCH 4/4] fix(test): avoid cross-running Android test binaries --- tests/cpu/CMakeLists.txt | 12 ++++++------ tests/models/qwen3_5/CMakeLists.txt | 8 ++++---- tests/nn/CMakeLists.txt | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 1b0e4ba27..70ba14ac4 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -27,9 +27,9 @@ target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCP target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) include(GoogleTest) -gtest_discover_tests( - Mllm-Test-CPUKernel - TEST_PREFIX "CPUKernelFocused." - TEST_FILTER - "CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*" - PROPERTIES LABELS cpu-kernel) +add_test( + NAME CPUKernelFocused + COMMAND + Mllm-Test-CPUKernel + --gtest_filter=CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*) +set_tests_properties(CPUKernelFocused PROPERTIES LABELS cpu-kernel) diff --git a/tests/models/qwen3_5/CMakeLists.txt b/tests/models/qwen3_5/CMakeLists.txt index 9880f3efc..683baf4b9 100644 --- a/tests/models/qwen3_5/CMakeLists.txt +++ b/tests/models/qwen3_5/CMakeLists.txt @@ -12,7 +12,7 @@ add_executable(Mllm-Test-Qwen35-Multimodal Qwen35MultimodalTest.cpp) target_link_libraries(Mllm-Test-Qwen35-Multimodal PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-Qwen35-Multimodal PRIVATE ${MLLM_INCLUDE_DIR}) -include(GoogleTest) -gtest_discover_tests(Mllm-Test-Qwen35-Config TEST_PREFIX "Qwen35ConfigFocused." PROPERTIES LABELS qwen35) -gtest_discover_tests(Mllm-Test-Qwen35-Tokenizer TEST_PREFIX "Qwen35TokenizerFocused." PROPERTIES LABELS qwen35) -gtest_discover_tests(Mllm-Test-Qwen35-Multimodal TEST_PREFIX "Qwen35MultimodalFocused." PROPERTIES LABELS qwen35) +add_test(NAME Qwen35ConfigFocused COMMAND Mllm-Test-Qwen35-Config) +add_test(NAME Qwen35TokenizerFocused COMMAND Mllm-Test-Qwen35-Tokenizer) +add_test(NAME Qwen35MultimodalFocused COMMAND Mllm-Test-Qwen35-Multimodal) +set_tests_properties(Qwen35ConfigFocused Qwen35TokenizerFocused Qwen35MultimodalFocused PROPERTIES LABELS qwen35) diff --git a/tests/nn/CMakeLists.txt b/tests/nn/CMakeLists.txt index aaf642e20..bd003caac 100644 --- a/tests/nn/CMakeLists.txt +++ b/tests/nn/CMakeLists.txt @@ -27,4 +27,4 @@ target_link_libraries(Mllm-Test-Nn-GatedDeltaRule PRIVATE gtest_main MllmRT Mllm target_include_directories(Mllm-Test-Nn-GatedDeltaRule PRIVATE ${MLLM_INCLUDE_DIR}) include(GoogleTest) -gtest_discover_tests(Mllm-Test-Nn-GatedDeltaRule TEST_PREFIX "GatedDeltaRuleFocused.") +add_test(NAME GatedDeltaRuleFocused COMMAND Mllm-Test-Nn-GatedDeltaRule)