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

// Per-definition constants for this moe Q8_0 (full quantization: activations AND
// weights) 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 + gate_proj/up_proj/down_proj are all int8 with a separate
// per-32-element-block fp16 scale tensor (Definition-level ABI, not GGML's
// packed block_q8_0 struct). router_weight stays plain float32 -- also never
// quantized (routing decisions are categorical top-k; quantization rounding
// could flip which experts get selected, so routing stays at the highest
// available precision).
//   <name>        [..., rows, cols]     int8
//   <name>_scales [..., rows, cols/32]  raw fp16 bits (uint16_t)
// where (rows, cols) = (n_ff, n_embd) for gate/up, (n_embd, n_ff) for down.
// Output accumulates/returns in fp32.
namespace moe_def {
constexpr int NEmbd       = 2048;
constexpr int NFf         = 1024;
constexpr int NExpert     = 64;
constexpr int NExpertUsed = 8;
constexpr int NEmbdBlk    = 64;
constexpr int NFfBlk      = 32;
} // namespace moe_def

#ifdef __cplusplus
extern "C" {
#endif
// LLM target: implement this in kernel.cpp.
// n_tokens is the only var dim.
void inner_moe(const int8_t* hidden_states, float* output,
               const uint16_t* hs_scales, const float* router_weight,
               const int8_t* gate_proj, const uint16_t* gate_scales,
               const int8_t* up_proj, const uint16_t* up_scales,
               const int8_t* down_proj, const uint16_t* down_scales,
               int n_tokens);
#ifdef __cplusplus
}
#endif