[Tracking] Audit and retire vLLM/vLLM-Omni patches for 0.28/0.29

[leviking98z-rgb]

Context

UniRL currently pins:

• vllm==0.27.0
• vllm-omni==0.27.0rc1

and carries a quarantined set of monkey patches under
unirl/rollout/engine/vllm_omni/patches/.

This issue tracks which patches are still necessary when upgrading the runtime,
which ones can be retired, and which implementations should be tightened before
an upgrade.

Related historical upgrade issue: #129.

Audit baselines

• UniRL origin/main: 7ec3683d581fe0e5e9543f6103024e94a20edc9a
• Runtime-tested stable pair:
  • vllm==0.28.0
  • vllm-omni==0.28.0
  • torch==2.13.0+cu130
  • transformers==5.6.0
  • NVIDIA H20
• Source-audited development pair:
  • vLLM b3b13c12923f0b68f36ee515ebbc3efc7679eb4f (roughly 0.29-era main)
  • vLLM-Omni 21d86ec9239d5aba08ea5fcac7e20a8f0679b0c9

vllm 0.29.0 + vllm-omni 0.28.0 should not be treated as a supported pair.
The stable paired upgrade target is 0.28/0.28; current vLLM-Omni main is the
appropriate source baseline for vLLM 0.29.

Summary

Patch / behavior	vLLM 0.28 + Omni 0.28	vLLM 0.29 + Omni main
Explicit register_capture_flush() call inside the bundle	Remove	Remove
vllm_omni.general_plugins capture-flush entry point	Keep: UniRL feature	Keep: UniRL feature
compat_hi3_lora tuple unwrapping	Keep	Remove / version-gate
DiT and AR in-memory tensor LoRA loaders	Keep	Keep
HI3 DiT namespace alias and GQA-interleaved QKV-B conversion	Keep	Keep
HI3 AR fused-QKV tensor conversion	Keep	Keep
fp32 LoRA wrapping skip	Keep for HI3 LoRA	Keep for HI3 LoRA
HV1.5 packed-module mapping	Keep	Keep
HV1.5 token-refiner nn.Linear LoRA support	Keep	Keep
lora_request passthrough through Omni.generate()	Keep	Keep
HI3 custom-sigmas passthrough	Keep for request-level/legacy execution	Keep for request-level/legacy execution
MoE workspace CuMem pool	Keep for colocated sleep mode	Keep for colocated sleep mode
HI3 Base tokenizer compatibility	Keep for Base; unnecessary for Instruct	Same
Per-request AR seed patch	Not needed for a single replica	Not needed for a single replica
Fate sharing/watchdog	Keep when hard parent-death cleanup is required	Same
Spawn-child bundle reinstall	Keep while any runtime patches remain	Same

Findings

Patches that can be retired

1. Remove the bundle's explicit capture-flush registration

Both vLLM-Omni 0.28 and current main call
load_omni_general_plugins() from the spawned diffusion stage/worker paths.
UniRL already publishes register_capture_flush through the
vllm_omni.general_plugins entry point.

The plugin itself is still required because capture metadata is a UniRL feature;
only the redundant direct call and its now-stale comment should be removed.

2. Retire compat_hi3_lora on current Omni main

In 0.28, HI3 returns:

(expert_mapping, expert_weights_remapping)

while the vLLM LoRA caller expects the flat expert-mapping list. The compatibility
patch is therefore still required on the stable 0.28 pair.

Current Omni main returns the flat list from get_expert_mapping() and exposes
the remapping through _get_expert_weights_remapping(), so this patch can be
removed or feature-gated for the 0.29-era stack.

3. Reclassify patch_per_request_ar_seed

The current retirement rationale says a shared SamplingParams object makes all
requests in a GRPO group produce identical tokens. That is not true for the
current v1 GPU model runner:

• seed=None selects the ordinary random-sampling path.
• Requests without an explicit seed use the runner's advancing global GPU RNG.
• A 0.28 GPU probe with 64 rows produced 55 distinct sampled tokens.

Therefore the patch is not required for UniRL's current single-replica stage
configs. It can still prevent correlated streams when multiple independent
replicas/engines are initialized with the same model seed: two independently
reset runners produced the same first sampled token in the probe.

If multi-replica protection is needed, prefer a deterministic seed derived from
the request/group identity over os.urandom, so rollout reproduction remains
possible.

Patches that remain necessary

• Native AR and diffusion LoRA loaders still load from local checkpoints and do
  not accept UniRL's generic in-memory tensor request.
• Native diffusion LoRA still lacks the HI3 transformer.layers.* to
  model.layers.* alias and does not convert HI3's GQA-interleaved fused QKV-B
  rows into packed [q, k, v] slices.
