From 4fe5fab396eb915ccb1e208cf4353b32249a08dc Mon Sep 17 00:00:00 2001
From: Simba Zhang <solderzzc@gmail.com>
Date: Mon, 2 Nov 2026 11:38:28 -0800
Subject: [PATCH 1/2] fix: KimiLinear MLP applies silu to gate_proj, not
 up_proj

SwiGLU is down(silu(gate(x)) * up(x)); the code had silu on up_proj. Adds a
hand-computed SwiGLU regression test (KimiMLP made internal for @testable access).

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 Sources/SwiftLM/KimiLinearDFlash.swift | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftLM/KimiLinearDFlash.swift b/Sources/SwiftLM/KimiLinearDFlash.swift
index 100a6496..82eeba98 100644
--- a/Sources/SwiftLM/KimiLinearDFlash.swift
+++ b/Sources/SwiftLM/KimiLinearDFlash.swift
@@ -140,7 +140,7 @@ public struct KimiLinearConfiguration: Codable, Sendable {
 
 // MARK: - KimiMLP
 
-private class KimiMLP: Module, UnaryLayer {
+class KimiMLP: Module, UnaryLayer {
     @ModuleInfo(key: "gate_proj") var gate: Linear
     @ModuleInfo(key: "up_proj") var up: Linear
     @ModuleInfo(key: "down_proj") var down: Linear
@@ -151,7 +151,7 @@ private class KimiMLP: Module, UnaryLayer {
         _down.wrappedValue = Linear(hiddenDimensions, dimensions, bias: false)
     }
 
-    func callAsFunction(_ x: MLXArray) -> MLXArray { down(gate(x) * silu(up(x))) }
+    func callAsFunction(_ x: MLXArray) -> MLXArray { down(silu(gate(x)) * up(x)) }
 }
 
 // MARK: - KimiMultiLinear

From 31830fe43618bf5b4f977d95463fbc5d12966939 Mon Sep 17 00:00:00 2001
From: Simba Zhang <solderzzc@gmail.com>
Date: Mon, 2 Nov 2026 11:38:57 -0800
Subject: [PATCH 2/2] test: add KimiLinear MLP SwiGLU regression test

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 tests/SwiftLMTests/KimiLinearMLPTests.swift | 52 +++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 tests/SwiftLMTests/KimiLinearMLPTests.swift

diff --git a/tests/SwiftLMTests/KimiLinearMLPTests.swift b/tests/SwiftLMTests/KimiLinearMLPTests.swift
new file mode 100644
index 00000000..f90992d0
--- /dev/null
+++ b/tests/SwiftLMTests/KimiLinearMLPTests.swift
@@ -0,0 +1,52 @@
+import XCTest
+import MLX
+import MLXNN
+@testable import SwiftLM
+
+/// KimiLinear's MLP is SwiGLU: down(silu(gate(x)) * up(x)). The activation must sit on
+/// gate_proj; putting it on up_proj silently produces wrong activations (no shape error).
+final class KimiLinearMLPTests: XCTestCase {
+
+    func testMLPIsSwiGLUWithSiluOnGate() {
+        let dim = 4, hidden = 6
+        let mlp = KimiMLP(dimensions: dim, hiddenDimensions: hidden)
+
+        // Fixed, non-symmetric weights so gate/up are distinguishable.
+        let g = (0..<(hidden * dim)).map { Float(sin(Double($0) * 0.7)) }
+        let u = (0..<(hidden * dim)).map { Float(cos(Double($0) * 0.4 + 1)) }
+        let d = (0..<(dim * hidden)).map { Float(sin(Double($0) * 0.3 + 2)) }
+        mlp.update(parameters: ModuleParameters.unflattened([
+            "gate_proj.weight": MLXArray(g, [hidden, dim]),
+            "up_proj.weight": MLXArray(u, [hidden, dim]),
+            "down_proj.weight": MLXArray(d, [dim, hidden]),
+        ]))
+
+        let xs: [Float] = [0.5, -1.0, 2.0, 0.25, -0.75, 1.5, 0.1, -0.3]
+        let x = MLXArray(xs, [2, dim])
+        let got = Device.withDefaultDevice(.cpu) { mlp(x).asArray(Float.self) }
+
+        // Hand-computed SwiGLU on the CPU.
+        func silu(_ v: Float) -> Float { v / (1 + exp(-v)) }
+        var expected = [Float]()
+        for r in 0..<2 {
+            var h = [Float](repeating: 0, count: hidden)
+            for j in 0..<hidden {
+                var gv: Float = 0, uv: Float = 0
+                for k in 0..<dim {
+                    gv += g[j * dim + k] * xs[r * dim + k]
+                    uv += u[j * dim + k] * xs[r * dim + k]
+                }
+                h[j] = silu(gv) * uv
+            }
+            for i in 0..<dim {
+                var o: Float = 0
+                for j in 0..<hidden { o += d[i * hidden + j] * h[j] }
+                expected.append(o)
+            }
+        }
+        XCTAssertEqual(got.count, expected.count)
+        for (a, b) in zip(got, expected) {
+            XCTAssertEqual(a, b, accuracy: 1e-4)
+        }
+    }
+}