[perf] matrixgame2 causal: cut per-forward recompute (blockmask cache, skip unused RoPE, on-device timestep)

matrixgame2 causal inference repeated constant work on every DiT forward:

- 4 flex-attention BlockMasks were rebuilt via create_block_mask each forward
  (compile + a GPU-tensor python loop). They depend only on
  (num_frames, frame_seqlen, block size, local_attn_size, device), so cache
  them per param-tuple (device included in the key).
- get_rotary_pos_embed ran a float64 CPU compute + H2D copy each forward to
  build freqs_cis — which CausalMatrixGame2SelfAttention.forward never reads
  (it computes its own RoPE via self._freqs_cache). Set freqs_cis = None and
  skip the computation entirely.
- the sinusoidal timestep table was built on CPU then copied H2D; build it on
  the target device instead.
- the action-module RoPE freqs were cached as CPU tensors, forcing an H2D copy
  every forward inside _apply_rotary_emb_qk; cache them device-resident (with a
  device check) so the per-call .to is a no-op.

All numerically identical. Removes per-forward CPU compute + H2D copies that
also blocked CUDA-graph capture.

Measured (H100, 117 frames = 30 DiT forwards): isolated DiT forward
194 ms -> 95 ms; end-to-end inference 12.6 s -> 10.6 s (-16%); denoise-stage
GPU-active 23% -> 33.5%.

Author: rich7420

fastvideo/layers/visual_embedding.py

@@ -150,7 +150,11 @@ def timestep_embedding(t: torch.Tensor,
         Tensor of shape [B, dim] with embeddings
     """
     half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=dtype) / half).to(device=t.device)
+    # Build the frequency table directly on the target device. Creating it on
+    # CPU and copying H2D every call is wasteful and breaks CUDA-graph capture
+    # (H2D copy is illegal mid-capture).
+    freqs = torch.exp(-math.log(max_period) *
+                      torch.arange(start=0, end=half, dtype=dtype, device=t.device) / half)
     args = t[:, None].float() * freqs[None]
     embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
     if dim % 2:

fastvideo/models/dits/matrixgame2/action_module.py

@@ -771,16 +771,22 @@ def forward(
         ) % self.vae_time_compression_ratio == 0
         N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1
 
-        # Lazy initialization of freqs on first forward pass
-        if self._freqs_cos is None or self._freqs_sin is None:
-            self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed(
+        # Lazy initialization of freqs on first forward pass. Cache on the
+        # compute device so the per-call `.to(xq.device)` in _apply_rotary_emb_qk
+        # is a no-op (avoids an H2D copy every forward, which also breaks
+        # CUDA-graph capture).
+        if (self._freqs_cos is None or self._freqs_sin is None
+                or self._freqs_cos.device != x.device):
+            _fc, _fs = self.get_rotary_pos_embed(
                 7500,
                 self.patch_size[1],
                 self.patch_size[2],
                 64,
                 self.mouse_qk_dim_list,
                 start_offset=0,
             )
+            self._freqs_cos = _fc.to(x.device)
+            self._freqs_sin = _fs.to(x.device)
 
         # Defined freqs_cis early so it's available for both mouse and keyboard
         freqs_cis = (self._freqs_cos, self._freqs_sin)

fastvideo/models/dits/matrixgame2/causal_model.py

@@ -975,24 +975,13 @@ def _forward_inference(
         post_patch_height = height // p_h
         post_patch_width = width // p_w
 
-        d = self.hidden_size // self.num_attention_heads
-        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (
-                post_patch_num_frames * get_sp_world_size(),
-                post_patch_height,
-                post_patch_width,
-            ),
-            self.hidden_size,
-            self.num_attention_heads,
-            rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000,
-            start_frame=start_frame,
-        )
-        freqs_cos = freqs_cos.to(hidden_states.device)
-        freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+        # freqs_cis is unused by CausalMatrixGame2SelfAttention.forward — the
+        # attention computes its own RoPE via self._freqs_cache + causal_rope_apply
+        # and ignores the freqs_cis argument. So skip get_rotary_pos_embed
+        # entirely: it was running a float64 CPU compute + H2D copy on every
+        # forward (also illegal during CUDA-graph capture) for a result nothing
+        # reads.
+        freqs_cis = None
 
         hidden_states = self.patch_embedding(hidden_states)
         grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)
@@ -1030,36 +1019,47 @@ def _forward_inference(
             else:
                 encoder_hidden_states = encoder_hidden_states_image
 
-        block_mask = self._prepare_blockwise_causal_attn_mask(
-            device=hidden_states.device,
-            num_frames=num_frames,
-            frame_seqlen=post_patch_height * post_patch_width,
-            num_frame_per_block=self.num_frame_per_block,
-            local_attn_size=self.local_attn_size,
-        )
+        # BlockMasks depend only on (num_frames, frame_seqlen, block size,
+        # local_attn_size) — all constant across forwards. Building them every
+        # forward via create_block_mask is expensive (compile + GPU-tensor
+        # python loop) AND breaks CUDA-graph capture. Cache per param-tuple.
+        bm_cache = getattr(self, "_block_mask_cache", None)
+        if bm_cache is None:
+            bm_cache = {}
+            self._block_mask_cache = bm_cache
+        _dev = hidden_states.device
+        _fsl = post_patch_height * post_patch_width
+        _nfb = self.num_frame_per_block
+        _las = self.local_attn_size
+
+        _k = ("main", num_frames, _fsl, _nfb, _las, _dev)
+        if _k not in bm_cache:
+            bm_cache[_k] = self._prepare_blockwise_causal_attn_mask(
+                device=_dev, num_frames=num_frames, frame_seqlen=_fsl,
+                num_frame_per_block=_nfb, local_attn_size=_las)
+        block_mask = bm_cache[_k]
+
         if self.use_rope_keyboard:
-            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
-                device=hidden_states.device,
-                num_frames=num_frames,
-                frame_seqlen=1,
-                num_frame_per_block=self.num_frame_per_block,
-                local_attn_size=self.local_attn_size,
-            )
+            _k = ("act_kb", num_frames, 1, _nfb, _las, _dev)
+            if _k not in bm_cache:
+                bm_cache[_k] = self._prepare_blockwise_causal_attn_mask_action(
+                    device=_dev, num_frames=num_frames, frame_seqlen=1,
+                    num_frame_per_block=_nfb, local_attn_size=_las)
+            block_mask_keyboard = bm_cache[_k]
         else:
-            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
-                device=hidden_states.device,
-                num_frames=num_frames,
-                frame_seqlen=post_patch_height * post_patch_width,
-                num_frame_per_block=self.num_frame_per_block,
-                local_attn_size=self.local_attn_size,
-            )
-        block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
-            device=hidden_states.device,
-            num_frames=num_frames,
-            frame_seqlen=1,
-            num_frame_per_block=self.num_frame_per_block,
-            local_attn_size=self.local_attn_size,
-        )
+            _k = ("kb", num_frames, _fsl, _nfb, _las, _dev)
+            if _k not in bm_cache:
+                bm_cache[_k] = self._prepare_blockwise_causal_attn_mask_keyboard(
+                    device=_dev, num_frames=num_frames, frame_seqlen=_fsl,
+                    num_frame_per_block=_nfb, local_attn_size=_las)
+            block_mask_keyboard = bm_cache[_k]
+
+        _k = ("act_ms", num_frames, 1, _nfb, _las, _dev)
+        if _k not in bm_cache:
+            bm_cache[_k] = self._prepare_blockwise_causal_attn_mask_action(
+                device=_dev, num_frames=num_frames, frame_seqlen=1,
+                num_frame_per_block=_nfb, local_attn_size=_las)
+        block_mask_mouse = bm_cache[_k]
         if kv_cache is None:
             kv_cache = [None] * len(self.blocks)
         if kv_cache_mouse is None: