[diffusion] comfyui: key cached conditioning on full content

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/adapter.py

@@ -38,6 +38,9 @@ class ComfyUIModelAdapter:
 
     model_types: tuple[str, ...] = ()
     pipeline_class_name: str = ""
+    # extra_req fields the worker caches per cond. Dropped on a cache hit and
+    # hashed into the cond key, so the two lists cannot drift apart.
+    cached_extra_keys: tuple[str, ...] = ("image_latent",)
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
@@ -71,7 +74,8 @@ def drop_cached_fields(self, packed: PackedForward) -> None:
         packed.prompt_embeds = []
         packed.prompt_seq_lens = None
         packed.pooled_embeds = None
-        packed.extra_req.pop("image_latent", None)
+        for key in self.cached_extra_keys:
+            packed.extra_req.pop(key, None)
 
 
 def get_adapter_class(model_type: str) -> type[ComfyUIModelAdapter]:

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py

@@ -2,6 +2,7 @@
 Base executor class for SGLang Diffusion ComfyUI integration.
 """
 
+import hashlib
 import uuid
 
 import torch
@@ -15,6 +16,27 @@
     )
 
 
+def _hash_value(digest, value) -> None:
+    """Hash ``value`` into ``digest``, framing each item so values cannot merge."""
+    if torch.is_tensor(value):
+        tensor = value.detach().contiguous().cpu()
+        digest.update(f"T{tensor.dtype}{tuple(tensor.shape)}".encode())
+        digest.update(tensor.reshape(-1).view(torch.uint8).numpy().tobytes())
+    elif isinstance(value, dict):
+        digest.update(f"D{len(value)}".encode())
+        for key, item in value.items():
+            _hash_value(digest, key)
+            _hash_value(digest, item)
+    elif isinstance(value, (list, tuple)):
+        digest.update(f"L{len(value)}".encode())
+        for item in value:
+            _hash_value(digest, item)
+    else:
+        text = repr(value).encode()
+        digest.update(f"V{len(text)}:".encode())
+        digest.update(text)
+
+
 class SGLDiffusionExecutor(torch.nn.Module):
     """Shared ComfyUI DiT-forward executor. Per-model logic lives on the adapter."""
 
@@ -36,7 +58,7 @@ def __init__(self, generator, model_path, model, config):
         self.adapter = self.adapter_cls()
         self.session_id = uuid.uuid4().hex
         self._run_id = 0
-        self._sent_conds: set[tuple] = set()
+        self._sent_conds: set[str] = set()
 
     @staticmethod
     def should_suppress_logs(timestep):
@@ -79,24 +101,34 @@ def sampler_sample_wrapper(self, executor, *args, **kwargs):
     def comfyui_session_id(self) -> str:
         return f"{self.session_id}:{self._run_id}"
 
-    def _cond_key(self, packed) -> tuple | None:
+    def _cond_key(self, packed) -> str | None:
         embeds = packed.prompt_embeds
         if not embeds:
             return None
         tensor = embeds[0]
         if not torch.is_tensor(tensor) or tensor.numel() == 0:
             return None
-        flat = tensor.reshape(-1)
-        return (
-            tuple(int(dim) for dim in tensor.shape),
-            float(flat[0].item()),
-            float(flat[-1].item()),
+        # Hash everything drop_cached_fields removes: a hit means the worker
+        # restores all of it, so a partial key would revive another cond.
+        digest = hashlib.blake2b(digest_size=16)
+        _hash_value(
+            digest,
+            (
+                embeds,
+                packed.pooled_embeds,
+                packed.prompt_seq_lens,
+                {
+                    key: packed.extra_req.get(key)
+                    for key in self.adapter.cached_extra_keys
+                },
+            ),
         )
+        return digest.hexdigest()
 
     def _mark_and_maybe_drop(self, packed) -> None:
         key = self._cond_key(packed)
         if key is not None:
-            packed.extra_req["comfyui_cond_key"] = repr(key)
+            packed.extra_req["comfyui_cond_key"] = key
             if key in self._sent_conds:
                 self.adapter.drop_cached_fields(packed)
             else:

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/minimax_h3.py

@@ -83,6 +83,17 @@ def _cache_fingerprint(video, sample_sigmas, serialized_layout) -> dict[str, Any
 class MiniMaxH3Adapter(ComfyUIModelAdapter):
     model_types = ("minimax_h3",)
     pipeline_class_name = "MiniMaxH3Pipeline"
+    # comfyui_cache_fp / comfyui_cond_key stay; later steps still send them.
+    cached_extra_keys = ComfyUIModelAdapter.cached_extra_keys + (
+        "h3_payload",
+        "h3_transformer_options",
+        "h3_sample_sigmas",
+        "h3_context",
+        "h3_layout",
+        "h3_denoise_mask",
+        "h3_audio_denoise_mask",
+        "sigmas",
+    )
 
     def pack(self, x, timestep, context, **kwargs) -> PackedForward:
         video, audio_src = _split_av(x)
@@ -195,18 +206,6 @@ def fill_req(self, req, packed: PackedForward) -> None:
                 extra[key] = value
         req.extra = extra
 
-    def drop_cached_fields(self, packed: PackedForward) -> None:
-        super().drop_cached_fields(packed)
-        packed.extra_req.pop("h3_payload", None)
-        packed.extra_req.pop("h3_transformer_options", None)
-        packed.extra_req.pop("h3_sample_sigmas", None)
-        packed.extra_req.pop("h3_context", None)
-        packed.extra_req.pop("h3_layout", None)
-        packed.extra_req.pop("h3_denoise_mask", None)
-        packed.extra_req.pop("h3_audio_denoise_mask", None)
-        packed.extra_req.pop("sigmas", None)
-        # Keep comfyui_cache_fp / comfyui_cond_key; later steps still send them.
-
 
 class MiniMaxH3Executor(SGLDiffusionExecutor):
     adapter_cls = MiniMaxH3Adapter

python/sglang/multimodal_gen/runtime/pipelines_core/comfyui_mode.py

@@ -124,8 +124,8 @@ def _evict_other_runs(session_id: str) -> None:
 def bind_comfyui_session(req):
     """Restore cached conditioning, then refresh the cache from whatever is set.
 
