DeepSeek V4.1

class ParallelHead(nn.Module):
    def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.norm_eps = norm_eps
        self.hc_eps = hc_eps
        self.part_vocab_size = vocab_size // world_size
        # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly
        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))

    def forward(self, x: torch.Tensor, full_logits=False):
        """x: [b, s, d]. Generation only needs the last position, so that is the default."""
        if not full_logits:
            x = x[:, -1]
        logits = F.linear(x.float(), self.weight)
        if world_size > 1:
            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
            dist.all_gather(all_logits, logits)
            logits = torch.cat(all_logits, dim=-1)
        return logits

    @torch.inference_mode()
    def forward(
        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
    ):
        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
        has already filled, so the padding it uses internally never reaches here. token_types /
        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
        inside the first (start_pos 0) chunk."""
        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
        engram_mask = None if image_mask is None else ~image_mask
        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
        h = self.embed(input_ids)
        if images is not None:
            assert start_pos == 0, "image spans must be prefilled in a single chunk"
            self.merge_image_embeddings(images, h)
        # Expand to hc_mult copies for Hyper-Connections
        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
        main_hiddens = []
        pre_mix = make_identity_pre_mix(h, self.hc_mult)
        for i, layer in enumerate(self.layers):
            if layer.engram is not None:
                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
            # the MTP head reads the attention input of its target layers, not their output
            if i in self.target_layer_ids:
                main_hiddens.append(h.mean(dim=2))
            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
        h = layer.hc_pre(h, pre_mix)
        logits = self.head(self.norm(h))
        output_ids = sample(logits, self.temperature)
        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
        return output_ids, logits, main_hidden

{
    "vocab_size": 129280,
    "dim": 5120,
    "moe_inter_dim": 2304,
    "n_layers": 40,
    "n_mtp_layers": 3,
    "dspark_block_size": 5,
    "dspark_noise_token_id": 128799,
    "dspark_target_layer_ids": [37, 38, 39],
    "dspark_markov_rank": 256,
    "dspark_n_routed_experts": 128,
    "dspark_n_activated_experts": 3,
    "n_heads": 64,
    "n_routed_experts": 384,
    "n_shared_experts": 1,
    "n_activated_experts": 6,
    "score_func": "sqrtsoftplus",
    "route_scale": 1.5,
    "swiglu_limit": 10.0,
    "q_lora_rank": 1280,
    "head_dim": 512,
    "rope_head_dim": 64,
    "norm_eps": 1e-20,
    "o_groups": 8,
    "o_lora_rank": 1024,
    "window_size": 128,
    "kv_source_layers": [2, 8, 14, 20],
    "index_source_layers": [2, 8, 14, 20, 24, 28, 32, 36],
    "original_seq_len": 65536,
    "rope_theta": 10000,
    "rope_factor": 16,
    "beta_fast": 32,
    "beta_slow": 1,
    "index_n_heads": 32,
    "index_head_dim": 128,
    "index_topk": 512,
    "candidate_source_layer": 20,
    "candidate_topk_blocks": 2048,
    "candidate_block_size": 8,
    "hc_mult": 4,
    "hc_sinkhorn_iters": 20,
    "hc_eps": 1e-06,
    "engram_layer_ids": [1, 14],
    "engram_vocab_size": 16000000,
    "engram_num_embeddings": [384006168, 384016682],
    "engram_max_ngram_size": 4,
    "engram_pad_id": 2,
    "engram_compressed_vocab_size": 99092,
    "dtype": "fp8",
    "expert_dtype": "fp4",
    "compress_rope_theta": 160000,
    "compress_ratios": [0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
    "vision_n_layers": 32,
    "vision_dim": 1024,
    "vision_n_heads": 16,
    "vision_inter_dim": 2816,
    "vision_patch_size": 14,
    "vision_downsample_ratio": 3,
    "vision_max_n_token": 1024,
    "vision_min_pixels": 295936,
    "vision_max_wh_ratio": null,
    "image_token_id": 129264,
    "engram_n_heads": 8,
    "engram_head_dim": 256,
    "vision_rope_theta": 10000
}