include/ling3/router.h

#pragma once

#include <array>
#include <cstddef>

namespace ling3 {

inline constexpr std::size_t kHiddenSize = 1536;
inline constexpr std::size_t kExpertCount = 128;
inline constexpr std::size_t kExpertsPerToken = 8;
inline constexpr std::size_t kExpertGroups = 8;
inline constexpr std::size_t kSelectedGroups = 4;

struct Route {
    std::array<int, kExpertsPerToken> experts{};
    std::array<float, kExpertsPerToken> weights{};
};

void RouterLogitsF32(const float * hidden, const float * weights, float * logits) noexcept;
Route SelectRoute(const float * logits, const float * expert_bias, float routed_scale = 2.5F) noexcept;

} // namespace ling3