[Model][Perf] Take the denoise step from the loop counter, not a host sync: Cosmos3, Wan 2.2, OmniVoice

mstar/model/cosmos3/submodules.py

@@ -112,6 +112,14 @@
 VIDEO_SOUND_GEN_LOOP = "video_sound_gen_loop"
 ACTION_GEN_LOOP = "action_gen_loop"
 ACTION_VIDEO_GEN_LOOP = "action_video_gen_loop"
+# Each denoise walk's loop; any other walk runs the image loop.
+DENOISE_LOOP_FOR_WALK = {
+    ACTION_GEN_WALK: ACTION_GEN_LOOP,
+    ACTION_VIDEO_GEN_WALK: ACTION_VIDEO_GEN_LOOP,
+    VIDEO_GEN_WALK: VIDEO_GEN_LOOP,
+    VIDEO_GEN_AR_WALK: VIDEO_GEN_AR_LOOP,
+    VIDEO_SOUND_GEN_WALK: VIDEO_SOUND_GEN_LOOP,
+}
 
 # Both action walks run the joint video+action denoise loop body; they differ
 # only in what they emit (the predicted action vs the predicted video).
@@ -152,6 +160,40 @@
 REASONER_LABEL = "main"
 
 
+def _loop_step(fwd_info, graph_walk: str, time_index: torch.Tensor) -> int:
+    """The denoise iteration a step runs, from the engine's loop counter.
+
+    ``time_index`` holds the same value (global over a windowed request's
+    windows and commits, like the counter), but it is the previous step's
+    device output: reading it on the host waits for that step to finish, which
+    stalls the speculative preparation of this one. Callers that drive the
+    submodule without a loop counter (tests, harnesses) fall back to the edge.
+    """
+    loop = DENOISE_LOOP_FOR_WALK.get(graph_walk, IMAGE_GEN_LOOP)
+    counts = getattr(fwd_info, "dynamic_loop_iter_counts", None) or {}
+    if loop in counts:
+        return int(counts[loop])
+    return int(time_index.reshape(-1)[0].item())
+
+
+# Attribute a request's scheduler carries its timesteps under as a host list,
+# filled once at build time. ``set_timesteps`` puts them on the device, and a
+# per-step read there (the guidance-interval test) would be a host sync.
+_HOST_TIMESTEPS = "mstar_host_timesteps"
+
+
+def _with_host_timesteps(scheduler):
+    setattr(scheduler, _HOST_TIMESTEPS, scheduler.timesteps.tolist())
+    return scheduler
+
+
+def _host_timesteps(scheduler) -> list:
+    """The scheduler's timesteps as Python numbers (the values ``.item()``
+    returns), from the host copy when it was built here."""
+    host = getattr(scheduler, _HOST_TIMESTEPS, None)
+    return host if host is not None else scheduler.timesteps.tolist()
+
+
 def native_flow_sigmas(num_inference_steps: int, num_train_timesteps: int):
     """The native flow-matching sigma grid: ``num_inference_steps`` values
     linearly spaced from ``1 - 1/T`` toward 0, the endpoint dropped. A numpy
@@ -430,7 +472,7 @@ def _new_scheduler(self, num_inference_steps: int, device, use_karras_sigma=None
             scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas)
         else:
             scheduler.set_timesteps(num_inference_steps, device=device)
-        return scheduler
+        return _with_host_timesteps(scheduler)
 
     def _new_distilled_scheduler(self, device):
         """The 4-step distilled sampler: a FlowMatchEuler scheduler with the
@@ -450,7 +492,7 @@ def _new_distilled_scheduler(self, device):
                 stochastic_sampling=bool(sc.stochastic_sampling),
             )
         scheduler.set_timesteps(sigmas=[float(x) for x in self.config.distilled_sigmas], device=device)
-        return scheduler
+        return _with_host_timesteps(scheduler)
 
     @staticmethod
     def _scheduler_step(st, velocity, t, latents):
@@ -495,7 +537,7 @@ def prepare_inputs(
         if graph_walk in SOUND_WALKS:
             return self._prepare_video_sound_gen(fwd_info, inputs, device)
         if graph_walk in ACTION_WALKS:
-            return self._prepare_action_gen(fwd_info, inputs, device)
+            return self._prepare_action_gen(graph_walk, fwd_info, inputs, device)
         raise ValueError(f"Unknown Cosmos3 DiT graph walk: {graph_walk!r}")
 
     def _prepare_prefill(self, fwd_info, inputs, device) -> ARNodeInputs:
@@ -1087,12 +1129,13 @@ def _prepare_image_gen(
                 # from the noise drawn above (the reference RNG order).
                 latents = st["vmask"] * st["cond_video_latents"] + (1.0 - st["vmask"]) * latents
             time_index = torch.zeros(1, dtype=torch.long, device=device)
+            step_index = 0
         else:
             latents = inputs["latents"][0]
             time_index = inputs["time_index"][0]
+            step_index = _loop_step(fwd_info, graph_walk, time_index)
 
         scheduler = st["scheduler"]
-        step_index = int(time_index.reshape(-1)[0].item())
         if windowed:
             # The loop counter is global over every window's iterations; the
             # schedule index is the within-window step, and in kv mode the
@@ -1168,13 +1211,14 @@ def _prepare_video_sound_gen(self, fwd_info, inputs, device) -> ARNodeInputs:
                 (1, self.config.sound_dim, st["num_sound"]), generator=gen, device=device, dtype=dtype
             )
             time_index = torch.zeros(1, dtype=torch.long, device=device)
+            step_index = 0
         else:
             latents = inputs["latents"][0]
             sound_latents = inputs["sound_latents"][0]
             time_index = inputs["time_index"][0]
+            step_index = _loop_step(fwd_info, VIDEO_SOUND_GEN_WALK, time_index)
 
         scheduler = st["scheduler"]
-        step_index = int(time_index.reshape(-1)[0].item())
         if step_index >= len(scheduler.timesteps):
             return None
         return ARNodeInputs(
@@ -1187,7 +1231,7 @@ def _prepare_video_sound_gen(self, fwd_info, inputs, device) -> ARNodeInputs:
             ),
         )
 
