sgl-project/sglang · error · ValueError

NPU packed attention requires matching Q/K head dimensions

Error message

NPU packed attention requires matching Q/K head dimensions

What it means

q.shape[-1] (query head dim) must equal k.shape[-1] (key head dim); the NPU fused attention kernel does not support different query and key head dimensions (no MQA-style head-dim projection inside the kernel).

Source

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

        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]:
        raise ValueError(
            "NPU packed attention requires matching K/V token and head counts"
        )
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("NPU packed attention requires matching Q/K head dimensions")

    q_boundaries = _packed_boundaries(
        cu_seqlens_q, cu_seqlens_q_host, q.shape[0], "cu_seqlens_q"
    )
    k_boundaries = _packed_boundaries(
        cu_seqlens_k, cu_seqlens_k_host, k.shape[0], "cu_seqlens_k"
    )
    if len(q_boundaries) != len(k_boundaries):
        raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch")

    q_nonempty = [
        stop > start for start, stop in zip(q_boundaries[:-1], q_boundaries[1:])
    ]
    k_nonempty = [
        stop > start for start, stop in zip(k_boundaries[:-1], k_boundaries[1:])
    ]
    if q_nonempty != k_nonempty:
        raise NotImplementedError(

View on GitHub (pinned to 0132848349)

Solutions

1. Check the model config: num_head_dims for query and key must match when using this backend
2. If the model genuinely has unequal q/k head dims, project q to k's head dim before the call or use a different backend
3. Fix accidental head-dim mismatches from a wrong checkpoint/config (e.g. wrong rope or qk layer weights)

Example fix

# before: q head_dim=128, k head_dim=64
out = fused_infer_attention_varlen(q, k, v, cu_q, cu_k)
# after
q_proj = q[..., : k.shape[-1]]  # only if truncation is semantically valid
# better: use a backend that supports asymmetric head dims

assert q.shape[-1] == k.shape[-1], "Q/K head dims must match on NPU kernel"

When it happens

Trigger: Passing q with head_dim 128 and k with head_dim 64 (e.g. a model with asymmetric q/k head dims), which some architectures and some backends support but this kernel does not.

Common situations: Multi-modal/diffusion models with per-layer head-dim tweaks; loading only some projections at a different head dim due to a config mismatch; porting a model that relied on a GPU backend's support for unequal head dims.

Related errors

• {name} must end at the packed token count {total_tokens}, go
• NPU packed attention requires q, k, and v in [T, N, D] layou
• NPU packed attention requires matching K/V token and head co
• D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE
• head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.

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