block_config#

Classes

BaseDataclass

SubblockConfig

SubblockRef

PrunableAxis

VariantAxis

MoEConfig

MambaConfig

Llama4AttentionConfig

AttentionConfig

MLAConfig

Per-layer multi-head latent-attention geometry and compression ranks.

FFNConfig

BlockConfig

Functions

maybe_cast_block_configs

iter_subblocks

class AttentionConfig#

Bases: SubblockConfig

__init__(*, kind='attention', name='attention', no_op=False, num_kv_heads=None, num_query_heads=None, qk_head_dim=None, v_head_dim=None, sliding_window_size=None, k_eq_v=None, kv_source_layer=None, llama4=None)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

  • num_kv_heads (int | None)

  • num_query_heads (int | None)

  • qk_head_dim (int | None)

  • v_head_dim (int | None)

  • sliding_window_size (int | str | None)

  • k_eq_v (bool | None)

  • kv_source_layer (int | None)

  • llama4 (Llama4AttentionConfig | None)

Return type:

None

k_eq_v: bool | None = None#
kind: str = 'attention'#
kv_source_layer: int | None = None#
llama4: Llama4AttentionConfig | None = None#
name: str = 'attention'#
num_kv_heads: int | None = None#
num_query_heads: int | None = None#
qk_head_dim: int | None = None#
sliding_window_size: int | str | None = None#
to_blockconfig()#
Return type:

BlockConfig

v_head_dim: int | None = None#
class BaseDataclass#

Bases: object

__init__()#
Return type:

None

to_dict()#
Return type:

dict[str, Any]

class BlockConfig#

Bases: BaseDataclass

__init__(*, subblock_configs)#
Parameters:

subblock_configs (tuple[SubblockConfig, ...])

Return type:

None

get_subblock(kind, name=None)#
Parameters:

  • kind (str)

  • name (str | None)

Return type:

SubblockConfig | None

prunable_axes()#
Return type:

list[PrunableAxis]

require_subblock(kind, name=None)#
Parameters:

  • kind (str)

  • name (str | None)

Return type:

SubblockConfig

subblock_configs: tuple[SubblockConfig, ...]#
subblocks()#
Return type:

list[SubblockRef]

to_dict()#
Return type:

dict[str, Any]

variant_axes()#
Return type:

list[VariantAxis]

with_subblock(subblock, *, replace_kinds=())#
Parameters:
Return type:

BlockConfig

without_subblocks(*kinds)#
Parameters:

kinds (str)

Return type:

BlockConfig

class FFNConfig#

Bases: SubblockConfig

__init__(*, kind='ffn', name='ffn', no_op=False, intermediate_size=None)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

  • intermediate_size (int | None)

Return type:

None

intermediate_size: int | None = None#
kind: str = 'ffn'#
name: str = 'ffn'#
to_blockconfig()#
Return type:

BlockConfig

class Llama4AttentionConfig#

Bases: BaseDataclass

__init__(*, attention_chunk_size=None, use_rope=None, use_qk_norm=None, attn_scale=None, floor_scale=None, attn_temperature_tuning=None, attention_dropout=None)#
Parameters:

  • attention_chunk_size (int | None)

  • use_rope (bool | None)

  • use_qk_norm (bool | None)

  • attn_scale (float | None)

  • floor_scale (float | None)

  • attn_temperature_tuning (bool | None)

  • attention_dropout (float | None)

Return type:

None

attention_chunk_size: int | None = None#
attention_dropout: float | None = None#
attn_scale: float | None = None#
attn_temperature_tuning: bool | None = None#
floor_scale: float | None = None#
use_qk_norm: bool | None = None#
use_rope: bool | None = None#
class MLAConfig#

Bases: SubblockConfig

Per-layer multi-head latent-attention geometry and compression ranks.

MLA decodes one non-RoPE key/value pair for every query head while sharing the rotary key component. num_heads is consequently one coupled axis, not separate GQA query-head and KV-group axes.

