// File: kernel.cpp
// Reference-scalar mla_decode bf16 (fp32 accumulate/output).
// LLM target: replace this file with an optimised inner_mla_decode.
// All per-definition constants live in mla_decode_def:: (mla_decode.h).
#include "mla_decode.h"
#include <cmath>
#include <cstring>
#include <vector>
using namespace mla_decode_def;

namespace {
// bf16 shares fp32's exponent field, so widening is an exact bit-shift -- no
// rounding/subnormal handling needed (unlike fp16->fp32).
inline float bf16_to_f32(uint16_t bits) {
    uint32_t b = (uint32_t)bits << 16;
    float f;
    std::memcpy(&f, &b, sizeof(f));
    return f;
}
} // namespace

extern "C" void inner_mla_decode(const uint16_t* Q, float* output, const uint16_t* K,
                                  const uint16_t* V, const uint16_t* v_mla, int M, int S)
{
    constexpr int QkTrueDim = QkNopeHeadDim + QkRopeHeadDim;
    const float scale = 1.0f / std::sqrt((float)QkTrueDim);
    const int delta = S - M;
    std::vector<float> scores(S);
    std::vector<float> ctx(KvLoraRank);

    for (int h = 0; h < NHeads; ++h) {
        for (int m = 0; m < M; ++m) {
            const uint16_t* q_row = Q + ((long)m * NHeads + h) * CQ;
            const int smax = m + delta; // inclusive causal bound: s <= m + delta

            float max_score = -INFINITY;
            for (int s = 0; s <= smax; ++s) {
                const uint16_t* k_row = K + (long)s * CQ; // K has no head axis
                float dot = 0.0f;
                for (int d = 0; d < CQ; ++d) {
                    dot += bf16_to_f32(q_row[d]) * bf16_to_f32(k_row[d]);
                }
                scores[s] = dot * scale;
                if (scores[s] > max_score) max_score = scores[s];
            }

            float sum = 0.0f;
            for (int s = 0; s <= smax; ++s) {
                scores[s] = std::exp(scores[s] - max_score);
                sum += scores[s];
            }

            for (int l = 0; l < KvLoraRank; ++l) ctx[l] = 0.0f;
            for (int s = 0; s <= smax; ++s) {
                const float w = scores[s] / sum;
                const uint16_t* v_row = V + (long)s * KvLoraRank; // V has no head axis
                for (int l = 0; l < KvLoraRank; ++l) {
                    ctx[l] += w * bf16_to_f32(v_row[l]);
                }
            }

            // Up-project the compressed context by this head's v_mla slab:
            // v_mla row-major (n_heads, v_head_dim, kv_lora_rank).
            const uint16_t* vm_head = v_mla + (long)h * VHeadDim * KvLoraRank;
            float* out_row = output + ((long)m * NHeads + h) * VHeadDim;
            for (int d = 0; d < VHeadDim; ++d) {
                const uint16_t* vm_row = vm_head + (long)d * KvLoraRank;
                float acc = 0.0f;
                for (int l = 0; l < KvLoraRank; ++l) {
                    acc += ctx[l] * bf16_to_f32(vm_row[l]);
                }
                out_row[d] = acc;
            }
        }
    }
}