{
  "architecture": {
    "attention_variant": "MLA",
    "family": "causal_lm",
    "is_moe": true,
    "mixer": "attention",
    "moe": {
      "experts_per_token": 6,
      "num_experts": 256,
      "num_shared_experts": 1
    },
    "positional": "rope",
    "sliding_window": 128,
    "tie_word_embeddings": false,
    "view": "decoder"
  },
  "capabilities": {
    "attention_backends": [
      "eager"
    ],
    "attention_patterns": [
      "compressed"
    ],
    "attention_schedule": [
      "heavily_compressed_attention",
      "heavily_compressed_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention",
      "compressed_sparse_attention",
      "heavily_compressed_attention"
    ],
    "kernels": {
      "RMSNorm": [
        "kernels-community/liger-kernels",
        "kernels-community/rmsnorm",
        "kernels-community/mlx_rmsnorm"
      ]
    },
    "task_heads": [
      "causal_lm"
    ],
    "tensor_parallel": false
  },
  "components": [
    {
      "children": [
        "embed_tokens",
        "decoder_layers",
        "norm",
        "rotary_emb",
        "hc_head"
      ],
      "class_name": "DeepseekV4Model",
      "id": "model",
      "kind": "model",
      "path_pattern": "model"
    },
    {
      "attributes": {
        "embedding_dim": "config.hidden_size",
        "num_embeddings": "config.vocab_size",
        "tied_lm_head": false
      },
      "class_name": "Embedding",
      "id": "embed_tokens",
      "kind": "embedding",
      "path_pattern": "model.embed_tokens"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "norm",
      "kind": "normalization",
      "path_pattern": "model.norm"
    },
    {
      "attributes": {
        "head_dim": "config.head_dim",
        "rope_theta": 10000.0,
        "scheme": "rope"
      },
      "class_name": "DeepseekV4RotaryEmbedding",
      "id": "rotary_emb",
      "kind": "position",
      "path_pattern": "model.rotary_emb"
    },
    {
      "children": [
        "hc_head.input_norm"
      ],
      "class_name": "DeepseekV4HyperHead",
      "id": "hc_head",
      "kind": "module",
      "path_pattern": "model.hc_head"
    },
    {
      "attributes": {
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4UnweightedRMSNorm",
      "id": "hc_head.input_norm",
      "kind": "normalization",
      "path_pattern": "model.hc_head.input_norm"
    }
  ],
  "config": {
    "class_name": "DeepseekV4Config",
    "model_type": "deepseek_v4",
    "module": "transformers.models.deepseek_v4.configuration_deepseek_v4",
    "referenced_fields": {
      "num_hidden_layers": 43
    },
    "salient_fields": {
      "head_dim": 512,
      "hidden_act": "silu",
      "hidden_size": 4096,
      "intermediate_size": 2048,
      "is_encoder_decoder": false,
      "max_position_embeddings": 1048576,
      "num_attention_heads": 64,
      "num_hidden_layers": 43,
      "num_key_value_heads": 1,
      "sliding_window": 128,
      "tie_word_embeddings": false,
      "vocab_size": 129280
    }
  },
  "edges": [
    {
      "kind": "data",
      "source": "embed_tokens",
      "target": "decoder_layers"
    },
    {
      "kind": "data",
      "source": "decoder_layers",
      "target": "norm"
    },
    {
      "kind": "data",
      "source": "decoder_layer.input_layernorm",
      "target": "decoder_layer.post_attention_layernorm"
    },
    {
      "kind": "mask",
      "source": "input:attention_mask",
      "target": "decoder_layer.self_attn"
    },
    {
      "kind": "residual",
      "source": "decoder_layer",
      "target": "decoder_layer.self_attn"
    },
    {
      "kind": "residual",
      "source": "decoder_layer",
      "target": "decoder_layer.ffn_hc"
    },
    {
      "kind": "position",
      "source": "decoder_layer.self_attn.compressor.rotary_emb",
      "target": "decoder_layer.self_attn"
    },
    {
      "kind": "position",
      "source": "rotary_emb",
      "target": "decoder_layer.self_attn"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.self_attn",
      "target": "decoder_layer.self_attn.q_a_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.self_attn",
      "target": "decoder_layer.self_attn.q_b_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.self_attn",
      "target": "decoder_layer.self_attn.kv_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.self_attn",
      "target": "decoder_layer.self_attn.o_a_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.self_attn",
      "target": "decoder_layer.self_attn.o_b_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.mlp.shared_experts",
      "target": "decoder_layer.mlp.shared_experts.gate_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.mlp.shared_experts",
      "target": "decoder_layer.mlp.shared_experts.up_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.mlp.shared_experts.gate_proj",
      "target": "decoder_layer.mlp.shared_experts.down_proj"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "decoder_layer.mlp.shared_experts.up_proj",
      "target": "decoder_layer.mlp.shared_experts.down_proj"
    },
    {
      "kind": "cache_read",
      "source": "state:kv_cache",
      "target": "decoder_layer.self_attn"
    },
    {
      "kind": "cache_write",
      "source": "decoder_layer.self_attn",
      "target": "state:kv_cache"
    },
    {
      "kind": "route",
      "source": "decoder_layer.mlp",
      "target": "decoder_layer.mlp.experts"
    },
    {
      "kind": "route",
      "source": "decoder_layer.mlp",
      "target": "decoder_layer.mlp.shared_experts"
    },
    {
      "kind": "route",
      "source": "decoder_layer.mlp.shared_experts",
      "target": "decoder_layer.mlp.shared_experts.gate_proj"
    },
    {
      "kind": "route",
      "source": "decoder_layer.mlp.shared_experts",
      "target": "decoder_layer.mlp.shared_experts.up_proj"
    },
    {
      "kind": "route",
      "source": "decoder_layer.mlp.shared_experts",
      "target": "decoder_layer.mlp.shared_experts.down_proj"
    }
  ],
  "extends": "mixtral",
  "model_type": "deepseek_v4",
  "patches": [
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": "normalization",
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4UnweightedRMSNorm"
    },
    {
      "component_kind": "position",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "LagunaRotaryEmbedding",
      "parent_model": "laguna",
      "relation": "inherits",
      "target_class": "DeepseekV4RotaryEmbedding"
    },
    {
      "added": {
        "attrs": [
          "layer_type"
        ],
        "methods": [
          "__init__",
          "store_compression_weights",
          "update",
          "update_compressor_states"
        ]
      },
      "component_kind": null,
      "parent_class": "DynamicSlidingWindowLayer",
      "relation": "new",
      "target_class": "DeepseekV4HCACache"
    },
    {
      "added": {
        "attrs": [
          "layer_type"
        ],
        "methods": [
          "__init__",
          "update_overlap_state"
        ]
      },
      "component_kind": null,
      "parent_class": "DeepseekV4HCACache",
      "relation": "new",
      "target_class": "DeepseekV4CSACache"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Linear",
      "relation": "new",
      "target_class": "DeepseekV4GroupedLinear"
    },
    {
      "added": {
        "attrs": [
          "rope_layer_type"
        ],
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4HCACompressor"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4IndexerScorer"
    },
    {
      "added": {
        "attrs": [
          "rope_layer_type"
        ],
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4Indexer"
    },
    {
      "added": {
        "attrs": [
          "rope_layer_type"
        ],
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4CSACompressor"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": "attention",
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4Attention"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4HyperConnection"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": null,
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4HyperHead"
    },
    {
      "component_kind": "feed_forward",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "LlamaMLP",
      "parent_model": "llama",
      "relation": "inherits",
      "target_class": "DeepseekV4MLP"
    },
    {
      "added": {
        "methods": [
          "_apply_gate"
        ]
      },
      "component_kind": "moe",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "MixtralExperts",
      "relation": "inherits",
      "target_class": "DeepseekV4Experts"
    },
    {
      "component_kind": "moe",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "MixtralTopKRouter",
      "relation": "inherits",
      "target_class": "DeepseekV4TopKRouter"
    },
    {
      "component_kind": "moe",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "MixtralTopKRouter",
      "relation": "inherits",
      "target_class": "DeepseekV4HashRouter"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": "moe",
      "parent_class": "Module",
      "relation": "new",
      "target_class": "DeepseekV4SparseMoeBlock"
    },
    {
      "added": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "component_kind": "transformer_block",
      "parent_class": "GradientCheckpointingLayer",
      "relation": "new",
      "target_class": "DeepseekV4DecoderLayer"
    },
    {
      "added": {
        "attrs": [
          "_is_stateful",
          "_keep_in_fp32_modules",
          "_keep_in_fp32_modules_strict",
          "_keys_to_ignore_on_load_unexpected",
          "config_class"
        ]
      },
      "component_kind": "model",
      "overridden": {
        "attrs": [
          "_can_compile_fullgraph",
          "_can_record_outputs",
          "_no_split_modules",
          "_supports_flash_attn",
          "_supports_flex_attn",
          "_supports_sdpa",
          "base_model_prefix"
        ],
        "methods": [
          "_init_weights"
        ]
      },
      "parent_class": "MixtralPreTrainedModel",
      "relation": "inherits",
      "target_class": "DeepseekV4PreTrainedModel"
    },
    {
      "component_kind": "model",
      "overridden": {
        "methods": [
          "__init__",
          "forward"
        ]
      },
      "parent_class": "LlamaModel",
      "parent_model": "llama",
      "relation": "inherits",
      "target_class": "DeepseekV4Model"
    }
  ],
  "provenance": {
    "config_class": "DeepseekV4Config",
    "config_module": "transformers.models.deepseek_v4.configuration_deepseek_v4",
    "model_class": "DeepseekV4Model",
    "model_module": "transformers.models.deepseek_v4.modeling_deepseek_v4"
  },
  "repeats": [
    {
      "body": "decoder_layer",
      "container_path_pattern": "model.layers",
      "count": 43,
      "count_expr": "config.num_hidden_layers",
      "count_source": "config",
      "id": "decoder_layers",
      "index_symbol": "i",
      "item_path_pattern": "model.layers.{i}",
      "kind": "symbolic_repeat",
      "provenance": {
        "class_name": "DeepseekV4DecoderLayer",
        "container_path_pattern": "model.layers",
        "item_path_pattern": "model.layers.{i}",
        "source": "module_tree_repeat_collapse"
      },
      "repeated_class_name": "DeepseekV4DecoderLayer"
    }
  ],
  "schema_version": "architecture-template-v0",
  "templates": [
    {
      "children": [
        "decoder_layer.self_attn",
        "decoder_layer.mlp",
        "decoder_layer.input_layernorm",
        "decoder_layer.post_attention_layernorm",
        "decoder_layer.attn_hc",
        "decoder_layer.ffn_hc"
      ],
      "class_name": "DeepseekV4DecoderLayer",
      "id": "decoder_layer",
      "kind": "transformer_block",
      "path_pattern": "model.layers.{i}"
    },
    {
      "attributes": {
        "head_dim": 512,
        "n_heads": 64,
        "n_kv_heads": 1,
        "pattern": "compressed",
        "rope": true,
        "sliding_window": 128,
        "variant": "MLA"
      },
      "children": [
        "decoder_layer.self_attn.q_a_proj",
        "decoder_layer.self_attn.q_a_norm",
        "decoder_layer.self_attn.q_b_proj",
        "decoder_layer.self_attn.q_b_norm",
        "decoder_layer.self_attn.kv_proj",
        "decoder_layer.self_attn.kv_norm",
        "decoder_layer.self_attn.o_a_proj",
        "decoder_layer.self_attn.o_b_proj",
        "decoder_layer.self_attn.compressor"
      ],
      "class_name": "DeepseekV4Attention",
      "id": "decoder_layer.self_attn",
      "kind": "attention",
      "path_pattern": "model.layers.{i}.self_attn"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": 1024
      },
      "class_name": "Linear",
      "id": "decoder_layer.self_attn.q_a_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.self_attn.q_a_proj"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "decoder_layer.self_attn.q_a_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.self_attn.q_a_norm"
    },
    {
      "attributes": {
        "in_features": 1024,
        "out_features": 32768
      },
      "class_name": "Linear",
      "id": "decoder_layer.self_attn.q_b_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.self_attn.q_b_proj"
    },
    {
      "attributes": {
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4UnweightedRMSNorm",
      "id": "decoder_layer.self_attn.q_b_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.self_attn.q_b_norm"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.head_dim"
      },
      "class_name": "Linear",
      "id": "decoder_layer.self_attn.kv_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.self_attn.kv_proj"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "decoder_layer.self_attn.kv_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.self_attn.kv_norm"
    },
    {
      "class_name": "DeepseekV4GroupedLinear",
      "id": "decoder_layer.self_attn.o_a_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.self_attn.o_a_proj"
    },
    {
      "attributes": {
        "in_features": 8192,
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "decoder_layer.self_attn.o_b_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.self_attn.o_b_proj"
    },
    {
      "children": [
        "decoder_layer.self_attn.compressor.kv_norm",
        "decoder_layer.self_attn.compressor.rotary_emb"
      ],
      "class_name": "DeepseekV4HCACompressor",
      "id": "decoder_layer.self_attn.compressor",
      "kind": "module",
      "path_pattern": "model.layers.{i}.self_attn.compressor"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "decoder_layer.self_attn.compressor.kv_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.self_attn.compressor.kv_norm"
    },
    {
      "attributes": {
        "head_dim": "config.head_dim",
        "rope_theta": 10000.0,
        "scheme": "rope"
      },
      "class_name": "DeepseekV4RotaryEmbedding",
      "id": "decoder_layer.self_attn.compressor.rotary_emb",
      "kind": "position",
      "path_pattern": "model.layers.{i}.self_attn.compressor.rotary_emb"
    },
    {
      "attributes": {
        "experts_per_token": 6,
        "num_experts": 256,
        "num_shared_experts": 1
      },
      "children": [
        "decoder_layer.mlp.gate",
        "decoder_layer.mlp.experts",
        "decoder_layer.mlp.shared_experts"
      ],
      "class_name": "DeepseekV4SparseMoeBlock",
      "id": "decoder_layer.mlp",
      "kind": "moe",
      "path_pattern": "model.layers.{i}.mlp"
    },
    {
      "attributes": {
        "experts_per_token": 6,
        "num_experts": 256,
        "num_shared_experts": 1
      },
      "class_name": "DeepseekV4HashRouter",
      "id": "decoder_layer.mlp.gate",
      "kind": "moe",
      "path_pattern": "model.layers.{i}.mlp.gate"
    },
    {
      "attributes": {
        "experts_per_token": 6,
        "num_experts": 256,
        "num_shared_experts": 1
      },
      "class_name": "DeepseekV4Experts",
      "id": "decoder_layer.mlp.experts",
      "kind": "moe",
      "path_pattern": "model.layers.{i}.mlp.experts"
    },
    {
      "attributes": {
        "experts_per_token": 6,
        "num_experts": 256,
        "num_shared_experts": 1
      },
      "children": [
        "decoder_layer.mlp.shared_experts.gate_proj",
        "decoder_layer.mlp.shared_experts.up_proj",
        "decoder_layer.mlp.shared_experts.down_proj"
      ],
      "class_name": "DeepseekV4MLP",
      "id": "decoder_layer.mlp.shared_experts",
      "kind": "moe",
      "path_pattern": "model.layers.{i}.mlp.shared_experts"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.intermediate_size"
      },
      "class_name": "Linear",
      "id": "decoder_layer.mlp.shared_experts.gate_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.mlp.shared_experts.gate_proj"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.intermediate_size"
      },
      "class_name": "Linear",
      "id": "decoder_layer.mlp.shared_experts.up_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.mlp.shared_experts.up_proj"
    },
    {
      "attributes": {
        "in_features": "config.intermediate_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "decoder_layer.mlp.shared_experts.down_proj",
      "kind": "projection",
      "path_pattern": "model.layers.{i}.mlp.shared_experts.down_proj"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "decoder_layer.input_layernorm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.input_layernorm"
    },
    {
      "attributes": {
        "kernel": "RMSNorm",
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4RMSNorm",
      "id": "decoder_layer.post_attention_layernorm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.post_attention_layernorm"
    },
    {
      "children": [
        "decoder_layer.attn_hc.input_norm"
      ],
      "class_name": "DeepseekV4HyperConnection",
      "id": "decoder_layer.attn_hc",
      "kind": "module",
      "path_pattern": "model.layers.{i}.attn_hc"
    },
    {
      "attributes": {
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4UnweightedRMSNorm",
      "id": "decoder_layer.attn_hc.input_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.attn_hc.input_norm"
    },
    {
      "attributes": {
        "activation": "silu",
        "hidden_size": 4096,
        "intermediate_size": 2048
      },
      "children": [
        "decoder_layer.ffn_hc.input_norm"
      ],
      "class_name": "DeepseekV4HyperConnection",
      "id": "decoder_layer.ffn_hc",
      "kind": "feed_forward",
      "path_pattern": "model.layers.{i}.ffn_hc"
    },
    {
      "attributes": {
        "norm_type": "rms"
      },
      "class_name": "DeepseekV4UnweightedRMSNorm",
      "id": "decoder_layer.ffn_hc.input_norm",
      "kind": "normalization",
      "path_pattern": "model.layers.{i}.ffn_hc.input_norm"
    }
  ]
}