diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 45df0afa4..8a43539d4 100644 --- a/.buildkite/test-nightly.yml +++ b/.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 diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index 4390793e9..009df77ff 100644 --- a/docs/design/module/dit_module.md +++ b/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) diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index b4c928936..81f2bac24 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/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") diff --git a/docs/mkdocs/hooks/search_index.py b/docs/mkdocs/hooks/search_index.py new file mode 100644 index 000000000..1b9e268d9 --- /dev/null +++ b/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 diff --git a/docs/mkdocs/javascript/mathjax.js b/docs/mkdocs/javascript/mathjax.js index eb89ace06..2238e32e4 100644 --- a/docs/mkdocs/javascript/mathjax.js +++ b/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(); + }); +}); diff --git a/docs/mkdocs/javascript/memory.js b/docs/mkdocs/javascript/memory.js new file mode 100644 index 000000000..1cc22a8ff --- /dev/null +++ b/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", () => {}); diff --git a/docs/mkdocs/javascript/mermaid.js b/docs/mkdocs/javascript/mermaid.js index 63676de9d..f8fae2d47 100644 --- a/docs/mkdocs/javascript/mermaid.js +++ b/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"), + }); + }); }); diff --git a/docs/user_guide/quantization/autoround.md b/docs/user_guide/quantization/autoround.md index 88fed3b62..3afc8b062 100644 --- a/docs/user_guide/quantization/autoround.md +++ b/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 diff --git a/examples/offline_inference/internvla_a1/end2end.py b/examples/offline_inference/internvla_a1/end2end.py index 6bc69f4e0..72cdbff0c 100755 --- a/examples/offline_inference/internvla_a1/end2end.py +++ b/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, diff --git a/mkdocs.yml b/mkdocs.yml index d19483002..74f2145fe 100644 --- a/mkdocs.yml +++ b/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 diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json index 82c2ca27d..1903bc670 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json +++ b/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", diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json index 8ae370df6..21695b752 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json +++ b/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", diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index 92dd73358..b64d6f738 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/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) diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py index f7ada6ebf..1365a3477 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py +++ b/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", ) diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py index 7aac36970..e5a1b6d83 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py +++ b/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( diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_pe.py b/tests/diffusion/models/ernie_image/test_ernie_image_pe.py index d1dc90715..a632ab3be 100644 --- a/tests/diffusion/models/ernie_image/test_ernie_image_pe.py +++ b/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 diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py new file mode 100644 index 000000000..8ab046f0c --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 5b254d50b..3930677d3 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/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) diff --git a/tests/diffusion/models/ovis_image/test_ovis_image.py b/tests/diffusion/models/ovis_image/test_ovis_image.py index 6e7ed2f58..8e357e6c6 100644 --- a/tests/diffusion/models/ovis_image/test_ovis_image.py +++ b/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, diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index dca9a001a..2d44f6a68 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a7b945eea..a8985f372 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index b1b63a1a2..852f6dc2e 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/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() diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index 56ec95d0c..9087a6e2f 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index ca28b2629..85d2e45c9 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/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), diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 2722de7e3..029fea120 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/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) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index aa2a74ac1..99d7a9c9c 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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: diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 36ceb6cfc..0419bb813 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/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"] diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index cea186019..2cc19f1b3 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/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 diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index 540444d42..338758cc0 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/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: diff --git a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py new file mode 100644 index 000000000..2ea4dc641 --- /dev/null +++ b/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() diff --git a/tests/e2e/online_serving/test_qwen_image_edit.py b/tests/e2e/online_serving/test_qwen_image_edit.py index bdf8c9122..738a2e50c 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit.py +++ b/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): diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index 632f75231..1457ad636 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/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']}" diff --git a/tests/model_executor/models/glm_image/test_glm_image_init.py b/tests/model_executor/models/glm_image/test_glm_image_init.py new file mode 100644 index 000000000..eb42e9649 --- /dev/null +++ b/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__") diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index d4b8e4a32..cd590161f 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/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 diff --git a/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py b/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py index 14ba86cc4..aa597f12a 100644 --- a/vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py +++ b/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}", ], } diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index f54cca6ef..c39ab792e 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index 37ea454de..f808b301b 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/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), diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8fd85fd9e..8cb0d5296 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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: diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 6727122d4..3aa409e14 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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, diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index b6f881d1c..d8eb2ac53 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/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, ) diff --git a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py index 7b5d80571..cd60cc93b 100644 --- a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py +++ b/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: diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index ddb32aa20..56a7066e3 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/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( diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 97cba18c2..102456d9f 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/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): diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 4d4328d25..79198c3ec 100644 --- a/vllm_omni/diffusion/request.py +++ b/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 diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 0bd16597f..812873a1c 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/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)) diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 1bfb69458..b50e3555a 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/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 diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index 5029e3593..ba9cda616 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/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) diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py index 58449281b..0bc2d4a3c 100644 --- a/vllm_omni/diffusion/sched/step_scheduler.py +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 19e890abb..07489d4c9 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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, ) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index ab6dbb79a..2720de038 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 782e79855..8f0577b49 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index 48c5bb57b..82033d9dc 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/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, diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 1e98a7784..a0bfcea9b 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/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: diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 372c2cb3d..3b234ad1c 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/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: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index ec00813b9..dbede17a3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/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 diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 64b7e60e2..6d991d0f2 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/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