Parity/coherence failure protocol

The model runs without errors but output is wrong. Scalar ops.print taps and recompile loops hide directional bugs and burn GPU time. Build a per-layer tensor-dump comparator first; every later check becomes a numpy read from disk.

Use this skill when MAX output disagrees with a PyTorch reference you can run and hook. The primary case is a custom-architecture port that serves but fails parity or coherence checks; the same protocol covers a quantized variant of a working port, a multi-GPU conversion of a working single-GPU port, and a regression after a MAX upgrade — anywhere a trusted reference exists.

Do not use this skill when:

• The server crashes on load → fix config, weights, graph (import-model)
• You have not finished implementing the graph → import-model Phase 2
• An already-verified model needs logit-comparison tolerances tuned → that is threshold calibration, not corruption

References

For MAX's built-in runtime debugging options (NaN checks, source tracebacks, op logging), see the MAX debugging tools. max.nn.hooks.PrintHook (covered in comparator-build.md) prints layer inputs and outputs for quick triage.

Protocol

Step 0: Sanity-check HF

Run model.generate(...) on the same HF repo, prompt, and checkpoint. If HF is incoherent, fix tokenizer/chat-template first; the MAX graph is not the problem.

pixi run python -c "
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained('<repo>')
model = AutoModelForCausalLM.from_pretrained('<repo>', torch_dtype=torch.bfloat16, device_map='auto')
text = tok.apply_chat_template([{'role':'user','content':'Hello!'}], tokenize=False, add_generation_prompt=True)
out = model.generate(**tok(text, return_tensors='pt').to(model.device), max_new_tokens=32, do_sample=False)
print(tok.decode(out[0]))
"

Step 1: Build the comparator

Follow comparator-build.md. You need three artifacts: HF dumper, MAX dumper (graph edits + standalone runner), comparator script. Cast dump tensors to FP32 in the MAX graph.

Guard: validate the dumpers before trusting them. Run both dumpers on a model MAX already serves correctly (any registered Llama works). Expect cos ≈ 0.999 at every layer, identical prompt_tokens.npy on both sides, and post_embed cos = 1.0. Anything less means the dumpers are broken — fix them before reading anything into a comparison on your port.

Step 2: Read comparator output, then branch

Follow comparator-output-patterns.md. Check false cliffs (wrong hidden_states indexing, missing attention_mask on decode-prefix dumps) before bisecting the graph.

The first trustworthy comparator run is a fork, not a checkpoint:

• Some layer diverges → graph hunt; continue with Steps 3 to 5.
• Every layer matches (cos ≥ 0.99) but generation still diverges → the graph is likely correct. Skip to Step 6; do not bisect layers.
• Pattern matches a false-cliff signature → fix the dumper, re-dump, re-read. Do not debug the graph against a broken comparator.

Compute per-token and per-dim cosine slices when global cos looks ambiguous:

cos_per_token = [cos(h[t], m[t]) for t in range(h.shape[0])]
cos_per_dim = [cos(h[:, d], m[:, d]) for d in range(h.shape[1])]

High max_diff where HF spikes and MAX is flat usually means HF formed an attention anchor your port did not, not "MAX exploding."

Step 3: Dispatch investigation agents

Follow agent-workflow.md. One lead agent analyzes dumps and ranks hypotheses with tensor evidence. Helpers run in parallel (weight stats, code diff, kernel inspection, sub-tap prep). Do not dispatch fix-attempt agents until the lead localizes.

Step 4: Verify numerically before recompiling

For each hypothesis: read dump tensors, compute what the fix would produce, compare to HF. Match → recompile. No match → next hypothesis.

Step 5: Apply fix, re-dump, re-compare

One compile, one smoke, full comparator pass (cos > 0.99 all layers). If verification still fails, see stacked-failures.md.

Step 6: Serve vs pipeline

When teacher-forced dumps at decode step K match HF but generated text diverges, the graph is likely correct. Bisect before re-bisecting layers:

Build if missing: pipeline decode compare, incremental layer dump, serve compare scripts. If teacher-forced and pipeline both pass but serve fails, do not edit the graph.