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

namespace {
constexpr int64_t kNHeads = 16;
constexpr int64_t kQkDim  = 192;
constexpr int64_t kVoDim  = 128;
} // namespace

extern "C" {
// inputs: [0]=Q bf16 [M,n_heads,qk_dim], [1]=K bf16 [S,n_heads,qk_dim],
// [2]=V bf16 [S,n_heads,vo_dim]; var_axes: [0]=M, [1]=S
int armbench_entry_mla_prefill(const void* const* inputs, void* output,
                               const int64_t* var_axes)
{
    return armbench_llamacpp_mla_prefill(
        inputs[0], inputs[1], inputs[2],
        reinterpret_cast<float*>(output),
        var_axes[0], var_axes[1], kNHeads, kQkDim, kVoDim);
}
} // extern "C"