// File: kernel.cpp
// Reference-scalar mla_prefill bf16 (fp32 accumulate/output).
// LLM target: replace this file with an optimised inner_mla_prefill.
// All per-definition constants live in mla_prefill_def:: (mla_prefill.h).
#include "mla_prefill.h"
#include <cmath>
#include <cstring>
#include <vector>
using namespace mla_prefill_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_prefill(const uint16_t* Q, float* output, const uint16_t* K,
                                   const uint16_t* V, int M, int S)
{
    const float scale = 1.0f / std::sqrt((float)QkDim);
    const int delta = S - M;
    std::vector<float> scores(S);

    for (int h = 0; h < NHeads; ++h) {
        for (int m = 0; m < M; ++m) {
            const uint16_t* q_row = Q + ((long)m * NHeads + h) * QkDim;
            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 * NHeads + h) * QkDim;
                float dot = 0.0f;
                for (int d = 0; d < QkDim; ++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];
            }

            float* out_row = output + ((long)m * NHeads + h) * VoDim;
            for (int d = 0; d < VoDim; ++d) out_row[d] = 0.0f;
            for (int s = 0; s <= smax; ++s) {
                const float w = scores[s] / sum;
                const uint16_t* v_row = V + ((long)s * NHeads + h) * VoDim;
                for (int d = 0; d < VoDim; ++d) {
                    out_row[d] += w * bf16_to_f32(v_row[d]);
                }
            }
        }
    }
}