-    Conditioning is stored per ``comfyui_cond_key`` so CFG pos/neg do not
-    overwrite each other. Packed extras are restored only when
+    Conditioning and packed extras are stored per ``comfyui_cond_key`` so CFG
+    pos/neg do not overwrite each other. Extras are restored only when
     ``comfyui_cache_fp`` matches; a mismatch must not revive a stale layout.
     """
     sid = session_id_from_req(req)
@@ -137,6 +137,7 @@ def bind_comfyui_session(req):
     cond_key = extra.get("comfyui_cond_key") or "_"
     cached = _SESSIONS.get(sid) or {}
     cond_snap = (cached.get("_conds") or {}).get(cond_key) or {}
+    extra_snap = (cached.get("_extra") or {}).get(cond_key) or {}
     current_fp = extra.get("comfyui_cache_fp")
     cached_fp = cached.get("_fp")
     extra_ok = current_fp is None or cached_fp is None or current_fp == cached_fp
@@ -145,7 +146,7 @@ def bind_comfyui_session(req):
         if _is_empty(getattr(req, name, None)):
             setattr(req, name, value)
     if extra_ok:
-        for key, value in (cached.get("_extra") or {}).items():
+        for key, value in extra_snap.items():
             if _is_empty(extra.get(key)):
                 extra[key] = value
     req.extra = extra
@@ -164,14 +165,17 @@ def bind_comfyui_session(req):
     conds = dict(cached.get("_conds") or {})
     if new_cond:
         conds[cond_key] = new_cond
-    snapshot: dict[str, Any] = {"_conds": conds}
+    extras = dict(cached.get("_extra") or {}) if extra_ok else {}
     if extra_snapshot:
-        snapshot["_extra"] = extra_snapshot
+        extras[cond_key] = extra_snapshot
+    snapshot: dict[str, Any] = {"_conds": conds}
+    if extras:
+        snapshot["_extra"] = extras
     if current_fp is not None:
         snapshot["_fp"] = current_fp
     elif extra_ok and cached_fp is not None:
         snapshot["_fp"] = cached_fp
-    if conds or extra_snapshot:
+    if conds or extras:
         _SESSIONS[sid] = snapshot
     return req
 

python/sglang/multimodal_gen/test/unit/test_comfyui_session.py

@@ -1,7 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import uuid
+
 import torch
 
+from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.adapter import (
+    ComfyUIModelAdapter,
+    PackedForward,
+)
+from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.base import (
+    SGLDiffusionExecutor,
+)
+from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.flux import FluxAdapter
+from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.minimax_h3 import (
+    MiniMaxH3Adapter,
+)
+from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.zimage import (
+    ZImageAdapter,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.comfyui_mode import (
     bind_comfyui_session,
     get_run_state,
@@ -188,6 +204,160 @@ def test_cond_keys_keep_positive_and_negative_apart() -> None:
     release_comfyui_session(sid)
 
 
+class _Executor(SGLDiffusionExecutor):
+    """Real executor bookkeeping, without a generator or model."""
+
+    def __init__(self, adapter):
+        torch.nn.Module.__init__(self)
+        self.adapter = adapter
+        self.session_id = uuid.uuid4().hex
+        self._run_id = 0
+        self._sent_conds = set()
+        self.begin_sampler_run()
+        self.sid = self.comfyui_session_id()
+
+    def send(self, packed) -> _Req:
+        """Executor half of _execute_packed, then the worker-side bind."""
+        self._mark_and_maybe_drop(packed)
+        req = _Req()
+        self.adapter.fill_req(req, packed)
+        req.extra = {
+            **(req.extra or {}),
+            "comfyui_session_id": self.sid,
+            "comfyui_cond_key": packed.extra_req["comfyui_cond_key"],
+        }
+        return bind_comfyui_session(req)
+
+
+def test_cond_key_flux_same_pooled_different_t5_not_conflated() -> None:
+    # ComfyUI's Flux pooled `y` comes from CLIP-L's first 77-token chunk only,
+    # so prompts that differ later share `y` but not the T5 context.
+    ex = _Executor(FluxAdapter())
+    x, t = torch.zeros(1, 16, 8, 8), torch.tensor([0.5])
+    y = torch.randn(1, 768)
+    ctx_a, ctx_b = torch.randn(1, 8, 4096), torch.randn(1, 8, 4096)
+    first = ex.send(ex.adapter.pack(x, t, ctx_a, y=y))
+    second = ex.send(ex.adapter.pack(x, t, ctx_b, y=y.clone()))
+    assert torch.equal(first.prompt_embeds[1], ctx_a)
+    assert torch.equal(second.prompt_embeds[1], ctx_b)
+    release_comfyui_session(ex.sid)
+
+
+def test_cond_key_same_first_last_scalar_not_conflated() -> None:
+    ex = _Executor(ZImageAdapter())
+    x, t = torch.zeros(1, 16, 8, 8), torch.tensor([0.5])
+    ctx_a, ctx_b = torch.randn(1, 6, 32), torch.randn(1, 6, 32)
+    ctx_b.view(-1)[0] = ctx_a.view(-1)[0]
+    ctx_b.view(-1)[-1] = ctx_a.view(-1)[-1]
+    ex.send(ex.adapter.pack(x, t, ctx_a))
+    second = ex.send(ex.adapter.pack(x, t, ctx_b))
+    assert torch.equal(second.prompt_embeds[0], ctx_b.squeeze(0))
+    release_comfyui_session(ex.sid)
+
+
+def test_cond_key_covers_image_latent() -> None:
+    def packed(image_latent):
+        return PackedForward(
+            latents=torch.zeros(1, 4, 8),
+            timesteps=torch.tensor([500.0]),
+            prompt_embeds=[torch.ones(3, 8)],
+            prompt_seq_lens=[[3]],
+            height=64,
+            width=64,
+            extra_req={"image_latent": image_latent},
+        )
+
+    ex = _Executor(ComfyUIModelAdapter())
+    ref_a, ref_b = torch.randn(1, 4, 8), torch.randn(1, 4, 8)
+    ex.send(packed(ref_a))
+    second = ex.send(packed(ref_b))
+    assert torch.equal(second.image_latent, ref_b)
+    release_comfyui_session(ex.sid)
+
+
+def test_cond_key_repeat_still_uses_cache() -> None:
+    ex = _Executor(FluxAdapter())
+    x, t = torch.zeros(1, 16, 8, 8), torch.tensor([0.5])
+    y, ctx = torch.randn(1, 768), torch.randn(1, 8, 4096)
+    ex.send(ex.adapter.pack(x, t, ctx, y=y))
+    # ComfyUI hands over fresh tensors each step; equal content must still hit.
+    repeat = ex.adapter.pack(x, t, ctx.clone(), y=y.clone())
+    restored = ex.send(repeat)
+    assert repeat.prompt_embeds == []
+    assert torch.equal(restored.prompt_embeds[1], ctx)
+    release_comfyui_session(ex.sid)
+
+
+def _h3_packed(text, payload):
+    return PackedForward(
+        latents=torch.zeros(1, 4, 2, 2),
+        timesteps=torch.tensor([500.0]),
+        prompt_embeds=[text],
+        prompt_seq_lens=[[int(text.shape[0])]],
+        height=2,
+        width=2,
+        extra_req={
+            "h3_payload": payload,
+            "h3_context": text,
+            "comfyui_cache_fp": {"spatial": (1, 2, 2)},
+        },
+    )
+
+
+def test_h3_cache_hit_restores_own_extras() -> None:
+    ex = _Executor(MiniMaxH3Adapter())
+    pos, neg = torch.ones(3, 4), torch.zeros(3, 4)
+    pos_payload, neg_payload = {"text_token_tags": [1, 2]}, {"text_token_tags": [3]}
+    ex.send(_h3_packed(pos, pos_payload))
+    ex.send(_h3_packed(neg, neg_payload))
+    later_pos = _h3_packed(pos.clone(), dict(pos_payload))
+    restored = ex.send(later_pos)
+    assert "h3_context" not in later_pos.extra_req  # cache hit, extras dropped
+    assert torch.equal(restored.extra["h3_context"], pos)
+    assert restored.extra["h3_payload"] == pos_payload
+    release_comfyui_session(ex.sid)
+
+
+def test_h3_same_text_different_payload_not_conflated() -> None:
+    ex = _Executor(MiniMaxH3Adapter())
+    text = torch.ones(3, 4)
+    ex.send(_h3_packed(text, {"refs": [torch.zeros(2, 2)]}))
+    second = ex.send(_h3_packed(text.clone(), {"refs": [torch.ones(2, 2)]}))
+    assert torch.equal(second.extra["h3_payload"]["refs"][0], torch.ones(2, 2))
+    release_comfyui_session(ex.sid)
+
+
+def test_cond_key_hashes_large_tensors_inside_lists() -> None:
+    ex = _Executor(MiniMaxH3Adapter())
+    text = torch.ones(3, 4)
+    latent_a = torch.zeros(4096)
+    latent_b = latent_a.clone()
+    latent_b[2048] = 1.0  # outside what repr() would print
+    key_a = ex._cond_key(_h3_packed(text, {"cond_video_latents": [latent_a]}))
+    key_b = ex._cond_key(_h3_packed(text, {"cond_video_latents": [latent_b]}))
+    assert key_a != key_b
+
+
+def test_extras_cached_per_cond_key() -> None:
+    sid = "exec4:1"
+    pos_ctx, neg_ctx = torch.ones(2, 4), torch.zeros(2, 4)
+    for key, ctx in (("pos", pos_ctx), ("neg", neg_ctx)):
+        req = _Req()
+        req.extra = {
+            "comfyui_session_id": sid,
+            "comfyui_cond_key": key,
+            "h3_context": ctx,
+        }
+        req.prompt_embeds = [ctx]
+        bind_comfyui_session(req)
+
+    later_pos = _Req()
+    later_pos.extra = {"comfyui_session_id": sid, "comfyui_cond_key": "pos"}
+    bind_comfyui_session(later_pos)
+    assert torch.equal(later_pos.extra["h3_context"], pos_ctx)
+    release_comfyui_session(sid)
+
+
 class _FakePipeline:
     def __init__(self):
         self.modules = {}