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

namespace {
constexpr int64_t kNHeads     = 16;
constexpr int64_t kCQ         = 576;
constexpr int64_t kKvLoraRank = 512;
constexpr int64_t kVHeadDim   = 128;
constexpr int64_t kQkTrueDim  = 128 + 64;
} // namespace

extern "C" {
// inputs: [0]=Q bf16 [M,n_heads,c_q], [1]=K bf16 [S,c_q], [2]=V bf16
// [S,kv_lora_rank], [3]=v_mla bf16 [n_heads,v_head_dim,kv_lora_rank];
// var_axes: [0]=M, [1]=S
int armbench_entry_mla_decode(const void* const* inputs, void* output,
                              const int64_t* var_axes)
{
    return armbench_llamacpp_mla_decode(
        inputs[0], inputs[1], inputs[2], inputs[3],
        reinterpret_cast<float*>(output),
        var_axes[0], var_axes[1],
        kNHeads, kCQ, kKvLoraRank, kVHeadDim, kQkTrueDim);
}
} // extern "C"