core.decompose

API reference for core.decompose, the model-agnostic module that turns a standard Hugging Face state dict into a attnquant theta bundle and verifies the reconstruction.

core/decompose.py converts pretrained weights into the theta bundles that attnquant serves: for each projection, a shared block W0, a per-row scale s, and a low-rank delta U, V, which together rebuild the weight as diag(s) @ W0 + U @ V.T. The conversion is analytic (a per-row least-squares scale against W0, then a truncated SVD of the residual), so no training or distillation runs here.

The module is model-agnostic. It reads plain numpy.ndarray values keyed by the standard MAX/Hugging Face state-dict layout (layers.{i}.self_attn.{fam}.weight, layers.{i}.mlp.{fam}.weight) and emits the flat attnquant keys that per-model weight_adapters.py remaps to nested names. Any decoder that exposes those keys decomposes without a code change. For the runtime math and the two reconstruction paths, see reconstruction.

Projection families

A projection family is one of the seven linear projections that repeats in every decoder layer. Three module constants group them:

• FAMILIES: all seven, q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj. These are the projections made attnquant.
• ATTN_FAMILIES: the attention subset q_proj, k_proj, v_proj, o_proj, selected by tie_families="attn".
• MLP_FAMILIES: the MLP subset gate_proj, up_proj, down_proj.

The order of FAMILIES matches the KIND_ORDER used by the native exporter and the model config, so keep it stable when extending the tuple.

decompose_standard_state_dict_to_attnquant()

The top-level entry point takes a standard state dict and returns a flat theta state dict:

decompose_standard_state_dict_to_attnquant(
    state_dict,          # Mapping[str, array-like]
    rank=32,             # rank r of the low-rank delta (SVD of the residual)
    shared_blocks=1,     # number of contiguous layer bands (G)
) -> dict[str, np.ndarray]

It groups projection weights by family, computes each band’s W0 as the mean of that band’s layers, and decomposes every layer into s, U, V with decompose_projection_weight(). Non-projection weights (embeddings, norms, biases) pass through unchanged. It drops lm_head.weight unconditionally, since MAX materializes it from embed_tokens.weight at runtime, which is correct for tied-embedding models like Gemma; an untied decoder would need adapter changes.

The shared_blocks argument (G) controls how many shared blocks each family gets:

• shared_blocks=1 (default, G=1): one W0.{fam} per family, the mean across all layers.
• shared_blocks>1 (G>1): layers split into G contiguous bands, each band getting its own W0.{band}.{fam}. Smaller bands lower the per-band variance and shrink the residual that the low-rank delta has to carry. See banded sharing.

At the passthrough operating point (G = num_layers, one W0 per layer), each W0 equals that layer’s own weight, the residual is zero, and the reconstructed weight equals the original. This is the production setting: output is native-equivalent to the built-in path under the logit-parity checks.

The returned dict uses these keys:

• W0.{fam} (G=1) or W0.{band}.{fam} (G>1): shared block [out_dim, in_dim].
• layers.{i}.{fam}.s: per-layer scale [out_dim].
• layers.{i}.{fam}.U: low-rank left factor [out_dim, rank].
• layers.{i}.{fam}.V: low-rank right factor [in_dim, rank].
• All passthrough (non-projection) keys from the input.

The per-layer s, U, V keys are identical across G, since the band is a property of W0 only.

decompose_projection_weight()

This helper decomposes a single layer’s projection against its shared block:

decompose_projection_weight(weight, w0, rank) -> (s, U, V)

It solves a per-output-row least-squares scale, s = (weight . W0) / (W0 . W0), then takes a truncated SVD of the residual weight - diag(s) @ W0. When rank == 0, U and V have zero columns, leaving a pure scaled shared block with no delta. When the residual’s effective rank is smaller than the requested rank, U and V are zero-padded up to the full rank, so every family produces consistently shaped tensors. That padding case fires for k_proj and v_proj in grouped-query attention (GQA) models, where out_dim is smaller than rank.

reconstruct_weight() and relative_frobenius_error()

Two helpers check how well a theta bundle reproduces the original weight:

w = reconstruct_weight(s, w0, U, V)          # diag(s) @ W0 + U @ V.T
err = relative_frobenius_error(original, w)  # ||orig - w||_F / ||orig||_F

reconstruct_weight() rebuilds the dense matrix from theta on the host in NumPy, matching the runtime reconstruction. relative_frobenius_error() returns the relative Frobenius error as a float, so 0.0 at the passthrough point and larger as G drops or rank shrinks. Use this pair to sweep rank and shared_blocks offline before you serve.

is_attnquant_theta_dict() and theta_bundle_bytes()

Two utilities support the auto-decompose adapter and size reporting:

• is_attnquant_theta_dict(state_dict) returns True when a state dict already uses theta keys, either the G=1 layout (W0.{fam}) or the banded G>1 layout (W0.{band}.{fam}). The adapter calls this to pass pre-converted checkpoints through without re-decomposing.
• theta_bundle_bytes(theta) sums the .nbytes of every array in a bundle, useful when you report the shipped theta size.

Related pages

To wire this module into a new decoder through the six-file per-model pattern, see extending attnquant. For the reconstruction paths and how the framework auto-selects between them, see reconstruction.