From 9a8ab2d22b4f85f840701e8ed717f663e1c4662b Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 2 Jun 2026 02:07:17 +0800 Subject: [PATCH] [diffusion] fix: align cosmos3 text packing with official pipeline (#26950) --- .../runtime/models/dits/cosmos3video.py | 11 ++++-- .../stages/model_specific_stages/cosmos3.py | 34 ++++++++++++++++--- .../sglang/multimodal_gen/test/test_utils.py | 2 +- 3 files changed, 38 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index d0a65bb97..46bdd8318 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -1051,6 +1051,7 @@ class Cosmos3OmniTransformer(CachableDiT): fps: float | None = None, cache_key: str = "default", noisy_frame_mask: torch.Tensor | None = None, + max_text_seq_len: int | None = None, **kwargs, ) -> torch.Tensor: """Forward pass for denoising. @@ -1069,6 +1070,8 @@ class Cosmos3OmniTransformer(CachableDiT): noisy frames (timestep embedding applied) and 0 marks conditioned frames (clean context, embedding skipped). ``None`` means every frame is noisy (T2V / T2I). + max_text_seq_len: Real text length already computed during + tokenization. When omitted it is derived from ``text_mask``. Returns: [B, C, T, H, W] velocity prediction @@ -1078,7 +1081,11 @@ class Cosmos3OmniTransformer(CachableDiT): batch_size, C, T, H, W = hidden_states.shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) - max_real_len = int(text_mask.sum(dim=1).max().item()) + if max_text_seq_len is None: + max_text_seq_len = int(text_mask.sum(dim=1).max().item()) + if max_text_seq_len < text_ids.shape[1]: + text_ids = text_ids[:, :max_text_seq_len] + text_mask = text_mask[:, :max_text_seq_len] # Check if sequence parallelism is enabled sequence_shard_enabled = self.sp_size > 1 @@ -1177,8 +1184,6 @@ class Cosmos3OmniTransformer(CachableDiT): residual: torch.Tensor | None = None for i, layer in enumerate(self.gen_layers): k_und, v_und = cached_kv_for_key[i] - k_und = k_und[:, :max_real_len] - v_und = v_und[:, :max_real_len] hidden_gen, residual = layer( hidden_gen, k_und, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 8a53b38f7..2082eaf96 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -126,10 +126,10 @@ class Cosmos3TokenizationStage(PipelineStage): device: torch.device, use_system_prompt: bool = False, system_prompt: str | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, int]: """Tokenize a prompt using Qwen2 chat template. - Returns (input_ids, attention_mask) as [1, S] tensors. + Returns (input_ids, attention_mask, seq_len) as [1, S] tensors. """ conversations = [] if use_system_prompt: @@ -180,7 +180,7 @@ class Cosmos3TokenizationStage(PipelineStage): input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=device) - return input_ids, attention_mask + return input_ids, attention_mask, seq_len def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Tokenize prompt and negative prompt.""" @@ -215,22 +215,30 @@ class Cosmos3TokenizationStage(PipelineStage): self.log_info(f"Prompt with duration: '{prompt}'") # Tokenize prompts - cond_ids, cond_mask = self._tokenize_prompt( + cond_ids, cond_mask, cond_seq_len = self._tokenize_prompt( prompt, max_sequence_length, device, use_system_prompt, system_prompt ) - uncond_ids, uncond_mask = self._tokenize_prompt( + uncond_ids, uncond_mask, uncond_seq_len = self._tokenize_prompt( negative_prompt, max_sequence_length, device, use_system_prompt, system_prompt, ) + # official Cosmos3 consumes packed text; keep a shared length for CFG batching + shared_seq_len = max(cond_seq_len, uncond_seq_len) + cond_ids = cond_ids[:, :shared_seq_len] + cond_mask = cond_mask[:, :shared_seq_len] + uncond_ids = uncond_ids[:, :shared_seq_len] + uncond_mask = uncond_mask[:, :shared_seq_len] # Store in batch.extra for denoising stage batch.extra["cond_text_ids"] = cond_ids batch.extra["cond_text_mask"] = cond_mask batch.extra["uncond_text_ids"] = uncond_ids batch.extra["uncond_text_mask"] = uncond_mask + batch.extra["cond_text_seq_len"] = cond_seq_len + batch.extra["uncond_text_seq_len"] = uncond_seq_len batch.extra["fps"] = fps # Mark as processed (even though we don't use standard embeddings) @@ -490,6 +498,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps: float, cache_key: str = "default", noisy_frame_mask: torch.Tensor | None = None, + max_text_seq_len: int | None = None, ) -> torch.Tensor: """Run transformer forward pass. @@ -517,6 +526,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key=cache_key, noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=max_text_seq_len, ) def _manage_device_placement(self, server_args: ServerArgs): @@ -629,6 +639,8 @@ class Cosmos3DenoisingStage(PipelineStage): guidance_scale=effective_scale, cfg_rank=cfg_rank, noisy_frame_mask=velocity_mask, + cond_text_seq_len=batch.extra["cond_text_seq_len"], + uncond_text_seq_len=batch.extra["uncond_text_seq_len"], ) elif effective_scale == 1.0: noise_pred = self._run_transformer( @@ -640,6 +652,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key="cond", noisy_frame_mask=velocity_mask, + max_text_seq_len=batch.extra["cond_text_seq_len"], ) else: noise_pred = self._predict_noise_cfg_batched( @@ -653,6 +666,10 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, guidance_scale=effective_scale, noisy_frame_mask=velocity_mask, + max_text_seq_len=max( + batch.extra["cond_text_seq_len"], + batch.extra["uncond_text_seq_len"], + ), ) else: noise_pred = self._run_transformer( @@ -664,6 +681,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key="cond", noisy_frame_mask=velocity_mask, + max_text_seq_len=batch.extra["cond_text_seq_len"], ) # I2V: zero-velocity at conditioned frames so the scheduler keeps @@ -698,6 +716,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps: float, guidance_scale: float, noisy_frame_mask: torch.Tensor | None = None, + max_text_seq_len: int | None = None, ) -> torch.Tensor: """Run CFG by stacking both branches into a batch_size=2 forward. @@ -724,6 +743,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key="cfg_batched", noisy_frame_mask=mask_batched, + max_text_seq_len=max_text_seq_len, ) noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0) @@ -745,6 +765,8 @@ class Cosmos3DenoisingStage(PipelineStage): guidance_scale: float, cfg_rank: int, noisy_frame_mask: torch.Tensor | None = None, + cond_text_seq_len: int | None = None, + uncond_text_seq_len: int | None = None, ) -> torch.Tensor: """Run CFG with one branch per CFG rank, combined by all-reduce. @@ -764,6 +786,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key="cond", noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=cond_text_seq_len, ) partial = guidance_scale * noise_pred else: @@ -776,6 +799,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, cache_key="uncond", noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=uncond_text_seq_len, ) partial = (1.0 - guidance_scale) * noise_pred diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 679c2d7bd..51ead5afd 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SGL_TEST_FILES_CI_DATA_REVISION = "10f3826199ae524b3af5026a57c8f817d207b2e5" +SGL_TEST_FILES_CI_DATA_REVISION = "20874fb9018d082c613a18ba92ab4f32479d3a32" SGL_TEST_FILES_CONSISTENCY_GT_ROOT = ( "https://raw.githubusercontent.com/" f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"