src/agiws_neural_quant/ssm_patterns.py

"""ssm_patterns — exclude patterns for State-Space Model / linear-attention
layers (Mamba, KDA, RWKV, gated linear attention, etc.).

SSM layers contain parameters that must NOT be quantized through the standard
weight-quantization walker:

  - `A_log`, `dt_bias`, `D` — nn.Parameter (not nn.Module): the walker does not
    touch them automatically. Documented here for clarity.
  - `ShortConvolution` / `conv1d` — custom 1D-conv wrappers (not nn.Conv1d):
    not in SUPPORTED_MODULE_TYPES, walker skips them.
  - Linear projections feeding the SSM core (`q_proj`, `k_proj`, `b_proj`,
    `f_a_proj`, `f_b_proj`) — distorting these breaks the delta rule / recurrent
    dynamics. Should be excluded for low-bit formats.
  - Output projections (`v_proj`, `o_proj`, `g_proj`, `g_a_proj`, `g_b_proj`)
    — less critical: quantizing them is generally safe.

Usage:
    from agiws_neural_quant.ssm_patterns import get_ssm_exclude_patterns

    excl = get_ssm_exclude_patterns("conservative")
    quantize_model(model, format="int4", exclude_modules=excl)

Levels:
    conservative — quantize only o_proj (output) inside SSM modules.
    moderate     — quantize o_proj + v_proj + g_* (output path), keep q/k/b/f.
    aggressive   — quantize all Linear inside SSM except q_proj/k_proj
                   (preserve SSM input, quantize output + gating).
    full         — quantize every Linear (no SSM exclusions; for fp8 / high-bit
                   where distortion is minimal).
    subtree      — skip whole SSM modules entirely (for Mamba/RWKV whose module
                   name contains 'mamba'/'rwkv'/'kda'/'delta_attn'). NOT suitable
                   for Kimi-K3 where the SSM layer is named `self_attn` (shared
                   with MLA layers) — use level-based exclude there.

All patterns use fnmatch glob syntax and are matched against the full dotted
module path (e.g. "layers.5.self_attn.q_proj"). Bare names without a leading
`*.` are intentionally omitted — they would never match the dotted path and
would be dead patterns.
"""

from __future__ import annotations


# Full exclusion: skip the entire SSM subtree (any Linear inside).
# Matched against the SSM module name itself. Use for Mamba/RWKV where the
# module has a distinctive name. For Kimi-K3 (SSM layer named `self_attn`,
# shared with MLA), use level-based exclude instead.
SSM_SUBTREE_PATTERNS: tuple[str, ...] = (
    "mamba", "*_mamba", "*mamba*",
    "rwkv", "*_rwkv", "*rwkv*",
    "*delta_attention*", "*delta_attn*",
    "*linear_attention*", "*linear_attn*",
    "*kda*", "*_kda",
)


# Conservative: quantize only the final output projection inside SSM.
# All other SSM Linear (q/k/v/b/f/g inputs and gating) are excluded.
SSM_CONSERVATIVE_EXCLUDE: tuple[str, ...] = (
    "*.self_attn.q_proj",
    "*.self_attn.k_proj",
    "*.self_attn.v_proj",
    "*.self_attn.b_proj",
    "*.self_attn.f_a_proj",
    "*.self_attn.f_b_proj",
    "*.self_attn.g_a_proj",
    "*.self_attn.g_b_proj",
    "*.self_attn.in_proj",
    "*.self_attn.dt_proj",
)


# Moderate: quantize output path (o/v/g), keep SSM input + gating (q/k/b/f).
SSM_MODERATE_EXCLUDE: tuple[str, ...] = (
    "*.self_attn.q_proj",
    "*.self_attn.k_proj",
    "*.self_attn.b_proj",
    "*.self_attn.f_a_proj",
    "*.self_attn.f_b_proj",
    "*.self_attn.in_proj",
    "*.self_attn.dt_proj",
)


# Aggressive: quantize all SSM Linear except q/k (preserve SSM input only).
SSM_AGGRESSIVE_EXCLUDE: tuple[str, ...] = (
    "*.self_attn.q_proj",
    "*.self_attn.k_proj",
    "*.self_attn.in_proj",
)


_LEVELS: dict[str, tuple[str, ...] | None] = {
    "full": None,  # no exclusions
    "aggressive": SSM_AGGRESSIVE_EXCLUDE,
    "moderate": SSM_MODERATE_EXCLUDE,
    "conservative": SSM_CONSERVATIVE_EXCLUDE,
    "subtree": SSM_SUBTREE_PATTERNS,  # skip whole SSM modules
}


def get_ssm_exclude_patterns(level: str = "conservative") -> list[str]:
    """Return glob patterns to pass to quantize_model(exclude_modules=...).

    Args:
        level: one of 'conservative', 'moderate', 'aggressive', 'full', 'subtree'.
            - conservative (default): keep SSM core + gating, quantize o_proj only.
            - moderate: quantize output path (o/v/g), keep SSM input + gating.
            - aggressive: quantize all SSM Linear except q/k (SSM input).
            - full: no SSM exclusions (empty list).
            - subtree: skip whole SSM modules entirely (most conservative,
                       for Mamba/RWKV; NOT for Kimi-K3 `self_attn`).

    Returns:
        list of fnmatch glob patterns; empty for 'full'.
    """
    if level not in _LEVELS:
        raise ValueError(
            f"Unknown SSM level {level!r}. Available: {sorted(_LEVELS)}"
        )
    pats = _LEVELS[level]
    return list(pats) if pats is not None else []


def get_ssm_subtree_patterns() -> list[str]:
    """Patterns to skip whole SSM modules (do not recurse into them at all).

    Use for Mamba/RWKV. For Kimi-K3 (SSM layer named `self_attn`, shared with
    MLA), prefer the level-based exclude instead.
    """
    return list(SSM_SUBTREE_PATTERNS)


__all__ = [
    "SSM_SUBTREE_PATTERNS",
    "SSM_CONSERVATIVE_EXCLUDE",
    "SSM_MODERATE_EXCLUDE",
    "SSM_AGGRESSIVE_EXCLUDE",
    "get_ssm_exclude_patterns",
    "get_ssm_subtree_patterns",
]