Kimi K2 · Causal mask · OpenModelAtlas

Causal mask

tensorattention_mask

Position-derived mask supplied to every block.

shape: $M$ [T queries, T keys]

Sources

Contextual evidence from Construct causal mask

The model embeds tokens, numbers positions from the cached length, builds a 4D causal mask, applies all blocks, normalizes and projects to vocabulary logits with an untied head, computing every position’s logits in float32.

transformer · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 1360–1616

Read code excerpt

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                DeepseekV3DecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

… lines 1371–1422 omitted …

        past_key_values_length = 0
        if use_cache:
            use_legacy_cache = not isinstance(past_key_values, Cache)
            if use_legacy_cache:
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
            past_key_values_length = past_key_values.get_usable_length(seq_length)

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0)

… lines 1439–1442 omitted …

        if self._use_flash_attention_2:
            # 2d mask is passed through the layers
            attention_mask = (
                attention_mask
                if (attention_mask is not None and 0 in attention_mask)
                else None
            )
        else:
            # 4d mask is passed through the layers
            attention_mask = _prepare_4d_causal_attention_mask(
                attention_mask,
                (batch_size, seq_length),
                inputs_embeds,
                past_key_values_length,
            )

… lines 1458–1521 omitted …

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

… lines 1523–1613 omitted …

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

Contextual evidence from Transformer blocks

A block applies input RMSNorm, attention and a residual addition, then post-attention RMSNorm, the FFN and a second residual addition; the first first_k_dense_replace blocks use the dense MLP, the rest MoE.

block · moonshotai/Kimi-K2-Instruct / modeling_deepseek.py
Pinned revision fd1984e2b7a3 ↗ · lines 1152–1218

Read code excerpt

        self.mlp = (
            DeepseekV3MoE(config)
            if (
                config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0
            )
            else DeepseekV3MLP(config)
        )
        self.input_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = DeepseekV3RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

… lines 1167–1197 omitted …

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states