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:

subblock (SubblockConfig)
replace_kinds (Iterable[str])

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:

index (int)
kind (str)
name (str)
config (SubblockConfig)

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