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

// Per-definition constants for this moe Q4_K (weight-only) 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 is a raw bf16 bit pattern (uint16_t) -- never quantized.
// router_weight is 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).
// gate_proj/up_proj/down_proj are each Q4_K-quantized, exposed as three flat
// tensors per weight (see gemm_q4_k_m.h.tmpl for the exact byte layout):
//   <name>_q4     [n_expert, rows, cols/2]   uint8
//   <name>_scales [n_expert, rows, cols/32]  raw fp16 bits (uint16_t)
//   <name>_mins   [n_expert, 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 NEmbdHalf   = 1024;
constexpr int NEmbdSub    = 64;
constexpr int NFfHalf     = 512;
constexpr int NFfSub      = 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 uint16_t* hidden_states, float* output,
               const float* router_weight,
               const uint8_t* gate_proj_q4, const uint16_t* gate_proj_scales,
               const uint16_t* gate_proj_mins,
               const uint8_t* up_proj_q4, const uint16_t* up_proj_scales,
               const uint16_t* up_proj_mins,
               const uint8_t* down_proj_q4, const uint16_t* down_proj_scales,
               const uint16_t* down_proj_mins,
               int n_tokens);
#ifdef __cplusplus
}
#endif