Allow CI tests to test all the pytests

[KEHANG]

Local Pytests

Note

Use mrt models init to download the necessary resources. Then use mrt checkpoints download to download mrt2_small.safetensors and mrt2_base.safetensors

I ran

mrt jax generate --model=mrt2_small
mrt mlx generate --model=mrt2_small
mrt mlx generate --model=mrt2_small --no-mlxfn --bits=8

pytest -s tests/test_musiccoca.py
pytest -s tests/test_prefill_correctness.py

python scripts/generate_test_reference.py
pytest -s tests/test_bitlevel_parity.py

and observed the following output:

all passed.

Benchmark Regression Test

I ran

python scripts/bench_track.py
python scripts/bench_show.py --samples

and observed the following output:

2026-06-07T11:28:13 mrt2_small_int8_rvq12_cfgs0_260607_ca3b8e9      ca3b8e9    11.3   11.3   11.5   11.6   88.2  ✅ WITHIN BUDGET
                      ↳ [M4 Pro]
                      L: [-0.01214600, -0.01123047, -0.01174927, -0.01290894, -0.01229858, -0.01040649, -0.01046753, -0.01254272, -0.01406860, -0.01489258]
                      R: [-0.01068115, -0.00979614, -0.01040649, -0.01205444, -0.01242065, -0.01083374, -0.01028442, -0.01165771, -0.01266479, -0.01324463]
  2026-06-09T13:45:45 mrt2_small_int8_rvq12_cfgs0_260609_ba118ee      ba118ee    11.4   11.2   12.2   13.1   88.1  ✅ WITHIN BUDGET
                      ↳ [M4 Pro]
                      L: [-0.01214600, -0.01123047, -0.01174927, -0.01290894, -0.01229858, -0.01040649, -0.01046753, -0.01254272, -0.01406860, -0.01489258]
                      R: [-0.01068115, -0.00979614, -0.01040649, -0.01205444, -0.01242065, -0.01083374, -0.01028442, -0.01165771, -0.01266479, -0.01324463]