Merge branch 'main' into codex/unskip-fp8-z-image-quality

.buildkite/test-nightly.yml

@@ -570,6 +570,7 @@ steps:
             tests/e2e/online_serving/test_longcat_image_expansion.py
             tests/e2e/online_serving/test_longcat_image_edit_expansion.py
             tests/e2e/online_serving/test_flux_2_dev_expansion.py
+            tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -842,8 +843,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
         timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -851,7 +852,7 @@ steps:
           - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -1287,7 +1288,7 @@ steps:
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - nightly-diffusion-x2iat-performance-qwen-image-layered
       - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
@@ -1304,7 +1305,7 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance

docs/design/module/dit_module.md

@@ -169,15 +169,13 @@ class DiffusionRequestStatus(enum.IntEnum):
 
 @dataclass
 class DiffusionRequestState:
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
 ```
 
 **Design Features**:
 
-- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
-
 - **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
 
 - **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
@@ -188,7 +186,6 @@ class DiffusionRequestState:
 class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self._request_states = {}
-        self._request_id_to_sched_req_id = {}
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
@@ -199,7 +196,8 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Common state storage**: Shared request maps and waiting/running sets live in the base class.
 
-- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
+- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
+  duplicated in each policy.
 
 - **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
 
@@ -219,12 +217,12 @@ class RequestScheduler(_BaseScheduler):
 
 - **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
 
-- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
 
 #### 2.5 Engine-Driven Execution Loop
 
 ```python
-sched_req_id = scheduler.add_request(request)
+request_id = scheduler.add_request(request)
 while True:
     sched_output = scheduler.schedule()
     output = executor.add_req(req)

docs/mkdocs/hooks/generate_examples.py

@@ -16,6 +16,7 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
 
 def fix_case(text: str) -> str:
@@ -191,18 +192,21 @@ class Example:
 
         return re.sub(link_pattern, replace_link, content)
 
+    def github_url(self, path: Path) -> str:
+        url = "https://github.com/vllm-project/vllm-omni/"
+        url += "tree/main" if path.is_dir() else "blob/main"
+        return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
+
     def generate(self) -> str:
         content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if self.path.is_dir() else "blob/main"
-        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
+        content += f"Source <{self.github_url(self.path)}>.\n\n"
 
         # Use long code fence to avoid issues with
         # included files containing code fences too
         code_fence = "``````"
 
         if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR)
+            main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
             content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
         else:
             with open(self.main_file, encoding="utf-8") as f:
@@ -216,10 +220,15 @@ class Example:
 
         content += "## Example materials\n\n"
         for file in sorted(self.other_files):
-            content += f'??? abstract "{file.relative_to(self.path)}"\n'
+            content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
+            if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
+                content += (
+                    f"    Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
+                )
+                continue
             if file.suffix != ".md":
                 content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}\n"
 
