Files
2026-10-20 12:09:42 +08:00
..
2026-10-20 12:09:42 +08:00
2026-10-20 12:09:42 +08:00
2026-10-20 12:09:42 +08:00

YoCo mask/compact equivalence verification

examples/yoco/yoco_equiv.py switches DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE within the current repository to verify that the mask and compact modes produce equivalent output.

Level Subcommand Cost What it verifies
L1 mask-compact Single device, runs in seconds, no weights loaded Compares SWA pivots, visible keys, store slots, and compressed key ranges
L2 logp + compare 8 devices, a compatible V4.1 HF checkpoint Uses a fixed teacher-forcing trajectory and compares logprob and top-K element by element

L1 metadata

# Compare in memory, writing no file
python examples/yoco/yoco_equiv.py mask-compact

# Optional arguments
python examples/yoco/yoco_equiv.py mask-compact --keep-tokens 140 --window 128

L2 end-to-end output

# Run mask and then compact, and compare automatically
examples/yoco/run_yoco_equiv_l2.sh --model_name /path/to/compatible-v41-hf

# The arguments are passed to both logp runs
examples/yoco/run_yoco_equiv_l2.sh --model_name /path/to/compatible-v41-hf --ep 8 --cp 1

# DFlash+YoCo ckpt: --mtp N enables MTP and additionally compares the confidence score bitwise (a direct probe of the DFlash main cache)
examples/yoco/run_yoco_equiv_l2.sh --model_name /path/to/compatible-v41-hf --mtp 5 --ep 8 --cp 1

L2 dumps are written to var/yoco_equiv.*/logp_{mask,compact}.rank*.json.

Notes and limitations

  • The .so compiled from the current repository must be used; when PrepareArgs.yoco_prefill_truncation_mode is missing, the script asks you to rebuild.
  • Must be run from the repository root.
  • The model must be a compatible V4.1 HF YoCo checkpoint (n_prefill_layer < n_layer); the script validates this constraint at runtime.
  • Apart from DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE, every DEEPINFER_* environment variable must be identical across the two L2 runs.
  • MTP is supported only with the DFlash head (--mtp N, which requires a DFlash+YoCo HF checkpoint). The confidence is attached to the last token emitted in each round and lags one round, so the first round is None, and the confidence produced by the last forward does not reach the output for comparison.