Skip to content

[Fix] Key compiled artifacts by the sizes they were specialised on - #50

Closed
cennn wants to merge 1 commit into
fix/compile-state-per-topologyfrom
fix/static-shape-in-cache-key
Closed

[Fix] Key compiled artifacts by the sizes they were specialised on#50
cennn wants to merge 1 commit into
fix/compile-state-per-topologyfrom
fix/static-shape-in-cache-key

Conversation

@cennn

@cennn cennn commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #49.

Problem

An artifact is keyed by runtime_shape alone, which is None for every dynamically compiled graph. But the dimensions nobody marked dynamic are not dynamic: _mark_static_shapes burns them into the graph as constants. So a graph traced for one input size is silently handed to another, and the mismatch surfaces far from its cause, as a stride assertion inside the generated kernel:

AssertionError: expected size 75==7, stride 8160==8160 at dim=2
This error most often comes from a incorrect fake (aka meta) kernel for a custom op.
  File ".cache/magi_compiler/magi_cache/model_1_decode_rank_0_cp1_dp1_pp1_tp1/5a9aca216c/artifact_shape_None_subgraph_0/...", line 6779, in call

A VAE decode graph traced for a 1s video clip (7 latent frames) was loaded for a 12s one (75 frames): both resolve to artifact_shape_None_subgraph_0 under the same model tag. The graph signature confirms only two dims generalise — submod_0(..., l_z_, s37, s38, ...) — the declared dynamic_arg_dims={"z": [3, 4]}, while dim 2 is a constant.

Nothing reports a miss; the load succeeds and the failure looks like a bad meta kernel.

Fix

CacheEntry and the artifact path now carry a digest of the sizes that were specialised. Traces differing only in those sizes stop colliding: the second one misses, compiles, and stores under its own key. Dimensions marked dynamic arrive as SymInt and are excluded from the digest, so genuinely shape-agnostic artifacts are reused exactly as before, and no extra compilation is introduced — dynamo already retraces on those shape changes, this only stops the disk cache from serving the wrong result across processes.

Verification

1s : 89e2381f22
12s: f8a6bb7f7a
1s again: 89e2381f22

End to end on athena's gaga4 pipeline: the 12s clip that reproduced the assertion above now decodes through, 8 requests at 73s each.

Same class of defect as #42 (parallel topology into the cache directory) and #49 (compile state per topology): a key that omits something the artifact actually depends on.

An artifact was keyed by runtime_shape alone, which is None for every dynamically compiled graph. The dimensions nobody marked dynamic are not dynamic at all -- _mark_static_shapes burns them into the graph as constants -- so one artifact was handed to inputs it was never traced for, and the mismatch surfaced far from its cause, as a stride assertion inside the generated kernel:

    AssertionError: expected size 75==7, stride 8160==8160 at dim=2

Here a graph traced for a 1s video clip (7 latent frames) was loaded for a 12s one (75), because both resolved to artifact_shape_None_subgraph_0 under the same model tag.

The entry and the artifact path now carry a digest of the sizes that were specialised. Traces that differ only in those sizes no longer collide: the second one misses, compiles, and stores under its own key. Sizes that were marked dynamic arrive as SymInt and are left out of the digest, so genuinely shape-agnostic artifacts keep being reused as before.
@cennn

cennn commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Closing: this targets the wrong layer.

compile_sizes defaults to an empty list and nothing sets it (athena has no occurrence; the only reference in this repo is a test passing []), so concrete_size_entries is always empty and PiecewiseBackend.call always falls through to compiled_graph_for_general_shape. There is only ever one graph in memory, so a cache key that distinguishes shapes only helps on the narrow reload-after-restart path.

Measured: with this change already applied, a VAE window traced at 8 latent frames was still handed a 9-frame input in the same process and died on the same kind of assertion (expected size 9==8 at dim=2). The key did not prevent it.

The real constraint is that a magi_compile'd callable admits exactly one combination of specialised sizes, because _mark_static_shapes burns every undeclared dimension into the graph while the backend treats the first trace as general. Missing dynamic_arg_dims should fail loudly and name the dimension, not be papered over by a key. The crash that motivated this was fixed on the athena side by declaring the time dimension dynamic (world-sim-dev/athena PR #955).

@cennn cennn closed this Jul 29, 2026
@cennn
cennn deleted the fix/static-shape-in-cache-key branch July 29, 2026 18:21
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.

1 participant