dynamic_block_prune

Generic (AnyModel) dynamic single-block pruning — no realized weights.

Makes one block behave as pruned during a forward pass, used by replace-1-block scoring (load the sorted teacher once, prune one block per candidate). The mechanism is descriptor-driven by module name:

  • Removal (FFN top-K, attention head/group removal): a forward-pre-hook masks the dropped channels/heads at the down_proj / o_proj input. Exact — the masked-forward output equals the pruned model’s output (the dropped channels/heads contribute zero downstream). Works on a plain tensor or a sharded DTensor (the mask is distributed to the activation’s placement).

The masks/keep-sets come from attention_ffn_surgery; the descriptor supplies the down_proj/o_proj/k_proj/v_proj module names, so this file is model-agnostic.

Classes

FFNRemovalSpec

Mask the down_proj input to keep keep_mask intermediate channels.

AttnRemovalSpec

Mask the o_proj input to keep keep_mask query-head columns.

Functions

apply_prune_hooks

Register removal masks for the given specs; returns hook handles to .remove().

register_mask_hook

Register the input-masking pre-hook on a resolved module object; returns the handle.

build_block_prune_specs

Map a per-block target onto the right D3 spec(s) for a sorted block.
class AttnRemovalSpec

Bases: object

Mask the o_proj input to keep keep_mask query-head columns.

__init__(module_name, keep_mask)
Parameters:

  • module_name (str)

  • keep_mask (Tensor)

Return type:

None

keep_mask: Tensor
module_name: str
class FFNRemovalSpec

Bases: object

Mask the down_proj input to keep keep_mask intermediate channels.

__init__(module_name, keep_mask)
Parameters:

  • module_name (str)

  • keep_mask (Tensor)

Return type:

None

keep_mask: Tensor
module_name: str
apply_prune_hooks(model, specs)

Register removal masks for the given specs; returns hook handles to .remove().

model.get_submodule(name) resolves each module (descriptor-supplied names), so this is model-agnostic.

Return type:

list

build_block_prune_specs(*, down_proj_name, o_proj_name, orig_intermediate, target_intermediate, orig_num_q, orig_num_kv, target_num_q, target_num_kv, head_dim)

Map a per-block target onto the right D3 spec(s) for a sorted block.

FFN target K (< orig) -> mask down_proj to the prefix [:K]. Attention removal keeps the first target_num_kv sorted groups and the first target_num_q/target_num_kv sorted query heads in each kept group. A q-preserving KV merge is intentionally not supported in the new Puzzletron attention contract. Module names are the loaded model’s paths (caller resolves them). Returns the spec list.

Parameters:

  • down_proj_name (str | None)

  • o_proj_name (str | None)

  • orig_intermediate (int | None)

  • target_intermediate (int | None)

  • orig_num_q (int | None)

  • orig_num_kv (int | None)

  • target_num_q (int | None)

  • target_num_kv (int | None)

  • head_dim (int | None)

Return type:

list

register_mask_hook(module, keep_mask)

Register the input-masking pre-hook on a resolved module object; returns the handle.