fix(rocm): explain the granite-4.0-h-tiny w1 perplexity gap against Metal #2154

Part of #1801

Problem / Background

In the ROCm correctness matrix (docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md:51,104, PR #2083), granite-4.0-h-tiny at w1 passes the top-1 gate (0 / 128 disagreements, 0 / 112 on decided positions) but its perplexity is +4.345% against Metal. The document measured a mean per-position NLL increase of 0.042 nats with a standard error of 0.017 (about 2.5 SE) and left it open as "weak evidence of a small bias". The other granite widths are within noise (w8 +0.302%, w256 +0.349%), and every other pair stays inside +/- 0.9%. The gap is unchanged after the SSM port: the 96cbce84 ROCm trace (PR #2099) gives the same +4.345%.

Current Behavior

w1 = logit_trace MODEL CORPUS 1 128 8 0 (benchmarks/logit_traces/rocm_gfx1151_c5fe9a16/README.md:17): 128 one-token forwards, each into a fresh cache with no context and no BOS (#1785), so every Mamba2 layer runs the ssm_step graph with no state on both backends (README line 53), attention sees one key, and every projection is an M = 1 quantized matvec. w8 and w256 take the multi-token paths.

Extra analysis of the committed traces (no GPU needed), Metal metal_m5_d1128266 against ROCm rocm_gfx1151_96cbce84, over the top-8 ids both sides list:

• Mean logit shift: +0.011 (SE 0.015), i.e. none.
• Top-1 minus top-8 logit spread: ROCm wider by 0.071 (SE 0.019, about 3.7 SE). The same statistic is +0.013 (SE 0.007) at w8 and -0.029 (SE 0.010) at w256.
• The target is in the top 8 at only 5 of 128 positions (mean Metal NLL 11.76 nats), so the NLL is set by where the target sits in the tail.

So the lead is a slightly sharper (wider-spread) logit distribution on ROCm in the single-token shape only, which raises NLL for low-probability targets without moving top-1. Granite scales logits by 1 / logits_scaling and the residual stream by embedding_multiplier and residual_multiplier (src/models/granitemoehybrid.rs:27-36), so a small per-layer difference can be amplified.

Proposed Solution

An investigation with a fixed decision rule, in this order:

1. Confirm or dismiss with more positions. Trace w1 with 1024 chunks on both backends from one commit: logit_trace models/mlx/granite-4.0-h-tiny-4bit tests/fixtures/wikitext2_excerpt.txt 1 1024 8 0. Compute the mean per-position NLL shift and the spread statistic above with their SEs (a short script reading both TSVs through load() in scripts/compare_logit_traces.py). If both are under 2 SE, close as noise: add a paragraph to the results page with the numbers and commit the traces under benchmarks/logit_traces/ with METADATA.txt, RUNS.txt and SHA256SUMS like the existing directories.
2. Find the outlier backend. If the gap holds, trace the same shape with MLXCEL_DEVICE=cpu on the ROCm host. The backend farther from the CPU trace (same two statistics) is the one to bisect.
3. Localize the op. The w1 path differs from w8 in exactly three places: M = 1 quantized matvec (dense projections and the MoE gather_qmm), ssm_step with an empty state, and single-key attention. Compare each against the CPU on the outlier backend with granite-4.0-h-tiny's real shapes and weights, and report the max and mean absolute error per op. Add the comparison as a test only if it exposes a defect.
4. Fix or document. If one op is out of line with its CPU result beyond its dtype's expected rounding, fix it in that op (an mlxcelverse overlay fix with a LOCAL_FIXES.md entry, if it is in the ROCm backend) and re-trace. If every op is within rounding and the gap is amplification of ordinary rounding by granite's multipliers, record that conclusion with the per-op numbers and close.

Scope

• In scope: the traces, the analysis, a results-page update, and a fix if step 4 finds a defective op.
• Out of scope: BOS handling in logit_trace (fix(examples): anchor logit_trace chunk 0 with the BOS token #1785) and the decided-position gate, which already passes.

Acceptance Criteria

• The 1024-position w1 comparison is published with mean NLL shift, spread difference and their SEs. (Against an f32 reference, since no Metal host was available; PR docs(rocm): explain the granite-4.0-h-tiny w1 perplexity gap #2240.)
• Either the issue closes as noise under the step 1 rule, or the op responsible is named with CPU-referenced error numbers. (Step 4 outcome: every op within bf16 rounding of an f32 reference; the CPU stream is unusable for bf16 here; PR docs(rocm): explain the granite-4.0-h-tiny w1 perplexity gap #2240.)
• (Not applicable, no defect found) If a defect was fixed: the re-traced w1 pair is under 2 SE on both statistics, w8 and w256 still pass --decided 2.0, and make verify-rocm passes.
• docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md no longer lists the gap as unexplained; it links the new evidence.

Verification

cargo build --release --features rocm --example logit_trace          # gfx1151
cargo build --release --features metal,accelerate --example logit_trace   # Metal
./target/release/examples/logit_trace models/mlx/granite-4.0-h-tiny-4bit tests/fixtures/wikitext2_excerpt.txt 1 1024 8 0 > w1_1024.tsv
MLXCEL_DEVICE=cpu ./target/release/examples/logit_trace models/mlx/granite-4.0-h-tiny-4bit tests/fixtures/wikitext2_excerpt.txt 1 1024 8 0 > w1_1024_cpu.tsv
python3 scripts/compare_logit_traces.py metal_w1_1024.tsv rocm_w1_1024.tsv --decided 2.0

Run GPU traces inside scripts/rocm_gpu_guard.sh on the shared gfx1151 host. Related: #1809, #2083, #2099, #1785.