Skip to content

[RFC] Split vLLM-Omni generation strategies and define typed multimodal rollout I/O #403

Description

@NancyFyong

[RFC] Split vLLM-Omni generation strategies and define typed multimodal rollout I/O

  • Status: Implementation delivered across the six-PR stack; strategy extraction and initial
    declarations are merged. The remaining request/consumer/wire/named-artifact changes are in review
    (PR5/PR6 remain drafts pending human review and upstream CI).
  • Scope:
    1. extract the AR and diffusion generation/configuration paths from
      vllm_omni_async_server.py behind an internal strategy interface; and
    2. replace diffusion's ad-hoc image/video/audio request and output conventions with a
      repo-owned typed contract.
  • Compatibility rule: the first implementation PR is behavior-preserving. Existing Ray RPC
    callers and existing TokenOutput consumers continue to work while the typed diffusion path is
    migrated incrementally.
  • Not in scope: changing sleep/wake/abort/drain semantics, deleting vLLM-Omni lifecycle shims,
    or changing the upstream vLLM-Omni wire protocol.
  • Last updated: 2026-09-07.
  • Implementation status:
    • #444 (strategy extraction) — merged;
    • #478 (adapter-owned diffusion media
      declarations) — merged;
    • #480 (typed internal request flow and
      ImageGenerationRequest removal and fail-closed input validation) — open / non-draft / under review;
    • #481 (consume declared media kind across
      V0/V1 trainer media paths, preserve generated media at reward entry, and make dumps best-effort) — draft / under review;
    • #556 (canonical private request wire schema,
      all-adapter token readers and stage-type-based AR entrance bridge) — draft / implemented;
    • #557 (all-adapter named media, native/decoded
      layout separation, explicit consumers and fail-fast schema validation) — draft / implemented; and
    • #508 (contributor documentation for
      DiffusionIOSpec) — merged.
  • Original design baseline: the initial investigation was performed against origin/main
    d8776e5 and the pinned vLLM-Omni environment available on 2026-08-22. The milestones below
    record the implementation as it evolved; they supersede that original schedule.

0. TL;DR

This RFC is being delivered as small, behavior-contained changes rather than one replacement of
vLLM-Omni's protocol.

  • Merged: vLLMOmniHttpServer now selects ARStrategy or DiffusionStrategy once during
    _init_config; the shared server retains lifecycle, engine ownership, LoRA cache, and replica
    behavior. Normal operation no longer branches on self._ar_mode.
  • Merged: every registered diffusion adapter declares a diffusion_io_spec using
    MediaSpec / DiffusionIOSpec. The shared diffusion strategy no longer supplies a hard-coded
    32 kHz audio default; MiniMax H3, LTX-2, and all image/video adapters own their declared media
    facts.
  • In review: OmniRolloutRequest gives both strategies one typed, server-internal request
    object while preserving the existing Ray generate(**kwargs) RPC and the pinned engine's legacy
    prompt keys. The obsolete ImageGenerationRequest abstraction is removed because its only two
    consumers needed the same centralized condition-image lookup.
  • In review: the diffusion strategy forwards declared media metadata through both V0 and V1
    trainer paths. V0/V1 dump and W&B paths use media_kind for image/video selection; V1 preserves
    metadata through TransferQueue, reuses the shared MP4/audio exporter, and treats background dump
    failures as best-effort observability rather than training failures.

PR5 #556 and PR6 #557 now implement wire alignment and the canonical named-media boundary for
all registered in-tree diffusion adapters. Serving media consumers no longer infer modality,
layout or representation from tensor rank/channel counts or tuple position. These changes are
on the draft stack, not yet merged into main. Public keyword RPC and explicit compatibility
projections remain; they do not add a second interpretation of the media.

                         shared lifecycle / transport
                                   |
                       vLLMOmniHttpServer
                                   |
                +------------------+------------------+
                |                                     |
           ARStrategy                         DiffusionStrategy
                |                                     |
          TokenOutput                  OmniRolloutRequest (in review)
                                              + DiffusionIOSpec (merged)
                                              + declared media consumers (in review)

The lifecycle findings in #377 remain owned by #377. This RFC adopts only its strategy extraction;
it does not duplicate or silently change sleep, wake, abort, drain, cache, profiling, or temporary
config behavior.

1. Relationship to #377 and #391

1.1 #377: same server, split ownership

#377 identifies two classes of work:

  1. separate AR and diffusion generation behavior; and
  2. repair lifecycle divergence from upstream verl (sleep, wake_up, abort, drain, cache
    management, profiling, and temporary patches).

This RFC takes ownership of the first item because the typed diffusion contract needs a stable
place to live. It does not take ownership of the second item.