@@ -291,7 +300,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: str(rel_path)})
+            category_items.append({example.title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")

docs/mkdocs/hooks/search_index.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reduce browser search-index memory for generated API reference pages."""
+
+from material.plugins.search.plugin import SearchIndex
+
+_add_entry_from_context = SearchIndex.add_entry_from_context
+
+
+def add_entry_from_context(self: SearchIndex, page) -> None:
+    search = page.meta.get("search") or {}
+    if search.get("exclude"):
+        return
+
+    if str(page.url).startswith("api/"):
+        entry = {
+            "location": page.url,
+            "title": str(page.meta.get("title", page.title)),
+            "text": "",
+        }
+        if "boost" in search:
+            entry["boost"] = search["boost"]
+        self.entries.append(entry)
+        return
+
+    _add_entry_from_context(self, page)
+
+
+SearchIndex.add_entry_from_context = add_entry_from_context

docs/mkdocs/javascript/mathjax.js

@@ -1,20 +1,44 @@
-// Enables MathJax rendering
+// Enables MathJax rendering only on pages that contain math.
 window.MathJax = {
-    tex: {
-      inlineMath: [["\\(", "\\)"]],
-      displayMath: [["\\[", "\\]"]],
-      processEscapes: true,
-      processEnvironments: true
-    },
-    options: {
-      ignoreHtmlClass: ".*|",
-      processHtmlClass: "arithmatex"
-    }
-  };
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true,
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex",
+  },
+};
 
-  document$.subscribe(() => {
-    MathJax.startup.output.clearCache()
-    MathJax.typesetClear()
-    MathJax.texReset()
-    MathJax.typesetPromise()
-  })
+let mathJaxLoading;
+
+function loadMathJax() {
+  if (window.MathJax.typesetPromise) {
+    return Promise.resolve();
+  }
+  if (!mathJaxLoading) {
+    mathJaxLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
+      script.async = true;
+      script.onload = resolve;
+      script.onerror = reject;
+      document.head.appendChild(script);
+    });
+  }
+  return mathJaxLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".arithmatex")) {
+    return;
+  }
+  loadMathJax().then(() => {
+    MathJax.startup.output.clearCache();
+    MathJax.typesetClear();
+    MathJax.texReset();
+    MathJax.typesetPromise();
+  });
+});

docs/mkdocs/javascript/memory.js

@@ -0,0 +1,2 @@
+// Opt out of bfcache so large generated API pages are released after navigation.
+window.addEventListener("unload", () => {});

docs/mkdocs/javascript/mermaid.js

@@ -1,21 +1,42 @@
-// Initialize Mermaid for diagram rendering
-mermaid.initialize({
-  startOnLoad: false,
-  theme: 'default',
-  securityLevel: 'loose',
-  flowchart: {
-    useMaxWidth: true,
-    htmlLabels: true
-  }
-});
+// Load Mermaid only on pages that contain diagrams.
+let mermaidLoading;
 
-// Render Mermaid diagrams when page content is ready
-document$.subscribe(() => {
-  const mermaidElements = document.querySelectorAll('.mermaid');
-  if (mermaidElements.length > 0) {
-    mermaid.run({
-      querySelector: '.mermaid',
-      nodes: mermaidElements
+function loadMermaid() {
+  if (window.mermaid) {
+    return Promise.resolve();
+  }
+  if (!mermaidLoading) {
+    mermaidLoading = new Promise((resolve, reject) => {
+      const script = document.createElement("script");
+      script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
+      script.async = true;
+      script.onload = () => {
+        mermaid.initialize({
+          startOnLoad: false,
+          theme: "default",
+          securityLevel: "loose",
+          flowchart: {
+            useMaxWidth: true,
+            htmlLabels: true,
+          },
+        });
+        resolve();
+      };
+      script.onerror = reject;
+      document.head.appendChild(script);
     });
   }
+  return mermaidLoading;
+}
+
+document$.subscribe(() => {
+  if (!document.querySelector(".mermaid")) {
+    return;
+  }
+  loadMermaid().then(() => {
+    mermaid.run({
+      querySelector: ".mermaid",
+      nodes: document.querySelectorAll(".mermaid"),
+    });
+  });
 });

docs/user_guide/quantization/autoround.md

@@ -52,8 +52,8 @@ vLLM-Omni's runtime module names.
 
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
+| GLM-Image | Diffusion transformer  | ✅ | `Intel/GLM-Image-int4-AutoRound` |
 | BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
-| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
 
 ## Configuration
 

examples/offline_inference/internvla_a1/end2end.py

@@ -127,7 +127,12 @@ def run_one_path(
             ),
             args.device,
         )
-        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-sample-{index}",
+        )
         pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
         results.append(
             {
@@ -150,6 +155,7 @@ def run_pipeline_forward(
     pipeline,
     batch_inputs: dict[str, torch.Tensor],
     noise: torch.Tensor,
+    request_id: str,
 ) -> torch.Tensor:
     output = pipeline.forward(
         OmniDiffusionRequest(
@@ -161,6 +167,7 @@ def run_pipeline_forward(
                     "decode_image": False,
                 }
             ),
+            request_id=request_id,
         )
     )
     if output.error:
@@ -212,23 +219,38 @@ def benchmark_forward(
 
     _synchronize(args.device)
     cold_start_begin = time.perf_counter()
-    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+    pred = run_pipeline_forward(
+        pipeline,
+        batch_inputs,
+        noise,
+        request_id=f"internvla-a1-benchmark-{index}-cold",
+    )
     _synchronize(args.device)
     cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
 
     warmup_ms: list[float] = []
-    for _ in range(args.warmup_iters):
+    for iter_idx in range(args.warmup_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
+        )
         _synchronize(args.device)
         warmup_ms.append((time.perf_counter() - begin) * 1000.0)
 
     benchmark_ms: list[float] = []
-    for _ in range(args.benchmark_iters):
+    for iter_idx in range(args.benchmark_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _ = run_pipeline_forward(
+            pipeline,
+            batch_inputs,
+            noise,
+            request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
+        )
         _synchronize(args.device)
         benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
 
@@ -295,7 +317,12 @@ def main() -> None:
             dataset=dataset,
             train_meta=train_meta,
             collate_samples=collate_open_loop_samples,
-            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
+                policy,
+                batch_inputs,
+                noise,
+                request_id="internvla-a1-open-loop",
+            ),
             num_episodes=args.num_episodes,
             seed=args.seed,
             device=args.device,

mkdocs.yml

@@ -38,13 +38,12 @@ theme:
   features:
     - content.action.edit
     - content.code.copy
-    - navigation.instant
-    - navigation.instant.progress
     - navigation.tracking
     - navigation.tabs
     - navigation.tabs.sticky
     - navigation.sections
     - navigation.indexes
+    - navigation.prune
     - navigation.top
     - search.suggest
     - search.highlight
@@ -56,6 +55,7 @@ theme:
   custom_dir: docs/mkdocs/overrides
 
 hooks:
+  - docs/mkdocs/hooks/search_index.py
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
@@ -117,6 +117,7 @@ plugins:
             separate_signature: true
             show_overloads: true
             signature_crossrefs: true
+            show_source: false
           inventories:
           - https://docs.python.org/3/objects.inv
           - https://typing-extensions.readthedocs.io/en/latest/objects.inv
@@ -167,9 +168,8 @@ extra_css:
   - mkdocs/stylesheets/extra.css
 
 extra_javascript:
+  - mkdocs/javascript/memory.js
   - mkdocs/javascript/mathjax.js
-  - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
-  - https://unpkg.com/mermaid@10/dist/mermaid.min.js
   - mkdocs/javascript/mermaid.js
   - mkdocs/javascript/edit_and_feedback.js
   - mkdocs/javascript/slack_and_forum.js

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -63,24 +63,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.15,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 55000,
-                    "peak_memory_mb_mean": 55000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i_2img",
                 "dataset": "random",

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -61,23 +61,6 @@
             }
         },
         "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_mean": 10,
-                    "peak_memory_mb_max": 56000,
-                    "peak_memory_mb_mean": 56000
-                }
-            },
             {
                 "name": "1536x1536_steps35_i2i",
                 "dataset": "random",

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -60,6 +60,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
             output_type="pil",
             generator_device="cpu",
         ),
+        "request_id": "diffusers-backend",
     }
     defaults.update(overrides)
     return OmniDiffusionRequest(**defaults)

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -48,6 +48,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=prompts or [{"prompt": "a cat dancing"}],
         sampling_params=sp,
+        request_id="dmd2-sanitize",
     )
 
 

tests/diffusion/models/dmd2/test_dmd2_scheduler.py

@@ -48,7 +48,7 @@ def _make_pipeline(cls):
 
 def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
     sp = OmniDiffusionSamplingParams(**sp_kwargs)
-    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
+    return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
 
 
 @pytest.fixture(

tests/diffusion/models/ernie_image/test_ernie_image_pe.py

@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
 
 
 def test_should_apply_pe_disables_dummy_warmup_request():
-    req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
+    req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
 
     assert ErnieImagePipeline._should_apply_pe(req) is False
 

tests/diffusion/models/glm_image/test_glm_image_quantization.py

@@ -0,0 +1,741 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
+
+These tests verify that the GLM-Image DiT transformer correctly accepts and uses
+quantization configs for W4A16/AutoRound quantization support.
+"""
+
+import os
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
+    ColumnParallelGELU,
+    ColumnParallelSiLU,
+    GlmImageAdaLayerNormContinuous,
+    GlmImageAdaLayerNormZero,
+    GlmImageAttention,
+    GlmImageFeedForward,
+    GlmImageImageProjector,
+    GlmImagePrepare,
+    GlmImageRotaryPosEmbed,
+    GlmImageTransformer2DModel,
+    GlmImageTransformerBlock,
+    _positive_divisors,
+    validate_glm_image_tp_constraints,
+)
+from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    ReplicatedLinear (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform
+    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
+    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
+    the dispatcher to always return the default (torch.nn.functional.linear)
+    implementation which works on any device."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_mocks(mocker: MockerFixture):
+    """Set up common mocks for all tests."""
+    mocker.patch(
+        "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
+        return_value=1,
+    )
+    mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 1
+    mock_get_tp_group.return_value = mock_tp_group
+    yield
+
+
+class TestPositiveDivisors:
+    """Test _positive_divisors helper function."""
+
+    def test_divisors_of_1(self):
+        assert _positive_divisors(1) == {1}
+
+    def test_divisors_of_12(self):
+        assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
+
+    def test_divisors_of_prime(self):
+        assert _positive_divisors(7) == {1, 7}
+
+    def test_divisors_of_0_returns_empty(self):
+        assert _positive_divisors(0) == set()
+
+    def test_divisors_of_negative_returns_empty(self):
+        assert _positive_divisors(-5) == set()
+
+
+class TestValidateGlmImageTpConstraints:
+    """Test TP constraint validation for GLM-Image."""
+
+    def test_valid_tp_size_1(self):
+        """TP=1 should always be valid."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=1,
+        )
+        assert 1 in result
+
+    def test_valid_tp_size_2_for_divisible_dim(self):
+        """TP=2 is valid when all dims are divisible by 2."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=2,
+        )
+        assert 1 in result
+        assert 2 in result
+
+    def test_valid_tp_size_4_for_divisible_dim(self):
+        """TP=4 is valid when all dims are divisible by 4."""
+        result = validate_glm_image_tp_constraints(
+            dim=2560,
+            num_heads=64,
+            ffn_hidden_dim=10240,
+            tensor_parallel_size=4,
+        )
+        assert 1 in result
+        assert 2 in result
+        assert 4 in result
+
+    def test_invalid_tp_size_3_for_divisible_dim(self):
+        """TP=3 is invalid when dim is not divisible by 3."""
+        with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,  # 2560 % 3 != 0
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=3,
+            )
+
+    def test_invalid_tp_size_zero(self):
+        """TP=0 should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=0,
+            )
+
+    def test_invalid_tp_size_negative(self):
+        """Negative TP size should raise error."""
+        with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
+            validate_glm_image_tp_constraints(
+                dim=2560,
+                num_heads=64,
+                ffn_hidden_dim=10240,
+                tensor_parallel_size=-1,
+            )
+
+
+class TestGlmImageAdaLayerNormZeroQuantization:
+    """Test GlmImageAdaLayerNormZero with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=mock_quant_config,
+            prefix="test.norm1",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormZero(
+            embedding_dim=512,
+            dim=2560,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        hidden_states = torch.randn(batch_size, seq_len, 2560)
+        encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
+        temb = torch.randn(batch_size, 512)
+
+        # This should work regardless of whether linear returns tuple or tensor
+        result = layer(hidden_states, encoder_hidden_states, temb)
+        assert len(result) == 10  # Should return 10 chunks
+
+
+class TestGlmImageAdaLayerNormContinuousQuantization:
+    """Test GlmImageAdaLayerNormContinuous with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the class accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=mock_quant_config,
+            prefix="test.norm_out",
+        )
+        assert layer.linear.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+        assert layer.linear.quant_config is None
+
+    def test_forward_handles_tuple_return_from_linear(self):
+        """Verify forward handles tuple returns from ReplicatedLinear."""
+        layer = GlmImageAdaLayerNormContinuous(
+            embedding_dim=2560,
+            conditioning_embedding_dim=512,
+            quant_config=None,
+        )
+
+        batch_size = 2
+        seq_len = 10
+        x = torch.randn(batch_size, seq_len, 2560)
+        conditioning_embedding = torch.randn(batch_size, 512)
+
+        result = layer(x, conditioning_embedding)
+        assert result.shape == (batch_size, seq_len, 2560)
+
+
+class TestGlmImageAttentionQuantization:
+    """Test GlmImageAttention with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageAttention accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=mock_quant_config,
+            prefix="test.attn1",
+        )
+        assert attn.to_qkv.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        attn = GlmImageAttention(
+            dim=2560,
+            num_heads=64,
+            head_dim=40,
+            quant_config=None,
+        )
+        assert attn.to_qkv.quant_config is None
+
+
+class TestColumnParallelModulesQuantization:
+    """Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
+
+    def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelGELU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelGELU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.gelu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+    def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
+        """Verify ColumnParallelSiLU accepts quant_config."""
+        mock_quant_config = mocker.MagicMock()
+        layer = ColumnParallelSiLU(
+            dim_in=2560,
+            dim_out=10240,
+            quant_config=mock_quant_config,
+            prefix="test.silu",
+        )
+        assert layer.proj.quant_config is mock_quant_config
+
+
+class TestGlmImageFeedForwardQuantization:
+    """Test GlmImageFeedForward with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageFeedForward accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=mock_quant_config,
+            prefix="test.ff",
+        )
+        # Check that the first layer (ColumnParallelGELU) has quant_config
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="gelu-approximate",
+            quant_config=None,
+        )
+        gelu_layer = ff.net[0]
+        assert gelu_layer.proj.quant_config is None
+
+    def test_linear_silu_activation(self):
+        """Test linear-silu activation function initialization works."""
+        # This test verifies the module can be instantiated with linear-silu activation.
+        # Full forward testing requires proper TP group setup, so we just verify construction.
+        ff = GlmImageFeedForward(
+            dim=2560,
+            dim_out=2560,
+            inner_dim=10240,
+            activation_fn="linear-silu",
+            quant_config=None,
+        )
+        # Verify the FFN has the correct structure
+        assert len(ff.net) == 3  # ColumnParallelSiLU, Identity, RowParallelLinear
+        assert isinstance(ff.net[0], ColumnParallelSiLU)
+
+
+class TestGlmImageTransformerBlockQuantization:
+    """Test GlmImageTransformerBlock with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify GlmImageTransformerBlock accepts quant_config parameter."""
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            parallel_config=parallel_config,
+            quant_config=mock_quant_config,
+            prefix="test.block",
+        )
+        # Check that inner modules have quant_config
+        assert block.norm1.linear.quant_config is mock_quant_config
+        assert block.attn1.to_qkv.quant_config is mock_quant_config
+        assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self):
+        """Verify quant_config=None is accepted."""
+        block = GlmImageTransformerBlock(
+            dim=2560,
+            num_attention_heads=64,
+            attention_head_dim=40,
+            time_embed_dim=512,
+            quant_config=None,
+        )
+        assert block.norm1.linear.quant_config is None
+
+
+class TestGlmImagePrepareModule:
+    """Test GlmImagePrepare module."""
+
+    def test_prepare_module_exists(self):
+        """Verify GlmImagePrepare module exists and works."""
+        projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
+        prepare = GlmImagePrepare(
+            image_projector=projector,
+            rope=rope,
+            patch_size=2,
+        )
+
+        # Test forward pass
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+        result = prepare(hidden_states)
+
+        assert len(result) == 5
+        hidden_out, rope_cos, rope_sin, height, width = result
+        assert hidden_out.shape[0] == 1  # batch size
+        assert hidden_out.shape[1] == 1024  # seq_len (32 * 32)
+        assert hidden_out.shape[2] == 2560  # hidden_dim
+
+
+class TestGlmImagePipelineConfig:
+    """Test GLM_IMAGE_PIPELINE configuration."""
+
+    def test_glm_image_pipeline_config_exists(self):
+        """Verify GLM_IMAGE_PIPELINE config exists."""
+        assert GLM_IMAGE_PIPELINE is not None
+
+    def test_pipeline_has_correct_model_type(self):
+        """Verify pipeline has correct model_type."""
+        assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
+
+    def test_pipeline_has_correct_model_arch(self):
+        """Verify pipeline has correct model_arch."""
+        assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
+
+    def test_pipeline_has_two_stages(self):
+        """Verify pipeline has two stages (AR + DiT)."""
+        assert len(GLM_IMAGE_PIPELINE.stages) == 2
+
+    def test_stage_0_is_ar_llm(self):
+        """Verify stage 0 is the AR LLM stage."""
+        stage_0 = GLM_IMAGE_PIPELINE.stages[0]
+        assert stage_0.stage_id == 0
+        assert stage_0.model_stage == "ar"
+        assert stage_0.execution_type.value == "llm_ar"
+        assert stage_0.owns_tokenizer is True
+
+    def test_stage_1_is_diffusion(self):
+        """Verify stage 1 is the DiT diffusion stage."""
+        stage_1 = GLM_IMAGE_PIPELINE.stages[1]
+        assert stage_1.stage_id == 1
+        assert stage_1.model_stage == "dit"
+        assert stage_1.execution_type.value == "diffusion"
+        assert stage_1.final_output is True
+        assert stage_1.final_output_type == "image"
+        assert stage_1.input_sources == (0,)  # Takes input from stage 0
+
+    def test_pipeline_has_diffusers_class_name(self):
+        """Verify pipeline has diffusers_class_name."""
+        assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
+
+
+class TestGlmImageImageProjector:
+    """Test GlmImageImageProjector module."""
+
+    def test_projector_output_shape(self):
+        """Verify projector produces correct output shape."""
+        projector = GlmImageImageProjector(
+            in_channels=16,
+            hidden_size=2560,
+            patch_size=2,
+        )
+
+        # Input: [B, C, H, W] = [1, 16, 64, 64]
+        hidden_states = torch.randn(1, 16, 64, 64)
+        output = projector(hidden_states)
+
+        # After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
+        assert output.shape == (1, 1024, 2560)
+
+    def test_projector_with_different_patch_sizes(self):
+        """Verify projector works with different patch sizes."""
+        for patch_size in [2, 4]:
+            projector = GlmImageImageProjector(
+                in_channels=16,
+                hidden_size=2560,
+                patch_size=patch_size,
+            )
+            h, w = 64, 64
+            hidden_states = torch.randn(1, 16, h, w)
+            output = projector(hidden_states)
+            expected_seq_len = (h // patch_size) * (w // patch_size)
+            assert output.shape[1] == expected_seq_len
+
+
+class TestGlmImageRotaryPosEmbed:
+    """Test GlmImageRotaryPosEmbed module."""
+
+    def test_rope_output_shape(self):
+        """Verify RoPE produces correct output shape."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
+
+        cos, sin = rope(hidden_states)
+
+        # After patchify: height=32, width=32
+        # Output: [32*32, dim] = [1024, 40]
+        assert cos.shape[0] == 1024
+        assert cos.shape[1] == 40
+        assert sin.shape == cos.shape
+
+    def test_rope_value_range(self):
+        """Verify RoPE produces values in valid range."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(2, 16, 32, 32)
+
+        cos, sin = rope(hidden_states)
+
+        # cos and sin should be in [-1, 1]
+        assert cos.min() >= -1.0 and cos.max() <= 1.0
+        assert sin.min() >= -1.0 and sin.max() <= 1.0
+
+    def test_rope_consistency(self):
+        """Verify RoPE is consistent for same input."""
+        rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
+        hidden_states = torch.randn(1, 16, 32, 32)
+
+        cos1, sin1 = rope(hidden_states)
+        cos2, sin2 = rope(hidden_states)
+
+        # Same input should produce same output
+        assert torch.allclose(cos1, cos2)
+        assert torch.allclose(sin1, sin2)
+
+
+class TestGlmImageTransformer2DModelQuantization:
+    """Test GlmImageTransformer2DModel with quantization config."""
+
+    def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
+        """Verify the model accepts quant_config parameter."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        # Create a minimal mock od_config
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2  # Small number for testing
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # Check that quantization config was passed to transformer blocks
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is mock_quant_config
+            assert block.attn1.to_qkv.quant_config is mock_quant_config
+            assert block.ff.net[0].proj.quant_config is mock_quant_config
+
+    def test_accepts_none_quant_config(self, mocker: MockerFixture):
+        """Verify quant_config=None is accepted."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Check that quantization config is None
+        for block in model.transformer_blocks:
+            assert block.norm1.linear.quant_config is None
+            assert block.attn1.to_qkv.quant_config is None
+
+    def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
+        """Verify norm_out (output layer) does NOT use quantization to preserve precision."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        mock_quant_config = mocker.MagicMock()
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=mock_quant_config,
+        )
+
+        # norm_out.linear should NOT have quant_config to preserve output precision
+        assert model.norm_out.linear.quant_config is None
+
+    def test_model_has_sp_plan(self, mocker: MockerFixture):
+        """Verify model has _sp_plan defined for sequence parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        # Verify _sp_plan exists
+        assert hasattr(model, "_sp_plan")
+        assert "prepare" in model._sp_plan
+        assert "proj_out" in model._sp_plan
+
+    def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
+        """Verify model has _hsdp_shard_conditions for HSDP parallelism."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        assert hasattr(model, "_hsdp_shard_conditions")
+        assert len(model._hsdp_shard_conditions) > 0
+
+    def test_model_creates_kv_cache(self, mocker: MockerFixture):
+        """Verify model can create KV cache for image editing."""
+        from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+        )
+
+        mock_tf_config = mocker.MagicMock()
+        mock_tf_config.patch_size = 2
+        mock_tf_config.in_channels = 16
+        mock_tf_config.out_channels = 16
+        mock_tf_config.num_attention_heads = 64
+        mock_tf_config.attention_head_dim = 40
+        mock_tf_config.time_embed_dim = 512
+        mock_tf_config.condition_dim = 256
+        mock_tf_config.prior_vq_quantizer_codebook_size = 16384
+        mock_tf_config.text_embed_dim = 1024
+        mock_tf_config.num_layers = 2
+
+        mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
+        mock_od_config.tf_model_config = mock_tf_config
+        mock_od_config.parallel_config = parallel_config
+
+        model = GlmImageTransformer2DModel(
+            od_config=mock_od_config,
+            quant_config=None,
+        )
+
+        kv_cache = model.create_kv_cache()
+        assert kv_cache is not None
+        assert len(kv_cache) == 2  # num_layers

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -339,6 +339,7 @@ class TestCFGParallelForwardPath:
                 audio_latents=audio_latents,
                 output_type="latent",
             ),
+            request_id="ltx23-cfg-parallel-forward-test",
         )
 
         output = pipe.forward(req)

tests/diffusion/models/ovis_image/test_ovis_image.py

@@ -167,6 +167,7 @@ def test_basic_generation(ovis_pipeline):
     # Setup request
     req = OmniDiffusionRequest(
         prompts=["A photo of a cat"],
+        request_id="ovis-basic",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -206,6 +207,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
                 "negative_prompt": "bad quality",
             }
         ],
+        request_id="ovis-guidance",
         sampling_params=OmniDiffusionSamplingParams(
             height=256,
             width=256,
@@ -225,6 +227,7 @@ def test_resolution_check(ovis_pipeline):
     # Pass odd resolution
     req = OmniDiffusionRequest(
         prompts=["test"],
+        request_id="ovis-resolution",
         sampling_params=OmniDiffusionSamplingParams(
             height=250,  # Not divisible by 16 (8*2)
             width=250,

tests/diffusion/test_diffusion_engine.py

@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int = 0
 
     @property
-    def scheduled_req_ids(self):
+    def scheduled_request_ids(self):
         ids = [req.request_id for req in self.scheduled_new_reqs]
         if self.scheduled_cached_reqs:
-            ids.extend(self.scheduled_cached_reqs.sched_req_ids)
+            ids.extend(self.scheduled_cached_reqs.request_ids)
         return ids
 
     @property
     def is_empty(self):
-        return len(self.scheduled_req_ids) == 0
+        return len(self.scheduled_request_ids) == 0
 
 
 class MockScheduler:
@@ -71,31 +71,36 @@ async def test_async_add_req_and_wait_for_response():
     engine.scheduler = MockScheduler()
     engine._out_queue = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
-    engine._rpc_lock = threading.Lock()
+    engine._rpc_queue = queue.Queue()
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._init_lock = asyncio.Lock()
+    engine._closed = False
+    engine._loop_started = False
+    engine.main_loop = None
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
     def mock_execute_batch(sched_output):
-        req_ids = sched_output.scheduled_req_ids
+        request_ids = sched_output.scheduled_request_ids
 
         time.sleep(1)
 
         class MockRunnerOutput:
             def __init__(self, ids):
-                self.req_id = ids
+                self.request_id = ids
                 self.step_index = 0
                 self.finished = True
                 self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
 
-            def get_req_output(self, rid):
+            def get_request_output(self, rid):
                 return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
 
-        return MockRunnerOutput(req_ids)
+        return MockRunnerOutput(request_ids)
 
     engine.execute_fn = mock_execute_batch
 
-    engine.stop_event = threading.Event()
-    engine.start_background_loop()
+    await engine._check_and_start_background_loop()
 
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
@@ -105,10 +110,13 @@ async def test_async_add_req_and_wait_for_response():
 
     task_ids = [f"req_{i}" for i in range(5)]
     tasks = [run_task(rid) for rid in task_ids]
-    results = await asyncio.gather(*tasks)
-
-    engine.stop_event.set()
-    engine.worker_thread.join(timeout=5)
+    try:
+        results = await asyncio.gather(*tasks)
+    finally:
+        with engine._cv:
+            engine.stop_event.set()
+            engine._cv.notify_all()
+        engine.worker_thread.join(timeout=5)
     assert len(results) == 5
     for rid, res, elapsed in results:
         assert rid in res.result_data

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -17,15 +17,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def _make_request(
-    request_id: str,
-    *,
-    request_ids: list[str] | None = None,
-) -> OmniDiffusionRequest:
+def _make_request(request_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{request_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=request_ids or [request_id],
         request_id=request_id,
     )
 
@@ -79,15 +74,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
     engine._finalize_finished_request.assert_not_called()
 
 
-def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
+def test_abort_request_id_aborts_scheduler_request() -> None:
     engine = _make_engine()
-    request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
-    sched_req_id = engine.scheduler.add_request(request)
+    request = _make_request("batch-parent")
+    request_id = engine.scheduler.add_request(request)
 
     engine.abort("batch-parent")
     engine._process_aborts_queue()
 
-    state = engine.scheduler.get_request_state(sched_req_id)
+    state = engine.scheduler.get_request_state(request_id)
     assert state is not None
     assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
 

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
-                tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
+                tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
                 return DiffusionOutput(error=f"result_for_{tag}")
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
             exec_all_ranks=True,
         )
         return RunnerOutput(
-            req_id=new_req.sched_req_id,
+            request_id=new_req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
 
 def _make_request(tag: str):
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
     )
@@ -148,6 +148,7 @@ def _make_engine_with_loop(
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._out_queue = {}
+    engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
 

tests/diffusion/test_diffusion_request.py

@@ -5,7 +5,7 @@ import random
 
 import pytest
 
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -15,9 +15,53 @@ def _make_request() -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[{"prompt": "a cup of coffee on a table"}],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="request-test",
     )
 
 
+def test_request_id_is_required():
+    with pytest.raises(TypeError, match="request_id"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        )
+
+
+def test_request_ids_identity_list_is_removed():
+    req = _make_request()
+
+    assert req.request_id == "request-test"
+    assert not hasattr(req, "request_ids")
+
+
+def test_request_id_must_be_non_empty():
+    with pytest.raises(ValueError, match="request_id must be a non-empty string"):
+        OmniDiffusionRequest(
+            prompts=[{"prompt": "a cup of coffee on a table"}],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_id="",
+        )
+
+
+def test_dummy_run_request_is_identified_by_reserved_request_id():
+    req = OmniDiffusionRequest(
+        prompts=[{"prompt": "dummy run"}],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id=DUMMY_DIFFUSION_REQUEST_ID,
+    )
+
+    assert req.is_dummy_run()
+    assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+
+
+def test_non_dummy_request_is_not_identified_as_dummy_run():
+    req = _make_request()
+
+    assert not req.is_dummy_run()
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
+    assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
+
+
 def test_tp_seed_same_across_ranks_and_varies_across_requests():
     random.seed(0)
     n_requests = 5

tests/diffusion/test_diffusion_scheduler.py

@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompts=[f"prompt_{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=None,
         finished=finished,
         result=DiffusionOutput(output=None, error=error),
@@ -52,7 +52,7 @@ def _make_step_output(
     error: str | None = None,
 ):
     return RunnerOutput(
-        req_id=req_id,
+        request_id=req_id,
         step_index=step_index,
         finished=finished,
         result=DiffusionOutput(output=None, error=error) if error is not None else None,
@@ -73,16 +73,16 @@ def _make_step_request(
             num_inference_steps=num_inference_steps,
             step_index=step_index,
         ),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
 def _new_ids(sched_output) -> list[str]:
-    return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
+    return [req.request_id for req in sched_output.scheduled_new_reqs]
 
 
 def _cached_ids(sched_output) -> list[str]:
-    return list(sched_output.scheduled_cached_reqs.sched_req_ids)
+    return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
 class _StubScheduler(SchedulerInterface):
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
         self._request = request
         self._output = output
         self.initialized_with = None
-        self._sched_req_id = request.request_ids[0]
+        self._request_id = request.request_id
         self._state = None
         self._scheduled = False
         self.max_num_running_reqs = 1
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
-        return self._sched_req_id
+        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        return self._request_id
 
     def schedule(self):
         if self._scheduled or self._state is None:
             return SimpleNamespace(
                 scheduled_new_reqs=[],
                 scheduled_cached_reqs=CachedRequestData.make_empty(),
-                scheduled_req_ids=[],
+                scheduled_request_ids=[],
                 is_empty=True,
             )
         self._scheduled = True
         return SimpleNamespace(
             scheduled_new_reqs=[NewRequestData.from_state(self._state)],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
-            scheduled_req_ids=[self._state.sched_req_id],
+            scheduled_request_ids=[self._state.request_id],
             is_empty=False,
         )
 
@@ -123,30 +123,25 @@ class _StubScheduler(SchedulerInterface):
         del sched_output
         assert output is self._output
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
-        return {self._sched_req_id}
+        return {self._request_id}
 
     def has_requests(self) -> bool:
         return not self._scheduled
 
-    def get_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def get_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        if request_id in self._request.request_ids:
-            return self._sched_req_id
-        return None
-
-    def pop_request_state(self, sched_req_id: str):
-        del sched_req_id
+    def pop_request_state(self, request_id: str):
+        del request_id
         return self._state
 
-    def preempt_request(self, sched_req_id: str) -> bool:
-        del sched_req_id
+    def preempt_request(self, request_id: str) -> bool:
+        del request_id
         return False
 
-    def finish_requests(self, sched_req_ids, status) -> None:
-        del sched_req_ids, status
+    def finish_requests(self, request_ids, status) -> None:
+        del request_ids, status
         return None
 
     def close(self) -> None:
@@ -171,7 +166,7 @@ class TestGetSamplingParamsKey:
         return OmniDiffusionRequest(
             prompts=["prompt"],
             sampling_params=sp,
-            request_ids=[f"req-{lora_int_id}-{lora_scale}"],
+            request_id=f"req-{lora_int_id}-{lora_scale}",
         )
 
     def test_distinguishes_lora_id(self) -> None:
@@ -291,7 +286,7 @@ class TestRequestScheduler:
             OmniDiffusionRequest(
                 prompts=["prompt_b"],
                 sampling_params=OmniDiffusionSamplingParams(width=768),
-                request_ids=["b"],
+                request_id="b",
             )
         )
         scheduler.add_request(_make_request("c"))
@@ -333,7 +328,7 @@ class TestRequestScheduler:
         assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
 
         assert self.scheduler.has_requests() is False
-        assert self.scheduler.schedule().scheduled_req_ids == []
+        assert self.scheduler.schedule().scheduled_request_ids == []
 
     def test_has_requests_state_transition(self) -> None:
         assert self.scheduler.has_requests() is False
@@ -348,60 +343,28 @@ class TestRequestScheduler:
         assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
         assert self.scheduler.has_requests() is False
 
-    def test_request_id_mapping_lifecycle(self) -> None:
+    def test_request_id_is_scheduler_key(self) -> None:
         request = OmniDiffusionRequest(
             prompts=["prompt_map_a", "prompt_map_b"],
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["map-a", "map-b"],
             request_id="map-parent",
         )
 
-        sched_req_id = self.scheduler.add_request(request)
+        request_id = self.scheduler.add_request(request)
 
-        assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
-        assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
+        assert request_id == "map-parent"
+        state = self.scheduler.get_request_state("map-parent")
+        assert state.request_id == "map-parent"
 
-        self.scheduler.pop_request_state(sched_req_id)
+        self.scheduler.pop_request_state("map-parent")
 
-        assert self.scheduler.get_sched_req_id("map-a") is None
-        assert self.scheduler.get_sched_req_id("map-b") is None
-        assert self.scheduler.get_sched_req_id("map-parent") is None
+        assert self.scheduler.get_request_state("map-parent") is None
 
-    def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
-        self.scheduler.add_request(
-            OmniDiffusionRequest(
-                prompts=["prompt_existing"],
-                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-                request_ids=["existing-child"],
-                request_id="duplicate-parent",
-            )
-        )
+    def test_duplicate_request_id_is_rejected(self) -> None:
+        self.scheduler.add_request(_make_request("dup"))
 
-        colliding_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="duplicate-parent",
-        )
-
-        with pytest.raises(ValueError, match="duplicate-parent"):
-            self.scheduler.add_request(colliding_request)
-
-        assert self.scheduler.get_request_state("unique-child") is None
-        assert self.scheduler.get_sched_req_id("unique-child") is None
-
-        valid_request = OmniDiffusionRequest(
-            prompts=["prompt_unique"],
-            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
-            request_ids=["unique-child"],
-            request_id="unique-parent",
-        )
-        sched_req_id = self.scheduler.add_request(valid_request)
-
-        assert sched_req_id == "unique-child"
-        assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
-        assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
+        with pytest.raises(ValueError, match="request_id 'dup' is already active"):
+            self.scheduler.add_request(_make_request("dup"))
 
 
 class TestDiffusionEngine:
@@ -662,7 +625,7 @@ class TestStepScheduler:
         finished = self.scheduler.update_from_output(
             sched_output,
             RunnerOutput(
-                req_id=req_id,
+                request_id=req_id,
                 step_index=None,
                 finished=True,
                 result=None,
@@ -934,7 +897,7 @@ class TestStepScheduler:
 
         for _ in range(expected_steps - 1):
             sched_output = self.scheduler.schedule()
-            assert sched_output.scheduled_req_ids == [req_id]
+            assert sched_output.scheduled_request_ids == [req_id]
             next_step = request.sampling_params.step_index + 1
             assert (
                 self.scheduler.update_from_output(
@@ -945,7 +908,7 @@ class TestStepScheduler:
             )
 
         final_output = self.scheduler.schedule()
-        assert final_output.scheduled_req_ids == [req_id]
+        assert final_output.scheduled_request_ids == [req_id]
         assert self.scheduler.update_from_output(
             final_output,
             _make_step_output(req_id, step_index=expected_steps, finished=True),

tests/diffusion/test_diffusion_step_pipeline.py

@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
 def _make_step_request(num_inference_steps: int = 2):
     return SimpleNamespace(
         prompts=["a prompt"],
-        request_ids=["req-1"],
+        request_id="req-1",
         sampling_params=SimpleNamespace(
             generator=None,
             seed=None,
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
     return OmniDiffusionRequest(
         prompts=[f"prompt-{req_id}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
 
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
     return runner
 
 
-def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
+def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
-        scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
+        scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
         scheduled_cached_reqs=CachedRequestData.make_empty(),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
@@ -287,7 +287,7 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
 
 def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     """Scheduler output for a homogeneous batch (one NewRequestData per req)."""
-    new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
+    new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=new_reqs,
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     )
 
 
-def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
+def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
         scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
+        scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
         num_waiting_reqs=0,
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
             runner,
             _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
         )
-        output = result.get_req_output("req-1")
+        output = result.get_request_output("req-1")
 
         assert output.finished is True
         assert output.result is not None
@@ -384,16 +384,16 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        first = result.get_req_output("req-1")
-        assert first.req_id == "req-1"
+        first = result.get_request_output("req-1")
+        assert first.request_id == "req-1"
         assert first.step_index == 1
         assert first.finished is False
         assert first.result is None
         assert "req-1" in runner.state_cache
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
-        second = result.get_req_output("req-1")
-        assert second.req_id == "req-1"
+        second = result.get_request_output("req-1")
+        assert second.request_id == "req-1"
         assert second.step_index == 2
         assert second.finished is True
         assert second.result is not None
@@ -410,13 +410,13 @@ class TestRunner:
         runner = _make_runner()
         req_1 = _make_step_request()
         req_2 = _make_step_request()
-        req_2.request_ids = ["req-2"]
+        req_2.request_id = "req-2"
 
         scheduler_output = DiffusionSchedulerOutput(
             step_id=0,
             scheduled_new_reqs=[
-                NewRequestData(sched_req_id="req-1", req=req_1),
-                NewRequestData(sched_req_id="req-2", req=req_2),
+                NewRequestData(request_id="req-1", req=req_1),
+                NewRequestData(request_id="req-2", req=req_2),
             ],
             scheduled_cached_reqs=CachedRequestData.make_empty(),
             finished_req_ids=set(),
@@ -431,7 +431,7 @@ class TestRunner:
         runner = _make_runner()
 
         with pytest.raises(ValueError, match="Missing cached state"):
-            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
+            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
 
     def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
         runner = _make_runner()
@@ -440,8 +440,8 @@ class TestRunner:
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
         result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-        output = result.get_req_output("req-1")
-        assert output.req_id == "req-1"
+        output = result.get_request_output("req-1")
+        assert output.request_id == "req-1"
         assert output.step_index == 0
         assert output.finished is True
         assert output.result is not None
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
     worker = object.__new__(DiffusionWorker)
     worker.lora_manager = lora_manager
     worker._step_lora_state = {}
-    output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
+    output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
     worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
     return worker
 
@@ -525,9 +525,9 @@ class TestWorker:
     """DiffusionWorker.execute_stepwise"""
 
     def test_delegates_to_model_runner(self):
-        expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
         worker = _make_step_worker(expected_output=expected)
-        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -536,7 +536,7 @@ class TestWorker:
     def test_deactivates_lora_when_request_has_no_adapter(self):
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -552,7 +552,7 @@ class TestWorker:
 
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
+        scheduler_output = _make_scheduler_output(request, request_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
@@ -568,8 +568,8 @@ class TestWorker:
 
         manager = _RecordingLoRAManager()
         worker = _make_step_worker(lora_manager=manager)
-        first = _make_scheduler_output(request, sched_req_id="req-1")
-        second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
+        first = _make_scheduler_output(request, request_id="req-1")
+        second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
 
         DiffusionWorker.execute_stepwise(worker, first)
         DiffusionWorker.execute_stepwise(worker, second)
@@ -610,10 +610,10 @@ class TestWorker:
         next_request.sampling_params.lora_request = lora_request
 
         worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
-        first = _make_scheduler_output(finishing, sched_req_id="req-1")
+        first = _make_scheduler_output(finishing, request_id="req-1")
         next_batch = _make_scheduler_output(
             next_request,
-            sched_req_id="req-2",
+            request_id="req-2",
             step_id=1,
             finished_req_ids={"req-1"},
         )
@@ -633,11 +633,11 @@ class TestExecutor:
     def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
         executor = object.__new__(MultiprocDiffusionExecutor)
         executor._ensure_open = lambda: None
-        expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
+        expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
         executor.collective_rpc = mocker.Mock(return_value=expected)
 
         request = _make_engine_request("req-step", num_inference_steps=2)
-        scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
+        scheduler_output = _make_scheduler_output(request, request_id="req-step")
 
         output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
 
@@ -653,7 +653,7 @@ class TestEngine:
         [
             (
                 lambda _: RunnerOutput(
-                    req_id="req-error",
+                    request_id="req-error",
                     step_index=1,
                     finished=True,
                     result=DiffusionOutput(error="boom"),
@@ -688,7 +688,7 @@ class TestEngine:
             call_count["n"] += 1
             finished = call_count["n"] == 2
             return RunnerOutput(
-                req_id="req-step",
+                request_id="req-step",
                 step_index=call_count["n"],
                 finished=finished,
                 result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
@@ -714,7 +714,7 @@ class TestEngine:
             step["n"] += 1
             engine.abort("req-stop")
             return RunnerOutput(
-                req_id="req-stop",
+                request_id="req-stop",
                 step_index=1,
                 finished=False,
                 result=None,
@@ -741,7 +741,7 @@ class TestEngine:
                 assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
                 engine.abort("req-mid")
             return RunnerOutput(
-                req_id="req-mid",
+                request_id="req-mid",
                 step_index=step["n"],
                 finished=False,
                 result=None,
@@ -760,7 +760,7 @@ class TestEngine:
         engine = _make_engine(
             scheduler,
             execute_fn=lambda _: RunnerOutput(
-                req_id="req-missing",
+                request_id="req-missing",
                 step_index=1,
                 finished=True,
                 result=None,
@@ -777,7 +777,7 @@ class TestEngine:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
 def _mock_request(tag: str):
     """Return a lightweight request object identifiable by *tag*."""
     return SimpleNamespace(
-        request_ids=[tag],
+        request_id=tag,
         prompts=[f"prompt_{tag}"],
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
     )
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
             req = req_q.get(timeout=10)
             method = req.get("method", "")
             args = req.get("args", ())
-            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
-                tag = f"result_for_{args[0].request_ids[0]}"
+            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
+                tag = f"result_for_{args[0].request_id}"
             elif args:
                 tag = f"result_for_{args[0]}"
             else:

tests/diffusion/test_stage_diffusion_proc.py

@@ -68,7 +68,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
 
         request = captured["request"]
         assert request.request_id == "req-parent"
-        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert not hasattr(request, "request_ids")
         assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
         assert result.request_id == "req-parent"
         assert result.images == ["img-1", "img-2"]

tests/distributed/omni_connectors/test_kv_flow.py

@@ -384,7 +384,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_recv"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
     # req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
     # But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
@@ -428,7 +428,6 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
         request_id=parent_req_id,
     )
 
@@ -454,7 +453,6 @@ def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_con
     req = OmniDiffusionRequest(
         prompts=["prompt-a", "prompt-b"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=["req-parent-0", "req-parent-1"],
         request_id="req-parent",
     )
     req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
@@ -516,7 +514,7 @@ def test_integration_flow(common_constants):
     req = OmniDiffusionRequest(
         prompts=["test_integ"],
         sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
+        request_id=req_id,
     )
 
     # Receive

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -253,37 +253,6 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
     )
 
 
-def _vllm_omni_output_multiple_image(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "vllm_omni_multiple.png"
-    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
-        output = _run_vllm_omni_image_edit(
-            omni_server=server,
-            prompt=PROMPT_MULTIPLE_IMAGE,
-            input_images=[qwen_bear_image, rabbit_image],
-            output_path=output_path,
-        )
-    return output
-
-
-def _diffusers_output_multiple_image(
-    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
-) -> Image.Image:
-    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
-    output_path = output_dir / "diffusers_multiple.png"
-    return _run_diffusers_image_edit(
-        model=MULTIPLE_MODEL,
-        pipeline_class=QwenImageEditPlusPipeline,
-        prompt=PROMPT_MULTIPLE_IMAGE,
-        input_images=[qwen_bear_image, rabbit_image],
-        output_path=output_path,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_single_matches_diffusers(
@@ -309,37 +278,6 @@ def test_qwen_image_edit_single_matches_diffusers(
     )
 
 
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(
-    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
-)
-def test_qwen_image_edit_multiple_matches_diffusers(
-    accuracy_artifact_root: Path,
-    qwen_bear_image: Image.Image,
-    rabbit_image: Image.Image,
-) -> None:
-    vllm_image = _vllm_omni_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    diffusers_image = _diffusers_output_multiple_image(
-        accuracy_artifact_root=accuracy_artifact_root,
-        qwen_bear_image=qwen_bear_image,
-        rabbit_image=rabbit_image,
-    )
-    assert_similarity(
-        model_name=MULTIPLE_MODEL,
-        vllm_image=vllm_image,
-        diffusers_image=diffusers_image,
-        width=WIDTH,
-        height=HEIGHT,
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
-    )
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:

tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py

@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for GLM-Image AutoRound W4A16 quantized inference.
+
+These tests cover text-to-image and image-to-image generation with
+the W4A16 quantized GLM-Image model.
+
+Requirements:
+  - 2 CUDA GPUs (H100 or equivalent)
+  - The quantized model checkpoint (Intel/GLM-Image-int4-AutoRound)
+"""
+
+import gc
+import math
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+from vllm import SamplingParams
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+QUANTIZED_MODEL = os.environ.get("GLM_IMAGE_AUTOROUND_MODEL", "Intel/GLM-Image-int4-AutoRound")
+
+# Small resolution to keep GPU memory & time manageable
+HEIGHT = 256
+WIDTH = 256
+NUM_STEPS = 2  # minimal for smoke-test
+
+# GLM-Image AR generation config (from generation_config.json)
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+_CI_DEPLOY = get_deploy_config_path("glm_image.yaml")
+
+
+def _get_stage_config():
+    """Build a CI-friendly stage config with eager mode for testing."""
+    return modify_stage_config(
+        _CI_DEPLOY,
+        updates={
+            "stages": {
+                0: {"enforce_eager": True},
+                1: {"enforce_eager": True},
+            },
+        },
+    )
+
+
+stage_config = _get_stage_config()
+
+# (model, stage_config_path) for ``omni_runner`` indirect parametrize
+_OMNI_RUNNER_PARAM = (QUANTIZED_MODEL, stage_config)
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """Compute max_new_tokens for GLM-Image AR text-to-image generation."""
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    return small_tokens + large_tokens + 1
+
+
+def _ar_sampling_params(max_tokens: int, height: int, width: int, seed: int = 42) -> SamplingParams:
+    """Build AR stage SamplingParams for GLM-Image."""
+    return SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=seed,
+        detokenize=False,
+        extra_args={
+            "target_h": height,
+            "target_w": width,
+        },
+    )
+
+
+def _diffusion_sampling_params(
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    num_steps: int = NUM_STEPS,
+    seed: int = 42,
+) -> OmniDiffusionSamplingParams:
+    """Build Diffusion stage OmniDiffusionSamplingParams."""
+    return OmniDiffusionSamplingParams(
+        height=height,
+        width=width,
+        num_inference_steps=num_steps,
+        guidance_scale=0.0,
+        seed=seed,
+    )
+
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.diffusion,
+]
+
+
+# ------------------------------------------------------------------
+# Test: text-to-image generation produces a valid image (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image model and verify it produces a valid image."""
+    gc.collect()
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
+    prompt_dict = {
+        "prompt": "A photo of a cat sitting on a laptop keyboard",
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    monitor.stop()
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
+    assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
+
+    # Sanity: image should not be blank (all-zero)
+    arr = np.array(img)
+    assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+
+# ------------------------------------------------------------------
+# Test: image-to-image generation (quantized)
+# ------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler):
+    """Load the W4A16 quantized GLM-Image and verify image-to-image generation works."""
+    ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"]
+
+    gc.collect()
+    current_omni_platform.empty_cache()
+    current_omni_platform.reset_peak_memory_stats()
+
+    prompt_dict = {
+        "prompt": "Make it look like winter",
+        "multi_modal_data": {"image": ref_image_arr},
+        "height": HEIGHT,
+        "width": WIDTH,
+        "mm_processor_kwargs": {
+            "target_h": HEIGHT,
+            "target_w": WIDTH,
+        },
+    }
+    ar_params = _ar_sampling_params(
+        max_tokens=compute_max_tokens(HEIGHT, WIDTH),
+        height=HEIGHT,
+        width=WIDTH,
+        seed=42,
+    )
+    diffusion_params = _diffusion_sampling_params(
+        height=HEIGHT,
+        width=WIDTH,
+        num_steps=NUM_STEPS,
+        seed=42,
+    )
+
+    outputs = omni_runner_handler.runner.generate(
+        [prompt_dict],
+        [ar_params, diffusion_params],
+    )
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    images = req_out.images
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert isinstance(img, Image.Image)
+    assert img.width == WIDTH
+    assert img.height == HEIGHT
+
+    gc.collect()
+    current_omni_platform.empty_cache()

tests/e2e/online_serving/test_qwen_image_edit.py

@@ -5,8 +5,8 @@
 Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
 
 - ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2509`` — two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2509`` only, ``advanced_model`` — mixed input
+- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
   resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
   scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
 
@@ -53,7 +53,6 @@ def _get_diffusion_feature_cases(model: str):
 
 
 @pytest.mark.advanced_model
-@pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
@@ -83,15 +82,16 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
     openai_client.send_diffusion_request(request_config)
 
 
+@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2509``."""
+    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
 
     image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
 
@@ -118,7 +118,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
     indirect=True,
 )
 def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -2,7 +2,7 @@
 Comprehensive tests of diffusion features that are available in online serving mode
 and are supported by the following models:
 - Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2509: single image input and two image inputs
+- Qwen-Image-Edit-2511: single image input and two image inputs
 """
 
 import pytest
@@ -14,7 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
-SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -34,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
+                    "cache_dit",
                     "--enable-layerwise-offload",
                 ],
             ),
