// File: mla_prefill.h
#pragma once
#include <cstdint>

// Per-definition constants for this mla_prefill bf16 specialisation
// (causal MLA prefill, decompressed/naive form).
namespace mla_prefill_def {
constexpr int NHeads = 16;
constexpr int QkDim  = 192;
constexpr int VoDim  = 128;
} // namespace mla_prefill_def

#ifdef __cplusplus
extern "C" {
#endif
// LLM target: implement this in kernel.cpp.
// M, S are the var dims. Q: (M, n_heads, qk_dim) bf16; K: (S, n_heads, qk_dim)
// bf16; V: (S, n_heads, vo_dim) bf16; output: (M, n_heads, vo_dim) float32.
// KV-cache-style causal softmax(Q.K^T * scale)*V per head: kv position s is
// visible to query position m iff s <= m + (S - M).
// Bytes are raw bf16 bit patterns -- reinterpret_cast<const __bf16*> if your
// target ISA has native bf16 support, or widen manually to fp32 first.
void inner_mla_prefill(const uint16_t* Q, float* output, const uint16_t* K,
                        const uint16_t* V, int M, int S);
#ifdef __cplusplus
}
#endif