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

// Per-definition constants for this gqa bf16 specialisation.
namespace gqa_def {
constexpr int NHeads  = 32;
constexpr int HeadDim = 128;
constexpr int KvHeads = 8;
} // namespace gqa_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.
// KV-cache-style causal softmax(Q.K^T * scale)*V per head, with K/V
// broadcast across each group of (n_heads / kv_heads) query heads
// (repeat_interleave along the head axis -- query head h reads kv head
// h / (n_heads / kv_heads)). 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_gqa(const uint16_t* Q, float* output, const uint16_t* K,
               const uint16_t* V, int M, int S);
#ifdef __cplusplus
}
#endif