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

// Per-definition constants for this mla_decode bf16 specialisation
// (causal MLA decode, weight-absorbed/compressed form, MQA-style shared K/V
// latent).
namespace mla_decode_def {
constexpr int NHeads       = 16;
constexpr int KvLoraRank   = 512;
constexpr int QkRopeHeadDim = 64;
constexpr int QkNopeHeadDim = 128;
constexpr int VHeadDim     = 128;
constexpr int CQ           = 576;
} // namespace mla_decode_def

#ifdef __cplusplus
extern "C" {
#endif
// LLM target: implement this in kernel.cpp.
// M, S are the var dims. Q: (M, n_heads, c_q) bf16; K: (S, c_q) bf16 (shared
// across heads); V: (S, kv_lora_rank) bf16 (shared across heads); v_mla:
// (n_heads, v_head_dim, kv_lora_rank) bf16; output: (M, n_heads, v_head_dim)
// float32.
// KV-cache-style causal softmax(Q.K^T * scale)*V per head, then up-projected
// per head by v_mla ([kv_lora_rank] -> [v_head_dim]). scale uses the true
// un-absorbed head dim (qk_nope_head_dim + qk_rope_head_dim), not c_q.
// 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_mla_decode(const uint16_t* Q, float* output, const uint16_t* K,
                       const uint16_t* V, const uint16_t* v_mla, int M, int S);
#ifdef __cplusplus
}
#endif