gated_delta_net

Qwen GatedDeltaNet geometry and coupled permutation helpers.

Classes

GDNPermutation

GDNShape

Functions

gated_delta_net_prefix_indices

Return coupled full-width indices for one legal nested GDN prefix.

permute_gated_delta_net_state_dict

Apply one function-preserving full-width Qwen GDN permutation in-place.

slice_gated_delta_net_state_dict

Physically slice every coupled Qwen GDN tensor to a nested target shape.
class GDNPermutation

Bases: object

__init__(key_groups, value_lanes, key_dim, value_dim)
Parameters:

  • key_groups (Tensor)

  • value_lanes (Tensor)

  • key_dim (Tensor)

  • value_dim (Tensor)

Return type:

None

classmethod from_score_payload(payload, shape)
Parameters:
Return type:

GDNPermutation

indices(shape)
Parameters:

shape (GDNShape)

Return type:

dict[str, Tensor]

key_dim: Tensor
key_groups: Tensor
value_dim: Tensor
value_lanes: Tensor
class GDNShape

Bases: object

__init__(num_key_heads, num_value_heads, key_head_dim, value_head_dim)
Parameters:

  • num_key_heads (int)

  • num_value_heads (int)

  • key_head_dim (int)

  • value_head_dim (int)

Return type:

None

classmethod from_mamba_config(config)
Return type:

GDNShape

classmethod from_module(module)
Return type:

GDNShape

key_head_dim: int
num_key_heads: int
num_value_heads: int
value_head_dim: int
property value_heads_per_group: int
gated_delta_net_prefix_indices(shape, target)

Return coupled full-width indices for one legal nested GDN prefix.

Parameters:
Return type:

dict[str, Tensor]

permute_gated_delta_net_state_dict(state_dict, *, prefix, shape, permutation)

Apply one function-preserving full-width Qwen GDN permutation in-place.

Parameters:
Return type:

dict[str, Tensor]

slice_gated_delta_net_state_dict(state_dict, *, prefix, shape, target)

Physically slice every coupled Qwen GDN tensor to a nested target shape.

Parameters:

  • state_dict (dict[str, Tensor])

  • prefix (str)

  • shape (GDNShape)

  • target (GDNShape)

Return type:

dict[str, Tensor]