[perf] matrixgame2 causal: cache per-forward RoPE/timestep/blockmask recompute on-device

matrixgame2 causal inference rebuilt several constant tables on EVERY DiT
forward, on CPU + H2D copy each time:
  - main RoPE tables (causal_model, float64 on CPU -> .to(device))
  - sinusoidal timestep embedding (visual_embedding, arange on CPU -> .to)
  - action-module RoPE freqs (cached but as CPU tensors -> .to every call)
  - 4x flex-attention BlockMasks via create_block_mask (compile + a
    GPU-tensor python loop) every forward

All of these depend only on constants (grid, start_frame, block params)
that are fixed within an AR block (and mostly across the whole gen), so
they are now memoized on-device keyed by their varying args. Numerically
identical to before (pure memoization of deterministic functions); the
sinusoid table is simply built on the target device instead of CPU+H2D.

Measured (Modal H100x1, num_frames=117 = 10 AR blocks x 3 DMD steps,
under nsys, vs the prior .item() fix):

  isolated DiT forward (microbench):  194 ms -> 95 ms   (~2.04x)
  end-to-end inference wall:          17.2 s -> 13.1 s  (-24%)
  denoising stage wall:               10.56 s -> 7.16 s (-32%)
  GPU active share (denoise):         23.0% -> 33.5%

The biggest single contributor is the BlockMask cache (create_block_mask
is expensive). Remaining GPU idle is now dominated by per-kernel Python
dispatch of the ~150k small kernels, which only CUDA Graph / persistent
kernel can address.

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,21 @@ 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
+        # 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:
-            self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed(
+            _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

@@ -977,21 +977,36 @@ def _forward_inference(
 
         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)
+        # RoPE tables depend only on (grid, start_frame); both are constant
+        # within an AR block's DMD steps. Recomputing them on CPU in float64
+        # and copying H2D every forward is wasteful AND breaks CUDA-graph
+        # capture (H2D copy is illegal mid-capture). Cache the device tensors
+        # keyed by (start_frame, grid) so the H2D happens once per block.
+        _rope_key = (start_frame, post_patch_num_frames, post_patch_height,
+                     post_patch_width)
+        _rope_cache = getattr(self, "_rope_device_cache", None)
+        if _rope_cache is None:
+            _rope_cache = {}
+            self._rope_device_cache = _rope_cache
+        if _rope_key in _rope_cache:
+            freqs_cos, freqs_sin = _rope_cache[_rope_key]
+        else:
+            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)
+            _rope_cache[_rope_key] = (freqs_cos, freqs_sin)
         freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
 
         hidden_states = self.patch_embedding(hidden_states)
@@ -1030,36 +1045,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)
+        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)
+            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)
+            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)
+        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: