sgl-project/sglang · error · ValueError

rope_pool_fused expects positions/slots to be 1-D

Error message

rope_pool_fused expects positions/slots to be 1-D

What it means

rope_pool_fused requires the per-token positions and pool slot indices to be 1-D tensors of length num_tokens. This error fires when either positions or slots has more than one dimension (or is scalars/2-D), because the kernel iterates tokens with a single flat index array.

Source

Thrown at python/sglang/kernels/aot/python/sgl_kernel/metal.py:66

    """Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.

    Args:
        q: Query tensor with shape `[num_tokens, num_qo_heads, head_dim]`.
        k: Key tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
        v: Value tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
        positions: int32 positions with shape `[num_tokens]`.
        slots: int32 KV-pool slots with shape `[num_tokens]`; values `< 0`
            skip the pool write for that token.
        k_pool: Existing K pool with shape `[pool_size, num_kv_heads, head_dim]`.
        v_pool: Existing V pool with shape `[pool_size, num_kv_heads, head_dim]`.

    Returns:
        `(q_rot, k_rot, k_pool_new, v_pool_new)`.
    """
    if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
        raise ValueError("rope_pool_fused expects q/k/v to be 3-D")
    if positions.ndim != 1 or slots.ndim != 1:
        raise ValueError("rope_pool_fused expects positions/slots to be 1-D")
    if k_pool.ndim != 3 or v_pool.ndim != 3:
        raise ValueError("rope_pool_fused expects pool tensors to be 3-D")
    q_shape = tuple(q.shape)
    k_shape = tuple(k.shape)
    v_shape = tuple(v.shape)
    positions_shape = tuple(positions.shape)
    slots_shape = tuple(slots.shape)
    k_pool_shape = tuple(k_pool.shape)
    v_pool_shape = tuple(v_pool.shape)

    if q_shape != (q_shape[0], num_qo_heads, head_dim):
        raise ValueError(
            "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
        )
    if k_shape != (q_shape[0], num_kv_heads, head_dim):
        raise ValueError(
            "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
        )

View on GitHub (pinned to 0132848349)

Solutions

1. Flatten positions/slots to 1-D: positions = positions.reshape(-1); slots = slots.reshape(-1)
2. If using scalars, wrap them: torch.tensor([pos], dtype=torch.int64)
3. Verify positions.ndim == 1 and slots.ndim == 1 before the call

Example fix

# before
metal.rope_pool_fused(q, k, v, positions, slots, ...)  # positions is [B, S]

# after
positions = positions.reshape(-1)
slots = slots.reshape(-1)
metal.rope_pool_fused(q, k, v, positions, slots, ...)

assert positions.ndim == 1 and slots.ndim == 1, (positions.shape, slots.shape)

def flat_index(t) -> bool:
    import torch
    return isinstance(t, torch.Tensor) and t.ndim == 1

When it happens

Trigger: Calling rope_pool_fused with positions or slots shaped [batch, seq] instead of [num_tokens], or passing scalar/0-D tensors for a single token instead of 1-D length-1 tensors.

Common situations: Porting code from attention backends that take batched position tensors; forgetting to flatten positions generated per batch; passing Python ints or 0-D tensors instead of torch tensors of shape [1].

Related errors

• rope_pool_fused expects q/k/v to be 3-D
• rope_pool_fused expects pool tensors to be 3-D
• q shape must be [num_tokens, num_qo_heads, head_dim], got {q
• k shape must be [num_tokens, num_kv_heads, head_dim], got {k
• v shape must match k shape, got {v.shape} vs {k.shape}

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