Concern Owner after this RFC
mode selection vLLMOmniHttpServer factory, once during initialization
AR config/input/generation/output ARStrategy
diffusion config/input/generation/output DiffusionStrategy
model-specific diffusion media declaration pipeline adapter DiffusionIOSpec
shared LoRA cache and request-independent engine state vLLMOmniHttpServer
sleep/wake/abort/drain/cache/profile semantics #377
removal of temp YAML/env/internal-queue patches #377 or dedicated follow-ups

Before implementation, the strategy item in #377 and this RFC should be cross-linked so there is
one implementation PR, not two competing refactors.

1.2 #391: engine-side protocol, not duplicated here

The installed vLLM-Omni already has a canonical diffusion post-process envelope:

{
    "payload": {"image" | "video" | "audio" | "text": value},
    "metadata": {
        "video": {"fps": ...},
        "audio": {"sample_rate": ...},
        ...,
    },
}

It also exposes DiffusionOutput.output, legacy trajectory_* fields, and
OmniRequestOutput.multimodal_output. #391 proposes further engine-side stability for scheduler
injection, trajectories, configurable outputs, and pre-tokenized input.

This RFC does not fork those types. DiffusionStrategy is the one translation boundary:

verl-omni typed request
    -> pinned vLLM-Omni prompt/sampling types
    -> OmniRequestOutput / payload+metadata
    -> verl-omni typed diffusion output

When #391 lands, only that translation layer changes. Agent loops, rewards, tracking, and trainers
remain insulated from the pinned engine representation.


2. Original problems (2026-08-22 investigation baseline)

This section records the motivating baseline, not the current main branch. Sections 0 and 6
track what has since landed or is in review.

2.1 One server contains two generation implementations

The 962-line vllm_omni_async_server.py branches on self._ar_mode in initialization, validation,
post-init, generation config, worker selection, engine kwargs, engine construction, request
preprocessing, engine invocation, and output processing.

The two paths use different concepts:

Concern AR path Diffusion path
model config OmniModelConfig DiffusionModelConfig
sampling vLLM SamplingParams OmniDiffusionSamplingParams list
prompt token prompt + multimodal data OmniCustomPrompt plus diffusion extras
output TokenOutput local DiffusionOutput
postprocess token IDs/logprobs pixels/latents/trajectory/media metadata
pipeline setup multi-stage deploy config registered custom diffusion pipeline

A change to one side therefore requires editing and testing a class that also owns the other side
and all lifecycle behavior.

2.2 Model-specific media knowledge is accumulating in the server

The diffusion output branch currently:

  • searches dict keys in the order video, image, output, audio;
  • interprets any tuple/list as (visual, audio);
  • writes audio into extra_fields;
  • defaults its sample rate to 32 kHz;
  • converts the remaining primary tensor as pixels or latents; and
  • flattens multimodal_output.metadata.prompt_embeddings and .rl into one dictionary.

Those are not generic HTTP-server responsibilities. The tuple rule and 32 kHz default are valid for
MiniMax H3, while LTX-2 derives its sample rate from its vocoder. Encoding those facts in the
server makes the next multi-output model another server patch.

A success criterion for this RFC is:

Adding a pipeline with a new combination of existing modalities changes its adapter declaration
and tests, not vllm_omni_async_server.py.

2.3 The request boundary is still a growing keyword list

The current path relays prompt_ids, sampling_params, image_data, video_data, audio_data,
negative_prompt_ids, masks, per-encoder token IDs, processor kwargs, and priority through the
agent loop, LLMServerClient, retry client, and server.

ImageGenerationRequest then scans five aliases for one condition image:

  • images;
  • image;
  • multi_modal_data.image;
  • extra_args.multi_modal_data.image; and
  • additional_information.condition_images.

The server currently writes multi_modal_data both at the prompt top level and under
extra_args. Other adapters contain copy-up and fallback lookups. Adding a second input with a
role or a second view of the same image has no stable representation.

The mismatch is visible in the installed vLLM-Omni version too: OmniCustomPrompt declares
prompt_ids, masks, negative IDs, and extra_args, while verl-omni currently also writes
prompt_token_ids, modalities, extra_prompt_ids, negative_extra_prompt_ids, and
multi_modal_data. These extensions work dynamically, but the contract is not represented by the
imported type.

2.4 The output boundary loses modality

The local DiffusionOutput is currently:

class DiffusionOutput(BaseModel):
    diffusion_output: Any
    log_probs: Any | None = None
    stop_reason: str | None = None
    num_preempted: int | None = None
    extra_fields: dict[str, Any] = {}

