// Fused packed W4A32 reference execution. No full decoded weight buffer.
// Sparse connectivity reduces dot products; INT4 alone does not reduce MAC count.
#ifdef __APPLE__
#define ACCELERATE_NEW_LAPACK
#define ACCELERATE_LAPACK_ILP64
#include <Accelerate/Accelerate.h>
#else
#include <cblas.h>
#endif
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <vector>
#if defined(__aarch64__) || defined(__arm64__)
#include <arm_neon.h>
#endif

struct Layer {
    int din, dout, width, group, rank, mode;
    const float *center, *unit, *bias, *weight, *scales, *a, *b;
    const uint8_t *codes;
    const uint32_t *indices;
};

static bool zero_center(const Layer& l) {
    if (!l.center) return true;
    for (int i = 0; i < l.din; ++i) if (l.center[i] != 0.0f) return false;
    return true;
}

extern "C" int relative_integer_dot_supported() {
#ifdef __ARM_FEATURE_DOTPROD
    return 1;
#else
    return 0;
#endif
}

extern "C" int relative_gated_delta_supported() {
#if defined(__aarch64__) || defined(__arm64__)
    return 1;
#else
    return 0;
#endif
}

// CPU streaming equivalent of mlx_lm.models.gated_delta.gated_delta_ops.
// State rows stay local to one (batch, head, value-dimension) sequence row.
extern "C" int relative_gated_delta(const float* q, const float* k,
                                     const float* v, const float* g,
                                     const float* beta, const float* state_in,
                                     const uint8_t* mask, int batch, int time,
                                     int hk_count, int hv_count, int dk, int dv,
                                     int vector_gate, int has_mask, float* y,
                                     float* state_out) {
    if (batch < 1 || time < 1 || hk_count < 1 || hv_count < 1 || dk < 1 || dv < 1 ||
        hv_count % hk_count || (vector_gate != 0 && vector_gate != 1) ||
        (has_mask != 0 && has_mask != 1) || !q || !k || !v || !g || !beta ||
        !state_in || !y || !state_out || (has_mask && !mask)) return 1;
    if (!relative_gated_delta_supported()) return 2;

#if !defined(__aarch64__) && !defined(__arm64__)
    return 2;
#else
    const int repeat = hv_count / hk_count;
    const size_t state_row_size = static_cast<size_t>(dk);
    std::vector<float> row(state_row_size), candidate(state_row_size);
    for (int b = 0; b < batch; ++b) {
        for (int h = 0; h < hv_count; ++h) {
            const int hk = h / repeat;
            for (int d = 0; d < dv; ++d) {
                const size_t state_offset =
                    ((static_cast<size_t>(b) * hv_count + h) * dv + d) * dk;
                std::copy(state_in + state_offset, state_in + state_offset + dk, row.begin());
                for (int t = 0; t < time; ++t) {
                    const size_t q_offset =
                        ((static_cast<size_t>(b) * time + t) * hk_count + hk) * dk;
                    const size_t v_offset =
                        (static_cast<size_t>(b) * time + t) * hv_count * dv +
                        static_cast<size_t>(h) * dv + d;
                    const size_t gate_offset =
                        (static_cast<size_t>(b) * time + t) * hv_count + h;
                    const float* kt = k + q_offset;
                    const float* qt = q + q_offset;
                    float kv_mem = 0.0f;
                    float32x4_t kv0 = vdupq_n_f32(0.0f), kv1 = kv0;
                    float32x4_t kv2 = kv0, kv3 = kv0, kv_tail = kv0;
                    int i = 0;
                    for (; i + 16 <= dk; i += 16) {
                        float32x4_t gate0, gate1, gate2, gate3;
                        if (vector_gate) {
                            const float* gt = g + gate_offset * dk + i;
                            gate0 = vld1q_f32(gt);
                            gate1 = vld1q_f32(gt + 4);
                            gate2 = vld1q_f32(gt + 8);
                            gate3 = vld1q_f32(gt + 12);
                        } else {
                            const float32x4_t gate = vdupq_n_f32(g[gate_offset]);
                            gate0 = gate1 = gate2 = gate3 = gate;
                        }
                        const float* rt = row.data() + i;
                        const float32x4_t c0 = vmulq_f32(vld1q_f32(rt), gate0);
                        const float32x4_t c1 = vmulq_f32(vld1q_f32(rt + 4), gate1);
                        const float32x4_t c2 = vmulq_f32(vld1q_f32(rt + 8), gate2);
                        const float32x4_t c3 = vmulq_f32(vld1q_f32(rt + 12), gate3);
                        float* ct = candidate.data() + i;
                        vst1q_f32(ct, c0); vst1q_f32(ct + 4, c1);
                        vst1q_f32(ct + 8, c2); vst1q_f32(ct + 12, c3);
                        kv0 = vaddq_f32(kv0, vmulq_f32(c0, vld1q_f32(kt + i)));
                        kv1 = vaddq_f32(kv1, vmulq_f32(c1, vld1q_f32(kt + i + 4)));
                        kv2 = vaddq_f32(kv2, vmulq_f32(c2, vld1q_f32(kt + i + 8)));
                        kv3 = vaddq_f32(kv3, vmulq_f32(c3, vld1q_f32(kt + i + 12)));
                    }
                    for (; i + 4 <= dk; i += 4) {
                        float32x4_t gate;
                        if (vector_gate) gate = vld1q_f32(g + gate_offset * dk + i);
                        else gate = vdupq_n_f32(g[gate_offset]);
                        const float32x4_t decayed =
                            vmulq_f32(vld1q_f32(row.data() + i), gate);
                        vst1q_f32(candidate.data() + i, decayed);
                        kv_tail = vaddq_f32(kv_tail,
                            vmulq_f32(decayed, vld1q_f32(kt + i)));
                    }
                    kv_mem = vaddvq_f32(vaddq_f32(vaddq_f32(kv0, kv1),
                                                  vaddq_f32(kv2, kv3))) +
                                vaddvq_f32(kv_tail);
                    for (; i < dk; ++i) {
                        const float gate = g[vector_gate ? gate_offset * dk + i : gate_offset];
                        candidate[i] = row[i] * gate;
                        kv_mem += candidate[i] * kt[i];
                    }
                    const float delta = (v[v_offset] - kv_mem) * beta[gate_offset];
                    float32x4_t out0 = vdupq_n_f32(0.0f), out1 = out0;
                    float32x4_t out2 = out0, out3 = out0, out_tail = out0;
                    i = 0;
                    const float32x4_t delta4 = vdupq_n_f32(delta);
                    for (; i + 16 <= dk; i += 16) {
                        float* ct = candidate.data() + i;
                        const float32x4_t c0 = vaddq_f32(vld1q_f32(ct),
                            vmulq_f32(vld1q_f32(kt + i), delta4));
                        const float32x4_t c1 = vaddq_f32(vld1q_f32(ct + 4),
                            vmulq_f32(vld1q_f32(kt + i + 4), delta4));
                        const float32x4_t c2 = vaddq_f32(vld1q_f32(ct + 8),
                            vmulq_f32(vld1q_f32(kt + i + 8), delta4));
                        const float32x4_t c3 = vaddq_f32(vld1q_f32(ct + 12),
                            vmulq_f32(vld1q_f32(kt + i + 12), delta4));
                        vst1q_f32(ct, c0); vst1q_f32(ct + 4, c1);
                        vst1q_f32(ct + 8, c2); vst1q_f32(ct + 12, c3);
                        out0 = vaddq_f32(out0, vmulq_f32(c0, vld1q_f32(qt + i)));
                        out1 = vaddq_f32(out1, vmulq_f32(c1, vld1q_f32(qt + i + 4)));
                        out2 = vaddq_f32(out2, vmulq_f32(c2, vld1q_f32(qt + i + 8)));
                        out3 = vaddq_f32(out3, vmulq_f32(c3, vld1q_f32(qt + i + 12)));
                    }
                    for (; i + 4 <= dk; i += 4) {
                        const float32x4_t c = vaddq_f32(vld1q_f32(candidate.data() + i),
                            vmulq_f32(vld1q_f32(kt + i), delta4));
                        vst1q_f32(candidate.data() + i, c);
                        out_tail = vaddq_f32(out_tail,
                            vmulq_f32(c, vld1q_f32(qt + i)));
                    }
                    float out = vaddvq_f32(vaddq_f32(vaddq_f32(out0, out1),
                                                     vaddq_f32(out2, out3))) +
                                vaddvq_f32(out_tail);
                    for (; i < dk; ++i) {
                        candidate[i] += kt[i] * delta;
                        out += candidate[i] * qt[i];
                    }
                    const size_t y_offset =
                        (static_cast<size_t>(b) * time + t) * hv_count * dv +
                        static_cast<size_t>(h) * dv + d;
                    y[y_offset] = out;
                    if (!has_mask || mask[static_cast<size_t>(b) * time + t])
                        row.swap(candidate);
                }
                std::copy(row.begin(), row.end(), state_out + state_offset);
            }
        }
    }
    return 0;
#endif
}

