wehub-resource-sync
297 lines
12 KiB
C++
297 lines
12 KiB
C++
//
|
||
// FusedGatedDeltaTest.cpp
|
||
// MNNTests
|
||
//
|
||
// End-to-end test for the MNNFusedGatedDelta kernel. Runs the
|
||
// LinearAttention op via Module::load with Qwen3-Next-style head
|
||
// dimensions (d_k=d_v=128) and Mamba-style (d_k=d_v=64) — these
|
||
// exercise the SIMD chunk path of the fused kernel that the existing
|
||
// LinearAttentionTest (d_v=4/16) does not.
|
||
//
|
||
// Compares the runtime output against a scalar Python-style reference
|
||
// that decomposes the gated_delta_rule recurrence step by step.
|
||
//
|
||
|
||
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
||
|
||
#include <MNN/expr/Expr.hpp>
|
||
#include <MNN/expr/ExprCreator.hpp>
|
||
#include <MNN/expr/Module.hpp>
|
||
#include "MNNTestSuite.h"
|
||
#include "TestUtils.h"
|
||
|
||
#include <cmath>
|
||
#include <cstring>
|
||
#include <random>
|
||
#include <vector>
|
||
|
||
using namespace MNN;
|
||
using namespace MNN::Express;
|
||
|
||
namespace {
|
||
|
||
// Reference implementation of one decode step of gated_delta_rule for a
|
||
// single attention head. Mirrors the math in CPULinearAttention.cpp:
|
||
// out_k = S^T @ k
|
||
// delta = beta * (v - decay * out_k)
|
||
// out = decay * (S^T @ q) + dot(k,q) * delta
|
||
// S = decay * S + k ⊗ delta
|
||
static void refOneStep(float* S, const float* k, const float* q, const float* v, float* out, float decay, float beta,
|
||
int dk, int dv) {
|
||
std::vector<float> outK(dv, 0.0f), outQ(dv, 0.0f), delta(dv, 0.0f);
|
||
for (int i = 0; i < dk; ++i) {
|
||
const float* row = S + i * dv;
|
||
for (int j = 0; j < dv; ++j) {
|
||
outK[j] += row[j] * k[i];
|
||
outQ[j] += row[j] * q[i];
|
||
}
|
||
}
|
||
float kq = 0.0f;
|
||
for (int i = 0; i < dk; ++i)
|
||
kq += k[i] * q[i];
|
||
for (int j = 0; j < dv; ++j) {
|
||
float vPred = decay * outK[j];
|
||
delta[j] = beta * (v[j] - vPred);
|
||
out[j] = decay * outQ[j] + kq * delta[j];
|
||
}
|
||
for (int i = 0; i < dk; ++i) {
|
||
float ki = k[i];
|
||
float* row = S + i * dv;
|
||
for (int j = 0; j < dv; ++j) {
|
||
row[j] = decay * row[j] + ki * delta[j];
|
||
}
|
||
}
|
||
}
|
||
|
||
// Apply L2-norm + scale (Q is also scaled by 1/sqrt(d_k)) — matches the
|
||
// useL2Norm=true path inside CPULinearAttention.
|
||
static void applyL2NormAndScale(float* q, float* k, int dk) {
|
||
const float eps = 1e-6f;
|
||
const float qScale = 1.0f / std::sqrt((float)dk);
|
||
float qSq = 0.0f, kSq = 0.0f;
|
||
for (int i = 0; i < dk; ++i) {
|
||
qSq += q[i] * q[i];
|
||
kSq += k[i] * k[i];
|
||
}
|
||
float qNS = qScale / std::sqrt(qSq + eps);
|
||
float kIN = 1.0f / std::sqrt(kSq + eps);
|
||
for (int i = 0; i < dk; ++i) {
|
||
q[i] *= qNS;
|
||
k[i] *= kIN;
|
||
}
|
||
}
|
||
|
||
// Build a LinearAttention op as a Module so we can run forward(). Uses the
|
||
// same FlatBuffers-based construction as LinearAttentionTest.
|
||
static std::shared_ptr<Module> makeModule(int numKHeads, int numVHeads, int dk, int dv) {
|
||
auto qkv = _Input();
|
||
auto gate = _Input();
|
||
auto beta = _Input();
|
||
auto convW = _Input();
|
||
|
||
std::shared_ptr<MNN::OpT> op(new MNN::OpT);
|
||
op->type = MNN::OpType_LinearAttention;
|
||
op->main.type = MNN::OpParameter_LinearAttentionParam;
|
||
op->main.value = new MNN::LinearAttentionParamT;
|
||
auto* p = op->main.AsLinearAttentionParam();
|
||
p->attn_type = "gated_delta_rule";
|
||
p->num_k_heads = numKHeads;
|
||
p->num_v_heads = numVHeads;
|
||
p->head_k_dim = dk;
|
||
p->head_v_dim = dv;
|
||
p->use_qk_l2norm = true;
|
||
|
||
auto out = Variable::create(Expr::create(op.get(), {qkv, gate, beta, convW}));
|
||
auto buffer = Variable::save({out});
|
||
|
||
MNN::ScheduleConfig config;
|
||
auto status = MNNTestSuite::get()->pStaus;
|
||
config.type = (MNNForwardType)status.forwardType;
|
||
MNN::BackendConfig bn;
|
||
bn.memory = (MNN::BackendConfig::MemoryMode)status.memory;
|
||
bn.precision = (MNN::BackendConfig::PrecisionMode)status.precision;
|
||
bn.power = (MNN::BackendConfig::PowerMode)status.power;
|
||
config.backendConfig = &bn;
|
||
config.numThread = 1;
|
||
|
||
std::shared_ptr<Executor::RuntimeManager> rt(Executor::RuntimeManager::createRuntimeManager(config));
|
||
return std::shared_ptr<Module>(Module::load({}, {}, (uint8_t*)buffer.data(), buffer.size(), rt));
|
||
}
|
||
|
||
struct Case {
|
||
const char* name;
|
||
int Hk;
|
||
int Hv;
|
||
int dk;
|
||
int dv;
|
||
int L; // 1 → exercises gated_delta_rule_decode; >1 → gated_delta_rule_mnn (prefill)
|
||
int numSteps; // number of forward() invocations; state accumulates across steps
|
||
};
|
||
|
||
// Run one Case: drives the LinearAttention module forward `numSteps` times,
|
||
// each invocation feeding a fresh random qkv tensor with seqLen=L. State (both
|
||
// the conv state and the recurrent S) accumulates across calls. The reference
|
||
// path replays the same conv1D + L2norm + gated_delta_rule pipeline as
|
||
// CPULinearAttention and compares element-wise.
|
||
static bool runCase(const Case& cs, std::mt19937& rng, float tolerance) {
|
||
const int B = 1;
|
||
const int Hk = cs.Hk;
|
||
const int Hv = cs.Hv;
|
||
const int dk = cs.dk;
|
||
const int dv = cs.dv;
|
||
const int L = cs.L;
|
||
const int K_conv = 4;
|
||
const int convStateSize = K_conv - 1;
|
||
const int key_dim = Hk * dk;
|
||
const int val_dim = Hv * dv;
|
||
const int D = 2 * key_dim + val_dim;
|
||
const int gqa_factor = (Hv > Hk) ? (Hv / Hk) : 1;
|
||
|
||
auto module = makeModule(Hk, Hv, dk, dv);
|
||
if (!module) {
|
||
MNN_PRINT("FusedGatedDeltaTest[%s]: module creation failed\n", cs.name);
|
||
return false;
|
||
}
|
||
|
||
auto convWVar = _Input({D, 1, K_conv}, NCHW, halide_type_of<float>());
|
||
{
|
||
float* w = convWVar->writeMap<float>();
|
||
std::uniform_real_distribution<float> dist(-0.15f, 0.15f);
|
||
for (int i = 0; i < D * K_conv; ++i)
|
||
w[i] = dist(rng);
|
||
}
|
||
|
||
// Reference state — module starts from zero on first call.
|
||
std::vector<float> refConvState(B * D * convStateSize, 0.0f);
|
||
std::vector<float> refS(B * Hv * dk * dv, 0.0f);
|
||
|
||
for (int step = 0; step < cs.numSteps; ++step) {
|
||
auto qkvVar = _Input({B, D, L}, NCHW, halide_type_of<float>());
|
||
auto gateVar = _Input({B, L, Hv}, NCHW, halide_type_of<float>());
|
||
auto betaVar = _Input({B, L, Hv}, NCHW, halide_type_of<float>());
|
||
std::uniform_real_distribution<float> qkvDist(-0.3f, 0.3f);
|
||
{
|
||
float* qkv = qkvVar->writeMap<float>();
|
||
for (int i = 0; i < B * D * L; ++i)
|
||
qkv[i] = qkvDist(rng);
|
||
float* g = gateVar->writeMap<float>();
|
||
for (int i = 0; i < B * L * Hv; ++i)
|
||
g[i] = -0.1f - 0.05f * (i % 3);
|
||
float* b = betaVar->writeMap<float>();
|
||
for (int i = 0; i < B * L * Hv; ++i)
|
||
b[i] = 0.4f + 0.1f * (i % 4);
|
||
}
|
||
|
||
// ── Reference path ──
|
||
std::vector<float> refOut(B * L * Hv * dv, 0.0f);
|
||
const float* convWPtr = convWVar->readMap<float>();
|
||
const float* qkvPtr = qkvVar->readMap<float>();
|
||
const float* gPtr = gateVar->readMap<float>();
|
||
const float* bPtr = betaVar->readMap<float>();
|
||
|
||
// Conv1D + SiLU across all L tokens, channel-by-channel.
|
||
// convOut layout matches the runtime: [B, D, L] (channel-major within batch).
|
||
std::vector<float> convOut(B * D * L, 0.0f);
|
||
for (int b = 0; b < B; ++b) {
|
||
for (int d = 0; d < D; ++d) {
|
||
float* mState = refConvState.data() + (b * D + d) * convStateSize;
|
||
const float* w = convWPtr + d * K_conv;
|
||
for (int l = 0; l < L; ++l) {
|
||
float xnew = qkvPtr[(b * D + d) * L + l];
|
||
float sum = xnew * w[convStateSize];
|
||
for (int kk = 0; kk < convStateSize; ++kk)
|
||
sum += mState[kk] * w[kk];
|
||
float sig = 1.0f / (1.0f + std::exp(-sum));
|
||
convOut[(b * D + d) * L + l] = sum * sig;
|
||
// Shift state and append xnew.
|
||
for (int kk = 0; kk < convStateSize - 1; ++kk)
|
||
mState[kk] = mState[kk + 1];
|
||
mState[convStateSize - 1] = xnew;
|
||
}
|
||
}
|
||
}
|
||
|
||
// gated_delta_rule across all timesteps and heads (GQA: k_head = h / gqa_factor).
|
||
for (int b = 0; b < B; ++b) {
|
||
for (int t = 0; t < L; ++t) {
|
||
for (int h = 0; h < Hv; ++h) {
|
||
const int k_head = h / gqa_factor;
|
||
std::vector<float> qLocal(dk), kLocal(dk), vLocal(dv);
|
||
for (int i = 0; i < dk; ++i) {
|
||
qLocal[i] = convOut[(b * D + k_head * dk + i) * L + t];
|
||
kLocal[i] = convOut[(b * D + key_dim + k_head * dk + i) * L + t];
|
||
}
|
||
for (int i = 0; i < dv; ++i) {
|
||
vLocal[i] = convOut[(b * D + 2 * key_dim + h * dv + i) * L + t];
|
||
}
|
||
applyL2NormAndScale(qLocal.data(), kLocal.data(), dk);
|
||
float decay = std::exp(gPtr[b * L * Hv + t * Hv + h]);
|
||
float beta_t = bPtr[b * L * Hv + t * Hv + h];
|
||
float* state = refS.data() + (b * Hv + h) * dk * dv;
|
||
float* outSlot = refOut.data() + ((b * L + t) * Hv + h) * dv;
|
||
refOneStep(state, kLocal.data(), qLocal.data(), vLocal.data(), outSlot, decay, beta_t, dk, dv);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── Module path ──
|
||
auto outputs = module->onForward({qkvVar, gateVar, betaVar, convWVar});
|
||
if (outputs.empty()) {
|
||
MNN_PRINT("FusedGatedDeltaTest[%s]: empty output at step %d\n", cs.name, step);
|
||
return false;
|
||
}
|
||
const float* res = outputs[0]->readMap<float>();
|
||
const int N = B * L * Hv * dv;
|
||
for (int i = 0; i < N; ++i) {
|
||
float diff = std::fabs(res[i] - refOut[i]);
|
||
if (diff > tolerance) {
|
||
MNN_PRINT(
|
||
"FusedGatedDeltaTest[%s] step %d MISMATCH idx=%d "
|
||
"ref=%.6f got=%.6f diff=%.4e (tol=%.4e)\n",
|
||
cs.name, step, i, refOut[i], res[i], diff, tolerance);
|
||
return false;
|
||
}
|
||
}
|
||
}
|
||
MNN_PRINT("FusedGatedDeltaTest[%s] Hk=%d Hv=%d dk=%d dv=%d L=%d × %d steps PASSED\n", cs.name, Hk, Hv, dk, dv, L,
|
||
cs.numSteps);
|
||
return true;
|
||
}
|
||
|
||
} // anonymous namespace
|
||
|
||
class FusedGatedDeltaTest : public MNNTestCase {
|
||
public:
|
||
virtual bool run(int precision) {
|
||
std::mt19937 rng(0x5A17u);
|
||
// Coverage matrix:
|
||
// - decode (L=1): single-head, multi-head, GQA, all production dk/dv
|
||
// - prefill (L>1): multi-head + GQA at dk=128/dv=128 to exercise the
|
||
// gated_delta_rule_mnn path with the shared per-thread buffer
|
||
std::vector<Case> cases = {
|
||
// {name, Hk, Hv, dk, dv, L, steps}
|
||
{"decode_1h_dk64_dv64", 1, 1, 64, 64, 1, 4},
|
||
{"decode_1h_dk64_dv128", 1, 1, 64, 128, 1, 4},
|
||
{"decode_1h_dk128_dv64", 1, 1, 128, 64, 1, 4},
|
||
{"decode_1h_dk128_dv128", 1, 1, 128, 128, 1, 4}, // ← Qwen3-Next shape
|
||
{"decode_4h_dk128_dv128", 4, 4, 128, 128, 1, 3}, // multi-head decode
|
||
{"decode_gqa2_dk128_dv128", 2, 4, 128, 128, 1, 3}, // GQA 2:1 decode
|
||
{"prefill_1h_dk128_dv128", 1, 1, 128, 128, 8, 2}, // L>1 single head
|
||
{"prefill_4h_dk128_dv128", 4, 4, 128, 128, 8, 2}, // L>1 multi-head
|
||
{"prefill_gqa2_dk128_dv128", 2, 4, 128, 128, 8, 2}, // L>1 + GQA
|
||
};
|
||
|
||
// fp16 path accumulates more round-off — loosen tolerance for low-precision.
|
||
float tol = (precision == BackendConfig::Precision_Low) ? 6e-2f : 5e-3f;
|
||
for (auto& cs : cases) {
|
||
if (!runCase(cs, rng, tol))
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
};
|
||
|
||
MNNTestSuiteRegister(FusedGatedDeltaTest, "op/fused_gated_delta");
|
||
|
||
#endif // MNN_SUPPORT_TRANSFORMER_FUSE
|