zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-25 10:22:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

Move KV cache end counter to Python int to avoid per-step host synchronization in AR sampling loops.

Browse Source
This commit is contained in:
Talmaj Marinc
2026-03-25 21:59:22 +01:00
parent 38514045ab
commit cac5120f96
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
comfy/k_diffusion/sampling.py
View File

@@ -1875,14 +1875,14 @@ def sample_ar_video(model, x, sigmas, extra_args=None, callback=None, disable=No
noisy_input = (1.0 - sigma_next) * denoised + sigma_next * fresh_noise
for cache in kv_caches:
cache["end"].fill_(cache["end"].item() - bf * frame_seq_len)
cache["end"] -= bf * frame_seq_len
step_count += 1
output[:, :, fs:fe] = noisy_input
for cache in kv_caches:
cache["end"].fill_(cache["end"].item() - bf * frame_seq_len)
cache["end"] -= bf * frame_seq_len
zero_sigma = sigmas.new_zeros([1])
_ = model(noisy_input, zero_sigma * s_in, **extra_args)