Skip to content

block_quant: quantize einsum weights at either input slot#2428

Open
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:fix/block-quant-weight-slot
Open

block_quant: quantize einsum weights at either input slot#2428
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:fix/block-quant-weight-slot

Conversation

@czoli1976

Copy link
Copy Markdown
Contributor

The block_quant transform assumed the constant weight was always EinSum input 0 (tract's canonical gmk,nk form), so on imported ONNX activation @ weight matmuls — where the weight is input 1 — it wired the block-quant tensor and the injected group axis to the wrong operands and tripped the EinSum input-rank check. This made block_quant unusable on essentially any imported transformer (a BERT/MiniLM encoder fails on the first attention matmul). The slot-0 assumption went unnoticed because the only coverage, the matmul_q40 suite, constructs its weights at slot 0.

Fix: derive the weight and activation slots from the matched input and wire them accordingly. Adds three regression tests covering both ONNX orientations (mk,kn, batched bmk,kn) and the canonical mk,nk, each asserting numerical equivalence to X @ Q4_0(W).

🍍

The block_quant rewrite assumed the constant weight was always einsum
input 0, so on imported ONNX `activation @ weight` matmuls (weight at
input 1) it wired the block-quant tensor and the injected group axis to
the wrong operands and tripped the EinSum rank check. Derive the weight
and activation slots from the matched input and wire them accordingly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant