feat: first-class dynamic tensor dimensions for streaming KV-cache decode (squashed) - #892
Closed
michalharakal wants to merge 1 commit into
Closed
feat: first-class dynamic tensor dimensions for streaming KV-cache decode (squashed)#892michalharakal wants to merge 1 commit into
michalharakal wants to merge 1 commit into
Conversation
Mirror of SKaiNET#891 (first-class Dim + dynamic-safe tracer/emitter) on top of the argMax fix, so `-PuseLocalSkainet=true` composite builds in SKaiNET-transformers compile the true-dynamic KV-cache decode graphs. Includes the SoftmaxCpuTest fix: VoidTensorOps only skips allocation for DYNAMIC shapes; static shapes keep readable zeros. Local build branch; the canonical change is SKaiNET#891. 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. |
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
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.