generic_decoder
Composable contracts for structurally standard decoder families.
Classes
A model-name-independent contract assembled from decoder components. |
|
ModelDescriptor adapter whose behavior is derived from a structural contract. |
|
Descriptor-owned MLA rank fields shared by conversion and runtime export. |
|
Per-layer embedding channel geometry shared by Gemma-style decoders. |
|
- class DecoderLayout
Bases:
object- __init__(language_config_path, language_prefix, layer_template, input_embedding, output_embedding, final_norm, layer_norm_names, hidden_size_field='hidden_size')
- Parameters:
language_config_path (tuple[str, ...])
language_prefix (str)
layer_template (str)
input_embedding (str)
output_embedding (str)
final_norm (str)
layer_norm_names (tuple[str, ...])
hidden_size_field (str)
- Return type:
None
- final_norm: str
- input_embedding: str
- language_config(config)
- Parameters:
config (Any)
- Return type:
Any
- language_config_path: tuple[str, ...]
- language_prefix: str
- layer_norm_names: tuple[str, ...]
- layer_path(layer_idx)
- Parameters:
layer_idx (int)
- Return type:
str
- property layer_pattern: str
- layer_template: str
- output_embedding: str
- class GatedDenseFFNContract
Bases:
object- __init__(module_name='mlp', gate_proj_name='gate_proj', up_proj_name='up_proj', down_proj_name='down_proj', intermediate_field='intermediate_size', double_wide_field=None, shared_kv_layers_field=None)
- Parameters:
module_name (str)
gate_proj_name (str)
up_proj_name (str)
down_proj_name (str)
intermediate_field (str)
double_wide_field (str | None)
shared_kv_layers_field (str | None)
- Return type:
None
- constructor_intermediate_size(config, *, layer_idx, actual_size)
Invert constructor-only expansion for a requested runtime width.
- Parameters:
config (Any)
layer_idx (int)
actual_size (int)
- Return type:
int
- double_wide_field: str | None = None
- down_proj_name: str = 'down_proj'
- gate_proj_name: str = 'gate_proj'
- intermediate_field: str = 'intermediate_size'
- layer_intermediate_size(config, *, layer_idx)
- Parameters:
config (Any)
layer_idx (int)
- Return type:
int
- module_name: str = 'mlp'
- up_proj_name: str = 'up_proj'
- class GenericContractModelDescriptor
Bases:
ModelDescriptorModelDescriptor adapter whose behavior is derived from a structural contract.
- DECODER_LAYER_CLS: type[nn.Module] | tuple[type[nn.Module], ...] | None = None
- static attn_no_op_post_init(decoder_layer)
- Parameters:
decoder_layer (nn.Module)
- classmethod block_config_to_layer_overrides(block_config)
- Parameters:
block_config (BlockConfig)
- classmethod decoder_layer_cls()
- classmethod embedding_pruning_spec(config, *, widths, alignment)
- Parameters:
alignment (int)
- classmethod final_norm_name()
- classmethod generic_decoder_contract(config)
- Return type:
- classmethod get_language_model_config(config)
- static init_rotary_embedding(model, runtime)
- classmethod input_embedding_name()
- classmethod layer_block_name(index)
- Parameters:
index (int)
- classmethod layer_name_predicates(num_layers)
- Parameters:
num_layers (int)
- classmethod local_kd_subblock_module_paths(block_config, *, layer_idx)
- Parameters:
block_config (BlockConfig)
layer_idx (int)
- Return type:
dict[tuple[str, str], str]
- static mlp_no_op_post_init(decoder_layer)
- Parameters:
decoder_layer (nn.Module)
- classmethod output_embedding_name()
- static passthrough_weight_name_predicates()
- classmethod patch_layer_config(layer_config, block_config, layer_idx)
- Parameters:
layer_config (Any)
block_config (BlockConfig)
layer_idx (int)
- Return type:
None
- classmethod ple_pruning_spec(config)
- classmethod puzzletron_capabilities(config)
- classmethod vision_module_names()
- Return type:
tuple[str, …]
- class GenericDecoderContract
Bases:
objectA model-name-independent contract assembled from decoder components.
- __init__(descriptor_name, model_family, layout, attention=None, latent_attention=None, dense_ffn=None, routed_moe=None, vision=None, mtp=None, ple=None, additional_tensor_rules=(), additional_exempt_patterns=(), checkpoint_key_rewrites=(), native_automodel_supported=True, ep_supported=False, sequence_parallel_supported=False, hidden_permutation_group_size=1, explicit_full_attention_window=False)
- Parameters:
descriptor_name (str)
model_family (str)
layout (DecoderLayout)
attention (StandardGQAAttentionContract | None)
latent_attention (LatentAttentionContract | None)
dense_ffn (GatedDenseFFNContract | None)
routed_moe (RoutedMoEContract | None)
vision (VisionLanguageContract | None)
mtp (MTPContract | None)
ple (PLEContract | None)
additional_tensor_rules (tuple[TensorAxisRule, ...])
additional_exempt_patterns (tuple[str, ...])
checkpoint_key_rewrites (tuple[tuple[str, str], ...])
native_automodel_supported (bool)
ep_supported (bool)
sequence_parallel_supported (bool)
hidden_permutation_group_size (int)
explicit_full_attention_window (bool)
- Return type:
None
- additional_exempt_patterns: tuple[str, ...] = ()
- additional_tensor_rules: tuple[TensorAxisRule, ...] = ()
- attention: StandardGQAAttentionContract | None = None
- capabilities(config=None)
- Parameters:
config (Any | None)
- checkpoint_key_rewrites: tuple[tuple[str, str], ...] = ()
- dense_ffn: GatedDenseFFNContract | None = None
- descriptor_name: str
- discover_prunable_modules(model)
- Parameters:
model (nn.Module)
- Return type:
dict[str, nn.Module]
- embedding_pruning_spec(config, *, widths, alignment)
- Parameters:
config (Any)
widths (tuple[int, ...] | list[int])
alignment (int)
- Return type:
- ep_supported: bool = False
- explicit_full_attention_window: bool = False
- language_config(config)
- Parameters:
config (Any)
- Return type:
Any
- latent_attention: LatentAttentionContract | None = None
- layout: DecoderLayout
- model_family: str
- mtp: MTPContract | None = None
- native_automodel_supported: bool = True
- ple: PLEContract | None = None
- reduced_axis_values(config, *, alignment)
- Parameters:
config (Any)
alignment (int)
- Return type:
dict[str, int]
- routed_moe: RoutedMoEContract | None = None
- sequence_parallel_supported: bool = False
- vision: VisionLanguageContract | None = None
- class LatentAttentionContract
Bases:
objectDescriptor-owned MLA rank fields shared by conversion and runtime export.
- __init__(module_name='self_attn', q_lora_rank_field='q_lora_rank', kv_lora_rank_field='kv_lora_rank')
- Parameters:
module_name (str)
q_lora_rank_field (str)
kv_lora_rank_field (str)
- Return type:
None
- kv_lora_rank_field: str = 'kv_lora_rank'
- module_name: str = 'self_attn'
- q_lora_rank_field: str = 'q_lora_rank'
- class MTPContract
Bases:
object- __init__(tensor_rules, exempt_patterns=())
- Parameters:
tensor_rules (tuple[TensorAxisRule, ...])
exempt_patterns (tuple[str, ...])
- Return type:
None
- exempt_patterns: tuple[str, ...] = ()
- tensor_rules: tuple[TensorAxisRule, ...]
- class PLEContract
Bases:
objectPer-layer embedding channel geometry shared by Gemma-style decoders.
- __init__(width_field='hidden_size_per_layer_input', layer_gate_name='per_layer_input_gate', layer_projection_name='per_layer_projection', model_embedding_name='embed_tokens_per_layer', model_projection_name='per_layer_model_projection', model_norm_name='per_layer_projection_norm')
- Parameters:
width_field (str)
layer_gate_name (str)
layer_projection_name (str)
model_embedding_name (str)
model_projection_name (str)
model_norm_name (str)
- Return type:
None
- layer_gate_name: str = 'per_layer_input_gate'
- layer_projection_name: str = 'per_layer_projection'
- model_embedding_name: str = 'embed_tokens_per_layer'
- model_norm_name: str = 'per_layer_projection_norm'
- model_projection_name: str = 'per_layer_model_projection'
- width_field: str = 'hidden_size_per_layer_input'
- class RoutedMoEContract
Bases:
object- __init__(module_name='mlp', experts_name='experts', router_name='gate', shared_expert_name='shared_expert', gate_proj_name='gate_proj', up_proj_name='up_proj', down_proj_name='down_proj', num_experts_field='num_experts', intermediate_field='moe_intermediate_size', shared_intermediate_field='shared_expert_intermediate_size', top_k_field='num_experts_per_tok', replaces_dense_ffn=True)
- Parameters:
module_name (str)
experts_name (str)
router_name (str)
shared_expert_name (str | None)
gate_proj_name (str)
up_proj_name (str)
down_proj_name (str)
num_experts_field (str)
intermediate_field (str)
shared_intermediate_field (str)
top_k_field (str)
replaces_dense_ffn (bool)
- Return type:
None
- down_proj_name: str = 'down_proj'
- experts_name: str = 'experts'
- gate_proj_name: str = 'gate_proj'
- intermediate_field: str = 'moe_intermediate_size'
- module_name: str = 'mlp'
- num_experts_field: str = 'num_experts'
- replaces_dense_ffn: bool = True
- router_name: str = 'gate'
- top_k_field: str = 'num_experts_per_tok'
- up_proj_name: str = 'up_proj'
- class StandardGQAAttentionContract
Bases:
object- __init__(module_name='self_attn', q_proj_name='q_proj', k_proj_name='k_proj', v_proj_name='v_proj', o_proj_name='o_proj', query_heads_field='num_attention_heads', kv_heads_field='num_key_value_heads', head_dim_field='head_dim', global_head_dim_field=None, global_kv_heads_field=None, k_eq_v_field=None, shared_kv_layers_field=None)
- Parameters:
module_name (str)
q_proj_name (str)
k_proj_name (str)
v_proj_name (str)
o_proj_name (str)
query_heads_field (str)
kv_heads_field (str)
head_dim_field (str)
global_head_dim_field (str | None)
global_kv_heads_field (str | None)
k_eq_v_field (str | None)
shared_kv_layers_field (str | None)
- Return type:
None
- global_head_dim_field: str | None = None
- global_kv_heads_field: str | None = None
- head_dim_field: str = 'head_dim'
- k_eq_v_field: str | None = None
- k_proj_name: str = 'k_proj'
- kv_heads_field: str = 'num_key_value_heads'
- layer_geometry(config, *, layer_idx, layer_types)
Resolve one layer’s true Q/K/V geometry and KV owner.
Standard models simply return their global GQA fields. Families with full/sliding heterogeneity can declare alternate full-attention fields, K=V coupling, and trailing KV-sharing layers without forking conversion.
- Parameters:
config (Any)
layer_idx (int)
layer_types (tuple[str, ...])
- Return type:
tuple[int, int, int, bool, int | None]
- layer_override_fields(attention)
Map one typed attention config to constructor fields.
Families with heterogeneous local/global geometry describe the alternate fields above. The generic descriptor can therefore target the correct constructor fields without naming the family or duplicating its rules.
- Parameters:
attention (Any)
- Return type:
dict[str, int]
- module_name: str = 'self_attn'
- o_proj_name: str = 'o_proj'
- q_proj_name: str = 'q_proj'
- query_heads_field: str = 'num_attention_heads'
- v_proj_name: str = 'v_proj'
- validate(config)
- Parameters:
config (Any)
- Return type:
tuple[int, int, int]
- class VisionLanguageContract
Bases:
object- __init__(module_names, projector_rules, projector_output_config_paths=(), exempt_patterns=())
- Parameters:
module_names (tuple[str, ...])
projector_rules (tuple[TensorAxisRule, ...])
projector_output_config_paths (tuple[tuple[str, ...], ...])
exempt_patterns (tuple[str, ...])
- Return type:
None
- exempt_patterns: tuple[str, ...] = ()
- module_names: tuple[str, ...]
- projector_output_config_paths: tuple[tuple[str, ...], ...] = ()
- projector_rules: tuple[TensorAxisRule, ...]