sgl-project/sglang · error · RuntimeError

Unexpected q_proj output shape {q_proj_output.shape} for {ty

Error message

Unexpected q_proj output shape {q_proj_output.shape} for {type(inner).__name__}

What it means

During batched decode the wrapper inspects q_proj's output width to detect plain vs gated (GQA-with-gate, e.g. Qwen-style) queries: it must equal either `q_width` or `2*q_width`. Any other last-dimension size means head counts/head_dim were mis-derived or the projection is non-standard, and reshaping would silently corrupt the tensor, so it raises.

Source

Thrown at python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py:263

        q_proj_output = inner.q_proj(x)
        keys = inner.k_proj(x)
        values = inner.v_proj(x)

        head_dim = self._head_dim
        if head_dim is None:
            head_dim = keys.shape[-1] // n_kv_heads

        q_width = n_heads * head_dim
        gate = None
        if q_proj_output.shape[-1] == q_width:
            queries = q_proj_output.reshape(B, 1, n_heads, head_dim)
        elif q_proj_output.shape[-1] == 2 * q_width:
            queries, gate = mx.split(
                q_proj_output.reshape(B, 1, n_heads, 2 * head_dim), 2, axis=-1
            )
            gate = gate.reshape(B, 1, q_width)
        else:
            raise RuntimeError(
                f"Unexpected q_proj output shape {q_proj_output.shape} for "
                f"{type(inner).__name__}"
            )

        keys = keys.reshape(B, 1, n_kv_heads, head_dim)
        values = values.reshape(B, 1, n_kv_heads, head_dim)

        if self._has_q_norm:
            queries = inner.q_norm(queries)
        if self._has_k_norm:
            keys = inner.k_norm(keys)

        queries = queries.transpose(0, 2, 1, 3)
        keys = keys.transpose(0, 2, 1, 3)
        values = values.transpose(0, 2, 1, 3)

        # Vectorized RoPE with per-batch offsets (cached on the context).
        offsets = ctx.offsets

View on GitHub (pinned to 0132848349)

Solutions

1. Check that the inner module reports n_heads/n_kv_heads/head_dim consistent with q_proj's actual output width.
2. Extend the head_dim inference (the code path for 'modules that expose head_dim only through a projection') to handle the new projection layout.
3. If widths genuinely differ, add an explicit branch for that layout instead of relying on inference.

Example fix

# before
# head_dim inferred wrongly -> shape[-1] == n_heads * head_dim * 3 / 2, raises
wrapper.decode(...)

# after
object.__setattr__(wrapper, "_head_dim", q_proj_weight_shape // n_heads)
wrapper.decode(...)

q_width = n_heads * head_dim
assert q_proj_output.shape[-1] in (q_width, 2 * q_width), q_proj_output.shape
keys, values, kq = wrapper.decode(...)

try:
    wrapper.decode(...)
except RuntimeError as e:
    if "Unexpected q_proj output shape" in str(e):
        raise ValueError(f"head_dim mis-derived for {model}") from e
    raise

When it happens

Trigger: Calling decode (or the wrappers/tests that funnel into `_batched_decode`) when `n_heads * head_dim != q_proj_output.shape[-1]` and not 2x either — mismatched head_dim detection for modules that only expose head_dim via a projection, or a gated-attention model with unusual width math.

Common situations: Supporting a new model whose head_dim must be inferred from projection weight shapes; partial-attribute modules where n_heads is right but head_dim is wrong; models with fused qkv projections producing unexpected widths.

Related errors

• {tensor_name}{context_clause} with shape {tensor.shape} cann
• q, k, and v must have the same 3D shape
• Cannot determine attention scale for {type(inner).__name__}
• Cannot determine attention head counts for {type(inner).__na
• num_heads must be divisible by num_epi_subtiles

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