generic_decoder

Composable contracts for structurally standard decoder families.

Classes

DecoderLayout

GatedDenseFFNContract

GenericDecoderContract

A model-name-independent contract assembled from decoder components.

GenericContractModelDescriptor

ModelDescriptor adapter whose behavior is derived from a structural contract.

LatentAttentionContract

Descriptor-owned MLA rank fields shared by conversion and runtime export.

MTPContract

PLEContract

Per-layer embedding channel geometry shared by Gemma-style decoders.

RoutedMoEContract

StandardGQAAttentionContract

VisionLanguageContract

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
hidden_size_field: str = 'hidden_size'
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'
shared_kv_layers_field: str | None = None
up_proj_name: str = 'up_proj'
class GenericContractModelDescriptor

Bases: ModelDescriptor

ModelDescriptor 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:

GenericDecoderContract

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:
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: object

A 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:
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:

EmbeddingPruningSpec

ep_supported: bool = False
explicit_full_attention_window: bool = False
hidden_permutation_group_size: int = 1
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: object

Descriptor-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: object

Per-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'
shared_expert_name: str | None = 'shared_expert'
shared_intermediate_field: str = 'shared_expert_intermediate_size'
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'
shared_kv_layers_field: str | None = None
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, ...]