Skip to content

[diffusion] comfyui: key cached conditioning on full content - #43109

Closed
OnePunchMonk wants to merge 2 commits into
sgl-project:mainfrom
OnePunchMonk:fix/comfyui-cond-key-collision
Closed

OnePunchMonk wants to merge 2 commits into
sgl-project:mainfrom
OnePunchMonk:fix/comfyui-cond-key-collision

Conversation

@OnePunchMonk

@OnePunchMonk OnePunchMonk commented Oct 8, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Fixes #43107.

SGLDiffusionExecutor._cond_key keyed cached conditioning on the shape plus the first and last scalar of prompt_embeds[0]. On a key hit the executor drops the embeddings and the worker restores whatever is cached under that key. For Flux, prompt_embeds[0] is the pooled CLIP-L y, which ComfyUI takes from the first 77-token chunk only, so two prompts that differ only past that point (or CLIPTextEncodeFlux with the same clip_l text) get the same key and the second prompt runs with the first one's T5 context. Qwen-Image-Edit had the same problem with image_latent, which is dropped but was not part of the key.

Modifications

  • _cond_key returns a blake2b digest over prompt_embeds, pooled_embeds, prompt_seq_lens and the adapter's cached extra_req fields. Nested dicts/lists are hashed by content with length framing (no repr() of tensors, which abbreviates large ones).
  • Adapters declare cached_extra_keys. drop_cached_fields and the key both read it, so they can't drift apart. For MiniMax-H3 this puts the h3_* fields and sigmas in the key (previously only the text was keyed).
  • bind_comfyui_session stores packed extras per comfyui_cond_key instead of in one shared slot. Before, on an H3 cache hit the positive cond could get the negative's h3_context / payload back (last write won). The comfyui_cache_fp check is unchanged.
  • Equal content still hits the cache, so later steps keep sending only latents and timestep.
  • Tests in test_comfyui_session.py run the real executor bookkeeping plus bind_comfyui_session: Flux same y with different T5, same first/last scalar, different image_latent, H3 pos/neg extras on a cache hit, H3 same text with a different payload, large tensors inside a payload list, per-key extras on the worker, and a repeat that must still use the cache.

Accuracy Tests

No model forward changes. Repro in a real ComfyUI KSampler run (ComfyUI master 58b176f, real CLIP-L weights, plugin FluxExecutor + SGLDModelPatcher wired like core/generator.py, scheduler round-trip replaced by the real bind_comfyui_session), prompt A for 0 to 45%, prompt B from 50%:

before: T5 context the worker used per DiT call: ['A', 'A', 'A', 'A', 'A', 'A']
after:  T5 context the worker used per DiT call: ['A', 'A', 'A', 'B', 'B', 'B']

Unit tests (CPU):

main source + new tests:  7 failed, 11 passed   test_comfyui_session.py
patch: 96 passed, 11 skipped (CUDA only), 3 xfailed   test_comfyui_{session,adapters,h3,generator,profile}.py, test_conditioning_cache.py
patch: 12 passed   apps/ComfyUI_SGLDiffusion/test/test_h3_request.py

Speed Tests and Profiling

The key now copies the embeddings to CPU and hashes them once per step (the old key already synced via .item()). Hashing measured on CPU: about 2.5 ms for a 512x4096 bf16 T5 context, 1.3 ms for a Qwen-Image context. For H3 the payload's cond latents are hashed too. Small next to a DiT step, but I did not measure it on GPU.

Checklist


CI States

Latest PR Test (Base): ❌ Run #37803977549
Latest PR Test (Extra): ❌ Run #37803977308
Latest PR Test (AMD ROCm 10): ❌ Run #37803977583

_cond_key used the shape plus the first and last scalar of prompt_embeds[0].
For Flux that is the pooled CLIP-L vector, so prompts that only differ in
T5 got the same key and later steps reused the first prompt's T5 context.
Hash every field drop_cached_fields removes instead.

Fixes sgl-project#43107
Adapters now declare cached_extra_keys; drop_cached_fields and the cond
key both read it, so H3's h3_* fields and sigmas are part of the key.
The worker stored packed extras in one shared slot, so on a cache hit
the positive cond could get the negative's h3_context/payload back.
Store them per comfyui_cond_key instead. Hash nested dicts/lists by
content with framing instead of repr().
@OnePunchMonk

Copy link
Copy Markdown
Contributor Author

Superseded by #43492, which consolidates this fix with 3 related executor conditioning-cache fixes per maintainer request to reduce PR overhead. Closing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

diffusion SGLang Diffusion

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] [diffusion] ComfyUI executor cond key collision reuses another prompt's conditioning

2 participants