sgl-project/sglang · error · ValueError

Z-Image caption tensor must have rank 2 or 3

Error message

Z-Image caption tensor must have rank 2 or 3

What it means

The Z-Image model padder must read the caption sequence length from a prompt tensor: rank 2 ([seq, dim]) yields shape[0], rank 3 ([batch, seq, dim]) yields shape[1]. Any other rank cannot be unambiguously padded, so _caption_seq_len (used by pad_zimage_prompt_kwargs) raises.

Source

Thrown at python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/zimage.py:41

def _first_caption_tensor(encoder_hidden_states: Any) -> torch.Tensor | None:
    tensor = bcg_utils.first_tensor(encoder_hidden_states)
    if not torch.is_tensor(tensor):
        return None
    if tensor.dim() == 2:
        return tensor
    if tensor.dim() == 3:
        return tensor[0]
    return None


def _caption_seq_len(tensor: torch.Tensor) -> int:
    if tensor.dim() == 2:
        return int(tensor.shape[0])
    if tensor.dim() == 3:
        return int(tensor.shape[1])
    raise ValueError("Z-Image caption tensor must have rank 2 or 3")


def _pad_caption(obj: Any, *, target: int) -> Any:
    if torch.is_tensor(obj):
        if obj.dim() == 2:
            return bcg_utils.pad_tensor_dim(obj, 0, target)
        if obj.dim() == 3:
            return bcg_utils.pad_tensor_dim(obj, 1, target)
        return obj
    if isinstance(obj, list):
        return [_pad_caption(item, target=target) for item in obj]
    if isinstance(obj, tuple):
        return tuple(_pad_caption(item, target=target) for item in obj)
    return obj


def _unwrap_model(current_model: Any) -> Any:
    for attr in ("module", "_orig_mod"):

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure caption embeddings are passed as [seq, hidden] or [batch, seq, hidden] before padding.
2. If you have 1D token ids, embed/expand them to rank 2 first (run the text encoder or add a batch dim).
3. If the tensor should not be treated as a caption, exclude it from the kwargs the Z-Image padder iterates.

Example fix

// before
caption = input_ids  # rank 1: [seq]
kwargs = pad_zimage_prompt_kwargs({"caption": caption}, target=128)
// after
caption = text_encoder(input_ids)  # rank 2: [seq, hidden]
kwargs = pad_zimage_prompt_kwargs({"caption": caption}, target=128)

def check_caption(t: torch.Tensor) -> None:
    assert t.dim() in (2, 3), f"caption must be rank 2/3, got rank {t.dim()}"

check_caption(caption)
out = pad_zimage_prompt_kwargs(prompt_kwargs, target=target)

def is_valid_caption(tensor: object) -> bool:
    return isinstance(tensor, torch.Tensor) and tensor.dim() in (2, 3)

try:
    padded = pad_zimage_prompt_kwargs(kwargs, target=t)
except ValueError as e:
    if "rank 2 or 3" in str(e):
        kwargs["caption"] = embed_to_2d(kwargs["caption"])
        padded = pad_zimage_prompt_kwargs(kwargs, target=t)
    else:
        raise

When it happens

Trigger: Calling pad_zimage_prompt_kwargs with a Z-Image caption tensor whose .dim() is neither 2 nor 3 — e.g. a flat 1D token-id tensor or a 4D input.

Common situations: Upstream prompt-encoding changes passing raw token ids instead of embedded 2D/3D captions; a new multimodal input key leaking into the caption padding path; batch-handling changes in the multimodal runtime altering tensor rank.

Related errors

• Z-Image text embeddings must have shape [seq, dim] or [batch
• Unable to infer Z-Image caption length for rotary embeddings
• f"Expected camera embedding shape [B, C, F, H, W], got {tupl
• f"Input shape {tuple(x.shape)} must be divisible by patch_si
• Unsupported dims: {self.dims}

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