diff --git a/Sources/SwiftLM/KimiLinearDFlash.swift b/Sources/SwiftLM/KimiLinearDFlash.swift index 100a649..82eeba9 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 diff --git a/tests/SwiftLMTests/KimiLinearMLPTests.swift b/tests/SwiftLMTests/KimiLinearMLPTests.swift new file mode 100644 index 0000000..f90992d --- /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..