[rollout, diffusion] fix: support prompt embedding cache for Qwen-Image

tests/workers/config/test_diffusion_config_on_cpu.py

@@ -168,6 +168,19 @@ def test_defaults(self):
 
 
 class TestDiffusionRolloutConfig:
+    def test_prompt_embed_cache_config(self):
+        cfg = DiffusionRolloutConfig(
+            name="vllm_omni",
+            enable_prompt_embed_cache=True,
+            prompt_embed_cache_size=64,
+        )
+        assert cfg.enable_prompt_embed_cache is True
+        assert cfg.prompt_embed_cache_size == 64
+
+    def test_invalid_prompt_embed_cache_size_raises(self):
+        with pytest.raises(ValueError, match="prompt_embed_cache_size must be positive"):
+            DiffusionRolloutConfig(name="vllm_omni", prompt_embed_cache_size=0)
+
     def test_invalid_rollout_adapter_raises(self):
         with pytest.raises(ValueError, match="Invalid diffusion rollout rollout_adapter"):
             DiffusionRolloutConfig(name="vllm_omni", rollout_adapter="bogus")

verl_omni/experimental/qwen_image_flow_grpo_stepwise/vllm_omni_rollout_adapter.py

@@ -117,6 +117,10 @@ def encode_prompt(
                 ``(B * num_images_per_prompt, L, D)`` and
                 ``(B * num_images_per_prompt, L)`` respectively.
         """
+        if isinstance(prompt_ids, list):
+            prompt_ids = torch.tensor(prompt_ids, device=self.device)
+        if isinstance(attention_mask, list):
+            attention_mask = torch.tensor(attention_mask, device=self.device)
         prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
         attention_mask = (
             attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask
@@ -173,7 +177,7 @@ def _tokenize_text_prompt(self, text: str | list[str]):
             truncation=True,
             return_tensors="pt",
         ).to(self.device)
-        return tokens.input_ids, tokens.attention_mask
+        return tokens.input_ids.tolist(), tokens.attention_mask.tolist()
 
     def prepare_encode(
         self,
@@ -191,12 +195,6 @@ def prepare_encode(
             state.prompts or []
         )
 
-        # Normalize list inputs to tensors on device.
-        if isinstance(prompt_ids, list):
-            prompt_ids = torch.tensor(prompt_ids, device=self.device)
-        if isinstance(negative_prompt_ids, list):
-            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
-
         if prompt_ids is None:
             raise ValueError(
                 "QwenImagePipelineWithLogProbStepwise.prepare_encode requires either "
@@ -221,10 +219,18 @@ def prepare_encode(
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt_ids is not None:
+        prompt_embed_cache = getattr(self, "_prompt_embed_cache", None)
+        prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled)
+        if not prompt_embed_cache_enabled:
+            if isinstance(prompt_ids, list):
+                prompt_ids = torch.tensor(prompt_ids, device=self.device)
+            if isinstance(negative_prompt_ids, list):
+                negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
+
+        if isinstance(prompt_ids, torch.Tensor):
             batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
         else:
-            batch_size = 1
+            batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1
 
         has_neg_prompt = negative_prompt_ids is not None
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt

verl_omni/pipelines/qwen_image_diffusion_nft/vllm_omni_rollout_adapter.py

@@ -72,13 +72,19 @@ def _prepare_token_id_generation_context(
         self._current_timestep = None
         self._interrupt = False
 
-        if isinstance(prompt_ids, list):
-            prompt_ids = torch.tensor(prompt_ids, device=self.device)
-        if isinstance(negative_prompt_ids, list):
-            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
+        prompt_embed_cache = getattr(self, "_prompt_embed_cache", None)
+        prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled)
+        if not prompt_embed_cache_enabled:
+            if isinstance(prompt_ids, list):
+                prompt_ids = torch.tensor(prompt_ids, device=self.device)
+            if isinstance(negative_prompt_ids, list):
+                negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
 
         if prompt_ids is not None:
-            batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+            if isinstance(prompt_ids, torch.Tensor):
+                batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+            else:
+                batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1
         elif prompt_embeds is not None:
             batch_size = prompt_embeds.shape[0]
         else:

verl_omni/pipelines/qwen_image_dpo/vllm_omni_rollout_adapter.py

@@ -82,6 +82,10 @@ def encode_prompt(
         prompt_embeds_mask: torch.Tensor | None = None,
         max_sequence_length: int = 1024,
     ):
+        if isinstance(prompt_ids, list):
+            prompt_ids = torch.tensor(prompt_ids, device=self.device)
+        if isinstance(attention_mask, list):
+            attention_mask = torch.tensor(attention_mask, device=self.device)
         prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
         attention_mask = (
             attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask
@@ -150,16 +154,21 @@ def forward(
         self._current_timestep = None
         self._interrupt = False
 
+        prompt_embed_cache = getattr(self, "_prompt_embed_cache", None)
+        prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled)
         if prompt_ids is not None:
-            if isinstance(prompt_ids, list):
+            if not prompt_embed_cache_enabled and isinstance(prompt_ids, list):
                 prompt_ids = torch.tensor(prompt_ids, device=self.device)
-            batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+            if isinstance(prompt_ids, torch.Tensor):
+                batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+            else:
+                batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1
         elif prompt_embeds is not None:
             batch_size = prompt_embeds.shape[0]
         else:
             return DiffusionOutput(output=None, custom_output={})
 
-        if isinstance(negative_prompt_ids, list):
+        if not prompt_embed_cache_enabled and isinstance(negative_prompt_ids, list):
             negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
 
         has_neg_prompt = negative_prompt_ids is not None or (

verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py

@@ -307,18 +307,23 @@ def forward(
         self._current_timestep = None
         self._interrupt = False
 
+        prompt_embed_cache = getattr(self, "_prompt_embed_cache", None)
+        prompt_embed_cache_enabled = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled)
         if prompt_token_ids is not None:
-            if isinstance(prompt_token_ids, list):
+            if not prompt_embed_cache_enabled and isinstance(prompt_token_ids, list):
                 prompt_token_ids = torch.tensor(prompt_token_ids, device=self.device)
-            batch_size = prompt_token_ids.shape[0] if prompt_token_ids.ndim == 2 else 1
+            if isinstance(prompt_token_ids, torch.Tensor):
+                batch_size = prompt_token_ids.shape[0] if prompt_token_ids.ndim == 2 else 1
+            else:
+                batch_size = len(prompt_token_ids) if prompt_token_ids and isinstance(prompt_token_ids[0], list) else 1
         elif prompt_embeds is not None:
             batch_size = prompt_embeds.shape[0]
         else:
             # Both prompt_token_ids and prompt_embeds are None (e.g. during warmup/dummy run).
             # Return a minimal dummy output to avoid crashing.
             return DiffusionOutput(output=None, custom_output={})
 
-        if isinstance(negative_prompt_ids, list):
+        if not prompt_embed_cache_enabled and isinstance(negative_prompt_ids, list):
             negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
 
         has_neg_prompt = negative_prompt_ids is not None or (

verl_omni/trainer/config/_generated_diffusion_trainer.yaml

@@ -207,6 +207,8 @@ actor_rollout_ref:
       enable: false
     engine_kwargs:
       vllm_omni: {}
+    enable_prompt_embed_cache: false
+    prompt_embed_cache_size: 32
     val_kwargs:
       _target_: verl_omni.workers.config.diffusion.DiffusionSamplingConfig
       'n': 1

verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml

@@ -248,6 +248,8 @@ actor_rollout_ref:
       enable: false
     engine_kwargs:
       vllm_omni: {}
+    enable_prompt_embed_cache: false
+    prompt_embed_cache_size: 32
     val_kwargs:
       _target_: verl_omni.workers.config.diffusion.DiffusionSamplingConfig
       'n': 1

verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml

@@ -122,6 +122,12 @@ engine_kwargs:
   # vllm-omni engine config
   vllm_omni: {}
 
+# Reuse text-encoder outputs for repeated diffusion prompts.
+enable_prompt_embed_cache: false
+
+# Maximum number of prompt embeddings cached per diffusion worker.
+prompt_embed_cache_size: 32
+
 # Sampling parameters used during validation (diffusion-specific).
 val_kwargs:
 

verl_omni/workers/config/diffusion/rollout.py

@@ -120,6 +120,9 @@ class DiffusionRolloutConfig(BaseConfig):
     # the *_stepwise variant of the pipeline (e.g. flow_grpo_stepwise).
     step_execution: bool = False
 
+    enable_prompt_embed_cache: bool = False
+    prompt_embed_cache_size: int = 32
+
     # note that the logprob computation should belong to the actor
     log_prob_micro_batch_size_per_gpu: Optional[int] = None
 
@@ -179,6 +182,8 @@ def __post_init__(self):
             raise ValueError(
                 f"Invalid diffusion rollout rollout_adapter: {self.rollout_adapter}. Must be one of ['default', 'old']."
             )
+        if self.prompt_embed_cache_size <= 0:
+            raise ValueError(f"prompt_embed_cache_size must be positive, got {self.prompt_embed_cache_size}.")
 
         if self.pipeline_model_parallel_size > 1:
             if self.name == "vllm_omni":

verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py

@@ -170,6 +170,9 @@ async def run_server(self, args: argparse.Namespace):
                 engine_args["enable_dummy_pipeline"] = True
                 engine_args["custom_pipeline_args"] = {"pipeline_class": pipeline_path}
 
+            engine_args["enable_prompt_embed_cache"] = self.config.enable_prompt_embed_cache
+            engine_args["prompt_embed_cache_size"] = self.config.prompt_embed_cache_size
+
         if getattr(self.config, "step_execution", False):
             engine_args["step_execution"] = True
 
@@ -360,7 +363,9 @@ def _preprocess_input(
 
         custom_prompt: OmniCustomPrompt = {"prompt_token_ids": prompt_ids}
         if prompt_mask is not None:
-            custom_prompt["prompt_mask"] = prompt_mask
+            custom_prompt["prompt_mask"] = (
+                prompt_mask.tolist() if isinstance(prompt_mask, torch.Tensor) else prompt_mask
+            )
         if len(default_params_list) > 1:
             # Multi-stage pipelines tag the diffusion stage so the orchestrator can route inputs correctly.
             custom_prompt["modalities"] = ["image"]