aimet_onnx.experimental.spinquant

Top level APIs

aimet_onnx.experimental.spinquant.apply_spinquant(model, visual_model=None, embedding=None, *, enable_r1=True, enable_r2=False, enable_r3=False, topology=None)[source]

Apply SpinQuant rotation transforms to an ONNX transformer model.

SpinQuant applies orthogonal Hadamard rotations to model weights to reduce quantization error. This function modifies the ONNX model(s) in-place by:

  1. Reading the backbone’s decoder-stack structure (block boundaries, role map, hidden size) off topology, and analyzing the optional visual encoder (PatchMerger output projection).

  2. Validating every selected rotation pass against the analysis.

  3. Applying every selected pass in order (R1 before R2 before R3).

  4. Serializing the result back onto the caller’s ModelProtos.

The rotations are performed on an onnx_ir copy of the graph and written back only once every pass has succeeded, so a failure part-way through leaves the caller’s model untouched rather than half-rotated. (The embedding tensor is the exception: it is a torch.Tensor rotated in place.)

Must be called on the float ONNX model BEFORE creating a QuantizationSimModel. The rotation modifies float weight initializers (R1 / R2) and may insert new nodes (R3); build the sim on the rotated graph and run compute_encodings afterward so quantizer scales are calibrated on the rotated weights.

Supported architectures:

  • LLaMA, Qwen2, Qwen3, Phi3 (backbone only)

  • Qwen2.5-VL, Qwen3-VL (backbone + visual)

Parameters:

  • model (ModelProto) – backbone.onnx ModelProto. Mutated in-place.

  • visual_model (Optional[ModelProto]) – Optional visual.onnx ModelProto (VLM only). Mutated in-place.

  • embedding (Optional[Tensor]) – Optional torch.Tensor of shape [vocab, hidden] loaded from embedding.pth (VLM only). Rotated in-place with R_L.

  • enable_r1 (bool) – If True (default), apply the R1 (residual-stream) rotation.

  • enable_r2 (bool) – If True, apply the R2 (per-head) rotation. Defaults to False. Not supported on architectures with fused QKV projections (e.g. Phi3).

  • enable_r3 (bool) – If True, apply the R3 online Hadamard rotation on Q and K paths into each block’s QK^T MatMul. Defaults to False. Inserts new MatMul nodes into the ONNX graph (does not mutate existing weights). MHA only — not supported on fused QKV or per-head split exports. The K-side rotation is placed upstream of the past-key Concat so K values entering the KV cache are already rotated; the model’s present_key output then carries rotated K (cache convention is self-consistent across steps).

  • topology (Optional[LlmTopology]) – Decoder-stack topology of model, from analyze_llm_topology(). Every rotation is placed by it — which linears read from and write to the residual stream (R1), which are the V/O projections (R2), which are the Q/K edges into QKᵀ (R3) — because discovering model structure belongs to llm_topology, not to SpinQuant. Analyze the same model this call rotates. Passing it explicitly is also the only way to override the analysis: this function would otherwise analyze with active_norms_per_block=2 and the default role patterns, which no argument here can change. Optional today: when omitted the topology is analyzed internally from model and a warning is raised. Passing it explicitly is recommended, and may become required in a future release.

Raises:

ValueError – If no rotation is enabled, if block detection or role classification fails, if topology does not describe model, or if any expected weight is missing / has the wrong shape.

Return type:

None

Example (LLM):

topology = analyze_llm_topology(model, model_type="llama")
apply_spinquant(model, topology=topology)
sim = QuantizationSimModel(model)   # built on the rotated graph
sim.compute_encodings(calibration_data)

Example (VLM):

embedding = torch.load("embedding.pth")   # torch.Tensor [vocab, hidden]
topology = analyze_llm_topology(backbone_model, model_type="qwen2_5_vl")
apply_spinquant(
    backbone_model,
    visual_model=visual_model,
    embedding=embedding,
    topology=topology,
)
torch.save(embedding, "embedding.pth")    # overwrite with rotated weights
backbone_sim = QuantizationSimModel(backbone_model)
visual_sim = QuantizationSimModel(visual_model)
backbone_sim.compute_encodings(backbone_calibration_data)
visual_sim.compute_encodings(visual_calibration_data)

Model topology

SpinQuant places every rotation using an LlmTopology, which describes the structure of the decoder stack: which linears read from and write to the residual stream (R1), which are the V and O projections (R2), and which Q/K edges feed each QKᵀ MatMul (R3). Analyze the model, then pass the result as topology:

