// ZayaGLA CPU inference engine — MTP + ExpertChoice + RL + VarDepth + INT8 KV + SelfDiff
#include "yasha.h"
#include <filesystem>
#include <iostream>
#include <unordered_set>

namespace fs = std::filesystem;

// === INT8 quantization ===
void TensorQ8::quantize(const float* src, int n) {
    q.resize(n);
    int bs = block_size;
    int n_blocks = (n + bs - 1) / bs;
    scales.resize(n_blocks);
    for (int b = 0; b < n_blocks; b++) {
        int start = b * bs;
        int end = std::min(start + bs, n);
        float max_abs = 0;
        for (int i = start; i < end; i++) max_abs = std::max(max_abs, std::abs(src[i]));
        scales[b] = max_abs / 127.0f;
        if (scales[b] < 1e-10f) scales[b] = 1e-10f;
        for (int i = start; i < end; i++)
            q[i] = (int8_t)std::round(src[i] / scales[b]);
    }
}

void TensorQ8::dequantize(float* dst, int n) const {
    int bs = block_size;
    for (int b = 0; b < (int)scales.size(); b++) {
        int start = b * bs;
        int end = std::min(start + bs, n);
        float s = scales[b];
        for (int i = start; i < end; i++)
            dst[i] = (float)q[i] * s;
    }
}

// === Q2 (2-bit) quantization ===
void TensorQ2::quantize(const float* src, int n) {
    int bs = block_size;
    int n_blocks = (n + bs - 1) / bs;
    scales.resize(n_blocks);
    q.resize((n + 3) / 4);
    for (int b = 0; b < n_blocks; b++) {
        int start = b * bs;
        int end = std::min(start + bs, n);
        float max_abs = 0;
        for (int i = start; i < end; i++) max_abs = std::max(max_abs, std::abs(src[i]));
        scales[b] = max_abs / 1.5f;  // 2-bit range: [-1.5, 1.5] mapped to [0, 3]
        if (scales[b] < 1e-10f) scales[b] = 1e-10f;
        for (int i = start; i < end; i++) {
            int idx = b * bs + i;
            int val = std::max(0, std::min(3, (int)std::round(src[i] / scales[b] + 1.5f)));
            int byte_idx = idx / 4;
            int shift = (idx % 4) * 2;
            q[byte_idx] = (q[byte_idx] & ~(3 << shift)) | (val << shift);
        }
    }
}

void TensorQ2::dequantize_block(float* dst, int block_idx) const {
    int bs = block_size;
    float s = scales[block_idx];
    int start = block_idx * bs;
    for (int i = 0; i < bs; i++) {
        int idx = start + i;
        int byte_idx = idx / 4;
        int shift = (idx % 4) * 2;
        int val = (q[byte_idx] >> shift) & 3;
        dst[i] = ((float)val - 1.5f) * s;
    }
}

// === Q3 (3-bit) quantization ===
void TensorQ3::quantize(const float* src, int n) {
    int bs = block_size;
    int n_blocks = (n + bs - 1) / bs;
    scales.resize(n_blocks);
    q.resize((n * 3 + 7) / 8);
    for (int b = 0; b < n_blocks; b++) {
        int start = b * bs;
        int end = std::min(start + bs, n);
        float max_abs = 0;
        for (int i = start; i < end; i++) max_abs = std::max(max_abs, std::abs(src[i]));
        scales[b] = max_abs / 3.5f;  // 3-bit range: [-3.5, 3.5] mapped to [0, 7]
        if (scales[b] < 1e-10f) scales[b] = 1e-10f;
        for (int i = start; i < end; i++) {
            int idx = b * bs + i;
            int val = std::max(0, std::min(7, (int)std::round(src[i] / scales[b] + 3.5f)));
            // Pack 8×3-bit = 24 bits = 3 bytes
            int byte_idx = (idx * 3) / 8;
            int bit_ofs = (idx * 3) % 8;
            if (bit_ofs <= 5) {
                q[byte_idx] = (q[byte_idx] & ~(7 << bit_ofs)) | (val << bit_ofs);
            } else {
                // Crosses byte boundary
                int low_bits = 8 - bit_ofs;
                q[byte_idx] = (q[byte_idx] & ~((1 << low_bits) - 1)) | (val << bit_ofs);
                q[byte_idx + 1] = (q[byte_idx + 1] & ~((1 << (3 - low_bits)) - 1)) | (val >> low_bits);
            }
        }
    }
}

void TensorQ3::dequantize_block(float* dst, int block_idx) const {
    int bs = block_size;
    float s = scales[block_idx];
    int start = block_idx * bs;
    for (int i = 0; i < bs; i++) {
        int idx = start + i;
        int byte_idx = (idx * 3) / 8;
        int bit_ofs = (idx * 3) % 8;
        int val;
        if (bit_ofs <= 5) {
            val = (q[byte_idx] >> bit_ofs) & 7;
        } else {
            int low_bits = 8 - bit_ofs;
            val = (q[byte_idx] >> bit_ofs) | ((int)q[byte_idx + 1] << low_bits);
            val &= 7;
        }
        dst[i] = ((float)val - 3.5f) * s;
    }
}

// === SIMD-packed weight (8×8 tiles) ===
void PackedWeight::pack(const float* src, int out_d, int in_d) {
    out_dim = out_d; in_dim = in_d;
    out_tiles = (out_d + 7) / 8;
    in_tiles = (in_d + 7) / 8;
    tiles.resize(out_tiles * in_tiles * 64, 0.0f);
    for (int to = 0; to < out_tiles; to++)
        for (int ti = 0; ti < in_tiles; ti++)
            for (int i = 0; i < 8; i++)
                for (int j = 0; j < 8; j++) {
                    int oi = to * 8 + i;
                    int ii = ti * 8 + j;
                    if (oi < out_d && ii < in_d)
                        tiles[(to * in_tiles + ti) * 64 + i * 8 + j] = src[oi * in_d + ii];
                }
}

void PackedWeight::matmul_tiled(float* out, const float* inp, int T, int out_d, int in_d) const {
    int ot = out_tiles, it = in_tiles;
    for (int t = 0; t < T; t++) {
        const float* xp = inp + t * in_d;
        float* op = out + t * out_d;
        std::fill(op, op + out_d, 0.0f);
        for (int to = 0; to < ot; to++) {
            for (int ti = 0; ti < it; ti++) {
                const float* tile = &tiles[(to * it + ti) * 64];
                const float* xi = xp + ti * 8;
                float* oi = op + to * 8;
                __m256 acc0 = _mm256_setzero_ps();
                __m256 acc1 = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= std::min(8, in_d - ti * 8); j += 8) {
                    __m256 xv = _mm256_loadu_ps(&xi[j]);
                    __m256 w0 = _mm256_loadu_ps(&tile[0 * 8 + j]);
                    __m256 w1 = _mm256_loadu_ps(&tile[1 * 8 + j]);
                    __m256 w2 = _mm256_loadu_ps(&tile[2 * 8 + j]);
                    __m256 w3 = _mm256_loadu_ps(&tile[3 * 8 + j]);
                    __m256 w4 = _mm256_loadu_ps(&tile[4 * 8 + j]);
                    __m256 w5 = _mm256_loadu_ps(&tile[5 * 8 + j]);
                    __m256 w6 = _mm256_loadu_ps(&tile[6 * 8 + j]);
                    __m256 w7 = _mm256_loadu_ps(&tile[7 * 8 + j]);
                    acc0 = _mm256_fmadd_ps(xv, w0, acc0);
                    acc1 = _mm256_fmadd_ps(xv, w1, acc1);
                    acc0 = _mm256_fmadd_ps(xv, w2, acc0);
                    acc1 = _mm256_fmadd_ps(xv, w3, acc1);
                    acc0 = _mm256_fmadd_ps(xv, w4, acc0);
                    acc1 = _mm256_fmadd_ps(xv, w5, acc1);
                    acc0 = _mm256_fmadd_ps(xv, w6, acc0);
                    acc1 = _mm256_fmadd_ps(xv, w7, acc1);
                }
                float vals[8] = {
                    hsum_ps(acc0), hsum_ps(acc1), 0, 0, 0, 0, 0, 0
                };
                for (int k = 0; k < 2 && to * 8 + k < out_d; k++)
                    oi[k] += vals[k];
            }
        }
    }
}

// === Constructor ===
YashaModel::YashaModel(const YashaConfig& c) : cfg(c) {
    int H=cfg.hidden_size, D=cfg.head_dim, Nh=cfg.num_heads, T=2048;
    x      = Tensor({T, H});
    x2     = Tensor({T, H});
    gate   = Tensor({T, H});
    up     = Tensor({T, cfg.ffn_hidden});
    down   = Tensor({T, H});
    q      = Tensor({T, Nh, D});
    k      = Tensor({T, 1, D});
    v      = Tensor({T, 1, D});
    g      = Tensor({T, Nh, D});
    attn_out = Tensor({T, H});
    attn_res = Tensor({T, H});
    router = Tensor({T, cfg.num_experts});
    route_w = Tensor({T, cfg.num_experts_per_tok});
    ffn_gate = Tensor({T, cfg.ffn_hidden});
    ffn_up   = Tensor({T, cfg.ffn_hidden});
    ffn_down = Tensor({T, H});
    logits = Tensor({1, cfg.vocab_size});
    diff_buffer = Tensor({T, H});
    conf_hidden = Tensor({cfg.vocab_size});
    conf_scores = Tensor({T});
    for (int k = 0; k < 4; k++)
        mtp_logits[k] = Tensor({1, cfg.vocab_size});
    ec_scores = Tensor({T, cfg.num_experts});
    ec_assign = Tensor({T});
    layer_conf = Tensor({T});
    if (cfg.kv_int8) {
        gla_state_q8 = TensorQ8({Nh, D, D}, cfg.kv_block_size);
    } else {
        gla_state = Tensor({Nh, D, D});
    }
}

// === NF4 dequant ===
void dequant_nf4(Tensor& out, const uint8_t* raw, const float* absmax, int n) {
    int n_blocks = (n + 63) / 64;
    for (int b = 0; b < n_blocks; b++) {
        float scale = absmax[b];
        int remaining = std::min(64, n - b * 64);
        for (int i = 0; i < remaining; i++) {
            int byte_idx = (b * 64 + i) / 2;
            bool hi = i & 1;
            int idx = (raw[byte_idx] >> (hi ? 4 : 0)) & 0x0F;
            out.d[b * 64 + i] = nf4_table_f32[idx] * scale;
        }
    }
}

