Banded sharing

How the shared_blocks count (G) controls how many layers reuse one shared block, and why production serving runs at the per-layer passthrough point.

A decoder stores one dense weight matrix per linear projection per layer, and those matrices dominate the checkpoint. Banded sharing is the dial that decides how many of those layers reuse a single shared block W0 from the theta bundle instead of each carrying a full matrix. The dial is shared_blocks, written G, and it applies per projection family (q_proj, k_proj, and so on).

G moves between two extremes:

• Full sharing (G=1): every layer reuses a single W0 per family, so the theta bundle stores one shared block plus a per-layer scale and delta. This is the smallest bundle.
• Per-layer (G=num_layers): each layer keeps its own W0, so the reconstructed weight equals the original. This is the passthrough operating point, and it’s the production setting.

Between those ends, attnquant partitions the decoder’s layers into G contiguous bands, each with its own W0 per family, the sharing scheme from the Relaxed Recursive Transformer paper. A narrower band sits closer to the layers inside it, so the per-layer low-rank delta shrinks. The cost is storing G blocks instead of one.

How layers map to bands

Bands are contiguous and as equal in size as possible, and earlier bands take the remainder. The model config’s band_index() method in models/gemma3/model_config.py produces the layer-to-band map, and W0 injection uses the correct band’s block for each layer. Here’s the assignment for a 26-layer model (Gemma 3 1B):

Theta-bundle key layout

Only the W0 key changes with G. The per-layer s, U, and V keys are identical either way, because a band is a property of the shared block alone:

G = 1:   W0.q_proj
G > 1:   W0.0.q_proj

For the full theta-bundle key layout, see reconstruction.

Set G for serving

shared_blocks (the config field), ATTNQUANT_BLOCKS (the env override), and G (the symbol) are the same value. The env variable applies to both the decomposition adapter and the model, so the two stay in sync.

For serving, run at the passthrough point G = num_layers. Set ATTNQUANT_BLOCKS to the layer count at runtime (ATTNQUANT_BLOCKS=26 for Gemma 3 1B, ATTNQUANT_BLOCKS=60 for Gemma 4 31B), or set shared_blocks to the layer count in config.json. A serve command looks like this:

ATTNQUANT_BLOCKS=26 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 \
  pixi run python run_max.py serve \
  --model google/gemma-3-1b-it \
  --custom-architectures attnquant \
  --port 3100

At passthrough the reconstructed weight equals the original, so output is native-equivalent, verified by the logit-parity checks on the quality page.

G is a bundle-size dial, not a speed dial: on Gemma 3 1B, G=1 versus G=26 measures 1.035x with capture off (CI [0.780, 1.345]) and 1.090x with capture on (CI [0.996, 1.245]), neither a resolved difference at that model’s ~20-25% minimum detectable effect. See throughput.

Settings below passthrough (G < num_layers) trade quality for size and are for decomposition-error analysis, not serving: making them production-quality requires a training or distillation step. If you want smaller weights today, use NVFP4 quantization, which compresses whole-model weight bytes 3.18x at passthrough with no sharing.