import type { Model } from "@/lib/schemas"

const models: Model[] = [
  {
    id: "deepseek-v3",
    name: "DeepSeek V3/R1",
    description: "DeepSeek V3 and R1 models.",
    modules: {
      DeepseekV3ForCausalLM: {
        count: 1,
        type: "block",
      },
      embed_tokens: {
        count: 1,
        parent: "DeepseekV3ForCausalLM",
        type: "layer",
        definitions: [],
      },
      DeepseekV3DecoderLayer: {
        count: 61,
        parent: "DeepseekV3ForCausalLM",
        type: "block",
        definitions: [],
      },
      input_layernorm: {
        count: 61,
        parent: "DeepseekV3DecoderLayer",
        type: "layer",
        definitions: ["rmsnorm_h7168", "fused_add_rmsnorm_h7168"],
      },
      self_attn: {
        count: 61,
        parent: "DeepseekV3DecoderLayer",
        type: "block",
        definitions: [],
      },
      fused_qkv_a_proj_with_mqa: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      q_a_layernorm: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: ["rmsnorm_h1536"],
      },
      q_b_proj: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      kv_a_layernorm: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: ["rmsnorm_h512"],
      },
      kv_b_proj: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      rotary_emb: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      attn_mla: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [
          "mla_ragged_prefill_causal_h16_qk192_vo128",
          "mla_paged_prefill_causal_h16_ckv512_kpe64_ps1",
          "mla_paged_decode_h16_ckv512_kpe64_ps1",
        ],
      },
      o_proj: {
        count: 61,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      post_attention_layernorm: {
        count: 61,
        parent: "DeepseekV3DecoderLayer",
        type: "layer",
        definitions: ["rmsnorm_h7168", "fused_add_rmsnorm_h7168"],
      },
      mlp: {
        count: 61,
        parent: "DeepseekV3DecoderLayer",
        type: "block",
        definitions: [],
      },
      mlp_dense: {
        count: 31,
        parent: "mlp",
        type: "block",
        definitions: [],
      },
      gate_up_proj: {
        count: 31,
        parent: "mlp_dense",
        type: "layer",
        definitions: [],
      },
      act_fn: {
        count: 31,
        parent: "mlp_dense",
        type: "layer",
        definitions: [],
      },
      down_proj: {
        count: 31,
        parent: "mlp_dense",
        type: "layer",
        definitions: [],
      },
      moe: {
        count: 30,
        parent: "mlp",
        type: "block",
        definitions: [],
      },
      moe_gate: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: ["moe_fp8_block_scale_ds_routing_topk8_ng8_kg4_e32_h7168_i2048"],
      },
      moe_topk: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: ["moe_fp8_block_scale_ds_routing_topk8_ng8_kg4_e32_h7168_i2048"],
      },
      moe_experts: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: ["moe_fp8_block_scale_ds_routing_topk8_ng8_kg4_e32_h7168_i2048"],
      },
      norm: {
        count: 1,
        parent: "DeepseekV3ForCausalLM",
        type: "layer",
        definitions: ["rmsnorm_h7168", "fused_add_rmsnorm_h7168"],
      },
      lm_head: {
        count: 1,
        parent: "DeepseekV3ForCausalLM",
        type: "layer",
        definitions: [],
      },
    },
  },
  {
    id: "llama-3.1-8b",
    name: "Llama 3.1 8B",
    description: "Meta's Llama 3.1 8B parameter model",
    modules: {
      LlamaForCausalLM: {
        count: 1,
        type: "block",
        definitions: [],
      },
      LlamaModel: {
        count: 1,
        parent: "LlamaForCausalLM",
        type: "block",
        definitions: [],
      },
      embed_tokens: {
        count: 1,
        parent: "LlamaModel",
        type: "layer",
        definitions: [],
      },
      layers: {
        count: 32,
        parent: "LlamaModel",
        type: "block",
        definitions: [],
      },
      LlamaDecoderLayer: {
        count: 32,
        parent: "layers",
        type: "block",
        definitions: [],
      },
      input_layernorm: {
        count: 32,
        parent: "LlamaDecoderLayer",
        type: "layer",
        definitions: ["rmsnorm_h4096", "fused_add_rmsnorm_h4096"],
      },
      self_attn: {
        count: 32,
        parent: "LlamaDecoderLayer",
        type: "block",
        definitions: [],
      },
      qkv_proj: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: ["gemm_n_6144_k_4096"],
      },
      rotary_emb: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      attn: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [
          "gqa_paged_prefill_causal_h32_kv8_d128_ps1",
          "gqa_paged_decode_h32_kv8_d128_ps1",
          "gqa_ragged_prefill_causal_h32_kv8_d128",
        ],
      },
      o_proj: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: ["gemm_n_4096_k_4096"],
      },
      post_attention_layernorm: {
        count: 32,
        parent: "LlamaDecoderLayer",
        type: "layer",
        definitions: ["fused_add_rmsnorm_h4096"],
      },
      mlp: {
        count: 32,
        parent: "LlamaDecoderLayer",
        type: "block",
        definitions: [],
      },
      gate_up_proj: {
        count: 32,
        parent: "mlp",
        type: "layer",
        definitions: ["gemm_n_28672_k_4096"],
      },
      act_fn: {
        count: 32,
        parent: "mlp",
        type: "layer",
        definitions: [],
      },
      down_proj: {
        count: 32,
        parent: "mlp",
        type: "layer",
        definitions: ["gemm_n_4096_k_14336"],
      },
      norm: {
        count: 1,
        parent: "LlamaModel",
        type: "layer",
        definitions: ["fused_add_rmsnorm_h4096"],
      },
      lm_head: {
        count: 1,
        parent: "LlamaForCausalLM",
        type: "layer",
        definitions: [],
      },
    },
  },
  {
    id: "qwen3-30b-a3b",
    name: "Qwen3 30 A3B",
    description: "Qwen3 MoE 30B a3b model.",
    modules: {
      Qwen3MoeModel: {
        count: 1,
        type: "block",
        definitions: [],
      },
      embed_tokens: {
        count: 1,
        parent: "Qwen3MoeModel",
        type: "layer",
        definitions: [],
      },
      layers: {
        count: 32,
        parent: "Qwen3MoeModel",
        type: "block",
        definitions: [],
      },
      Qwen3MoeDecoderLayer: {
        count: 32,
        parent: "layers",
        type: "block",
        definitions: [],
      },
      self_attn: {
        count: 32,
        parent: "Qwen3MoeDecoderLayer",
        type: "block",
        definitions: [],
      },
      qkv_proj: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      q_norm: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: ["rmsnorm_h128"],
      },
      k_norm: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: ["rmsnorm_h128"],
      },
      rotary_emb: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      attn: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [
          "gqa_paged_prefill_causal_h32_kv4_d128_ps1",
          "gqa_ragged_prefill_causal_h32_kv4_d128",
          "gqa_paged_decode_h32_kv4_d128_ps1",
        ],
      },
      o_proj: {
        count: 32,
        parent: "self_attn",
        type: "layer",
        definitions: [],
      },
      mlp: {
        count: 32,
        parent: "Qwen3MoeDecoderLayer",
        type: "block",
        definitions: [],
      },
      moe: {
        count: 30,
        parent: "mlp",
        type: "block",
        definitions: [],
      },
      moe_gate: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: [],
      },
      moe_topk: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: [],
      },
      moe_experts: {
        count: 30,
        parent: "moe",
        type: "layer",
        definitions: [],
      },
      input_layernorm: {
        count: 32,
        parent: "Qwen3MoeDecoderLayer",
        type: "layer",
        definitions: ["rmsnorm_h2048", "fused_add_rmsnorm_h2048"],
      },
      post_attention_layernorm: {
        count: 32,
        parent: "Qwen3MoeDecoderLayer",
        type: "layer",
        definitions: ["fused_add_rmsnorm_h2048"],
      },
      norm: {
        count: 1,
        parent: "Qwen3MoeModel",
        type: "layer",
        definitions: ["fused_add_rmsnorm_h2048", "rmsnorm_h2048"],
      },
    },
  },
] satisfies Model[]

export default models
