sgl-project/sglang · error · ValueError

pad_masks and att_masks must be [batch, seq]

Error message

pad_masks and att_masks must be [batch, seq]

What it means

make_att_2d_masks builds a 2-D [batch, seq, seq] attention mask from a padding mask and an attention mask via cumulative-sum logic that only works when both inputs are 2-D [batch, seq]. The guard rejects any other rank. It is used by encode_prefix and denoise_step, so the offending tensors usually come from the caller's prefix/action preparation code.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py:854

    fraction = torch.linspace(
        0.0,
        1.0,
        dimension // 2,
        dtype=torch.float64,
        device=time.device,
    )
    period = min_period * (max_period / min_period) ** fraction
    scaling = 1.0 / period * 2 * math.pi
    sin_input = scaling[None, :] * time[:, None].to(torch.float64)
    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)


def make_att_2d_masks(
    pad_masks: torch.Tensor,
    att_masks: torch.Tensor,
) -> torch.Tensor:
    if att_masks.ndim != 2 or pad_masks.ndim != 2:
        raise ValueError("pad_masks and att_masks must be [batch, seq]")
    cumsum = torch.cumsum(att_masks, dim=1)
    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
    return att_2d_masks & pad_2d_masks


def trim_trailing_padding_tokens(
    tokens: torch.Tensor,
    token_masks: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    token_len = int(token_masks.sum(dim=1).max().item())
    if token_len <= 0 or token_len >= tokens.shape[1]:
        return tokens, token_masks
    return tokens[:, :token_len], token_masks[:, :token_len]


def prepare_optional_full_attention_mask(
    att_2d_masks: torch.Tensor,

View on GitHub (pinned to 0132848349)

Solutions

1. Squeeze extra dims: pad_masks = pad_masks.squeeze(-1), att_masks = att_masks.squeeze(1)
2. Ensure masks are boolean/integer of shape [batch, seq] before calling encode_prefix/denoise_step
3. Add asserts pad_masks.ndim == 2 and att_masks.ndim == 2 in your mask-prep helper

Example fix

# before
att_2d = make_att_2d_masks(pad_masks[:, :, 0], att_masks)  # 1-D pad mask

# after
att_2d = make_att_2d_masks(pad_masks, att_masks)  # both [batch, seq]

for name, m in (("pad_masks", pad_masks), ("att_masks", att_masks)):
    assert m.ndim == 2, f"{name} must be [batch, seq], got {tuple(m.shape)}"

def is_2d_mask(m: torch.Tensor) -> bool:
    return isinstance(m, torch.Tensor) and m.ndim == 2

try:
    att2d = make_att_2d_masks(pad_masks, att_masks)
except ValueError:
    pad_masks = pad_masks.reshape(pad_masks.shape[0], -1)
    att_masks = att_masks.reshape(att_masks.shape[0], -1)
    att2d = make_att_2d_masks(pad_masks, att_masks)

When it happens

Trigger: Calling make_att_2d_masks with pad_masks or att_masks of rank != 2 — e.g. a 3-D mask [batch, heads, seq], a 1-D mask [seq], or a mask that still has a trailing singleton dim like [batch, seq, 1].

Common situations: Passing masks straight from a tokenizer that returns [batch, 1, seq, seq]; per-head attention masks from another backend; forgetting .squeeze(-1) after an unsqueeze elsewhere.

Related errors

• time must have shape [batch]
• hd256 forward varlen expects q rank 3 or 5, got rank {q_rank
• hd256 forward non-varlen expects q rank 4 or 5, got rank {q_
• unsupported input for causal Conv3D cat/pad CUDA
• unsupported input for usp_merge_heads CUDA

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