vllm.model_executor.models.muse_glimmer ¶

def _muse_glimmer_query_prescale(config) -> float:
    """Post-QK-norm query pre-scale (``scale_query_by``), normalized across the
    two config schemas so the net query scaling matches the native reference.

    HF native modeling computes ``scale_query_by = qk_scale_factor / sqrt(head_dim)``
    where the NATIVE ``qk_scale_factor`` is the raw ``params.json`` value
    (~43.784). The modular HF ``text_config`` PRE-FOLDS the ``1/sqrt(head_dim)``
    factor and ships ``qk_scale_factor = 43.784 / sqrt(128) = 3.87`` already,
    expecting it applied directly. Both must yield the SAME ``scale_query_by``
    (~3.87), then softmax uses ``scaling = head_dim**-0.5``.

    Precedence:
      1. explicit ``scale_query_by`` (already the final factor) -> use as-is.
      2. else derive from ``qk_scale_factor``:
         - if it is already the folded value (``~= qk_scale_factor/sqrt(hd)`` is
           NOT what we want; detect the native form and divide) — we decide by
           magnitude: the native raw value is ``folded * sqrt(head_dim)``. If
           ``qk_scale_factor`` is close to ``folded_expected * sqrt(hd)`` we treat
           it as native and divide; otherwise it is already folded, use directly.
    """
    head_dim = config.head_dim
    sqrt_hd = head_dim**0.5

    explicit = getattr(config, "scale_query_by", None)
    if explicit is not None:
        return float(explicit)

    qk_scale = getattr(config, "qk_scale_factor", None)
    if qk_scale is None:
        # No scale info at all: fall back to the plain 1/sqrt(head_dim) identity
        # (net query scaling then just the softmax scaling); should not happen
        # for real MuseGlimmer checkpoints, which always carry qk_scale_factor.
        return 1.0

    qk_scale = float(qk_scale)
    # Disambiguate native (raw, ~43.78) vs modular (folded, ~3.87). The native
    # form, when divided by sqrt(head_dim), yields the folded target; the folded
    # form is already the target. Native values are ~sqrt(head_dim)x larger than
    # folded. Use a threshold at sqrt(head_dim) (with margin): if qk_scale is
    # comparable to or larger than sqrt(head_dim), it is the native raw value and
    # must be divided; otherwise it is already folded and used directly.
    #   head_dim=128 -> sqrt=11.31; native 43.78 > 11.31 (divide -> 3.87),
    #   folded 3.87 < 11.31 (use as-is).
    if qk_scale >= sqrt_hd:
        return qk_scale / sqrt_hd
    return qk_scale

def _muse_glimmer_use_attn_output_gate(config) -> bool:
    """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS
    applies it; the modular HF ``text_config`` omits ``use_attn_output_gate``
    (reads as ``None``). Missing/None -> True; only explicit ``False`` disables."""
    val = getattr(config, "use_attn_output_gate", None)
    return True if val is None else bool(val)

def _muse_glimmer_use_qk_norm(config) -> bool:
    """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF
    ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``).
    Treat a missing/None flag as True — only an explicit ``False`` disables it."""
    val = getattr(config, "use_qk_norm", None)
    return True if val is None else bool(val)

def _text_config(config):
    """MuseGlimmer checkpoints may nest the text config under ``text_config``
    (multimodal ``MuseGlimmerConfig``) or expose it directly
    (``MuseGlimmerTextConfig``)."""
    return getattr(config, "text_config", config)