attention_ffn_surgery

Model-agnostic surgery primitives for FFN-channel and attention-head pruning.

These are pure tensor functions — no model, descriptor, or parallelism dependency — that the descriptor wires up by passing the standard HF projection weights (gate/up/down, q/k/v/o) and a head layout. They back the whole “sorted teacher + dynamic prune” flow:

• permute_*: reorder a contracted dim by importance. Because the reorder is applied consistently to every weight that touches that dim, the module’s output is unchanged — this is how the sorted teacher stays functionally identical to the teacher (unit-tested).

• slice_*: keep a subset of channels/heads -> physically smaller weights (bypass / realize).

• *_keep_mask: a boolean mask over the projection’s input features used to prune by zeroing activations at the down_proj / o_proj input during scoring (exact for removal: the masked forward equals the sliced model’s forward).

Conventions (standard HF GQA layout): gate/up are [intermediate, hidden] (rows = intermediate channels), down is [hidden, intermediate] (cols = intermediate channels); q is usually [num_q*head_dim, hidden], k/v are [num_kv*head_dim, hidden], and o is [hidden, num_q*head_dim]. Some families store multiple row groups per query head in q (for example Qwen3.5 stores query and gate rows); those row groups are carried together when query heads are permuted or sliced. Query heads are grouped by KV head (group-major: q_index = group*n_heads_in_group + head_in_group).

Functions

ffn_permutation	Descending-importance permutation of the intermediate channels ([intermediate] -> idx).
permute_ffn_weights	Reorder the intermediate dim consistently.
slice_ffn_weights	Keep the keep intermediate channels (LongTensor / slice) -> smaller weights.
ffn_keep_mask	Boolean [intermediate] mask (True = keep) for masking the down_proj input.
attention_permutations	Group-major query-head permutation + KV-head permutation, both by importance.
grouped_attention_permutations	Sort KV groups and query heads from explicit grouped-attention scores.
attention_permutations_from_scores	(q_perm [num_q], kv_perm [num_kv]) from either a per-query or per-KV importance vector.
permute_query_rows_by_head	Permute query-head rows, preserving any extra per-head row groups.
permute_attention_weights	Reorder query heads (q rows, o cols) and KV heads (k/v rows).
slice_query_rows_by_head	Slice query-head rows, preserving any extra per-head row groups.
slice_attention_weights	Keep the given query heads (q rows, o cols) and KV heads (k/v rows) -> smaller weights.
sorted_attention_keep_indices	Heads to keep when pruning a sorted attention layer to (target q, target kv).
attention_keep_mask	Boolean [num_q*head_dim] mask (True = keep) for masking the o_proj input.
aggregate_query_scores_to_kv	Sum a per-query-head [num_q] importance into per-KV-head [num_kv] (group-major).

aggregate_query_scores_to_kv(query_importance, num_kv_heads)

Sum a per-query-head [num_q] importance into per-KV-head [num_kv] (group-major).

Parameters:

• query_importance (Tensor)

• num_kv_heads (int)

Return type:

Tensor

attention_keep_mask(num_q_heads, keep_q_heads, head_dim, device=None)

Boolean [num_q*head_dim] mask (True = keep) for masking the o_proj input.

Zeroing the removed query heads’ columns at the o_proj input reproduces the pruned attention output exactly (each query head contributes independently to o_proj).

Parameters:

• num_q_heads (int)

• head_dim (int)

Return type:

Tensor

attention_permutations(query_importance, num_kv_heads, n_heads_in_group)

Group-major query-head permutation + KV-head permutation, both by importance.

Groups are ordered by aggregate (summed) query-head importance; query heads within each group are ordered by their own importance. Returns (q_perm [num_q], kv_perm [num_kv]) such that applying them to q/o (q_perm) and k/v (kv_perm) yields the sorted, functionally-identical layout.

Parameters:

• query_importance (Tensor)

• num_kv_heads (int)

• n_heads_in_group (int)

Return type:

tuple[Tensor, Tensor]

attention_permutations_from_scores(attn_importance, num_q_heads, num_kv_heads)

(q_perm [num_q], kv_perm [num_kv]) from either a per-query or per-KV importance vector.

Legacy compatibility helper for callers that still pass one flat attention score vector. New Puzzletron attention scoring should prefer grouped_attention_permutations(), which takes explicit KV-group and within-group query-head scores.

Parameters:

• attn_importance (Tensor)

• num_q_heads (int)

• num_kv_heads (int)

Return type:

tuple[Tensor, Tensor]

ffn_keep_mask(intermediate_size, keep, device=None)

Boolean [intermediate] mask (True = keep) for masking the down_proj input.

Parameters:

intermediate_size (int)

Return type:

Tensor

ffn_permutation(channel_importance)

Descending-importance permutation of the intermediate channels ([intermediate] -> idx).

Parameters:

channel_importance (Tensor)

Return type:

Tensor

grouped_attention_permutations(kv_group_importance, query_importance, num_kv_heads, n_heads_in_group)

Sort KV groups and query heads from explicit grouped-attention scores.

kv_group_importance ranks whole KV groups. query_importance ranks query heads within each original KV group. The returned query permutation is group-major after the KV group sort and prefix-sliceable for both supported attention operations: drop whole KV groups and keep the same number of query heads per surviving group.

Parameters:

• kv_group_importance (Tensor)

• query_importance (Tensor)

• num_kv_heads (int)

• n_heads_in_group (int)

Return type:

tuple[Tensor, Tensor]

permute_attention_weights(q, k, v, o, q_perm, kv_perm, head_dim)

Reorder query heads (q rows, o cols) and KV heads (k/v rows). Output unchanged (invariant).

permute_ffn_weights(gate, up, down, perm)

Reorder the intermediate dim consistently. Output is unchanged (permutation invariant).

permute_query_rows_by_head(q, q_perm, head_dim, num_q_heads)

Permute query-head rows, preserving any extra per-head row groups.

Parameters:

• q (Tensor)

• q_perm (Tensor)

• head_dim (int)

• num_q_heads (int)

Return type:

Tensor

slice_attention_weights(q, k, v, o, keep_q_heads, keep_kv_heads, head_dim)

Keep the given query heads (q rows, o cols) and KV heads (k/v rows) -> smaller weights.

slice_ffn_weights(gate, up, down, keep)

Keep the keep intermediate channels (LongTensor / slice) -> smaller weights.

slice_query_rows_by_head(q, keep_q_heads, head_dim, num_q_heads)

Slice query-head rows, preserving any extra per-head row groups.

Parameters:

• q (Tensor)

• keep_q_heads (Tensor)

• head_dim (int)

• num_q_heads (int)

Return type:

Tensor

sorted_attention_keep_indices(target_kv_heads, target_heads_in_group, orig_heads_in_group)

Heads to keep when pruning a sorted attention layer to (target q, target kv).

On the sorted teacher, keeping the most important (q, kv) = the first target_kv_heads groups and, within each, its first target_heads_in_group query heads. Returns (keep_q_heads, keep_kv_heads) index tensors (the query set is blocked, not a flat prefix).

Parameters:

• target_kv_heads (int)

• target_heads_in_group (int)

• orig_heads_in_group (int)

Return type:

tuple[Tensor, Tensor]