// === RoPE ===
void YashaModel::rope_partial(float* qp, float* kp, int pos, int D) {
    int hk = cfg.hk;
    if (hk <= 0) return;
    for (int d = 0; d < hk/2; d++) {
        float freq = 1.0f / std::pow(cfg.rope_theta, (2.0f * d) / cfg.head_dim);
        float ang = pos / cfg.rope_scaling * freq;
        float c = std::cos(ang), s = std::sin(ang);
        float q1 = qp[d], q2 = qp[d + hk/2];
        qp[d] = q1 * c - q2 * s;
        qp[d + hk/2] = q1 * s + q2 * c;
        if (kp) {
            float k1 = kp[d], k2 = kp[d + hk/2];
            kp[d] = k1 * c - k2 * s;
            kp[d + hk/2] = k1 * s + k2 * c;
        }
    }
}

// === GLA (AVX2 + threaded across heads) ===
void YashaModel::gla(int T) {
    int Nh = cfg.num_heads, D = cfg.head_dim;
    int64_t state_sz = (int64_t)Nh * D * D;
    if ((int64_t)gla_state.d.size() < state_sz)
        gla_state.d.assign(state_sz, 0.0f);
    else
        std::fill(gla_state.d.begin(), gla_state.d.begin() + state_sz, 0.0f);

    int n_threads = std::min(Nh, (int)std::thread::hardware_concurrency());
    auto work = [&](int h0, int h1) {
        for (int h = h0; h < h1; h++) {
            float* sp = gla_state.data() + (int64_t)h * D * D;
            for (int t = 0; t < T; t++) {
                float* qp = q.data() + ((int64_t)t * Nh + h) * D;
                float* kp = k.data() + (int64_t)t * D;
                float* vp = v.data() + (int64_t)t * D;
                float* gp = g.data() + ((int64_t)t * Nh + h) * D;
                float* op = attn_out.data() + ((int64_t)t * Nh + h) * D;
                float gate_val = 1.0f / (1.0f + std::exp(-gp[0]));
                __m256 gv = _mm256_set1_ps(gate_val);

                // State = State * gate + outer(k, v)
                for (int i = 0; i < D; i++) {
                    __m256 ki = _mm256_set1_ps(kp[i]);
                    int j = 0;
                    for (; j + 8 <= D; j += 8) {
                        __m256 s = _mm256_loadu_ps(&sp[(int64_t)i * D + j]);
                        __m256 vj = _mm256_loadu_ps(&vp[j]);
                        s = _mm256_mul_ps(s, gv);
                        s = _mm256_fmadd_ps(ki, vj, s);
                        _mm256_storeu_ps(&sp[(int64_t)i * D + j], s);
                    }
                    for (; j < D; j++)
                        sp[(int64_t)i * D + j] = sp[(int64_t)i * D + j] * gate_val + kp[i] * vp[j];
                }

                // Output: op[i] = sum_j qp[j] * sp[j][i]  (cache-friendly: accumulate row-wise)
                for (int i = 0; i < D; i++) op[i] = 0.0f;
                for (int j = 0; j < D; j++) {
                    __m256 qj = _mm256_set1_ps(qp[j]);
                    int i = 0;
                    for (; i + 8 <= D; i += 8) {
                        __m256 s = _mm256_loadu_ps(&sp[(int64_t)j * D + i]);
                        __m256 o = _mm256_loadu_ps(&op[i]);
                        o = _mm256_fmadd_ps(qj, s, o);
                        _mm256_storeu_ps(&op[i], o);
                    }
                    for (; i < D; i++) op[i] += qp[j] * sp[(int64_t)j * D + i];
                }

                // Residual: op += qp
                int i = 0;
                for (; i + 8 <= D; i += 8) {
                    __m256 o = _mm256_loadu_ps(&op[i]);
                    __m256 q = _mm256_loadu_ps(&qp[i]);
                    o = _mm256_add_ps(o, q);
                    _mm256_storeu_ps(&op[i], o);
                }
                for (; i < D; i++) op[i] += qp[i];
            }
        }
    };

    if (n_threads <= 1) { work(0, Nh); return; }
    std::vector<std::thread> threads;
    int heads_per = Nh / n_threads;
    for (int th = 0; th < n_threads; th++) {
        int h0 = th * heads_per;
        int h1 = (th == n_threads - 1) ? Nh : (th + 1) * heads_per;
        threads.emplace_back(work, h0, h1);
    }
    for (auto& th : threads) th.join();
}

// === GLA with INT8 quantized state (4x less memory) ===
void YashaModel::gla_quantized(int T) {
    int Nh = cfg.num_heads, D = cfg.head_dim;
    int64_t state_sz = (int64_t)Nh * D * D;
    Tensor float_state({Nh, D, D});
    // If existing state in Q8, dequant first
    if (gla_state_q8.numel() > 0)
        gla_state_q8.dequantize(float_state.data(), (int)state_sz);
    else
        std::fill(float_state.data(), float_state.data() + state_sz, 0.0f);

    int n_threads = std::min(Nh, (int)std::thread::hardware_concurrency());
    auto work = [&](int h0, int h1) {
        for (int h = h0; h < h1; h++) {
            float* sp = float_state.data() + (int64_t)h * D * D;
            for (int t = 0; t < T; t++) {
                float* qp = q.data() + ((int64_t)t * Nh + h) * D;
                float* kp = k.data() + (int64_t)t * D;
                float* vp = v.data() + (int64_t)t * D;
                float* gp = g.data() + ((int64_t)t * Nh + h) * D;
                float* op = attn_out.data() + ((int64_t)t * Nh + h) * D;
                float gate_val = 1.0f / (1.0f + std::exp(-gp[0]));
                __m256 gv = _mm256_set1_ps(gate_val);

                for (int i = 0; i < D; i++) {
                    __m256 ki = _mm256_set1_ps(kp[i]);
                    int j = 0;
                    for (; j + 8 <= D; j += 8) {
                        __m256 s = _mm256_loadu_ps(&sp[(int64_t)i * D + j]);
                        __m256 vj = _mm256_loadu_ps(&vp[j]);
                        s = _mm256_mul_ps(s, gv);
                        s = _mm256_fmadd_ps(ki, vj, s);
                        _mm256_storeu_ps(&sp[(int64_t)i * D + j], s);
                    }
                    for (; j < D; j++)
                        sp[(int64_t)i * D + j] = sp[(int64_t)i * D + j] * gate_val + kp[i] * vp[j];
                }
                for (int i = 0; i < D; i++) op[i] = 0.0f;
                for (int j = 0; j < D; j++) {
                    __m256 qj = _mm256_set1_ps(qp[j]);
                    int i = 0;
                    for (; i + 8 <= D; i += 8) {
                        __m256 s = _mm256_loadu_ps(&sp[(int64_t)j * D + i]);
                        __m256 o = _mm256_loadu_ps(&op[i]);
                        o = _mm256_fmadd_ps(qj, s, o);
                        _mm256_storeu_ps(&op[i], o);
                    }
                    for (; i < D; i++) op[i] += qp[j] * sp[(int64_t)j * D + i];
                }
                int i = 0;
                for (; i + 8 <= D; i += 8) {
                    __m256 o = _mm256_loadu_ps(&op[i]);
                    __m256 q = _mm256_loadu_ps(&qp[i]);
                    o = _mm256_add_ps(o, q);
                    _mm256_storeu_ps(&op[i], o);
                }
                for (; i < D; i++) op[i] += qp[i];
            }
        }
    };
    if (n_threads <= 1) { work(0, Nh); }
    else {
        std::vector<std::thread> threads;
        int heads_per = Nh / n_threads;
        for (int th = 0; th < n_threads; th++) {
            int h0 = th * heads_per;
            int h1 = (th == n_threads - 1) ? Nh : (th + 1) * heads_per;
            threads.emplace_back(work, h0, h1);
        }
        for (auto& th : threads) th.join();
    }
    // Re-quantize back to INT8
    gla_state_q8.quantize(float_state.data(), (int)state_sz);
}

// === MoE single merged expert (for merged model, 2x speed) ===
void YashaModel::moe_merged(int li, int T) {
    int H = cfg.hidden_size, F = cfg.ffn_hidden;
    std::string p = "model.layers." + std::to_string(li) + ".";
    // For merged model: single expert at experts.0
    auto& gw = w[p + "experts.0.w1.weight"];
    auto& uw = w[p + "experts.0.w3.weight"];
    auto& dw = w[p + "experts.0.w2.weight"];
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        for (int j = 0; j < F; j++) {
            __m256 gs = _mm256_setzero_ps();
            __m256 us = _mm256_setzero_ps();
            int i = 0;
            for (; i + 8 <= H; i += 8) {
                __m256 xv = _mm256_loadu_ps(&xp[i]);
                gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)j * H + i]), gs);
                us = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&uw.data()[(int64_t)j * H + i]), us);
            }
            ffn_gate.data()[j] = hsum_ps(gs);
            ffn_up.data()[j] = hsum_ps(us);
            for (; i < H; i++) {
                ffn_gate.data()[j] += xp[i] * gw.data()[(int64_t)j * H + i];
                ffn_up.data()[j] += xp[i] * uw.data()[(int64_t)j * H + i];
            }
        }
        for (int j = 0; j < F; j++) {
            float gi = ffn_gate.data()[j];
            ffn_up.data()[j] = (gi / (1.0f + std::exp(-gi))) * ffn_up.data()[j];
        }
        for (int i = 0; i < H; i++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= F; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&ffn_up.data()[j]),
                    _mm256_loadu_ps(&dw.data()[(int64_t)i * F + j]), sum);
            float s = hsum_ps(sum);
            for (; j < F; j++) s += ffn_up.data()[j] * dw.data()[(int64_t)i * F + j];
            ffn_down.data()[(int64_t)t * H + i] = s;
        }
    }
    for (int t = 0; t < T * H; t++) ffn_down.d[t] += x.d[t];
}

