nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: cache the cast direction too, not just the zero fallback

Browse Source 
Adversarial review caught this: zero_dir_cache only cached the zero-vector
path. The real-steered block's direction.to(out.dtype) ran fresh every call
with no caching -- a genuine allocation every time, since direction is
always float32 (h3_repr_steering.capture()) while H3 runs bf16/fp16, so the
cast is never a no-op. From step 2 of a scene onward that put the steered
block back to allocating one more tensor per call than the other 49 --
exactly the per-call allocation-count asymmetry the uniform-clone commit
was meant to close, just narrowed rather than eliminated.

Renamed zero_dir_cache -> dir_t_cache and cache the real cast-direction
result the same way the zero fallback already was: built once per
(dtype, device) combo, reused every later call for that block.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-09-19 19:42:35 +03:00
co-authored by Claude Sonnet 5
parent 3f459038ab
commit 04f66dbc48
1 changed files with 19 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samplers.py
+19 -10
View File
@@ -6910,7 +6910,10 @@ class FunPackLTXAVSceneChainSampler:
def _make_hook(block, direction):
inner = dit_patches.get(("double_block", block))
mask_cache = {}
zero_dir_cache = {} # (dtype, device) -> zeros, built once per combo seen
dir_t_cache = {} # (dtype, device) -> the cast direction (or zeros), built
# once per combo seen -- real direction and zero-direction blocks must both
# cache, or the real one keeps re-allocating a fresh cast every call. See
# inside _hook.
def _hook(args, extra):
out = extra["original_block"](args)["img"] if inner is None else \
@@ -6958,15 +6961,21 @@ class FunPackLTXAVSceneChainSampler:
# tradeoff, see the 2026-09-19 session.
rows = out[mask]
row_norm = rows.detach().float().norm(dim=-1).mean()
if direction is not None:
dir_t = direction.to(out.dtype).to(out.device)
else:
key = (out.dtype, out.device)
dir_t = zero_dir_cache.get(key)
if dir_t is None:
dir_t = torch.zeros(out.shape[-1], dtype=out.dtype,
device=out.device)
zero_dir_cache[key] = dir_t
# Cached by (dtype, device), same as the zero-direction case below
# -- direction is a FIXED value for this block's whole install, so
# re-running .to(out.dtype) every call (a real allocation: direction
# is always float32 from capture(), H3 runs bf16/fp16) would put
# this block back to allocating once more per call than the other
# 49 from the second step onward, reopening the exact per-call
# asymmetry this fix exists to close.
key = (out.dtype, out.device)
dir_t = dir_t_cache.get(key)
if dir_t is None:
dir_t = (direction.to(out.dtype).to(out.device)
if direction is not None
else torch.zeros(out.shape[-1], dtype=out.dtype,
device=out.device))
dir_t_cache[key] = dir_t
# GPU-only guard (no .item()/bool() host sync) -- a host sync here
# forces the CPU thread to wait mid-block-loop, perturbing the
# async weight-prefetch pipeline's stream-wait timing.