sgl-project/sglang · error · RuntimeError

Cannot determine attention head counts for {type(inner).__na

Error message

Cannot determine attention head counts for {type(inner).__name__}

What it means

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.

Source

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

    def __init__(
        self, inner: nn.Module, layer_idx: int, window_size: int | None = None
    ):
        super().__init__()
        object.__setattr__(self, "_inner", inner)
        object.__setattr__(self, "_layer_idx", layer_idx)
        object.__setattr__(self, "_window_size", window_size)
        # Resolved once at patch time (weights are loaded before patching and
        # the inner module is never swapped afterwards), keeping the decode
        # hot path free of attribute scans and failing fast on a bad module.
        scale = get_attention_scale(inner)
        if scale is None:
            raise RuntimeError(
                f"Cannot determine attention scale for {type(inner).__name__}"
            )
        n_heads = get_num_heads(inner)
        n_kv_heads = get_num_kv_heads(inner)
        if n_heads is None or n_kv_heads is None:
            raise RuntimeError(
                f"Cannot determine attention head counts for {type(inner).__name__}"
            )
        object.__setattr__(self, "_scale", scale)
        object.__setattr__(self, "_n_heads", n_heads)
        object.__setattr__(self, "_n_kv_heads", n_kv_heads)
        # None for modules that expose head_dim only through a projection
        # shape; _batched_decode falls back to the runtime K shape.
        object.__setattr__(self, "_head_dim", get_head_dim(inner))
        object.__setattr__(self, "_has_q_norm", hasattr(inner, "q_norm"))
        object.__setattr__(self, "_has_k_norm", hasattr(inner, "k_norm"))
        # Only pass sinks when the module has them: the kwarg requires a
        # recent mlx and must not constrain models without sinks.
        sinks = getattr(inner, "sinks", None)
        object.__setattr__(self, "_sinks", sinks)
        object.__setattr__(
            self, "_sink_kwargs", {} if sinks is None else {"sinks": sinks}
        )

View on GitHub (pinned to 0132848349)

Solutions

1. Expose standard `n_heads` / `n_kv_heads` attributes on the inner module.
2. Teach get_num_heads/get_num_kv_heads the new attribute names in the MLX backend.
3. Pin mlx_lm to a compatible version or disable the MLX KV-cache wrapper for this model.

Example fix

# before
class MyAttention(nn.Module):
    def __init__(self):
        self.num_query_heads = 32  # unrecognized

# after
class MyAttention(nn.Module):
    def __init__(self):
        self.n_heads = 32
        self.n_kv_heads = 8

if get_num_heads(module) is None or get_num_kv_heads(module) is None:
    module.n_heads, module.n_kv_heads = config.n_heads, config.n_kv_heads

def has_head_counts(module) -> bool:
    return get_num_heads(module) is not None and get_num_kv_heads(module) is not None

When it happens

Trigger: Patching an attention module whose head-count attributes are absent or named unconventionally (custom MQA/GQA implementation, renamed attrs in a newer mlx_lm).

Common situations: New mlx_lm model releases with refactored attention classes; custom attention modules; models with head counts only inferable from config rather than the module.

Related errors

• Cannot determine attention scale for {type(inner).__name__}
• Unexpected q_proj output shape {q_proj_output.shape} for {ty
• MLX model has no supported attention layers
• num_heads must be divisible by num_epi_subtiles
• num_heads // num_epi_subtiles must be divisible by 4 (FMA un

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