registry

Registry indexing the per-model ModelSpec by HF model type.

Model modules register their spec at import time (see models/); one spec per model type. Lookups return None (or an empty list) when nothing matches, so callers can fail loudly or fall back per their own policy.

Matching is by model-type and class-name strings only, so this package stays dependency-free (any nn.Module — or any object — can be passed to the lookups without importing torch here).

Functions

get_spec	Return the spec registered for model_type, or None.
get_specs	Return all registered specs, in registration order (aggregators, tests).
hf_model_type	Return the root HF model type (model.config.model_type), or None.
iter_gate_up_pairs	Yield the distinct (gate, up) projection-name pairs across all MoE variants.
iter_pqs_fuse_rules	Yield every (module_class_substrings, fuse_into, fuse_from) AWQ fusion rule.
match_moe_block	Return the MoE layout variant for module, resolved by model type.
register	Register a model spec and return it.
weight_plus_one_norm_names	All norm class names whose stored weight is w - 1, across all specs.

get_spec(model_type)

Return the spec registered for model_type, or None.

Parameters:

model_type (str)

Return type:

ModelSpec | None

get_specs()

Return all registered specs, in registration order (aggregators, tests).

Return type:

list[ModelSpec]

hf_model_type(model)

Return the root HF model type (model.config.model_type), or None.

Accepts a model or a config object (duck-typed, no transformers import): a model contributes via its config attribute, a config via its own model_type. This is the key for get_spec / match_moe_block.

Return type:

str | None

iter_gate_up_pairs()

Yield the distinct (gate, up) projection-name pairs across all MoE variants.

GLOBAL-VOCABULARY semantics: consumers (currently only calibration sibling grouping in quantization/model_calib.py, which also walks dense MLPs that no MoE variant can match) try every pair opportunistically on every module, getattr-guarded. Adding a pair to any spec therefore changes behavior for ALL models whose modules happen to carry those attribute names — prefer per-module resolution (match_moe_block(module).gate_up_pair) wherever the module is an identifiable MoE block. The dense-MLP case moves to a fusion-group topic section in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5).

Return type:

Iterator[tuple[str, str]]

iter_pqs_fuse_rules()

Yield every (module_class_substrings, fuse_into, fuse_from) AWQ fusion rule.

Aggregated across all registered specs (the consumer matches each model module against the substrings, so the order across specs does not matter).

match_moe_block(module, model_type=None)

Return the MoE layout variant for module, resolved by model type.

model_type is the model’s root HF type (model.config.model_type) and is a strict filter: only that model’s own spec is consulted. A model whose model_type has no spec resolves to None even if its module class names coincide with another model’s — register a spec instead of inheriting a neighbor’s data. None (no config available: unit tests, the TRT-LLM path) searches all specs. Sub-model types of composite models are not walked today; a composite whose MoE lives under a tower type registers the root type too (see gemma4.py).

Within the spec, variant block_names identifies the block and disambiguates same-model layout variants (MoESpec.match_moe_variant); quantized wrapper classes match through their original base class in the MRO.

Parameters:

• module (nn.Module)

• model_type (str | None)

Return type:

MoEVariant | None

register(spec)

Register a model spec and return it. One spec per model type.

Parameters:

spec (ModelSpec)

Return type:

ModelSpec

weight_plus_one_norm_names()

All norm class names whose stored weight is w - 1, across all specs.

Return type:

tuple[str, …]