// === Fused QKV projection (single matmul for Q, K, V, G) ===
void YashaModel::fused_qkv_proj(int T, int li) {
    int H = cfg.hidden_size, Nh = cfg.num_heads, D = cfg.head_dim;
    int QD = Nh * D, KD = D;
    std::string p = "model.layers." + std::to_string(li) + ".";
    // Fused weight: [QD + QD + KD + KD] × H  (Q, G, K, V stacked)
    // If not available, fall back to individual projections
    auto it_fused = w.find(p + "self_attn.fused_qkv.weight");
    if (it_fused != w.end()) {
        float* fused = it_fused->second.data();
        int total = QD + QD + KD + KD;
        for (int t = 0; t < T; t++) {
            float* xp = x2.data() + t * H;
            // One batched matmul over all projections
            for (int i = 0; i < total; i++) {
                __m256 sum = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= H; j += 8)
                    sum = _mm256_fmadd_ps(_mm256_loadu_ps(&xp[j]),
                        _mm256_loadu_ps(&fused[i * H + j]), sum);
                float s = hsum_ps(sum);
                for (; j < H; j++) s += xp[j] * fused[i * H + j];
                if (i < QD) q.data()[t * QD + i] = s;
                else if (i < 2 * QD) g.data()[t * QD + (i - QD)] = s;
                else if (i < 2 * QD + KD) k.data()[t * KD + (i - 2 * QD)] = s;
                else v.data()[t * KD + (i - 2 * QD - KD)] = s;
            }
        }
    } else {
        // Fallback: packed matmul for each projection
        auto& qw = w[p + "self_attn.q_proj.weight"];
        auto& kw = w[p + "self_attn.k_proj.weight"];
        auto& vw = w[p + "self_attn.v_proj.weight"];
        auto& gw = w[p + "self_attn.g_proj.weight"];
        for (int t = 0; t < T; t++) {
            float* xp = x2.data() + t * H;
            float* qp = q.data() + t * QD;
            float* gp = g.data() + t * QD;
            float* kp = k.data() + t * KD;
            float* vp = v.data() + t * KD;
            auto proj = [&](float* out, const Tensor& wt, int d) {
                for (int i = 0; i < d; i++) {
                    __m256 sum = _mm256_setzero_ps();
                    int j = 0;
                    for (; j + 8 <= H; j += 8)
                        sum = _mm256_fmadd_ps(_mm256_loadu_ps(&xp[j]),
                            _mm256_loadu_ps(&wt.data()[i * H + j]), sum);
                    out[i] = hsum_ps(sum);
                    for (; j < H; j++) out[i] += xp[j] * wt.data()[i * H + j];
                }
            };
            proj(qp, qw, QD); proj(gp, gw, QD);
            proj(kp, kw, KD); proj(vp, vw, KD);
        }
    }
}

// === Adaptive expert: dynamic top-k based on router entropy ===
void YashaModel::moe_adaptive(int li, int T) {
    int H = cfg.hidden_size, F = cfg.ffn_hidden, E = cfg.num_experts;
    std::string p = "model.layers." + std::to_string(li) + ".";
    auto& rw = w[p + "router.weight"];
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        float* rp = router.data() + (int64_t)t * E;
        for (int e = 0; e < E; e++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= H; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&xp[j]),
                    _mm256_loadu_ps(&rw.data()[(int64_t)e * H + j]), sum);
            rp[e] = hsum_ps(sum);
            for (; j < H; j++) rp[e] += xp[j] * rw.data()[(int64_t)e * H + j];
        }
        float max_r = *std::max_element(rp, rp + E);
        float sum_exp = 0;
        for (int e = 0; e < E; e++) { rp[e] = std::exp(rp[e] - max_r); sum_exp += rp[e]; }
        for (int e = 0; e < E; e++) rp[e] /= sum_exp;
        // Adaptive: if top-1 weight > threshold, use just 1 expert (2× speedup)
        // Otherwise use 2 experts. If entropy is very high, use 3.
        int K = cfg.num_experts_per_tok;
        if (cfg.adaptive_expert) {
            float top1 = *std::max_element(rp, rp + E);
            if (top1 > cfg.adaptive_expert_threshold) {
                K = 1;
            } else {
                // Compute entropy
                float entropy = 0;
                for (int e = 0; e < E; e++) if (rp[e] > 0) entropy -= rp[e] * std::log(rp[e]);
                float max_ent = std::log((float)E);
                if (entropy > max_ent * 0.8f) K = 3;  // very uncertain → 3 experts
            }
        }
        std::vector<std::pair<float,int>> idx;
        for (int e = 0; e < E; e++) idx.push_back({rp[e], e});
        std::partial_sort(idx.begin(), idx.begin()+K, idx.end(),
                          [](auto& a, auto& b){ return a.first > b.first; });
        for (int k = 0; k < K; k++) route_w.data()[(int64_t)t * K + k] = (float)idx[k].second;
        // Expert compute
        std::fill(ffn_down.data() + (int64_t)t * H, ffn_down.data() + (int64_t)(t+1) * H, 0.0f);
        for (int k = 0; k < K; k++) {
            int e = (int)route_w.data()[(int64_t)t * K + k];
            float wgt = rp[e];
            std::string e_pre = p + "experts." + std::to_string(e) + ".";
            auto& gw = w[e_pre + "w1.weight"];
            auto& uw = w[e_pre + "w3.weight"];
            auto& dw = w[e_pre + "w2.weight"];
            for (int j = 0; j < F; j++) {
                __m256 gs = _mm256_setzero_ps(); __m256 us = _mm256_setzero_ps();
                int i = 0;
                for (; i + 8 <= H; i += 8) {
                    __m256 xv = _mm256_loadu_ps(&xp[i]);
                    gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)j * H + i]), gs);
                    us = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&uw.data()[(int64_t)j * H + i]), us);
                }
                ffn_gate.data()[j] = hsum_ps(gs); ffn_up.data()[j] = hsum_ps(us);
                for (; i < H; i++) {
                    ffn_gate.data()[j] += xp[i] * gw.data()[(int64_t)j * H + i];
                    ffn_up.data()[j] += xp[i] * uw.data()[(int64_t)j * H + i];
                }
            }
            for (int j = 0; j < F; j++) {
                float gi = ffn_gate.data()[j];
                ffn_up.data()[j] = (gi / (1.0f + std::exp(-gi))) * ffn_up.data()[j];
            }
            for (int i = 0; i < H; i++) {
                __m256 sum = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= F; j += 8)
                    sum = _mm256_fmadd_ps(_mm256_loadu_ps(&ffn_up.data()[j]),
                        _mm256_loadu_ps(&dw.data()[(int64_t)i * F + j]), sum);
                float s = hsum_ps(sum);
                for (; j < F; j++) s += ffn_up.data()[j] * dw.data()[(int64_t)i * F + j];
                ffn_down.data()[(int64_t)t * H + i] += wgt * s;
            }
        }
    }
    for (int t = 0; t < T * H; t++) ffn_down.d[t] += x.d[t];
}

// === Pack all weights for SIMD-optimized matmul ===
void YashaModel::pack_all_weights() {
    pw.clear();
    for (auto& [key, tensor] : w) {
        if (tensor.ndim() == 2 && key.find("weight") != std::string::npos) {
            PackedWeight p;
            p.pack(tensor.data(), tensor.sh[0], tensor.sh[1]);
            pw[key] = std::move(p);
        }
    }
}

// === Load quantized weights (Q2/Q3) from safetensors ===
void YashaModel::load_quantized_weights(const std::string& dir) {
    std::cerr << "Loading quantized weights from " << dir << "...\n";
    for (auto& entry : fs::directory_iterator(dir)) {
        if (entry.path().extension() == ".safetensors") {
            std::unordered_map<std::string, Tensor> temp_w;
            load_safetensors(entry.path().string(), temp_w);
            for (auto& [key, t] : temp_w) {
                if (cfg.weight_bits == 2) {
                    TensorQ2 q2(t.sh); q2.quantize(t.data(), (int)t.numel());
                    w_q2[key] = std::move(q2);
                } else if (cfg.weight_bits == 3) {
                    TensorQ3 q3(t.sh); q3.quantize(t.data(), (int)t.numel());
                    w_q3[key] = std::move(q3);
                } else {
                    w[key] = std::move(t);
                }
            }
        }
    }
    std::cerr << "Loaded " << w.size() << " unquantized + " << w_q2.size() << " Q2 + "
              << w_q3.size() << " Q3 tensors\n";
}

// === Self-diffusion: re-run AR model on own hidden states ===
void YashaModel::diffuse_self(float* h, int T, int D) {
    // Level 1: Single-step refinement through one AR layer
    // Re-run the final layer using h as input to refine hidden states
    int li = cfg.num_layers - 1;
    std::string p = "model.layers." + std::to_string(li) + ".";
    int H = cfg.hidden_size, Nh = cfg.num_heads;
    float* xp = x2.data();

    // Copy h into x buffer for layer processing
    std::memcpy(xp, h, T * H * sizeof(float));
    rmsnorm(x2, x, w["model.norm.weight"]);

    // QKV projection for this layer
    auto& qw = w[p + "self_attn.q_proj.weight"];
    auto& kw = w[p + "self_attn.k_proj.weight"];
    auto& vw = w[p + "self_attn.v_proj.weight"];
    auto& gw = w[p + "self_attn.g_proj.weight"];
    auto& ow = w[p + "self_attn.o_proj.weight"];

    for (int t = 0; t < T; t++) {
        float* xt = x2.data() + t * H;
        float* qt = q.data() + t * Nh * D;
        float* gt = g.data() + t * Nh * D;
        float* kt = k.data() + t * D;
        float* vt = v.data() + t * D;
        for (int i = 0; i < Nh * D; i++) {
            __m256 qs = _mm256_setzero_ps(); __m256 gs = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= H; j += 8) {
                __m256 xv = _mm256_loadu_ps(&xt[j]);
                qs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&qw.data()[(int64_t)i * H + j]), qs);
                gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)i * H + j]), gs);
            }
            qt[i] = hsum_ps(qs); gt[i] = hsum_ps(gs);
        }
        for (int i = 0; i < D; i++) {
            __m256 ks = _mm256_setzero_ps(); __m256 vs = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= H; j += 8) {
                __m256 xv = _mm256_loadu_ps(&xt[j]);
                ks = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&kw.data()[(int64_t)i * H + j]), ks);
                vs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&vw.data()[(int64_t)i * H + j]), vs);
            }
            kt[i] = hsum_ps(ks); vt[i] = hsum_ps(vs);
        }
        rope_partial(qt, kt, t, D);
    }

    // Run GLA
    if (cfg.kv_int8) gla_quantized(T); else gla(T);

    // Output projection + residual
    for (int t = 0; t < T; t++) {
        float* attn_t = attn_out.data() + t * Nh * D;
        float* res_t = attn_res.data() + t * H;
        float* ht = h + t * H;
        for (int i = 0; i < H; i++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= Nh * D; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&attn_t[j]),
                    _mm256_loadu_ps(&ow.data()[(int64_t)i * Nh * D + j]), sum);
            res_t[i] = hsum_ps(sum);
            for (; j < Nh * D; j++) res_t[i] += attn_t[j] * ow.data()[(int64_t)i * Nh * D + j];
        }
        for (int i = 0; i < H; i++) ht[i] += res_t[i];  // residual to h
    }
}

