anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
WeiQing Chen authored and GitHub committed 2026-05-28 19:18:52 +08:00
2 parents 5ab2932ce6 74258b18b7
commit 4bbe81308b
56 files changed
+1821 -736
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/diffusion_backend/test_diffusers_backend.py
+1
View File
@@ -60,6 +60,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
output_type="pil",
generator_device="cpu",
),
"request_id": "diffusers-backend",
}
defaults.update(overrides)
return OmniDiffusionRequest(**defaults)
tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+1
View File
@@ -48,6 +48,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=prompts or [{"prompt": "a cat dancing"}],
sampling_params=sp,
request_id="dmd2-sanitize",
)
tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+1 -1
View File
@@ -48,7 +48,7 @@ def _make_pipeline(cls):
def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(**sp_kwargs)
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
@pytest.fixture(
tests/diffusion/models/ernie_image/test_ernie_image_pe.py
+1 -1
View File
@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
def test_should_apply_pe_disables_dummy_warmup_request():
req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
assert ErnieImagePipeline._should_apply_pe(req) is False
tests/diffusion/models/glm_image/test_glm_image_quantization.py
+741
View File
@@ -0,0 +1,741 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
These tests verify that the GLM-Image DiT transformer correctly accepts and uses
quantization configs for W4A16/AutoRound quantization support.
"""
import os
import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
ColumnParallelGELU,
ColumnParallelSiLU,
GlmImageAdaLayerNormContinuous,
GlmImageAdaLayerNormZero,
GlmImageAttention,
GlmImageFeedForward,
GlmImageImageProjector,
GlmImagePrepare,
GlmImageRotaryPosEmbed,
GlmImageTransformer2DModel,
GlmImageTransformerBlock,
_positive_divisors,
validate_glm_image_tp_constraints,
)
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture(scope="function", autouse=True)
def _init_distributed():
"""Initialize the minimal distributed environment required by
ReplicatedLinear (tensor-parallel group must exist)."""
from vllm.distributed.parallel_state import (
cleanup_dist_env_and_memory,
init_distributed_environment,
initialize_model_parallel,
)
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29502")
init_distributed_environment(
world_size=1,
rank=0,
local_rank=0,
distributed_init_method="env://",
)
initialize_model_parallel()
yield
cleanup_dist_env_and_memory()
@pytest.fixture(scope="function", autouse=True)
def _force_default_gemm(monkeypatch):
"""Force CPU-compatible GEMM dispatch for tests using CPU tensors.
vLLM's dispatch_unquantized_gemm() selects the backend by platform
(e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
CPU test tensors crash with NotImplementedError on ROCm. Monkeypatch
the dispatcher to always return the default (torch.nn.functional.linear)
implementation which works on any device."""
from vllm.model_executor.layers.utils import default_unquantized_gemm
monkeypatch.setattr(
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
lambda: default_unquantized_gemm,
)
@pytest.fixture(scope="function", autouse=True)
def setup_mocks(mocker: MockerFixture):
"""Set up common mocks for all tests."""
mocker.patch(
"vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
return_value=1,
)
mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
mock_tp_group = mocker.MagicMock()
mock_tp_group.world_size = 1
mock_get_tp_group.return_value = mock_tp_group
yield
class TestPositiveDivisors:
"""Test _positive_divisors helper function."""
def test_divisors_of_1(self):
assert _positive_divisors(1) == {1}
def test_divisors_of_12(self):
assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
def test_divisors_of_prime(self):
assert _positive_divisors(7) == {1, 7}
def test_divisors_of_0_returns_empty(self):
assert _positive_divisors(0) == set()
def test_divisors_of_negative_returns_empty(self):
assert _positive_divisors(-5) == set()
class TestValidateGlmImageTpConstraints:
"""Test TP constraint validation for GLM-Image."""
def test_valid_tp_size_1(self):
"""TP=1 should always be valid."""
result = validate_glm_image_tp_constraints(
dim=2560,
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=1,
)
assert 1 in result
def test_valid_tp_size_2_for_divisible_dim(self):
"""TP=2 is valid when all dims are divisible by 2."""
result = validate_glm_image_tp_constraints(
dim=2560,
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=2,
)
assert 1 in result
assert 2 in result
def test_valid_tp_size_4_for_divisible_dim(self):
"""TP=4 is valid when all dims are divisible by 4."""
result = validate_glm_image_tp_constraints(
dim=2560,
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=4,
)
assert 1 in result
assert 2 in result
assert 4 in result
def test_invalid_tp_size_3_for_divisible_dim(self):
"""TP=3 is invalid when dim is not divisible by 3."""
with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
validate_glm_image_tp_constraints(
dim=2560, # 2560 % 3 != 0
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=3,
)
def test_invalid_tp_size_zero(self):
"""TP=0 should raise error."""
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
validate_glm_image_tp_constraints(
dim=2560,
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=0,
)
def test_invalid_tp_size_negative(self):
"""Negative TP size should raise error."""
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
validate_glm_image_tp_constraints(
dim=2560,
num_heads=64,
ffn_hidden_dim=10240,
tensor_parallel_size=-1,
)
class TestGlmImageAdaLayerNormZeroQuantization:
"""Test GlmImageAdaLayerNormZero with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify the class accepts quant_config parameter."""
mock_quant_config = mocker.MagicMock()
layer = GlmImageAdaLayerNormZero(
embedding_dim=512,
dim=2560,
quant_config=mock_quant_config,
prefix="test.norm1",
)
assert layer.linear.quant_config is mock_quant_config
def test_accepts_none_quant_config(self):
"""Verify quant_config=None is accepted."""
layer = GlmImageAdaLayerNormZero(
embedding_dim=512,
dim=2560,
quant_config=None,
)
assert layer.linear.quant_config is None
def test_forward_handles_tuple_return_from_linear(self):
"""Verify forward handles tuple returns from ReplicatedLinear."""
layer = GlmImageAdaLayerNormZero(
embedding_dim=512,
dim=2560,
quant_config=None,
)
batch_size = 2
seq_len = 10
hidden_states = torch.randn(batch_size, seq_len, 2560)
encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
temb = torch.randn(batch_size, 512)
# This should work regardless of whether linear returns tuple or tensor
result = layer(hidden_states, encoder_hidden_states, temb)
assert len(result) == 10 # Should return 10 chunks
class TestGlmImageAdaLayerNormContinuousQuantization:
"""Test GlmImageAdaLayerNormContinuous with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify the class accepts quant_config parameter."""
mock_quant_config = mocker.MagicMock()
layer = GlmImageAdaLayerNormContinuous(
embedding_dim=2560,
conditioning_embedding_dim=512,
quant_config=mock_quant_config,
prefix="test.norm_out",
)
assert layer.linear.quant_config is mock_quant_config
def test_accepts_none_quant_config(self):
"""Verify quant_config=None is accepted."""
layer = GlmImageAdaLayerNormContinuous(
embedding_dim=2560,
conditioning_embedding_dim=512,
quant_config=None,
)
assert layer.linear.quant_config is None
def test_forward_handles_tuple_return_from_linear(self):
"""Verify forward handles tuple returns from ReplicatedLinear."""
layer = GlmImageAdaLayerNormContinuous(
embedding_dim=2560,
conditioning_embedding_dim=512,
quant_config=None,
)
batch_size = 2
seq_len = 10
x = torch.randn(batch_size, seq_len, 2560)
conditioning_embedding = torch.randn(batch_size, 512)
result = layer(x, conditioning_embedding)
assert result.shape == (batch_size, seq_len, 2560)
class TestGlmImageAttentionQuantization:
"""Test GlmImageAttention with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify GlmImageAttention accepts quant_config parameter."""
mock_quant_config = mocker.MagicMock()
attn = GlmImageAttention(
dim=2560,
num_heads=64,
head_dim=40,
quant_config=mock_quant_config,
prefix="test.attn1",
)
assert attn.to_qkv.quant_config is mock_quant_config
def test_accepts_none_quant_config(self):
"""Verify quant_config=None is accepted."""
attn = GlmImageAttention(
dim=2560,
num_heads=64,
head_dim=40,
quant_config=None,
)
assert attn.to_qkv.quant_config is None
class TestColumnParallelModulesQuantization:
"""Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
"""Verify ColumnParallelGELU accepts quant_config."""
mock_quant_config = mocker.MagicMock()
layer = ColumnParallelGELU(
dim_in=2560,
dim_out=10240,
quant_config=mock_quant_config,
prefix="test.gelu",
)
assert layer.proj.quant_config is mock_quant_config
def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
"""Verify ColumnParallelSiLU accepts quant_config."""
mock_quant_config = mocker.MagicMock()
layer = ColumnParallelSiLU(
dim_in=2560,
dim_out=10240,
quant_config=mock_quant_config,
prefix="test.silu",
)
assert layer.proj.quant_config is mock_quant_config
class TestGlmImageFeedForwardQuantization:
"""Test GlmImageFeedForward with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify GlmImageFeedForward accepts quant_config parameter."""
mock_quant_config = mocker.MagicMock()
ff = GlmImageFeedForward(
dim=2560,
dim_out=2560,
inner_dim=10240,
activation_fn="gelu-approximate",
quant_config=mock_quant_config,
prefix="test.ff",
)
# Check that the first layer (ColumnParallelGELU) has quant_config
gelu_layer = ff.net[0]
assert gelu_layer.proj.quant_config is mock_quant_config
def test_accepts_none_quant_config(self):
"""Verify quant_config=None is accepted."""
ff = GlmImageFeedForward(
dim=2560,
dim_out=2560,
inner_dim=10240,
activation_fn="gelu-approximate",
quant_config=None,
)
gelu_layer = ff.net[0]
assert gelu_layer.proj.quant_config is None
def test_linear_silu_activation(self):
"""Test linear-silu activation function initialization works."""
# This test verifies the module can be instantiated with linear-silu activation.
# Full forward testing requires proper TP group setup, so we just verify construction.
ff = GlmImageFeedForward(
dim=2560,
dim_out=2560,
inner_dim=10240,
activation_fn="linear-silu",
quant_config=None,
)
# Verify the FFN has the correct structure
assert len(ff.net) == 3 # ColumnParallelSiLU, Identity, RowParallelLinear
assert isinstance(ff.net[0], ColumnParallelSiLU)
class TestGlmImageTransformerBlockQuantization:
"""Test GlmImageTransformerBlock with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify GlmImageTransformerBlock accepts quant_config parameter."""
mock_quant_config = mocker.MagicMock()
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
block = GlmImageTransformerBlock(
dim=2560,
num_attention_heads=64,
attention_head_dim=40,
time_embed_dim=512,
parallel_config=parallel_config,
quant_config=mock_quant_config,
prefix="test.block",
)
# Check that inner modules have quant_config
assert block.norm1.linear.quant_config is mock_quant_config
assert block.attn1.to_qkv.quant_config is mock_quant_config
assert block.ff.net[0].proj.quant_config is mock_quant_config
def test_accepts_none_quant_config(self):
"""Verify quant_config=None is accepted."""
block = GlmImageTransformerBlock(
dim=2560,
num_attention_heads=64,
attention_head_dim=40,
time_embed_dim=512,
quant_config=None,
)
assert block.norm1.linear.quant_config is None
class TestGlmImagePrepareModule:
"""Test GlmImagePrepare module."""
def test_prepare_module_exists(self):
"""Verify GlmImagePrepare module exists and works."""
projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
prepare = GlmImagePrepare(
image_projector=projector,
rope=rope,
patch_size=2,
)
# Test forward pass
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
result = prepare(hidden_states)
assert len(result) == 5
hidden_out, rope_cos, rope_sin, height, width = result
assert hidden_out.shape[0] == 1 # batch size
assert hidden_out.shape[1] == 1024 # seq_len (32 * 32)
assert hidden_out.shape[2] == 2560 # hidden_dim
class TestGlmImagePipelineConfig:
"""Test GLM_IMAGE_PIPELINE configuration."""
def test_glm_image_pipeline_config_exists(self):
"""Verify GLM_IMAGE_PIPELINE config exists."""
assert GLM_IMAGE_PIPELINE is not None
def test_pipeline_has_correct_model_type(self):
"""Verify pipeline has correct model_type."""
assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
def test_pipeline_has_correct_model_arch(self):
"""Verify pipeline has correct model_arch."""
assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
def test_pipeline_has_two_stages(self):
"""Verify pipeline has two stages (AR + DiT)."""
assert len(GLM_IMAGE_PIPELINE.stages) == 2
def test_stage_0_is_ar_llm(self):
"""Verify stage 0 is the AR LLM stage."""
stage_0 = GLM_IMAGE_PIPELINE.stages[0]
assert stage_0.stage_id == 0
assert stage_0.model_stage == "ar"
assert stage_0.execution_type.value == "llm_ar"
assert stage_0.owns_tokenizer is True
def test_stage_1_is_diffusion(self):
"""Verify stage 1 is the DiT diffusion stage."""
stage_1 = GLM_IMAGE_PIPELINE.stages[1]
assert stage_1.stage_id == 1
assert stage_1.model_stage == "dit"
assert stage_1.execution_type.value == "diffusion"
assert stage_1.final_output is True
assert stage_1.final_output_type == "image"
assert stage_1.input_sources == (0,) # Takes input from stage 0
def test_pipeline_has_diffusers_class_name(self):
"""Verify pipeline has diffusers_class_name."""
assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
class TestGlmImageImageProjector:
"""Test GlmImageImageProjector module."""
def test_projector_output_shape(self):
"""Verify projector produces correct output shape."""
projector = GlmImageImageProjector(
in_channels=16,
hidden_size=2560,
patch_size=2,
)
# Input: [B, C, H, W] = [1, 16, 64, 64]
hidden_states = torch.randn(1, 16, 64, 64)
output = projector(hidden_states)
# After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
assert output.shape == (1, 1024, 2560)
def test_projector_with_different_patch_sizes(self):
"""Verify projector works with different patch sizes."""
for patch_size in [2, 4]:
projector = GlmImageImageProjector(
in_channels=16,
hidden_size=2560,
patch_size=patch_size,
)
h, w = 64, 64
hidden_states = torch.randn(1, 16, h, w)
output = projector(hidden_states)
expected_seq_len = (h // patch_size) * (w // patch_size)
assert output.shape[1] == expected_seq_len
class TestGlmImageRotaryPosEmbed:
"""Test GlmImageRotaryPosEmbed module."""
def test_rope_output_shape(self):
"""Verify RoPE produces correct output shape."""
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
cos, sin = rope(hidden_states)
# After patchify: height=32, width=32
# Output: [32*32, dim] = [1024, 40]
assert cos.shape[0] == 1024
assert cos.shape[1] == 40
assert sin.shape == cos.shape
def test_rope_value_range(self):
"""Verify RoPE produces values in valid range."""
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
hidden_states = torch.randn(2, 16, 32, 32)
cos, sin = rope(hidden_states)
# cos and sin should be in [-1, 1]
assert cos.min() >= -1.0 and cos.max() <= 1.0
assert sin.min() >= -1.0 and sin.max() <= 1.0
def test_rope_consistency(self):
"""Verify RoPE is consistent for same input."""
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
hidden_states = torch.randn(1, 16, 32, 32)
cos1, sin1 = rope(hidden_states)
cos2, sin2 = rope(hidden_states)
# Same input should produce same output
assert torch.allclose(cos1, cos2)
assert torch.allclose(sin1, sin2)
class TestGlmImageTransformer2DModelQuantization:
"""Test GlmImageTransformer2DModel with quantization config."""
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
"""Verify the model accepts quant_config parameter."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
mock_quant_config = mocker.MagicMock()
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
# Create a minimal mock od_config
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2 # Small number for testing
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=mock_quant_config,
)
# Check that quantization config was passed to transformer blocks
for block in model.transformer_blocks:
assert block.norm1.linear.quant_config is mock_quant_config
assert block.attn1.to_qkv.quant_config is mock_quant_config
assert block.ff.net[0].proj.quant_config is mock_quant_config
def test_accepts_none_quant_config(self, mocker: MockerFixture):
"""Verify quant_config=None is accepted."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=None,
)
# Check that quantization config is None
for block in model.transformer_blocks:
assert block.norm1.linear.quant_config is None
assert block.attn1.to_qkv.quant_config is None
def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
"""Verify norm_out (output layer) does NOT use quantization to preserve precision."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
mock_quant_config = mocker.MagicMock()
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=mock_quant_config,
)
# norm_out.linear should NOT have quant_config to preserve output precision
assert model.norm_out.linear.quant_config is None
def test_model_has_sp_plan(self, mocker: MockerFixture):
"""Verify model has _sp_plan defined for sequence parallelism."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=None,
)
# Verify _sp_plan exists
assert hasattr(model, "_sp_plan")
assert "prepare" in model._sp_plan
assert "proj_out" in model._sp_plan
def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
"""Verify model has _hsdp_shard_conditions for HSDP parallelism."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=None,
)
assert hasattr(model, "_hsdp_shard_conditions")
assert len(model._hsdp_shard_conditions) > 0
def test_model_creates_kv_cache(self, mocker: MockerFixture):
"""Verify model can create KV cache for image editing."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
parallel_config = DiffusionParallelConfig(
tensor_parallel_size=1,
sequence_parallel_size=1,
)
mock_tf_config = mocker.MagicMock()
mock_tf_config.patch_size = 2
mock_tf_config.in_channels = 16
mock_tf_config.out_channels = 16
mock_tf_config.num_attention_heads = 64
mock_tf_config.attention_head_dim = 40
mock_tf_config.time_embed_dim = 512
mock_tf_config.condition_dim = 256
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
mock_tf_config.text_embed_dim = 1024
mock_tf_config.num_layers = 2
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
mock_od_config.tf_model_config = mock_tf_config
mock_od_config.parallel_config = parallel_config
model = GlmImageTransformer2DModel(
od_config=mock_od_config,
quant_config=None,
)
kv_cache = model.create_kv_cache()
assert kv_cache is not None
assert len(kv_cache) == 2 # num_layers
tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
+1
View File
@@ -339,6 +339,7 @@ class TestCFGParallelForwardPath:
audio_latents=audio_latents,
output_type="latent",
),
request_id="ltx23-cfg-parallel-forward-test",
)
output = pipe.forward(req)
tests/diffusion/models/ovis_image/test_ovis_image.py
+3
View File
@@ -167,6 +167,7 @@ def test_basic_generation(ovis_pipeline):
# Setup request
req = OmniDiffusionRequest(
prompts=["A photo of a cat"],
request_id="ovis-basic",
sampling_params=OmniDiffusionSamplingParams(
height=256,
width=256,
@@ -206,6 +207,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
"negative_prompt": "bad quality",
}
],
request_id="ovis-guidance",
sampling_params=OmniDiffusionSamplingParams(
height=256,
width=256,
@@ -225,6 +227,7 @@ def test_resolution_check(ovis_pipeline):
# Pass odd resolution
req = OmniDiffusionRequest(
prompts=["test"],
request_id="ovis-resolution",
sampling_params=OmniDiffusionSamplingParams(
height=250, # Not divisible by 16 (8*2)
width=250,
tests/diffusion/test_diffusion_engine.py
+22 -14
View File
@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
num_waiting_reqs: int = 0
@property
def scheduled_req_ids(self):
def scheduled_request_ids(self):
ids = [req.request_id for req in self.scheduled_new_reqs]
if self.scheduled_cached_reqs:
ids.extend(self.scheduled_cached_reqs.sched_req_ids)
ids.extend(self.scheduled_cached_reqs.request_ids)
return ids
@property
def is_empty(self):
return len(self.scheduled_req_ids) == 0
return len(self.scheduled_request_ids) == 0
class MockScheduler:
@@ -71,31 +71,36 @@ async def test_async_add_req_and_wait_for_response():
engine.scheduler = MockScheduler()
engine._out_queue = {}
engine.abort_queue: queue.Queue[str] = queue.Queue()
engine._rpc_lock = threading.Lock()
engine._rpc_queue = queue.Queue()
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._init_lock = asyncio.Lock()
engine._closed = False
engine._loop_started = False
engine.main_loop = None
engine._finalize_finished_request = lambda rid, out, err: out.result
def mock_execute_batch(sched_output):
req_ids = sched_output.scheduled_req_ids
request_ids = sched_output.scheduled_request_ids
time.sleep(1)
class MockRunnerOutput:
def __init__(self, ids):
self.req_id = ids
self.request_id = ids
self.step_index = 0
self.finished = True
self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
def get_req_output(self, rid):
def get_request_output(self, rid):
return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
return MockRunnerOutput(req_ids)
return MockRunnerOutput(request_ids)
engine.execute_fn = mock_execute_batch
engine.stop_event = threading.Event()
engine.start_background_loop()
await engine._check_and_start_background_loop()
async def run_task(rid):
req = SimpleNamespace(request_id=rid)
@@ -105,10 +110,13 @@ async def test_async_add_req_and_wait_for_response():
task_ids = [f"req_{i}" for i in range(5)]
tasks = [run_task(rid) for rid in task_ids]
results = await asyncio.gather(*tasks)
engine.stop_event.set()
engine.worker_thread.join(timeout=5)
try:
results = await asyncio.gather(*tasks)
finally:
with engine._cv:
engine.stop_event.set()
engine._cv.notify_all()
engine.worker_thread.join(timeout=5)
assert len(results) == 5
for rid, res, elapsed in results:
assert rid in res.result_data
tests/diffusion/test_diffusion_engine_cleanup.py
+5 -10
View File
@@ -17,15 +17,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _make_request(
request_id: str,
*,
request_ids: list[str] | None = None,
) -> OmniDiffusionRequest:
def _make_request(request_id: str) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=[f"prompt_{request_id}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=request_ids or [request_id],
request_id=request_id,
)
@@ -79,15 +74,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
engine._finalize_finished_request.assert_not_called()
def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
def test_abort_request_id_aborts_scheduler_request() -> None:
engine = _make_engine()
request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
sched_req_id = engine.scheduler.add_request(request)
request = _make_request("batch-parent")
request_id = engine.scheduler.add_request(request)
engine.abort("batch-parent")
engine._process_aborts_queue()
state = engine.scheduler.get_request_state(sched_req_id)
state = engine.scheduler.get_request_state(request_id)
assert state is not None
assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
tests/diffusion/test_diffusion_engine_rpc_routing.py
+4 -3
View File
@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
if method in {"execute_model", "execute_stepwise", "generate"}:
# args[0] is the request-like object for execute_model.
req = args[0] if args else None
tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
return DiffusionOutput(error=f"result_for_{tag}")
tag = args[0] if args else method
return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
exec_all_ranks=True,
)
return RunnerOutput(
req_id=new_req.sched_req_id,
request_id=new_req.request_id,
step_index=None,
finished=True,
result=result,
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
def _make_request(tag: str):
return SimpleNamespace(
request_ids=[tag],
request_id=tag,
prompts=[f"prompt_{tag}"],
sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
)
@@ -148,6 +148,7 @@ def _make_engine_with_loop(
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._out_queue = {}
engine._closed = False
engine.abort_queue = queue.Queue()
engine._rpc_queue = queue.Queue()
tests/diffusion/test_diffusion_request.py
+45 -1
View File
@@ -5,7 +5,7 @@ import random
import pytest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -15,9 +15,53 @@ def _make_request() -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=[{"prompt": "a cup of coffee on a table"}],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="request-test",
)
def test_request_id_is_required():
with pytest.raises(TypeError, match="request_id"):
OmniDiffusionRequest(
prompts=[{"prompt": "a cup of coffee on a table"}],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
)
def test_request_ids_identity_list_is_removed():
req = _make_request()
assert req.request_id == "request-test"
assert not hasattr(req, "request_ids")
def test_request_id_must_be_non_empty():
with pytest.raises(ValueError, match="request_id must be a non-empty string"):
OmniDiffusionRequest(
prompts=[{"prompt": "a cup of coffee on a table"}],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="",
)
def test_dummy_run_request_is_identified_by_reserved_request_id():
req = OmniDiffusionRequest(
prompts=[{"prompt": "dummy run"}],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id=DUMMY_DIFFUSION_REQUEST_ID,
)
assert req.is_dummy_run()
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
def test_non_dummy_request_is_not_identified_as_dummy_run():
req = _make_request()
assert not req.is_dummy_run()
assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
def test_tp_seed_same_across_ranks_and_varies_across_requests():
random.seed(0)
n_requests = 5
tests/diffusion/test_diffusion_scheduler.py
+37 -74
View File
@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=[f"prompt_{req_id}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=[req_id],
request_id=req_id,
)
def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
return RunnerOutput(
req_id=req_id,
request_id=req_id,
step_index=None,
finished=finished,
result=DiffusionOutput(output=None, error=error),
@@ -52,7 +52,7 @@ def _make_step_output(
error: str | None = None,
):
return RunnerOutput(
req_id=req_id,
request_id=req_id,
step_index=step_index,
finished=finished,
result=DiffusionOutput(output=None, error=error) if error is not None else None,
@@ -73,16 +73,16 @@ def _make_step_request(
num_inference_steps=num_inference_steps,
step_index=step_index,
),
request_ids=[req_id],
request_id=req_id,
)
def _new_ids(sched_output) -> list[str]:
return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
return [req.request_id for req in sched_output.scheduled_new_reqs]
def _cached_ids(sched_output) -> list[str]:
return list(sched_output.scheduled_cached_reqs.sched_req_ids)
return list(sched_output.scheduled_cached_reqs.request_ids)
class _StubScheduler(SchedulerInterface):
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
self._request = request
self._output = output
self.initialized_with = None
self._sched_req_id = request.request_ids[0]
self._request_id = request.request_id
self._state = None
self._scheduled = False
self.max_num_running_reqs = 1
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
def add_request(self, request: OmniDiffusionRequest) -> str:
assert request is self._request
self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
return self._sched_req_id
self._state = SimpleNamespace(request_id=self._request_id, req=request)
return self._request_id
def schedule(self):
if self._scheduled or self._state is None:
return SimpleNamespace(
scheduled_new_reqs=[],
scheduled_cached_reqs=CachedRequestData.make_empty(),
scheduled_req_ids=[],
scheduled_request_ids=[],
is_empty=True,
)
self._scheduled = True
return SimpleNamespace(
scheduled_new_reqs=[NewRequestData.from_state(self._state)],
scheduled_cached_reqs=CachedRequestData.make_empty(),
scheduled_req_ids=[self._state.sched_req_id],
scheduled_request_ids=[self._state.request_id],
is_empty=False,
)
@@ -123,30 +123,25 @@ class _StubScheduler(SchedulerInterface):
del sched_output
assert output is self._output
self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
return {self._sched_req_id}
return {self._request_id}
def has_requests(self) -> bool:
return not self._scheduled
def get_request_state(self, sched_req_id: str):
del sched_req_id
def get_request_state(self, request_id: str):
del request_id
return self._state
def get_sched_req_id(self, request_id: str) -> str | None:
if request_id in self._request.request_ids:
return self._sched_req_id
return None
def pop_request_state(self, sched_req_id: str):
del sched_req_id
def pop_request_state(self, request_id: str):
del request_id
return self._state
def preempt_request(self, sched_req_id: str) -> bool:
del sched_req_id
def preempt_request(self, request_id: str) -> bool:
del request_id
return False
def finish_requests(self, sched_req_ids, status) -> None:
del sched_req_ids, status
def finish_requests(self, request_ids, status) -> None:
del request_ids, status
return None
def close(self) -> None:
@@ -171,7 +166,7 @@ class TestGetSamplingParamsKey:
return OmniDiffusionRequest(
prompts=["prompt"],
sampling_params=sp,
request_ids=[f"req-{lora_int_id}-{lora_scale}"],
request_id=f"req-{lora_int_id}-{lora_scale}",
)
def test_distinguishes_lora_id(self) -> None:
@@ -291,7 +286,7 @@ class TestRequestScheduler:
OmniDiffusionRequest(
prompts=["prompt_b"],
sampling_params=OmniDiffusionSamplingParams(width=768),
request_ids=["b"],
request_id="b",
)
)
scheduler.add_request(_make_request("c"))
@@ -333,7 +328,7 @@ class TestRequestScheduler:
assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
assert self.scheduler.has_requests() is False
assert self.scheduler.schedule().scheduled_req_ids == []
assert self.scheduler.schedule().scheduled_request_ids == []
def test_has_requests_state_transition(self) -> None:
assert self.scheduler.has_requests() is False
@@ -348,60 +343,28 @@ class TestRequestScheduler:
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
assert self.scheduler.has_requests() is False
def test_request_id_mapping_lifecycle(self) -> None:
def test_request_id_is_scheduler_key(self) -> None:
request = OmniDiffusionRequest(
prompts=["prompt_map_a", "prompt_map_b"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=["map-a", "map-b"],
request_id="map-parent",
)
sched_req_id = self.scheduler.add_request(request)
request_id = self.scheduler.add_request(request)
assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
assert request_id == "map-parent"
state = self.scheduler.get_request_state("map-parent")
assert state.request_id == "map-parent"
self.scheduler.pop_request_state(sched_req_id)
self.scheduler.pop_request_state("map-parent")
assert self.scheduler.get_sched_req_id("map-a") is None
assert self.scheduler.get_sched_req_id("map-b") is None
assert self.scheduler.get_sched_req_id("map-parent") is None
assert self.scheduler.get_request_state("map-parent") is None
def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
self.scheduler.add_request(
OmniDiffusionRequest(
prompts=["prompt_existing"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=["existing-child"],
request_id="duplicate-parent",
)
)
def test_duplicate_request_id_is_rejected(self) -> None:
self.scheduler.add_request(_make_request("dup"))
colliding_request = OmniDiffusionRequest(
prompts=["prompt_unique"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=["unique-child"],
request_id="duplicate-parent",
)
with pytest.raises(ValueError, match="duplicate-parent"):
self.scheduler.add_request(colliding_request)
assert self.scheduler.get_request_state("unique-child") is None
assert self.scheduler.get_sched_req_id("unique-child") is None
valid_request = OmniDiffusionRequest(
prompts=["prompt_unique"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_ids=["unique-child"],
request_id="unique-parent",
)
sched_req_id = self.scheduler.add_request(valid_request)
assert sched_req_id == "unique-child"
assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
with pytest.raises(ValueError, match="request_id 'dup' is already active"):
self.scheduler.add_request(_make_request("dup"))
class TestDiffusionEngine:
@@ -662,7 +625,7 @@ class TestStepScheduler:
finished = self.scheduler.update_from_output(
sched_output,
RunnerOutput(
req_id=req_id,
request_id=req_id,
step_index=None,
finished=True,
result=None,
@@ -934,7 +897,7 @@ class TestStepScheduler:
for _ in range(expected_steps - 1):
sched_output = self.scheduler.schedule()
assert sched_output.scheduled_req_ids == [req_id]
assert sched_output.scheduled_request_ids == [req_id]
next_step = request.sampling_params.step_index + 1
assert (
self.scheduler.update_from_output(
@@ -945,7 +908,7 @@ class TestStepScheduler:
)
final_output = self.scheduler.schedule()
assert final_output.scheduled_req_ids == [req_id]
assert final_output.scheduled_request_ids == [req_id]
assert self.scheduler.update_from_output(
final_output,
_make_step_output(req_id, step_index=expected_steps, finished=True),
tests/diffusion/test_diffusion_step_pipeline.py
+35 -35
View File
@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
def _make_step_request(num_inference_steps: int = 2):
return SimpleNamespace(
prompts=["a prompt"],
request_ids=["req-1"],
request_id="req-1",
sampling_params=SimpleNamespace(
generator=None,
seed=None,
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
return OmniDiffusionRequest(
prompts=[f"prompt-{req_id}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
request_ids=[req_id],
request_id=req_id,
)
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
return runner
def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
return DiffusionSchedulerOutput(
step_id=step_id,
scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
scheduled_cached_reqs=CachedRequestData.make_empty(),
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
num_running_reqs=1,
@@ -287,7 +287,7 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
"""Scheduler output for a homogeneous batch (one NewRequestData per req)."""
new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
return DiffusionSchedulerOutput(
step_id=step_id,
scheduled_new_reqs=new_reqs,
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
)
def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
return DiffusionSchedulerOutput(
step_id=step_id,
scheduled_new_reqs=[],
scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
num_running_reqs=1,
num_waiting_reqs=0,
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
runner,
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
)
output = result.get_req_output("req-1")
output = result.get_request_output("req-1")
assert output.finished is True
assert output.result is not None
@@ -384,16 +384,16 @@ class TestRunner:
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
first = result.get_req_output("req-1")
assert first.req_id == "req-1"
first = result.get_request_output("req-1")
assert first.request_id == "req-1"
assert first.step_index == 1
assert first.finished is False
assert first.result is None
assert "req-1" in runner.state_cache
result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
second = result.get_req_output("req-1")
assert second.req_id == "req-1"
second = result.get_request_output("req-1")
assert second.request_id == "req-1"
assert second.step_index == 2
assert second.finished is True
assert second.result is not None
@@ -410,13 +410,13 @@ class TestRunner:
runner = _make_runner()
req_1 = _make_step_request()
req_2 = _make_step_request()
req_2.request_ids = ["req-2"]
req_2.request_id = "req-2"
scheduler_output = DiffusionSchedulerOutput(
step_id=0,
scheduled_new_reqs=[
NewRequestData(sched_req_id="req-1", req=req_1),
NewRequestData(sched_req_id="req-2", req=req_2),
NewRequestData(request_id="req-1", req=req_1),
NewRequestData(request_id="req-2", req=req_2),
],
scheduled_cached_reqs=CachedRequestData.make_empty(),
finished_req_ids=set(),
@@ -431,7 +431,7 @@ class TestRunner:
runner = _make_runner()
with pytest.raises(ValueError, match="Missing cached state"):
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
runner = _make_runner()
@@ -440,8 +440,8 @@ class TestRunner:
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
output = result.get_req_output("req-1")
assert output.req_id == "req-1"
output = result.get_request_output("req-1")
assert output.request_id == "req-1"
assert output.step_index == 0
assert output.finished is True
assert output.result is not None
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
worker = object.__new__(DiffusionWorker)
worker.lora_manager = lora_manager
worker._step_lora_state = {}
output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
return worker
@@ -525,9 +525,9 @@ class TestWorker:
"""DiffusionWorker.execute_stepwise"""
def test_delegates_to_model_runner(self):
expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
worker = _make_step_worker(expected_output=expected)
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
@@ -536,7 +536,7 @@ class TestWorker:
def test_deactivates_lora_when_request_has_no_adapter(self):
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
DiffusionWorker.execute_stepwise(worker, scheduler_output)
@@ -552,7 +552,7 @@ class TestWorker:
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
scheduler_output = _make_scheduler_output(request, request_id="req-1")
DiffusionWorker.execute_stepwise(worker, scheduler_output)
@@ -568,8 +568,8 @@ class TestWorker:
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
first = _make_scheduler_output(request, sched_req_id="req-1")
second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
first = _make_scheduler_output(request, request_id="req-1")
second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
DiffusionWorker.execute_stepwise(worker, first)
DiffusionWorker.execute_stepwise(worker, second)
@@ -610,10 +610,10 @@ class TestWorker:
next_request.sampling_params.lora_request = lora_request
worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
first = _make_scheduler_output(finishing, sched_req_id="req-1")
first = _make_scheduler_output(finishing, request_id="req-1")
next_batch = _make_scheduler_output(
next_request,
sched_req_id="req-2",
request_id="req-2",
step_id=1,
finished_req_ids={"req-1"},
)
@@ -633,11 +633,11 @@ class TestExecutor:
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
executor = object.__new__(MultiprocDiffusionExecutor)
executor._ensure_open = lambda: None
expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
executor.collective_rpc = mocker.Mock(return_value=expected)
request = _make_engine_request("req-step", num_inference_steps=2)
scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
scheduler_output = _make_scheduler_output(request, request_id="req-step")
output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
@@ -653,7 +653,7 @@ class TestEngine:
[
(
lambda _: RunnerOutput(
req_id="req-error",
request_id="req-error",
step_index=1,
finished=True,
result=DiffusionOutput(error="boom"),
@@ -688,7 +688,7 @@ class TestEngine:
call_count["n"] += 1
finished = call_count["n"] == 2
return RunnerOutput(
req_id="req-step",
request_id="req-step",
step_index=call_count["n"],
finished=finished,
result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
@@ -714,7 +714,7 @@ class TestEngine:
step["n"] += 1
engine.abort("req-stop")
return RunnerOutput(
req_id="req-stop",
request_id="req-stop",
step_index=1,
finished=False,
result=None,
@@ -741,7 +741,7 @@ class TestEngine:
assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
engine.abort("req-mid")
return RunnerOutput(
req_id="req-mid",
request_id="req-mid",
step_index=step["n"],
finished=False,
result=None,
@@ -760,7 +760,7 @@ class TestEngine:
engine = _make_engine(
scheduler,
execute_fn=lambda _: RunnerOutput(
req_id="req-missing",
request_id="req-missing",
step_index=1,
finished=True,
result=None,
@@ -777,7 +777,7 @@ class TestEngine:
class TestIPC:
def test_pack_unpack_runner_output_shm(self):
tensor = torch.zeros(300_000, dtype=torch.float32)
output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
packed = pack_diffusion_output_shm(output)
assert isinstance(packed.result.output, dict)
tests/diffusion/test_multiproc_engine_concurrency.py
+3 -3
View File
@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
def _mock_request(tag: str):
"""Return a lightweight request object identifiable by *tag*."""
return SimpleNamespace(
request_ids=[tag],
request_id=tag,
prompts=[f"prompt_{tag}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
)
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
req = req_q.get(timeout=10)
method = req.get("method", "")
args = req.get("args", ())
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
tag = f"result_for_{args[0].request_ids[0]}"
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
tag = f"result_for_{args[0].request_id}"
elif args:
tag = f"result_for_{args[0]}"
else:
tests/diffusion/test_stage_diffusion_proc.py
+1 -1
View File
@@ -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"]