[rollout, diffusion] fix: support prompt embedding cache for tokenized prompts

[Sky-Trigger]

Keep token IDs and attention masks hashable until encode_prompt(), convert them to tensors inside token-ID encoders, and expose vLLM-Omni prompt embedding cache controls through the diffusion rollout configuration.

What does this PR do?

This PR adapts verl-omni's tokenized diffusion prompt path to the prompt embedding cache introduced by vLLM-Omni PR #2962.

Previously, rollout adapters converted token ID lists and masks to torch.Tensor before calling encode_prompt(). Since the vLLM-Omni cache cannot build keys from tensor arguments, these calls bypassed the cache.

This change keeps prompt IDs and masks list-based across the cache boundary and converts them to tensors only inside the underlying token-ID encoder. It covers:

• Qwen-Image FlowGRPO and MixGRPO
• Qwen-Image DiffusionNFT
• Qwen-Image online DPO
• Qwen-Image stepwise execution
• Wan2.2 DanceGRPO

It also adds rollout configuration for enabling and sizing the cache without requiring environment variables.

Checklist Before Starting

• Searched for similar PRs: open PR search
• Formatted the PR title as:
  [rollout, diffusion] fix: support prompt embedding cache for tokenized prompts

Test

Static checks completed:

• ruff check
• ruff format --check
• Python compileall
• git diff --check

Manual validation:

• examples/diffusionnft_trainer/run_qwen_image_ocr_lora.sh completed successfully with the updated DiffusionNFT adapter.

Configuration unit tests were added to tests/workers/config/test_diffusion_config_on_cpu.py.

Full pytest and all adapter-specific GPU/NPU experiments were not run in the current environment. The remaining independent rollout paths are FlowGRPO, online DPO, and Wan2.2 DanceGRPO.

API and Usage Example

The cache is disabled by default and uses the upstream default capacity of 32.

It can be enabled through Hydra configuration:

python3 -m verl_omni.trainer.main_diffusion \
    actor_rollout_ref.rollout.enable_prompt_embed_cache=True \
    actor_rollout_ref.rollout.prompt_embed_cache_size=64 \
    ...

Equivalent YAML:

actor_rollout_ref:
  rollout:
    enable_prompt_embed_cache: true
    prompt_embed_cache_size: 64

No example shell scripts are modified.

Design & Code Changes

Prompt cache boundary

Before:

adapter
  -> list to torch.Tensor
  -> cached encode_prompt()
  -> tensor argument is not hashable
  -> cache bypass

After:

adapter
  -> list-based prompt IDs and masks
  -> cached encode_prompt()
  -> cache lookup
  -> on miss, token-ID encoder converts lists to tensors
  -> text encoder

Adapter changes

• Added list-compatible batch-size calculation.
• Moved list-to-tensor conversion into Qwen-Image and Wan token-ID encoders.
• Kept positive and negative prompt IDs hashable until encode_prompt().
• Normalized rollout prompt_mask to a list before constructing the diffusion request.
• Preserved the existing DPO and stepwise class structures and inheritance relationships.
• Left SD3 and BAGEL unchanged because they do not use this token-ID cache path.

Configuration changes

Added to DiffusionRolloutConfig:

enable_prompt_embed_cache: bool = False
prompt_embed_cache_size: int = 32

These values are passed directly to AsyncOmni. Cache size must be a positive integer.

Checklist Before Submitting

• Read the [Contribute Guide](https://github.com/verl-project/verl-omni/blob/main/CONTRIBUTING.md).
• Run full pre-commit checks:
  pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update documentation.
  • No standalone documentation was added; the new configuration fields are documented in the rollout YAML.
• Added configuration unit tests to the existing CPU config test suite.
  • Adapter behavior was validated through DiffusionNFT end-to-end execution.
  • Additional GPU/NPU coverage is not added to CI because it requires model weights and accelerator resources.

AI Assistance

AI assistance was used to inspect the adapter call paths, implement the changes, and prepare this PR description. All changed lines and validation results were reviewed by the human submitter.

[gemini-code-assist]

Code Review

This pull request introduces a prompt embedding cache mechanism for diffusion workers by adding enable_prompt_embed_cache and prompt_embed_cache_size configurations across trainer and rollout configs. It also refactors several rollout adapters (such as Qwen and Wan22) to support list-based inputs for prompt IDs and masks, standardizing batch size calculation via a new helper function get_prompt_batch_size. There are no review comments provided, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[zhtmike]

@knlnguyen1802 , prompt cache related, PTAL

[knlnguyen1802]

@Sky-Trigger Could you help to run a verification and post a curve

[Sky-Trigger]

I have verified this change with DiffusionNFT, FlowGRPO, and related models. This modification only delays the point at which prompt_ids and other inputs are converted into tensors, so that the code path is compatible with vLLM-Omni PR #2962. It does not introduce any numerical changes, precision-related behavior, or new model structure.

[knlnguyen1802]

@Sky-Trigger Thanks for verify it, could you run several step to compare performance gain on this branch ?