// === Self-diffusion level 1: refine each token during AR ===
void YashaModel::apply_self_diffusion_level1(float* h, int T, int D) {
    diffuse_self(h, T, D);
}

// === Self-diffusion level 2: refine full sequence after generation ===
void YashaModel::apply_self_diffusion_level2(int T, int D) {
    // Re-run full AR on the generated tokens to get refined hidden states
    // h already contains the final hidden states from the initial forward
    float* h = x2.data() + (T - 1) * D;
    for (int step = 0; step < 2; step++) {
        // Add small noise
        for (int i = 0; i < T * D; i++) h[i] += randn() * 0.05f;
        diffuse_self(h, T, D);
    }
}

// === Self-diffusion level 3: regenerate with correction prompt ===
void YashaModel::apply_self_diffusion_level3(std::vector<int>& result, const std::vector<int>& prompt,
                                              int n_pred, float temp, float top_p) {
    std::vector<int> correction_prompt = prompt;
    std::string fix_str = "Check your work carefully. Fix any mistakes and improve your answer.";
    auto fix_ids = encode_text(fix_str);
    correction_prompt.insert(correction_prompt.end(), fix_ids.begin(), fix_ids.end());
    // Add current result as context
    for (int id : result) correction_prompt.push_back(id);
    // Regenerate
    Tensor r = forward(correction_prompt, n_pred, temp, top_p);
    result.clear();
    for (size_t i = correction_prompt.size(); i < (size_t)r.numel(); i++)
        result.push_back((int)r.d[i]);
}

// === Confidence ensemble (multi-head) ===
float YashaModel::score_confidence_ensemble(const float* h, int D) {
    // Load ensemble weights if available
    float conf = score_confidence(h, D);
    auto it = w.find("confidence_head.1.proj.0.weight");
    if (it != w.end()) {
        float c2 = 0;
        for (int j = 0; j < D; j++) c2 += it->second.data()[j] * h[j];
        conf = (conf + 1.0f / (1.0f + std::exp(-c2))) * 0.5f;
    }
    return conf;
}

// === MoE top-2 (AVX2) ===
void YashaModel::moe(int li, int T) {
    int H = cfg.hidden_size, F = cfg.ffn_hidden, E = cfg.num_experts, K = cfg.num_experts_per_tok;
    std::string p = "model.layers." + std::to_string(li) + ".";
    auto& rw = w[p + "router.weight"];
    // Router: AVX2 dot product per expert
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        float* rp = router.data() + (int64_t)t * E;
        for (int e = 0; e < E; e++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= H; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&xp[j]),
                    _mm256_loadu_ps(&rw.data()[(int64_t)e * H + j]), sum);
            rp[e] = hsum_ps(sum);
            for (; j < H; j++) rp[e] += xp[j] * rw.data()[(int64_t)e * H + j];
        }
        float max_r = *std::max_element(rp, rp + E);
        float sum = 0;
        for (int e = 0; e < E; e++) { rp[e] = std::exp(rp[e] - max_r); sum += rp[e]; }
        for (int e = 0; e < E; e++) rp[e] /= sum;
        std::vector<std::pair<float,int>> idx;
        for (int e = 0; e < E; e++) idx.push_back({rp[e], e});
        std::partial_sort(idx.begin(), idx.begin()+K, idx.end(),
                          [](auto& a, auto& b){ return a.first > b.first; });
        for (int k = 0; k < K; k++) route_w.data()[(int64_t)t * K + k] = (float)idx[k].second;
    }
    // Expert compute: AVX2 gate/up + down
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        std::fill(ffn_down.data() + (int64_t)t * H, ffn_down.data() + (int64_t)(t+1) * H, 0.0f);
        for (int k = 0; k < K; k++) {
            int e = (int)route_w.data()[(int64_t)t * K + k];
            float wgt = router.data()[(int64_t)t * E + e];
            std::string e_pre = p + "experts." + std::to_string(e) + ".";
            auto& gw = w[e_pre + "w1.weight"];
            auto& uw = w[e_pre + "w3.weight"];
            auto& dw = w[e_pre + "w2.weight"];
            // Gate + Up projection (AVX2)
            for (int j = 0; j < F; j++) {
                __m256 gs = _mm256_setzero_ps();
                __m256 us = _mm256_setzero_ps();
                int i = 0;
                for (; i + 8 <= H; i += 8) {
                    __m256 xv = _mm256_loadu_ps(&xp[i]);
                    gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)j * H + i]), gs);
                    us = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&uw.data()[(int64_t)j * H + i]), us);
                }
                ffn_gate.data()[j] = hsum_ps(gs); ffn_up.data()[j] = hsum_ps(us);
                for (; i < H; i++) {
                    ffn_gate.data()[j] += xp[i] * gw.data()[(int64_t)j * H + i];
                    ffn_up.data()[j] += xp[i] * uw.data()[(int64_t)j * H + i];
                }
            }
            // SiLU activation (minor, not worth AVX2)
            for (int j = 0; j < F; j++) {
                float gi = ffn_gate.data()[j];
                ffn_up.data()[j] = (gi / (1.0f + std::exp(-gi))) * ffn_up.data()[j];
            }
            // Down projection (AVX2)
            for (int i = 0; i < H; i++) {
                __m256 sum = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= F; j += 8)
                    sum = _mm256_fmadd_ps(_mm256_loadu_ps(&ffn_up.data()[j]),
                        _mm256_loadu_ps(&dw.data()[(int64_t)i * F + j]), sum);
                float s = hsum_ps(sum);
                for (; j < F; j++) s += ffn_up.data()[j] * dw.data()[(int64_t)i * F + j];
                ffn_down.data()[(int64_t)t * H + i] += wgt * s;
            }
        }
    }
    for (int t = 0; t < T * H; t++) ffn_down.d[t] += x.d[t];
}

// === Expert Choice routing (global capacity-based) ===
void YashaModel::compute_router_scores(int T) {
    int H = cfg.hidden_size, E = cfg.num_experts;
    // Compute router prob for all tokens × all experts and store
    for (int li = 0; li < cfg.num_layers; li++) {
        std::string p = "model.layers." + std::to_string(li) + ".";
        auto* rwp = w.count(p + "router.weight") ? w[p + "router.weight"].data() : nullptr;
        if (!rwp) continue;
        (void)T; // placeholder — scores computed per-layer inside moe_expert_choice
    }
}

void YashaModel::assign_experts_global(int T) {
    // Global assignment: capacity = ceil(T / E) per expert
    int E = cfg.num_experts;
    int cap = (T + E - 1) / E;
    std::vector<int> expert_load(E, 0);
    // Greedy: assign each token to highest-scoring expert with remaining capacity
    for (int t = 0; t < T; t++) {
        int best_e = 0;
        float best_s = -1e9;
        for (int e = 0; e < E; e++) {
            if (expert_load[e] < cap && ec_scores.data()[t * E + e] > best_s) {
                best_s = ec_scores.data()[t * E + e];
                best_e = e;
            }
        }
        ec_assign.d[t] = (float)best_e;
        expert_load[best_e]++;
    }
}

void YashaModel::moe_expert_choice(int li, int T) {
    int H = cfg.hidden_size, F = cfg.ffn_hidden, E = cfg.num_experts;
    std::string p = "model.layers." + std::to_string(li) + ".";
    auto& rw = w[p + "router.weight"];

    // Compute all router scores (AVX2)
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        float* rp = ec_scores.data() + (int64_t)t * E;
        for (int e = 0; e < E; e++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= H; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&xp[j]),
                    _mm256_loadu_ps(&rw.data()[(int64_t)e * H + j]), sum);
            rp[e] = hsum_ps(sum);
            for (; j < H; j++) rp[e] += xp[j] * rw.data()[(int64_t)e * H + j];
        }
    }

    assign_experts_global(T);

    // Route tokens to their assigned experts (AVX2)
    for (int t = 0; t < T; t++) {
        float* xp = x.data() + (int64_t)t * H;
        int e = (int)ec_assign.d[t];
        std::fill(ffn_down.data() + (int64_t)t * H, ffn_down.data() + (int64_t)(t+1) * H, 0.0f);
        std::string e_pre = p + "experts." + std::to_string(e) + ".";
        auto& gw = w[e_pre + "w1.weight"];
        auto& uw = w[e_pre + "w3.weight"];
        auto& dw = w[e_pre + "w2.weight"];
        for (int j = 0; j < F; j++) {
            __m256 gs = _mm256_setzero_ps();
            __m256 us = _mm256_setzero_ps();
            int i = 0;
            for (; i + 8 <= H; i += 8) {
                __m256 xv = _mm256_loadu_ps(&xp[i]);
                gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)j * H + i]), gs);
                us = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&uw.data()[(int64_t)j * H + i]), us);
            }
            ffn_gate.data()[j] = hsum_ps(gs); ffn_up.data()[j] = hsum_ps(us);
            for (; i < H; i++) {
                ffn_gate.data()[j] += xp[i] * gw.data()[(int64_t)j * H + i];
                ffn_up.data()[j] += xp[i] * uw.data()[(int64_t)j * H + i];
            }
        }
        for (int j = 0; j < F; j++) {
            float gi = ffn_gate.data()[j];
            ffn_up.data()[j] = (gi / (1.0f + std::exp(-gi))) * ffn_up.data()[j];
        }
        for (int i = 0; i < H; i++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= F; j += 8)
                sum = _mm256_fmadd_ps(_mm256_loadu_ps(&ffn_up.data()[j]),
                    _mm256_loadu_ps(&dw.data()[(int64_t)i * F + j]), sum);
            float s = hsum_ps(sum);
            for (; j < F; j++) s += ffn_up.data()[j] * dw.data()[(int64_t)i * F + j];
            ffn_down.data()[(int64_t)t * H + i] = s;
        }
    }
    for (int t = 0; t < T * H; t++) ffn_down.d[t] += x.d[t];
}