@@ -59,7 +59,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "cache_dit",
+                    "tea_cache",
                     "--ring",
                     "2",
                 ],
@@ -145,7 +145,7 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
     indirect=True,
 )
 def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with a single image input.
+    """Test Qwen-Image-Edit-2511 with a single image input.
 
     Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
     postprocess closure used the doubled temb (shape 2*batch) without halving
@@ -156,7 +156,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
-    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
 
     request_config = {
         "model": omni_server.model,
@@ -180,7 +180,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
     indirect=True,
 )
 def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2509 with two image inputs."""
+    """Test Qwen-Image-Edit-2511 with two image inputs."""
     image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
     image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 

tests/model_executor/models/glm_image/test_glm_image_init.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for GLM-Image model __init__.py lazy import pattern.
+
+The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing
+``transformers.models.glm_image`` at module init, which may not be available
+in all environments.
+"""
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestLazyImport:
+    """Test the __getattr__ lazy import pattern in glm_image __init__.py."""
+
+    def test_getattr_exists_and_is_callable(self):
+        """Verify __getattr__ exists and is callable."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        assert hasattr(glm_image_pkg, "__getattr__")
+        assert callable(glm_image_pkg.__getattr__)
+
+    def test_getattr_returns_class_for_known_attribute(self):
+        """Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Call __getattr__ directly to test the lazy import logic
+        result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration")
+
+        # Verify we get a class (the actual GlmImageForConditionalGeneration)
+        assert result is not None
+        assert isinstance(result, type)
+        assert result.__name__ == "GlmImageForConditionalGeneration"
+
+    def test_getattr_raises_for_unknown_attribute(self):
+        """Verify __getattr__ raises AttributeError for unknown attributes."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Test unknown attribute via __getattr__ directly
+        with pytest.raises(AttributeError, match="has no attribute"):
+            glm_image_pkg.__getattr__("UnknownClass")
+
+    def test___all___exports_correct_symbols(self):
+        """Verify __all__ contains the expected exported symbols."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        assert hasattr(glm_image_pkg, "__all__")
+        assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__
+
+
+class TestLazyImportDoesNotImportTransformersAtInit:
+    """Verify that importing the package does not eagerly load transformers."""
+
+    def test_glm_image_module_has_getattr(self):
+        """Test that the module has __getattr__ for lazy loading."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # The module should have __getattr__
+        assert hasattr(glm_image_pkg, "__getattr__")
+        assert callable(glm_image_pkg.__getattr__)
+
+        # And __all__
+        assert hasattr(glm_image_pkg, "__all__")
+
+    def test_module_has_proper_structure(self):
+        """Test that the module has proper Python module structure."""
+        import vllm_omni.model_executor.models.glm_image as glm_image_pkg
+
+        # Should be a module
+        assert hasattr(glm_image_pkg, "__name__")
+        assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image"
+
+        # Should have __file__ (even if it's a package)
+        assert hasattr(glm_image_pkg, "__file__")

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -267,8 +267,8 @@ def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
     model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
 
     ref_audio_ids = []
