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

// Per-definition constants for this moe bf16 specialisation.
// SwiGLU MoE with softmax-then-top-k routing and normalized top-k weights:
//   probs = softmax(hidden @ router^T); top-k experts; w /= sum(w)
//   out   = sum_k w_k * down_e @ (silu(gate_e @ x) * (up_e @ x))
// hidden_states/router_weight/gate_proj/up_proj/down_proj are raw bf16 bit
// patterns (uint16_t); output accumulates/returns in fp32.
namespace moe_def {
constexpr int NEmbd       = 2048;
constexpr int NFf         = 1024;
constexpr int NExpert     = 64;
constexpr int NExpertUsed = 8;
} // namespace moe_def

#ifdef __cplusplus
extern "C" {
#endif
// LLM target: implement this in kernel.cpp.
// n_tokens is the only var dim. hidden_states: (n_tokens, n_embd) bf16;
// router_weight: (n_expert, n_embd) bf16; gate/up_proj: (n_expert, n_ff,
// n_embd) bf16; down_proj: (n_expert, n_embd, n_ff) bf16;
// output: (n_tokens, n_embd) float32.
// 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_moe(const uint16_t* hidden_states, float* output,
               const uint16_t* router_weight, const uint16_t* gate_proj,
               const uint16_t* up_proj, const uint16_t* down_proj,
               int n_tokens);
#ifdef __cplusplus
}
#endif