// === Single layer with variable depth support ===
void YashaModel::layer(int li, int T) {
    int H = cfg.hidden_size, D = cfg.head_dim, Nh = cfg.num_heads;
    std::string p = "model.layers." + std::to_string(li) + ".";

    rmsnorm(x2, x, w["model.norm.weight"]);
    auto it_in = w.find(p + "input_layernorm.weight");
    if (it_in != w.end()) rmsnorm(x2, x, it_in->second);

    auto& ow = w[p + "self_attn.o_proj.weight"];

    if (cfg.fused_qkv) {
        fused_qkv_proj(T, li);
    } else {
        auto& qw = w[p + "self_attn.q_proj.weight"];
        auto& kw = w[p + "self_attn.k_proj.weight"];
        auto& vw = w[p + "self_attn.v_proj.weight"];
        auto& gw = w[p + "self_attn.g_proj.weight"];
        int QD = Nh * D, KD = D;
        for (int t = 0; t < T; t++) {
            float* xp = x2.data() + (int64_t)t * H;
            float* qp = q.data() + (int64_t)t * QD;
            float* gp = g.data() + (int64_t)t * QD;
            float* kp = k.data() + (int64_t)t * KD;
            float* vp = v.data() + (int64_t)t * KD;
            for (int i = 0; i < QD; i++) {
                __m256 qs = _mm256_setzero_ps(); __m256 gs = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= H; j += 8) {
                    __m256 xv = _mm256_loadu_ps(&xp[j]);
                    qs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&qw.data()[(int64_t)i * H + j]), qs);
                    gs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&gw.data()[(int64_t)i * H + j]), gs);
                }
                qp[i] = hsum_ps(qs); gp[i] = hsum_ps(gs);
                for (; j < H; j++) {
                    qp[i] += xp[j] * qw.data()[(int64_t)i * H + j];
                    gp[i] += xp[j] * gw.data()[(int64_t)i * H + j];
                }
            }
            for (int i = 0; i < KD; i++) {
                __m256 ks = _mm256_setzero_ps(); __m256 vs = _mm256_setzero_ps();
                int j = 0;
                for (; j + 8 <= H; j += 8) {
                    __m256 xv = _mm256_loadu_ps(&xp[j]);
                    ks = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&kw.data()[(int64_t)i * H + j]), ks);
                    vs = _mm256_fmadd_ps(xv, _mm256_loadu_ps(&vw.data()[(int64_t)i * H + j]), vs);
                }
                kp[i] = hsum_ps(ks); vp[i] = hsum_ps(vs);
                for (; j < H; j++) {
                    kp[i] += xp[j] * kw.data()[(int64_t)i * H + j];
                    vp[i] += xp[j] * vw.data()[(int64_t)i * H + j];
                }
            }
            rope_partial(qp, kp, t, D);
        }
    }

    if (cfg.kv_int8) gla_quantized(T); else gla(T);

    // Output projection (AVX2)
    for (int t = 0; t < T; t++) {
        float* attn_t = attn_out.data() + (int64_t)t * Nh * D;
        float* res_t = attn_res.data() + (int64_t)t * H;
        for (int i = 0; i < H; i++) {
            __m256 sum = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= Nh * D; j += 8) {
                sum = _mm256_fmadd_ps(
                    _mm256_loadu_ps(&attn_t[j]),
                    _mm256_loadu_ps(&ow.data()[(int64_t)i * Nh * D + j]),
                    sum);
            }
            res_t[i] = hsum_ps(sum);
            for (; j < Nh * D; j++)
                res_t[i] += attn_t[j] * ow.data()[(int64_t)i * Nh * D + j];
        }
    }
    for (int t = 0; t < T * H; t++) x.d[t] += attn_res.d[t];

    auto it_post = w.find(p + "post_attention_layernorm.weight");
    if (it_post != w.end()) rmsnorm(x2, x, it_post->second);

    if (cfg.merged_experts) {
        moe_merged(li, T);
    } else if (cfg.adaptive_expert) {
        moe_adaptive(li, T);
    } else if (cfg.expert_choice && li % 2 == 1) {
        moe_expert_choice(li, T);
    } else {
        moe(li, T);
    }

    for (int t = 0; t < T * H; t++) x.d[t] = ffn_down.d[t];

    // Variable depth: score confidence after each layer
    if (cfg.var_depth_threshold < 1.0f) {
        for (int t = 0; t < T; t++) {
            float* hp = x.data() + t * H;
            layer_conf.d[t] = score_confidence(hp, H);
        }
    }
}

// === Confidence (uses learned head, no separate diffuser MLP needed) ===
float YashaModel::score_confidence(const float* h, int D) {
    auto it = w.find("confidence_head.weight");
    if (it == w.end()) return 0.5f;
    const float* cw = it->second.data();
    float s = w.count("confidence_head.bias") ? w["confidence_head.bias"].data()[0] : 0;
    for (int j = 0; j < D; j++) s += cw[j] * h[j];
    return 1.0f / (1.0f + std::exp(-s));
}

// === MTP — predict K tokens from one hidden state ===
void YashaModel::predict_mtp(const float* h, int D, int* out_ids, int K) {
    int V = cfg.vocab_size;
    auto& lm = w["lm_head.weight"];

    // Token 0 = normal LM head
    for (int i = 0; i < V; i++) {
        float s = 0;
        for (int j = 0; j < D; j++) s += h[j] * lm.data()[i * D + j];
        mtp_logits[0].d[i] = s;
    }
    out_ids[0] = sample(mtp_logits[0].data(), V, 0.7f, 0.9f);

    // Tokens 1..K-1 = MTP heads (if available) or reuse main head
    float* embed = w.count("model.embed_tokens.weight") ? w["model.embed_tokens.weight"].data() : nullptr;
    for (int k = 1; k < K; k++) {
        std::string hname = "model.mtp_head." + std::to_string(k-1) + ".weight";
        auto it = w.find(hname);
        if (it != w.end()) {
            float* mtpw = it->second.data();
            float* mtpb = w.count("model.mtp_head." + std::to_string(k-1) + ".bias")
                          ? w["model.mtp_head." + std::to_string(k-1) + ".bias"].data() : nullptr;
            for (int i = 0; i < V; i++) {
                float s = mtpb ? mtpb[i] : 0;
                for (int j = 0; j < D; j++) s += mtpw[i * D + j] * h[j];
                mtp_logits[k].d[i] = s;
            }
        } else {
            // Fallback: use previous token's embedding to refine
            if (embed && out_ids[k-1] >= 0 && out_ids[k-1] < (int)w["model.embed_tokens.weight"].sh[0]) {
                float* prev_emb = embed + out_ids[k-1] * D;
                for (int i = 0; i < V; i++) {
                    float s = 0;
                    for (int j = 0; j < D; j++) s += prev_emb[j] * lm.data()[i * D + j];
                    mtp_logits[k].d[i] = s;
                }
            } else {
                std::memcpy(mtp_logits[k].d.data(), mtp_logits[k-1].d.data(), V * sizeof(float));
            }
        }
        out_ids[k] = sample(mtp_logits[k].data(), V, 0.7f, 0.9f);
    }
}

// === Forward with MTP (chunk prediction) ===
float YashaModel::forward_mtp_chunk(const std::vector<int>& tokens, int start, int* out_chunk, int K) {
    int H = cfg.hidden_size;
    // Run AR for the prefix to get the last hidden state
    Tensor r = forward(tokens, 0, 0.7f, 0.9f); // n_pred=0 means just get logits
    float* h = x2.data() + ((int)tokens.size() - 1) * H;
    predict_mtp(h, H, out_chunk, K);
    // Score the chunk's first token for confidence
    return score_confidence(h, H);
}

// === Variable depth forward ===
void YashaModel::forward_vardepth(std::vector<int>& result, int n_pred, float temp, float top_p) {
    int H = cfg.hidden_size, V = cfg.vocab_size;
    auto& lm = w["lm_head.weight"];
    int T = (int)result.size();

    // Embedding
    auto& emb = w["model.embed_tokens.weight"];
    for (int t = 0; t < T; t++) {
        int id = result[t];
        if (id >= 0 && id < emb.sh[0])
            std::memcpy(x.data() + t * H, emb.data() + id * H, H * sizeof(float));
    }

    // Layers with early exit
    for (int li = 0; li < cfg.num_layers; li++) {
        layer(li, T);
        // Check if all tokens have high confidence → exit
        bool all_confident = true;
        for (int t = 0; t < T; t++) {
            if (layer_conf.d[t] < cfg.var_depth_threshold) {
                all_confident = false;
                break;
            }
        }
        if (all_confident && li >= cfg.num_layers / 2) {
            std::cerr << "  early exit at layer " << li << "/" << cfg.num_layers << "\n";
            break;
        }
    }

    // Final norm + LM head with self-diffusion
    rmsnorm(x2, x, w["model.norm.weight"]);
    float* h = x2.data() + (T - 1) * H;
    if (cfg.self_diffusion_level >= SD_LEVEL1_TOKEN)
        apply_self_diffusion_level1(h, 1, H);
    for (int i = 0; i < V; i++) {
        float s = 0;
        for (int j = 0; j < H; j++) s += h[j] * lm.data()[i * H + j];
        logits.data()[i] = s;
    }
    int id = sample(logits.data(), V, temp, top_p);
    result.push_back(id);
}

