feat: first-class dynamic tensor dimensions for streaming KV-cache decode - #891
Merged
Conversation
…e decode Render a -1 tensor extent as MLIR `?` and emit op forms iree-compile accepts under a dynamic dim, so one compiled vmfb serves every autoregressive decode step (growing KV cache) instead of a single fixed cache length. - TypeMapper: DYNAMIC_DIM=-1 marker + List<Int>.hasDynamic() predicate; every dynamic path is gated behind hasDynamic() so static graphs are byte-for-byte unchanged. - AttentionOperationsConverter: fold the attention scale into Q before the QK dot_general ((q·s)@kᵀ ≡ scores·s, exact) — removes the scores-sized splat constant that is invalid under a dynamic key dim. Broadcast reduced max/sum back with dynamic_broadcast_in_dim (runtime shape via get_dimension_size + concatenate) when the scores shape is dynamic. - ActivationOperationsConverter: same dynamic_broadcast_in_dim treatment for the standalone softmax. - ShapeOperationsConverter: elide identity reshapes (input type == result type) — the KV-cache cache-as-output-sink pattern, invalid under a dynamic result. CHLO implicit-broadcast ops are rejected as illegal by IREE's stablehlo input pipeline, so the dynamic path uses stablehlo.dynamic_broadcast_in_dim instead. Verified: full static regression suite unchanged; new dynamic unit tests pass; the emitted dynamic SDPA iree-compiles and one vmfb runs at key lengths 3 and 17. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Complete the dynamic-shape path so a real dynamic (-1) extent threads through
a whole autoregressive decode trace and lowers to a compilable vmfb — no more
sentinel-dimension + post-emit text substitution.
- VoidTensorOps: shape-only tracing via ShapeOnlyTensorData (carries a Shape,
allocates NO backing buffer). A -1 extent previously threw
NegativeArraySizeException when zeros(shape) allocated a negative-size buffer.
- VoidTensorOps.calculateConcatShape + ShapeOperationsConverter concat: keep the
concatenated axis dynamic when any operand is dynamic there, instead of summing
(a growing cache `? ++ 1` was becoming a bogus static `0`).
- ShapeOperationsConverter: elide full-extent identity slice/narrow. Besides being
a no-op copy, a static stablehlo.slice cannot express a full-extent bound on a
dynamic axis (limit would be the -1 extent, e.g. `0:-1:1`).
Verified end-to-end: the real FunctionGemma with_past decode graph (dynamic
1x{nKV}x?x256 cache, real weights) self-compiles from the DSL to a 150KB CPU
vmfb — a graph the previous sentinel-hack path could not compile. Full static
regression (compile-hlo + lang-core) unchanged.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…reshape-infer Replace the overloaded `-1` dynamic-extent sentinel with an explicit Dim vocabulary so "dynamic" stops colliding with reshape's `-1` = infer and the dynamic-aware shape arithmetic lives in one place instead of scattered `extent < 0` guards. - New sk.ainet.lang.tensor.Dim: DYNAMIC = Int.MIN_VALUE (reserved, distinct from -1), plus isDynamic/isStatic/concat/compatible/render + List/IntArray.hasDynamic(). - Shape: hasDynamic()/isDynamic(axis)/dynamicAxes; volume now THROWS on a dynamic shape (no materializable element count) instead of a corrupt product; toString renders `?` and omits the undefined volume. - Slice DSL is dynamic-aware: `all()` over a dynamic axis is a valid symbolic full-axis (getResultSize passes the extent through); partial Range/Step/At require concrete non-negative bounds (can't resolve from-end indices against unknown size). - VoidTensorOps: concat routes through Dim.concat; reshape passes a dynamic target through unchanged (distinct from `-1` infer, which stays volume-based). - compile-hlo emitter shares the one sentinel: TypeMapper.DYNAMIC_DIM aliases Dim.DYNAMIC, formatShape/converters render+detect dynamic via Dim (no local `< 0` checks or duplicate hasDynamic). Verified: full lang-core + compile-hlo suites pass; the synthetic dynamic SDPA still iree-compiles and runs at key lengths 3 and 17 under the new sentinel. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ace finalize TraceToGraphBuilder.extractFloatArray fell through to `if (tensor.volume > 0)` for non-FloatArrayTensorData tensors and then always returned null anyway — dead code that now (correctly) throws on a dynamic shape, whose volume is undefined. A dynamic graph input (e.g. a `?` KV-cache tensor) is not a constant to embed, so return null without probing volume. Verified: the real Moonshine v2 with_past decoder — dynamic self AND cross caches (1x8x?x40, one vmfb serving every decode position and every encoder length) — self-compiles from the SKaiNET DSL to a vmfb (EXIT 0). compile-dag suite green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…reshape - DimTest: reshape passes a dynamic target through (not mistaken for -1 infer). - DynamicShapeHloExportTest: concat of a dynamic cache stays `?` (never `?+1=0`); full-extent narrow on a dynamic axis is elided (no invalid `0:-1:1` slice). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
…location-free The allocation-free VoidTensorOps change made ALL shape-only tensors non-readable, but existing code creates static void tensors and reads their zeros (e.g. a CPU softmax forward), which then threw IllegalStateException. Scope it: a static shape delegates to DenseTensorDataFactory (real readable zeros, prior behavior); only a DYNAMIC shape — which cannot be allocated — gets the allocation-free ShapeOnlyTensorData. The dynamic KV-cache decode path is unchanged. Fixes the SoftmaxCpuTest regression on feat/emitter-dynamic-shapes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
Collaborator
|
You are doing a lot in this |
Contributor
Author
|
Yes @MacOS this is a quite a heavy feature, blocking SKaiNET on JVM |
Collaborator
|
Reviewing it is a nightmare :D. |
Contributor
Author
|
Some features are still at an early stage and require changes across multiple layers. Missing capabilities continue to surface while building two serious, real-world applications with the library. What gives me confidence is that the core foundation and architecture are consistently being validated. At this point, we are mainly iterating, refining, and fine-tuning the details. |
aharakal
approved these changes
Jul 30, 2026
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
3 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.
Summary
Makes dynamic tensor dimensions a first-class, correctly-compiling concept end-to-end, so a single compiled
vmfbserves every autoregressive decode position — unbounded streaming KV-cache decode from the pure SKaiNET DSL.Previously "dynamic decode" existed only as a downstream sentinel-prime + post-emit text hack (
x7919x → x?x) that did not actually compile: the emitter produced static op forms thatiree-compilerejects once a dim is?. This PR fixes that at the source and removes the overloading of-1(which meant "dynamic" in the emitter, "infer" in reshape, and "last index" in the slice DSL).Closes #890.
What changed
First-class
Dim(skainet-lang-core)Dim.DYNAMIC = Int.MIN_VALUE— a reserved sentinel distinct from reshape's-1= infer — plus centralized dynamic arithmetic (concat,compatible,render,isDynamic/isStatic).Shape:hasDynamic(),isDynamic(axis),dynamicAxes;volumenow throws on a dynamic shape instead of returning a corrupt product.all()over a dynamic axis is a valid symbolic full-axis.Shape-only, dynamic-safe tracing (
skainet-lang-core,skainet-compile-dag)VoidTensorOpspropagates shapes through aShapeOnlyTensorDatathat allocates no backing buffer, so a-1extent threads through a whole decode trace (wasNegativeArraySizeException). Concat routes throughDim.concat; reshape passes a dynamic target through.TraceToGraphBuilder.extractFloatArrayno longer probesvolumeof a non-dense/dynamic input.Dynamic-shape-safe StableHLO emission (
skainet-compile-hlo)(q·s)@kᵀ ≡ scores·s) — removes the invalid scores-sized dynamic splat constant.stablehlo.dynamic_broadcast_in_dim+ a runtimeget_dimension_sizeshape operand (CHLO is rejected by IREE's stablehlo pipeline).? ++ 1stays?, never a bogus0).TypeMapper.DYNAMIC_DIMaliasesDim.DYNAMICso tracer and emitter agree by construction.Every dynamic path is gated behind
hasDynamic(); static graphs are emitted byte-for-byte unchanged.Verification
skainet-lang-core+skainet-compile-hlo+skainet-compile-dagsuites pass; new tests cover Dim arithmetic, guardedvolume, dynamic concat/reshape/slice, and dynamic SDPA emission.iree-compiles and runs at key length 3 and 17 from one module.with_past(dynamic1x{nKV}x?x256, real weights) self-compiles from the DSL → 151 KB vmfb (previously uncompilable).with_past(dynamic self and cross caches1x8x?x40, weights baked) self-compiles → 86 MB vmfb.decoder_kv.onnx, cache 1→12: cos = 1.000000 and argmax match at every step.