#ifdef __ARM_FEATURE_DOTPROD
// Decode одного токена: четыре независимых аккумулятора без внутреннего
// динамического цикла по выходам, с тем же порядком сумм по группам.
static void integer_single_int8(const Layer& l, const int8_t* x,
                                const float* scales, float* y) {
    const int groups=(l.din+l.group-1)/l.group;
    int j=0;
    for (; j+4<=l.dout; j+=4) {
        const auto* w0=reinterpret_cast<const int8_t*>(l.codes)+j*l.din;
        const auto* w1=w0+l.din;
        const auto* w2=w1+l.din;
        const auto* w3=w2+l.din;
        float sums[4]={y[j],y[j+1],y[j+2],y[j+3]};
        for (int g=0; g<groups; ++g) {
            int k=g*l.group, end=std::min(k+l.group,l.din);
            int32x4_t a0=vdupq_n_s32(0), a1=a0, a2=a0, a3=a0;
            for (; k+16<=end; k+=16) {
                const auto a=vld1q_s8(x+k);
                a0=vdotq_s32(a0,vld1q_s8(w0+k),a);
                a1=vdotq_s32(a1,vld1q_s8(w1+k),a);
                a2=vdotq_s32(a2,vld1q_s8(w2+k),a);
                a3=vdotq_s32(a3,vld1q_s8(w3+k),a);
            }
            int s0=vaddvq_s32(a0), s1=vaddvq_s32(a1);
            int s2=vaddvq_s32(a2), s3=vaddvq_s32(a3);
            for (; k<end; ++k) {
                s0+=w0[k]*x[k]; s1+=w1[k]*x[k];
                s2+=w2[k]*x[k]; s3+=w3[k]*x[k];
            }
            sums[0]+=s0*(scales[g]*l.scales[j*groups+g]);
            sums[1]+=s1*(scales[g]*l.scales[(j+1)*groups+g]);
            sums[2]+=s2*(scales[g]*l.scales[(j+2)*groups+g]);
            sums[3]+=s3*(scales[g]*l.scales[(j+3)*groups+g]);
        }
        y[j]=sums[0]; y[j+1]=sums[1]; y[j+2]=sums[2]; y[j+3]=sums[3];
    }
    for (; j<l.dout; ++j) {
        const auto* w=reinterpret_cast<const int8_t*>(l.codes)+j*l.din;
        float sum_out=y[j];
        for (int g=0; g<groups; ++g) {
            int k=g*l.group, end=std::min(k+l.group,l.din);
            auto a=vdupq_n_s32(0);
            for (; k+16<=end; k+=16) a=vdotq_s32(a,vld1q_s8(w+k),vld1q_s8(x+k));
            int sum=vaddvq_s32(a);
            for (; k<end; ++k) sum+=w[k]*x[k];
            sum_out+=sum*(scales[g]*l.scales[j*groups+g]);
        }
        y[j]=sum_out;
    }
}