// === Forward (standard AR with diffusion refine) ===
// === Diffusion-only generation (uses self-diffusion = re-run AR) ===
Tensor YashaModel::generate_diffusion(int n_pred, float temp, float top_p) {
    int H = cfg.hidden_size, V = cfg.vocab_size;
    auto& lm = w["lm_head.weight"];
    auto& emb = w["model.embed_tokens.weight"];
    Tensor result({n_pred});
    // Use diff_buffer for state
    if (diff_buffer.numel() < (int64_t)H) diff_buffer = Tensor({H});
    float* state = diff_buffer.data();

    for (int p = 0; p < n_pred; p++) {
        // Initialize state with noise
        for (int i = 0; i < H; i++)
            state[i] = randn() * 0.5f;
        // Apply self-diffusion refinement (re-run final AR layer)
        diffuse_self(state, 1, H);
        for (int i = 0; i < V; i++) {
            float s = 0;
            for (int j = 0; j < H; j++) s += state[j] * lm.data()[i * H + j];
            logits.data()[i] = s;
        }
        int id = sample(logits.data(), V, temp, top_p);
        result.d[p] = (float)id;
        if (id == 0 || id == 2) break;
        if (id >= 0 && id < emb.sh[0]) {
            float* next_emb = emb.data() + id * H;
            std::memcpy(state, next_emb, H * sizeof(float));
        }
    }
    return result;
}

// === Hard task heuristic ===
bool YashaModel::is_hard(const std::vector<int>& tokens) {
    if ((int)tokens.size() >= cfg.hard_threshold) return true;
    std::unordered_set<int> uniq(tokens.begin(), tokens.end());
    return (float)uniq.size() / std::max(1, (int)tokens.size()) > 0.6f;
}

// === RL rejection sampling (test-time compute) ===
void YashaModel::forward_rl_once(std::vector<int>& result, int n_pred, float temp, float top_p,
                                 float& best_score, std::mutex& mtx) {
    int H = cfg.hidden_size;
    std::vector<int> attempt(n_pred + 1, 0);

    // Randomly choose strategy for this attempt
    float roll = std::uniform_real_distribution<float>(0, 1)(::rng());
    Tensor r;
    if (roll < cfg.diffusion_prob) {
        r = generate_diffusion(n_pred, temp, top_p);
        for (int i = 0; i < (int)r.numel(); i++) attempt[i] = (int)r.d[i];
    } else {
        // Short AR prefix + MTP chunk
        r = forward({0}, n_pred, temp, top_p);
        for (size_t i = 1; i < (size_t)r.numel(); i++) attempt[i-1] = (int)r.d[i];
    }

    // Score the result
    float score = 0;
    if ((int)r.numel() > 0) {
        float* h = x2.data() + (std::min((int)r.numel(), n_pred) - 1) * H;
        score = score_confidence(h, H);
    }

    std::lock_guard<std::mutex> lk(mtx);
    if (score > best_score) {
        best_score = score;
        result = attempt;
    }
}

// === Parallel beam expansion ===
void YashaModel::expand_beam(const Beam& b, int depth, int max_d, int n_pred, float temp, float top_p,
                            std::vector<Beam>& results, std::mutex& mtx) {
    if (depth >= max_d) {
        Tensor r = forward(b.ids, n_pred, temp, top_p);
        Beam nb; nb.ids = b.ids; nb.score = b.score;
        for (size_t i = b.ids.size(); i < (size_t)r.numel(); i++) nb.ids.push_back((int)r.d[i]);
        float* h = x2.data() + ((int)nb.ids.size() - 1) * cfg.hidden_size;
        nb.score = score_confidence(h, cfg.hidden_size);
        std::lock_guard<std::mutex> lk(mtx); results.push_back(nb);
        return;
    }

    int nf = std::min(cfg.n_beams, std::max(1, n_pred));
    if (cfg.rl_samples > nf) nf = cfg.rl_samples;
    std::vector<Beam> forks(nf);
    for (int i = 0; i < nf; i++) { forks[i].ids = b.ids; forks[i].score = b.score; }

    std::vector<std::thread> thr;
    std::mutex fm;
    std::vector<Beam> ex;
    for (int i = 0; i < nf; i++)
        thr.emplace_back([this, &forks, i, depth, max_d, n_pred, temp, top_p, &ex, &fm]() {
            // Per-fork: randomly pick diffusion or AR based on diffusion_prob
            bool use_diff = std::bernoulli_distribution(cfg.diffusion_prob)(::rng());
            Tensor r;
            if (use_diff && depth >= max_d - 1) {
                r = generate_diffusion(std::max(1, n_pred / 2), temp, top_p);
                for (int k = 0; k < (int)r.numel(); k++)
                    forks[i].ids.push_back((int)r.d[k]);
            } else if (cfg.mtp_heads > 1 && !use_diff) {
                // Use MTP chunk prediction for this fork
                std::vector<int> prefix = forks[i].ids;
                int chunk[4];
                float cscore = forward_mtp_chunk(prefix, (int)prefix.size(), chunk, cfg.mtp_heads);
                for (int k = 0; k < cfg.mtp_heads; k++)
                    forks[i].ids.push_back(chunk[k]);
                forks[i].score = cscore;
            } else {
                r = forward(forks[i].ids, std::max(1, n_pred / 2), temp, top_p);
                for (size_t j = forks[i].ids.size(); j < (size_t)r.numel(); j++)
                    forks[i].ids.push_back((int)r.d[j]);
            }
            if (use_diff || !(cfg.mtp_heads > 1 && !use_diff)) {
                // Score normally if not already scored by MTP
                float* h = x2.data() + ((int)forks[i].ids.size() - 1) * cfg.hidden_size;
                forks[i].score = score_confidence(h, cfg.hidden_size);
            }
            expand_beam(forks[i], depth + 1, max_d, n_pred, temp, top_p, ex, fm);
        });
    for (auto& t : thr) t.join();

    std::sort(ex.begin(), ex.end(), [](const Beam& a, const Beam& b) { return a.score > b.score; });
    for (int i = 0; i < std::min(1, (int)ex.size()); i++) {
        std::lock_guard<std::mutex> lk(mtx); results.push_back(ex[i]);
    }
}

// === Parallel generation (entry point) ===
Tensor YashaModel::generate_parallel(const std::vector<int>& tokens, int n_pred, float temp, float top_p) {
    // Phase 1: RL rejection sampling over N candidates
    if (cfg.rl_samples > 1) {
        std::vector<int> best_result;
        float best_score = -1e9;
        std::mutex mtx;
        std::vector<std::thread> thr;
        for (int i = 0; i < cfg.rl_samples; i++)
            thr.emplace_back([this, &best_result, n_pred, temp, top_p, &best_score, &mtx]() {
                std::vector<int> attempt;
                forward_rl_once(attempt, n_pred, temp, top_p, best_score, mtx);
            });
        for (auto& t : thr) t.join();

        if (!best_result.empty()) {
            Tensor out({(int)best_result.size()});
            for (size_t i = 0; i < best_result.size(); i++) out.d[i] = (float)best_result[i];
            return out;
        }
    }

    // Phase 2: Tree search over beams
    int H = cfg.hidden_size, V = cfg.vocab_size;
    auto& lm = w["lm_head.weight"];
    Beam seed; seed.ids = tokens; seed.score = 0.5f;
    std::vector<Beam> results;
    std::mutex mtx;
    expand_beam(seed, 0, cfg.max_depth, n_pred, temp, top_p, results, mtx);

    if (results.empty()) {
        Tensor r = forward(tokens, n_pred, temp, top_p);
        Tensor out({(int)r.numel() - (int)tokens.size()});
        for (size_t i = tokens.size(); i < (size_t)r.numel(); i++) out.d[i - tokens.size()] = r.d[i];
        return out;
    }

    auto best = std::max_element(results.begin(), results.end(),
        [](const Beam& a, const Beam& b) { return a.score > b.score; });
    float* h = x2.data() + ((int)best->ids.size() - 1) * H;
    for (int i = 0; i < V; i++) {
        float s = 0;
        for (int j = 0; j < H; j++) s += h[j] * lm.data()[i * H + j];
        logits.data()[i] = s;
    }
    Tensor out({(int)best->ids.size() - (int)tokens.size()});
    for (size_t i = tokens.size(); i < best->ids.size(); i++) out.d[i - tokens.size()] = (float)best->ids[i];
    return out;
}

// === Sampling ===
int sample(const float* logits, int n, float temp, float top_p) {
    if (n <= 0) return 0;
    std::vector<std::pair<float,int>> p;
    float max_l = *std::max_element(logits, logits + n);
    float sum = 0;
    for (int i = 0; i < n; i++) {
        float v = std::exp((logits[i] - max_l) / std::max(temp, 0.001f));
        p.push_back({v, i}); sum += v;
    }
    std::sort(p.begin(), p.end(), [](auto& a, auto& b) { return a.first > b.first; });
    float cum = 0; int cutoff = n;
    for (int i = 0; i < n; i++) { cum += p[i].first / sum; if (cum > top_p) { cutoff = i + 1; break; } }
    std::uniform_real_distribution<float> dist(0, cum);
    float r = dist(::rng()); cum = 0;
    for (int i = 0; i < cutoff; i++) { cum += p[i].first / sum; if (r < cum) return p[i].second; }
    return p[0].second;
}

void softmax(float* p, int n) {
    float max_v = *std::max_element(p, p + n);
    float sum = 0;
    for (int i = 0; i < n; i++) { p[i] = std::exp(p[i] - max_v); sum += p[i]; }
    for (int i = 0; i < n; i++) p[i] /= sum;
}

void rmsnorm(Tensor& o, const Tensor& x, const Tensor& w, float eps) {
    int D = x.sh.back(); int N = (int)x.numel() / D;
    for (int i = 0; i < N; i++) {
        const float* xp = x.data() + i * D; float* op = o.data() + i * D;
        float ss = 0; for (int j = 0; j < D; j++) ss += xp[j] * xp[j];
        float s = 1.0f / std::sqrt(ss / D + eps);
        for (int j = 0; j < D; j++) op[j] = xp[j] * s * (j < (int)w.numel() ? w.data()[j] : 1.0f);
    }
}

void gelu(Tensor& o, const Tensor& x) {
    int N = (int)x.numel();
    for (int i = 0; i < N; i++) o.d[i] = 0.5f * x.d[i] * (1.0f + std::erf(x.d[i] / 1.41421356f));
}

