{"record":{"id":"353424b27737f8e8","repo":"sgl-project/sglang","slug":"cannot-determine-attention-head-counts-for-type-i","errorCode":null,"errorMessage":"Cannot determine attention head counts for {type(inner).__name__}","messagePattern":"Cannot determine attention head counts for (.+?)","errorType":"error_code","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py","lineNumber":213,"sourceCode":"    def __init__(\n        self, inner: nn.Module, layer_idx: int, window_size: int | None = None\n    ):\n        super().__init__()\n        object.__setattr__(self, \"_inner\", inner)\n        object.__setattr__(self, \"_layer_idx\", layer_idx)\n        object.__setattr__(self, \"_window_size\", window_size)\n        # Resolved once at patch time (weights are loaded before patching and\n        # the inner module is never swapped afterwards), keeping the decode\n        # hot path free of attribute scans and failing fast on a bad module.\n        scale = get_attention_scale(inner)\n        if scale is None:\n            raise RuntimeError(\n                f\"Cannot determine attention scale for {type(inner).__name__}\"\n            )\n        n_heads = get_num_heads(inner)\n        n_kv_heads = get_num_kv_heads(inner)\n        if n_heads is None or n_kv_heads is None:\n            raise RuntimeError(\n                f\"Cannot determine attention head counts for {type(inner).__name__}\"\n            )\n        object.__setattr__(self, \"_scale\", scale)\n        object.__setattr__(self, \"_n_heads\", n_heads)\n        object.__setattr__(self, \"_n_kv_heads\", n_kv_heads)\n        # None for modules that expose head_dim only through a projection\n        # shape; _batched_decode falls back to the runtime K shape.\n        object.__setattr__(self, \"_head_dim\", get_head_dim(inner))\n        object.__setattr__(self, \"_has_q_norm\", hasattr(inner, \"q_norm\"))\n        object.__setattr__(self, \"_has_k_norm\", hasattr(inner, \"k_norm\"))\n        # Only pass sinks when the module has them: the kwarg requires a\n        # recent mlx and must not constrain models without sinks.\n        sinks = getattr(inner, \"sinks\", None)\n        object.__setattr__(self, \"_sinks\", sinks)\n        object.__setattr__(\n            self, \"_sink_kwargs\", {} if sinks is None else {\"sinks\": sinks}\n        )\n","sourceCodeStart":195,"sourceCodeEnd":231,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py#L195-L231","documentation":"Companion to the scale check: the wrapper also resolves head counts via `get_num_heads(inner)` / `get_num_kv_heads(inner)`. If either returns None (module exposes no recognized n_heads/n_kv_heads attributes), it raises at patch time so grouping math can't proceed with unknown dimensions.","triggerScenarios":"Patching an attention module whose head-count attributes are absent or named unconventionally (custom MQA/GQA implementation, renamed attrs in a newer mlx_lm).","commonSituations":"New mlx_lm model releases with refactored attention classes; custom attention modules; models with head counts only inferable from config rather than the module.","solutions":["Expose standard `n_heads` / `n_kv_heads` attributes on the inner module.","Teach get_num_heads/get_num_kv_heads the new attribute names in the MLX backend.","Pin mlx_lm to a compatible version or disable the MLX KV-cache wrapper for this model."],"exampleFix":"# before\nclass MyAttention(nn.Module):\n    def __init__(self):\n        self.num_query_heads = 32  # unrecognized\n\n# after\nclass MyAttention(nn.Module):\n    def __init__(self):\n        self.n_heads = 32\n        self.n_kv_heads = 8","handlingStrategy":"validation","validationCode":"if get_num_heads(module) is None or get_num_kv_heads(module) is None:\n    module.n_heads, module.n_kv_heads = config.n_heads, config.n_kv_heads","typeGuard":"def has_head_counts(module) -> bool:\n    return get_num_heads(module) is not None and get_num_kv_heads(module) is not None","tryCatchPattern":null,"preventionTips":["Ensure custom attention modules expose n_heads/n_kv_heads attributes.","Run a patch-time compatibility check per model family."],"tags":["mlx","attention","model-compatibility","patching"],"backgroundTag":"unsupported-model-architecture","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}