diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
index 2426e96c4..2a825570e 100644
--- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx
+++ b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
@@ -39,7 +39,7 @@ Rows are grouped when a family shares the same runtime path or optimization supp
| LongCat-Image |
- meituan-longcat/LongCat-Image
|
+ meituan-longcat/LongCat-Imagemeituan-longcat/LongCat-Image-Editmeituan-longcat/LongCat-Image-Edit-Turbo
|
| SD3 / SD3.5 |
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py
index 8f620e9c0..2d0daf4a9 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py
@@ -1,3 +1,4 @@
+import math
from dataclasses import dataclass, field
from typing import Callable
@@ -20,6 +21,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
ModelDeploymentConfig,
)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_sp_world_size,
+ model_parallel_is_initialized,
+)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
@@ -252,6 +257,28 @@ class LongCatImagePipelineConfig(ImagePipelineConfig):
keep_resident_components=("text_encoder", "vae"),
)
+ def expand_conditioning_to_sample_batch(self, batch):
+ # Noise/reference latents are built at batch_size = prompts * num_outputs,
+ # but text encoding stays per-prompt; repeat the embeds to match. No-op
+ # for num_outputs == 1. Shared by T2I and Edit.
+ from sglang.multimodal_gen.runtime.utils.condition_expansion import (
+ PromptToSampleBatchExpander,
+ )
+
+ expander = PromptToSampleBatchExpander.from_batch(batch)
+ if expander is None:
+ return batch
+ for field_name in (
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ "prompt_embeds_mask",
+ "negative_prompt_embeds_mask",
+ "prompt_seq_lens",
+ "negative_prompt_seq_lens",
+ ):
+ expander.expand_field(batch, field_name)
+ return batch
+
# --- LatentPreparationStage hooks ---
def prepare_latent_shape(self, batch, batch_size, num_frames):
@@ -416,3 +443,164 @@ class LongCatImagePipelineConfig(ImagePipelineConfig):
noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0)
return noise_pred * scale
+
+
+def _calculate_edit_dimensions(target_area, ratio):
+ """Output size for LongCat-Image-Edit: fit `target_area`, ceil to /16.
+
+ Copied from diffusers pipeline_longcat_image_edit.calculate_dimensions.
+ Note this intentionally differs from sglang.multimodal_gen.utils
+ calculate_dimensions (which rounds to /32).
+ """
+ width = math.sqrt(target_area * ratio)
+ height = width / ratio
+
+ width = width if width % 16 == 0 else (width // 16 + 1) * 16
+ height = height if height % 16 == 0 else (height // 16 + 1) * 16
+
+ return int(width), int(height)
+
+
+@dataclass
+class LongCatImageEditPipelineConfig(LongCatImagePipelineConfig):
+ """Configuration for the LongCat-Image-Edit I2I pipeline.
+
+ Mirrors diffusers LongCatImageEditPipeline: the reference image is resized
+ to the output resolution, VAE-encoded with argmax sampling, packed, and
+ concatenated after the noisy latents along the sequence dim. RoPE ids use
+ modality 1 for noisy tokens and modality 2 for reference tokens, both
+ offset by the full text sequence length (VL image tokens + 512 body).
+ """
+
+ task_type: ModelTaskType = ModelTaskType.I2I
+ # diffusers reference draws latent noise with a CPU generator
+ generator_device: str = "cpu"
+
+ # --- InputValidationStage hooks ---
+
+ def calculate_condition_image_size(self, image, width, height):
+ return _calculate_edit_dimensions(1024 * 1024, width / height)
+
+ # --- LatentPreparationStage hooks ---
+
+ def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
+ # The edit reference draws noise directly in the prompt-embeds dtype
+ # (bf16), unlike the T2I reference which draws in float32.
+ return prompt_dtype
+
+ def maybe_prepare_latent_ids(self, latents):
+ # img_ids (noisy + reference) are built per-step in
+ # prepare_*_cond_kwargs because the text start offset depends on the
+ # VL image token count, which is unknown at latent preparation time.
+ return None
+
+ # --- ImageVAEEncodingStage hooks ---
+
+ def preprocess_vae_encode(self, image, vae):
+ # AutoencoderKL is a 2D image VAE; drop the frames dim added by
+ # ImageVAEEncodingStage ([B, C, 1, H, W] -> [B, C, H, W]).
+ if image.dim() == 5 and image.shape[2] == 1:
+ image = image.squeeze(2)
+ return image
+
+ def postprocess_image_latent(self, latent_condition, batch):
+ if latent_condition.dim() == 5 and latent_condition.shape[2] == 1:
+ latent_condition = latent_condition.squeeze(2)
+ batch_size = batch.batch_size
+ if batch_size > latent_condition.shape[0]:
+ if batch_size % latent_condition.shape[0] != 0:
+ raise ValueError(
+ f"Cannot duplicate reference image of batch size "
+ f"{latent_condition.shape[0]} to {batch_size} prompts."
+ )
+ latent_condition = latent_condition.repeat(
+ batch_size // latent_condition.shape[0], 1, 1, 1
+ )
+ _, num_channels_latents, height, width = latent_condition.shape
+ return _pack_latents(
+ latent_condition, batch_size, num_channels_latents, height, width
+ )
+
+ # --- Denoising hooks ---
+
+ def shard_latents_for_sp(self, batch, latents):
+ # (h/2)*(w/2) is odd at most ~1MP edit resolutions, so SP has to pad, and
+ # the pads stay unmasked (USPAttention rejects a mask alongside the
+ # replicated text prefix). Repeat the last token instead of the base
+ # class's zeros, which would carry the RoPE of text token 0.
+ if latents.dim() == 3 and model_parallel_is_initialized():
+ sp_world_size = get_sp_world_size()
+ remainder = latents.shape[1] % sp_world_size
+ if remainder:
+ pad = latents[:, -1:].expand(-1, sp_world_size - remainder, -1)
+ latents = torch.cat([latents, pad], dim=1)
+ return super().shard_latents_for_sp(batch, latents)
+
+ def _maybe_shard_pos_ids_for_sp(self, batch, pos_ids):
+ # RoPE ids must be sharded exactly like their latents, so reuse the same
+ # helper. It reads the SP group, which only exists under model parallelism.
+ if not model_parallel_is_initialized() or get_sp_world_size() == 1:
+ return pos_ids
+ sharded, _ = self.shard_latents_for_sp(batch, pos_ids.unsqueeze(0))
+ return sharded.squeeze(0)
+
+ def _edit_img_ids(self, batch, num_token, device):
+ """Position ids for [noisy | reference] packed latent tokens.
+
+ The reference image is resized to the output resolution, so both grids
+ share the same shape; they differ only in modality id (1 vs 2).
+ """
+ vae_scale_factor = self.vae_config.get_vae_scale_factor()
+ h = 2 * (int(batch.height) // (vae_scale_factor * 2))
+ w = 2 * (int(batch.width) // (vae_scale_factor * 2))
+ noisy_ids = _prepare_pos_ids(
+ modality_id=1,
+ token_type="image",
+ start=(num_token, num_token),
+ height=h // 2,
+ width=w // 2,
+ )
+ ref_ids = _prepare_pos_ids(
+ modality_id=2,
+ token_type="image",
+ start=(num_token, num_token),
+ height=h // 2,
+ width=w // 2,
+ )
+ noisy_ids = self._maybe_shard_pos_ids_for_sp(batch, noisy_ids)
+ ref_ids = self._maybe_shard_pos_ids_for_sp(batch, ref_ids)
+ return torch.cat([noisy_ids, ref_ids], dim=0).to(device)
+
+ def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+ num_token = batch.prompt_embeds[0].shape[1]
+ return {
+ "txt_ids": _prepare_pos_ids(
+ modality_id=0, token_type="text", start=(0, 0), num_token=num_token
+ ).to(device),
+ "img_ids": self._edit_img_ids(batch, num_token, device),
+ }
+
+ def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+ num_token = batch.negative_prompt_embeds[0].shape[1]
+ return {
+ "txt_ids": _prepare_pos_ids(
+ modality_id=0, token_type="text", start=(0, 0), num_token=num_token
+ ).to(device),
+ "img_ids": self._edit_img_ids(batch, num_token, device),
+ }
+
+ def slice_noise_pred(self, noise, latents):
+ # Drop predictions over the appended reference-image tokens.
+ return noise[:, : latents.size(1)]
+
+ def post_denoising_loop(self, latents, batch):
+ # The SP gather leaves the noisy latents at their padded length; trim the
+ # trailing pad tokens before unpacking to the (h/2)*(w/2) grid.
+ if latents.dim() == 3:
+ vae_scale_factor = self.vae_config.get_vae_scale_factor()
+ h = 2 * (int(batch.height) // (vae_scale_factor * 2))
+ w = 2 * (int(batch.width) // (vae_scale_factor * 2))
+ expected = (h // 2) * (w // 2)
+ if latents.shape[1] > expected:
+ latents = latents[:, :expected, :]
+ return super().post_denoising_loop(latents, batch)
diff --git a/python/sglang/multimodal_gen/configs/sample/longcat_image.py b/python/sglang/multimodal_gen/configs/sample/longcat_image.py
index c9999436b..91d35b5c2 100644
--- a/python/sglang/multimodal_gen/configs/sample/longcat_image.py
+++ b/python/sglang/multimodal_gen/configs/sample/longcat_image.py
@@ -12,3 +12,28 @@ class LongCatImageSamplingParams(SamplingParams):
# Override base class defaults to enable LongCat-specific features by default
enable_cfg_renorm: bool = True
enable_prompt_rewrite: bool = True
+
+
+@dataclass
+class LongCatImageEditSamplingParams(SamplingParams):
+ """Defaults for LongCat-Image-Edit (mirrors diffusers LongCatImageEditPipeline).
+
+ Output height/width are derived from the condition image (~1MP), so no
+ defaults are set here. The reference uses an empty negative prompt and
+ plain CFG (no renorm, no prompt rewrite).
+ """
+
+ num_frames: int = 1
+ num_inference_steps: int = 50
+ guidance_scale: float = 4.5
+ negative_prompt: str = ""
+ enable_cfg_renorm: bool = False
+ enable_prompt_rewrite: bool = False
+
+
+@dataclass
+class LongCatImageEditTurboSamplingParams(LongCatImageEditSamplingParams):
+ """LongCat-Image-Edit-Turbo: distilled, 8 steps, CFG disabled."""
+
+ num_inference_steps: int = 8
+ guidance_scale: float = 1.0
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index e78679527..9d09359ad 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/python/sglang/multimodal_gen/registry.py
@@ -74,6 +74,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.lingbot_video_moe import (
LingBotVideoMoEPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+ LongCatImageEditPipelineConfig,
LongCatImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig
@@ -150,6 +151,8 @@ from sglang.multimodal_gen.configs.sample.lingbot_world import (
LingBotWorldSamplingParams,
)
from sglang.multimodal_gen.configs.sample.longcat_image import (
+ LongCatImageEditSamplingParams,
+ LongCatImageEditTurboSamplingParams,
LongCatImageSamplingParams,
)
from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams
@@ -1327,6 +1330,34 @@ def _register_configs():
],
)
+ # LongCat-Image-Edit-Turbo (registered before Edit so its detector wins)
+ register_configs(
+ sampling_param_cls=LongCatImageEditTurboSamplingParams,
+ pipeline_config_cls=LongCatImageEditPipelineConfig,
+ hf_model_paths=[
+ "meituan-longcat/LongCat-Image-Edit-Turbo",
+ ],
+ model_detectors=[
+ lambda hf_id: "longcat" in hf_id.lower()
+ and "edit" in hf_id.lower()
+ and "turbo" in hf_id.lower(),
+ ],
+ )
+
+ # LongCat-Image-Edit
+ register_configs(
+ sampling_param_cls=LongCatImageEditSamplingParams,
+ pipeline_config_cls=LongCatImageEditPipelineConfig,
+ hf_model_paths=[
+ "meituan-longcat/LongCat-Image-Edit",
+ ],
+ model_detectors=[
+ lambda hf_id: "longcat" in hf_id.lower()
+ and "edit" in hf_id.lower()
+ and "turbo" not in hf_id.lower(),
+ ],
+ )
+
_register_configs()
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
index f86f80204..c7383ccf2 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
@@ -12,6 +12,10 @@ from transformers import PretrainedConfig
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
from sglang.multimodal_gen.configs.models import EncoderConfig
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+ LongCatImageEditPipelineConfig,
+ LongCatImagePipelineConfig,
+)
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
QwenImageEditPipelineConfig,
)
@@ -869,14 +873,16 @@ class TextEncoderLoader(ComponentLoader):
with model_device, skip_init_modules():
architectures = getattr(model_config, "architectures", [])
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
- enable_image_understanding = (
- True
- if isinstance(
- server_args.pipeline_config, QwenImageEditPipelineConfig
- )
- else False
+ enable_image_understanding = isinstance(
+ server_args.pipeline_config,
+ (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
)
model_config.enable_image_understanding = enable_image_understanding
+ # LongCat feeds its padded body to the DiT, so it must mask
+ # padding on the cache-free path; scoped so others are unchanged.
+ model_config.honor_cache_free_padding_mask = isinstance(
+ server_args.pipeline_config, LongCatImagePipelineConfig
+ )
model = model_cls(model_config)
if not isinstance(model, EncoderTensorParallelMixin):
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
index 9d22fe7fa..f7c2dbc08 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
@@ -443,6 +443,7 @@ class _LongCatSingleAttention(nn.Module):
self,
hidden_states: torch.Tensor,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ num_replicated_prefix: int = 0,
cos_sin_cache: Optional[torch.Tensor] = None,
positions: Optional[torch.Tensor] = None,
) -> torch.Tensor:
@@ -464,7 +465,7 @@ class _LongCatSingleAttention(nn.Module):
positions,
)
- x = self.attn(q, k, v)
+ x = self.attn(q, k, v, num_replicated_prefix=num_replicated_prefix)
return x.flatten(2, 3).to(q.dtype)
@@ -564,6 +565,8 @@ class _SingleTransformerBlock(nn.Module):
attn_output = self.attn(
hidden_states=norm_hidden_states,
image_rotary_emb=image_rotary_emb,
+ # Text is replicated per SP rank; keep it out of the all-to-all.
+ num_replicated_prefix=text_seq_len,
cos_sin_cache=cos_sin_cache,
positions=positions,
)
diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
index d6bf3dbc4..ab8b74f0e 100644
--- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
+++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
@@ -194,6 +194,9 @@ class Qwen2_5_VLAttention(nn.Module):
def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None):
super().__init__()
self.config = config
+ self.honor_cache_free_padding_mask = getattr(
+ config, "honor_cache_free_padding_mask", False
+ )
self.layer_idx = layer_idx
if layer_idx is None:
logger.warning(
@@ -315,14 +318,14 @@ class Qwen2_5_VLAttention(nn.Module):
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
- # Diffusion text encoding is cache-free and historically uses the native
- # causal kernel; its trailing padding is removed during postprocessing.
- # Cached generation still needs the explicit mask for padded batches.
+ # LongCat masks padding on the cache-free path too; others keep the
+ # original mask-free fast path unchanged.
+ honor_mask = use_cache or self.honor_cache_free_padding_mask
attn_output = self.attn(
query_states,
key_states,
value_states,
- attn_mask=attention_mask if use_cache else None,
+ attn_mask=attention_mask if honor_mask else None,
)
attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
@@ -425,10 +428,27 @@ class Qwen2_5_VLDecoderLayer(nn.Module):
return hidden_states
+def _build_causal_padding_mask(
+ attention_mask: torch.Tensor, inputs_embeds: torch.Tensor
+) -> torch.Tensor:
+ """Build a causal+padding bool mask ``[batch, 1, q_len, kv_len]``; True attends."""
+ q_len = inputs_embeds.shape[1]
+ kv_len = attention_mask.shape[-1]
+ device = inputs_embeds.device
+ q_idx = torch.arange(kv_len - q_len, kv_len, device=device).unsqueeze(-1)
+ kv_idx = torch.arange(kv_len, device=device).unsqueeze(0)
+ causal = (kv_idx <= q_idx).unsqueeze(0)
+ padding = attention_mask.to(device=device, dtype=torch.bool).unsqueeze(1)
+ return (causal & padding).unsqueeze(1)
+
+
class Qwen2_5_VLTextModel(nn.Module):
def __init__(self, config: PretrainedConfig):
super().__init__()
self.config = config
+ self.honor_cache_free_padding_mask = getattr(
+ config, "honor_cache_free_padding_mask", False
+ )
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
@@ -552,6 +572,19 @@ class Qwen2_5_VLTextModel(nn.Module):
create_sliding_window_causal_mask(**mask_kwargs)
)
+ # create_causal_mask returns None for this _attn_implementation, so
+ # build the causal+padding mask here. LongCat-only.
+ if (
+ self.honor_cache_free_padding_mask
+ and causal_mask_mapping["full_attention"] is None
+ and isinstance(attention_mask, torch.Tensor)
+ and attention_mask.dim() == 2
+ and not bool(attention_mask.all())
+ ):
+ causal_mask_mapping["full_attention"] = _build_causal_padding_mask(
+ attention_mask, inputs_embeds
+ )
+
hidden_states = inputs_embeds
# decoder layers
@@ -1123,7 +1156,12 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder):
super().__init__(config)
enable_image_understanding = config.enable_image_understanding
generation_config = config.generation_config
+ # LongCat-only; propagate to the text config (see TextEncoderLoader).
+ honor_cache_free_padding_mask = getattr(
+ config, "honor_cache_free_padding_mask", False
+ )
config = config.arch_config
+ config.text_config.honor_cache_free_padding_mask = honor_cache_free_padding_mask
self.model = Qwen2_5_VLModel(
config, enable_image_understanding=enable_image_understanding
)
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py b/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py
index 87222a1e7..b843026ad 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py
@@ -1,12 +1,19 @@
-"""LongCat-Image pipeline for SGLang."""
+"""LongCat-Image pipelines (T2I and Edit) for SGLang."""
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase,
)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+ ImageVAEEncodingStage,
+ InputValidationStage,
+)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.longcat_image import (
LongCatPromptRewriteStage,
)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.longcat_image_edit import (
+ LongCatImageEditTextEncodingStage,
+)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
@@ -68,4 +75,62 @@ class LongCatImagePipeline(LoRAPipeline, ComposedPipelineBase):
self.add_standard_decoding_stage()
-EntryClass = [LongCatImagePipeline]
+class LongCatImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
+ """Pipeline for LongCat-Image-Edit image editing (I2I).
+
+ Mirrors diffusers LongCatImageEditPipeline. The output resolution is
+ derived from the condition image (~1MP, /16); the reference image is
+ VAE-encoded (argmax) and concatenated after the noisy latents; the
+ edit instruction is encoded jointly with the image via Qwen2.5-VL.
+ """
+
+ pipeline_name = "LongCatImageEditPipeline"
+
+ _required_config_modules = [
+ "text_encoder",
+ "tokenizer",
+ "text_processor",
+ "vae",
+ "transformer",
+ "scheduler",
+ ]
+
+ def create_pipeline_stages(self, server_args: ServerArgs):
+ # 1. Load the condition image, resize it to the calculated output
+ # resolution, and set batch.height/width (pipeline config hooks).
+ self.add_stage(InputValidationStage())
+
+ # 2. Joint text+image (VL) prompt encoding. Also encodes the negative
+ # prompt against the same image when CFG is enabled.
+ self.add_stage(
+ LongCatImageEditTextEncodingStage(
+ text_encoder=self.get_module("text_encoder"),
+ tokenizer=self.get_module("tokenizer"),
+ text_processor=self.get_module("text_processor"),
+ text_encoder_dtype=PRECISION_TO_TYPE[
+ server_args.pipeline_config.text_encoder_precisions[0]
+ ],
+ )
+ )
+
+ # 3. Reference-image VAE encoding -> packed batch.image_latent, which
+ # DenoisingStage concatenates after the noisy latents (dim=1).
+ self.add_stage(ImageVAEEncodingStage(vae=self.get_module("vae")))
+
+ # 4. Latent preparation (noise drawn in prompt-embeds dtype).
+ self.add_standard_latent_preparation_stage()
+
+ # 5. Timestep preparation (mu from the packed noisy token count only).
+ self.add_standard_timestep_preparation_stage(
+ prepare_extra_kwargs=[_prepare_mu],
+ )
+
+ # 6. Standard denoising loop; slice_noise_pred drops the reference
+ # tokens from each prediction before CFG/scheduler step.
+ self.add_standard_denoising_stage()
+
+ # 7. Standard VAE decoding
+ self.add_standard_decoding_stage()
+
+
+EntryClass = [LongCatImagePipeline, LongCatImageEditPipeline]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
index 5c6b0c201..3f0ada45f 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
@@ -257,6 +257,7 @@ class ComposedPipelineBase(ABC):
"QwenImageEditPipeline": {"vae"},
"QwenImageEditPlusPipeline": {"vae"},
"QwenImageLayeredPipeline": {"vae", "transformer"},
+ "LongCatImageEditPipeline": {"vae"},
"GlmImagePipeline": {"vae", "transformer"},
"WanImageToVideoPipeline": {"vae"},
"WanImageToVideoDmdPipeline": {"vae"},
@@ -1063,6 +1064,10 @@ class ComposedPipelineBase(ABC):
main_process_only=True,
)
+ self.component_residency_manager = get_global_component_residency_manager(
+ self, server_args
+ )
+ self.executor.component_residency_manager = self.component_residency_manager
return self.executor.execute_group_with_profiling(
self.stages, batches, server_args
)
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py
new file mode 100644
index 000000000..5a3532339
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py
@@ -0,0 +1,256 @@
+"""Text+image (VL) prompt encoding stage for LongCat-Image-Edit (I2I).
+
+Mirrors diffusers ``LongCatImageEditPipeline._encode_prompt``: the condition
+image (already resized to the output resolution by ``InputValidationStage``)
+is downscaled by 2x and fed to the Qwen2.5-VL vision tower; the prompt body is
+quote-aware tokenized to a fixed 512 tokens; the edit system prefix has its
+``<|image_pad|>`` placeholder expanded to the per-image token count. The
+resulting hidden states are sliced from the ``<|vision_start|>`` token through
+the 512-token body, so the DiT conditioning includes the VL image tokens.
+"""
+
+import PIL.Image
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+ _tokenize_prompt_for_encode,
+)
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+ ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+IMAGE_TOKEN = "<|image_pad|>"
+
+# Copied from diffusers LongCatImageEditPipeline.prompt_template_encode_prefix/suffix.
+PROMPT_TEMPLATE_ENCODE_PREFIX = (
+ "<|im_start|>system\nAs an image editing expert, first analyze the content and "
+ "attributes of the input image(s). Then, based on the user's editing instructions, "
+ "clearly and precisely determine how to modify the given image(s), ensuring that "
+ "only the specified parts are altered and all other aspects remain consistent with "
+ "the original(s).<|im_end|>\n<|im_start|>user\n"
+ "<|vision_start|><|image_pad|><|vision_end|>"
+)
+PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+
+
+class LongCatImageEditTextEncodingStage(PipelineStage):
+ """Encode the edit instruction together with the reference image.
+
+ The stage fills ``batch.prompt_embeds`` (and negative counterparts under
+ CFG) with hidden states sliced as ``[vision_start ... vision_end, 512-token
+ body]``. Both branches share the same VL image inputs, matching the
+ diffusers reference which encodes positive and negative prompts against
+ the same ``prompt_image``.
+ """
+
+ deduplicated_output_fields = (
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ "prompt_embeds_mask",
+ "negative_prompt_embeds_mask",
+ "prompt_seq_lens",
+ "negative_prompt_seq_lens",
+ )
+
+ def __init__(
+ self,
+ text_encoder,
+ tokenizer,
+ text_processor,
+ text_encoder_dtype: torch.dtype,
+ ):
+ super().__init__()
+ self.text_encoder = text_encoder
+ self.tokenizer = tokenizer
+ self.text_processor = text_processor
+ self.text_encoder_dtype = text_encoder_dtype
+ self._suffix_ids: list[int] | None = None
+
+ def component_uses(
+ self, server_args: ServerArgs, stage_name: str | None = None
+ ) -> list[ComponentUse]:
+ stage_name = self._component_stage_name(stage_name)
+ return [
+ ComponentUse(
+ stage_name,
+ "text_encoder",
+ target_dtype=self.text_encoder_dtype,
+ memory_intensive=True,
+ ),
+ ]
+
+ def _get_suffix_ids(self) -> list[int]:
+ if self._suffix_ids is None:
+ self._suffix_ids = self.tokenizer(
+ PROMPT_TEMPLATE_ENCODE_SUFFIX, add_special_tokens=False
+ )["input_ids"]
+ return self._suffix_ids
+
+ def _build_expanded_prefix_ids(
+ self, image_grid_thw: torch.Tensor
+ ) -> tuple[list[int], int]:
+ """Expand ``<|image_pad|>`` to the per-image token count and tokenize.
+
+ Returns (prefix_token_ids, prefix_len) where prefix_len is the index of
+ ``<|vision_start|>`` — the slice start for the DiT conditioning.
+ """
+ merge_length = self.text_processor.image_processor.merge_size**2
+ num_image_tokens = int(image_grid_thw.prod().item()) // merge_length
+ text = PROMPT_TEMPLATE_ENCODE_PREFIX
+ while IMAGE_TOKEN in text:
+ text = text.replace(IMAGE_TOKEN, "<|placeholder|>" * num_image_tokens, 1)
+ text = text.replace("<|placeholder|>", IMAGE_TOKEN)
+
+ prefix_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
+ vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+ prefix_len = prefix_ids.index(vision_start_id)
+ return prefix_ids, prefix_len
+
+ def _encode(
+ self,
+ prompt: list[str],
+ pixel_values: torch.Tensor,
+ image_grid_thw: torch.Tensor,
+ prefix_ids: list[int],
+ prefix_len: int,
+ device: torch.device,
+ ) -> torch.Tensor:
+ body = _tokenize_prompt_for_encode(prompt, self.tokenizer)
+ suffix_ids = self._get_suffix_ids()
+ suffix_len = len(suffix_ids)
+ batch_size = body.input_ids.size(0)
+
+ prefix_ids_t = (
+ torch.tensor(prefix_ids, dtype=body.input_ids.dtype)
+ .unsqueeze(0)
+ .expand(batch_size, -1)
+ )
+ suffix_ids_t = (
+ torch.tensor(suffix_ids, dtype=body.input_ids.dtype)
+ .unsqueeze(0)
+ .expand(batch_size, -1)
+ )
+ prefix_mask_t = torch.ones(
+ batch_size, len(prefix_ids), dtype=body.attention_mask.dtype
+ )
+ suffix_mask_t = torch.ones(
+ batch_size, suffix_len, dtype=body.attention_mask.dtype
+ )
+
+ input_ids = torch.cat((prefix_ids_t, body.input_ids, suffix_ids_t), dim=-1).to(
+ device
+ )
+ attention_mask = torch.cat(
+ (prefix_mask_t, body.attention_mask, suffix_mask_t), dim=-1
+ ).to(device)
+
+ with set_forward_context(current_timestep=0, attn_metadata=None):
+ outputs = self.text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+
+ hidden_states = outputs.hidden_states[-1]
+ # Keep [vision_start ... vision_end, 512-token body].
+ return hidden_states[:, prefix_len:-suffix_len, :]
+
+ @staticmethod
+ def _all_ones_conditioning(prompt_embeds: torch.Tensor):
+ batch_size, seq_len = prompt_embeds.shape[:2]
+ mask = torch.ones(
+ batch_size, seq_len, dtype=torch.bool, device=prompt_embeds.device
+ )
+ return mask, [seq_len] * batch_size
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ device = get_local_torch_device()
+
+ image = batch.condition_image
+ if isinstance(image, list):
+ if len(image) != 1:
+ raise ValueError(
+ "LongCat-Image-Edit supports exactly one condition image, "
+ f"got {len(image)}."
+ )
+ image = image[0]
+ if image is None:
+ # Mirrors ImageEncodingStage: nothing to encode without an image
+ # (e.g. warmup requests). Downstream stages will validate inputs.
+ logger.debug(
+ "LongCatImageEditTextEncodingStage skipped: no condition image."
+ )
+ return batch
+
+ prompt = batch.prompt if isinstance(batch.prompt, list) else [batch.prompt]
+ if len(prompt) != 1:
+ raise ValueError(
+ "LongCat-Image-Edit supports a single prompt per request, got "
+ f"{len(prompt)}."
+ )
+
+ # The VL tower sees the condition image at half the output resolution
+ # (matching diffusers: image_processor.resize(image, h // 2, w // 2)).
+ prompt_image = image.resize(
+ (int(batch.width) // 2, int(batch.height) // 2),
+ PIL.Image.Resampling.LANCZOS,
+ )
+ vl_inputs = self.text_processor.image_processor(
+ images=prompt_image, return_tensors="pt"
+ )
+ image_grid_thw = vl_inputs["image_grid_thw"].to(device)
+ prefix_ids, prefix_len = self._build_expanded_prefix_ids(image_grid_thw)
+
+ with self.use_declared_component(
+ component_name="text_encoder",
+ module=self.text_encoder,
+ ) as text_encoder:
+ assert text_encoder is not None
+ self.text_encoder = text_encoder
+
+ pixel_values = vl_inputs["pixel_values"].to(
+ device=device, dtype=self.text_encoder_dtype
+ )
+
+ prompt_embeds = self._encode(
+ prompt, pixel_values, image_grid_thw, prefix_ids, prefix_len, device
+ )
+ batch.prompt_embeds.append(prompt_embeds)
+ mask, seq_lens = self._all_ones_conditioning(prompt_embeds)
+ batch.prompt_embeds_mask = [mask]
+ batch.prompt_seq_lens = [seq_lens]
+
+ if batch.do_classifier_free_guidance:
+ negative_prompt = batch.negative_prompt or ""
+ if isinstance(negative_prompt, list):
+ negative_prompt = negative_prompt[:1]
+ else:
+ negative_prompt = [negative_prompt]
+ negative_prompt_embeds = self._encode(
+ negative_prompt,
+ pixel_values,
+ image_grid_thw,
+ prefix_ids,
+ prefix_len,
+ device,
+ )
+ batch.negative_prompt_embeds.append(negative_prompt_embeds)
+ neg_mask, neg_seq_lens = self._all_ones_conditioning(
+ negative_prompt_embeds
+ )
+ batch.negative_prompt_embeds_mask = [neg_mask]
+ batch.negative_prompt_seq_lens = [neg_seq_lens]
+
+ return batch
diff --git a/python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py b/python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py
new file mode 100644
index 000000000..3069a8630
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py
@@ -0,0 +1,157 @@
+"""Unit tests for LongCat-Image-Edit pipeline config hooks (CPU only)."""
+
+import types
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+ LongCatImageEditPipelineConfig,
+ LongCatImagePipelineConfig,
+ _calculate_edit_dimensions,
+)
+
+
+@pytest.fixture
+def edit_config():
+ return LongCatImageEditPipelineConfig()
+
+
+def _make_batch(**kwargs):
+ return types.SimpleNamespace(**kwargs)
+
+
+def test_edit_dimensions_match_diffusers_formula():
+ import math
+
+ for w, h in [(1000, 1500), (1024, 1024), (1920, 1080), (333, 777)]:
+ ratio = w / h
+ got_w, got_h = _calculate_edit_dimensions(1024 * 1024, ratio)
+ # diffusers reference (ceil to /16)
+ ref_w = math.sqrt(1024 * 1024 * ratio)
+ ref_h = ref_w / ratio
+ ref_w = ref_w if ref_w % 16 == 0 else (ref_w // 16 + 1) * 16
+ ref_h = ref_h if ref_h % 16 == 0 else (ref_h // 16 + 1) * 16
+ assert (got_w, got_h) == (int(ref_w), int(ref_h))
+ assert got_w % 16 == 0 and got_h % 16 == 0
+
+
+def test_edit_config_task_type_and_generator(edit_config):
+ assert edit_config.task_type == ModelTaskType.I2I
+ assert edit_config.generator_device == "cpu"
+
+
+def test_slice_noise_pred_drops_reference_tokens(edit_config):
+ latents = torch.zeros(1, 100, 64)
+ noise = torch.arange(200, dtype=torch.float32).view(1, 200, 1).expand(1, 200, 64)
+ sliced = edit_config.slice_noise_pred(noise, latents)
+ assert sliced.shape == (1, 100, 64)
+ assert torch.equal(sliced, noise[:, :100])
+
+
+def test_maybe_prepare_latent_ids_returns_none(edit_config):
+ assert edit_config.maybe_prepare_latent_ids(torch.zeros(1, 16, 8, 8)) is None
+
+
+def test_preprocess_vae_encode_squeezes_frames_dim(edit_config):
+ image_5d = torch.zeros(1, 3, 1, 64, 64)
+ assert edit_config.preprocess_vae_encode(image_5d, None).shape == (1, 3, 64, 64)
+ image_4d = torch.zeros(1, 3, 64, 64)
+ assert edit_config.preprocess_vae_encode(image_4d, None).shape == (1, 3, 64, 64)
+
+
+def test_postprocess_image_latent_packs_2x2(edit_config):
+ latent = torch.randn(1, 16, 8, 8)
+ batch = _make_batch(batch_size=1)
+ packed = edit_config.postprocess_image_latent(latent, batch)
+ assert packed.shape == (1, 16, 64) # (8//2 * 8//2, 16*4)
+
+
+def test_edit_img_ids_modalities_and_offset(edit_config):
+ # 1024x1024 -> latent 128x128 -> packed grid 64x64 -> 4096 tokens per image
+ batch = _make_batch(height=1024, width=1024)
+ num_token = 600 # text length including VL image tokens
+ img_ids = edit_config._edit_img_ids(batch, num_token, torch.device("cpu"))
+ assert img_ids.shape == (2 * 64 * 64, 3)
+ noisy_ids, ref_ids = img_ids[: 64 * 64], img_ids[64 * 64 :]
+ assert torch.all(noisy_ids[:, 0] == 1)
+ assert torch.all(ref_ids[:, 0] == 2)
+ # same grid layout, offset by text length
+ assert torch.equal(noisy_ids[:, 1:], ref_ids[:, 1:])
+ assert noisy_ids[0, 1].item() == num_token
+ assert noisy_ids[0, 2].item() == num_token
+ assert noisy_ids[-1, 1].item() == num_token + 63
+ assert noisy_ids[-1, 2].item() == num_token + 63
+
+
+def test_prepare_pos_cond_kwargs_uses_prompt_embeds_length(edit_config):
+ prompt_embeds = torch.zeros(1, 700, 3072)
+ batch = _make_batch(height=512, width=512, prompt_embeds=[prompt_embeds])
+ kwargs = edit_config.prepare_pos_cond_kwargs(
+ batch, torch.device("cpu"), None, torch.float32
+ )
+ txt_ids, img_ids = kwargs["txt_ids"], kwargs["img_ids"]
+ assert txt_ids.shape == (700, 3)
+ assert torch.all(txt_ids[:, 0] == 0)
+ # 512 -> latent 64 -> packed grid 32x32
+ assert img_ids.shape == (2 * 32 * 32, 3)
+ assert img_ids[0, 1].item() == 700
+
+
+def test_t2i_config_unchanged():
+ t2i = LongCatImagePipelineConfig()
+ assert t2i.task_type == ModelTaskType.T2I
+ # T2I still slices nothing and builds latent ids at preparation time
+ noise = torch.zeros(1, 10, 64)
+ assert t2i.slice_noise_pred(noise, torch.zeros(1, 5, 64)).shape == (1, 10, 64)
+ ids = t2i.maybe_prepare_latent_ids(torch.zeros(1, 16, 8, 8))
+ assert ids is not None and ids.shape == (16, 3)
+
+
+def test_expand_conditioning_repeats_embeds_for_num_outputs():
+ # Single prompt, num_outputs_per_prompt=2: text encoding produces per-prompt
+ # (batch 1) embeds while latents are built at batch 2, so the conditioning
+ # must be repeated to match or the DiT sees mismatched batch dims.
+ config = LongCatImageEditPipelineConfig()
+ batch = _make_batch(
+ prompt=["edit it"],
+ num_outputs_per_prompt=2,
+ prompt_embeds=[torch.randn(1, 850, 3584)],
+ negative_prompt_embeds=[torch.randn(1, 850, 3584)],
+ prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+ negative_prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+ prompt_seq_lens=[[850]],
+ negative_prompt_seq_lens=[[850]],
+ )
+ pos0 = batch.prompt_embeds[0]
+
+ config.expand_conditioning_to_sample_batch(batch)
+
+ assert batch.prompt_embeds[0].shape == (2, 850, 3584)
+ assert batch.negative_prompt_embeds[0].shape == (2, 850, 3584)
+ assert batch.prompt_embeds_mask[0].shape == (2, 850)
+ assert batch.negative_prompt_embeds_mask[0].shape == (2, 850)
+ assert batch.prompt_seq_lens[0] == [850, 850]
+ assert batch.negative_prompt_seq_lens[0] == [850, 850]
+ # Each sample is the original prompt, not garbage.
+ assert torch.equal(batch.prompt_embeds[0][0], pos0[0])
+ assert torch.equal(batch.prompt_embeds[0][1], pos0[0])
+
+
+def test_expand_conditioning_noop_for_single_output():
+ config = LongCatImagePipelineConfig()
+ pe = torch.randn(1, 850, 3584)
+ batch = _make_batch(
+ prompt=["a photo"],
+ num_outputs_per_prompt=1,
+ prompt_embeds=[pe],
+ negative_prompt_embeds=None,
+ prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+ negative_prompt_embeds_mask=None,
+ prompt_seq_lens=[[850]],
+ negative_prompt_seq_lens=None,
+ )
+ config.expand_conditioning_to_sample_batch(batch)
+ assert batch.prompt_embeds[0] is pe # untouched
+ assert batch.prompt_embeds[0].shape == (1, 850, 3584)
diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py b/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py
index b5d9d20ee..13b6f444a 100644
--- a/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py
+++ b/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py
@@ -148,7 +148,7 @@ def test_text_mlp_uses_single_rank_when_intermediate_size_is_not_tp_divisible(
assert isinstance(layer.mlp.down_proj, ReplicatedLinear)
-def test_explicit_attention_mask_is_limited_to_cached_generation(monkeypatch):
+def test_explicit_attention_mask_is_honored_without_a_cache(monkeypatch):
attention = Qwen2_5_VLAttention.__new__(Qwen2_5_VLAttention)
nn.Module.__init__(attention)
attention.q_proj = nn.Identity()
@@ -174,9 +174,19 @@ def test_explicit_attention_mask_is_limited_to_cached_generation(monkeypatch):
"position_ids": torch.zeros(3, 1, 2, dtype=torch.long),
}
+ # LongCat opts into masking the padded body on the cache-free path.
+ attention.honor_cache_free_padding_mask = True
attention(**kwargs, use_cache=False)
attention(**kwargs, use_cache=True)
+ assert attention.attn.masks[0] is explicit_mask
+ assert attention.attn.masks[1] is explicit_mask
+ # Every other pipeline keeps the original behavior: mask dropped when
+ # cache-free, honored only under cached generation.
+ attention.attn.masks.clear()
+ attention.honor_cache_free_padding_mask = False
+ attention(**kwargs, use_cache=False)
+ attention(**kwargs, use_cache=True)
assert attention.attn.masks[0] is None
assert attention.attn.masks[1] is explicit_mask