zhangshuai/mlx-swift-examples
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adopt MLXFast.scaledDotProductAttention (#23)

Browse Source
This commit is contained in:
David Koski
2024-03-12 14:04:43 -07:00
committed by GitHub
parent a94bf79d7e
commit 0fb74cbfdc
7 changed files with 39 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
.circleci/config.yml
View File

@@ -43,12 +43,17 @@ workflows:
when:
and:
- matches:
pattern: "^(?!pull/)[-\\w]+$"
pattern: "^(?!pull/)[-_./\\w]+$"
value: << pipeline.git.branch >>
- not: << pipeline.parameters.nightly_build >>
- not: << pipeline.parameters.weekly_build >>
jobs:
- mac_build_and_test
- hold:
type: approval
- apple/authenticate:
context: pr-approval
- mac_build_and_test:
requires: [ hold ]
prb:
when:

22
Libraries/LLM/Gemma.swift
View File

@@ -2,6 +2,7 @@
import Foundation
import MLX
import MLXFast
import MLXNN
// port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/gemma.py
@@ -31,7 +32,6 @@ private class RMSNorm: Module, UnaryLayer {
private class Attention: Module {
let args: GemmaConfiguration
let repeats: Int
let scale: Float
@ModuleInfo(key: "q_proj") var wq: Linear
@@ -48,8 +48,6 @@ private class Attention: Module {
let heads = args.attentionHeads
let kvHeads = args.kvHeads
self.repeats = heads / kvHeads
let headDim = args.headDimensions
self.scale = pow(Float(headDim), -0.5)
@@ -76,11 +74,6 @@ private class Attention: Module {
keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
if repeats > 1 {
keys = MLXArray.repeat(keys, count: repeats, axis: 1)
values = MLXArray.repeat(values, count: repeats, axis: 1)
}
if let (keyCache, valueCache) = cache {
queries = rope(queries, offset: keyCache.dim(2))
keys = rope(keys, offset: keyCache.dim(2))
@@ -91,14 +84,11 @@ private class Attention: Module {
keys = rope(keys)
}
var scores = (queries * self.scale).matmul(keys.transposed(0, 1, 3, 2))
if let mask {
scores = scores + mask
}
scores = softMax(scores.asType(.float32), axis: -1).asType(scores.dtype)
let output = matmul(scores, values).transposed(0, 2, 1, 3).reshaped(B, L, -1)
let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask
)
.transposed(0, 2, 1, 3)
.reshaped(B, L, -1)
return (wo(output), (keys, values))
}

22
Libraries/LLM/Llama.swift
View File

@@ -2,6 +2,7 @@
import Foundation
import MLX
import MLXFast
import MLXNN
// port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/llama.py
@@ -9,7 +10,6 @@ import MLXNN
private class Attention: Module {
let args: LlamaConfiguration
let repeats: Int
let scale: Float
@ModuleInfo(key: "q_proj") var wq: Linear
@@ -26,8 +26,6 @@ private class Attention: Module {
let heads = args.attentionHeads
let kvHeads = args.kvHeads
self.repeats = heads / kvHeads
let headDim = args.hiddenSize / heads
self.scale = pow(Float(headDim), -0.5)
@@ -69,11 +67,6 @@ private class Attention: Module {
keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
if repeats > 1 {
keys = MLXArray.repeat(keys, count: repeats, axis: 1)
values = MLXArray.repeat(values, count: repeats, axis: 1)
}
if let (keyCache, valueCache) = cache {
queries = rope(queries, offset: keyCache.dim(2))
keys = rope(keys, offset: keyCache.dim(2))
@@ -84,14 +77,11 @@ private class Attention: Module {
keys = rope(keys)
}
var scores = (queries * self.scale).matmul(keys.transposed(0, 1, 3, 2))
if let mask {
scores = scores + mask
}
scores = softMax(scores.asType(.float32), axis: -1).asType(scores.dtype)
let output = matmul(scores, values).transposed(0, 2, 1, 3).reshaped(B, L, -1)
let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask
)
.transposed(0, 2, 1, 3)
.reshaped(B, L, -1)
return (wo(output), (keys, values))
}

22
Libraries/LLM/Phi.swift
View File

@@ -2,6 +2,7 @@
import Foundation
import MLX
import MLXFast
import MLXNN
// https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/phi.py
@@ -17,7 +18,6 @@ private class PhiAttention: Module {
let args: PhiConfiguration
let heads: Int
let headDim: Int
let repeats: Int
@ModuleInfo(key: "q_proj") var wq: Linear
@ModuleInfo(key: "k_proj") var wk: Linear
@@ -33,7 +33,6 @@ private class PhiAttention: Module {
self.heads = args.attentionHeads
self.headDim = args.hiddenSize / heads
let kvHeads = args.kvHeads
self.repeats = heads / kvHeads
if headDim * heads != hiddenSize {
fatalError("hidden_size must be divisible by num_heads")
@@ -63,11 +62,6 @@ private class PhiAttention: Module {
keys = keys.reshaped(B, L, args.kvHeads, headDim).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, headDim).transposed(0, 2, 1, 3)
if repeats > 1 {
keys = MLXArray.repeat(keys, count: repeats, axis: 1)
values = MLXArray.repeat(values, count: repeats, axis: 1)
}
// Add RoPE to the queries and keys and combine them with the cache
if let (keyCache, valueCache) = cache {
queries = rope(queries, offset: keyCache.dim(2))
@@ -84,15 +78,13 @@ private class PhiAttention: Module {
// Finally perform the attention computation
let scale = sqrt(1 / Float(queries.dim(-1)))
var scores = (queries * scale).matmul(keys.transposed(0, 1, 3, 2))
if let mask {
scores = scores + mask
}
let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask
)
.transposed(0, 2, 1, 3)
.reshaped(B, L, -1)
scores = softMax(scores, axis: -1).asType(values.dtype)
let valuesHat = matmul(scores, values).transposed(0, 2, 1, 3).reshaped(B, L, -1)
return (dense(valuesHat), (keys, values))
return (dense(output), (keys, values))
}
}

22
Libraries/LLM/Qwen2.swift
View File

@@ -7,13 +7,13 @@
import Foundation
import MLX
import MLXFast
import MLXNN
// port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/qwen2.py
private class Attention: Module {
let args: Qwen2Configuration
let repeats: Int
let scale: Float
@ModuleInfo(key: "q_proj") var wq: Linear
@@ -30,8 +30,6 @@ private class Attention: Module {
let heads = args.attentionHeads
let kvHeads = args.kvHeads
self.repeats = heads / kvHeads
let headDim = args.hiddenSize / heads
self.scale = pow(Float(headDim), -0.5)
@@ -73,11 +71,6 @@ private class Attention: Module {
keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
if repeats > 1 {
keys = MLXArray.repeat(keys, count: repeats, axis: 1)
values = MLXArray.repeat(values, count: repeats, axis: 1)
}
if let (keyCache, valueCache) = cache {
queries = rope(queries, offset: keyCache.dim(2))
keys = rope(keys, offset: keyCache.dim(2))
@@ -88,14 +81,11 @@ private class Attention: Module {
keys = rope(keys)
}
var scores = (queries * scale).matmul(keys.transposed(0, 1, 3, 2))
if let mask {
scores = scores + mask
}
scores = softMax(scores.asType(.float32), axis: -1).asType(scores.dtype)
let output = matmul(scores, values).transposed(0, 2, 1, 3).reshaped(B, L, -1)
let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask
)
.transposed(0, 2, 1, 3)
.reshaped(B, L, -1)
return (wo(output), (keys, values))
}

22
Libraries/LLM/Starcoder2.swift
View File

@@ -7,13 +7,13 @@
import Foundation
import MLX
import MLXFast
import MLXNN
// port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/starcoder2.py
private class Attention: Module {
let args: Starcoder2Configuration
let repeats: Int
let scale: Float
@ModuleInfo(key: "q_proj") var wq: Linear
@@ -30,8 +30,6 @@ private class Attention: Module {
let heads = args.attentionHeads
let kvHeads = args.kvHeads
self.repeats = heads / kvHeads
let headDim = args.hiddenSize / heads
self.scale = pow(Float(headDim), -0.5)
@@ -57,11 +55,6 @@ private class Attention: Module {
keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
if repeats > 1 {
keys = MLXArray.repeat(keys, count: repeats, axis: 1)
values = MLXArray.repeat(values, count: repeats, axis: 1)
}
if let (keyCache, valueCache) = cache {
queries = rope(queries, offset: keyCache.dim(2))
keys = rope(keys, offset: keyCache.dim(2))
@@ -72,14 +65,11 @@ private class Attention: Module {
keys = rope(keys)
}
var scores = (queries * scale).matmul(keys.transposed(0, 1, 3, 2))
if let mask {
scores = scores + mask
}
scores = softMax(scores.asType(.float32), axis: -1).asType(scores.dtype)
let output = matmul(scores, values).transposed(0, 2, 1, 3).reshaped(B, L, -1)
let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask
)
.transposed(0, 2, 1, 3)
.reshaped(B, L, -1)
return (wo(output), (keys, values))
}

2
mlx-swift-examples.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
View File

@@ -15,7 +15,7 @@
"location" : "https://github.com/ml-explore/mlx-swift",
"state" : {
"branch" : "main",
"revision" : "eb249b04b1188b72c122223e5c640a41745a61b9"
"revision" : "948000ceaa27c343f4dd5ce40f727f221bf45c6e"
}
},
{