Fix rope：cache to each request

lightx2v/models/networks/ernie_image/infer/pre_infer.py

@@ -50,10 +50,13 @@ def __init__(self, config):
         self.rope_axes_dim = config.get("rope_axes_dim", (32, 48, 48))
         self.rotary_dim = sum(self.rope_axes_dim)
         self.rope = None
+        self.scheduler = None
         self.clear_rope_cache()
 
     def clear_rope_cache(self):
-        self._rope_cache = {}
+        self._rope_cache_request_id = None
+        # True：infer_condition=True； False：infer_condition=False
+        self._rope_cache = {True: None, False: None}
 
     def set_rope(self, rope):
         self.rope = rope
@@ -63,6 +66,10 @@ def set_scheduler(self, scheduler):
         self.scheduler = scheduler
         self.clear_rope_cache()
 
+    @staticmethod
+    def _device_key(device):
+        return device.type, device.index
+
     def prepare_rope_cache(self, image_hw, text_len, device):
         if self.rope is None:
             raise RuntimeError("ErnieImagePreInfer RoPE is not initialized.")
@@ -74,12 +81,28 @@ def prepare_rope_cache(self, image_hw, text_len, device):
         return rotary_freqs, rotary_positions
 
     def get_rope_cache(self, image_hw, text_len, device):
-        cache_key = (*image_hw, text_len, str(device))
-        cached = self._rope_cache.get(cache_key)
-        if cached is None:
-            cached = self.prepare_rope_cache(image_hw, text_len, device)
-            self._rope_cache[cache_key] = cached
-        return cached
+        if self.scheduler is None:
+            raise RuntimeError("ErnieImagePreInfer scheduler is not initialized.")
+
+        request_id = self.scheduler.rope_request_id
+        if request_id != self._rope_cache_request_id:
+            self._rope_cache_request_id = request_id
+            self._rope_cache = {True: None, False: None}
+
+        cache_key = (self._device_key(device), *image_hw, text_len)
+        branch = bool(self.scheduler.infer_condition)
+        cached = self._rope_cache[branch]
+        if cached is not None and cached[0] == cache_key:
+            return cached[1]
+
+        shared = self._rope_cache[not branch]
+        if shared is not None and shared[0] == cache_key:
+            self._rope_cache[branch] = shared
+            return shared[1]
+
+        value = self.prepare_rope_cache(image_hw, text_len, device)
+        self._rope_cache[branch] = (cache_key, value)
+        return value
 
     def _pos_embed(self, ids: torch.Tensor) -> torch.Tensor:
         emb = torch.cat([_rope(ids[..., i], self.rope_axes_dim[i], self.rope_theta) for i in range(3)], dim=-1)

lightx2v/models/networks/motus/infer/pre_infer.py

@@ -11,16 +11,45 @@ def __init__(self, model, config):
         super().__init__(config)
         self.model = model
         self.scheduler = None
-        self._grid_cache = {}
+        self.clear_request_cache()
+
+    def clear_request_cache(self):
+        self._rope_cache_request_id = None
+        self._grid_cache = None
+        self.cos_sin = None
+        self.rope_positions = None
+        self.grid_sizes = (0, 0, 0)
+
+    def set_rope(self, rope):
+        super().set_rope(rope)
+        self.clear_request_cache()
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
+        self.clear_request_cache()
+
+    def _begin_request(self):
+        request_id = self.scheduler.rope_request_id
+        if request_id != self._rope_cache_request_id:
+            self.clear_request_cache()
+            self._rope_cache_request_id = request_id
+
+    def _get_grid_output(self, batch_size, grid_tuple, device):
+        grid_key = (batch_size, grid_tuple, device.type, device.index)
+        if self._grid_cache is not None and self._grid_cache[0] == grid_key:
+            return self._grid_cache[1]
+
+        grid_sizes = torch.tensor([grid_tuple], dtype=torch.long, device=device).expand(batch_size, -1)
+        grid_output = GridOutput(tensor=grid_sizes, tuple=grid_tuple)
+        self._grid_cache = (grid_key, grid_output)
+        return grid_output
 
     @torch.no_grad()
     def infer(self, weights, inputs, kv_start=0, kv_end=0):
         del weights, kv_start, kv_end
         if self.scheduler is None:
             raise RuntimeError("MotusPreInfer requires a scheduler before infer().")
+        self._begin_request()
 
         first_frame = inputs["motus_first_frame"]
         state = inputs["motus_state"]
