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

// Harness contract for the llama.cpp (ggml) mla_prefill baseline.
// Causal MLA prefill (decompressed/naive form). Called by
// armbench_entry_mla_prefill (binding.cpp); implemented by kernel.cpp.
// Q is row-major bf16 [M, n_heads, qk_dim]; K/V are row-major bf16
// [S, n_heads, qk_dim]/[S, n_heads, vo_dim]; out is row-major fp32
// [M, n_heads, vo_dim].
// KV-cache-style causal softmax(Q.K^T * scale)*V per head: kv position s is
// visible to query position m iff s <= m + (S - M).
int armbench_llamacpp_mla_prefill(const void* Q, const void* K, const void* V,
                                  float* out, int64_t M, int64_t S,
                                  int64_t n_heads, int64_t qk_dim, int64_t vo_dim);