-    def _prepare_action_gen(self, fwd_info, inputs, device) -> ARNodeInputs:
+    def _prepare_action_gen(self, graph_walk, fwd_info, inputs, device) -> ARNodeInputs:
         st = self.request_states[fwd_info.rid_handle]
         if "latents" not in inputs or len(inputs["latents"]) == 0:
             self._ingest_cond_latents(st, inputs, device)
@@ -1211,13 +1255,14 @@ def _prepare_action_gen(self, fwd_info, inputs, device) -> ARNodeInputs:
             v_noise = randn_tensor(st["latent_shape"], generator=gen, device=device, dtype=dtype)
             latents = st["vmask"] * st["cond_video_latents"] + (1.0 - st["vmask"]) * v_noise
             time_index = torch.zeros(1, dtype=torch.long, device=device)
+            step_index = 0
         else:
             latents = inputs["latents"][0]
             action_latents = inputs["action_latents"][0]
             time_index = inputs["time_index"][0]
+            step_index = _loop_step(fwd_info, graph_walk, time_index)
 
         scheduler = st["scheduler"]
-        step_index = int(time_index.reshape(-1)[0].item())
         if step_index >= len(scheduler.timesteps):
             return None
         return ARNodeInputs(
@@ -1567,10 +1612,10 @@ def _cfg_active(self, st, step_index: int) -> bool:
         gi = st.get("guidance_interval")
         if gi is None:
             return True
-        sched = st["scheduler"]
-        if step_index >= len(sched.timesteps):
+        timesteps = _host_timesteps(st["scheduler"])
+        if step_index >= len(timesteps):
             return False
-        t = float(sched.timesteps[step_index].item())
+        t = float(timesteps[step_index])
         return gi[0] <= t <= gi[1]
 
     def _forward_image_gen(self, attn, st, latents, time_index, **kwargs) -> dict:
@@ -2385,13 +2430,7 @@ def check_stop(self, request_id, request_info, outputs) -> set[str]:
         st = self.request_states.get(request_id)
         if st is None:
             return set()
-        loop = {
-            ACTION_GEN_WALK: ACTION_GEN_LOOP,
-            ACTION_VIDEO_GEN_WALK: ACTION_VIDEO_GEN_LOOP,
-            VIDEO_GEN_WALK: VIDEO_GEN_LOOP,
-            VIDEO_GEN_AR_WALK: VIDEO_GEN_AR_LOOP,
-            VIDEO_SOUND_GEN_WALK: VIDEO_SOUND_GEN_LOOP,
-        }.get(request_info.graph_walk, IMAGE_GEN_LOOP)
+        loop = DENOISE_LOOP_FOR_WALK.get(request_info.graph_walk, IMAGE_GEN_LOOP)
         iter_idx = request_info.dynamic_loop_iter_counts.get(loop, 0)
         # Windowed requests run every window's iterations in one loop; others
         # stop at their scheduler's step count.

mstar/model/omnivoice/submodules.py

@@ -43,6 +43,21 @@
 UNMASK_LOOP_NAME = "unmask_loop"
 
 
+def _loop_step(fwd_info: CurrentForwardPassInfo, step_index: torch.Tensor) -> int:
+    """The unmask iteration a step runs, from the engine's loop counter.
+
+    ``step_index`` holds the same value, but it is the previous step's device
+    output: reading it on the host waits for that step to finish, which stalls
+    the speculative preparation of the next one. The counter is per batch and
+    counts speculative iterations too (#351). Callers that drive the
+    submodule without a loop counter (tests, harnesses) fall back to the edge.
+    """
+    counts = getattr(fwd_info, "dynamic_loop_iter_counts", None) or {}
+    if UNMASK_LOOP_NAME in counts:
+        return int(counts[UNMASK_LOOP_NAME])
+    return int(step_index.reshape(-1)[0].item())
+
+
 class _SingleRequestMixin:
     """Serve one request per step through the engine's batched entry point.
 
@@ -173,9 +188,6 @@ def prepare_inputs(
 
         prefix_ids = inputs["prefix_ids"][0].to(device)
         prefix_audio_mask = inputs["prefix_audio_mask"][0].to(device)
-        # target_len rides an edge rather than step_metadata: computing it needs
-        # the tokenizer and the duration estimator, both data-worker assets.
-        target_len = int(inputs["target_len"][0].reshape(-1)[0].item())
 
         # Cloning walk only: the codec's reference tokens close out the prefix.
         # They are joined here rather than on the data worker because the encode
@@ -197,7 +209,11 @@ def prepare_inputs(
             # Iteration 0: the loop-back edges arrive empty. Seed a fully
             # masked canvas; there is no noise and no seed to honour, because
             # the stochastic part of this model is the reveal order, not the
-            # starting state.
+            # starting state. target_len rides an edge rather than
+            # step_metadata: computing it needs the tokenizer and the duration
+            # estimator, both data-worker assets. Later iterations read it off
+            # the canvas shape instead, which costs no host read.
+            target_len = int(inputs["target_len"][0].reshape(-1)[0].item())
             tokens = torch.full(
                 (self.config.num_audio_codebook, target_len),
                 self.config.audio_mask_id,
@@ -209,11 +225,7 @@ def prepare_inputs(
             step_index = torch.zeros(1, dtype=torch.int64, device=device)
         else:
             step_index = inputs["step_index"][0]
-            # Read the counter off the edge tensor, not from
-            # fwd_info.dynamic_loop_iter_counts: the worker sets that wrong on
-            # the speculative path, so the edge is the reliable source until
-            # that is fixed upstream. Costs a .item() per iteration.
-            k = int(step_index.reshape(-1)[0].item())
+            k = _loop_step(fwd_info, step_index)
             num_step = int(meta["num_step"])
             if k >= num_step:
                 # Async scheduling dispatched an iteration past this request's
@@ -225,6 +237,7 @@ def prepare_inputs(
                 )
                 return None
             tokens = inputs["audio_tokens"][0].to(device)
+            target_len = tokens.shape[-1]
 
         return NodeInputs(
             tensor_inputs={
@@ -350,13 +363,12 @@ def postprocess(
 
         Scoring, the reveal schedule and the in-place write are per-request and
         data-dependent, which is what keeps them out of ``forward_batched``.
-        The iteration index is read off the ``step_index`` edge rather than
-        ``dynamic_loop_iter_counts``, which the worker sets wrong on the
-        speculative path; everything else comes from step metadata.
+        The iteration index is the engine's loop counter (see ``_loop_step``);
+        everything else comes from step metadata.
         """
         assert inputs is not None, "OmniVoice backbone postprocess needs its inputs"
         meta = request_info.step_metadata
-        k = int(inputs.tensor_inputs["step_index"].reshape(-1)[0].item())
+        k = _loop_step(request_info, inputs.tensor_inputs["step_index"])
 
         # Cloned because apply_reveal writes in place and this tensor is the
         # edge the engine routed in from the previous iteration; mutating it

mstar/model/wan22/submodules.py

@@ -154,6 +154,20 @@ def _set_buffer(module: nn.Module, dotted_name: str, value: torch.Tensor):
     setattr(owner, attr, value)
 
 
+def _loop_step(fwd_info: CurrentForwardPassInfo, time_index: torch.Tensor) -> int:
+    """The denoise iteration this step runs, from the engine's loop counter.
+
+    ``time_index`` holds the same value, but it is the previous step's device
+    output: reading it on the host waits for that step to finish, which stalls
+    the speculative preparation of this one. Callers that drive the submodule
+    without a loop counter (tests, harnesses) fall back to the edge.
+    """
+    counts = getattr(fwd_info, "dynamic_loop_iter_counts", None) or {}
+    if DENOISE_LOOP_NAME in counts:
+        return int(counts[DENOISE_LOOP_NAME])
+    return int(time_index.reshape(-1)[0].item())
+
+
 def _latent_grid(config: Wan22Config, step_metadata: dict) -> tuple[int, int, int]:
     """(T_lat, H_lat, W_lat) for the request's pixel-space dims (diffusers
     ``prepare_latents``: T=(F-1)//4+1, spatial /16)."""
@@ -404,7 +418,7 @@ def prepare_inputs(
             tensor_inputs["unipc_last_sample"] = torch.zeros(shape, dtype=torch.float32, device=device)
         else:
             time_index = inputs["time_index"][0]
-            k = int(time_index.reshape(-1)[0].item())
+            k = _loop_step(fwd_info, time_index)
             num_steps = int(fwd_info.step_metadata["num_inference_steps"])
             if k >= num_steps:
                 # Async scheduling dispatched an iteration past this request's

test/cpu-core.txt

@@ -34,6 +34,7 @@ test/modular/test_kv_publish_retrieve.py
 test/modular/test_parallel_groups.py
 test/modular/test_preplan_error_paths.py
 test/modular/test_diffusion_denoise_loop.py
+test/modular/test_step_prep_host_syncs.py
 test/modular/test_diffusion_noise_stager.py
 test/modular/test_flux2_klein_model.py
 test/modular/test_z_image_model.py