Consequences in the current tree include:

  • tracking interprets a rank-4 tensor as video;
  • the HTTP scorer interprets a rank-4 tensor as a batch and keeps the first item;
  • the trainer uses rank 5 to identify batched video;
  • audio and its sample rate travel separately through extra_fields / tool_extra_fields; and
  • tensor-valued extras and Python-valued extras are routed into different DataProto containers by
    runtime type.

Rank may still be used to validate or normalize layout, but it should not answer “which modality is
this?” once the producer has declared that fact.

2.5 The upstream envelope exists, but the local boundary does not preserve it

In the verified vLLM-Omni environment:

  • DiffusionOutput.output accepts a scalar/tensor, tuple, dict, or canonical payload/metadata
    envelope;
  • normalize_diffusion_postprocess_output() recognizes that envelope;
  • known video/audio metadata is validated;
  • OmniRequestOutput exposes primary output through images and side data through
    multimodal_output;
  • final_output_type can represent text/image/audio/video, although not every formatter path
    currently sets video explicitly.

verl-omni then collapses this back into one Any plus extra_fields. The new local interface should
adapt and retain the useful structure, not invent a second engine protocol.


3. Goals and non-goals

Goals

  • G1 — One mode dispatch. Select an AR or diffusion strategy once. Normal initialization and
    generation contain no repeated self._ar_mode branching.
  • G2 — Thin server. vLLMOmniHttpServer owns shared transport, engine lifetime, shared LoRA
    state, and lifecycle hooks; mode-specific config/input/output behavior lives elsewhere.
  • G3 — One internal request envelope. Positive/negative/per-encoder tokens and image/video/audio
    conditions are assembled once inside the server strategy template. M3 keeps the public keyword
    RPC; typed agent-loop/client transport is not delivered or required by this six-PR plan. Any
    future public typed transport needs a separately agreed migration.
  • G4 — Declared diffusion output. M2/M4 declare and preserve the current primary/optional-audio
    contract. M6 replaces it with named artifacts: each declares modality, representation and layout,
    so latent reward/training data and decoded logging previews can coexist without substitution.
  • G5 — Adapter-owned model facts. Prompt rendering, required input views, primary output kind,
    layout, and required metadata are declared next to the pipeline adapter.
  • G6 — One engine compatibility boundary. Only the diffusion strategy knows how the pinned
    vLLM-Omni prompt and output objects are shaped.
  • G7 — Incremental migration. Preserve the existing generate(**kwargs) RPC and legacy
    diffusion_output / extra_fields view until all consumers migrate.
  • G8 — CPU-testable. Contracts and strategy processing are testable with fake engines and do
    not require model weights or a GPU.

Non-goals

  • N1. Fixing or changing sleep, wake, abort, drain, cache reset, profiling, DP, or headless
    semantics. Those remain under [RFC][tracking] Refactor vllm_omni_async_server.py — Align with Upstream verl, Remove Temp Patches #377.
  • N2. Changing vLLM-Omni's OmniCustomPrompt, DiffusionOutput, output formatter, or IPC
    protocol. [RFC] Stable vLLM-Omni Rollout Interface for Diffusion RL #391 owns engine-side changes.
  • N3. Replacing TokenOutput or redesigning AR agent-loop output.
  • N4. Typing every sampling parameter. The request carries the current mapping and each
    strategy lowers it to its engine-specific type.
  • N5. Creating a third-party strategy plugin framework. There are exactly two private
    strategies in this RFC: AR and diffusion.
  • N6. Rewriting every pipeline in one PR.
  • N7. Changing training math, reward formulas, or tensor values while extracting the strategy.

4. Proposed architecture

4.1 Internal generation strategy

The strategy seam is implementation structure, not a third-party plugin API. The landed classes use public names for readability, but there are exactly two in-tree strategies:

class OmniStrategyBase(ABC):
    def init_config(self, config: Any) -> Any: ...
    def init_model_config(self, model_config: Any) -> Any: ...
    def validate_configs(self) -> None: ...
    def post_init(self, cuda_visible_devices: str) -> None: ...
    def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: ...
    async def generate(self, ...) -> TokenOutput | DiffusionOutput: ...

Selection happens once in _init_config, before the superclass invokes its remaining initialization
hooks. The server still delegates the existing keyword RPC; PR #480 constructs
OmniRolloutRequest inside the shared strategy template:

class vLLMOmniHttpServer(vLLMHttpServer):
    def _init_config(self, config):
        omni_kwargs = (getattr(config, "engine_kwargs", {}) or {}).get("vllm_omni", {}) or {}
        strategy_cls = ARStrategy if omni_kwargs.get("output_mode", "diffusion") == "ar" else DiffusionStrategy
        self._generate_strategy = strategy_cls(self)
        return self._generate_strategy.init_config(config)

    async def generate(self, **legacy_kwargs):
        return await self._generate_strategy.generate(**legacy_kwargs)

