Skip to content

Commit e51c063

Browse files
committed
[bugfix] matrixgame2 causal: build block-mask cache with effective block size
The cached BlockMask was keyed/built from self.num_frame_per_block, but the transformer blocks run with effective_num_frame_per_block (the per-call num_frame_per_block override). Align the mask with the block size the blocks actually use; it is part of the cache key so different sizes get distinct entries. Verified on H100: the size-1 and size-3 masks genuinely differ, so any forward whose effective size diverges from the default (e.g. a boundary block of size 1) would otherwise apply the wrong attention pattern. For the current shipped matrixgame2 configs the two sizes coincide, so output is unchanged — this removes the latent mismatch.
1 parent 69ff7ba commit e51c063

1 file changed

Lines changed: 3 additions & 1 deletion

File tree

fastvideo/models/dits/matrixgame2/causal_model.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1029,7 +1029,9 @@ def _forward_inference(
10291029
self._block_mask_cache = bm_cache
10301030
_dev = hidden_states.device
10311031
_fsl = post_patch_height * post_patch_width
1032-
_nfb = self.num_frame_per_block
1032+
# Use the block size the blocks actually run with (the per-call override),
1033+
# not the model default, so the mask matches. _nfb is in the cache key.
1034+
_nfb = effective_num_frame_per_block
10331035
_las = self.local_attn_size
10341036

10351037
_k = ("main", num_frames, _fsl, _nfb, _las, _dev)

0 commit comments

Comments
 (0)