Feat/sa3 lora

acestep/engine/sa3_context.py

@@ -85,6 +85,7 @@ def latent_rate_hz(self) -> float:
 
     def make_dit(
         self, *, latent_frames: int, seconds_total: float, backend: str = "eager",
+        prefer_refittable: bool = False,
     ):
         """The per-step velocity callable for one session: with
         ``backend="tensorrt"``, the built TRT engine when one covers
@@ -98,7 +99,10 @@ def make_dit(
         ``backend`` is the resolved acceleration value the session
         creator threads through from the serving layer's accel param
         (compile is already normalized to eager there: SA3 has no
-        torch.compile path)."""
+        torch.compile path). ``prefer_refittable`` is the LoRA-session
+        preference (notes/SA3_LORA_PLAN.md D6b): pick a refit-built
+        engine when one covers the window, and avoid fp8 (whose refit
+        story is unproven) otherwise."""
         from acestep.engine.sa3_trt import SA3TRTDit, find_dit_engine
 
         if backend != "tensorrt":
@@ -108,7 +112,10 @@ def make_dit(
             )
             return self.dit
 
-        engine_path = find_dit_engine(self.model_id, int(latent_frames))
+        engine_path = find_dit_engine(
+            self.model_id, int(latent_frames),
+            want_refittable=bool(prefer_refittable),
+        )
         if engine_path is None:
             logger.info(
                 "sa3_dit_eager model_id={} latent_frames={} reason=no_trt_engine",

acestep/engine/sa3_trt.py

@@ -85,11 +85,24 @@
 _DIT_FP8_DIR_RE = re.compile(
     r"^(?P<prefix>.+_dit)_fp8_l(?P<lo>\d+)_(?P<opt>\d+)_(?P<hi>\d+)$"
 )
+# Refittable fp16mixed DiT engines (BuilderFlag.REFIT; built by
+# ``sa3_build --refit``): ``sa3_m_dit_refit_l{min}_{opt}_{max}``. The
+# LoRA refit path (notes/SA3_LORA_PLAN.md D6b) mutates these in place,
+# so they are EXCLUDED from the shared deserialization cache below —
+# every consumer gets an exclusively-owned instance that dies with its
+# session (which is also the base-weight rollback guarantee: a mutated
+# engine can never be handed to a later session).
+_DIT_REFIT_DIR_RE = re.compile(
+    r"^(?P<prefix>.+_dit)_refit_l(?P<lo>\d+)_(?P<opt>\d+)_(?P<hi>\d+)$"
+)
 _SAME_L_DIR_RE = re.compile(r"^same_l_decode_window_t(?P<lo>\d+)_(?P<opt>\d+)_(?P<hi>\d+)$")
 
 # Deserialized-engine process cache. Engines are immutable post-load and
 # support multiple execution contexts, so sharing one deserialization
 # across sessions is safe; the per-wrapper state is the context+buffers.
+# The IMMUTABILITY ASSUMPTION is the cache's contract: refittable
+# engines (see _DIT_REFIT_DIR_RE) must never enter it — they are
+# deserialized per-session with exclusive ownership instead.
 _ENGINE_CACHE: dict = {}
 _ENGINE_CACHE_LOCK = threading.Lock()
 _SAME_PLUGIN_REGISTERED = False
@@ -99,9 +112,31 @@ def trt_engines_dir() -> Path:
     return paths.models_dir() / "sa3" / "trt_engines"
 
 
-def _deserialize_engine(path: Path):
+def is_refittable_engine_path(path: Path) -> bool:
+    """Whether ``path`` names a refit-built DiT engine (by the naming
+    marker ``sa3_build --refit`` stamps)."""
+    return _DIT_REFIT_DIR_RE.match(Path(path).parent.name) is not None
+
+
+def _deserialize_engine(path: Path, *, exclusive: bool = False):
     import tensorrt as trt
 
+    exclusive = exclusive or is_refittable_engine_path(path)
+    if exclusive:
+        # Refittable (or explicitly exclusive) engines bypass the cache:
+        # in-place refit would otherwise mutate every consumer and
+        # outlive the session that made it. The load cost (~seconds for
+        # the 2.8 GB DiT) is the accepted price of exclusive ownership.
+        logger.info(
+            "sa3_trt_engine_load path={} size_gb={:.1f} ownership=exclusive",
+            path, path.stat().st_size / 1e9,
+        )
+        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
+        engine = runtime.deserialize_cuda_engine(path.read_bytes())
+        if engine is None:
+            raise RuntimeError(f"failed to deserialize TRT engine {path}")
+        return engine
+
     with _ENGINE_CACHE_LOCK:
         engine = _ENGINE_CACHE.get(str(path))
         if engine is None:
@@ -144,19 +179,44 @@ def _register_same_plugin() -> None:
 # ---------------------------------------------------------------------------
 
 
-def find_dit_engine(model_id: str, latent_frames: int) -> Optional[Path]:
+def find_dit_engine(
+    model_id: str, latent_frames: int, *, want_refittable: bool = False,
+) -> Optional[Path]:
     """Smallest-profile built DiT engine covering ``latent_frames`` for