The existing RPC signature remains during migration; it becomes a compatibility adapter instead of
being copied through every layer.

Suggested file ownership:

workers/rollout/vllm_rollout/
├── vllm_omni_async_server.py       # shared engine/server/lifecycle + replica
├── vllm_omni_ar_strategy.py        # AR config, prompt, invocation, TokenOutput
└── vllm_omni_diffusion_strategy.py # diffusion config, prompt, invocation, media output

pipelines/
├── rollout_media.py              # MediaSpec / DiffusionIOSpec (merged)
└── rollout_request.py            # OmniRolloutRequest (PR #480)

No registry is needed for strategies. Existing pipeline registries remain the extension point for
model-specific behavior.

4.2 Responsibility split

ARStrategy

Moves, without semantic changes:

  • OmniModelConfig conversion and max_model_len handling;
  • multi-stage deploy config and rollout flags;
  • AR engine-argument cleanup and logprob mode;
  • multimodal pad-token deduplication;
  • vLLM SamplingParams construction;
  • AR engine.generate() arguments; and
  • token IDs, token logprobs, stop reason, and TokenOutput construction.

DiffusionStrategy

Moves, without semantic changes in the first PR:

  • DiffusionModelConfig conversion and PIL/tensor conversion;
  • external pipeline loading and request-batching capability checks;
  • OmniDiffusionSamplingParams construction;
  • lowering the repo-owned request to the currently pinned OmniCustomPrompt shape;
  • diffusion engine.generate() invocation;
  • image/video/audio/text/latent normalization;
  • trajectory and rollout metadata extraction; and
  • typed local DiffusionOutput construction.

vLLMOmniHttpServer

Keeps:

  • AsyncOmni construction and app startup;
  • shared request admission and global-step state;
  • shared LoRA cache ownership;
  • sleep/wake/abort/drain hooks, unchanged by this RFC; and
  • Ray replica integration.

Mode-specific engine-argument mutation is delegated, while the actual engine lifecycle remains in
the server.

4.3 Repo-owned internal request envelope (PR3)

The implemented scope is server-internal, not a new Ray transport object:

@dataclass(frozen=True)
class PromptBundle:
    token_ids: list[int]
    mask: Any | None = None
    negative_token_ids: list[int] | None = None
    extra_token_ids: Mapping[str, list[int]] | None = None
    negative_extra_token_ids: Mapping[str, list[int]] | None = None
    mm_processor_kwargs: Mapping[str, Any] | None = None

@dataclass(frozen=True)
class MediaInput:
    modality: Literal["image", "video", "audio"]
    data: Any

@dataclass(frozen=True)
class OmniRolloutRequest:
    prompt: PromptBundle
    media: tuple[MediaInput, ...] = ()

    @classmethod
    def from_generate_kwargs(cls, **kwargs) -> "OmniRolloutRequest": ...

The public request ID, sampling arguments and priority remain on the existing RPC. Each strategy
lowers this internal object to its own prompt type; AR remains token-centric and diffusion remains
media-centric. They do not acquire one interchangeable output contract.

PR3 rejects duplicate modalities, malformed/conflicting condition-image aliases, unsupported AR
prompt fields and nonzero diffusion priority. Ref2VA processor metadata is forwarded, not rejected.
Equivalent NumPy/Torch aliases are accepted only with exact shape/value equality, never broadcasting.
Empty media lists remain distinct from absent (None) media.

ImageGenerationRequest is deleted; its two former consumers use the shared image parser directly.
The private prompt_token_ids keys and dual multi_modal_data placement remain until the all-adapter
M5 migration. No typed public client transport or canonical wire alignment is claimed by PR3.

4.4 Adapter-owned diffusion I/O declaration (merged PR2, hardened in PR4)

Current declarations are deliberately small:

@dataclass(frozen=True)
class MediaSpec:
    modality: Literal["image", "video", "audio"]
    sample_rate: int | None = None
    fps: float | None = None

@dataclass(frozen=True)
class DiffusionIOSpec:
    primary: MediaSpec
    auxiliary: tuple[MediaSpec, ...] = ()

# Joint video/audio example; image pipelines declare only MediaSpec("image").
io_spec = DiffusionIOSpec(
    primary=MediaSpec("video"),
    auxiliary=(MediaSpec("audio", sample_rate=32000),),
)

The current transport supports a single primary output or (visual, audio) with audio at position

  1. It does not support arbitrary auxiliary[i] streams. PR4 rejects unsupported declarations and
    extra/undeclared tuple outputs instead of silently treating a preview as audio. Earlier wording in
    the merged guide overstated that capability; PR4 corrects it.