• The AR merged-QKV path has more packed-tensor handling in newer vLLM, but does
  not implement the HI3 interleaved layout conversion.
• vllm.lora.utils.from_layer() still has no dtype guard, while the HI3 AR
  router gate remains fp32.
• The diffusion LoRA manager derives packed mappings from a module attribute
  named stacked_params_mapping. HV1.5 exposes packed_modules_mapping at class
  level and only creates stacked_params_mapping as a local variable inside
  load_weights().
• HV1.5 token-refiner projections are ordinary torch.nn.Linear modules, which
  the native diffusion LoRA wrapper does not support.
• AsyncOmniEngine.add_request() accepts lora_request, but Omni.generate()
  still does not expose/forward it.
• The HI3 inner pipeline can consume custom sigmas, but the outer request path
  still does not forward sampling_params.sigmas. Step execution explicitly
  rejects custom timesteps/sigmas, so this patch should not be described as
  enabling every execution mode.
• WorkspaceManager still allocates its workspace with ordinary torch.empty,
  outside a sleep-managed CuMem pool.
• HI3 still computes ratio_36 + 1; the Base tokenizer lacks ratio tokens
  33-36. Instruct checkpoints do not need this compatibility behavior.
• Upstream death signals and death_pipe handling do not provide the same
  root-anchored watchdog coverage for a parent SIGKILL, nested spawn chains,
  or a worker stuck inside CUDA/NCCL.

Required cleanup before/with the upgrade

• Remove the explicit register_capture_flush() call from
  VLLMOmniHijack.hijack() and update the stale plugin-loading comments.
• Feature/version-gate compat_hi3_lora, retaining it for 0.28 and removing
  it for current Omni main.
• Remove patch_per_request_ar_seed for single-replica configurations, or
  replace it with request-derived deterministic seeding when stage replicas
  are enabled.
• Make the whole patch bundle genuinely idempotent. Repeated install() on
  0.28 currently re-wraps DiffusionLoRAManager._load_adapter and
  AsyncOmniEngine.add_request, despite the README claiming repeat installs
  are safe.
• Change patch_dit_lora_loader to intercept only
  OmniTensorLoRARequest; delegate every other request to the original
  _load_adapter. The current full replacement bypasses newer upstream
  model-specific _load_diffusion_lora_adapter hooks.
• Update the MoE workspace pool key from only ubatch to (ubatch, lane);
  newer WorkspaceManager supports both dimensions.
• Scope tokenizer fallback to HI3 Base initialization instead of globally
  changing every slow and fast Hugging Face tokenizer's missing-token
  behavior.
• Keep tests that distinguish “the patch can still install” from “upstream
  still lacks the behavior”.
• Upgrade vLLM and vLLM-Omni as a tested pair rather than mixing the 0.29
  vLLM wheel with the 0.28 Omni wheel.

Validation result

The 0.28/0.28 GPU harness ran 19 independent checks. Eighteen patch behaviors
passed. The only failing check was repeated bundle installation:

SUMMARY total=19 failures=1

identity changed after the second install:
- DiffusionLoRAManager._load_adapter
- AsyncOmniEngine.add_request

This means most patches are mechanically compatible with 0.28, but compatibility
alone should not be confused with continued necessity; the retirement decisions
above are based on the corresponding native upstream paths.

[leviking98z-rgb]

Implementation update:

• fix(vllm-omni): retire per-request AR seed patch #474 removes patch_per_request_ar_seed from current main. The vLLM 0.27/0.28 sampler keeps seed=None on the advancing batched RNG path, so the patch is unnecessary for current single-replica recipes and its os.urandom override prevents rollout reproduction.
• feat(rollout): migrate vLLM-Omni to stable 0.28.0 #413 already removes the redundant direct register_capture_flush() call as part of the 0.28 migration.
• compat_hi3_lora remains required for the stable 0.28 pair and is intentionally not removed.

[MinhaoLi0318]

I'd like to take two of the cleanup items:

1. Make the bundle genuinely idempotent and narrow patch_dit_lora_loader to intercept only OmniTensorLoRARequest, delegating every other request to the original _load_adapter (same shape as patch_ar_lora_loader already has). That covers the one failing check in your 19-check harness and stops bypassing newer upstream _load_diffusion_lora_adapter hooks. One PR, CPU tests for the sentinel/delegation, plus a rollout regression on SD3 via vllm_omni.
2. patch_lora_request_passthrough: I confirmed sync Omni.generate still lacks lora_request on both vllm-omni v0.28.0 and current main (only AsyncOmni.generate has it), so I'll open the small upstream PR the patch README asks for, and retire the patch here once a pinned release carries it.

Leaving compat_hi3_lora version-gating and the MoE workspace pool key to you or whoever picks them up, unless you'd like me to take those too.