-    ``model_id``'s weights, or None (caller falls back to eager)."""
+    ``model_id``'s weights, or None (caller falls back to eager).
+
+    ``want_refittable`` is the LoRA-session preference (plan D6b):
+
+    * A covering refit-built engine wins when one exists — LoRA strength
+      changes then run as in-place refits instead of the eager-DiT swap.
+    * Otherwise selection prefers **fp16mixed over fp8** (D6c: FP8 refit
+      on Stability's graph is its own unproven investigation; until then
+      a LoRA session must not land on an engine the refit path can't
+      serve at higher fidelity than the fp16mixed one), logged, and the
+      interim eager swap covers actual enables.
+
+    Without it, selection is unchanged: fp8 preferred when covering,
+    else fp16mixed; refit-built engines are ignored (their refit
+    support costs a little optimization freedom, and non-LoRA sessions
+    shouldn't pay it).
+    """
     prefix = DIT_ENGINE_PREFIX.get(model_id)
     base = trt_engines_dir()
     if prefix is None or not base.is_dir():
         return None
-    best = None       # smallest-covering fp16mixed engine
-    best_fp8 = None   # smallest-covering fp8 engine (preferred when present)
+    best = None        # smallest-covering fp16mixed engine
+    best_fp8 = None    # smallest-covering fp8 engine
+    best_refit = None  # smallest-covering refit-built engine
     for sub in base.iterdir():
         f = sub / f"{sub.name}.trt"
         if not f.is_file():
             continue
+        mr = _DIT_REFIT_DIR_RE.match(sub.name)
+        if mr and mr.group("prefix") == prefix:
+            lo, hi = int(mr.group("lo")), int(mr.group("hi"))
+            if lo <= latent_frames <= hi and (best_refit is None or hi < best_refit[0]):
+                best_refit = (hi, f)
+            continue
         mf = _DIT_FP8_DIR_RE.match(sub.name)
         if mf and mf.group("prefix") == prefix:
             lo, hi = int(mf.group("lo")), int(mf.group("hi"))
@@ -168,6 +228,20 @@ def find_dit_engine(model_id: str, latent_frames: int) -> Optional[Path]:
             lo, hi = int(m.group("lo")), int(m.group("hi"))
             if lo <= latent_frames <= hi and (best is None or hi < best[0]):
                 best = (hi, f)
+    if want_refittable:
+        if best_refit is not None:
+            logger.info(
+                "sa3_dit_refit_selected engine={} latent_frames={}",
+                best_refit[1].parent.name, latent_frames,
+            )
+            return best_refit[1]
+        if best_fp8 is not None and best is not None:
+            logger.info(
+                "sa3_dit_fp8_skipped_for_lora engine={} using={} "
+                "reason=fp8_refit_unproven",
+                best_fp8[1].parent.name, best[1].parent.name,
+            )
+        return best[1] if best else (best_fp8[1] if best_fp8 else None)
     # fp8 is ~1.8x faster; prefer it when one covers the window.
     if best_fp8 is not None:
         logger.info(
@@ -192,7 +266,11 @@ def max_dit_engine_latents(model_id: str) -> Optional[int]:
     # install would otherwise report no cap and skip the clamp.
     his = []
     for sub in base.iterdir():
-        m = _DIT_DIR_RE.match(sub.name) or _DIT_FP8_DIR_RE.match(sub.name)
+        m = (
+            _DIT_REFIT_DIR_RE.match(sub.name)
+            or _DIT_DIR_RE.match(sub.name)
+            or _DIT_FP8_DIR_RE.match(sub.name)
+        )
         if (
             m
             and m.group("prefix") == prefix
@@ -259,6 +337,14 @@ class SA3TRTDit:
 
     def __init__(self, engine_path: Path, *, latent_frames: int, seconds_total: float):
         engine = _deserialize_engine(engine_path)
+        # Exposed for the LoRA refit mirror: on a refit-built engine
+        # (exclusively owned, never cached) the mirror attaches an
+        # IRefitter to this same instance. None of the wrapper's own
+        # state cares about refits — buffers and context stay valid
+        # across refit_cuda_engine().
+        self.engine = engine
+        self.engine_path = Path(engine_path)
+        self.refittable = is_refittable_engine_path(engine_path)
         self._ctx = engine.create_execution_context()
         self._stream = torch.cuda.Stream()
         L = int(latent_frames)