Skip to content

Commit 00bfd4e

Browse files
aselimcclaude
andcommitted
Fix the five merge blockers and the minors from the 7.2 review
MAJOR-1 (crash). jepa_ac_context_update/_trim allocated a fresh ggml_view_3d into a pool sized for four tensor overheads, so the third _update walked off the end and aborted inside ggml_new_object() — in Release, with no error path. The documented receding-horizon loop was exactly what crashed. The handle now creates the view once, immediately after the backend buffer exists, and afterwards only moves its extents (ne[1], nb[2], nb[3]); nothing allocates per frame. MANDATORY regression, tests/test-predictor.cpp::run_ac_context_loop: six _updates (1 -> 6 frames) then a _trim, each step checked against jepa_ac_rollout_ex over the same frames — max|d| 0.000e+00 at every step, bit-identical, on a handle whose capacity is 32. tests/forge-gguf.h raises the forged jepa.pred.n_frames 3 -> 5 so the capacity guard no longer masks the object budget in test-errors either. MAJOR-2. gen_benchmarks_md.py keyed the GPU->CPU cross-reference on the raw jepa-bench shape string while runs_from_results_json feeds back the rendered label, so the join missed for video models at one frame. Both sides now key on shape_label(); vjepa2_1-vitb-384 f16 and q4_k have their cpu_f16_t32_ms and speedup cells back (60.335 ms, 13.77x and 17.52x). MAJOR-3. results.svg and results-latency.svg were stale after benchmarks.json grew; regenerated with --split. All four figures --check green. MAJOR-4 (published false claim). vjepa2-ac-vitg-q4_k.gguf does NOT report q4_0: it is file_type 12 (q4_k), 121 q4_0 / 137 q4_K, because the AC predictor's 24 blocks are 1024-wide and keep the K-quant. Only the encoder-only ViT-g file goes the other way (158/52, file_type 2). docs/quantization.md now tabulates both and says why they differ; both hub cards corrected. MAJOR-5. python bindings: added jepa_ac_cem_default_params and jepa_ac_plan plus the jepa_ac_cem_params ctypes struct, and restored header order (the new binds had been appended). 92 declared prototypes, 92 bound, none extra. Full python suite green against a wheel built from this branch: 52 passed, 1 skipped. MINORs. 1: the churn accounting was wrong — measured against main it is 26 document cells (not 16), each at most one unit in the last printed digit, and tokens_per_s in the JSON moves by up to 1.5 at its 0.1 resolution; the bound describes the rendered document, not the JSON, and CHANGELOG + the generator now say so. 2/3: the planning table is re-taken FROM tests/results/benchmarks-gpu.json (the grid gained a K:H syntax so H=4 rows are measured and recorded rather than quoted ad hoc), and cached-vs-explicit is re-measured back-to-back in ONE session: +0.11 / +0.07 / +0.14 / -0.18 / -0.20 % across K=16..256 and H=2..4 — plus or minus 0.2 %, straddling zero. The earlier "0.3-1.3 %" compared two sessions; the conclusion "not a speed-up" survives and is now properly evidenced, and every mention of the K=64 row traces to the artifact. 4: dead parity anchor fixed. 5: 118 -> 126 cases. 6: the rollout no longer allocates a tail buffer at horizon 1 (144 MiB at the CEM defaults). 7: generated_from_dir is relative again and the GPU artifact's ambiguous top-level device_index is dropped — the per-row device field is authoritative and the protocol block says so. NITs. 1: the 11 611.1 MiB figure is quoted as the library prints it. 2: the encoder-cosine sentence is scoped, and the global max |d| (4.18e-2 on enc.blk.31.attn_qkv.bias, values to 8.2) is given. My reproducible count below 0.999 is 38, not 39 — nothing sits within 4e-6 of the boundary, so it is not a rounding artefact; the doc now states the mapping method so the difference can be resolved. 3: the CEM label says "final <= initial", which is what is checked. 4: rollout failures name the entry point the caller used. 5: TODO left for a strided jepa_ac_energy. Also: jepa.h notes that axis={} and close_gripper=None are reference defaults the struct cannot express. ctest 16/16 on build/ and build-cuda/, test-errors 126/0, python 52 passed, all five --check generators green, mkdocs strict, README 3 tables. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzdJ2iSxVUKxBkkJ1By1nm
1 parent d53ab48 commit 00bfd4e

19 files changed

Lines changed: 2608 additions & 1353 deletions

