Request: publish TriAttention calibration script (TRIA v2/v3 stats.bin producer)

[minhtcai]

Summary

The feature/triattention-scoring branch wires up the --triattention <path> flag and ships a complete loader (src/triattention.c::tria_load, magic 0x54524941, versions 1/2/3 supported), but the repository does not include a calibration script that produces the stats.bin file the loader expects.

The README/TURBOQUANT.md reports results like "Qwen3.5-27B turbo3 + TriAtt 75%" — so a calibration pipeline clearly exists internally. Could it be published?

What's documented

src/triattention-integration.md describes the algorithm (CPU-side scoring every 128 tokens, GQA aggregation, z-norm, top-K). The header documents the on-disk schema (layer count, kv-head count, fc, head_dim, optional layer_budget_scales, per-head MRL and non-rot data). Sufficient to reverse-engineer the format, but not the calibration math.

Why it matters

Without the producer, third parties (including users on non-AMD hardware trying to validate the approach — sm_75 / sm_72 in our case) can build the binary but not exercise the TriAttention path end-to-end. We end up only able to compare TurboQuant variants.

Asks (in order of preference)

1. Publish the script as-is (even if it's research-grade / undocumented).
2. Document the math sufficient for a third-party reimplementation: how Q/K concentration centers are computed, what calibration corpus you used, MRL bin definitions, what nonrot_alpha modulates.
3. Confirm format details for a community reimplementation: byte order, per-head record layout, how layer_budget_scales is computed, version differences (v1 vs v2 vs v3).

Context

Trying to evaluate feasibility of the TriAttention + TurboQuant path on NVIDIA Volta-class hardware (Jetson Xavier AGX, sm_72) for an LLM cluster. Happy to share results back if we get it working on Turing/Volta.

Thanks for the work — the turbo3 + TriAtt numbers are striking.

[minhtcai]

Working calibration tool: llama-triattention-calibrate

I built a self-contained calibration tool as a llama.cpp subcommand that produces a TRIA v2 stats.bin accepted by tria_load. Verified end-to-end on Xavier AGX (sm_72) with Qwen3.5-0.8B + turbo3 KV.

How it works

Hooks the Qcur_normed (post-norm, pre-RoPE) tensor at every attention layer via llama_set_callback_eval, accumulates per-(layer × head × freq) statistics in RoPE complex space:

• q_mean_real[f] = E[ Re(q_complex[f]) ]
• q_mean_imag[f] = E[ Im(q_complex[f]) ]
• q_abs_mean[f] = E[ |q_complex[f]| ]
• mrl[f] = stub zeros (loader skips this field via fseek(fp, fc*4, SEEK_CUR))

RoPE complex pairing matches the model's rope_type:

• LLAMA_ROPE_TYPE_NEOX: q_c[f] = q[f] + i*q[f+head_dim/2]
• LLAMA_ROPE_TYPE_NORM / IMROPE / MROPE: q_c[f] = q[2f] + i*q[2f+1]

Layer budget scales default to 1.0 (uniform).

Validation

tria_load: OK — 24 layers, 8 heads, 2 kv_heads, fc=128 (v2)
TriAttention runtime initialized (budget=50% window=256 interval=128 sink=0)
tria_score: n_kv=2102 n_old=1846 budget=1846 (pct) new=1846 mode=full (pass 0)
tria_score: value-aware boost applied to 850/1846 tokens (lambda=0.25, vstd=0.0445)

Quality A/B on a comprehension prompt (Qwen3.5-0.8B, 8K ctx, turbo3):

Setup	TG (t/s)	Output
turbo3 baseline	47.41	"A photon is a quantum of light with zero rest mass…"
turbo3 + TriAtt budget=50%	47.31	Longer with reasoning trace
turbo3 + TriAtt budget=25%	36.26	"A photon is a quantum particle of light, possessing zero rest mass…"

All three coherent. Pruning to 25% budget preserves answer quality.

Code

The whole tool is ~280 lines (tools/triattention-calibrate/calibrate.cpp). Plus a 4-line CMakeLists.txt and a 1-line addition to tools/CMakeLists.txt.

Usage:

llama-triattention-calibrate \
    -m model.gguf \
    -f calib.txt \
    -o stats.bin \
    -c 2048 -b 2048 \
    --n-tokens 16384 \
    -ngl 99

Calibration corpus used: ~500 KB plain text (Alice in Wonderland + Tom Sawyer + Moby Dick). 8K tokens (4× 2048-token chunks with KV reset between) was enough to produce stable Q means on the 0.8B model. Larger models likely want more.

Happy to submit a PR if useful — could you confirm the math is right and any ask about the layer_budget_scales / mrl fields you'd want populated differently? Specifically:

1. Is mrl[f] supposed to be the mean of complex magnitude radial length at offset f? The current loader skips it (fseek(fp, fc*4, SEEK_CUR)), so I'm writing zeros. If runtime ever reads it, what should it be?
2. Layer budget scales — you presumably profile per-layer attention sparsity offline. What's the recipe? Is it just relative attention entropy normalized, or something more?
3. Calibration corpus size — what did you use for the published Qwen3.5-27B turbo3 + TriAtt 75% number? Curious whether 1K, 10K, 1M tokens.

Even at 50% budget the small-model numbers above show the path works — just want to align before scaling up to the 35B prod model.

[jasonriedy]

IIRC, https://github.com/domvox/triattention-ggml is the author's calibration code.