sgl-project/sglang · error · ValueError

Unexpected RoPE rank: {cos.ndim}

Error message

Unexpected RoPE rank: {cos.ndim}

What it means

_slice_rope accepts RoPE cos/sin of rank 3 (B, L, D) or rank 4 (B, H, L, D) and slices along the sequence axis. Any other rank (e.g. rank 2 plain (L, D)) raises this ValueError.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py:79

        patch_size,
        patch_size,
    )
    return (
        tokens.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
    )


def _slice_rope(
    rope: tuple[torch.Tensor, torch.Tensor], start: int, end: Optional[int] = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """Slice along token axis for either interleaved (rank-3) or split (rank-4)."""
    cos, sin = rope
    end_ = end if end is not None else cos.shape[-2 if cos.ndim == 4 else 1]
    if cos.ndim == 3:
        return cos[:, start:end_], sin[:, start:end_]
    if cos.ndim == 4:
        return cos[:, :, start:end_, :], sin[:, :, start:end_, :]
    raise ValueError(f"Unexpected RoPE rank: {cos.ndim}")


def _streaming_self_attention(
    attn: LTX2Attention,
    hidden_states: torch.Tensor,
    video_rotary_emb: tuple[torch.Tensor, torch.Tensor],
    n_context_tokens: int,
) -> torch.Tensor:
    """Streaming SLA: context attends to context only, current attends to context+current.

    Mirrors NVlabs `inference_sana_wm.py::_streaming_self_attention`.
    """
    seq_len = hidden_states.shape[1]
    if n_context_tokens <= 0 or n_context_tokens >= seq_len:
        return attn(hidden_states, context=None, pe=video_rotary_emb)

    ctx_rope = _slice_rope(video_rotary_emb, 0, n_context_tokens)
    out_ctx = attn(

View on GitHub (pinned to 0132848349)

Solutions

  1. Unsqueeze rank-2 tables to rank 3: cos[None], sin[None] (or broadcast to (B, L, D)) before passing
  2. Use the RoPE helper expected by the refiner to build video_rotary_emb
  3. Add an assert cos.ndim in (3, 4) at your call site

Example fix

# before
rope = (cos, sin)  # each (L, D)
# after
rope = (cos.unsqueeze(0), sin.unsqueeze(0))  # (1, L, D)
Defensive patterns

Strategy: type-guard

Validate before calling

assert cos.ndim in (3, 4) and sin.ndim == cos.ndim

Type guard

def valid_rope(rope) -> bool:
    cos, sin = rope
    return cos.ndim in (3, 4) and sin.ndim == cos.ndim and cos.shape[-2 if cos.ndim==4 else -2] == sin.shape[-2]

Prevention

When it happens

Trigger: Passing video_rotary_emb to the refiner's streaming self-attention where cos/sin are 2-D (L, D) — as produced by some RoPE helpers that omit batch/head dims — or a mistakenly unsqueezed 5-D tensor.

Common situations: Swapping in a different rotary embedding utility that returns unbatched tables; refiner fed RoPE computed for the base SANA-WM transformer with a different layout.

Related errors


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