Skip to content

feat: first-class dynamic tensor dimensions for streaming KV-cache decode (squashed) - #892

Closed
michalharakal wants to merge 1 commit into
developfrom
feat/dynamic-shapes-local
Closed

feat: first-class dynamic tensor dimensions for streaming KV-cache decode (squashed)#892
michalharakal wants to merge 1 commit into
developfrom
feat/dynamic-shapes-local

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Note: This is the single-commit variant of the dynamic-shape capability, cut from the sibling checkout used for -transformers composite builds (on top of the argMax fix, with the SoftmaxCpuTest fix folded in). It is functionally identical to #891 (which has the same change as 5 reviewable commits). Only one of the two should be merged — pick whichever base/history you prefer and close the other.

Summary

Makes dynamic tensor dimensions a first-class, correctly-compiling concept end-to-end, so a single compiled vmfb serves 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 that iree-compile rejects 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; volume now throws on a dynamic shape instead of returning a corrupt product.
  • Slice DSL: all() over a dynamic axis is a valid symbolic full-axis.

Shape-only, dynamic-safe tracing (skainet-lang-core, skainet-compile-dag)

  • VoidTensorOps propagates shapes through a ShapeOnlyTensorData that allocates no backing buffer, so a -1 extent threads through a whole decode trace (was NegativeArraySizeException). Concat routes through Dim.concat; reshape passes a dynamic target through.
  • TraceToGraphBuilder.extractFloatArray no longer probes volume of a non-dense/dynamic input.

Dynamic-shape-safe StableHLO emission (skainet-compile-hlo)

  • Scale folded into Q ((q·s)@kᵀ ≡ scores·s) — removes the invalid scores-sized dynamic splat constant.
  • Softmax reduce/broadcast uses stablehlo.dynamic_broadcast_in_dim + a runtime get_dimension_size shape operand (CHLO is rejected by IREE's stablehlo pipeline).
  • Identity reshape, full-extent slice/narrow elided; concat keeps the axis dynamic (? ++ 1 stays ?, never a bogus 0).
  • TypeMapper.DYNAMIC_DIM aliases Dim.DYNAMIC so tracer and emitter agree by construction.

Every dynamic path is gated behind hasDynamic(); static graphs are emitted byte-for-byte unchanged.

Verification

  • Full skainet-lang-core + skainet-compile-hlo + skainet-compile-dag suites pass; new tests cover Dim arithmetic, guarded volume, dynamic concat/reshape/slice, and dynamic SDPA emission.
  • Synthetic dynamic SDPA iree-compiles and runs at key length 3 and 17 from one module.
  • FunctionGemma with_past (dynamic 1x{nKV}x?x256, real weights) self-compiles from the DSL → 151 KB vmfb (previously uncompilable).
  • Moonshine v2 with_past (dynamic self and cross caches 1x8x?x40, weights baked) self-compiles → 86 MB vmfb.
  • Growing-cache greedy decode of the Moonshine v2 dynamic vmfb vs onnxruntime decoder_kv.onnx, cache 1→12: cos = 1.000000 and argmax match at every step.

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>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-892 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal
michalharakal deleted the feat/dynamic-shapes-local branch July 27, 2026 15:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: First-class dynamic tensor dimensions for streaming KV-cache decode

1 participant