CHANGELOG.md

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ across releases.
3535
`jepa_ac_context` holds the observed frames' latents on the compute device and the AC graph takes
3636
its context as a shared prefix plus a per-candidate tail, so the observed frames are neither
3737
replicated across the K candidates nor re-uploaded per step (`jepa_ac_context_new` / `_update` /
38-
`_trim` / `_free`, `jepa_ac_rollout_cached`). Measured, it is worth almost nothing in time —
39-
786.11 vs 787.32 ms cached against explicit at K = 64 — and everything in shape: one object across
40-
CEM iterations and receding-horizon steps. `jepa_ac_plan` is Meta's own `mpc_utils.py::cem` over
38+
`_trim` / `_free`, `jepa_ac_rollout_cached`). Measured back-to-back in one session, it is worth
39+
**nothing** in time — cached against explicit is ±0.2 % and straddles zero across K = 16…256 and
40+
H = 2…4 — and everything in shape: one object across CEM iterations and receding-horizon steps. `jepa_ac_plan` is Meta's own `mpc_utils.py::cem` over
4141
it, and replaying the random draws their loop made it returns the same plan to **2.98e-08**.
4242
`jepa-worldmodel --plan` and `jepa-bench --mode ac-plan` expose both.
4343
- `jepa_ac_rollout_ex`, which takes the actions between the observed frames — the thing
@@ -53,9 +53,15 @@ across releases.
5353
Nothing changes for the existing single-sequence callers.
5454
- `scripts/gen_benchmarks_md.py --merge-json` seeds the CPU rows from the committed artifact and
5555
overlays a sweep on top, so re-measuring part of the grid adds rows instead of dropping every other
56-
model's — the CPU half now has the artifact fallback the GPU half always had.
56+
model's — the CPU half now has the artifact fallback the GPU half always had. Session provenance is
57+
concatenated, not replaced. The artifact also stores raw `weight_bytes` / `peak_rss_bytes` /
58+
`ms_mean_raw` / `load_ms_raw` so a future rebuild is exact; rows written before this change carry
59+
only the rounded fields, and rebuilding those moved **26 numeric cells of `docs/benchmarks.md`, each
60+
by at most one unit in its last printed digit**, and `tokens_per_s` in
61+
`tests/results/benchmarks.json` **by up to 1.5** at its 0.1 resolution — the "one unit in the last
62+
printed digit" bound describes the rendered document, not the JSON. Nothing was re-measured.
5763
- Eight new loader checks for the AC keys and nine new `jepa_ac_*` argument guards in `test-errors`
58-
(118 cases, up from 101), forged from a complete tiny AC model.
64+
(126 cases, up from 101), forged from a complete tiny AC model.
5965

6066
## [0.1.1] — 2026-09-01
6167

docs/api.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,11 @@ typedef struct {
362362
uint32_t seed; // RNG seed, used only when `noise` is NULL
363363
} jepa_ac_cem_params;
364364

365-
// Every field at the reference's default (mpc_utils.py's own signature defaults).
365+
// Every field this struct HAS, at the reference's default (mpc_utils.py's own signature defaults).
366+
// Two of cem()'s parameters have no field here because the released planner never varies them:
367+
// `axis={}` (pin a sampled dimension to a constant) and `close_gripper=None` (force the gripper shut
368+
// from step h onwards). Both are reference defaults, not omissions from the algorithm; a caller that
369+
// needs them can pin the corresponding lanes in `noise` or post-process the plan.
366370
jepa_ac_cem_params jepa_ac_cem_default_params(void);
367371

368372
// Plan against `goal` ([tokens_per_frame, enc_dim], normalised like the handle's latents).

docs/architecture.md

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -333,22 +333,24 @@ per step, and `pred.embed` runs on them once per graph instead of K times. `jepa
333333
appends a newly observed frame (the receding-horizon step) and `jepa_ac_context_trim` slides the
334334
window; capacity is `jepa.pred.n_frames`, allocated once.
335335

336-
**What that is worth, measured: almost nothing in time, and that is the honest answer.** Cached
337-
against explicit context on CUDA1 at f16 is 786.11 vs 787.32 ms (K = 64, H = 2) and 318.02 vs 317.79
338-
ms (K = 16, H = 4) — inside run-to-run noise. The shared-prefix broadcast against a fully replicated
339-
context is 527.95 vs 529.43 ms on the GPU and 2181 vs 2200 ms on the CPU at 16 threads, 0.3–0.9 %.
340-
The reason is arithmetic: `pred.embed` is one [1408 → 1024] matmul over 256 rows, against 24 blocks
341-
of 1024-d attention and FFN over K × 258 rows, and the upload is 1.44 MB against a half-second graph.
342-
The handle earns its place as an **API**, not as a speed-up — one object held across CEM iterations
343-
and receding-horizon steps, one device allocation for the observed frames — and the parity suite
344-
gates it as bit-identical to the explicit path.
336+
**What that is worth, measured: nothing in time, and that is the honest answer.** Back-to-back in one
337+
session on CUDA1 at f16, cached against explicit context is +0.11 % at K = 16 / H = 2, +0.07 % at
338+
K = 64, +0.14 % at K = 256, −0.18 % at K = 16 / H = 4 and −0.20 % at K = 64 / H = 4: **±0.2 %,
339+
straddling zero**. The shared-prefix broadcast against a fully replicated context is the same story
340+
(527.95 vs 529.43 ms on the GPU, 2181 vs 2200 ms on the CPU at 16 threads). The reason is arithmetic:
341+
`pred.embed` is one [1408 → 1024] matmul over 256 rows, against 24 blocks of 1024-d attention and FFN
342+
over K × 258 rows, and the upload is 1.44 MB against a half-second graph. The handle earns its place
343+
as an **API**, not as a speed-up — one object held across CEM iterations and receding-horizon steps,
344+
one device allocation for the observed frames — and the parity suite gates it as bit-identical to the
345+
explicit path, including a six-update receding-horizon loop
346+
([performance](performance.md#planning-what-a-cem-decision-costs) has the table).
345347

346348
The planner on top of it, `jepa_ac_plan`, is a port of Meta's `mpc_utils.py::cem`. Two details of
347349
that loop decide whether an implementation is right: only four of the seven action dimensions are
348350
sampled (translation and the gripper; rotation is hard zeros), and `selected.std(0)` is torch's
349351
**unbiased** estimator, a factor 1.155 at topk = 4 applied every iteration. Replaying the random
350352
draws their loop made, jepa.cpp returns the same plan to max |Δ| 2.98e-08
351-
([parity](parity.md#v-jepa-2-ac-jepapredkind--ac--the-action-conditioned-world-model)).
353+
([parity](parity.md#v-jepa-2-ac-jepapredkind-ac-the-action-conditioned-world-model)).
352354

353355
## Robustness
354356

0 commit comments

Comments
 (0)