-    model._encode_ref_audio_to_code = lambda wav, _sr: (
-        ref_audio_ids.append(id(wav)) or torch.ones((2, 2), dtype=torch.long)
+    model._encode_ref_audio_batch = lambda wavs, sr, *, device: (
+        [ref_audio_ids.append(id(w)) or torch.ones((2, 2), dtype=torch.long) for w in wavs]
     )
     model._extract_speaker_embedding = lambda wav, _sr: (
         ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
@@ -303,21 +303,16 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
         )
     )
 
-    class FakeSpeechTokenizer:
-        def __init__(self):
-            self.calls = []
+    encode_calls = []
 
-        def encode(self, audios, *, sr, return_dict):
-            self.calls.append((audios, sr, return_dict))
-            return SimpleNamespace(
-                audio_codes=[
-                    torch.full((2, 2), 11, dtype=torch.long),
-                    torch.full((3, 2), 22, dtype=torch.long),
-                ]
-            )
+    def _fake_encode_batch(wavs, sr, *, device):
+        encode_calls.append((wavs, sr))
+        return [
+            torch.full((2, 2), 11, dtype=torch.long),
+            torch.full((3, 2), 22, dtype=torch.long),
+        ]
 
-    tok = FakeSpeechTokenizer()
-    model._ensure_speech_tokenizer_loaded = lambda: tok
+    model._encode_ref_audio_batch = _fake_encode_batch
 
     class FakeTextTokenizer:
         def __init__(self):
@@ -353,12 +348,11 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
     )
 
     assert normalize_calls == ["a.wav", "b.wav"]
-    assert len(tok.calls) == 1
-    audios, sr, return_dict = tok.calls[0]
-    assert audios[0] is wav1
-    assert audios[1] is wav2
+    assert len(encode_calls) == 1
+    wavs, sr = encode_calls[0]
+    assert wavs[0] is wav1
+    assert wavs[1] is wav2
     assert sr == 16000
-    assert return_dict is True
     assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
     assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
     assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1

vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py

@@ -123,10 +123,6 @@ class OmniRandomMultiModalDataset(RandomMultiModalDataset):
                 "23",
                 "-movflags",
                 "+faststart",
-                "-pix_fmt",
-                "yuv420p",
-                "-vf",
-                f"scale={width}:{height}",
             ],
         }
 

vllm_omni/deploy/qwen3_tts.yaml

@@ -28,6 +28,11 @@ connectors:
       codec_left_context_frames: 72
       # Emit only the first audio chunk early, then return to codec_chunk_frames.
       initial_codec_chunk_frames: 1
+      # Common Stage1 decode buckets:
+          #   no-ref first/steady chunks: 25 / 97 frames
+          #   Base ref-context first/steady chunks: 73 / 169 frames
+          #   decoder internal non-streaming chunks: 325 frames
+      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
 
 stages:
   - stage_id: 0

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -77,7 +77,7 @@ class DefaultAdapter:
                 os.path.join("transformer", "config.json"),
                 od_config.model,
             )
-            od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+            od_config.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
 
         loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
         # load_model will handle dtypes / device placement, put in .eval() mode
