minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: align cosmos3 text packing with official pipeline (#26950)

Browse Source
This commit is contained in:
Mick
2026-06-02 02:07:17 +08:00
committed by GitHub
parent 86afa21ca7
commit 9a8ab2d22b
3 changed files with 38 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
+8 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
+29 -5
View File
@@ -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
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -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}/"