Fix streaming-inference alignment with training-time contracts

wan_va/modules/model.py

@@ -361,6 +361,12 @@ def init_kv_cache(self, cache_name, total_tolen, num_head, head_dim,
             torch.zeros((total_tolen, ), dtype=torch.bool, device=device),
             "is_pred":
             torch.zeros((total_tolen, ), dtype=torch.bool, device=device),
+            "frame_id":
+            torch.full((total_tolen, ), -1, device=device, dtype=torch.long),
+            "current_frame_id":
+            -1,
+            "window_size":
+            999999,  # default: no window limit
         }
 
     def allocate_slots(self, cache_name, key_size):
@@ -406,8 +412,49 @@ def update_cache(self, cache_name, key, value, is_pred):
         cache['mask'][slots] = True
         cache['id'][slots] = new_id
         cache['is_pred'][slots] = is_pred
+        cache['frame_id'][slots] = cache.get('current_frame_id', -1)
         return slots
 
+    def set_frame_context(self, cache_name, frame_id, window_size):
+        """Set current frame_id and window_size for subsequent cache updates and attention filtering."""
+        if self.attn_caches is None or cache_name not in self.attn_caches:
+            return
+        cache = self.attn_caches[cache_name]
+        if cache is not None:
+            cache['current_frame_id'] = frame_id
+            cache['window_size'] = window_size
+
+    @staticmethod
+    def _select_valid_cache_slots(cache, update_cache):
+        """Filter cache slots to match training-time attention visibility.
+
+        During training, a noisy query may attend to:
+          - past clean tokens (real observations / real actions from earlier frames)
+          - same-frame predicted tokens (within the same chunk)
+          - strictly-past predicted tokens
+        and only within the attention window. The original streaming-inference
+        code instead surfaced every allocated slot, which let noisy queries
+        attend to same-frame clean tokens and to predicted tokens from future
+        frames -- a train/inference distribution mismatch.
+        """
+        valid = cache['mask'].nonzero(as_tuple=False).squeeze(-1)
+        if update_cache not in (0, 1) or valid.numel() == 0:
+            return valid
+
+        is_pred = cache['is_pred'][valid]
+        frame_ids = cache['frame_id'][valid]
+        current_frame_id = cache.get('current_frame_id', -1)
+        window_size = cache.get('window_size', 999999)
+
+        is_clean = ~is_pred
+        is_past = (frame_ids >= 0) & (frame_ids < current_frame_id)
+        is_same_frame = frame_ids == current_frame_id
+        within_window = (frame_ids - current_frame_id).abs() <= window_size
+        # Training mask: noisy query -> clean key must be past; noisy query -> noisy
+        # key must be same-frame or past; all subject to the attention window.
+        allowed = (is_clean & is_past) | (is_pred & (is_same_frame | is_past))
+        return valid[allowed & within_window]
+
     def restore_cache(self, cache_name, slots):
         self.attn_caches[cache_name]['mask'][slots] = False
 
@@ -447,8 +494,10 @@ def apply_rotary_emb(x, freqs):
                                       is_pred=(update_cache == 1))
             key_pool = self.attn_caches[cache_name]['k']
             value_pool = self.attn_caches[cache_name]['v']
-            mask = self.attn_caches[cache_name]['mask']
-            valid = mask.nonzero(as_tuple=False).squeeze(-1)
+            valid = self._select_valid_cache_slots(
+                self.attn_caches[cache_name], update_cache
+            )
+
             key = key_pool[:, valid]
             value = value_pool[:, valid]
 
@@ -658,6 +707,11 @@ def clear_pred_cache(self, cache_name):
         for block in self.blocks:
             block.attn1.clear_pred_cache(cache_name)
 
+    def set_frame_context(self, cache_name, frame_id, window_size):
+        """Set frame context for inference attention masking across all layers."""
+        for block in self.blocks:
+            block.attn1.set_frame_context(cache_name, frame_id, window_size)
+
     def create_empty_cache(self, cache_name, attn_window,
                            latent_token_per_chunk, action_token_per_chunk,
                            device, dtype, batch_size):

wan_va/wan_va_server.py

@@ -250,7 +250,37 @@ def postprocess_action(self, action):
             raise NotImplementedError
         action = action.squeeze(0).detach().cpu().numpy()
         return action[self.job_config.used_action_channel_ids]