// === safetensors loader ===
bool load_safetensors(const std::string& path, std::unordered_map<std::string, Tensor>& w) {
    std::ifstream f(path, std::ios::binary);
    if (!f) return false;
    uint64_t hlen; f.read((char*)&hlen, 8);
    std::string hdr((size_t)hlen, 0); f.read(hdr.data(), hlen);
    size_t pos = 0;
    auto skip_ws = [&]() { while (pos < hdr.size() && (hdr[pos]==' '||hdr[pos]=='\n'||hdr[pos]=='\t'||hdr[pos]=='\r')) pos++; };
    auto expect = [&](char c) { skip_ws(); if (hdr[pos] != c) return false; pos++; return true; };
    if (!expect('{')) return false;
    while (pos < hdr.size()) {
        skip_ws();
        if (hdr[pos] == '}') break;
        if (hdr[pos] == ',') { pos++; continue; }
        if (hdr[pos] != '"') break; pos++;
        size_t endk = hdr.find('"', pos);
        std::string key = hdr.substr(pos, endk - pos); pos = endk + 1;
        if (!expect(':')) break; if (!expect('{')) break;
        auto find_field = [&](const std::string& name) -> std::string {
            size_t p = hdr.find(name, pos); if (p == std::string::npos) return "";
            p = hdr.find('"', p + name.size() + 2); if (p == std::string::npos) return "";
            size_t e = hdr.find('"', p+1); return hdr.substr(p+1, e-p-1);
        };
        auto find_offsets = [&]() -> std::pair<uint64_t, uint64_t> {
            size_t p = hdr.find("data_offsets", pos); if (p == std::string::npos) return {0,0};
            p = hdr.find('[', p); if (p == std::string::npos) return {0,0}; p++;
            char* end; uint64_t s = strtoull(hdr.c_str() + p, &end, 10);
            p = end - hdr.c_str() + 1; uint64_t e = strtoull(hdr.c_str() + p, &end, 10);
            return {s, e};
        };
        auto find_shape = [&]() -> std::vector<int> {
            std::vector<int> s; size_t p = hdr.find("shape", pos); if (p == std::string::npos) return s;
            p = hdr.find('[', p); if (p == std::string::npos) return s; p++;
            while (p < hdr.size() && hdr[p] != ']') {
                if (hdr[p] >= '0' && hdr[p] <= '9') {
                    char* end; int64_t v = strtoll(hdr.c_str() + p, &end, 10);
                    s.push_back((int)v); p = end - hdr.c_str();
                } else p++;
            } return s;
        };
        auto shape = find_shape(); auto [dstart, dend] = find_offsets(); (void)dend;
        uint64_t dsize = 1; for (int s : shape) dsize *= s; dsize *= 4;
        Tensor t(shape); f.seekg(8 + hlen + dstart); f.read((char*)t.d.data(), dsize);
        w[key] = std::move(t);
        int brace = 1;
        while (brace > 0 && pos < hdr.size()) { if (hdr[pos] == '{') brace++; else if (hdr[pos] == '}') brace--; pos++; }
    }
    return true;
}

// === Self-consistency: generate one candidate ===
void YashaModel::generate_one_answer(const std::vector<int>& prompt, int n_pred, float temp, float top_p,
                                     std::vector<int>& out, std::mutex& mtx) {
    float roll = std::uniform_real_distribution<float>(0, 1)(::rng());
    Tensor r;
    if (roll < cfg.diffusion_prob) {
        r = generate_diffusion(n_pred, temp, top_p);
    } else {
        r = forward(prompt, n_pred, temp, top_p);
    }
    std::vector<int> ans;
    for (size_t i = r.numel() > (int)prompt.size() ? prompt.size() : 0; i < (size_t)r.numel(); i++)
        ans.push_back((int)r.d[i]);
    std::lock_guard<std::mutex> lk(mtx);
    out = ans;
}

std::vector<int> YashaModel::longest_common_prefix(const std::vector<std::vector<int>>& answers) {
    if (answers.empty()) return {};
    // Count votes for each prefix position
    int max_len = 0;
    for (auto& a : answers) if ((int)a.size() > max_len) max_len = (int)a.size();
    std::vector<int> result;
    for (int pos = 0; pos < max_len; pos++) {
        std::unordered_map<int, int> votes;
        for (auto& a : answers) {
            if (pos < (int)a.size()) votes[a[pos]]++;
        }
        int best_tok = -1, best_votes = 0;
        for (auto& [tok, v] : votes) {
            if (v > best_votes) { best_votes = v; best_tok = tok; }
        }
        if (best_votes < (int)answers.size() / 2 + 1) break; // no majority
        result.push_back(best_tok);
    }
    return result;
}

// === Self-consistency generation (majority voting) ===
Tensor YashaModel::forward(const std::vector<int>& tokens, int n_pred, float temp, float top_p) {
    // If self-consistency is active and this is a hard task, do majority voting
    if (cfg.sc_samples > 1 && is_hard(tokens) && n_pred > 0) {
        std::vector<std::vector<int>> answers(cfg.sc_samples);
        std::mutex mtx;
        std::vector<std::thread> thr;
        for (int i = 0; i < cfg.sc_samples; i++)
            thr.emplace_back([this, &tokens, n_pred, temp, top_p, &answers, i, &mtx]() {
                generate_one_answer(tokens, n_pred, temp, top_p, answers[i], mtx);
            });
        for (auto& t : thr) t.join();

        auto consensus = longest_common_prefix(answers);
        if (consensus.empty()) consensus = answers[0];
        Tensor out({(int)consensus.size()});
        for (size_t i = 0; i < consensus.size(); i++) out.d[i] = (float)consensus[i];
        return out;
    }

    // Normal forward (existing code follows)
    int T = (int)tokens.size();
    int H = cfg.hidden_size, V = cfg.vocab_size;
    auto& emb = w["model.embed_tokens.weight"];
    for (int t = 0; t < T; t++) {
        int id = tokens[t];
        if (id >= 0 && id < emb.sh[0])
            std::memcpy(x.data() + t * H, emb.data() + id * H, H * sizeof(float));
    }
    for (int li = 0; li < cfg.num_layers; li++) {
        layer(li, T);
        if (li % 10 == 0) std::cerr << "\r  layer " << li << "/" << cfg.num_layers;
    }
    std::cerr << "\r  layers done      \n";
    rmsnorm(x2, x, w["model.norm.weight"]);
    // Self-diffusion Level 1: refine final hidden state (always on)
    if (cfg.self_diffusion_level >= SD_LEVEL1_TOKEN) {
        float* h_final = x2.data() + (T - 1) * H;
        apply_self_diffusion_level1(h_final, 1, H);
    }

    int last_T = T;
    auto& lm = w["lm_head.weight"];
    if (n_pred <= 0) {
        // Compute logits for final token
        float* h = x2.data() + (T - 1) * H;
        for (int i = 0; i < V; i++) {
            float s = 0;
            for (int j = 0; j < H; j++) s += h[j] * lm.data()[i * H + j];
            logits.data()[i] = s;
        }
        return logits;
    }

    std::vector<int> result = tokens;
    for (int p = 0; p < n_pred; p++) {
        float* h = x2.data() + ((int)result.size() - 1) * H;
        for (int i = 0; i < V; i++) {
            float s = 0;
            for (int j = 0; j < H; j++) s += h[j] * lm.data()[i * H + j];
            logits.data()[i] = s;
        }
        int id = sample(logits.data(), V, temp, top_p);
        result.push_back(id);
        if (id == 0 || id == 2) break;
        if ((int)result.size() > cfg.max_seq) result.erase(result.begin());
        if (p < n_pred - 1) {
            auto next = forward(result, 0, temp, top_p);
            logits = next;
        }
    }

    // Self-diffusion Level 2: refine full sequence if mean confidence is low
    if (cfg.self_diffusion_level >= SD_LEVEL2_SEQUENCE && n_pred > 0) {
        int ar_T = (int)result.size();
        float mean_conf = 0;
        for (int t = 0; t < ar_T; t++) {
            float* ht = x2.data() + t * H;
            mean_conf += score_confidence_ensemble(ht, H);
        }
        mean_conf /= ar_T;
        if (mean_conf < cfg.self_diffusion_threshold) {
            std::cerr << "\r  self-diffusion L2 (conf=" << mean_conf << ")\n";
            float* h_all = x2.data();
            apply_self_diffusion_level2(ar_T, H);
        }
    }

    // Self-diffusion Level 3: full regeneration if still low confidence
    if (cfg.self_diffusion_level >= SD_LEVEL3_SELFCORRECT && n_pred > 0) {
        int ar_T = (int)result.size();
        float mean_conf = 0;
        for (int t = 0; t < ar_T; t++) {
            float* ht = x2.data() + t * H;
            mean_conf += score_confidence_ensemble(ht, H);
        }
        mean_conf /= ar_T;
        if (mean_conf < cfg.self_diffusion_correction_threshold) {
            std::cerr << "\r  self-diffusion L3 (conf=" << mean_conf << ")\n";
            apply_self_diffusion_level3(result, tokens, n_pred, temp, top_p);
        }
    }
    Tensor r;
    r.d.resize(result.size());
    for (size_t i = 0; i < result.size(); i++) r.d[i] = (float)result[i];
    return r;
}

// === Prompt cache: reuse GLA state across multi-turn ===
void YashaModel::clear_cache() {
    has_cache = false;
    cached_prefix.clear();
    cached_gla_state = Tensor();
    cached_x = Tensor();
}

