include/ling3/model_format.h

#pragma once

#include <array>
#include <cstddef>
#include <cstdint>

namespace ling3 {

inline constexpr std::array<char, 8> kPackageMagic = {'L', '3', 'R', 'K', 'N', 'N', '1', '\0'};
inline constexpr std::uint32_t kPackageVersion = 1;
inline constexpr std::uint32_t kEndianTag = 0x01020304U;
inline constexpr std::size_t kPackageAlignment = 4096;
inline constexpr std::uint32_t kPackageOfficialInt4 = 0x100U;
// Per-entry W4 or BF16 source weights; BF16 entries execute through W8A8.
inline constexpr std::uint32_t kPackageMixedW4W8 = 0x200U;

enum class DataType : std::uint32_t {
    kUnknown = 0,
    kBFloat16 = 1,
    kFloat16 = 2,
    kFloat32 = 3,
    kInt8 = 4,
    kInt4Low = 5,
    kInt16 = 6,
    kInt32 = 7,
    kRknnModel = 8,
};

enum class TensorRole : std::uint32_t {
    kUnknown = 0,
    kEmbedding = 1,
    kNorm = 2,
    kLinearWeight = 3,
    kLinearScale = 4,
    kBias = 5,
    kRknnIsland = 6,
    kTokenizer = 7,
};

enum class QuantType : std::uint32_t {
    kNone = 0,
    kPerTensor = 1,
    kPerOutputChannel = 2,
    kFixedActivation = 3,
};

enum class TensorLayout : std::uint32_t {
    kRowMajor = 0,
    kRknnNativeB = 1,
    kPackedInt4Low = 2,
    kOpaque = 3,
};

struct PackageHeader {
    char magic[8];
    std::uint32_t version;
    std::uint32_t header_bytes;
    std::uint32_t endian_tag;
    std::uint32_t tensor_count;
    std::uint32_t tensor_entry_bytes;
    std::uint32_t flags;
    std::uint64_t tensor_table_offset;
    std::uint64_t string_table_offset;
    std::uint64_t string_table_bytes;
    std::uint64_t payload_offset;
    std::uint64_t payload_bytes;
    std::uint64_t file_bytes;
    std::uint8_t source_revision[20];
    std::uint32_t header_crc32;
    std::uint32_t vocab_size;
    std::uint32_t hidden_size;
    std::uint32_t layer_count;
    std::uint32_t attention_heads;
    std::uint32_t head_dim;
    std::uint32_t kv_lora_rank;
    std::uint32_t q_lora_rank;
    std::uint32_t qk_nope_dim;
    std::uint32_t qk_rope_dim;
    std::uint32_t value_head_dim;
    std::uint32_t dense_ffn_dim;
    std::uint32_t expert_ffn_dim;
    std::uint32_t shared_ffn_dim;
    std::uint32_t expert_count;
    std::uint32_t experts_per_token;
    std::uint32_t expert_group_count;
    std::uint32_t selected_group_count;
    std::uint32_t layer_group_size;
    std::uint32_t leading_dense_layers;
    std::uint32_t convolution_kernel;
    std::uint32_t max_context;
    std::uint32_t eos_token;
    std::uint32_t pad_token;
    std::uint32_t bos_token;
    std::uint32_t mla_layer_count;
    std::uint32_t kda_layer_count;
    std::uint32_t default_weight_bits;
    std::uint32_t default_activation_bits;
    float rms_epsilon;
    float rope_theta;
    float routed_scale;
    float kda_lower_bound;
    std::uint8_t reserved[24];
};

struct TensorEntry {
    std::uint32_t name_offset;
    std::uint32_t name_bytes;
    std::uint32_t dtype;
    std::uint32_t rank;
    std::uint32_t role;
    std::uint32_t quant;
    std::uint32_t layout;
    std::uint32_t flags;
    std::uint32_t dims[4];
    std::uint64_t data_offset;
    std::uint64_t data_bytes;
    std::uint64_t aux_offset;
    std::uint64_t aux_bytes;
    float scale;
    float zero_point;
    std::uint32_t layer;
    std::uint32_t op;
    std::uint32_t expert;
    std::uint32_t core;
    std::uint8_t sha256[32];
    std::uint8_t reserved[24];
};

static_assert(sizeof(PackageHeader) == 256);
static_assert(sizeof(TensorEntry) == 160);

} // namespace ling3