From af8501dcc791e1f7b3775e6d154dbe41896239d9 Mon Sep 17 00:00:00 2001 From: Taichang Zhou Date: Thu, 22 Jan 2026 22:46:06 +0800 Subject: [PATCH 1/3] debug Qwen TTS Signed-off-by: Taichang Zhou --- .../qwen3_tts/configuration_qwen3_tts.py | 11 ++++++++-- .../models/qwen3_tts/modeling_qwen3_tts.py | 22 ++++++++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index dcf9d54ac..c7db33a59 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -504,12 +504,19 @@ class Qwen3TTSConfig(PretrainedConfig): self.tts_bos_token_id = tts_bos_token_id self.tts_eos_token_id = tts_eos_token_id + #TODO: remove these dummy values after + self.image_token_id = 0 # dummy image token id + self.video_token_id = 0 # dummy video token id + self.vision_start_token_id = 0 # dummy vision start token id + self.vision_config = PretrainedConfig() # dummy vision config + self.vision_config.spatial_merge_size = 1 + def get_text_config(self, **kwargs): # vLLM expects text config to expose hidden_size/num_attention_heads. # For Qwen3 TTS, the talker config is the text model config. config = self.talker_config - if hasattr(config, "rope_parameters"): - delattr(config, "rope_parameters") + # if hasattr(config, "rope_parameters"): + # delattr(config, "rope_parameters") return config diff --git a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py index ac9844f89..aac14574a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py @@ -426,7 +426,20 @@ def mel_spectrogram( mel_spec = dynamic_range_compression_torch(mel_spec) return mel_spec +def _compute_default_rope_parameters( + config, + device, +): + base = config.rope_theta + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, attention_factor class Qwen3TTSPreTrainedModel(PreTrainedModel): config_class = Qwen3TTSConfig @@ -500,7 +513,9 @@ class Qwen3TTSTalkerRotaryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + self.rope_init_fn: Callable = _compute_default_rope_parameters + if self.rope_type != "default": + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -536,7 +551,9 @@ class Qwen3TTSRotaryEmbedding(nn.Module): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + self.rope_init_fn: Callable = _compute_default_rope_parameters + if self.rope_type != "default": + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -1387,7 +1404,6 @@ class Qwen3TTSTalkerModel(Qwen3TTSTalkerTextPreTrainedModel): def __init__(self, config): super().__init__(config) - self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.layers = nn.ModuleList( [Qwen3TTSTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] From dc6177520e358fef6034b3871d7e5eec3e705dd7 Mon Sep 17 00:00:00 2001 From: Rustam Khadipash <16683750+hadipash@users.noreply.github.com> Date: Mon, 2 Feb 2026 18:28:25 +0800 Subject: [PATCH 2/3] [Feature] add Tensor Parallelism to LongCat-Image(-Edit) (#926) Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com> --- .../diffusion/parallelism_acceleration.md | 27 +++++++------ .../longcat_image_transformer.py | 38 +++++++++++-------- 2 files changed, 35 insertions(+), 30 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 7334df1d4..6e2c18d64 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -18,20 +18,19 @@ The following table shows which models are currently supported by parallelism me ### ImageGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | -|-------|------------------|------------|---------|--------------|--------------------------| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | -| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ | - +| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | +|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ | !!! note "TP Limitations for Diffusion Models" We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP. diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 9d216363d..2f6d54aef 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -9,10 +9,11 @@ import torch.nn as nn from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle +from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather from vllm.logger import init_logger from vllm.model_executor.layers.activation import get_act_fn from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear +from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata @@ -35,9 +36,9 @@ class FeedForward(nn.Module): inner_dim = int(dim * mult) dim_out = dim_out if dim_out is not None else dim - self.w_in = ReplicatedLinear(dim, inner_dim, bias=bias, return_bias=False) + self.w_in = ColumnParallelLinear(dim, inner_dim, bias=bias, return_bias=False) self.act = get_act_fn("gelu_pytorch_tanh") - self.w_out = ReplicatedLinear(inner_dim, dim_out, bias=bias, return_bias=False) + self.w_out = RowParallelLinear(inner_dim, dim_out, bias=bias, return_bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.w_in(hidden_states) @@ -83,12 +84,11 @@ class LongCatImageAttention(nn.Module): hidden_size=query_dim, head_size=self.head_dim, total_num_heads=self.heads, - disable_tp=True, bias=bias, ) if not self.pre_only: - self.to_out = torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias) + self.to_out = RowParallelLinear(self.inner_dim, self.out_dim, bias=out_bias) if self.added_kv_proj_dim is not None: self.norm_added_q = RMSNorm(dim_head, eps=eps) @@ -98,11 +98,10 @@ class LongCatImageAttention(nn.Module): hidden_size=self.added_kv_proj_dim, head_size=self.head_dim, total_num_heads=self.heads, - disable_tp=True, bias=added_proj_bias, ) - self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) + self.to_add_out = RowParallelLinear(self.inner_dim, query_dim, bias=out_bias) self.attn = Attention( num_heads=heads, @@ -185,22 +184,26 @@ class LongCatImageAttention(nn.Module): """ qkv, _ = self.to_qkv(hidden_states) - query, key, value = qkv.chunk(3, dim=-1) + 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) - query = query.unflatten(-1, (self.heads, -1)) - key = key.unflatten(-1, (self.heads, -1)) - value = value.unflatten(-1, (self.heads, -1)) + query = query.unflatten(-1, (self.to_qkv.num_heads, -1)) + key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) + value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) query = self.norm_q(query) key = self.norm_k(key) if self.added_kv_proj_dim is not None: encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) - encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1) + q_size = self.add_kv_proj.num_heads * self.head_dim + kv_size = self.add_kv_proj.num_kv_heads * self.head_dim + encoder_query, encoder_key, encoder_value = encoder_qkv.split([q_size, kv_size, kv_size], dim=-1) - encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) + encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1)) + encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1)) + encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1)) # Apply RMSNorm to text Q/K encoder_query = self.norm_added_q(encoder_query) @@ -284,11 +287,14 @@ class LongCatImageAttention(nn.Module): encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 ) - hidden_states = self.to_out(hidden_states) + hidden_states, _ = self.to_out(hidden_states) encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) return hidden_states, encoder_hidden_states else: + # For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction + if get_tensor_model_parallel_world_size() > 1: + hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1) return hidden_states From 47b11c6b130ca70aee5f434c069824b5341ae3b5 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Mon, 2 Feb 2026 20:36:56 +0800 Subject: [PATCH 3/3] [CI] Temporarily remove slow tests. (#1143) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Signed-off-by: princepride Co-authored-by: princepride --- .buildkite/pipeline.yml | 142 +++++++++--------- .../stage_configs/bagel_mooncake_ci.yaml | 4 +- .../stage_configs/bagel_sharedmemory_ci.yaml | 4 +- .../offline_inference/test_bagel_text2img.py | 4 +- 4 files changed, 77 insertions(+), 77 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index f7f3e4b05..402e62518 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -233,43 +233,43 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Omni Model Test with H100" - timeout_in_minutes: 30 - depends_on: image-build - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export VLLM_TEST_CLEAN_GPU_MEMORY="1" - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py - - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py - - pytest -s -v tests/e2e/online_serving/test_async_omni.py - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + # - label: "Omni Model Test with H100" + # timeout_in_minutes: 30 + # depends_on: image-build + # commands: + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py + # - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py + # - pytest -s -v tests/e2e/online_serving/test_async_omni.py + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 2 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate - label: "Diffusion Image Edit Test with H100 (1 GPU)" timeout_in_minutes: 20 @@ -306,37 +306,37 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: "Bagel Text2Img Model Test with H100" - timeout_in_minutes: 30 - depends_on: image-build - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + # - label: "Bagel Text2Img Model Test with H100" + # timeout_in_minutes: 30 + # depends_on: image-build + # commands: + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 1 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 4c6a67e8b..8c170dffd 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -12,7 +12,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.60 + gpu_memory_utilization: 0.35 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -44,7 +44,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.55 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index 67ff7309f..ef605a3a3 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -12,7 +12,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.65 + gpu_memory_utilization: 0.35 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -43,7 +43,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.4 + gpu_memory_utilization: 0.55 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 777af52a7..9aef47a74 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -161,7 +161,7 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima def test_bagel_text2img_shared_memory_connector(): """Test Bagel text2img with shared memory connector.""" config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") - omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path) + omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300) try: generated_image = _generate_bagel_image(omni) @@ -288,7 +288,7 @@ def test_bagel_text2img_mooncake_connector(): http_port=MOONCAKE_HTTP_PORT, ) - omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file) + omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300) generated_image = _generate_bagel_image(omni) _validate_pixels(generated_image)