Runtime audio sample rate takes precedence over an adapter default (MiniMax H3 32000, LTX-2 24000).
MediaSpec.fps is currently vocabulary only; exporters still use trainer.video_fps. Layout is not
part of PR4. Future prompt rendering/views should be introduced only with an actual consumer.

4.5 Named typed diffusion artifacts (PR6 target, not implemented)

PR6 evolves the local output instead of treating latent as another modality. Its design must support
one rollout returning model-native latents for reward/training and a decoded preview for logging.
The following is an illustrative target, not an already available API or a finalized class name:

@dataclass(frozen=True)
class MediaArtifact:
    modality: Literal["image", "video", "audio"]
    representation: Literal["latent", "decoded"]
    data: Any
    layout: str
    metadata: Mapping[str, Any] = field(default_factory=dict)

class DiffusionOutput(BaseModel):
    artifacts: dict[str, MediaArtifact]  # declared names, not tuple positions
    primary_artifact: str | None        # compatibility projection, None for abort/error
    trajectory: Any | None = None
    training: dict[str, Any] = Field(default_factory=dict)
    stop_reason: str | None = None
    num_preempted: int | None = None

Example artifact names are video_latent, video_preview, and audio. Each artifact independently
declares its modality, representation, layout and metadata. Consumers request/select the artifact
they need; a missing decoded preview must never be replaced by latent data. Request-side
requested_outputs can avoid decoding/transferring a preview when it is not needed.

Required rules:

  1. data.dtype is authoritative; do not add a separately configurable dtype field. Latents retain
    their actual floating dtype. Decoded image/video is uint8 [0,255]; decoded audio is a floating
    waveform. Pixel versus latent is a runtime representation, not a static pipeline modality.
  2. Preserve model-native latent layouts for training/replay, including packed latents. The decoded
    vocabulary CHW, TCHW, CTHW, THWC, T, CT is not an exhaustive latent-layout catalog.
  3. Reject declared-but-missing, returned-but-undeclared, duplicate and requested-but-absent artifacts.
    Validate modality, representation, layout, shape, dtype and required sample rate with pipeline,
    artifact name, expected/actual values and request ID in errors.
  4. Keep trajectories and algorithm-owned training tensors separately named. Do not silently flatten
    metadata, overwrite fields, or choose the first payload key and discard the remaining artifacts.
  5. Temporarily project explicitly selected artifacts to legacy diffusion_output, all_latents
    and audio fields while migrating the agent-loop/DataProto and consumers. Remove the projections
    after migration; AR's TokenOutput is unchanged.

Request/schema/strategy capability and reward/training mismatches fail fast. FFmpeg, Pillow,
W&B and dump I/O failures warn and record a fallback or skip without terminating training. These
are distinct error boundaries, not a blanket catch around correctness checks.

4.6 One compatibility adapter for the installed vLLM-Omni version (PR5/PR6 target)

The target diffusion normalizer will consume, in order:

  1. canonical payload/metadata exposed by OmniRequestOutput.multimodal_output;
  2. primary images plus final_output_type;
  3. temporary legacy dict/tuple forms emitted by in-tree adapters during migration.

Legacy handling is isolated and emits a deprecation warning identifying the adapter. Once all
in-tree adapters emit canonical payload/metadata and declare io_spec, tuple and alias fallbacks are
removed together.

This is deliberately different from putting another if MiniMax... or tuple convention in the
server.


5. Invariants

The implementation must preserve these invariants:

  • AR and diffusion select exactly one strategy per server instance.
  • No lifecycle behavior changes in the strategy-extraction PR.
  • No strategy calls a lifecycle method through a mode-specific code path.
  • Decoded image/video pixels remain uint8 [0, 255]; latent tensors retain their actual floating
    dtype, and decoded audio remains a floating waveform.
  • Audio is never passed through visual quantization.
  • The current transport has one primary plus optional audio; PR6 has independently named artifacts
    with an explicit compatibility-primary selection.
  • Every artifact has an explicit modality and representation; latent is not a modality.
  • Adapter-owned FPS/sample-rate defaults are declared next to the media stream; runtime values retain
    their legacy projections until a typed media container has a real consumer.
  • Every registered diffusion adapter declares a primary media modality; joint audio/video adapters
    also declare their auxiliary audio default.
  • Declared media_kind is propagated through V0/V1 media consumers during M4; full layout
    normalization and typed-output compatibility cleanup are deferred to M6.
  • Existing Ray callers can use the keyword API until the request-object migration completes.
  • Adding an existing-modality combination does not edit vllm_omni_async_server.py.

6. Implementation status and revised timeline