Tensor YashaModel::forward_cached(const std::vector<int>& tokens, int n_pred, float temp, float top_p) {
    int H = cfg.hidden_size, L = cfg.num_layers, Nh = cfg.num_heads, D = cfg.head_dim;
    // Find longest prefix match
    int common = 0;
    if (has_cache) {
        size_t min_len = std::min(cached_prefix.size(), tokens.size());
        while (common < (int)min_len && cached_prefix[common] == tokens[common]) common++;
        if (common > 0) std::cerr << "\r  cache hit: " << common << "/" << tokens.size() << " tokens\n";
    }
    if (common > 0 && has_cache) {
        // Restore cached GLA state & last hidden state
        int gla_sz = L * Nh * D * D;
        std::memcpy(gla_state.data(), cached_gla_state.data(), gla_sz * sizeof(float));
        int TS = cached_prefix.size();
        std::memcpy(x.data(), cached_x.data(), TS * H * sizeof(float));
    } else {
        common = 0;
    }
    // Embed new (uncached) suffix tokens
    int T = (int)tokens.size();
    auto& emb = w["model.embed_tokens.weight"];
    for (int t = common; t < T; t++) {
        int id = tokens[t];
        if (id >= 0 && id < emb.sh[0])
            std::memcpy(x.data() + t * H, emb.data() + id * H, H * sizeof(float));
    }
    // Process only suffix layers
    for (int li = 0; li < L; li++) {
        int batch_T = (li == 0 && common > 0) ? T : T; // re-process all if cache invalid
        // For first layers where we have cache, only run new tokens
        if (common > 0) {
            // Run layer on full sequence (needed for proper residual)
            layer(li, T);
        } else {
            layer(li, T);
        }
        if (li % 10 == 0) std::cerr << "\r  layer " << li << "/" << L;
    }
    std::cerr << "\r  layers done      \n";
    rmsnorm(x2, x, w["model.norm.weight"]);

    // Update cache
    cached_prefix = tokens;
    int gla_sz = L * Nh * D * D;
    cached_gla_state = Tensor({gla_sz});
    std::memcpy(cached_gla_state.data(), gla_state.data(), gla_sz * sizeof(float));
    cached_x = Tensor({T, H});
    std::memcpy(cached_x.data(), x.data(), T * H * sizeof(float));
    has_cache = true;

    // Self-diffusion L1
    if (cfg.self_diffusion_level >= SD_LEVEL1_TOKEN) {
        float* h_final = x2.data() + (T - 1) * H;
        apply_self_diffusion_level1(h_final, 1, H);
    }

    // Generation loop
    auto& lm = w["lm_head.weight"];
    int V = cfg.vocab_size;
    std::vector<int> result = tokens;
    for (int p = 0; p < n_pred; p++) {
        float* h = x2.data() + ((int)result.size() - 1) * H;
        for (int i = 0; i < V; i++) {
            float s = 0;
            for (int j = 0; j < H; j++) s += h[j] * lm.data()[i * H + j];
            logits.data()[i] = s;
        }
        int id = sample(logits.data(), V, temp, top_p);
        result.push_back(id);
        if (id == 0 || id == 2) break;
        if ((int)result.size() > cfg.max_seq) result.erase(result.begin());
        if (p < n_pred - 1) {
            // Single-token forward for next step
            auto next = forward(result, 0, temp, top_p);
            logits = next;
        }
    }

    // Self-diffusion L2/L3
    if (cfg.self_diffusion_level >= SD_LEVEL2_SEQUENCE && n_pred > 0) {
        int ar_T = (int)result.size();
        float mean_conf = 0;
        for (int t = 0; t < ar_T; t++) {
            mean_conf += score_confidence_ensemble(x2.data() + t * H, H);
        }
        mean_conf /= ar_T;
        if (mean_conf < cfg.self_diffusion_threshold) {
            std::cerr << "\r  self-diffusion L2 (conf=" << mean_conf << ")\n";
            apply_self_diffusion_level2(ar_T, H);
        }
    }
    if (cfg.self_diffusion_level >= SD_LEVEL3_SELFCORRECT && n_pred > 0) {
        int ar_T = (int)result.size();
        float mean_conf = 0;
        for (int t = 0; t < ar_T; t++) {
            mean_conf += score_confidence_ensemble(x2.data() + t * H, H);
        }
        mean_conf /= ar_T;
        if (mean_conf < cfg.self_diffusion_correction_threshold) {
            std::cerr << "\r  self-diffusion L3 (conf=" << mean_conf << ")\n";
            apply_self_diffusion_level3(result, tokens, n_pred, temp, top_p);
        }
    }

    Tensor r;
    r.d.resize(result.size());
    for (size_t i = 0; i < result.size(); i++) r.d[i] = (float)result[i];
    return r;
}

// === Speculative decoding: n-gram draft ===
void YashaModel::draft_ngram(const int* ctx, int ctx_len, int* draft, int K,
                            const float* orig_logits, const float* emb, int H, int V) {
    // Build simple unigram + bigram probs from the logits distribution
    // Draft by sampling from a smoothed mix of unigram (from logits) and bigram (repetition penalty)
    for (int k = 0; k < K; k++) {
        // Use the model's own logits distribution with temperature annealing
        float temp_k = 0.8f + k * 0.05f; // slight temp increase for later positions
        float max_l = *std::max_element(orig_logits, orig_logits + V);
        std::vector<std::pair<float,int>> cand;
        float sum = 0;
        for (int i = 0; i < V; i++) {
            float v = std::exp((orig_logits[i] - max_l) / temp_k);
            // Bigram penalty: reduce prob of recently generated tokens
            for (int r = std::max(0, ctx_len + k - 3); r < ctx_len + k; r++) {
                if (r < ctx_len + k && (r < ctx_len ? ctx[r] : draft[r - ctx_len]) == i)
                    v *= 0.5f;
            }
            cand.push_back({v, i});
            sum += v;
        }
        std::sort(cand.begin(), cand.end(), [](auto& a, auto& b) { return a.first > b.first; });
        float cum = 0;
        float r = std::uniform_real_distribution<float>(0, sum)(::rng());
        for (auto& [v, id] : cand) {
            cum += v;
            if (r < cum) { draft[k] = id; break; }
        }
    }
}

Tensor YashaModel::forward_speculative(const std::vector<int>& tokens, int n_pred, float temp, float top_p) {
    int H = cfg.hidden_size, V = cfg.vocab_size;
    auto& lm = w["lm_head.weight"];
    std::vector<int> result = tokens;
    int K = std::min(cfg.spec_draft, n_pred);

    while ((int)result.size() - (int)tokens.size() < n_pred) {
        int rem = n_pred - ((int)result.size() - (int)tokens.size());
        K = std::min(K, rem);

        // Run forward to get hidden state + logits
        Tensor r = forward(result, 0, temp, top_p); // n_pred=0 => just logits

        float* h = x2.data() + ((int)result.size() - 1) * H;

        // Draft K tokens from the logits distribution
        int draft[16];
        float* emb_ptr = w.count("model.embed_tokens.weight") ? w["model.embed_tokens.weight"].data() : nullptr;
        draft_ngram(result.data(), (int)result.size(), draft, K,
                    logits.data(), emb_ptr, H, V);

        // Verify all K at once by appending drafts and running forward
        std::vector<int> verify_seq = result;
        for (int k = 0; k < K; k++) verify_seq.push_back(draft[k]);
        Tensor vr = forward(verify_seq, 0, temp, top_p);
        float* vh = x2.data() + ((int)verify_seq.size() - 1) * H;

        // Speculatively accept: score the drafted path, if confident accept all
        float conf = score_confidence(vh, H);
        int accept;
        if (conf > 0.7f) {
            accept = K; // accept all
        } else if (conf > 0.4f) {
            accept = std::max(1, K / 2); // accept half
        } else {
            accept = 1; // accept just first
        }

        for (int k = 0; k < accept; k++) {
            result.push_back(draft[k]);
            if (draft[k] == 0 || draft[k] == 2) break;
        }

        if (accept < K) {
            // Roll back remaining and sample one normally
            int id = sample(logits.data(), V, temp, top_p);
            result.push_back(id);
            if (id == 0 || id == 2) break;
        }
    }

    Tensor out;
    for (size_t i = tokens.size(); i < result.size(); i++) out.d.push_back((float)result[i]);
    out.sh = {(int)out.d.size()};
    return out;
}

// === Iterative refinement ===
Tensor YashaModel::forward_refined(const std::vector<int>& tokens, int n_pred, float temp, float top_p) {
    // Generate initial answer
    Tensor initial = forward(tokens, n_pred, temp, top_p);

    // Build critique prompt: encode "Check your work carefully. What did you miss?"
    std::vector<int> critique_prompt = tokens;
    std::string prompt_str = "Check your work carefully. What did you miss?";
    auto critique_ids = encode_text(prompt_str);
    // Append critique + initial output as context, then regenerate
    for (int id : critique_ids) critique_prompt.push_back(id);
    for (size_t i = tokens.size(); i < (size_t)initial.numel(); i++)
        critique_prompt.push_back((int)initial.d[i]);

    // Generate refined answer
    Tensor refined = forward(critique_prompt, n_pred, temp, top_p);

    // Score both
    float initial_score = 0;
    float refined_score = 0;
    int H = cfg.hidden_size;
    if ((int)initial.numel() > (int)tokens.size()) {
        // Get last hidden state for initial
        Tensor initial_forward = forward(tokens, 0, temp, top_p);
        (void)initial_forward;
        float* hi = x2.data() + ((int)tokens.size() - 1) * H;
        initial_score = score_confidence(hi, H);
    }
    if ((int)refined.numel() > (int)critique_prompt.size()) {
        float* hr = x2.data() + ((int)critique_prompt.size() - 1) * H;
        refined_score = score_confidence(hr, H);
    }

    // Return whichever scored higher
    if (refined_score > initial_score + 0.05f) {
        Tensor out;
        for (size_t i = critique_prompt.size(); i < (size_t)refined.numel(); i++)
            out.d.push_back((float)refined.d[i]);
        out.sh = {(int)out.d.size()};
        return out;
    }
    return initial;
}

// === Model loading ===
bool YashaModel::load(const std::string& dir) {
    std::cerr << "Loading model from " << dir << "...\n";
    for (auto& entry : fs::directory_iterator(dir)) {
        if (entry.path().extension() == ".safetensors") {
            std::cerr << "  " << entry.path().filename() << "\n";
            load_safetensors(entry.path().string(), w);
        }
    }
    std::cerr << "Loaded " << w.size() << " tensors\n";
    return !w.empty();
}

// === BPE tokenizer stub ===
std::vector<int> encode_text(const std::string& text) {
    std::vector<int> ids;
    for (char c : text) ids.push_back((int)(unsigned char)c + 3);
    return ids;
}

std::string decode_ids(const std::vector<int>& ids) {
    std::string s;
    for (int id : ids) {
        if (id >= 3 && id < 259) s += (char)(id - 3);
        else if (id == 0 || id == 1 || id == 2) {}
        else s += "\xef\xbf\xbd";
    }
    return s;
}