TurboQuant roles inverted, plus scale-dtype seam and missing signed sub-8-bit decode

[ehartford]

Summary

Our TurboQuant implementation has the key and value quantizer roles inverted with
respect to the algorithm. Keys get plain uniform quantization; values get the
rotated Lloyd-Max path. It is the other way around.

This traces to version 1 of the shared contract, specs/formats/turboquant.md,
which described the inverted assignment. This backend implemented that contract
faithfully, which is why our backends agree with each other and disagree with the
published algorithm. The contract has been corrected to version 2 (QuixiAI/QuixiCore@613c73c).

What the algorithm specifies

TurboQuant (Zandieh et al., Online Vector Quantization with Near-optimal
Distortion Rate, arXiv:2504.19874, ICLR 2026)
defines two variants:

• TurboQuant_mse — random rotation, Lloyd-Max optimal scalar quantization,
  rotate back.
• TurboQuant_prod — the same plus a 1-bit QJL transform on the residual,
  giving an unbiased inner-product estimator.

Attention scores are inner products between a query and a key, so keys take
the rotated Lloyd-Max path — that is the tensor whose fidelity sets the score
distribution, and the tensor the coordinate-concentration argument is about. The
attention output is a convex combination of values, a reconstruction rather than
an inner product, so values take uniform quantization.

The reference implementation agrees: 3-bit keys via "random orthogonal rotation
→ Lloyd-Max optimal scalar quantization (b−1 bits) on Beta-distributed rotated
values", with values on group quantization with per-group scales and zeros
(https://github.com/0xsero/turboquant).

Impact

Keys are the tensor that most needs the rotation, and they are the one getting
vanilla uniform quantization. Key quantization error biases every attention
score, which is the failure mode the rotated path exists to prevent. Expect
quality loss concentrated at low bit widths and long contexts — exactly where
this format is meant to be used.

This is a correctness defect, not a performance or tidiness issue.

Remediation

The stored codes change meaning, so this is a wire-format break, not an in-place
fix:

1. Swap the two paths: keys take sign vector → unnormalized FWHT →
   1/sqrt(head_size) → per-group FP16 RMS → centroid table with midpoint
   boundaries, lower centroid winning ties. Values take per-group uniform
   quantization with FP16 scale and zero, decode (code + zero) * scale.
2. Add format_version to cache metadata and refuse a version the
   implementation does not support. Version 1 caches cannot be migrated in place.
3. Add the unrotated FP8 key path (key_bits == 8) and norm correction as
   declared metadata, per version 2 of the contract.
4. Update the correctness oracle — it currently encodes the inverted assignment,
   so the existing tests pass against the wrong behaviour.

Note that QJL is deliberately out of contract in version 2; independent
evaluations report it amplifies variance through softmax. Do not add it as part
of this fix.

Affected files in this backend

• kernels/quantization/turboquant/variants/rocm_cdna3/turboquant.cuh —
  tq_encode applies tg_forward_fwht + centroid selection to the value and
  the per-group fp16 scale/zero chain to the key.
• kernels/quantization/followups/variants/rocm_cdna3/turboquant.cuh — a
  byte-identical copy of the above (diff is empty). Both need the same fix, or
  the duplication should be resolved first.
• kernels/serving/phase2_quant_decode/variants/rocm_cdna3/phase2_quant_decode.cu
  — paged_attention_turboquant_kernel reads the inverted layout: it dequantizes
  keys through (code + zero) * scale and accumulates values in the rotated
  domain with one deferred inverse FWHT. Under the corrected contract the
  deferred-transform trick moves to the key/score side and values decode
  directly into the weighted sum.

Two further defects worth fixing in the same pass

1. Scale-plane dtype seam. tq_encode writes __half scale/zero caches
   (turboquant.cuh:218-219) but paged_attention_turboquant_kernel reads
   const float* (phase2_quant_decode.cu:687-688). The two ops are not
   directly composable today; a conversion step is required between them.
   Version 2 of the contract makes scale precision required metadata, so this
   should be declared rather than implied.
2. Signed sub-8-bit decode is missing. turbo_dequant_key sign-extends only
   when bits == 8; for signed widths below 8 it treats the code as unsigned
   (phase2_quant_decode.cu:664-667). The host reference in
   followups_test.cu:90 does sign-extend, but no test exercises that
   combination — every signed case uses 8 bits, so the divergence is untested.

Registration gaps

• turboquant_query (kernels/quantization/turboquant_query/) has no entry in
  .quixicore/kernels.yaml and no canonical OperationId. It is an orphan
  directory with its own Makefile.
• The registered turboquant entry is a bundle (status: imported, covering
  FWHT rotate + permute_cols + moe_lora_align + mixed-format primitives), not the
  codec. Actual tq_encode coverage lives under kernels/quantization/followups/.
• .quixicore/quant-formats.yaml has no turboquant entry at all; tracked in
  docs/capability-gaps.md:124.