State handling for static-shape inference engine - #160
Open
carinapeng wants to merge 1 commit into
Open
Conversation
carinapeng
force-pushed
the
carina/static-engine-runner-design
branch
3 times, most recently
from
August 10, 2026 22:57
50693ed to
35b1494
Compare
carinapeng
requested review from
alejandro-isaza,
blevine1,
kevchengcodes,
stikves and
tjia1818
August 11, 2026 19:04
| let desc = try Self.requireDescriptor(model: model, functionName: name) | ||
| if Self.contextLength(descriptor: desc, config: configuration) == configuration.maxContextLength { | ||
| let ctx = | ||
| Self.parseContextLength(functionName: name) |
Contributor
There was a problem hiding this comment.
is there an advantage to introducing this method to get context length from the function name? Or was the original Self.contextLength(descriptor: desc, config: configuration) not sufficient anymore?
| let ctx = | ||
| Self.parseContextLength(functionName: name) | ||
| ?? Self.contextLength(descriptor: desc, config: configuration) | ||
| guard descriptorsByCtx[ctx] == nil else { continue } |
Contributor
There was a problem hiding this comment.
can we instead make sure extendFunctionNames is a set at this point, so there's no need for this check? Maybe we can just make the input a set instead of a list.
| throw InferenceRuntimeError.invalidState("No extend function to size the state caches from") | ||
| } | ||
| let ctxLadder = descriptorsByCtx.keys.sorted() | ||
| self.stateCtxLadder = ctxLadder |
Contributor
There was a problem hiding this comment.
Why store this? I don't really see it used
Comment on lines
+172
to
+173
| let smallestCtx = ctxLadder.first! | ||
| let largestCtx = ctxLadder.last! |
Contributor
There was a problem hiding this comment.
could also just do descriptorsByCtx.keys.min() and .max()
| case .float16, .bfloat16: run(Float16.self) | ||
| case .float32: run(Float.self) | ||
| case .int8: run(Int8.self) | ||
| default: run(Float16.self) |
Contributor
There was a problem hiding this comment.
is this ok default behavior if a different scalar type is passed?
| var outputs = try await fn.run( | ||
| inputs: inputs, | ||
| states: consume states, | ||
| states: _unsafeEscapeMutableViews(consume states), |
Contributor
There was a problem hiding this comment.
is this unavoidable?
|
|
||
| var tokenArray = NDArray(descriptor: tokenNDDesc) | ||
| let tokenView = tokenArray.mutableView(as: Int32.self) | ||
| var tokenView = tokenArray.mutableView(as: Int32.self) |
Contributor
There was a problem hiding this comment.
is tokenView modified downstream?
| let inputs = Set(desc.inputNames) | ||
| let states = Set(desc.stateNames) | ||
|
|
||
| // Gemma4: PLE + dual-RoPE + sliding-window KV cache. |
Contributor
There was a problem hiding this comment.
Are Gemma4 and Qwen3.5 the only models supported in this workflow?
Comment on lines
+119
to
+144
| // metadata.json sits beside the .aimodel. Try the dir + its parent, then scan for it. | ||
| var candidates = [ | ||
| bundleURL.appendingPathComponent("metadata.json"), | ||
| bundleURL.deletingLastPathComponent().appendingPathComponent("metadata.json"), | ||
| ] | ||
| for dir in [bundleURL, bundleURL.deletingLastPathComponent()] { | ||
| let items = | ||
| (try? FileManager.default.contentsOfDirectory( | ||
| at: dir, includingPropertiesForKeys: nil)) ?? [] | ||
| candidates.append(contentsOf: items.filter { $0.lastPathComponent == "metadata.json" }) | ||
| } | ||
| // A model bundle has TWO metadata.json files: the bundle *manifest* (with the | ||
| // `language` block we need) beside the .aimodel, and the .aimodel's own internal | ||
| // metadata (assetVersion/creationDate/producer, no `language`). When `bundleURL` | ||
| // is the .aimodel path, `bundleURL/metadata.json` resolves to the internal one | ||
| // first — so pick the first candidate whose JSON actually carries a `language` | ||
| // block rather than the first that merely exists. | ||
| var lang: [String: Any]? = nil | ||
| for url in candidates where FileManager.default.fileExists(atPath: url.path) { | ||
| guard let data = try? Data(contentsOf: url), | ||
| let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], | ||
| let l = json["language"] as? [String: Any] | ||
| else { continue } | ||
| lang = l | ||
| break | ||
| } |
Contributor
There was a problem hiding this comment.
can this be simplified? Every LLM has the same ModelBundle structure so we can just look for the correct metadata.json instead of parsing through candidates?
| /// Sliding sub-range `[0, slidingHeadDim)` = GPT-NeoX full rotary (freqs repeated in two halves). | ||
| /// Global sub-range = partial rotary: only the first `floor(partialRotaryFactor·globalHeadDim/2)` | ||
| /// frequency pairs rotate; the rest are NoPE (θ = 0 ⇒ cos = 1, sin = 0). | ||
| final class RoPEProvider: SyncInputHandler { |
Contributor
There was a problem hiding this comment.
could use some better test coverage of the classes defined in this file
- Right-size KV cache per context bucket - Extend shared SyncInputHandler / InputContext (drop Static* fork) - Unit tests + gemma4/qwen static verified
carinapeng
force-pushed
the
carina/static-engine-runner-design
branch
from
August 12, 2026 03:52
35b1494 to
2bdb457
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Static-shape LLM inference engine (StaticShapeEngine) built on the
bind(into:)state binding from #156. States are discovered by name and the KV cache is right-sized per context bucket; ctx is parsed from the function name. Input preparation is pluggable per model family via StaticInputProvider + StaticModelProfile; StaticInputContext composes the shared InputContext and the provider mirrors SyncInputHandler.Testing: LanguageModelsTests (317) pass; correct output on static assets.