sgl-project/sglang · error · ValueError

NPU packed attention requires q, k, and v in [T, N, D] layou

Error message

NPU packed attention requires q, k, and v in [T, N, D] layout; invalid tensors: {', '.join(invalid_layouts)}

What it means

fused_infer_attention_varlen (NPU packed attention) requires q, k, v as rank-3 tensors in [T, N, D] (tokens, heads, head_dim) layout. Any of the three with ndim != 3 is rejected, and the message names the offending tensors.

Source

Thrown at python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py:68

    return boundaries


def fused_infer_attention_varlen(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    cu_seqlens_q: torch.Tensor,
    cu_seqlens_k: torch.Tensor,
    *,
    cu_seqlens_q_host: Sequence[int] | None = None,
    cu_seqlens_k_host: Sequence[int] | None = None,
    softmax_scale: float | None = None,
    return_softmax_lse: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    tensors = {"q": q, "k": k, "v": v}
    invalid_layouts = [name for name, tensor in tensors.items() if tensor.ndim != 3]
    if invalid_layouts:
        raise ValueError(
            "NPU packed attention requires q, k, and v in [T, N, D] layout; "
            f"invalid tensors: {', '.join(invalid_layouts)}"
        )
    invalid_devices = [
        name
        for name, tensor in tensors.items()
        if tensor.device.type != "npu" or tensor.device != q.device
    ]
    if invalid_devices:
        raise ValueError(
            "NPU packed attention requires q, k, and v on the same NPU; "
            f"invalid tensors: {', '.join(invalid_devices)}"
        )
    if not (q.dtype == k.dtype == v.dtype):
        raise ValueError(
            "NPU packed attention requires q, k, and v with the same dtype"
        )
    if k.shape[:2] != v.shape[:2]:

View on GitHub (pinned to 0132848349)

Solutions

  1. Reshape to packed layout: q = q.reshape(-1, num_heads, head_dim) (same for k, v) with cu_seqlens covering the packed T dimension
  2. If your tensors are [B, S, N, D], flatten batch and sequence dims: q.reshape(q.shape[0]*q.shape[1], q.shape[2], q.shape[3])
  3. Check the error's invalid list to see which tensor(s) came in wrong and fix those call sites

Example fix

# before: [B, S, N, D]
out = fused_infer_attention_varlen(q, k, v, cu_q, cu_k)
# after: packed [T, N, D]
q3 = q.reshape(-1, q.shape[-2], q.shape[-1])
k3 = k.reshape(-1, k.shape[-2], k.shape[-1])
v3 = v.reshape(-1, v.shape[-2], v.shape[-1])
out = fused_infer_attention_varlen(q3, k3, v3, cu_q, cu_k)
Defensive patterns

Strategy: type-guard

Validate before calling

q3 = q.reshape(-1, num_heads, head_dim)
k3 = k.reshape(-1, num_kv_heads, head_dim)
v3 = v.reshape(-1, num_kv_heads, head_dim)

Type guard

def is_tnd(t: torch.Tensor) -> bool:
    return t.ndim == 3

Prevention

When it happens

Trigger: Passing [B, S, N, D] batched tensors, [B, N, S, D] transposed attention-layout tensors, or 2D/4D projections; invalid_layouts lists which of q/k/v are wrong.

Common situations: Porting code from a backend that takes [B, S, N, D] (like the base AITerImpl.forward or flash_attn batched API); forgetting to flatten a batched tensor to packed tokens before the varlen call.

Related errors


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