@@ -41,12 +70,7 @@ def infer(self, weights, inputs, kv_start=0, kv_end=0):
             latent_h // self.model.video_backbone.patch_size[1],
             latent_w // self.model.video_backbone.patch_size[2],
         )
-        grid_key = (batch_size, grid_tuple, state.device.type, state.device.index)
-        grid_output = self._grid_cache.get(grid_key)
-        if grid_output is None:
-            grid_sizes = torch.tensor([grid_tuple], dtype=torch.long, device=state.device).expand(batch_size, -1)
-            grid_output = GridOutput(tensor=grid_sizes, tuple=grid_tuple)
-            self._grid_cache[grid_key] = grid_output
+        grid_output = self._get_grid_output(batch_size, grid_tuple, state.device)
 
         if self.cos_sin is None or self.grid_sizes != grid_output.tuple:
             self.grid_sizes = grid_output.tuple

lightx2v/models/networks/motus/infer/transformer_infer.py

@@ -47,14 +47,25 @@ def __init__(self, model, config):
         self.cos_sin = None
         self.rope_positions = None
         self.weights = None
-        self._cu_seqlens_cache = {}
-        self.reset_infer_states()
+        self.clear_request_cache()
+
+    def clear_request_cache(self):
+        self._cu_seqlens_cache_request_id = None
+        self._cu_seqlens_cache = {
+            "joint_attn_cu_seqlens_q": None,
+            "cross_attn_cu_seqlens_q": None,
+            "cross_attn_cu_seqlens_kv": None,
+        }
+
+    def set_scheduler(self, scheduler):
+        super().set_scheduler(scheduler)
+        self.clear_request_cache()
 
-    def reset_infer_states(self):
-        self.joint_attn_cu_seqlens_q = None
-        self.joint_attn_cu_seqlens_kv = None
-        self.cross_attn_cu_seqlens_q = None
-        self.cross_attn_cu_seqlens_kv = None
+    def _begin_request(self):
+        request_id = self.scheduler.rope_request_id
+        if request_id != self._cu_seqlens_cache_request_id:
+            self.clear_request_cache()
+            self._cu_seqlens_cache_request_id = request_id
 
     def _maybe_empty_cache(self):
         if self.clean_cuda_cache:
@@ -70,12 +81,17 @@ def _get_und_block(self, layer_idx):
         return self.weights.und.blocks[layer_idx]
 
     def _get_cu_seqlens(self, cache_name, batch, seq_len, device, attn_type):
+        if cache_name not in self._cu_seqlens_cache:
+            raise ValueError(f"Unsupported Motus attention metadata cache: {cache_name}")
+
         cache_key = (cache_name, batch, seq_len, device.type, device.index, attn_type)
-        cu_seqlens = self._cu_seqlens_cache.get(cache_key)
-        if cu_seqlens is None:
-            tensor = torch.arange(0, (batch + 1) * seq_len, seq_len, dtype=torch.int32)
-            cu_seqlens = tensor.to(device, non_blocking=True) if attn_type in ["flash_attn2", "flash_attn3"] else tensor
-            self._cu_seqlens_cache[cache_key] = cu_seqlens
+        cached = self._cu_seqlens_cache[cache_name]
+        if cached is not None and cached[0] == cache_key:
+            return cached[1]
+
+        tensor = torch.arange(0, (batch + 1) * seq_len, seq_len, dtype=torch.int32)
+        cu_seqlens = tensor.to(device, non_blocking=True) if attn_type in ["flash_attn2", "flash_attn3"] else tensor
+        self._cu_seqlens_cache[cache_name] = (cache_key, cu_seqlens)
         return cu_seqlens
 
     def _normalize_attention_dtype(self, tensor):
@@ -332,7 +348,7 @@ def infer(self, weights, pre_infer_out):
         self.weights = weights
         self.cos_sin = pre_infer_out.cos_sin
         self.rope_positions = pre_infer_out.rope_positions
-        self.reset_infer_states()
+        self._begin_request()
 
         processed_t5_context = pre_infer_out.context
         und_tokens = pre_infer_out.und_tokens.clone()

lightx2v/models/networks/motus/model.py

@@ -233,7 +233,6 @@ def __init__(self, config, device):
         logger.info("[Motus] Loading VLM processor")
         self.vlm_processor = AutoProcessor.from_pretrained(self.config["vlm_path"], trust_remote_code=True)
         self._load_normalization_stats()
-        self._rope_cos_sin_cache = {}
         self._patch_qwen3_vl_rope_index(self.vlm_model)
         logger.info("[Motus] Building Motus backbone helpers")
         self.video_backbone = MotusVideoBackbone(self.config, self.pre_weight.video, self.transformer_weights.video)
@@ -502,31 +501,6 @@ def denormalize_actions(self, actions):
         restored = flat * self.action_range.unsqueeze(0) + self.action_min.unsqueeze(0)
         return restored.reshape(shape)
 
-    def get_wan_freqs(self):
-        return self.pre_infer.freqs
-
-    def get_wan_rotary_cos_sin(self, grid_size):
-        if grid_size in self._rope_cos_sin_cache:
-            return self._rope_cos_sin_cache[grid_size]
-        freqs = self.get_wan_freqs()
-        head_dim_half = freqs.shape[1]
-        c_f = head_dim_half - 2 * (head_dim_half // 3)
-        c_h = head_dim_half // 3
-        c_w = head_dim_half // 3
-        fpart, hpart, wpart = freqs.split([c_f, c_h, c_w], dim=1)
-        f, h, w = grid_size
-        freq_grid = torch.cat(
-            [
-                fpart[:f].view(f, 1, 1, -1).expand(f, h, w, -1),
-                hpart[:h].view(1, h, 1, -1).expand(f, h, w, -1),
-                wpart[:w].view(1, 1, w, -1).expand(f, h, w, -1),
-            ],
-            dim=-1,
-        ).reshape(f * h * w, -1)
-        cos_sin = (freq_grid.real.contiguous(), freq_grid.imag.contiguous())
-        self._rope_cos_sin_cache[grid_size] = cos_sin
-        return cos_sin
-
     def prepare_frame(self, image_path):
         image = Image.open(image_path).convert("RGB")
         image_np = np.asarray(image).astype(np.float32) / 255.0

lightx2v/models/networks/wan/dreamzero_model.py

@@ -29,9 +29,10 @@ def _init_infer_class(self):
     def cfg_cache_name(cache_name, infer_condition):
         return f"{cache_name}_{'cond' if infer_condition else 'uncond'}"
 
-    def clear_cache(self, cache_name=None):
+    def clear_cache(self, cache_name=None, clear_pre_infer=True):
         self.transformer_infer.clear_cache(cache_name)
-        self.pre_infer.clear_cache()
+        if clear_pre_infer:
+            self.pre_infer.clear_cache()
 
     @staticmethod
     def _gather_cfg_tensor(tensor, group):

lightx2v/models/networks/wan/infer/causvid/pre_infer.py

@@ -12,12 +12,36 @@ class WanCausVidPreInfer(WanPreInfer):
 
     def __init__(self, config):
         super().__init__(config)
-        self._causvid_start_frame = 0
+        self.scheduler = None
         self._causvid_rope_cache: Dict[tuple, Tuple[torch.Tensor, torch.Tensor | None]] = {}
+        self._reset_request_cache()
+
+    def set_scheduler(self, scheduler):
+        super().set_scheduler(scheduler)
+        self._reset_request_cache()
+        self._rope_cache_request_id = scheduler.rope_request_id
 
     def set_rope(self, rope):
         super().set_rope(rope)
+        self._reset_request_cache()
+        if self.scheduler is not None:
+            self._rope_cache_request_id = self.scheduler.rope_request_id
+
+    def _reset_request_cache(self):
+        self._rope_cache_request_id = None
         self._causvid_rope_cache.clear()
+        self._causvid_start_frame = 0
+        self.cos_sin = None
+        self.rope_positions = None
+        self.grid_sizes = (0, 0, 0)
+
+    def _sync_request_cache(self):
+        request_id = self.scheduler.rope_request_id
+        if request_id == self._rope_cache_request_id:
+            return
+
+        self._reset_request_cache()
+        self._rope_cache_request_id = request_id
 
     def _rope_cache_key(self, grid_sizes, start_frame):
         device = self.freqs.device
@@ -44,6 +68,10 @@ def infer(self, weights, inputs, kv_start=0, kv_end=0):
             raise NotImplementedError("Sequence parallel inference is not implemented for CausVid.")
         if kv_start < 0 or kv_end <= kv_start:
             raise ValueError(f"Invalid CausVid KV range: [{kv_start}, {kv_end}).")
+        if self.scheduler is None:
+            raise RuntimeError("WanCausVidPreInfer scheduler is not initialized.")
+
+        self._sync_request_cache()
 
         # The previous local grid is the most accurate source of the number of
         # spatial tokens. On the first chunk, fall back to the configured value;