__init__(*, kind='mla', name='mla', no_op=False, num_heads=None, q_lora_rank=None, kv_lora_rank=None)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

  • num_heads (int | None)

  • q_lora_rank (int | None)

  • kv_lora_rank (int | None)

Return type:

None

kind: str = 'mla'#
kv_lora_rank: int | None = None#
name: str = 'mla'#
num_heads: int | None = None#
q_lora_rank: int | None = None#
class MambaConfig#

Bases: SubblockConfig

__init__(*, kind='mamba', name='mamba', no_op=False, state_dim=None, num_heads=None, head_dim=None, num_groups=None, conv_kernel_size=4)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

  • state_dim (int | None)

  • num_heads (int | None)

  • head_dim (int | None)

  • num_groups (int | None)

  • conv_kernel_size (int | None)

Return type:

None

conv_kernel_size: int | None = 4#
head_dim: int | None = None#
kind: str = 'mamba'#
name: str = 'mamba'#
num_groups: int | None = None#
num_heads: int | None = None#
state_dim: int | None = None#
class MoEConfig#

Bases: SubblockConfig

__init__(*, kind='moe', name='moe', no_op=False, num_experts=None, expert_intermediate_size=None, shared_expert_intermediate_size=None, top_k=None, latent_dim=None)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

  • num_experts (int | None)

  • expert_intermediate_size (int | None)

  • shared_expert_intermediate_size (int | None)

  • top_k (int | None)

  • latent_dim (int | None)

Return type:

None

expert_intermediate_size: int | None = None#
kind: str = 'moe'#
latent_dim: int | None = None#
name: str = 'moe'#
num_experts: int | None = None#
shared_expert_intermediate_size: int | None = None#
top_k: int | None = None#
class PrunableAxis#

Bases: object

__init__(*, axis_id, layer_idx, subblock, field, size=None, group_size=None, sort_kind='permutation', tensor_bindings=(), parallel_axis=None)#
Parameters:

  • axis_id (str)

  • layer_idx (int | None)

  • subblock (str)

  • field (str)

  • size (int | None)

  • group_size (int | None)

  • sort_kind (str)

  • tensor_bindings (tuple[str, ...])

  • parallel_axis (str | None)

Return type:

None

axis_id: str#
field: str#
group_size: int | None = None#
layer_idx: int | None#
parallel_axis: str | None = None#
size: int | None = None#
sort_kind: str = 'permutation'#
subblock: str#
tensor_bindings: tuple[str, ...] = ()#
class SubblockConfig#

Bases: BaseDataclass

__init__(*, kind, name, no_op=False)#
Parameters:

  • kind (str)

  • name (str)

  • no_op (bool)

Return type:

None

kind: str#
name: str#
no_op: bool = False#
class SubblockRef#

Bases: object

__init__(*, index, kind, name, config)#
Parameters:
Return type:

None

config: SubblockConfig#
index: int#
kind: str#
name: str#
class VariantAxis#

Bases: object

__init__(*, axis_id, layer_idx, subblock, field, values=(), requires_rpc=False, requires_vllm=False)#
Parameters:

  • axis_id (str)

  • layer_idx (int | None)

  • subblock (str)

  • field (str)

  • values (tuple[Any, ...])

  • requires_rpc (bool)

  • requires_vllm (bool)

Return type:

None

axis_id: str#
field: str#
layer_idx: int | None#
requires_rpc: bool = False#
requires_vllm: bool = False#
subblock: str#
values: tuple[Any, ...] = ()#
iter_subblocks(block_configs)#
Parameters:

block_configs (Iterable[BlockConfig])

Return type:

Iterable[SubblockRef]

maybe_cast_block_configs(block_configs)#
Parameters:

block_configs (list[BlockConfig | dict[str, Any]] | None)

Return type:

list[BlockConfig] | None