// File: binding.cpp
#include "gqa_bf16_contract.h"

namespace {
constexpr int64_t kNHeads  = 32;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kKvHeads = 8;
} // namespace

extern "C" {
// inputs: [0]=Q bf16 [M,n_heads,head_dim], [1]=K bf16 [S,kv_heads,head_dim],
// [2]=V bf16 [S,kv_heads,head_dim]; var_axes: [0]=M, [1]=S
int armbench_entry_gqa(const void* const* inputs, void* output,
                       const int64_t* var_axes)
{
    return armbench_llamacpp_gqa(
        inputs[0], inputs[1], inputs[2],
        reinterpret_cast<float*>(output),
        var_axes[0], var_axes[1], kNHeads, kHeadDim, kKvHeads);
}
} // extern "C"