Two changes to _extended_self_attn dispatch (per-block, per-step hot path):
- The flash-attn LSE path now auto-enables whenever flash_attn is
importable, instead of requiring ComfyUI's --use-flash-attention launch
flag (that flag only governs comfy's own attention calls; gating our
internal path on it left the fast path dark on stock `comfy launch`).
A runtime kernel failure permanently drops to the SDPA fallback.
- The SDPA fallback no longer passes a float attn_mask (which forced the
~2.6x slower mem-efficient backend). The uniform b_cond bias is folded
into one appended head dimension (q extra = sqrt(D), target-k extra = 0,
cond-k extra = b_cond, explicit scale=1/sqrt(D)), head dim zero-padded
to a multiple of 8, padded value channels sliced off the output — exact,
and flash-backend eligible.
Measured on RTX 5070 Ti at ~1MP shapes (B=2, 16h, 4200+4200 tokens, bf16),
per call: old float-mask 17.3 ms -> fallback 11.2 ms / flash 6.9 ms
(native target-only attention: 3.3 ms). x28 blocks this was ~300 ms/step
of recoverable overhead. Cross-path agreement at the bf16 noise floor;
gate-off invariant unchanged (7.55e-08 in test_smoke).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TSkULyWmisdHfrxHNc7J37
The ~0.9GB per-apply cond KV cache had two OOM paths on tight cards:
- swapping the reference image left the previous MODEL output (and its
cache) alive in ComfyUI's execution cache while the new one built —
_build_cond_kv now evicts every other state's cache first (they rebuild
lazily if that model is sampled again);
- the cache is allocated lazily inside the first DiT forward, where
comfy's memory manager can no longer make room — a PREPARE_SAMPLING
wrapper now enforces a free-VRAM floor (cache*1.25 + 768MB) before
sampling, partially unloading model weights instead of OOMing.
Verified: headless A->B->A swap run logs 'evicted 1 stale cond KV
cache(s)' per swap, no OOM, outputs clean.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- adaln_lora_{self_attn,cross_attn,mlp} deltas merge onto
adaln_modulation_*.2.weight as ModelPatcher LoRA patches (strength =
node strength, matching training's delta_scale=multiplier); the
cond-stream prefill subtracts them so the reference stream keeps the
frozen modulation exactly as training did.
- channel-scaled checkpoints: *.inv_scale now pre-scales the cond-LoRA
input (trained function is up(down(x*inv)) — was silently skipped).
- the loader now refuses checkpoints carrying tensors it doesn't map,
instead of silently dropping them (un-applied adaln = garbage output
with no error).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
