[perf]: magi-human denoising: precompute step-invariant packed layout

build_static_packed_inputs was called every step inside the denoising
loop, redoing the meshgrid/torch.full() coords + modality-map work even
though those depend only on latent shape, audio length, and channel
widths — all fixed for a single generation.

Add StaticPackedLayout + precompute_static_packed_layout. The latent prep
stage stashes the layout on the batch; the denoise / sr-denoise loops
pass it back through the new layout= arg, which short-circuits the
invariant work and only rebuilds the per-step token tensors. The slow
path (layout=None) is kept bit-exact for build_packed_inputs callers in
parity tests.

Bit-exact verified slow vs fast path. dit / distill_dit / pipeline_smoke
/ sr540p / sr1080p / vae parity tests pass.

fastvideo/pipelines/basic/magi_human/stages/denoising.py

@@ -162,6 +162,7 @@ class MagiHumanDenoisingStage(PipelineStage):
                 audio_feat_len=audio_feat_len,
                 patch_size=self.patch_size,
                 coords_style=self.coords_style,
+                layout=getattr(batch, "magi_static_packed_layout", None),
             )
             with trace_step(idx), set_forward_context(
                     current_timestep=int(t.item()) if torch.is_tensor(t) else int(t),

fastvideo/pipelines/basic/magi_human/stages/latent_preparation.py

@@ -213,6 +213,18 @@ class MagiHumanLatentPreparationStage(PipelineStage):
         batch.magi_latent_T = latent_T
         batch.magi_latent_H = latent_H
         batch.magi_latent_W = latent_W
+        # Precompute the step-invariant packed layout (coords / modality
+        # maps / channel-padding width) once; the denoise loop reuses it
+        # every step instead of rebuilding meshgrids on each call.
+        batch.magi_static_packed_layout = precompute_static_packed_layout(
+            latent_shape=tuple(video_latent.shape),  # type: ignore[arg-type]
+            audio_feat_len=int(audio_latent.shape[1]),
+            z_dim=self.z_dim,
+            audio_in_channels=self.audio_in_channels,
+            patch_size=self.patch_size,
+            coords_style=self.coords_style,
+            device=video_latent.device,
+        )
         return batch
 
 
@@ -256,12 +268,117 @@ class StaticPackedInputs:
         self.max_ch = max_ch
 
 
+class StaticPackedLayout:
+    """Step- and value-invariant portion of the static packed inputs.
+
+    Coords, modality maps, and the channel-padding width depend only on the
+    latent shape, audio length, channel widths, and patch sizes — all fixed
+    for a single generation. Precompute once before the denoise loop and
+    reuse on every step. Only the per-step token tensors must be rebuilt.
+    """
+
+    __slots__ = (
+        "video_coords",
+        "audio_coords",
+        "video_mm",
+        "audio_mm",
+        "max_ch",
+        "video_token_num",
+        "audio_feat_len",
+    )
+
+    def __init__(
+        self,
+        video_coords: torch.Tensor,
+        audio_coords: torch.Tensor,
+        video_mm: torch.Tensor,
+        audio_mm: torch.Tensor,
+        max_ch: int,
+        video_token_num: int,
+        audio_feat_len: int,
+    ) -> None:
+        self.video_coords = video_coords
+        self.audio_coords = audio_coords
+        self.video_mm = video_mm
+        self.audio_mm = audio_mm
+        self.max_ch = max_ch
+        self.video_token_num = video_token_num
+        self.audio_feat_len = audio_feat_len
+
+
+def precompute_static_packed_layout(
+    latent_shape: tuple[int, int, int, int, int],
+    audio_feat_len: int,
+    z_dim: int,
+    audio_in_channels: int,
+    patch_size: tuple[int, int, int],
+    coords_style: Literal["v1", "v2"] = "v2",
+    device: torch.device | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> StaticPackedLayout:
+    """Precompute the invariant fields used by ``build_static_packed_inputs``.
+
+    Arguments are derived from configs and latent shape — none depend on
+    the current denoising-step values. Call this once in the latent
+    preparation stage (or any pre-loop site) and pass the result via the
+    ``layout=`` arg of ``build_static_packed_inputs`` to skip the
+    meshgrid/full() work on every step.
+    """
+    pT, pH, pW = patch_size
+    _, _, T, H, W = latent_shape
+    if device is None:
+        device = torch.device("cpu")
+
+    video_token_num = (T // pT) * (H // pH) * (W // pW)
+    # `_img2tokens` packs to channel `z_dim * pT * pH * pW`; audio tokens
+    # are `audio_in_channels` wide — both are config constants.
+    max_ch = max(z_dim * pT * pH * pW, audio_in_channels)
+
+    video_ref_shape = (T // pT, H // pH, W // pW)
+    video_coords = _build_coords(
+        shape=video_ref_shape,
+        ref_feat_shape=video_ref_shape,
+        device=device,
+        dtype=dtype,
+    )
+
+    if coords_style == "v2":
+        audio_ref_t = (audio_feat_len - 1) // _AUDIO_TEMPORAL_COMPRESSION + 1
+        audio_coords = _build_coords(
+            shape=(audio_feat_len, 1, 1),
+            ref_feat_shape=(audio_ref_t // pT, 1, 1),
+            device=device,
+            dtype=dtype,
+        )
+    else:
+        audio_coords = _build_coords(
+            shape=(audio_feat_len, 1, 1),
+            ref_feat_shape=(T // pT, 1, 1),
+            device=device,
+            dtype=dtype,
+        )
+
+    video_mm = torch.full((video_token_num, ), MODALITY_VIDEO, dtype=torch.int64, device=device)
+    audio_mm = torch.full((audio_feat_len, ), MODALITY_AUDIO, dtype=torch.int64, device=device)
+
+    return StaticPackedLayout(
+        video_coords=video_coords,
+        audio_coords=audio_coords,
+        video_mm=video_mm,
+        audio_mm=audio_mm,
+        max_ch=max_ch,
+        video_token_num=video_token_num,
+        audio_feat_len=audio_feat_len,
+    )
+
+
 def build_static_packed_inputs(
     video_latent: torch.Tensor,
     audio_latent: torch.Tensor,
     audio_feat_len: int,
     patch_size: tuple[int, int, int],
     coords_style: Literal["v1", "v2"] = "v2",
+    layout: StaticPackedLayout | None = None,
 ) -> StaticPackedInputs:
     """Build the step-invariant portion of the packed token stream.
 
@@ -271,14 +388,37 @@ def build_static_packed_inputs(
 
     Mirrors SingleData.token_sequence / coords_mapping / modality_mapping in
     inference/pipeline/data_proxy.py, minus the text portion.
+
+    When ``layout`` is provided, coords / modality maps / max_ch are taken
+    from the precomputed values and only the per-step token tensors are
+    rebuilt; this is the hot-path call from the denoising loop. When
+    ``layout`` is None the function recomputes everything from scratch
+    (e.g. for one-shot tests via ``build_packed_inputs``).
     """
     pT, pH, pW = patch_size
-    _, z_dim, T, H, W = video_latent.shape
     assert video_latent.size(0) == 1, "batch size 1 required for MagiHuman base"
 
     video_tokens = _img2tokens(video_latent, t_patch=pT, patch=pH)[0]
     audio_tokens = audio_latent[0, :audio_feat_len].contiguous()
 
+    if layout is not None:
+        max_ch = layout.max_ch
+        video_tokens = F.pad(video_tokens, (0, max_ch - video_tokens.size(-1)))
+        audio_tokens = F.pad(audio_tokens, (0, max_ch - audio_tokens.size(-1)))
+        return StaticPackedInputs(
+            video_tokens=video_tokens,
+            audio_tokens=audio_tokens,
+            video_coords=layout.video_coords,
+            audio_coords=layout.audio_coords,
+            video_mm=layout.video_mm,
+            audio_mm=layout.audio_mm,
+            max_ch=max_ch,
+        )
+
+    # Slow path: rebuild every invariant from scratch. Kept for the
+    # ``build_packed_inputs`` one-shot wrapper used by tests/parity helpers.
+    _, z_dim, T, H, W = video_latent.shape
+
     max_ch = max(video_tokens.size(-1), audio_tokens.size(-1))
     video_tokens = F.pad(video_tokens, (0, max_ch - video_tokens.size(-1)))
     audio_tokens = F.pad(audio_tokens, (0, max_ch - audio_tokens.size(-1)))

fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py

@@ -108,6 +108,7 @@ class MagiHumanSRDenoisingStage(PipelineStage):
                 audio_feat_len=audio_feat_len,
                 patch_size=self.patch_size,
                 coords_style=self.coords_style,
+                layout=getattr(batch, "magi_static_packed_layout", None),
             )
             with trace_step(idx), set_forward_context(
                     current_timestep=int(t.item()) if torch.is_tensor(t) else int(t),