// Четыре выхода используют общую загрузку активаций; пакет делит загрузки весов.
template<int Rows>
static void integer_base_quantized(const Layer& l, const int8_t* activations,
                                   const float* scales, float* y) {
    if constexpr (Rows==1) {
        if (l.mode==3) {
            integer_single_int8(l,activations,scales,y);
            return;
        }
    }
    const int groups=(l.din+l.group-1)/l.group;
    const int stride=l.mode==3 ? l.din : (l.din+1)/2;
    for (int j=0; j<l.dout; j+=4) {
        const int outputs=std::min(4,l.dout-j);
        for (int g=0; g<groups; ++g) {
            int k=g*l.group, end=std::min(k+l.group,l.din);
            int32x4_t accumulator[Rows][4];
            int tails[Rows][4]={};
            for (int row=0; row<Rows; ++row)
                for (int out=0; out<4; ++out) accumulator[row][out]=vdupq_n_s32(0);
            if (l.mode!=3 && (k&1)) {
                for (int out=0; out<outputs; ++out) {
                    int q=int(l.codes[(j+out)*stride+k/2]>>4)-7;
                    for (int row=0; row<Rows; ++row) tails[row][out]+=q*activations[row*l.din+k];
                }
                ++k;
            }
            for (; k+16<=end; k+=16) {
                int8x16_t a[Rows];
                for (int row=0; row<Rows; ++row) a[row]=vld1q_s8(activations+row*l.din+k);
                for (int out=0; out<outputs; ++out) {
                    const uint8_t* weights=l.codes+(j+out)*stride;
                    int8x16_t w;
                    if (l.mode==3) w=vld1q_s8(reinterpret_cast<const int8_t*>(weights+k));
                    else {
                        uint8x8_t packed=vld1_u8(weights+k/2);
                        uint8x8x2_t q=vzip_u8(vand_u8(packed,vdup_n_u8(15)),vshr_n_u8(packed,4));
                        w=vsubq_s8(vreinterpretq_s8_u8(vcombine_u8(q.val[0],q.val[1])),vdupq_n_s8(7));
                    }
                    for (int row=0; row<Rows; ++row) accumulator[row][out]=vdotq_s32(accumulator[row][out],w,a[row]);
                }
            }
            for (; k<end; ++k)
                for (int out=0; out<outputs; ++out) {
                    const uint8_t* weights=l.codes+(j+out)*stride;
                    int q=l.mode==3 ? reinterpret_cast<const int8_t*>(weights)[k] : int((weights[k/2]>>((k&1)*4))&15)-7;
                    for (int row=0; row<Rows; ++row) tails[row][out]+=q*activations[row*l.din+k];
                }
            for (int row=0; row<Rows; ++row)
                for (int out=0; out<outputs; ++out)
                    y[row*l.dout+j+out]+=(vaddvq_s32(accumulator[row][out])+tails[row][out])*
                        (scales[row*groups+g]*l.scales[(j+out)*groups+g]);
        }
    }
}

template<int Rows>
static void integer_base(const Layer& l, const float* x, float* y) {
    const int groups=(l.din+l.group-1)/l.group;
    std::vector<int8_t> activations(Rows*l.din);
    std::vector<float> scales(Rows*groups);
    for (int row=0; row<Rows; ++row)
        for (int g=0; g<groups; ++g) {
            int start=g*l.group, end=std::min(start+l.group,l.din);
            float maximum=0;
            for (int k=start; k<end; ++k) maximum=std::max(maximum,std::abs(x[row*l.din+k]));
            float scale=maximum ? maximum/127 : 1.0f;
            scales[row*groups+g]=scale;
            for (int k=start; k<end; ++k)
                activations[row*l.din+k]=int8_t(std::max(-127.0f,std::min(127.0f,std::round(x[row*l.din+k]/scale))));
        }
    integer_base_quantized<Rows>(l,activations.data(),scales.data(),y);
}
#endif

static void linear(const Layer& l, const float* x, float* y) {
    std::vector<float> normalized;
    const float* r = x;
    if (l.center) {
        normalized.resize(l.din);
        for (int i = 0; i < l.din; ++i) normalized[i] = (x[i] - l.center[i]) * l.unit[i];
        r = normalized.data();
    }
    std::copy(l.bias, l.bias + l.dout, y);
    if (l.mode >= 2) {
#ifdef __ARM_FEATURE_DOTPROD
        integer_base<1>(l, r, y);
#endif
    } else if (l.mode == 0) {
        if (l.indices) {
            for (int j = 0; j < l.dout; ++j)
                for (int k = 0; k < l.width; ++k)
                    y[j] += r[l.indices[j*l.width+k]] * l.weight[j*l.width+k];
        } else cblas_sgemv(CblasRowMajor, CblasNoTrans, l.dout, l.din, 1.0f,
                    l.weight, l.din, r, 1, 1.0f, y, 1);
    } else {
        int stride = (l.width + 1) / 2;
        int groups = (l.width + l.group - 1) / l.group;
        for (int j = 0; j < l.dout; ++j) {
            const uint8_t* row = l.codes + j * stride;
            for (int g = 0; g < groups; ++g) {
                float sum = 0;
                int end = std::min((g + 1) * l.group, l.width);
                int k = g * l.group;
#ifdef __aarch64__
                // SIMD packed decode and FMA, with scalar tails for odd groups.
                if (k & 1) {
                    int q = (row[k/2] >> 4) - 7;
                    int i = l.indices ? l.indices[j*l.width+k] : k;
                    sum += r[i]*q; ++k;
                }
                float32x4_t acc = vdupq_n_f32(0);
                for (; k + 16 <= end; k += 16) {
                    uint8x8_t packed = vld1_u8(row+k/2);
                    uint8x8x2_t nibbles = vzip_u8(vand_u8(packed,vdup_n_u8(15)),vshr_n_u8(packed,4));
                    for (int h = 0; h < 2; ++h) {
                        int16x8_t q = vmovl_s8(vsub_s8(vreinterpret_s8_u8(nibbles.val[h]),vdup_n_s8(7)));
                        float32x4_t low, high;
                        if (l.indices) {
                            float values[8];
                            for (int z = 0; z < 8; ++z)
                                values[z] = r[l.indices[j*l.width+k+h*8+z]];
                            low = vld1q_f32(values);
                            high = vld1q_f32(values+4);
                        } else {
                            low = vld1q_f32(r+k+h*8);
                            high = vld1q_f32(r+k+h*8+4);
                        }
                        acc = vfmaq_f32(acc,low,vcvtq_f32_s32(vmovl_s16(vget_low_s16(q))));
                        acc = vfmaq_f32(acc,high,vcvtq_f32_s32(vmovl_s16(vget_high_s16(q))));
                    }
                }
                sum += vaddvq_f32(acc);
#endif
                for (; k < end; ++k) {
                    int q = ((row[k / 2] >> ((k & 1) * 4)) & 15) - 7;
                    int i = l.indices ? l.indices[j * l.width + k] : k;
                    sum += r[i] * q;
                }
                y[j] += sum * l.scales[j * groups + g];
            }
        }
    }
    if (l.rank) {
        std::vector<float> h(l.rank);
        cblas_sgemv(CblasRowMajor, CblasNoTrans, l.rank, l.din, 1.0f,
                    l.a, l.din, r, 1, 0.0f, h.data(), 1);
        cblas_sgemv(CblasRowMajor, CblasNoTrans, l.dout, l.rank, 1.0f,
                    l.b, l.rank, h.data(), 1, 1.0f, y, 1);
    }
}

extern "C" int relative_lookup_rows(const Layer* layer, const int32_t* rows,
                                     int count, float* output) {
    if (!layer || count < 0 || (count && (!rows || !output))) return 1;
    const auto& l = *layer;
    if (l.din < 1 || l.dout < 1 || l.group < 1 || l.rank || !zero_center(l) ||
        l.indices || (l.mode != 1 && l.mode != 2 && l.mode != 3)) return 2;
    if (!l.codes || !l.scales) return 3;
    const int groups = (l.din + l.group - 1) / l.group;
    const int stride = l.mode == 3 ? l.din : (l.din + 1) / 2;
    for (int n = 0; n < count; ++n) {
        const int j = rows[n];
        if (j < 0 || j >= l.dout) return 4;
        for (int k = 0; k < l.din; ++k) {
            float q;
            if (l.mode == 3) {
                q = reinterpret_cast<const int8_t*>(l.codes + j*stride)[k];
            } else {
                q = ((l.codes[j*stride+k/2] >> ((k&1)*4)) & 15) - 7;
            }
            const float unit = l.unit ? l.unit[k] : 1.0f;
            output[n*l.din+k] = q*l.scales[j*groups+k/l.group]*unit;
        }
    }
    return 0;
}

extern "C" int relative_forward_many(const Layer* layers, int count,
                                      const float* input, float** outputs,
                                      int rows) {
    if (!layers || count < 1 || rows < 0 || (rows && (!input || !outputs))) return 1;
    if (!relative_integer_dot_supported()) return 5;
    const auto& first = layers[0];
    if (first.din < 1 || first.group < 1 || first.rank || !zero_center(first) ||
        (first.mode != 2 && first.mode != 3)) return 2;
    const int groups = (first.din + first.group - 1) / first.group;
    for (int k = 0; k < count; ++k) {
        const auto& l = layers[k];
        if (l.din != first.din || l.group != first.group || l.dout < 1 ||
            l.rank || !zero_center(l) || l.indices || !l.bias || !l.codes || !l.scales ||
            (l.mode != 2 && l.mode != 3) ||
            std::min(l.group,l.din) > INT32_MAX/(127*127) || (rows && !outputs[k])) return 3;
        for (int i = 0; i < first.din; ++i) {
            const float a = first.unit ? first.unit[i] : 1.0f;
            const float b = l.unit ? l.unit[i] : 1.0f;
            if (a != b) return 4;
        }
    }
    std::vector<int8_t> activations(rows*first.din);
    std::vector<float> activation_scales(rows*groups);
    for (int row = 0; row < rows; ++row) {
        for (int g = 0; g < groups; ++g) {
            const int start = g*first.group, end = std::min(start+first.group, first.din);
            float maximum = 0.0f;
            for (int i = start; i < end; ++i) {
                const float unit = first.unit ? first.unit[i] : 1.0f;
                maximum = std::max(maximum, std::abs(input[row*first.din+i]*unit));
            }
            const float scale = maximum ? maximum/127.0f : 1.0f;
            activation_scales[row*groups+g] = scale;
            for (int i = start; i < end; ++i) {
                const float unit = first.unit ? first.unit[i] : 1.0f;
                activations[row*first.din+i] = int8_t(std::max(-127.0f, std::min(127.0f,
                    std::round(input[row*first.din+i]*unit/scale))));
            }
        }
    }
    for (int n = 0; n < count; ++n) {
        const auto& l = layers[n];
        for (int row = 0; row < rows; ++row)
            std::copy(l.bias, l.bias+l.dout, outputs[n]+row*l.dout);
#ifdef __ARM_FEATURE_DOTPROD
        int row = 0;
        for (; row+4 <= rows; row += 4)
            integer_base_quantized<4>(l,activations.data()+row*l.din,
                activation_scales.data()+row*groups,outputs[n]+row*l.dout);
        for (; row < rows; ++row)
            integer_base_quantized<1>(l,activations.data()+row*l.din,
                activation_scales.data()+row*groups,outputs[n]+row*l.dout);
#endif
    }
    return 0;
}

