Skip to content

Commit 50693ed

Browse files
author
Carina Peng
committed
Add static-shape inference engine
Static-shape LLM inference engine (StaticShapeEngine). States are discovered by name and the KV cache is right-sized per context bucket (each bucket graph is compiled with its own per-ctx strides, so a max-ctx buffer sliced down corrupts KV); ctx is parsed from the function name. Input preparation extends Sukru's shared SyncInputHandler / InputContext (apple#147): concrete handlers (position ids, causal mask, step, RoPE, PLE, sliding) conform to SyncInputHandler; InputContext gains a per-graph descriptors map so handlers size their own buffers. No Static input protocol/context — one input family for dynamic and static engines. State binding uses bind(into:) (apple#156). Testing: LanguageModelsTests pass incl. new unit tests (ctx-bucket parsing, causal mask fill); correct output on gemma4-E2B, qwen2.5-1.5B, qwen3-0.6B static. Depends on apple#156, apple#147.
1 parent 1677713 commit 50693ed

6 files changed

Lines changed: 864 additions & 183 deletions

File tree

swift/Sources/CoreAILanguageModels/Handlers/InputHandler.swift

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ public struct InputContext: Sendable {
1818
public let batchSize: Int
1919
/// Sliding window size (nil for models without sliding attention).
2020
public let slidingWindow: Int?
21+
/// The active graph's input descriptors, keyed by input name. Lets a handler size its
22+
/// own output buffer to the running bucket. Empty for engines that don't vary by bucket.
23+
public let descriptors: [String: NDArrayDescriptor]
2124

2225
/// For dynamic-shape engines (Sequential, Pipelined).
2326
/// alignedStep = processedTokenCount, batchSize = tokens.count.
@@ -30,22 +33,25 @@ public struct InputContext: Sendable {
3033
processedTokenCount: processedTokenCount,
3134
alignedStep: processedTokenCount,
3235
batchSize: tokens.count,
33-
slidingWindow: nil)
36+
slidingWindow: nil,
37+
descriptors: [:])
3438
}
3539

3640
/// For static-shape engines. Batch is fixed-size and aligned.
3741
public static func `static`(
3842
tokens: ArraySlice<Int32>,
3943
alignedStep: Int,
4044
batchSize: Int,
41-
slidingWindow: Int?
45+
slidingWindow: Int?,
46+
descriptors: [String: NDArrayDescriptor] = [:]
4247
) -> InputContext {
4348
InputContext(
4449
tokens: tokens,
4550
processedTokenCount: alignedStep,
4651
alignedStep: alignedStep,
4752
batchSize: batchSize,
48-
slidingWindow: slidingWindow)
53+
slidingWindow: slidingWindow,
54+
descriptors: descriptors)
4955
}
5056
}
5157

0 commit comments

Comments
 (0)