The original M0–M6 sequence was a design-time plan. The milestones below are the actual delivery
order and are deliberately narrower where a broader change would alter the pinned engine contract.
Each item remains separately reviewable.

Completed — strategy seam and output declarations

Milestone Status Actual result
M0/M1 — characterize and split generation modes #444 merged Introduced OmniStrategyBase, ARStrategy, and DiffusionStrategy; selection occurs once in _init_config. Shared lifecycle remains on the server. The public RPC and output objects are unchanged.
M2 — declare diffusion media output facts #478 merged Added CPU-importable rollout_media.py with MediaSpec and DiffusionIOSpec. All registered diffusion adapters declare their primary modality; MiniMax H3 and LTX-2 additionally declare auxiliary audio and sample-rate defaults. The shared 32 kHz fallback was removed.
Documentation #508 merged Documented adapter declarations, auxiliary tuple mapping, metadata defaults, and the contributor checklist.

M2 intentionally did not add a MediaOutput container or expand the local DiffusionOutput.
The current agent-loop DataProto.responses projection has a real tensor consumer, so adding an
unused second output representation would not yet reduce a compatibility boundary.

In review — finish the current MVP layers

M3 / PR3 — unify the internal request flow

#480 introduces PromptBundle, MediaInput,
and OmniRolloutRequest in pipelines/rollout_request.py.

  • OmniStrategyBase.generate() constructs one request object; both AR and diffusion strategies
    consume it in preprocess_input().
  • The public generate(**kwargs) RPC is unchanged.
  • Each strategy lowers the object back to the byte-compatible prompt keys consumed by the currently
    pinned engine and adapters, including the temporary duplicated multi_modal_data placement.
  • The five historical condition-image aliases are parsed and conflict-checked in one function, and
    the now-dead ImageGenerationRequest class is deleted. NumPy/Torch equality is shape-exact.
  • Unsupported strategy fields and duplicate modalities fail early; Ref2VA processor kwargs survive.

This is intentionally a server-side unification. It does not yet migrate every adapter's private
custom_prompt reads or rename prompt_token_ids to upstream prompt_ids.

M4 / PR4 — consume declared media metadata across V0 and V1

#481 covers the specified V0/V1 media consumers
and reliability boundaries, rather than being a V0-only change:

  1. carry the adapter-declared media_kind to V0 rollout dump and validation/W&B consumers;
  2. preserve media_kind, generated audio, and audio sample rate through the V1 agent-loop →
    TransferQueue → DataProto path;
  3. make V1 use the shared image/video/MP4-with-audio exporter;
  4. retain per-sample video .pt fallback and JSONL error metadata;
  5. log background dump failures instead of rethrowing them into training through
    Future.result(); and
  6. honor rollout save frequency and validation/rollout max-sample settings;
  7. forward generated audio/sample rate/media kind to both single and multi colocated reward scorers,
    rejecting conflicting generated sources before scoring;
  8. contain final W&B/table submission and directory/image/JSONL write failures, not just encoder or
    background-future failures; and
  9. reject runtime/declaration modality conflicts and unsupported auxiliary tuples. Correctness
    checks remain outside best-effort observability catches.

Tensor rank remains a compatibility fallback only when no declaration is present. Layout vocabulary
and normalization are not part of PR4: W&B/reward may still normalize axes heuristically until
PR6 verifies each pipeline's real VAE layout.

This also closes the V1 gap exposed by #501,
which adds MiniMax H3 V1 recipes but does not modify the trainer. It does not duplicate
#340, which addresses a V0 Wan2.2 layout case.

Implemented in draft — wire alignment and canonical media

M5 / PR5 — align the private engine request wire schema

After M3 proves the typed internal request boundary, migrate all in-tree adapters together:

  • replace private prompt_token_ids reads with the selected canonical engine field;
  • remove the top-level / extra_args duplicate multi_modal_data write and adapter workarounds;
  • regularize negative and extra-encoder prompt fields; and
  • validate every affected pipeline through GPU CI.

Implemented in #556. All in-tree readers use the canonical prompt boundary; diffusion media is
written once at top level. The AR entrance bridge follows the first stage's declared stage_type,
not the number of stages. CPU regressions cover both stage types with one/two stages; local GPU
coverage of the affected adapters is recorded in #557. Upstream CI remains a merge gate.

M6 / PR6 — establish the canonical typed media boundary