extern "C" int relative_forward(const Layer* layers, int count, const float* input,
                                 float* output, int rows) {
    if (!layers || count < 1 || rows < 0) return 1;
    for (int k = 0; k < count; ++k) {
        const auto& l = layers[k];
        if (l.din < 1 || l.dout < 1 || l.width < 1 || l.group < 1 ||
            (k && layers[k-1].dout != l.din)) return 2;
        if (l.mode>=2 && (!relative_integer_dot_supported() || l.indices)) return 3;
    }
    // Prefill uses the same optimized dense BLAS path as ordinary FP32 models.
    if (count == 1 && rows > 1 && !layers[0].mode && !layers[0].indices &&
        !layers[0].center && !layers[0].rank) {
        const auto& l = layers[0];
        for (int row = 0; row < rows; ++row)
            std::copy(l.bias, l.bias+l.dout, output+row*l.dout);
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, rows, l.dout,
                    l.din, 1.0f, input, l.din, l.weight, l.din,
                    1.0f, output, l.dout);
        return 0;
    }
    if (count==1 && rows>1 && layers[0].mode>=2) {
#ifdef __ARM_FEATURE_DOTPROD
        const auto& l=layers[0];
        std::vector<float> normalized;
        const float* x=input;
        if (l.center) {
            normalized.resize(rows*l.din);
            for (int row=0; row<rows; ++row)
                for (int k=0; k<l.din; ++k)
                    normalized[row*l.din+k]=(input[row*l.din+k]-l.center[k])*l.unit[k];
            x=normalized.data();
        }
        for (int row=0; row<rows; ++row) std::copy(l.bias,l.bias+l.dout,output+row*l.dout);
        int row=0;
        for (; row+4<=rows; row+=4) integer_base<4>(l,x+row*l.din,output+row*l.dout);
        for (; row<rows; ++row) integer_base<1>(l,x+row*l.din,output+row*l.dout);
        if (l.rank) {
            std::vector<float> correction(rows*l.rank);
            cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasTrans,rows,l.rank,l.din,1.0f,x,l.din,l.a,l.din,0.0f,correction.data(),l.rank);
            cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasTrans,rows,l.dout,l.rank,1.0f,correction.data(),l.rank,l.b,l.rank,1.0f,output,l.dout);
        }
        return 0;
