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

// Per-definition constants for this mha bf16 specialisation.
// Non-causal softmax(Q·K^T / sqrt(head_dim))·V per head.
// Q/K/V are raw bf16 bit patterns (uint16_t); output accumulates/returns in
// fp32.
namespace mha_def {
constexpr int NHeads  = 16;
constexpr int HeadDim = 128;
constexpr int KvHeads = 16;
} // namespace mha_def

#ifdef __cplusplus
extern "C" {
#endif
// LLM target: implement this in kernel.cpp.
// M, S are the var dims. Q: (M, n_heads, head_dim) bf16; K/V: (S, kv_heads,
// head_dim) bf16; output: (M, n_heads, head_dim) float32.
// 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_mha(const uint16_t* Q, float* output, const uint16_t* K, const uint16_t* V,
               int M, int S);
#ifdef __cplusplus
}
#endif