Jenga is (somewhat) slower than the baseline (Kijai integration)

[kabachuha]

So, I'm implementing Jenga into kijai's WanVideoWrapper kijai/ComfyUI-WanVideoWrapper#658.

However, from my experiments it seems that Jenga is 2x+ times (10 seconds) slower than my baseline kijai workflow. Maybe it's because Jenga doesn't support sage attention yet or has too high precision?

Anyway, I'd like to have a second eye on this integration, if possible. The list of the changes can be seen in the pull request's files list.

[JulianJuaner]

Thank you for helping to implement the Jenga ComfyUI Node here! Since I'm not very familiar with ComfyUI, I looked at your code implementation and found it reasonable. However, my personal development machine is currently under maintenance, so it may take some time before I can test this implementation on my machine. I have a few questions that may need confirmation:

1. Your specific testing settings (such as video size, model size) and GPU model (if it's a Hopper architecture, it might be compatible with FA3+SageAttention, which is not implemented in Jenga (FA2-style FP16)).
2. 2. Your specific inference time, settings, and whether you used timestep skip (i.e., Teacache).

[kabachuha]

Hi, thanks for the comment!

I have the following testing settings 512x512x49 video size, 25 steps, the GPU is 4090 (Hopper, I think). I use practically all the optimizations: TeaCache, fp16_fast base precision, SageAttn2, SLG, torch compile and fp8 quantization.

With Jenga I use 0.6xsteps transition, 0.7 initial drop rate, 0.8 final drop rate, and p_remain_rates of 0.3.

For example, my specific settings (Comfy workflows) are at the following Github gists:

Base settings - 53 seconds: https://gist.github.com/kabachuha/e8450470c352681d9867a38815907b0e

Settings with Jenga (w/o turbo): https://gist.github.com/kabachuha/a6e17257b2ef8cde8c498055b2f33b58

---

Important edit ❗

At the first run Jenga took 2.5 minutes. However, notably, now Jenga launched at only 1:03 minutes (63 seconds) at the second run. Maybe, it was related to torch compilation?

[kabachuha]

I added @torch.compiler.disable() to block_sparse_attention and now it took 01:06 (66 seconds) even in the first run

[JulianJuaner]

Thank you for your reply! If the turbo mode is off, it seems this issue might be related to the comparisons with quantized attentions like Sage Attention. We will also look into whether we can better adapt for the relevant optimizations later.

The advantages of Jenga are more evident when using larger models or context lengths (for example, 14B, 720P, 81f), as it can achieve a better speedup ratio through turbo mode (progressive res + sparse attention).

We have also found that when the context length is short, block selection in attention becomes a bottleneck in the computation, especially when adapting to AR generation with Self-Forcing. We are currently working on optimizing this bottleneck.