#endif
    }
    if (count == 1 && rows > 1 && layers[0].mode==1 && !layers[0].indices) {
        const auto& l = layers[0];
        std::vector<float> normalized;
        const float* x = input;
        if (l.center) {
            normalized.resize(rows*l.din);
            for (int row = 0; row < rows; ++row)
                for (int k = 0; k < l.din; ++k)
                    normalized[row*l.din+k] = (input[row*l.din+k]-l.center[k])*l.unit[k];
            x = normalized.data();
        }
        for (int row = 0; row < rows; ++row)
            std::copy(l.bias, l.bias+l.dout, output+row*l.dout);
        // Only a bounded input tile is decoded, never the full weight matrix.
        const int tile_width = std::min(32, l.din);
        std::vector<float> tile(tile_width*l.dout);
        const int stride = (l.din+1)/2;
        const int groups = (l.din+l.group-1)/l.group;
        for (int start = 0; start < l.din; start += tile_width) {
            int width = std::min(tile_width, l.din-start);
            for (int j = 0; j < l.dout; ++j)
                for (int k = 0; k < width; ++k) {
                    int ix = start+k;
                    int q = ((l.codes[j*stride+ix/2] >> ((ix&1)*4))&15)-7;
                    tile[j*width+k] = q*l.scales[j*groups+ix/l.group];
                }
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, rows, l.dout,
                        width, 1.0f, x+start, l.din, tile.data(), width,
                        1.0f, output, l.dout);
        }
        if (l.rank) {
            std::vector<float> correction(rows*l.rank);
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, rows, l.rank,
                        l.din, 1.0f, x, l.din, l.a, l.din,
                        0.0f, correction.data(), l.rank);
            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, rows, l.dout,
                        l.rank, 1.0f, correction.data(), l.rank, l.b, l.rank,
                        1.0f, output, l.dout);
        }
        return 0;
    }
    for (int row = 0; row < rows; ++row) {
        const float* x = input + row * layers[0].din;
        std::vector<float> current, next;
        for (int k = 0; k < count; ++k) {
            next.resize(layers[k].dout);
            linear(layers[k], x, next.data());
            if (k + 1 < count) for (float& v : next) v = std::max(0.0f, v);
            current.swap(next);
            x = current.data();
        }
        std::copy(current.begin(), current.end(), output + row * layers[count-1].dout);
    }
    return 0;
}