Commit 50693ed
Carina Peng
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
- InferenceEngines
- Tests/LanguageModelsTests
Lines changed: 9 additions & 3 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
18 | 18 | | |
19 | 19 | | |
20 | 20 | | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
21 | 24 | | |
22 | 25 | | |
23 | 26 | | |
| |||
30 | 33 | | |
31 | 34 | | |
32 | 35 | | |
33 | | - | |
| 36 | + | |
| 37 | + | |
34 | 38 | | |
35 | 39 | | |
36 | 40 | | |
37 | 41 | | |
38 | 42 | | |
39 | 43 | | |
40 | 44 | | |
41 | | - | |
| 45 | + | |
| 46 | + | |
42 | 47 | | |
43 | 48 | | |
44 | 49 | | |
45 | 50 | | |
46 | 51 | | |
47 | 52 | | |
48 | | - | |
| 53 | + | |
| 54 | + | |
49 | 55 | | |
50 | 56 | | |
51 | 57 | | |
| |||
0 commit comments