@@ -202,6 +202,7 @@ class TeaCacheCoefficientEstimator:
         self.hook.start_collection()
         req = OmniDiffusionRequest(
             prompts=[prompt],
+            request_id="teacache-coefficient-estimator",
             sampling_params=OmniDiffusionSamplingParams(
                 num_inference_steps=generate_kwargs.get("num_inference_steps", 20),
                 seed=generate_kwargs.get("seed", 42),

vllm_omni/diffusion/diffusion_engine.py

@@ -30,10 +30,10 @@ from vllm_omni.diffusion.registry import (
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
 from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
-from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -217,7 +217,7 @@ class DiffusionEngine:
             logger.warning("Output is None, returning empty OmniRequestOutput")
             return [
                 OmniRequestOutput.from_diffusion(
-                    request_id=request.request_ids[i] if i < len(request.request_ids) else "",
+                    request_id=request.request_id,
                     images=[],
                     prompt=prompt,
                     metrics={},
@@ -304,7 +304,7 @@ class DiffusionEngine:
         if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
             prompt = request.prompts[0]
-            request_id = request.request_ids[0] if request.request_ids else ""
+            request_id = request.request_id
 
             if is_text_output:
                 return [
@@ -368,10 +368,9 @@ class DiffusionEngine:
             # Split images based on num_outputs_per_prompt for each request
             results = []
             output_idx = 0
+            request_id = request.request_id
 
             for i, prompt in enumerate(request.prompts):
-                request_id = request.request_ids[i] if i < len(request.request_ids) else ""
-
                 # Get images for this request
                 num_outputs = request.sampling_params.num_outputs_per_prompt
                 start_idx = output_idx
@@ -467,17 +466,17 @@ class DiffusionEngine:
                 runner_output = self.execute_fn(sched_output)
             except Exception as exc:
                 logger.error(
-                    "Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
+                    "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
                 )
                 runner_output = BatchRunnerOutput.from_list(
                     [
                         RunnerOutput(
-                            req_id=req_id,
+                            request_id=request_id,
                             step_index=None,
                             finished=True,
                             result=DiffusionOutput(error=str(exc)),
                         )
-                        for req_id in sched_output.scheduled_req_ids
+                        for request_id in sched_output.scheduled_request_ids
                     ]
                 )
 
@@ -544,7 +543,7 @@ class DiffusionEngine:
     def _handle_finished_requests(
         self,
         finished_ids: set[str],
-        runner_output: RunnerOutput | None = None,
+        runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
     ):
         for rid in finished_ids:
@@ -553,7 +552,7 @@ class DiffusionEngine:
             if fut is None:
                 continue
             if runner_output is not None:
-                _output = runner_output.get_req_output(rid)
+                _output = runner_output.get_request_output(rid)
             else:
                 _output = None
             out = self._finalize_finished_request(rid, _output, missing_result_error)
@@ -579,11 +578,11 @@ class DiffusionEngine:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
             fut = self.main_loop.create_future()
-            sched_req_id = self.scheduler.add_request(request)
-            self._out_queue[sched_req_id] = fut
+            request_id = self.scheduler.add_request(request)
+            self._out_queue[request_id] = fut
             self._cv.notify_all()
 
-        return sched_req_id
+        return request_id
 
     async def get_result(self, request_id: str) -> DiffusionOutput:
         fut = self._out_queue.get(request_id)
@@ -599,37 +598,37 @@ class DiffusionEngine:
     async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         # No lock needed: add_request is already protected by self._cv, and
         # all executor calls are serialized inside the busy loop.
-        sched_req_id = self.add_request(request)
-        return await self.get_result(sched_req_id)
+        request_id = self.add_request(request)
+        return await self.get_result(request_id)
 
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         with self._rpc_lock:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
-            target_sched_req_id = self.scheduler.add_request(request)
+            target_request_id = self.scheduler.add_request(request)
 
             # keep scheduling and executing until the target request is finished
             while True:
                 self._process_aborts_queue()
                 sched_output = self.scheduler.schedule()
                 if sched_output.is_empty:
-                    if target_sched_req_id in sched_output.finished_req_ids:
-                        return self._finalize_finished_request(target_sched_req_id)
+                    if target_request_id in sched_output.finished_req_ids:
+                        return self._finalize_finished_request(target_request_id)
                     if not self.scheduler.has_requests():
                         raise RuntimeError("Diffusion scheduler has no runnable requests.")
                     continue
 
                 # NOTE: add_req_and_wait_for_response() is synchronous, will be only called
                 # within _dummy_run, only one request will be scheduled
-                sched_req_id = sched_output.scheduled_req_ids[0]
+                request_id = sched_output.scheduled_request_ids[0]
                 try:
                     runner_output = self.execute_fn(sched_output)
                 except EngineDeadError:
                     raise
                 except Exception as exc:
-                    logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
+                    logger.error("Execution failed for diffusion request %s", request_id, exc_info=True)
                     runner_output = RunnerOutput(
-                        req_id=sched_req_id,
+                        request_id=request_id,
                         step_index=None,
                         finished=True,
                         result=DiffusionOutput(error=str(exc)),
@@ -642,10 +641,10 @@ class DiffusionEngine:
                 # sync func should receive one result
                 if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
                     raise ValueError("Sync func should receive one result at one time")
-                if target_sched_req_id in finished_req_ids:
-                    runner_output = runner_output.get_req_output(target_sched_req_id)
+                if target_request_id in finished_req_ids:
+                    runner_output = runner_output.get_request_output(target_request_id)
                     return self._finalize_finished_request(
-                        target_sched_req_id,
+                        target_request_id,
                         runner_output=runner_output,
                         missing_result_error="Diffusion execution finished without a final output.",
                     )
@@ -697,7 +696,7 @@ class DiffusionEngine:
         num_frames = 2 if supports_audio_input or supports_audio_output(self.od_config.model_class_name) else 1
         req = OmniDiffusionRequest(
             prompts=[prompt],
-            request_ids=["dummy_req_id"],
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
             sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
@@ -904,34 +903,27 @@ class DiffusionEngine:
     def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
         request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
 
-        sched_req_ids: list[str] = []
         for request_id in dict.fromkeys(request_ids):
-            sched_req_id = self.scheduler.get_sched_req_id(request_id)
-            if sched_req_id is not None:
-                sched_req_ids.append(sched_req_id)
-
-        for sched_req_id in dict.fromkeys(sched_req_ids):
-            if self.scheduler.get_request_state(sched_req_id) is not None:
-                self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+            if self.scheduler.get_request_state(request_id) is not None:
+                self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
 
     def _finalize_finished_request(
         self,
-        sched_req_id: str,
+        request_id: str,
         runner_output: RunnerOutput | None = None,
         missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
     ) -> DiffusionOutput:
-        state = self.scheduler.get_request_state(sched_req_id)
-        popped_state = self.scheduler.pop_request_state(sched_req_id)
+        state = self.scheduler.get_request_state(request_id)
+        popped_state = self.scheduler.pop_request_state(request_id)
         state = state or popped_state
 
         if state is None:
-            raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
+            raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")
 
         if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
-            request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
             return DiffusionOutput(
                 aborted=True,
-                abort_message=f"Request {request_id} aborted.",
+                abort_message=f"Request {state.req.request_id} aborted.",
             )
 
         if runner_output is not None and runner_output.result is not None:

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -311,7 +311,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise RuntimeError(f"Unexpected response type for execute_request: {type(result)!r}")
 
         return RunnerOutput(
-            req_id=new_req.sched_req_id,
+            request_id=new_req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -334,9 +334,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
         # failed multi-request step batches.
         if isinstance(result, DiffusionOutput):
-            req_id = scheduler_output.scheduled_req_ids[0] if scheduler_output.scheduled_req_ids else ""
+            request_id = scheduler_output.scheduled_request_ids[0] if scheduler_output.scheduled_request_ids else ""
             return RunnerOutput(
-                req_id=req_id,
+                request_id=request_id,
                 step_index=None,
                 finished=True,
                 result=result,

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -128,7 +128,6 @@ class InlineStageDiffusionClient(StageClientBase):
             request = OmniDiffusionRequest(
                 prompts=[prompt],
                 sampling_params=sampling_params,
-                request_ids=[request_id],
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
             )
@@ -187,7 +186,6 @@ class InlineStageDiffusionClient(StageClientBase):
             request = OmniDiffusionRequest(
                 prompts=prompts,
                 sampling_params=sampling_params,
-                request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
             )

vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py

@@ -247,8 +247,10 @@ class ErnieImagePipeline(
 
     @staticmethod
     def _is_warmup_request(req: OmniDiffusionRequest) -> bool:
-        request_ids = getattr(req, "request_ids", None) or ()
-        return len(request_ids) == 1 and request_ids[0] == "dummy_req_id"
+        is_dummy_run = getattr(req, "is_dummy_run", None)
+        if callable(is_dummy_run):
+            return bool(is_dummy_run())
+        return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
 
     @staticmethod
     def _should_apply_pe(req: OmniDiffusionRequest) -> bool:

vllm_omni/diffusion/models/glm_image/glm_image_transformer.py

@@ -15,6 +15,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
+    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -33,6 +34,9 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = init_logger(__name__)
 
 
@@ -230,11 +234,24 @@ class GlmImagePrepare(nn.Module):
 class GlmImageAdaLayerNormZero(nn.Module):
     """Adaptive LayerNorm with zero initialization for both image and text streams."""
 
-    def __init__(self, embedding_dim: int, dim: int) -> None:
+    def __init__(
+        self,
+        embedding_dim: int,
+        dim: int,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
         self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True)
+        self.linear = ReplicatedLinear(
+            embedding_dim,
+            12 * dim,
+            bias=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
 
     def forward(
         self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
@@ -244,6 +261,8 @@ class GlmImageAdaLayerNormZero(nn.Module):
         norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype)
 
         emb = self.linear(temb)
+        if isinstance(emb, tuple):
+            emb = emb[0]
         (
             shift_msa,
             c_shift_msa,
@@ -286,14 +305,25 @@ class GlmImageAdaLayerNormContinuous(nn.Module):
         elementwise_affine: bool = True,
         eps: float = 1e-5,
         bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
-        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        self.linear = ReplicatedLinear(
+            conditioning_embedding_dim,
+            embedding_dim * 2,
+            bias=bias,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
         self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
 
     def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
         # NO SiLU here
         emb = self.linear(conditioning_embedding.to(x.dtype))
+        if isinstance(emb, tuple):
+            emb = emb[0]
         scale, shift = torch.chunk(emb, 2, dim=1)
         x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
         return x
@@ -469,6 +499,7 @@ class GlmImageAttention(nn.Module):
         out_bias: bool = True,
         eps: float = 1e-5,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.dim = dim
@@ -485,6 +516,7 @@ class GlmImageAttention(nn.Module):
             bias=True,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
         )
 
         # QK normalization (LayerNorm, not RMSNorm for GLM-Image)
@@ -501,6 +533,7 @@ class GlmImageAttention(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
                 ),
                 nn.Dropout(0.0),
             ]
@@ -554,7 +587,7 @@ class GlmImageAttention(nn.Module):
         hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
         # QKV projection
-        qkv = self.to_qkv(hidden_states_combined)
+        qkv = self.to_qkv(hidden_states_combined.contiguous())
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
         query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
@@ -604,8 +637,12 @@ class GlmImageAttention(nn.Module):
             # Project combined [text, image] outputs, then split.
             # This keeps SP numerically aligned with the non-SP path.
             joint_hidden_states_out = joint_hidden_states_out.flatten(2, 3).to(dtype)
+            # Contiguous for FP8/W4A16 quantized RowParallelLinear
             for module in self.to_out:
-                joint_hidden_states_out = module(joint_hidden_states_out)
+                if isinstance(module, RowParallelLinear):
+                    joint_hidden_states_out = module(joint_hidden_states_out.contiguous())
+                else:
+                    joint_hidden_states_out = module(joint_hidden_states_out)
 
             encoder_hidden_states_out = joint_hidden_states_out[:, :text_seq_length, :]
             hidden_states_out = joint_hidden_states_out[:, text_seq_length:, :]
@@ -645,7 +682,10 @@ class GlmImageAttention(nn.Module):
 
             # Output projection
             for module in self.to_out:
-                hidden_states_out = module(hidden_states_out)
+                if isinstance(module, RowParallelLinear):
+                    hidden_states_out = module(hidden_states_out.contiguous())
+                else:
+                    hidden_states_out = module(hidden_states_out)
 
             # Split back to text and image
             encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :]
@@ -663,6 +703,7 @@ class ColumnParallelGELU(nn.Module):
         approximate: str = "none",
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
@@ -672,6 +713,7 @@ class ColumnParallelGELU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.proj",
         )
         self.approximate = approximate
 
@@ -688,6 +730,7 @@ class ColumnParallelSiLU(nn.Module):
         *,
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
@@ -697,6 +740,7 @@ class ColumnParallelSiLU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.proj",
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -714,6 +758,7 @@ class GlmImageFeedForward(nn.Module):
         bias: bool = True,
         activation_fn: str = "gelu",
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         inner_dim = inner_dim or int(dim * mult)
@@ -721,7 +766,7 @@ class GlmImageFeedForward(nn.Module):
 
         if activation_fn == "linear-silu":
             layers: list[nn.Module] = [
-                ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config),
+                ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0"),
                 nn.Identity(),
                 RowParallelLinear(
                     inner_dim,
@@ -730,12 +775,20 @@ class GlmImageFeedForward(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.net.2",
                 ),
             ]
         else:
             approximate = "tanh" if activation_fn == "gelu-approximate" else "none"
             layers = [
-                ColumnParallelGELU(dim, inner_dim, approximate=approximate, bias=bias, quant_config=quant_config),
+                ColumnParallelGELU(
+                    dim,
+                    inner_dim,
+                    approximate=approximate,
+                    bias=bias,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.0",
+                ),
                 nn.Identity(),
                 RowParallelLinear(
                     inner_dim,
@@ -744,6 +797,7 @@ class GlmImageFeedForward(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=f"{prefix}.net.2",
                 ),
             ]
 
@@ -751,7 +805,12 @@ class GlmImageFeedForward(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         for module in self.net:
-            hidden_states = module(hidden_states)
+            if isinstance(module, ColumnParallelLinear):
+                hidden_states, _ = module(hidden_states)
+            elif isinstance(module, RowParallelLinear):
+                hidden_states = module(hidden_states.contiguous())
+            else:
+                hidden_states = module(hidden_states)
         return hidden_states
 
 
@@ -765,19 +824,21 @@ class GlmImageTransformerBlock(nn.Module):
         attention_head_dim: int = 40,
         time_embed_dim: int = 512,
         ffn_hidden_dim: int | None = None,
-        quant_config: "QuantizationConfig | None" = None,
         parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
         # 1. Attention with AdaLN
-        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim, quant_config=quant_config, prefix=f"{prefix}.norm1")
         self.attn1 = GlmImageAttention(
             dim=dim,
             num_heads=num_attention_heads,
             head_dim=attention_head_dim,
-            quant_config=quant_config,
             parallel_config=parallel_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn1",
         )
 
         # 2. Feedforward
@@ -789,6 +850,7 @@ class GlmImageTransformerBlock(nn.Module):
             inner_dim=ffn_hidden_dim,
             activation_fn="gelu-approximate",
             quant_config=quant_config,
+            prefix=f"{prefix}.ff",
         )
 
     def forward(
@@ -959,6 +1021,7 @@ class GlmImageTransformer2DModel(CachedTransformer):
             inner_dim=inner_dim,
             activation_fn="gelu",
             quant_config=quant_config,
+            prefix="glyph_projector",
         )
         self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim)
         self.prior_projector = GlmImageFeedForward(
@@ -967,6 +1030,7 @@ class GlmImageTransformer2DModel(CachedTransformer):
             inner_dim=inner_dim,
             activation_fn="linear-silu",
             quant_config=quant_config,
+            prefix="prior_projector",
         )
 
         # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan)
@@ -988,15 +1052,20 @@ class GlmImageTransformer2DModel(CachedTransformer):
                     attention_head_dim,
                     time_embed_dim,
                     ffn_hidden_dim=ffn_hidden_dim,
-                    quant_config=quant_config,
                     parallel_config=self.parallel_config,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
         # 4. Output projection
-        self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False)
+        # Final modulation feeds proj_out; quant_config is NOT applied here
+        # to avoid precision degradation in the final projection layer.
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, time_embed_dim, elementwise_affine=False, quant_config=None, prefix="norm_out"
+        )
         self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True)
 
     def forward(

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -301,10 +301,7 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         # Load transformer (DiT)
         logger.info("Loading GlmImageTransformer2DModel (DiT)...")
-        self.transformer = GlmImageTransformer2DModel(
-            od_config=od_config,
-            quant_config=od_config.quantization_config,
-        )
+        self.transformer = GlmImageTransformer2DModel(od_config=od_config, quant_config=od_config.quantization_config)
 
         # Weight sources for DiT loading
         self.weights_sources = [
@@ -689,15 +686,11 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
         # NOTE: DiffusionEngine does an internal warmup "dummy run" during
-        # initialization. That request has no request_id and does not carry
+        # initialization. That request carries the fixed dummy request_id and does not carry
         # Stage-0 (AR) outputs via req.extra. For GLM-Image, we allow that
         # specific warmup request to proceed by synthesizing minimal prior
         # tokens, while still raising a clear error for real requests.
-        is_dummy_warmup = (
-            not getattr(req, "request_id", None)
-            and prompt == "dummy run"
-            and (req.sampling_params.num_inference_steps == 1)
-        )
+        is_dummy_warmup = req.is_dummy_run()
 
         # Get pre-computed prompt embeddings if provided
         if isinstance(first_prompt, str):

vllm_omni/diffusion/request.py

@@ -3,10 +3,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import random
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
+DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id"
+
 
 @dataclass
 class OmniDiffusionRequest:
@@ -24,13 +26,14 @@ class OmniDiffusionRequest:
 
     prompts: list[OmniPromptType]  # Actually supporting str-based prompts
     sampling_params: OmniDiffusionSamplingParams
-
-    request_ids: list[str] = field(default_factory=list)
-    request_id: str | None = None
+    request_id: str
     kv_sender_info: dict | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
+        if not isinstance(self.request_id, str) or not self.request_id:
+            raise ValueError("OmniDiffusionRequest.request_id must be a non-empty string.")
+
         # When neither a generator nor a seed is provided, assign a random seed
         # so that all ranks derive the same generator state.
         if self.sampling_params.generator is None and self.sampling_params.seed is None:
@@ -56,3 +59,10 @@ class OmniDiffusionRequest:
         # so downstream code always has a valid value.
         if self.sampling_params.guidance_scale_2 is None:
             self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale
+
+    def is_dummy_run(self) -> bool:
+        return self.is_dummy_run_request_id(self.request_id)
+
+    @classmethod
+    def is_dummy_run_request_id(cls, request_id: str | None) -> bool:
+        return request_id == DUMMY_DIFFUSION_REQUEST_ID

vllm_omni/diffusion/sched/base_scheduler.py

@@ -46,7 +46,6 @@ class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self.od_config: OmniDiffusionConfig | None = None
         self._request_states: dict[str, DiffusionRequestState] = {}
-        self._request_id_to_sched_req_id: dict[str, str] = {}
         self._step_id: int = 0
         self._waiting: deque[str] = deque()
         self._running: list[str] = []
@@ -57,7 +56,6 @@ class _BaseScheduler(SchedulerInterface):
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
         self._request_states.clear()
-        self._request_id_to_sched_req_id.clear()
         self._step_id = 0
         self._waiting.clear()
         self._running.clear()
@@ -71,32 +69,31 @@ class _BaseScheduler(SchedulerInterface):
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
-        sched_req_id = self._make_sched_req_id(request)
-        return self._add_request_with_sched_req_id(sched_req_id, request)
+        return self._add_request_with_request_id(request.request_id, request)
 
-    def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
-        state = self._make_request_state(sched_req_id, request)
-        request_ids = self._request_ids_for_mapping(request)
-        self._register_request_ids(request_ids, sched_req_id)
-        self._request_states[sched_req_id] = state
-        self._waiting.append(sched_req_id)
-        logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, sched_req_id, len(self._waiting))
-        return sched_req_id
+    def _add_request_with_request_id(self, request_id: str, request: OmniDiffusionRequest) -> str:
+        if request_id in self._request_states:
+            raise ValueError(f"request_id {request_id!r} is already active.")
+        state = self._make_request_state(request_id, request)
+        self._request_states[request_id] = state
+        self._waiting.append(request_id)
+        logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting))
+        return request_id
 
     def schedule(self) -> DiffusionSchedulerOutput:
         scheduled_new_reqs: list[NewRequestData] = []
-        scheduled_cached_req_ids: list[str] = []
+        scheduled_cached_request_ids: list[str] = []
 
         # First, schedule the RUNNING request(s)
-        for sched_req_id in self._running:
-            state = self._request_states.get(sched_req_id)
+        for request_id in self._running:
+            state = self._request_states.get(request_id)
             if state is not None:
-                scheduled_cached_req_ids.append(sched_req_id)
+                scheduled_cached_request_ids.append(request_id)
 
         # Second, schedule WAITING requests while capacity remains.
         while self._waiting and len(self._running) < self.max_num_running_reqs:
-            sched_req_id = self._waiting[0]
-            state = self._request_states.get(sched_req_id)
+            request_id = self._waiting[0]
+            state = self._request_states.get(request_id)
             if state is None:
                 self._waiting.popleft()
                 continue
@@ -108,16 +105,16 @@ class _BaseScheduler(SchedulerInterface):
             if not self._running:
                 self._running_sampling_params_key = state.sampling_params_key
             state.status = DiffusionRequestStatus.RUNNING
-            self._running.append(sched_req_id)
+            self._running.append(request_id)
             if was_new_request:
                 scheduled_new_reqs.append(NewRequestData.from_state(state))
             else:
-                scheduled_cached_req_ids.append(sched_req_id)
+                scheduled_cached_request_ids.append(request_id)
 
         scheduler_output = DiffusionSchedulerOutput(
             step_id=self._step_id,
             scheduled_new_reqs=scheduled_new_reqs,
-            scheduled_cached_reqs=CachedRequestData(sched_req_ids=scheduled_cached_req_ids),
+            scheduled_cached_reqs=CachedRequestData(request_ids=scheduled_cached_request_ids),
             finished_req_ids=set(self._finished_req_ids),
             num_running_reqs=len(self._running),
             num_waiting_reqs=len(self._waiting),
@@ -131,40 +128,33 @@ class _BaseScheduler(SchedulerInterface):
     def has_requests(self) -> bool:
         return bool(self._waiting or self._running)
 
-    def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
-        return self._request_states.get(sched_req_id)
+    def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
+        return self._request_states.get(request_id)
 
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        return self._request_id_to_sched_req_id.get(request_id)
+    def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
+        self._pop_extra_request_state(request_id)
+        return self._request_states.pop(request_id, None)
 
-    def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
-        self._pop_extra_request_state(sched_req_id)
-        state = self._request_states.pop(sched_req_id, None)
-        if state is not None:
-            self._unregister_request_ids(self._request_ids_for_mapping(state.req), sched_req_id)
-        return state
-
-    def preempt_request(self, sched_req_id: str) -> bool:
-        if sched_req_id not in self._request_states:
+    def preempt_request(self, request_id: str) -> bool:
+        if request_id not in self._request_states:
             return False
-        if sched_req_id in self._running:
-            self._running.remove(sched_req_id)
+        if request_id in self._running:
+            self._running.remove(request_id)
             if not self._running:
                 self._running_sampling_params_key = None
-            self._waiting.appendleft(sched_req_id)
-            self._request_states[sched_req_id].status = DiffusionRequestStatus.PREEMPTED
+            self._waiting.appendleft(request_id)
+            self._request_states[request_id].status = DiffusionRequestStatus.PREEMPTED
             return True
         return False
 
-    def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
+    def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
         assert DiffusionRequestStatus.is_finished(status)
-        if isinstance(sched_req_ids, str):
-            sched_req_ids = [sched_req_ids]
-        self._finish_requests({sched_req_id: status for sched_req_id in sched_req_ids})
+        if isinstance(request_ids, str):
+            request_ids = [request_ids]
+        self._finish_requests({request_id: status for request_id in request_ids})
 
     def close(self) -> None:
         self._request_states.clear()
-        self._request_id_to_sched_req_id.clear()
         self._waiting.clear()
         self._running.clear()
         self._running_sampling_params_key = None
@@ -183,33 +173,31 @@ class _BaseScheduler(SchedulerInterface):
         running_to_remove: set[str] = set()
         waiting_to_remove: set[str] = set()
 
-        for sched_req_id, status in statuses.items():
+        for request_id, status in statuses.items():
             assert DiffusionRequestStatus.is_finished(status)
-            state = self._request_states.get(sched_req_id)
+            state = self._request_states.get(request_id)
             if state is None or state.is_finished():
                 continue
 
-            finished_req_ids.add(sched_req_id)
-            if sched_req_id in self._running:
-                running_to_remove.add(sched_req_id)
-            if sched_req_id in self._waiting:
-                waiting_to_remove.add(sched_req_id)
+            finished_req_ids.add(request_id)
+            if request_id in self._running:
+                running_to_remove.add(request_id)
+            if request_id in self._waiting:
+                waiting_to_remove.add(request_id)
 
         if running_to_remove:
-            self._running = [sched_req_id for sched_req_id in self._running if sched_req_id not in running_to_remove]
+            self._running = [request_id for request_id in self._running if request_id not in running_to_remove]
             if not self._running:
                 self._running_sampling_params_key = None
         if waiting_to_remove:
-            self._waiting = deque(
-                sched_req_id for sched_req_id in self._waiting if sched_req_id not in waiting_to_remove
-            )
+            self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove)
 
-        for sched_req_id in finished_req_ids:
-            state = self._request_states[sched_req_id]
-            status = statuses[sched_req_id]
+        for request_id in finished_req_ids:
+            state = self._request_states[request_id]
+            status = statuses[request_id]
             state.status = status
             if status == DiffusionRequestStatus.FINISHED_ERROR:
-                state.error = None if errors is None else errors.get(sched_req_id)
+                state.error = None if errors is None else errors.get(request_id)
             else:
                 state.error = None
 
@@ -227,7 +215,7 @@ class _BaseScheduler(SchedulerInterface):
         # marked finished at that point, but we still need to surface its id
         # in this update so the engine can observe the terminal state.
         finished_req_ids = {
-            sched_req_id for sched_req_id in sched_output.scheduled_req_ids if sched_req_id in self._finished_req_ids
+            request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids
         }
         finished_req_ids |= self._finish_requests(statuses, errors)
         return finished_req_ids
@@ -235,12 +223,12 @@ class _BaseScheduler(SchedulerInterface):
     def _reset_scheduler_state(self) -> None:
         """Reset subclass-owned state during initialize()/close()."""
 
-    def _pop_extra_request_state(self, sched_req_id: str) -> None:
+    def _pop_extra_request_state(self, request_id: str) -> None:
         """Remove subclass-owned per-request state before popping request state."""
 
-    def _make_request_state(self, sched_req_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
+    def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
         return DiffusionRequestState(
-            sched_req_id=sched_req_id,
+            request_id=request_id,
             req=request,
             sampling_params_key=get_sampling_params_key(request),
         )
@@ -258,24 +246,3 @@ class _BaseScheduler(SchedulerInterface):
         state = self._request_states.get(self._running[0])
         self._running_sampling_params_key = None if state is None else state.sampling_params_key
         return self._running_sampling_params_key
-
-    def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
-        unique_request_ids = list(dict.fromkeys(request_ids))
-        for request_id in unique_request_ids:
-            existing = self._request_id_to_sched_req_id.get(request_id)
-            if existing is not None and existing != sched_req_id:
-                raise ValueError(f"request_id {request_id!r} is already mapped to active sched_req_id {existing!r}.")
-        for request_id in unique_request_ids:
-            self._request_id_to_sched_req_id[request_id] = sched_req_id
-
-    def _unregister_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
-        for request_id in request_ids:
-            if self._request_id_to_sched_req_id.get(request_id) == sched_req_id:
-                self._request_id_to_sched_req_id.pop(request_id, None)
-
-    @staticmethod
-    def _request_ids_for_mapping(request: OmniDiffusionRequest) -> list[str]:
-        request_ids = list(request.request_ids)
-        if request_id := getattr(request, "request_id", None):
-            request_ids.append(request_id)
-        return list(dict.fromkeys(request_ids))

vllm_omni/diffusion/sched/interface.py

@@ -4,22 +4,17 @@
 from __future__ import annotations
 
 import enum
-import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from vllm.logger import init_logger
-
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.worker.utils import RunnerOutput
 
-logger = init_logger(__name__)
-
 
 class DiffusionRequestStatus(enum.IntEnum):
     """Request status tracked by diffusion scheduler."""
@@ -75,10 +70,7 @@ class SamplingParamsKey:
 class DiffusionRequestState:
     """Scheduler-owned state for one queued OmniDiffusionRequest."""
 
-    # Unique scheduler-owned request ID.
-    # NOTE: This identifies one OmniDiffusionRequest, which may contain multiple request_ids.
-    # TODO: Align this with OmniDiffusionRequest.request_ids once scheduler batching is supported.
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
     sampling_params_key: SamplingParamsKey | None = None
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
@@ -92,23 +84,23 @@ class DiffusionRequestState:
 class NewRequestData:
     """Full request payload for a newly scheduled diffusion request."""
 
-    sched_req_id: str
+    request_id: str
     req: OmniDiffusionRequest
 
     @classmethod
     def from_state(cls, state: DiffusionRequestState) -> NewRequestData:
-        return cls(sched_req_id=state.sched_req_id, req=state.req)
+        return cls(request_id=state.request_id, req=state.req)
 
 
 @dataclass
 class CachedRequestData:
-    """Cached diffusion requests that only need their scheduler ids resent."""
+    """Cached diffusion requests that only need their request ids resent."""
 
-    sched_req_ids: list[str]
+    request_ids: list[str]
 
     @classmethod
     def make_empty(cls) -> CachedRequestData:
-        return cls(sched_req_ids=[])
+        return cls(request_ids=[])
 
 
 @dataclass
@@ -123,21 +115,18 @@ class DiffusionSchedulerOutput:
     num_waiting_reqs: int
 
     @cached_property
-    def scheduled_req_ids(self) -> list[str]:
+    def scheduled_request_ids(self) -> list[str]:
         """
         All scheduled request ids in this cycle, including both new and cached ones.
-        NOTE:
-            This id is generated and owned by the scheduler,
-            and may be different from the OmniDiffusionRequest.request_ids.
         """
         return [
-            *(req.sched_req_id for req in self.scheduled_new_reqs),
-            *self.scheduled_cached_reqs.sched_req_ids,
+            *(req.request_id for req in self.scheduled_new_reqs),
+            *self.scheduled_cached_reqs.request_ids,
         ]
 
     @property
     def num_scheduled_reqs(self) -> int:
-        return len(self.scheduled_req_ids)
+        return len(self.scheduled_request_ids)
 
     @property
     def is_empty(self) -> bool:
@@ -147,25 +136,6 @@ class DiffusionSchedulerOutput:
 class SchedulerInterface(ABC):
     """Abstract lifecycle contract for diffusion schedulers."""
 
-    def _make_sched_req_id(self, request: OmniDiffusionRequest) -> str:
-        """
-        Generate a unique scheduler request ID for the given request.
-            The default implementation uses the first request_id from the request if available,
-            otherwise generates a random one.
-        """
-        if request.request_ids:
-            base = request.request_ids[0]
-        else:
-            logger.warning("Request has no request_ids, generating a random one. Request: %s", request)
-            base = f"req_{uuid.uuid4().hex[:8]}"
-
-        sched_req_id = base
-        suffix = 1
-        while self.get_request_state(sched_req_id) is not None:
-            sched_req_id = f"{base}#{suffix}"
-            suffix += 1
-        return sched_req_id
-
     @abstractmethod
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         """Initialize or reset scheduler state."""
@@ -183,7 +153,7 @@ class SchedulerInterface(ABC):
         """Update scheduler state from executor output."""
 
     @abstractmethod
-    def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
+    def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
         """Return request state if present."""
 
     @abstractmethod
@@ -191,19 +161,15 @@ class SchedulerInterface(ABC):
         """Return whether the scheduler still owns runnable requests."""
 
     @abstractmethod
-    def get_sched_req_id(self, request_id: str) -> str | None:
-        """Resolve a public request_id to the active scheduler request id."""
-
-    @abstractmethod
-    def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
+    def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
         """Remove and return request state if present."""
 
     @abstractmethod
-    def preempt_request(self, sched_req_id: str) -> bool:
+    def preempt_request(self, request_id: str) -> bool:
         """Preempt a running request back to waiting."""
 
     @abstractmethod
-    def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
+    def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
         """Mark one or more requests finished."""
 
     @abstractmethod

vllm_omni/diffusion/sched/request_scheduler.py

@@ -26,26 +26,26 @@ class RequestScheduler(_BaseScheduler):
         return super().schedule()
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
-        scheduled_req_ids = sched_output.scheduled_req_ids
-        if not scheduled_req_ids:
+        scheduled_request_ids = sched_output.scheduled_request_ids
+        if not scheduled_request_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        for sched_req_id in scheduled_req_ids:
-            state = self._request_states.get(sched_req_id)
+        for request_id in scheduled_request_ids:
+            state = self._request_states.get(request_id)
             if state is None or state.is_finished():
                 continue
-            req_output = output.get_req_output(sched_req_id)
+            req_output = output.get_request_output(request_id)
             result = req_output.result if req_output is not None else None
             if result is None:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "No output result"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "No output result"
             elif result.error:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = result.error
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = result.error
             else:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
-                terminal_errors[sched_req_id] = None
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                terminal_errors[request_id] = None
 
         return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)

vllm_omni/diffusion/sched/step_scheduler.py

@@ -38,84 +38,84 @@ class StepScheduler(_BaseScheduler):
         self._request_progress.clear()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
-        sched_req_id = self._make_sched_req_id(request)
+        request_id = request.request_id
         total_steps = self._get_total_steps(request)
         if total_steps <= 0:
-            raise ValueError(f"Diffusion request {sched_req_id} must have positive total_steps, got {total_steps}")
+            raise ValueError(f"Diffusion request {request_id} must have positive total_steps, got {total_steps}")
 
         current_step = request.sampling_params.step_index or 0
         if current_step < 0 or current_step >= total_steps:
             raise ValueError(
-                f"Diffusion request {sched_req_id} has invalid initial step_index {current_step} "
+                f"Diffusion request {request_id} has invalid initial step_index {current_step} "
                 f"for total_steps={total_steps}"
             )
 
         request.sampling_params.step_index = current_step
-        sched_req_id = self._add_request_with_sched_req_id(sched_req_id, request)
-        self._request_progress[sched_req_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
+        request_id = self._add_request_with_request_id(request_id, request)
+        self._request_progress[request_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
         logger.debug(
             "StepScheduler add_request: %s (step=%d/%d, waiting=%d)",
-            sched_req_id,
+            request_id,
             current_step,
             total_steps,
             len(self._waiting),
         )
-        return sched_req_id
+        return request_id
 
     def schedule(self) -> DiffusionSchedulerOutput:
         return super().schedule()
 
     def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
-        scheduled_req_ids = sched_output.scheduled_req_ids
-        if not scheduled_req_ids:
+        scheduled_request_ids = sched_output.scheduled_request_ids
+        if not scheduled_request_ids:
             return set()
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        for sched_req_id in scheduled_req_ids:
-            state = self._request_states.get(sched_req_id)
-            progress = self._request_progress.get(sched_req_id)
+        for request_id in scheduled_request_ids:
+            state = self._request_states.get(request_id)
+            progress = self._request_progress.get(request_id)
             if state is None or progress is None or state.is_finished():
                 continue
-            req_output = output.get_req_output(sched_req_id)
+            req_output = output.get_request_output(request_id)
             if req_output is None:
                 logger.warning(
                     "No RunnerOutput for request %s, treating as error",
-                    sched_req_id,
+                    request_id,
                 )
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "No output for request"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "No output for request"
                 continue
 
             req_result = req_output.result
             output_error = req_result.error if req_result is not None else None
             if output_error is not None:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = output_error
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = output_error
                 continue
 
             if req_output.step_index is None:
                 logger.warning(
                     "Received RunnerOutput with no step_index for request %s, treating as error",
-                    sched_req_id,
+                    request_id,
                 )
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = "Missing step_index in RunnerOutput"
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[request_id] = "Missing step_index in RunnerOutput"
                 continue
 
             # We assume that the decoding stage is executed immediately after the denoising stage completes.
             progress.current_step = req_output.step_index
             state.req.sampling_params.step_index = req_output.step_index
             if req_output.finished:
-                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
-                terminal_errors[sched_req_id] = None
+                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                terminal_errors[request_id] = None
             else:
                 state.error = None
 
         return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
 
-    def _pop_extra_request_state(self, sched_req_id: str) -> None:
-        self._request_progress.pop(sched_req_id, None)
+    def _pop_extra_request_state(self, request_id: str) -> None:
+        self._request_progress.pop(request_id, None)
 
     def _get_total_steps(self, request: OmniDiffusionRequest) -> int:
         sampling = request.sampling_params

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -159,7 +159,6 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=[prompt],
             sampling_params=sampling_params,
-            request_ids=[request_id],
             request_id=request_id,
             kv_sender_info=kv_sender_info,
         )
@@ -189,7 +188,6 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=prompts,
             sampling_params=sampling_params,
-            request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
             request_id=request_id,
             kv_sender_info=kv_sender_info,
         )

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -322,50 +322,44 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self, scheduler_output: DiffusionSchedulerOutput
     ) -> tuple[list[DiffusionRequestState], list[str]]:
         """Step-before update: cleanup finished requests and get/create one running state."""
-        for req_id in scheduler_output.finished_req_ids:
-            self.state_cache.pop(req_id, None)
+        for request_id in scheduler_output.finished_req_ids:
+            self.state_cache.pop(request_id, None)
 
         resolved: list[DiffusionRequestState] = []
-        new_req_id: list[str] = []
+        new_request_ids: list[str] = []
         try:
             # process new requests
             for sched_new_req in scheduler_output.scheduled_new_reqs:
-                # new_req_data = scheduler_output.scheduled_new_reqs[0]
-                req_id = sched_new_req.sched_req_id
+                request_id = sched_new_req.request_id
                 req = sched_new_req.req
-                new_req_id.append(req_id)
-                if req_id in self.state_cache:
-                    raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
-                request_ids = req.request_ids or [req_id]
-                if len(request_ids) != len(req.prompts):
-                    raise ValueError(
-                        f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
-                    )
+                new_request_ids.append(request_id)
+                if request_id in self.state_cache:
+                    raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.")
                 new_state = DiffusionRequestState(
-                    req_id=req_id,
+                    request_id=request_id,
                     sampling=copy.deepcopy(req.sampling_params),
                     prompts=req.prompts,
                 )
-                self.state_cache[req_id] = new_state
+                self.state_cache[request_id] = new_state
                 resolved.append(new_state)
 
             # process cached requests
-            for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
-                state = self.state_cache.get(req_id)
+            for request_id in scheduler_output.scheduled_cached_reqs.request_ids:
+                state = self.state_cache.get(request_id)
                 if state is None:
-                    raise ValueError(f"Missing cached state for request {req_id}.")
+                    raise ValueError(f"Missing cached state for request {request_id}.")
                 resolved.append(state)
         except Exception:
-            for req_id in new_req_id:
-                self.state_cache.pop(req_id, None)
+            for request_id in new_request_ids:
+                self.state_cache.pop(request_id, None)
             raise
 
-        return resolved, new_req_id
+        return resolved, new_request_ids
 
     def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
         # process new reqs
         for state in states:
-            if state.req_id in new_request_ids:
+            if state.request_id in new_request_ids:
                 # set generator
                 if state.sampling.generator is None and state.sampling.seed is not None:
                     if state.sampling.generator_device is not None:
@@ -407,7 +401,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         for state in states:
             if interrupted or state.denoise_completed:
-                self.state_cache.pop(state.req_id, None)
+                self.state_cache.pop(state.request_id, None)
 
     def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
         model_state = getattr(self, "model_state", None)
@@ -449,7 +443,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     for state in states:
                         runner_output_list.append(
                             RunnerOutput(
-                                req_id=state.req_id,
+                                request_id=state.request_id,
                                 step_index=state.step_index,
                                 finished=True,
                                 result=DiffusionOutput(error="stepwise denoise interrupted"),
@@ -470,7 +464,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             result = None
                         runner_output_list.append(
                             RunnerOutput(
-                                req_id=req.req_id,
+                                request_id=req.request_id,
                                 step_index=req.step_index,
                                 finished=req.denoise_completed,
                                 result=result,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -189,7 +189,7 @@ class DiffusionWorker:
         self.lora_manager: DiffusionLoRAManager | None = None
         # Worker-side cache of (lora_request, lora_scale) per scheduled
         # request id. Used by step mode to recover LoRA identity for cached
-        # requests, which only carry their sched_req_id in subsequent ticks.
+        # requests, which only carry their request_id in subsequent ticks.
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
@@ -391,12 +391,12 @@ class DiffusionWorker:
         homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
         so any scheduled request id resolves to the active LoRA identity.
         """
-        for sched_req_id in scheduler_output.finished_req_ids:
-            self._step_lora_state.pop(sched_req_id, None)
+        for request_id in scheduler_output.finished_req_ids:
+            self._step_lora_state.pop(request_id, None)
 
         for new_req in scheduler_output.scheduled_new_reqs:
             sampling = new_req.req.sampling_params
-            self._step_lora_state[new_req.sched_req_id] = (
+            self._step_lora_state[new_req.request_id] = (
                 sampling.lora_request,
                 sampling.lora_scale,
             )
@@ -406,8 +406,8 @@ class DiffusionWorker:
 
         lora_request: LoRARequest | None = None
         lora_scale = 1.0
-        for sched_req_id in scheduler_output.scheduled_req_ids:
-            entry = self._step_lora_state.get(sched_req_id)
+        for request_id in scheduler_output.scheduled_request_ids:
+            entry = self._step_lora_state.get(request_id)
             if entry is not None:
                 lora_request, lora_scale = entry
                 break

vllm_omni/diffusion/worker/input_batch.py

@@ -81,8 +81,8 @@ def _select_states(
     return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
 
 
-def _prepare_req_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
-    return [state.req_id for state in states]
+def _prepare_request_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
+    return [state.request_id for state in states]
 
 
 def _prepare_prompt_field_on_state(
@@ -192,11 +192,13 @@ def _get_request_prompt_seq_lens(
         mask_attr=mask_attr,
     )
     if embeds is None:
-        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
 
     if mask is not None:
         if mask.shape[0] != embeds.shape[0]:
-            raise ValueError(f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.req_id}.")
+            raise ValueError(
+                f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.request_id}."
+            )
         return _get_seq_lens_from_mask(mask)
 
     seq_lens = getattr(state, seq_lens_attr)
@@ -220,7 +222,7 @@ def _prepare_request_prompt_field(
         mask_attr=mask_attr,
     )
     if embeds is None:
-        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
 
     actual_seq_lens = _get_request_prompt_seq_lens(
         state,
@@ -248,7 +250,7 @@ def _prepare_request_prompt_field(
     if current_seq_len > target_seq_len:
         if max_actual_seq_len > target_seq_len:
             raise ValueError(
-                f"{embeds_attr} for request {state.req_id} requires seq_len "
+                f"{embeds_attr} for request {state.request_id} requires seq_len "
                 f"{max_actual_seq_len}, got target {target_seq_len}."
             )
         return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
@@ -349,7 +351,7 @@ def _require_state_latents(
 ) -> torch.Tensor:
     latents = state.latents
     if latents is None:
-        raise ValueError(f"Request {state.req_id} has no latents while preparing {for_field}.")
+        raise ValueError(f"Request {state.request_id} has no latents while preparing {for_field}.")
     return latents
 
 
@@ -526,12 +528,12 @@ def _prepare_negative_prompt_embeds(
 
 def _same_composition(
     cached_batch: InputBatch | None,
-    req_ids: list[str],
+    request_ids: list[str],
     idx_mapping_np: np.ndarray,
 ) -> bool:
     if cached_batch is None:
         return False
-    if cached_batch.req_ids != req_ids:
+    if cached_batch.request_ids != request_ids:
         return False
     return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
 
@@ -581,7 +583,7 @@ class InputBatch:
     unchanged.
     """
 
-    req_ids: list[str]
+    request_ids: list[str]
     num_reqs: int
     num_reqs_after_padding: int
     idx_mapping: torch.Tensor
@@ -604,9 +606,9 @@ class InputBatch:
     negative_txt_seq_lens: list[int] | None = None
 
     def __post_init__(self) -> None:
-        if len(self.req_ids) != int(self.idx_mapping.numel()):
-            raise ValueError("`req_ids` and `idx_mapping` must have the same length.")
-        if self.num_reqs != len(self.req_ids):
+        if len(self.request_ids) != int(self.idx_mapping.numel()):
+            raise ValueError("`request_ids` and `idx_mapping` must have the same length.")
+        if self.num_reqs != len(self.request_ids):
             raise ValueError("`num_reqs` must match the number of request ids.")
         if self.num_reqs_after_padding < self.num_reqs:
             raise ValueError("`num_reqs_after_padding` must be >= `num_reqs`.")
@@ -653,11 +655,11 @@ class InputBatch:
         selected_states: Sequence[DiffusionRequestState],
         idx_mapping: torch.Tensor,
         idx_mapping_np: np.ndarray,
-        req_ids: list[str],
+        request_ids: list[str],
     ) -> InputBatch:
-        self.req_ids = req_ids
-        self.num_reqs = len(req_ids)
-        self.num_reqs_after_padding = len(req_ids)
+        self.request_ids = request_ids
+        self.num_reqs = len(request_ids)
+        self.num_reqs_after_padding = len(request_ids)
         self.idx_mapping = idx_mapping
         self.idx_mapping_np = idx_mapping_np
         self.latents = _prepare_latents(selected_states, out=self.latents)
@@ -675,9 +677,9 @@ class InputBatch:
     ) -> InputBatch:
         """Build a temporary step-local batch view from request states."""
         selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
-        req_ids = _prepare_req_ids(selected_states)
+        request_ids = _prepare_request_ids(selected_states)
 
-        if _same_composition(cached_batch, req_ids, idx_mapping_np):
+        if _same_composition(cached_batch, request_ids, idx_mapping_np):
             assert cached_batch is not None
             cached_batch._repack_dynamic_fields(selected_states)
             return cached_batch
@@ -687,14 +689,14 @@ class InputBatch:
                 selected_states,
                 idx_mapping,
                 idx_mapping_np,
-                req_ids,
+                request_ids,
             )
 
         prompt_embeds, prompt_embeds_mask = _prepare_prompt_embeds(selected_states)
         negative_prompt_embeds, negative_prompt_embeds_mask = _prepare_negative_prompt_embeds(selected_states)
         do_true_cfg, true_cfg_scale, cfg_normalize = _prepare_cfg_scalars(selected_states)
         return cls(
-            req_ids=req_ids,
+            request_ids=request_ids,
             num_reqs=len(selected_states),
             num_reqs_after_padding=len(selected_states),
             idx_mapping=idx_mapping,

vllm_omni/diffusion/worker/utils.py

@@ -38,7 +38,7 @@ class DiffusionRequestState:
     """
 
     # ── Identity / request-level inputs ──
-    req_id: str
+    request_id: str
     sampling: OmniDiffusionSamplingParams
     prompts: list[OmniPromptType] | None = None
 
@@ -112,7 +112,7 @@ class DiffusionRequestState:
 
 class BaseRunnerOutput(ABC):
     @abstractmethod
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
         pass
 
 
@@ -125,13 +125,13 @@ class RunnerOutput(BaseRunnerOutput):
     _request_state_cache.
     """
 
-    req_id: str
+    request_id: str
     step_index: int | None = None
     finished: bool = False
     result: DiffusionOutput | None = None
 
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
-        return self if self.req_id == sched_req_id else None
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
+        return self if self.request_id == request_id else None
 
 
 @dataclass
@@ -140,18 +140,18 @@ class BatchRunnerOutput(BaseRunnerOutput):
     _id_to_idx: dict[str, int] = field(init=False, repr=False)
 
     def __post_init__(self) -> None:
-        self._id_to_idx = {out.req_id: i for i, out in enumerate(self.runner_outputs)}
+        self._id_to_idx = {out.request_id: i for i, out in enumerate(self.runner_outputs)}
 
-    def __getitem__(self, req_id: str) -> RunnerOutput | None:
-        """access single RunnerOutput by req_id"""
-        idx = self._id_to_idx.get(req_id)
+    def __getitem__(self, request_id: str) -> RunnerOutput | None:
+        """access single RunnerOutput by request_id"""
+        idx = self._id_to_idx.get(request_id)
         return self.runner_outputs[idx] if idx is not None else None
 
-    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
-        return self[sched_req_id]
+    def get_request_output(self, request_id: str) -> RunnerOutput | None:
+        return self[request_id]
 
     @property
-    def req_ids(self) -> list[str]:
+    def request_ids(self) -> list[str]:
         return list(self._id_to_idx.keys())
 
     def __len__(self) -> int:

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -12,6 +12,8 @@ from typing import Any
 import torch
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
 from .factory import OmniConnectorFactory
 from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec
 from .utils.env import expand_env_int
@@ -1014,7 +1016,7 @@ class OmniKVTransferManager:
             return None, 0
 
         # Skip during warmup dummy run — no sender is available.
-        if request_id == "dummy_req_id":
+        if OmniDiffusionRequest.is_dummy_run_request_id(request_id):
             logger.info("Skip receiving KV cache for dummy warmup request")
             return None, 0
 
@@ -1171,12 +1173,7 @@ class OmniKVTransferManager:
     @staticmethod
     def _resolve_request_id(req: Any) -> str | None:
         """Resolve the logical request ID used for KV transfer lookups."""
-        request_id = getattr(req, "request_id", None)
-        if request_id:
-            return request_id
-        if hasattr(req, "request_ids") and req.request_ids:
-            return req.request_ids[0]
-        return None
+        return getattr(req, "request_id", None)
 
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -4,7 +4,6 @@ import base64
 import copy
 import hashlib
 import io
-import os
 from collections import OrderedDict
 from collections.abc import Callable, Iterable, Mapping
 from functools import lru_cache
@@ -16,14 +15,14 @@ import soundfile as sf
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import AutoTokenizer
+from transformers import AutoFeatureExtractor, AutoTokenizer
 from transformers.activations import ACT2FN
-from transformers.utils.hub import cached_file
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
 from vllm.model_executor.models.qwen3 import Qwen3Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
 from vllm.multimodal.audio import AudioResampler
@@ -36,7 +35,8 @@ from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
-from .qwen3_tts_tokenizer import Qwen3TTSTokenizer
+from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
+from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
 
 logger = init_logger(__name__)
 
@@ -459,7 +459,23 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         # Tokenizer for prompt building.
         self._tokenizer = None
-        self._speech_tokenizer: Qwen3TTSTokenizer | None = None
+
+        tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
+            self.model_path,
+            subfolder="speech_tokenizer",
+        )
+        self.encoder = Qwen3TTSTokenizerV2Encoder._from_config(
+            tokenizer_config.encoder_config,
+        )
+        self.encoder.eval()
+        self.encoder.to(dtype=torch.bfloat16)
+        self._encoder_valid_num_quantizers = int(tokenizer_config.encoder_valid_num_quantizers)
+        self._encoder_downsample_rate = int(tokenizer_config.encode_downsample_rate)
+
+        self._encoder_feature_extractor = AutoFeatureExtractor.from_pretrained(
+            self.model_path,
+            subfolder="speech_tokenizer",
+        )
 
         self._speaker_cache = get_speaker_cache()
         self._ref_audio_artifact_cache_max_entries = 256
@@ -1393,53 +1409,35 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         cache[key] = cached
         return cached
 
-    def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
-        if self._speech_tokenizer is not None:
-            return self._speech_tokenizer
-        speech_tokenizer_path = cached_file(self.model_path, "speech_tokenizer/config.json")
-        if speech_tokenizer_path is None:
-            raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found")
-        # Ensure the HF feature extractor config is present. Transformers'
-        # AutoFeatureExtractor does not proactively fetch this file.
-        preprocessor_config_path = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
-        if preprocessor_config_path is None:
-            raise ValueError(f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found")
-        speech_tokenizer_dir = os.path.dirname(speech_tokenizer_path)
-        tok = Qwen3TTSTokenizer.from_pretrained(
-            speech_tokenizer_dir,
-            torch_dtype=torch.bfloat16,
-        )
-        # Only move encoder to GPU; the decoder is unused by Talker (which
-        # only calls tok.encode()) and would otherwise waste bf16 VRAM.
-        # NOTE: after this point the tokenizer instance is encode-only;
-        # calling tok.decode() will fail because tok.model.decoder is None.
-        dev = next(self.parameters()).device
-        if dev.type != "cpu":
-            try:
-                del tok.model.decoder
-                tok.model.decoder = None
-                tok.model.encoder.to(dev)
-                tok.device = dev
-            except Exception as e:
-                raise RuntimeError(f"Failed to move speech tokenizer encoder to {dev}: {e}") from e
-        else:
-            tok.device = dev
-        self._speech_tokenizer = tok
-        return tok
+    def _encode_ref_audio_batch(self, wavs: list[np.ndarray], sr: int, *, device: torch.device) -> list[torch.Tensor]:
+        fe = self._encoder_feature_extractor
+        target_sr = int(fe.sampling_rate)
+        if int(sr) != target_sr:
+            resampler = AudioResampler(target_sr=target_sr)
+            wavs = [resampler.resample(w.astype(np.float32), orig_sr=int(sr)) for w in wavs]
 
-    def _encode_ref_audio_to_code(self, wav: np.ndarray, sr: int) -> torch.Tensor:
-        tok = self._ensure_speech_tokenizer_loaded()
-        enc = tok.encode(wav, sr=int(sr), return_dict=True)
-        ref_code = getattr(enc, "audio_codes", None)
-        if isinstance(ref_code, list):
-            ref_code = ref_code[0] if ref_code else None
-        if isinstance(ref_code, torch.Tensor):
-            # 12Hz: likely [T, Q] or [B, T, Q]
-            if ref_code.ndim == 3:
-                ref_code = ref_code[0]
-            out = ref_code.to(device=next(self.parameters()).device, dtype=torch.long)
-            return out
-        raise ValueError("SpeechTokenizer.encode did not return audio_codes tensor")
+        inputs = fe(
+            raw_audio=wavs,
+            sampling_rate=target_sr,
+            return_tensors="pt",
+        )
+        inputs = inputs.to(device).to(torch.bfloat16)
+
+        input_values = inputs["input_values"].squeeze(1)
+        padding_mask = inputs["padding_mask"].squeeze(1)
+
+        with torch.inference_mode():
+            encoded = self.encoder.encode(
+                input_values=input_values.unsqueeze(1),
+                return_dict=True,
+            )
+
+        audio_codes = encoded.audio_codes[:, : self._encoder_valid_num_quantizers]
+        downsample = self._encoder_downsample_rate
+        return [
+            code[..., : -(-mask.sum() // downsample)].transpose(0, 1).to(device=device, dtype=torch.long)
+            for code, mask in zip(audio_codes, padding_mask)
+        ]
 
     @staticmethod
     def _make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str:
@@ -1525,18 +1523,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             return None
         return ids.to(device=device, dtype=torch.long).contiguous()
 
-    @staticmethod
-    def _split_ref_code_batch(ref_codes: object, expected: int) -> list[object] | None:
-        if isinstance(ref_codes, torch.Tensor):
-            if ref_codes.ndim == 3 and int(ref_codes.shape[0]) == expected:
-                return [ref_codes[i] for i in range(expected)]
-            if expected == 1:
-                return [ref_codes]
-            return None
-        if isinstance(ref_codes, (list, tuple)) and len(ref_codes) == expected:
-            return list(ref_codes)
-        return None
-
     @staticmethod
     def _voice_clone_prompt_dict(raw: object) -> dict[str, object] | None:
         raw = Qwen3TTSTalkerForConditionalGeneration._first_value(raw)
@@ -1663,17 +1649,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         if not groups:
             return
 
-        tok = self._ensure_speech_tokenizer_loaded()
         for sr, items in groups.items():
             wavs = [wav for _, wav, _ in items]
             try:
-                enc = tok.encode(wavs, sr=int(sr), return_dict=True)
-                ref_codes = self._split_ref_code_batch(getattr(enc, "audio_codes", None), len(items))
+                ref_codes = self._encode_ref_audio_batch(wavs, int(sr), device=device)
             except Exception as exc:
                 logger.debug("Qwen3-TTS batched ref_code encode failed; falling back to serial path: %s", exc)
                 continue
-            if ref_codes is None:
-                continue
             for (info_dict, wav, item_sr), ref_code in zip(items, ref_codes, strict=True):
                 ref_code_t = self._coerce_ref_code_tensor(ref_code, device=device)
                 if ref_code_t is None:
@@ -1982,7 +1964,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     ref_code_len = int(ref_code_t.shape[0])
                 else:
                     wav_np, sr = _get_ref_audio()
-                    ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
+                    ref_code_t = self._encode_ref_audio_batch([wav_np], int(sr), device=input_ids.device)[0]
                     ref_code_len = int(ref_code_t.shape[0])
                     self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
             if isinstance(ref_code_t, torch.Tensor):
@@ -2240,6 +2222,43 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             # Some checkpoints do not include speaker_encoder weights; keep the
             # eagerly initialized module and satisfy the strict loader check.
             loaded |= {name for name, _ in self.named_parameters() if name.startswith("speaker_encoder.")}
+        # Load speech tokenizer encoder weights from speech_tokenizer/
+        # subfolder.  Skip decoder weights — the Talker only uses the
+        # encoder for ref_audio encoding.
+        model_loader = DefaultModelLoader(self.vllm_config.load_config)
+        source = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            subfolder="speech_tokenizer",
+        )
+        subfolder_weights = model_loader._get_weights_iterator(source)
+        enc_loaded = AutoWeightsLoader(
+            self,
+            skip_prefixes=["decoder."],
+        ).load_weights(subfolder_weights)
+        loaded |= enc_loaded
+
+        # AutoWeightsLoader only loads parameters; the encoder's VQ
+        # codebook state (embed, embed_sum, cluster_usage, initialized)
+        # are registered as buffers. Load them from the checkpoint
+        # directly so the quantizer produces correct codes.
+        encoder_buffers = dict(self.encoder.named_buffers())
+        source2 = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=self.vllm_config.model_config.revision,
+            subfolder="speech_tokenizer",
+        )
+        for name, tensor in model_loader._get_weights_iterator(source2):
+            if not name.startswith("encoder."):
+                continue
+            buf_name = name[len("encoder.") :]
+            if buf_name in encoder_buffers:
+                encoder_buffers[buf_name].copy_(tensor)
+                loaded.add(name)
+
+        device = self.vllm_config.device_config.device
+        self.encoder.to(device=device, dtype=torch.bfloat16)
+
         logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded))
         self._build_stacked_codec_embed()
         return loaded

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -1245,9 +1245,7 @@ class OmniConnectorModelRunnerMixin:
         if self._kv_transfer_manager is None:
             return False
 
-        request_id = getattr(req, "request_id", None) or (
-            req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None
-        )
+        request_id = getattr(req, "request_id", None)
         if not request_id:
             logger.warning("Request has no ID, cannot receive KV cache")
             return False