Skip to main content

ClikaRT::nn::MoE

class

Header: ClikaRT/nn/moe.h
Inherits: ClikaRT::nn::Module

Dense mixture-of-experts module: a router picks top-k experts per token, each expert runs its gated FFN, and the outputs combine under the router weights.

The expert weights live as stacked [E, ...] tensors; the optional trailing MoeOptions carries the routing variants (grouped routing, sparse-mixer, the clamped-SwiGLU knobs, per-expert combine scales). The module owns its weights: construct with make(...), or declare shapes and bind a checkpoint via load_state_dict. After the first forward (or initialize()) the weight lives ONLY in the backend's packed form (one resident copy); to(dtype) restores, casts, and repacks on the next forward.

auto moe = ClikaRT::nn::MoE::make(router_w, gate_up_w, down_w, opts);
auto y   = moe->forward(x);

Static member functions​

make()​

static std::shared_ptr<MoE> make(
    std::int64_t num_experts,
    std::int64_t hidden_size,
    std::int64_t intermediate_size,
    std::int64_t top_k,
    bool gate_up_bias = false,
    bool down_bias = false,
    DataType dtype = DataType::Float32,
    Device device = Device::cpu(),
    MoeOptions options = {}
)

Defaults: gate_up_bias / down_bias = false (no bias slots), dtype = Float32, device = CPU, options = MoeOptions{} (each field's own default; see the struct). The routing variants ride the optional trailing options; the four counts are the whole required surface.

Throws

  • ClikaRT::Error: as stated above.

Declared in ClikaRT/nn/moe.h, line 113

Member functions​

~MoE()​

~MoE() override

Releases the module's packed expert weights.

Declared in ClikaRT/nn/moe.h, line 125

set_weights()​

void set_weights(
    Tensor gate_up_experts,
    Tensor down_experts,
    OptionalTensor gate_up_bias = {},
    OptionalTensor down_bias = {},
    OptionalTensor gate_experts = {},
    OptionalTensor gate_bias = {},
    OptionalTensor e_score_correction_bias = {}
)

Bind the declared slots positionally: gate_up_experts [E, F·I, H] and down_experts [E, H, I] (required), the declared per-expert biases, the Unfused gate_experts [E, I, H] (+ gate_bias), and the aux-free selection bias [E] (bindable even when not declared; it is routing data, not a declared parameter). The declared placement wins; shapes must match the declaration. Re-binding after a pack drops the pack; the next forward re-packs. Raises ClikaRT::Error on a geometry mismatch or a bias without its declared slot.

Declared in ClikaRT/nn/moe.h, line 139

initialize_impl()​

virtual Result<void> initialize_impl() override

Optional warm-up: run the first-forward pack NOW (idempotent, thread-safe). Every declared slot must hold a real (loaded) tensor; a still-fake slot is a clean error. After the pack the expert weight slots are released (the packed form is the single resident copy); biases stay enumerable.

Declared in ClikaRT/nn/moe.h, line 156

to_impl(StreamOrDevice)​

virtual Result<void> to_impl(StreamOrDevice where) override

Move to a placement (the pack rebuilds on the target) or cast to a dtype (the raw experts restore from the pack, cast, and re-pack on the next forward).

Declared in ClikaRT/nn/moe.h, line 161

to_impl(DataType)​

virtual Result<void> to_impl(DataType dtype) override

Declared in ClikaRT/nn/moe.h, line 162

forward()​

Tensor forward(
    Tensor x,
    Tensor router_logits,
    OptionalTensor shared_output = {}
) const

One fused MoE step: x [T, H] (packed/ragged rows are native, no cu_seqlens; MoE is per-token) and the caller-computed router_logits [T, E] → [T, H]. The first call packs (once, thread-safe); every later call reuses the pack. Raises ClikaRT::Error on a shape/dtype mismatch or a still-fake slot. An optional shared_output (the shared-expert branch's [T, H] at x's dtype, EXACT shape) folds into the routed combine (the fused routed+shared epilogue; no separate add pass).

Declared in ClikaRT/nn/moe.h, line 174