sgl-project/sglang · error · RuntimeError

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

Error message

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

What it means

When patching an mlx_lm attention module, the wrapper resolves the attention softmax scale once via `get_attention_scale(inner)`. If the inner module exposes none of the recognized attributes (e.g. `scale`, `head_dim`), the helper returns None and the wrapper raises at patch time — deliberately, so the decode hot path never has to guess a scale.

Source

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

    the trailing window of the cached keys only, which is numerically
    identical to a rotating cache.  Both cache kinds keep KV in temporal
    order and report absolute offsets, so the same trailing-window slice
    works whether the cache holds full history or only the window.
    """

    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.

View on GitHub (pinned to 0132848349)

Solutions

1. Extend `get_attention_scale` to recognize the new module's scale attribute (or expose `scale`/`head_dim` on your custom module).
2. Pin mlx_lm to a supported version for your SGLang release.
3. Skip MLX KV-cache patching for unsupported architectures (fall back to default path).

Example fix

# before
class MyAttention(nn.Module):
    def __init__(self, d):
        self.softmax_scale = d ** -0.5  # unrecognized name

# after
class MyAttention(nn.Module):
    def __init__(self, d):
        self.scale = d ** -0.5  # recognized by get_attention_scale

from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import get_attention_scale
if get_attention_scale(module) is None:
    module.scale = module.head_dim ** -0.5  # or skip patching

def is_patchable_attention(module) -> bool:
    return get_attention_scale(module) is not None

When it happens

Trigger: Wrapping an attention module whose type doesn't expose a known scale attribute — a new/custom attention implementation, an upstream mlx_lm rename, or a model class not yet supported by the MLX backend.

Common situations: Loading a brand-new mlx_lm model architecture; upgrading mlx_lm so attention internals are renamed; using a custom attention subclass with non-standard attribute names.

Related errors

• Cannot determine attention head counts for {type(inner).__na
• 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/3871c357d26fa6af. Report an issue: GitHub.