[Perf] Support request-level batching for Wan2.2 pipelines (#5676)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>

vllm_omni/diffusion/sched/interface.py

@@ -88,9 +88,10 @@ class RequestBatchSamplingParamsKey:
     """Request level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``.
 
     Only request-batch-wide fields belong here. Request-local values such as
-    seeds, generators, latent tensors, timesteps, and pipeline-specific
+    seeds, generators, latent tensors, timesteps, and request-local
     ``extra_args`` are read per request from
-    ``DiffusionRequestBatch.sampling_params_list``.
+    ``DiffusionRequestBatch.sampling_params_list``. Structural ``extra_args``
+    that affect batching are listed explicitly below.
     """
 
     # Spatial / temporal shape.
@@ -131,6 +132,16 @@ class RequestBatchSamplingParamsKey:
     # the model and must remain distinct from explicit ``lossless``.
     quality: str | None = None
 
+    # Wan scheduler structure is carried through extra_args. Requests using
+    # different solvers or flow shifts must not share a request batch.
+    sample_solver: str | None = None
+    flow_shift: float | None = None
+
+    # Pipeline-specific condition structure populated during preprocessing.
+    # It prevents independently valid requests with incompatible conditions
+    # from being admitted to the same request batch.
+    condition_key: tuple[Any, ...] | None = None
+
     # LoRA identity.
     lora_int_id: int | None = None
     lora_scale: float = 1.0

vllm_omni/diffusion/sched/request_scheduler.py

@@ -22,7 +22,21 @@ if TYPE_CHECKING:
 # on sampling params, so it must be resolved separately from the bulk lookup.
 _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(
     field.name for field in fields(RequestBatchSamplingParamsKey)
-) - {"lora_int_id"}
+) - {"condition_key", "flow_shift", "lora_int_id", "sample_solver"}
+
+
+def _normalize_explicit_sample_solver(value: object | None) -> str | None:
+    """Normalize an explicitly provided solver without selecting a default."""
+    if value is None:
+        return None
+    return str(value).strip().lower()
+
+
+def _normalize_explicit_flow_shift(value: object | None) -> float | None:
+    """Normalize an explicitly provided flow shift without selecting a default."""
+    if value is None:
+        return None
+    return float(value)
 
 
 def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
@@ -31,6 +45,13 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re
     # LoRA identity is optional on sampling params (and on test stubs).
     lora_request = getattr(sampling, "lora_request", None)
     key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
+    extra_args = sampling.extra_args or {}
+    # Match pipeline resolution for explicit overrides, but preserve None:
+    # pipeline/engine defaults are configuration-dependent and must not be
+    # inferred while building the request-batch key.
+    key_kwargs["sample_solver"] = _normalize_explicit_sample_solver(extra_args.get("sample_solver"))
+    key_kwargs["flow_shift"] = _normalize_explicit_flow_shift(extra_args.get("flow_shift"))
+    key_kwargs["condition_key"] = getattr(request, "batch_compatibility_key", None)
     key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
     return RequestBatchSamplingParamsKey(**key_kwargs)