veloxquant_mlx.profiling.model_profiler module

Model-architecture extraction for the auto-selection pipeline.

Profiles a language model’s attention geometry without loading weights: from a HuggingFace-style config.json (a dict), a loaded HF/MX PretrainedConfig (object with attribute access), or a loaded mlx_lm model. The derived ModelProfile is what the candidate filter and memory estimator feed on.

class veloxquant_mlx.profiling.model_profiler.ModelProfile(model_id='unknown', architecture='unknown', num_layers=0, num_query_heads=0, num_kv_heads=0, head_dim=0, attention_type='gqa', dtype='float16', parameter_count=None, hidden_size=None, _extra=<factory>)

Bases: object

Normalized model-architecture description used by the planner.

Variables:

• model_id – Model identifier (any string; the caller’s model_id).

• architecture – Short architecture slug ("llama", "qwen", …).

• num_layers – Number of attention-bearing layers.

• num_query_heads – Attention heads per layer.

• num_kv_heads – Key/value heads per layer (== num_query_heads for MHA).

• head_dim – Attention head dimension.

• attention_type – "mha", "gqa", or "mqa".

• dtype – Expected compute dtype ("float16", …).

• parameter_count – Total weight count, when the config reports one.

• hidden_size – Model hidden width, when known (informational).

Parameters:

• model_id (str)

• architecture (str)

• num_layers (int)

• num_query_heads (int)

• num_kv_heads (int)

• head_dim (int)

• attention_type (str)

• dtype (str)

• parameter_count (int | None)

• hidden_size (int | None)

• _extra (dict[str, Any])

model_id: str = 'unknown'

architecture: str = 'unknown'

num_layers: int = 0

num_query_heads: int = 0

num_kv_heads: int = 0

head_dim: int = 0

attention_type: str = 'gqa'

dtype: str = 'float16'

parameter_count: int | None = None

hidden_size: int | None = None

property baseline_kv_bytes_per_token: int

Full K+V footprint in bytes for one token, one layer (fp16).

baseline_kv_bytes(num_layers=None, batch=1)

Total fp16 K+V footprint across layers/batch for one token.

Return type:

int

Parameters:

• num_layers (int | None)

• batch (int)

to_dict()

Serialize all fields to a plain dict.

Return type:

dict[str, Any]

veloxquant_mlx.profiling.model_profiler.profile_model_from_config(config=None, model_id=None, *, num_layers=None, num_query_heads=None, num_kv_heads=None, head_dim=None, dtype=None, parameter_count=None, hidden_size=None, architecture=None)

Build a ModelProfile from a HF-style config without loading weights.

Process: 1. Extract num_layers from config (keys: num_hidden_layers, num_layers) 2. Extract num_query_heads (keys: num_attention_heads, num_heads) 3. Extract head_dim (key: head_dim) OR infer from hidden_size / num_query_heads 4. Extract num_kv_heads (key: num_key_value_heads); default to num_query_heads (MHA) 5. Detect attention type from head ratio (MHA/GQA/MQA) 6. Extract dtype, parameter_count, architecture 7. Override all extracted values with keyword arguments (if provided)

Config formats supported: - dict: HuggingFace config.json as dictionary - object: Loaded PretrainedConfig (attribute access) - None: Allowed if all critical kwargs are provided

Keyword overrides: All kwargs override config values (even if config present) Useful for partial configs or CLI inputs (–num-layers, –model-class).

Model families: - Llama, Qwen, Mistral, Mixtral, Phi, Gemma, OLMo, Command-R Auto-detected via architecture name or model_type field

Attention types: - MHA (Multi-Head Attention): num_kv_heads == num_query_heads - GQA (Grouped Query Attention): 1 < num_kv_heads < num_query_heads - MQA (Multi-Query Attention): num_kv_heads == 1

Return type:

ModelProfile

Parameters:

• config (dict[str, Any] | Any | None)

• model_id (str | None)

• num_layers (int | None)

• num_query_heads (int | None)

• num_kv_heads (int | None)

• head_dim (int | None)

• dtype (str | None)

• parameter_count (int | None)

• hidden_size (int | None)

• architecture (str | None)

Arguments:

config: HF config dict/object (optional; can be None if overrides given) model_id: Model identifier (e.g., “Qwen/Qwen2.5-7B”). Extracted from

config._name_or_path if not provided.

num_layers: Number of transformer layers (override) num_query_heads: Number of query attention heads (override) num_kv_heads: Number of key/value heads (override) head_dim: Attention head dimension (override) dtype: Compute dtype string (override) parameter_count: Total model parameters (override) hidden_size: Model hidden dimension (override) architecture: Architecture slug (override; auto-detected if not given)

Returns:

ModelProfile with all extracted/derived fields.

Raises:

ValueError: If critical geometry is missing (num_layers, num_query_heads,

head_dim cannot be inferred). Keyword overrides help resolve this.

Cost: ~1ms (pure config parsing, no weights loaded)

Example:

>>> config = {"num_hidden_layers": 32, "hidden_size": 4096,
...           "num_attention_heads": 32, "architectures": ["QwenForCausalLM"]}
>>> profile = profile_model_from_config(config, model_id="Qwen/Qwen2.5-7B")
>>> print(f"Architecture: {profile.architecture}")  # "qwen"
>>> print(f"Attention: {profile.attention_type}")  # "gqa" or "mha"

veloxquant_mlx.profiling.model_profiler.attention_type_from_heads(num_query_heads, num_kv_heads)

Classify MHA / GQA / MQA from the head counts.

Return type:

str

Parameters:

• num_query_heads (int)

• num_kv_heads (int)

veloxquant_mlx.profiling.model_profiler.architecture_alias(raw_name)

Map an HF class name to a short architecture slug, else pass through.

Return type:

str

Parameters:

raw_name (str)