From 1975cf3c5cdb958235770787d77d727dc5093533 Mon Sep 17 00:00:00 2001
From: cenzhiyao <2523403608@qq.com>
Date: Thu, 30 Jul 2026 00:53:23 +0800
Subject: [PATCH] [Fix] Key compiled artifacts by the sizes they were
 specialised 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.
---
 magi_compiler/magi_backend/_cache_data_cls.py |  6 +++++
 magi_compiler/magi_backend/magi_backend.py    | 24 +++++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/magi_compiler/magi_backend/_cache_data_cls.py b/magi_compiler/magi_backend/_cache_data_cls.py
index d8c3e75..3066832 100644
--- a/magi_compiler/magi_backend/_cache_data_cls.py
+++ b/magi_compiler/magi_backend/_cache_data_cls.py
@@ -27,3 +27,9 @@ class CacheEntry:
     runtime_shape: int | None
     graph_index: int
     backend_name: str
+    # A digest of the sizes dynamo turned into constants for this trace. runtime_shape
+    # alone does not identify an artifact: it is None for every dynamically compiled
+    # graph, while the dimensions nobody marked dynamic are baked in and differ between
+    # traces. Without this, a graph traced for one input size is silently handed to
+    # another and fails inside the kernel on a stride assertion.
+    static_sizes: str = ""
diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py
index 8ef38bd..2aa02bc 100644
--- a/magi_compiler/magi_backend/magi_backend.py
+++ b/magi_compiler/magi_backend/magi_backend.py
@@ -16,7 +16,7 @@
 import dataclasses
 import pprint
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -75,6 +75,21 @@ def make_compiler(compile_config: CompileConfig) -> CompilerInterface:
         return EagerAdaptor()
 
 
+def _static_sizes_digest(example_inputs: Iterable[Any]) -> str:
+    """The sizes this trace was specialised on.
+
+    A dimension marked dynamic arrives as a SymInt and the artifact generalises over it.
+    Every other dimension is a plain int that is now a constant in the graph, so two
+    traces that differ only in those ints are different programs.
+    """
+    sizes = [
+        (i, tuple(d for d in x.shape if isinstance(d, int)))
+        for i, x in enumerate(example_inputs)
+        if isinstance(x, torch.Tensor)
+    ]
+    return compute_hash([sizes])[:10]
+
+
 class CompilerManager:
     """
     Manage the compilation process, including graph compilation, compile artifacts caching and loading.
@@ -151,7 +166,7 @@ def save_to_file(self):
             return
         # serialize to a literal-friendly dict
         serializable = {
-            (e.runtime_shape, e.graph_index, e.backend_name): (h.key, h.path, h.restart_analysis_count)
+            (e.runtime_shape, e.graph_index, e.backend_name, e.static_sizes): (h.key, h.path, h.restart_analysis_count)
             for e, h in self.cache.items()
         }
         printer = pprint.PrettyPrinter(indent=4)
@@ -209,13 +224,14 @@ def compile(
             compilation_start_time = time.time()
 
         # Step1: Try loading from the cache
-        cache_entry = CacheEntry(runtime_shape, graph_index, self.compiler.name)
+        static_sizes = _static_sizes_digest(example_inputs)
+        cache_entry = CacheEntry(runtime_shape, graph_index, self.compiler.name, static_sizes)
         compiled_graph = self.load(graph, example_inputs, cache_entry)
         if compiled_graph is not None:
             return compiled_graph
 
         # Step2: Compile the graph
-        key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}"
+        key = f"artifact_shape_{runtime_shape}_static_{static_sizes}_subgraph_{graph_index}"
 
         with self.compile_context(runtime_shape, graph_index):
             compiled_graph, cache_handle = self.compiler.compile(