Combine typed output migration and layout normalization in one end-to-end change so neither lands as
an unused abstraction:

  1. GPU-verify decoded output axis order for each pipeline; separately document native/packed latent
    layouts without permuting training/replay tensors into decoded-media layouts.
  2. Replace the primary-plus-ordered-auxiliary convention with named artifacts carrying modality,
    representation, layout and data. Support latent reward/training plus optional decoded previews.
  3. Normalize decoded image/video/audio once to CHW / TCHW / CT, preserving each artifact's
    sample rate/FPS. Treat data dtype as authoritative.
  4. Migrate agent-loop/DataProto, trainer, W&B, ImageBind, CLAP and latent HTTP scoring to explicit
    artifacts. Revisit diffusion_rollout_output payload helpers, which currently select one key.
  5. Validate missing/extra/duplicate/requested artifacts and metadata with contextual errors; never
    silently replace a preview with latents. Remove tuple, first-payload-key, rank, channel-count and
    representation guessing after their consumers migrate.

Implemented in #557. Every in-tree adapter declares available names and emits explicit primary,
preview and audio selectors. Qwen/SD3/Boogu/Wan/H3/LTX/Bagel native or packed latents remain separate
from decoded CHW/TCHW/CT media and algorithm-owned replay tensors. Both ordinary/TQ transport,
reward managers, pixel/audio/latent scorers and V0/V1 media paths consume declarations. Missing
requested outputs, conflicting projections and schema errors fail; absent unrequested previews
skip export, never fall back to latents. V1 dump queue/storage retention is bounded and tested.

GPU validation also exposed and repaired Qwen-Edit raw-versus-preprocessed image view confusion,
non-serving warmup media validation and the pinned LTX forward-context sampler field. Generalized
prompt-rendering/condition-view abstractions remain separate follow-ups; no new abstraction is
needed to fix these concrete boundary consumers.

Delivery count