topology = analyze_llm_topology(onnx_model, model_type="llama")
apply_spinquant(onnx_model, topology=topology)
sim = QuantizationSimModel(onnx_model, ...)   # built on the rotated graph
aimet_onnx.experimental.llm_topology.analyze_llm_topology(model, model_type, *, ir_model=None)[source]

Analyze model end-to-end and return a name-based LlmTopology.

How the decoder blocks and their layers are found depends on model_type:

  • For ACTIVE_NORM_MODEL_TYPES, by graph structure, as analyze_llm_topology_by_norm_count() does with its defaults: every active RMSNorm is detected, and the norms are grouped two per decoder block.

  • For every other supported model type, by HuggingFace module name (see hf_patterns): the per-block projections and norms, embed_tokens, the final norm and lm_head. Blocks are the decoder layers those names belong to. The remaining fields are derived from the graph around the named layers:

    • residual_input — the tensor entering the block’s input norm; residual_output — the next block’s residual_input, or for the last block the tensor entering the final norm (the final residual Add’s output for a headless backbone).

    • qk_matmul / attn_v_matmul — the dynamic MatMuls between the named q/k/v and o_proj (best-effort: empty when attention is a fused op).

    • KV-cache names, hidden_size and head_dim — as before.

    The named layers are then cross-checked against the graph, so a name that points at the wrong node is an error rather than a wrong topology: each block’s q/k/v (gate/up) must be exactly the weighted linears its input (post-attention) norm feeds, o_proj and down_proj must be the linears writing the residual stream, and the final norm must feed exactly lm_head.

Parameters:

  • model (ModelProto) – ONNX ModelProto to analyze. May be a float export or a QuantizationSimModel graph. Not mutated. A dynamo export must first go through fix_node_names_pass().

  • model_type (str) – HuggingFace PretrainedConfig.model_type of the model, e.g. "llama". For a VLM, that of its top-level config (e.g. "qwen3_vl"); the name tables of VLM backbones not analyzed by active norms are keyed by the text config’s instead (e.g. "qwen2_5_vl_text").

  • ir_model (Optional[Model]) – Pre-built analysis IR for model — as returned by build_analysis_ir(), i.e. quantizer-stripped and RMSNorm-fused. Built here when None. Pass one only to avoid a second from_proto of a large model when the caller already holds it; a faithful (unfused) IR will not detect norms and must not be passed.

Return type:

LlmTopology

Returns:

LlmTopology with every field populated. head_dim is None when the export exposes no past_value graph input to derive it from.

Raises:

ValueError – If model_type is unsupported, if the layers cannot be identified by name (NamedLayerMatchError), if the named layers disagree with the graph, or if the decoder blocks cannot be found by active norms.

class aimet_onnx.experimental.llm_topology.LlmTopology(embed_tokens=<factory>, lm_head=<factory>, blocks=<factory>, past_key_input_names=<factory>, past_key_output_names=<factory>, past_value_input_names=<factory>, past_value_output_names=<factory>, active_norms=None, hidden_size=None, head_dim=None)[source]

Topology of an ONNX decoder-stack model: blocks + backbone-level roles + dims.

Parameters:

  • embed_tokens (List[str]) – Names of the token-embedding ``Gather``(s) that produce the initial residual-stream activations.

  • lm_head (List[str]) – Names of the vocabulary-projection linear(s) downstream of the final norm.

  • blocks (List[BlockTopology]) – Per-decoder-block topology in topological order.

  • past_key_input_names (List[str]) – Raw past_key_* graph inputs in declaration order, collected tolerantly (empty for prefill-only exports without a KV-cache). Pairing these to blocks and validating that their count matches the block count are the consumer’s responsibility (e.g. R3) — R1-only and prefill-only flows do not require KV-cache inputs.

  • past_key_output_names (List[str]) – Key-cache graph outputs in declaration order.

  • past_value_input_names (List[str]) – Value-cache graph inputs in declaration order.

  • past_value_output_names (List[str]) – Value-cache graph outputs in declaration order.

  • active_norms (Optional[List[ActiveNorm]]) – Active norms in topological order used to build the topology.

  • hidden_size (Optional[int]) – Residual-stream hidden dimension (None if not inferred; analyze_llm_topology() fills it).

  • head_dim (Optional[int]) – Per-head dimension (None when it could not be derived, e.g. an export without KV-cache inputs).