Skip to content

Commit 4828fe0

Browse files
authored
Merge branch 'main' into feature/gpu-constrained-sampling
2 parents 6e7f799 + aa3bbf6 commit 4828fe0

16 files changed

Lines changed: 654 additions & 175 deletions

File tree

Package.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,7 @@ let package = Package(
131131
name: "llm-runner",
132132
dependencies: [
133133
"CoreAILanguageModels",
134+
"CoreAIShared",
134135
.product(name: "ArgumentParser", package: "swift-argument-parser"),
135136
],
136137
path: "swift/Sources/Tools/llm-runner",
@@ -154,6 +155,7 @@ let package = Package(
154155
name: "object-detector",
155156
dependencies: [
156157
"CoreAIObjectDetector",
158+
"CoreAIShared",
157159
.product(name: "ArgumentParser", package: "swift-argument-parser"),
158160
],
159161
path: "swift/Sources/Tools/object-detector",
@@ -165,6 +167,7 @@ let package = Package(
165167
name: "diffusion-runner",
166168
dependencies: [
167169
"CoreAIDiffusionPipeline",
170+
"CoreAIShared",
168171
.product(name: "ArgumentParser", package: "swift-argument-parser"),
169172
],
170173
path: "swift/Sources/Tools/diffusion-runner",
@@ -176,6 +179,7 @@ let package = Package(
176179
name: "speech-runner",
177180
dependencies: [
178181
"CoreAISpeech",
182+
"CoreAIShared",
179183
.product(name: "ArgumentParser", package: "swift-argument-parser"),
180184
],
181185
path: "swift/Sources/Tools/speech-runner",
@@ -213,6 +217,7 @@ let package = Package(
213217
name: "LanguageModelsTests",
214218
dependencies: [
215219
"CoreAILanguageModels",
220+
"CoreAIShared",
216221
"TestUtilities",
217222
.product(name: "Transformers", package: "swift-transformers"),
218223
],

python/src/coreai_models/export/mlir_ops.py

Lines changed: 263 additions & 74 deletions
Large diffs are not rendered by default.

python/src/coreai_models/primitives/_ops.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,69 @@ def mutable_slice_update_meta( # type: ignore
4949
):
5050
"""Fake implementation for tracing/meta operations."""
5151
return torch.empty(x.shape, dtype=x.dtype)
52+
53+
54+
@torch.library.custom_op("coreai::mutable_cache_update_and_fetch", mutates_args=["x"])
55+
def mutable_cache_update_and_fetch(
56+
x: Tensor,
57+
update: Tensor,
58+
begin: Tensor,
59+
end: Tensor,
60+
layer_idx: int,
61+
seq_dim: int,
62+
seq_len: int | None,
63+
) -> Tensor:
64+
"""
65+
Fused KV-cache update-and-fetch operation.
66+
67+
Writes ``update`` into the slice ``x[begin:end]`` (mutating ``x``), then
68+
fetches one layer: ``x.narrow(0, layer_idx, 1)`` optionally followed by
69+
``narrow(seq_dim, 0, seq_len)``, then ``squeeze(0)``. ``seq_len=None``
70+
skips the seq-prefix narrow (iOS path returns the full cache row).
71+
During export ``remove_functionalization`` rewrites this into
72+
``unsqueeze`` + ``immutable_slice_update`` + slice (+ optional slice) +
73+
``squeeze`` for MLIR lowering.
74+
75+
Layouts:
76+
- macOS: x is (L, 1, n_kv_heads, max_seq, head_dim), seq_dim=-2,
77+
seq_len=offset+query_len. Output: (1, n_kv_heads, seq_len, head_dim).
78+
- iOS: x is (L, 1, hidden, 1, max_seq), seq_dim=-1, seq_len=None.
79+
Output: (1, hidden, 1, max_seq).
80+
81+
Args:
82+
x: 5D cache tensor to update.
83+
update: 4D values to write into the slice (op unsqueezes internally).
84+
begin: 5-elem tensor of per-dim start indices.
85+
end: 5-elem tensor of per-dim end indices.
86+
layer_idx: Layer to fetch after the update.
87+
seq_dim: Dim along which to truncate the fetched layer (negative idx ok).
88+
seq_len: Populated seq length to fetch, or None to skip the truncation.
89+
"""
90+
update = update.unsqueeze(0)
91+
begin = torch.split(begin, 1, dim=0) # type: ignore
92+
end = torch.split(end, 1, dim=0) # type: ignore
93+
slices = tuple(slice(b.item(), e.item()) for b, e in zip(begin, end, strict=False))
94+
x[slices] = update
95+
fetched = x.narrow(0, layer_idx, 1)
96+
if seq_len is not None:
97+
fetched = fetched.narrow(seq_dim, 0, seq_len)
98+
# Clone because the returned slice aliases the mutated cache.
99+
return fetched.squeeze(0).clone()
100+
101+
102+
@mutable_cache_update_and_fetch.register_fake
103+
def mutable_cache_update_and_fetch_meta( # type: ignore
104+
x: Tensor,
105+
update: Tensor,
106+
begin: Tensor,
107+
end: Tensor,
108+
layer_idx: int,
109+
seq_dim: int,
110+
seq_len: int | None,
111+
):
112+
"""Fake implementation for tracing/meta operations."""
113+
out_shape = list(x.shape)
114+
if seq_len is not None:
115+
out_shape[seq_dim] = seq_len
116+
out_shape.pop(0) # squeeze layer dim
117+
return torch.empty(out_shape, dtype=x.dtype)

python/src/coreai_models/primitives/ios/cache.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from torch import nn
88
from typing_extensions import Self
99

10-
from coreai_models.primitives._ops import mutable_slice_update
10+
from coreai_models.primitives._ops import mutable_cache_update_and_fetch
1111

1212

1313
class KVCacheHandler:
@@ -155,23 +155,29 @@ def update_and_fetch(
155155

156156
begin, end = self.gen_slice_args(layer_idx, offset, num_token_updates)
157157

158-
# update k - note that iOS updates on dimension 4 (the last dimension)
159-
mutable_slice_update(
158+
# update k and fetch the full layer row in a single fused op.
159+
k_out = mutable_cache_update_and_fetch(
160160
x=self._k_cache,
161-
update=k.unsqueeze(0),
161+
update=k,
162162
begin=begin,
163163
end=end,
164+
layer_idx=layer_idx,
165+
seq_dim=-1,
166+
seq_len=None,
164167
)
165168

166-
# update v - note that iOS updates on dimension 4 (the last dimension)
167-
mutable_slice_update(
169+
# update v and fetch the full layer row in a single fused op
170+
v_out = mutable_cache_update_and_fetch(
168171
x=self._v_cache,
169-
update=v.unsqueeze(0),
172+
update=v,
170173
begin=begin,
171174
end=end,
175+
layer_idx=layer_idx,
176+
seq_dim=-1,
177+
seq_len=None,
172178
)
173179

174-
return self._k_cache[layer_idx], self._v_cache[layer_idx]
180+
return k_out, v_out
175181

176182
@property
177183
def k_cache(self) -> torch.Tensor:

python/src/coreai_models/primitives/macos/cache.py

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import torch
77
from typing_extensions import Self
88

9-
from coreai_models.primitives._ops import mutable_slice_update
9+
from coreai_models.primitives._ops import mutable_cache_update_and_fetch, mutable_slice_update
1010

1111

1212
class KVCache:
@@ -117,10 +117,10 @@ def update_and_fetch(
117117
layer_index = torch.tensor((layer_idx,), dtype=torch.int32, device=device)
118118
layer_index_end = torch.tensor((layer_idx + 1,), dtype=torch.int32, device=device)
119119

120-
# update k
121-
mutable_slice_update(
120+
# update k and fetch its populated prefix in a single fused op
121+
k_out = mutable_cache_update_and_fetch(
122122
x=self._k_cache,
123-
update=k.unsqueeze(0),
123+
update=k,
124124
begin=torch.concatenate(
125125
[
126126
layer_index,
@@ -135,16 +135,19 @@ def update_and_fetch(
135135
layer_index_end,
136136
torch.tensor((self._k_cache.size(1),), dtype=torch.int32, device=device),
137137
torch.tensor((self._k_cache.size(2),), dtype=torch.int32, device=device),
138-
torch.tensor((offset + k.size(2),), dtype=torch.int32, device=device),
138+
torch.tensor((offset + k.size(-2),), dtype=torch.int32, device=device),
139139
torch.tensor((self._k_cache.size(4),), dtype=torch.int32, device=device),
140140
]
141141
),
142+
layer_idx=layer_idx,
143+
seq_dim=-2,
144+
seq_len=seq_len,
142145
)
143146

144-
# update v
145-
mutable_slice_update(
147+
# update v and fetch its populated prefix in a single fused op
148+
v_out = mutable_cache_update_and_fetch(
146149
x=self._v_cache,
147-
update=v.unsqueeze(0),
150+
update=v,
148151
begin=torch.cat(
149152
[
150153
layer_index,
@@ -159,17 +162,15 @@ def update_and_fetch(
159162
layer_index_end,
160163
torch.tensor((int(self._v_cache.size(1)),), dtype=torch.int32, device=device),
161164
torch.tensor((int(self._v_cache.size(2)),), dtype=torch.int32, device=device),
162-
torch.tensor((offset + v.size(2),), dtype=torch.int32, device=device),
165+
torch.tensor((offset + v.size(-2),), dtype=torch.int32, device=device),
163166
torch.tensor((int(self._v_cache.size(4)),), dtype=torch.int32, device=device),
164167
]
165168
),
169+
layer_idx=layer_idx,
170+
seq_dim=-2,
171+
seq_len=seq_len,
166172
)
167173

168-
# return the slice k, v
169-
k = self._k_cache.narrow(0, layer_idx, 1).narrow(-2, 0, seq_len)
170-
v = self._v_cache.narrow(0, layer_idx, 1).narrow(-2, 0, seq_len)
171-
k_out = k.squeeze(0)
172-
v_out = v.squeeze(0)
173174
if cross_device:
174175
return k_out.to(compute_device), v_out.to(compute_device)
175176
return k_out, v_out

swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
// Use of this source code is governed by a BSD-3-clause license that can
44
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
55

6+
import CoreAIShared
67
import Foundation
78

89
/// A protocol abstracting time measurement for testability.

swift/Sources/CoreAILanguageModels/Profiling/Timing.swift

Lines changed: 0 additions & 34 deletions
This file was deleted.
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import Foundation
7+
8+
// MARK: - Duration Time Conversions
9+
10+
/// Convenience conversions from `Duration` to floating-point time units, shared
11+
/// across the package's command-line tools and profiling code for elapsed-time reporting.
12+
///
13+
/// Example usage:
14+
/// ```swift
15+
/// let start = ContinuousClock.now
16+
/// // ... do work ...
17+
/// let elapsed = (ContinuousClock.now - start).inSeconds
18+
/// print("Elapsed: \(elapsed)s")
19+
/// ```
20+
extension Duration {
21+
/// Duration in seconds as a `Double`.
22+
package var inSeconds: Double {
23+
let (secs, attoseconds) = components
24+
return Double(secs) + Double(attoseconds) / 1e18
25+
}
26+
27+
/// Duration in milliseconds as a `Double`.
28+
package var inMilliseconds: Double {
29+
inSeconds * 1000.0
30+
}
31+
}

swift/Sources/CoreAIShared/Runtime/ModelStructure.swift

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,86 @@ public struct PreparedModel: Sendable {
131131
return url
132132
}
133133

134+
// MARK: - Cache Inspection
135+
136+
/// File extensions that identify a Core AI model asset (source or compiled).
137+
private static let assetExtensions: Set<String> = ["aimodel", "aimodelc"]
138+
139+
/// Enumerates the Core AI model asset(s) reachable from `url`.
140+
///
141+
/// A model directory (e.g. an LLM bundle) contains one or more asset components alongside
142+
/// other files (tokenizer, metadata); we can't assume specific component filenames, so scan
143+
/// the directory for every `.aimodel`/`.aimodelc` entry. If `url` is itself an asset, it is
144+
/// returned as the sole component. This filename-agnostic approach stays correct as new model
145+
/// families add differently-named components.
146+
///
147+
/// - Parameter url: Either a bundle directory containing asset components, or a single asset.
148+
/// - Returns: The asset URLs to operate on, sorted for stable output. Empty only if `url` is a
149+
/// directory with no asset components.
150+
public static func modelAssetURLs(at url: URL) throws -> [URL] {
151+
// A path ending in a known asset extension IS the asset (asset bundles are themselves
152+
// directories, so this must be checked before treating `url` as a container to scan).
153+
if assetExtensions.contains(url.pathExtension) {
154+
return [url]
155+
}
156+
let entries = try FileManager.default.contentsOfDirectory(
157+
at: url,
158+
includingPropertiesForKeys: nil
159+
)
160+
return
161+
entries
162+
.filter { assetExtensions.contains($0.pathExtension) }
163+
.sorted { $0.path < $1.path }
164+
}
165+
166+
/// Clears the Core AI specialization cache for every model asset reachable from `url`,
167+
/// forcing re-specialization on the next load.
168+
///
169+
/// Discovers components via ``modelAssetURLs(at:)`` — pass either a bundle directory or a
170+
/// single asset. Used by CLI tools implementing `--clear-coreai-cache`.
171+
///
172+
/// - Parameter url: A bundle directory containing asset components, or a single asset.
173+
/// - Returns: The asset URLs whose cache entries were cleared.
174+
@discardableResult
175+
public static func clearCache(at url: URL) throws -> [URL] {
176+
let assetURLs = try modelAssetURLs(at: url)
177+
for assetURL in assetURLs {
178+
let coreaiURL = resolveCoreAIModelURL(from: assetURL)
179+
try AIModelCache.default.deleteEntries(for: coreaiURL)
180+
}
181+
return assetURLs
182+
}
183+
184+
/// Reports whether the default Core AI cache already holds a specialized asset for `url`
185+
/// under the given `options`.
186+
///
187+
/// This only inspects the cache via `AIModelCache.model(for:options:)`; it never triggers
188+
/// specialization. Returns `false` if no entry exists, or if an entry exists but fails to load.
189+
///
190+
/// - Important: `options` must match the options the loader will use for `url`, otherwise a
191+
/// real cached specialization won't be found. Callers that load via `prepare(at:)` should use
192+
/// the ``isCached(at:)`` overload; callers that load via `AIModel(contentsOf:)` or a custom
193+
/// `SpecializationOptions` must pass the same value here.
194+
public static func isCached(at url: URL, options: SpecializationOptions) -> Bool {
195+
let coreaiURL = resolveCoreAIModelURL(from: url)
196+
do {
197+
return try AIModelCache.default.model(for: coreaiURL, options: options) != nil
198+
} catch {
199+
return false
200+
}
201+
}
202+
203+
/// Reports whether the default Core AI cache already holds a specialized asset for `url`,
204+
/// using the same structure-derived `SpecializationOptions` that ``prepare(at:)`` uses.
205+
///
206+
/// Use this only for models loaded through ``prepare(at:)``. For other loaders, use
207+
/// ``isCached(at:options:)`` with the matching options.
208+
public static func isCached(at url: URL) -> Bool {
209+
let coreaiURL = resolveCoreAIModelURL(from: url)
210+
let options = probeStructure(at: coreaiURL).specializationOptions
211+
return isCached(at: coreaiURL, options: options)
212+
}
213+
134214
// MARK: - Asset Preparation
135215

136216
/// Prepares a Core AI model asset by loading via `AIModel` and detecting its structure.

0 commit comments

Comments
 (0)