The RFC therefore consists of six implementation PRs: two merged (#444, #478) and four in
review (#480, #481, #556, #557). The remaining implementation is now present in the draft stack;
review, upstream CI and merging are not claimed complete.
Documentation PR #508 is complete and is not part of the numbered implementation series.

7. Verification

Current evidence (2026-09-07)

  • PR3 00ce99a: full CPU-pattern suite 1137 passed.
  • PR4 34da705, stacked on PR3: full CPU-pattern suite 1178 passed.
  • Regression tests reproduce alias broadcasting/equality errors, TQ-to-scorer media loss, modality
    conflicts, unsupported auxiliary outputs, and escaped W&B/Pillow/filesystem failures; the fixes
    are covered by repository tests. CPU failure injection is not real-engine validation.
  • Stack head 34da705: AR GSPO PASS (two training steps + validation, rc=0) and MiniMax H3
    T2VA DiffusionNFT V0 PASS on retry
    (one step + CLAP/ImageBind + actual MP4/audio export, rc=0).
    Export check: one capped sample, 384×256, 24 fps, 107 frames (96 requested, aligned by H3's
    17n+5 rule), stereo 32000 Hz audio, no fallback. [3/N][diffusion, rollout, tests] refactor: unify diffusion rollout request contract #480/[4/N][diffusion, rollout, trainer, tests] refactor: consume declared media kind #481 record the commands and limits.
  • The first MiniMax attempt exposed an outdated local synthetic VAE fixture: vllm-omni 0.28 now
    queries vae_ratio and split_tiles. An isolated tiny-fixture copy was corrected and CPU-checked;
    no PR or installed-package changes were needed for the passing retry.
  • PR5 9c28924 stage-entrance repair: 1196 CPU tests passed. PR6 beb2498 complete stack,
    including the preserved upstream main/merged-LoRA changes: 1320 CPU tests passed. Config/sanity/lint checks pass;
    full pre-commit launcher remains blocked by host Git 2.29.2's missing --deduplicate support.
  • PR6 local GPU matrix passes actual request lowering, registered adapters, named output parsing
    and export for Qwen FlowGRPO/DPO/NFT (request and step paths), dual/mix, SD3, Boogu, Wan, LTX,
    H3 NFT/FlowGRPO and Bagel. Qwen/SD3/Boogu packed tests verify multiple request IDs in producer
    context. Image-Edit completes conditioned LoRA training with latent primary, explicit decoded
    preview, JPEG reward and export.
  • Nine independent GPU probes load actual VAE implementations/local weights, including the
    2.6B-parameter H3 video VAE, and verify decoded/native axis contracts. LTX BWE audio is verified
    at 48000 Hz. This is distinct from the H3 tiny-model trainer fixture.
  • H3 NFT V0 and real V1/TQ complete a training step with latent primary, actual CLAP/ImageBind,
    weight synchronization and capped MP4/audio export. The V1 run explicitly enables
    trainer.use_v1=True and transfer_queue.enable=True, rather than counting the V0 fallback.
  • Reports, commands, source snapshot hashes and caveats are recorded in [6/N][BREAKING][diffusion, rollout, trainer, reward, tests] feat: add named media artifact path #557. These are interface,
    layout and plumbing checks, not convergence/quality claims or a GPU pass for every possible
    conditioning/parallelism combination. The AR entrance topology guard is CPU-tested; no new
    multistage-AR GPU pass is claimed. Upstream CI and human review remain required.

Run from the repository root with the project virtual environment activated:

source .venv/bin/activate

python -m pytest -q \
  tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py \
  tests/pipelines/test_diffusion_io_spec_on_cpu.py \
  tests/pipelines/test_rollout_request_on_cpu.py \
  tests/pipelines/test_minimax_h3_rollout_contract_on_cpu.py \
  tests/pipelines/test_minimax_h3_flow_grpo_on_cpu.py \
  tests/utils/test_tracking_media_on_cpu.py

Add focused tests for:

  • strategy selection and delegation;
  • no mode-dependent mutation leaking between strategies;
  • request lowering against the installed OmniCustomPrompt behavior;
  • canonical payload/metadata normalization;
  • legacy compatibility during M2-M4;
  • explicit image/video/audio modalities and separate latent/decoded representations;
  • video layout normalization;
  • required FPS/sample-rate validation;
  • abort/error output with no primary media; and
  • metadata collision rejection.

GPU evidence before removing compatibility paths:

  • one AR Qwen3-Omni rollout;
  • one image diffusion rollout;
  • one video rollout;
  • one joint video+audio rollout; and
  • one LoRA rollout with sleep/wake across an actor update.

The final lifecycle run is regression evidence only; lifecycle semantics are not changed here.


8. Risks and mitigations

Risk Mitigation
strategy extraction changes behavior while moving code M0 characterization tests; M1 contains no semantic fixes
strategy objects become another inheritance framework private protocol, exactly two implementations, no registry
request object breaks verl's shared LLMServerClient API retain the keyword RPC adapter through M3 and the later wire-schema migration
local media types duplicate upstream envelope use one translator; mirror upstream payload/metadata semantics rather than modifying the engine protocol
final_output_type is incomplete for some video paths adapter io_spec is authoritative; cross-check upstream declarations and fail on conflict
legacy prompt aliases or output projections become permanent remove request aliases in M5 and output projections in M6 after their respective GPU coverage
typed training becomes a renamed extra_fields reserve known groups (trajectory, prompt_embeddings, algorithm rl); reject collisions; forbid cross-pipeline consumers of private extras
scope duplicates #377 own only #377 §4.1 strategy extraction; leave lifecycle work in #377 and link the single implementation PR to both issues
scope becomes too large for one PR completed/in-review layers remain independently useful; V1 reliability, wire-schema, and layout work stay separate

9. Open questions

  1. Final class name: evolve the existing local DiffusionOutput in place, or introduce
    OmniMediaOutput and retain DiffusionOutput as a temporary adapter? In-place evolution avoids
    a third output representation and is the default proposal.
  2. Public request migration: should LLMServerClient.generate() eventually accept
    OmniRolloutRequest directly, or should the object remain verl-omni-specific behind
    DiffusionWholeSampleRetryLLMServerClient? M3 does not require this decision.
  3. Primary video declaration: should vLLM-Omni always propagate its normalized
    primary_key="video" into final_output_type? If accepted upstream, the local io_spec remains
    a validator instead of the only source of truth.
  4. Layout: M6 will verify each pipeline's real VAE layout before declaring it, then normalize
    to canonical TCHW; no axis order is inferred solely from model family or tensor rank.
  5. Batch cardinality: should num_outputs_per_prompt become a first-class request field rather
    than remain inside sampling_params? This should follow the request-batching work, not block M1.
  6. Metadata defaults: require every producer to emit FPS/sample rate, or permit adapter-owned
    defaults? The proposed rule permits declared adapter defaults during migration and requires
    emitted metadata in the end state.

10. Success criteria

This RFC is complete when:

  • vllm_omni_async_server.py is a shared server/lifecycle shell rather than an AR+diffusion
    implementation file;
  • mode-specific generation tests instantiate their strategies directly;
  • adding a model using existing modalities changes its adapter, I/O declaration, recipe and tests,
    not server, trainer, reward, W&B or dump code;
  • the standard diffusion path has one request envelope and one canonical condition lookup;
  • downstream code reads declared media kind instead of inferring it from rank;
  • named joint/latent/preview artifacts retain their metadata and are never silently dropped,
    relabeled or substituted;
  • legacy aliases and tuple conventions are removed after migration; and
  • [RFC][tracking] Refactor vllm_omni_async_server.py — Align with Upstream verl, Remove Temp Patches #377 remains the single source of truth for lifecycle behavior fixes.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

refactorCode transformations that improve a system's maintainability

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions