{"record":{"id":"c895cffbed5974ae","repo":"sgl-project/sglang","slug":"unexpected-q-proj-output-shape-q-proj-output-shap","errorCode":null,"errorMessage":"Unexpected q_proj output shape {q_proj_output.shape} for {type(inner).__name__}","messagePattern":"Unexpected q_proj output shape (.+?) for (.+?)","errorType":"error_code","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py","lineNumber":263,"sourceCode":"        q_proj_output = inner.q_proj(x)\n        keys = inner.k_proj(x)\n        values = inner.v_proj(x)\n\n        head_dim = self._head_dim\n        if head_dim is None:\n            head_dim = keys.shape[-1] // n_kv_heads\n\n        q_width = n_heads * head_dim\n        gate = None\n        if q_proj_output.shape[-1] == q_width:\n            queries = q_proj_output.reshape(B, 1, n_heads, head_dim)\n        elif q_proj_output.shape[-1] == 2 * q_width:\n            queries, gate = mx.split(\n                q_proj_output.reshape(B, 1, n_heads, 2 * head_dim), 2, axis=-1\n            )\n            gate = gate.reshape(B, 1, q_width)\n        else:\n            raise RuntimeError(\n                f\"Unexpected q_proj output shape {q_proj_output.shape} for \"\n                f\"{type(inner).__name__}\"\n            )\n\n        keys = keys.reshape(B, 1, n_kv_heads, head_dim)\n        values = values.reshape(B, 1, n_kv_heads, head_dim)\n\n        if self._has_q_norm:\n            queries = inner.q_norm(queries)\n        if self._has_k_norm:\n            keys = inner.k_norm(keys)\n\n        queries = queries.transpose(0, 2, 1, 3)\n        keys = keys.transpose(0, 2, 1, 3)\n        values = values.transpose(0, 2, 1, 3)\n\n        # Vectorized RoPE with per-batch offsets (cached on the context).\n        offsets = ctx.offsets","sourceCodeStart":245,"sourceCodeEnd":281,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py#L245-L281","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Check that the inner module reports n_heads/n_kv_heads/head_dim consistent with q_proj's actual output width.","Extend the head_dim inference (the code path for 'modules that expose head_dim only through a projection') to handle the new projection layout.","If widths genuinely differ, add an explicit branch for that layout instead of relying on inference."],"exampleFix":"# before\n# head_dim inferred wrongly -> shape[-1] == n_heads * head_dim * 3 / 2, raises\nwrapper.decode(...)\n\n# after\nobject.__setattr__(wrapper, \"_head_dim\", q_proj_weight_shape // n_heads)\nwrapper.decode(...)","handlingStrategy":"validation","validationCode":"q_width = n_heads * head_dim\nassert q_proj_output.shape[-1] in (q_width, 2 * q_width), q_proj_output.shape\nkeys, values, kq = wrapper.decode(...)","typeGuard":null,"tryCatchPattern":"try:\n    wrapper.decode(...)\nexcept RuntimeError as e:\n    if \"Unexpected q_proj output shape\" in str(e):\n        raise ValueError(f\"head_dim mis-derived for {model}\") from e\n    raise","preventionTips":["Unit-test wrapper decode per model architecture with real weights.","Verify head_dim inference against q_proj weight shape at patch time."],"tags":["mlx","decode","attention","shape-mismatch"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}