onnx: import MatMulNBits at bits=2

[czoli1976]

Stacked on #2648 — review that one first; this adds bits=2 on top of the same zero-point split.

MatMulNBits rejected anything but bits=4, so the 2-bit ORT-GenAI exports could not be imported at all. These are the natural width for a ternary-trained model: the Bonsai 1.7B export is 482 MB at 2 bits against 1.1 GB at 4, for logits that agree to 6e-5.

What changed

The code and zero-point unpacking is now written once for either width — 4-bit reads two values per byte, 2-bit four, both low-order bits first — rather than hardcoding nibbles.

For the weight itself, tract already has a 2-bit block quant in Q2_0_T, and it fits: its dequant is (code - 1) * scale for any of the four codes, not just the ternary three it quantizes to. So the split #2648 uses for Q4_0 applies unchanged with a base of 1 instead of 8:

(q - z) * s  ==  (q - 1) * s  +  (1 - z) * s

The exports block by 128 while both formats block by 32. A wider quantization block holds a single scale for several 32-blocks, so it is carried by repeating that scale — the dequantized weight is identical, at 2.5 bits/weight against the export's 2.14, and against 32 for the f32 fallback.

One consequence worth noting: the correction is no longer tied to the zero-point input being present. A symmetric 2-bit export defaults its zero to 2, which is not Q2_0_T's 1, so it needs the offset too; the code now decides on the offsets themselves rather than on which inputs exist.

Q2_0_T gains a pack_prequantized mirroring Q4_0's, so an importer can hand over codes it already has.

Testing

Checked against onnxruntime across block sizes 32 and 128, with and without zero points, and at both bit widths so the 4-bit path is covered too: relative error ~2.5e-4, which is the f16 scale rounding. End to end, onnx-community/Bonsai-1.7B-ONNX/model_q2.onnx (Qwen3, 28 layers, 2-bit) matches onnxruntime at 4.6e-6 relative with identical argmax at every position.

The added case uses integer activations and power-of-two scales so every product and partial sum is exactly representable, making it an equality test of the unpacking rather than of accumulation order. Its nnef pass is skipped: Q2_0_T has no NNEF tensor serializer today (only Q4_0 and Q8_1 do), which is untouched by this change.

🍍