SwiGLU-style fused MLP: down(activation(gate(x)) * up(x)).
(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor, activation: str = "silu")
| 43 | |
| 44 | # Fused MLP (SwiGLU-style) |
| 45 | def fused_mlp_ref(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor, activation: str = "silu") -> torch.Tensor: |
| 46 | """SwiGLU-style fused MLP: down(activation(gate(x)) * up(x)).""" |
| 47 | gate = x @ w_gate.T |
| 48 | up = x @ w_up.T |
| 49 | if activation == "silu": |
| 50 | gate = F.silu(gate) |
| 51 | elif activation == "gelu": |
| 52 | gate = F.gelu(gate) |
| 53 | elif activation == "relu2": |
| 54 | gate = F.relu(gate) ** 2 |
| 55 | return (gate * up) @ w_down.T |
| 56 | |
| 57 | # Cross Entropy Loss |
| 58 | def cross_entropy_ref(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: |
nothing calls this directly
no outgoing calls
no test coverage detected