-
+
+    def _make_zero_action_condition(self):
+        """Build the cold-start action conditioning token.
+
+        At the very first inference step (``frame_st_id == 0``) there is no
+        real action history to feed the model, so we construct an all-zero
+        action in *physical* space. Crucially, the model operates in
+        quantile-normalized space, where physical zero is *not* the tensor
+        ``0`` -- it is ``(0 - q01) / (q99 - q01) * 2 - 1`` per channel.
+        Directly feeding ``torch.zeros`` therefore decoded to the per-channel
+        statistical midpoint ``(q01 + q99) / 2`` rather than to physical zero,
+        seeding the diffusion loop with a spurious non-zero action. That
+        erroneous conditioning entered both the action denoising loop and the
+        KV cache, corrupting every subsequent chunk.
+        """
+        action_cond = torch.zeros(
+            [1, self.job_config.action_dim, 1, self.action_per_frame, 1],
+            device=self.device,
+            dtype=torch.float32,
+        )
+        if self.action_norm_method == 'quantiles':
+            q01 = self.actions_q01.to(self.device).view(1, -1, 1, 1, 1)
+            q99 = self.actions_q99.to(self.device).view(1, -1, 1, 1, 1)
+            action_cond = (action_cond - q01) / (q99 - q01 + 1e-6) * 2. - 1.
+        else:
+            raise NotImplementedError
+
+        action_mask = self.action_mask.to(self.device)
+        action_cond[:, ~action_mask] = 0
+        return action_cond.to(self.dtype)
+
     def _repeat_input_for_cfg(self, input_dict):
         if self.use_cfg:
             input_dict['noisy_latents'] = input_dict['noisy_latents'].repeat(2, 1, 1, 1, 1)
@@ -486,6 +516,11 @@ def _infer(self, obs, frame_st_id=0):
                 torch.no_grad(),
         ):
             # 1. Video Generation Loop
+            # Set frame context so noisy latent queries cannot attend to
+            # predicted tokens of future frames in the KV cache.
+            latent_frame_id = frame_st_id
+            self.transformer.set_frame_context(self.cache_name, latent_frame_id,
+                                               self.job_config.attn_window)
             for i, t in enumerate(tqdm(timesteps)):
                 last_step = i == len(timesteps) - 1
                 latent_cond = init_latent[:, :, 0:1].to(
@@ -521,15 +556,16 @@ def _infer(self, obs, frame_st_id=0):
 
                 latents[:, :, 0:1] = latent_cond if frame_st_id == 0 else latents[:, :, 0:1]
 
+            # 2. Action Generation Loop
+            # Action tokens live at the *next* frame id relative to the video
+            # tokens they accompany; advance the frame context accordingly.
+            action_frame_id = frame_st_id + 1
+            self.transformer.set_frame_context(self.cache_name, action_frame_id,
+                                               self.job_config.attn_window)
             for i, t in enumerate(tqdm(action_timesteps)):
                 last_step = i == len(action_timesteps) - 1
-                action_cond = torch.zeros(
-                    [
-                        1, self.job_config.action_dim, 1,
-                        self.action_per_frame, 1
-                    ],
-                    device=self.device,
-                    dtype=self.dtype) if frame_st_id == 0 else None
+                action_cond = self._make_zero_action_condition(
+                ) if frame_st_id == 0 else None
 
                 input_dict = self._prepare_latent_input(
                     None,
@@ -591,11 +627,17 @@ def _compute_kv_cache(self, obs):
         with (
                 torch.no_grad(),
         ):
+            # Set latent frame context (clean observations, update_cache=2)
+            self.transformer.set_frame_context(self.cache_name, self.frame_st_id,
+                                               self.job_config.attn_window)
             self.transformer(self._repeat_input_for_cfg(input_dict['latent_res_lst']),
                              update_cache=2,
                              cache_name=self.cache_name,
                              action_mode=False)
 
+            # Set action frame context (clean observations, update_cache=2)
+            self.transformer.set_frame_context(self.cache_name, self.frame_st_id + 1,
+                                               self.job_config.attn_window)
             self.transformer(self._repeat_input_for_cfg(input_dict['action_res_lst']),
                              update_cache=2,
                              cache_name=self.cache_name,