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]