Purpose of the block-causal mask when the prefill is recomputed on every AR step?
Thanks for releasing the model and inference code. I have a question about the block-causal attention mask in model_for_inference.py.
The README shows the 15-frame window using a block-causal mask, so each frame attends only to itself and earlier frames. As far as I can tell, though, the context KV cache is only reused across the Euler steps of a single prediction, not across autoregressive steps:
- example.py calls world.generate(...) once per frame with a freshly packed window (packed_ids(t, future_ids)).
- generate() calls setup_caches(...), which reuses only the buffers, then _prefill(...). That writes all 14 context frames again at cache_pos = arange(0, prefix_tokens).
- forward_n_steps(...) then reuses those K/Vs only for the target frame's denoising steps.
I don't think cross-step reuse is possible with the current conditioning, because it's window-relative:
- History frames get fidxs = 0..8 by slot, so a frame's index changes as it slides.
- Future anchors get PAST + remaining …, which counts down every step, and every later frame attends to them.
- Token positions are the packed index (input_pos = arange(prefix_tokens)).
Within one prediction, any mask that hides the target frame from the context frames would allow the same prefill/decode split, so the block-causal structure doesn't seem to be what enables caching here.
My questions:
- Is the block-causal mask mainly a training/modeling choice, for example robustness to corrupted history or matching the rollout setting? Or is it meant to enable a streaming KV cache across AR steps?
- If it's the latter, is there a variant with time-invariant conditioning (absolute frame time instead of window slots or a countdown) where old context K/Vs can be kept and only the newest frame and target recomputed?
- Did you compare against full bidirectional attention within the context frames?
Thanks!
1.
(i) Yes, "any mask that hides the target frame from the context frames would allow the same prefill/decode split," we called that "LAST_FRAME_CAUSAL." We compared "LAST_FRAME_CAUSAL" mask and "BLOCK_CAUSAL" and the latter was more robust to long video generation. But the differences weren't huge.
(ii) It has the nice property that we can stream, or generate n target frames at once, even though we don't implement that
2.
No, but I am currently working on that, with self-forcing for long video generation.
3.
Yes, see 1.
Got it, thx!
On second thought, causal blocking also lacks streaming because it discards the oldest frame to keep a fixed context length; only frame-local masking enables streaming KV, at the cost of no inter-frame interaction within conditioning.