Kimi K3 · Activated gate · OpenModelAtlas

93 layers › 68 KDA layers › Latent mixture of experts › Shared experts

Activated gate

tensorsitu_a · 79

The gate features after SiTU.

shape: $a$ [B batch, T positions, 6144 shared d_ff]

Sources

Contextual evidence from SiTU gate

The dense feed-forward network: down(SiTU(gate x, up x)).

mlp · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 273–301

Read code excerpt

class KimiMLP(nn.Module):
    def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
        self.gate_proj = nn.Linear(
            self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(
            self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(
            self.intermediate_size, self.hidden_size, bias=False)
        if config.hidden_act == "situ":
            beta, linear_beta = _get_situ_activation_params(config)
            self.act_fn = SituAndMul(
                beta=beta,
                linear_beta=linear_beta,
            )
        else:
            self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        if self.config.hidden_act == "situ":
            gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1)
            down_proj = self.down_proj(self.act_fn(gate_up))
        else:
            down_proj = self.down_proj(self.act_fn(
                self.gate_proj(x)) * self.up_proj(x))
        return down_proj

Contextual evidence from SiTU gate

SiTU: beta·tanh(gate/beta)·sigmoid(gate), multiplied by the up branch passed through linear_beta·tanh(up/linear_beta), computed in float32.

situ · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 64–91

Read code excerpt

class SituAndMul(nn.Module):
    """
    SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up
    When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta).
    """

    def __init__(self, beta: float = 1.0, linear_beta: float | None = None):
        super().__init__()
        self.beta = beta
        self.linear_beta = linear_beta

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        d = x.shape[-1] // 2
        gate = x[..., :d].to(torch.float32)
        up = x[..., d:].to(torch.float32)
        situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
        if self.linear_beta is not None:
            up = self.linear_beta * torch.tanh(up / self.linear_beta)
        return (situ_a * up).to(x.dtype)


ACT2FN["situ"] = SituAndMul


def _get_situ_activation_params(config: KimiLinearConfig):
    beta = getattr(config, "activation_situ_beta", None)
    linear_beta = getattr(config, "activation_situ_linear_beta", None)
    return beta or 1.0, linear_beta