Kimi K3 · Softmax over keys

93 layers › 24 MLA layers › Gated MLA › Causal attention

Softmax over keys

no parametersF.softmax(scores, dim=-1, dtype=torch.float32) · 328

Turn each query’s scores over its own and earlier positions into weights that sum to one.

{\htmlData{light=weights softmax}{A_{t,s}}}=\frac{\exp {\htmlData{light=scale-mask}{S_{t,s}}}}{\sum_{u\le t}\exp S_{t,u}}\;(s\le t)

{\htmlData{light=weights softmax}{A_{t,s}}}=0\;(s>t)

in: $S$ [B batch, 96 heads, 1 query, 1001 keys]
out: $A$ [B batch, 96 heads, 1 query, 1001 keys]

The eager function computes it in float32 and casts back (modeling_kimi_linear.py 328); the flash-attention kernel the model runs computes it internally.

Codemodeling_kimi_linear.py 311–332, 450–463

modeling_kimi_linear.py 311–332

def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)

    scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling
    if attention_mask is not None:
        scores = scores + attention_mask[:, :, :, : key.shape[-2]]

    probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
    probs = F.dropout(probs, p=dropout, training=module.training)
    out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous()

    return out, probs

modeling_kimi_linear.py 450–463

        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, _ = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

Sources1

Direct evidence

The eager function scales the dot product of queries and keys, adds a mask when one is given, runs softmax in float32 and mixes the values; the model itself calls the flash-attention function with is_causal.

mla-attend · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 311–463

Read code excerpt

def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)

    scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling
    if attention_mask is not None:
        scores = scores + attention_mask[:, :, :, : key.shape[-2]]

    probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
    probs = F.dropout(probs, p=dropout, training=module.training)
    out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous()

    return out, probs

… lines 333–449 omitted …

        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, _ = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )