sgl-project/sglang · error · ValueError

Z-Image text embeddings must have shape [seq, dim] or [batch

Error message

Z-Image text embeddings must have shape [seq, dim] or [batch, seq, dim]

What it means

Z-Image pipeline conditioning code accepts text embeddings only as a 2-D [seq, dim] tensor or a 3-D [batch, seq, dim] tensor. The splitter _split_text_embeds_for_dit raises this when embeds.ndim is neither 2 nor 3, i.e. the tensor coming from the text encoder has an unexpected rank (1-D, 4-D, etc.).

Source

Thrown at python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py:232

    def _split_text_embeds_for_dit(self, batch, *, negative: bool = False):
        """Return per-request text tensors, trimming padded batched embeddings."""
        embeds = batch.negative_prompt_embeds if negative else batch.prompt_embeds
        if embeds is None:
            return None

        if isinstance(embeds, (list, tuple)):
            if not embeds:
                return []
            embeds = embeds[0]

        if not torch.is_tensor(embeds):
            return embeds

        if embeds.ndim == 2:
            return [self._pad_text_embed_for_dit(embeds)]

        if embeds.ndim != 3:
            raise ValueError(
                "Z-Image text embeddings must have shape [seq, dim] or [batch, seq, dim]"
            )

        seq_lens = self.require_text_seq_lens(
            batch,
            0,
            negative=negative,
            expected_batch_size=int(embeds.shape[0]),
        )
        return [
            self._pad_text_embed_for_dit(embeds[idx, :seq_len].contiguous())
            for idx, seq_len in enumerate(seq_lens)
        ]

    def _caption_rope_length(self, prompt_embeds, batch, *, negative: bool = False):
        """Return the shared caption RoPE length for current text embeddings."""
        if torch.is_tensor(prompt_embeds):
            if prompt_embeds.ndim == 2:

View on GitHub (pinned to 0132848349)

Solutions

  1. Check the tensor shape fed to the pipeline: print text_embeds.shape before calling; ensure it is [seq, dim] (single prompt) or [batch, seq, dim] (batched)
  2. If your encoder returns [batch, dim] pooled output, re-export per-token sequence embeddings instead — pooled vectors are not supported here
  3. If you have [batch*seq, dim] flattened output, reshape to [batch, seq, dim] before passing it in
  4. Add an assertion/reshape in your encoder wrapper: embeds = embeds.reshape(-1, embeds.shape[-1]) for 2-D or embeds[:, None, :] style fix depending on your data

Example fix

# before
embeds = text_encoder(prompt)  # returns [batch, dim] pooled -> ValueError
pipe.get_pos_prompt_embeds(batch)

# after
embeds = text_encoder(prompt)          # [seq, dim] or [batch, seq, dim]
assert embeds.ndim in (2, 3), embeds.shape
pipe.get_pos_prompt_embeds(batch)
Defensive patterns

Strategy: validation

Validate before calling

embeds = get_text_embeds(...)
assert embeds.ndim in (2, 3), f"expected [seq, dim] or [batch, seq, dim], got {embeds.shape}"
if embeds.ndim == 1:
    embeds = embeds.unsqueeze(0)  # [dim] -> [1, dim] only if that is semantically a 1-token seq

Type guard

def is_valid_text_embeds(t) -> bool:
    return hasattr(t, "ndim") and t.ndim in (2, 3)

Prevention

When it happens

Trigger: Calling get_pos_prompt_embeds(batch) or get_neg_prompt_embeds(batch) on the Z-Image pipeline config when the text encoder returned embeddings with ndim not in (2, 3) — e.g. a pooled [batch, dim] vector, a 1-D [dim] tensor, or a 4-D tensor from a custom encoder wrapper.

Common situations: Swapping in a custom text encoder or encoder wrapper that squeezes/unsqueezes dims incorrectly; passing pooled prompt embeddings instead of sequence embeddings; upstream changes in the multimodal encoder output shape after a version bump.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/ee92e9312f85c763. Report an issue: GitHub.