Repository navigation
Phase 8.2 WP2 (Java): instrumentation, host observations and the FFM calibration - #90
Merged
Merged
Conversation
Move BindingCounterAgent from jmlx-core's test source to testFixtures so jmlx-models' baseline worker classpath (testImplementation testFixtures) resolves the instrumented entry point, and generalize it for the decode baseline diagnostics: - property prefix jmlx.qwen35.binding.* -> jmlx.binding.* - record each instrumented downcall's JVM descriptor (MethodModel methodType) so consumers can group counts by signature class without re-reading the generated bytecode; new signatures()/failureReason() - rename the qwen35BindingAgent Jar task to bindingAgent (testFixtures output, binding-agent classifier); deltaBenchmark and qwen35DecodeCounter follow - new opt-in bindingCoverage task + gate test: one counted window that allocates (mlx_array_new_data), evaluates a matmul (mlx_matmul) and frees through scope close (mlx_array_free) -- the counters must cover the direct allocation/free paths as well as the op facades
…13.3 contract The WP1 adapters deviated from the frozen boundary semantics (findings section 13.3: 'the prompt cache is created before the timer starts' and 'the per-step s stops at the evaluated logits'), a bug not covered by the section 14.1 intentional-deviation list: - prefill_s included the per-request freshCaches creation: t0 was taken before the caches, while the pinned Python loops (controlled.py run / run_fixed_continuation) create the prompt cache before t0. t0 now sits after freshCaches in both runControlled and runFixedContinuation; it is also the event-time origin, matching the Python token_rel_s origin. step_region_s is unchanged: it is t_end - t_after_prefill on both definitions. - the per-step s was computed after the argmax (and, for fixed continuation, after the full-vocabulary readback and sort), while the Python loops stop the step timer immediately after the step evaluation. The elapsed is now captured right after evalStep, before any host-side selection work, and the misleading comments claiming the opposite are corrected. The section 14.6 session 'd' local-evidence numbers were measured under the old Java semantics (recorded in the WP2 findings section).
Adds the diagnostic (instrumented) run shape on top of the WP1 worker; headline job files and reports stay byte-identical (the job field is absent by default): - BaselineAdapters gains the package-private PhaseObserver seam (window(phase, position, work) + ModelWork + NONE): the controlled and fixed-continuation loops wrap prefill-construction / prefill-sync / prefill-selection (position 0) and per cached step construction / sync / selection at promptLength + k + 1 in observer windows; the stock adapter gets one whole-request window (the documented controlled-first limitation, plan section 5.1). Windows never change the work, the scope lifetimes or the frozen section 13.3 timer boundaries; the no-observer overloads are unchanged. - BaselineWorker gains the job field instrumentation (none | binding-count | tracing) with the diagnostics gates before model load: binding-count requires the installed jmlx-core BindingCounterAgent and a clean failureReason, tracing requires the thread CPU time API. Intentional instrumentation failure fails the run with a failed report, never counts silently (plan section 5.1). - identity.instrumented is set for instrumented workers (the launcher excludes them from headline aggregation, frozen section 13.3) and carries a per-worker host block (hw model, OS, disk free, load average, vm_stat, top processes) that the frozen Python workers never carry. - Per-request binding_counts (phases + total over the phase windows, non-empty-total gate) in binding-count mode, per-request tracing (request wall/CPU spans, per-window wall/CPU, and the unassigned residual = span - sum of windows) in tracing mode; the report's worker block carries the agent-recorded binding_signatures for the FFM transition estimate. - The section 10 memory block is filled for instrumented workers: active/peak active/cached MLX allocator bytes (peak reset per request so the load is excluded) plus rss_peak_bytes, the peak of the after-request RSS samples (a lower bound on the true peak). Protocol tests: the headline report now pins instrumentation 'none'; binding-count without the -javaagent fails the run with a failed report naming the binding-agent requirement; a tracing run asserts the exact 24-window phase/position sequence, the residual identity against the request span, and the populated memory block.
…e isolation Extends the phase 8.2 baseline launcher for the instrumented (diagnostic) worker shape; headline specs and sessions keep their WP1 form (the spec field is optional, the job file field is written only when not none, so headline job files stay byte-identical): - BaselineJob gains the instrumentation component (none | binding-count | tracing), an -<instrumentation> suffix on the cell directory name so a diagnostic run never collides with the headline cell it measures, and two launcher-side components: the binding-agent jar (the fork command adds -javaagent:<jar> for java + binding-count) and the session's isolated native-cache directory (every forked Java worker gets -Djmlx.native.cache.path=<out_root>/native-cache; the user's default extraction cache is never touched). - The spec gains the optional per-worker instrumentation field (python workers are never instrumented: the frozen Python workers carry no diagnostics) and the session-level orchestration field (spec-order, the default, or alternating, which requires consecutive workers to use different runtimes). A binding-count java worker without the jmlx.baseline.binding.agent.jar system property is a spec failure — an instrumented worker that cannot count fails loudly, never silently (plan section 5.1). benchmarkBaseline sets the property to the :jmlx-core:bindingAgent archive and depends on the task. - The session block records orchestration and the native-cache path; fork_order now carries the full cell identifiers (worker directory names) in fork order instead of bare runtimes. - Frozen section 13.3 implemented in the launcher: an ok report whose identity is instrumented is excluded from headline aggregation by construction, with the exclusion recorded (reason 'instrumented'); jobViolations cross-checks the report's instrumentation against the launched job. - The task-owned isolated native-cache directory is retained for a successful session (session evidence) and deleted, best-effort, on any session failure. - BaselineReportValidator: the WP1 'memory is null' rule becomes instrumented-aware (the four section 10 counters, exactly, for instrumented requests; null otherwise), and the new per-request binding_counts / tracing blocks, the worker-level instrumentation / binding_signatures fields and the per-worker identity.host record are validated structurally, including the cross-mode absence rules. Launcher tests: fork_order cell identifiers, the orchestration validation (consecutive-same-runtime and unknown value refused) and a valid alternating java+python session (fake venv, stock-shaped python report), python instrumentation refused, binding-count without the agent jar refused, a full binding-count session (agent on the fork command, suffixed cell, job-file field, by-construction exclusion), native-cache retention/deletion, the job command's new flags, and validator fixtures/tests for the instrumented report shapes.
Adds the phase 8.2 WP2 FFM calibration (plan section 5.1): a calibrated estimate of the Project Panama downcall transition cost per representative shape of the generated mlx-c bindings, measured against no-op C functions in fresh JVMs. - src/main/resources/ffm_cal.c: one non-elidable no-op per representative downcall-shape class (the 16 dominant generated-wrapper shapes: (M)I (M)J (M)Z (M)V (M)M (AM)M (MM)I (MMM)I (MMMM)I (MMMMM)I (MI)I (MMIM)I (MMIIM)I (MZM)I (MMMJM)I (MMIMM)I). Each body folds its arguments into a volatile sink (an observable side effect, so -O2 cannot eliminate the call) and never dereferences an input. Carrier fidelity is documented in the source: every opaque mlx-c struct is a single 8-byte pointer, register-passed exactly like a void* on AArch64, and 8-byte struct returns use the same register convention. - FfmCalWorker: fresh-JVM worker. Binds each shape through Linker.nativeLinker().downcallHandle exactly like the generated holders and calls it through the generated wrappers' own invokeExact form (one dedicated typed call site per shape, including the SegmentAllocator first argument of struct-return shapes). One sample is the wall time of 'calls' consecutive downcalls, with the shape order rotated per iteration and an empty Java loop measured as the timing floor. Raw samples are written to JSON. - FfmCalibration: extracts the C source (a classpath resource) under the output directory, builds the dylib with cc -shared -O2 (clear error when cc is absent), records compiler/flags and source/binary SHA-256s in build-info.json, forks fresh JVMs (the Tf32Benchmark pattern), and pools the raw samples into per-shape p10/median/p90 with min/max and every raw sample retained (ffm-calibration.json). The aggregate step is a pure method over the fork outputs. - FfmTransitionEstimate: consumes a binding-count worker report (the agent binding_signatures plus the requests[].binding_counts phase windows) and a calibration summary. Normalizes the generated wrappers' JVM descriptors to the calibration shape keys and produces the per-decode- step budget (sum over each cached step's construction/sync/selection windows of count x calibrated nanos) with its p10/p90 range, the prefill and request-total budgets, the per-shape call/budget table, the coverage of classified calls, the uncovered signatures with their descriptors, and the explicit assumptions (no-op transition only; carrier fidelity; counts are attempted downcalls, not Metal dispatches; calibrated estimate, not a causal partition). - benchmarkFfmCalibration: opt-in JavaExec task in the benchmark group (output-directory [iterations calls forks warmup], defaults 200 2000 5 2000); not part of check. Tests: FfmCalibrationTest runs the full pipeline (extract, cc build, in-process worker run, aggregation) and asserts every shape's ordered quantiles sit above the empty-loop floor; it is skipped when cc is not on the PATH so ordinary check gains no compiler prerequisite. FfmTransitionEstimateTest is pure Java: descriptor normalization (including the unparsable and missing-descriptor cases) and hand-computed budget/coverage/uncovered values on synthetic binding-count reports. Local measurement baseline (M5 Max, Apple clang 21, this host, 2 forks x 20 iterations x 500 calls): empty-loop floor ~3 ns/iteration; (M)I median ~32 ns/call, 4-5 carrier shapes ~40-47 ns/call, the (AM)M struct-return shape ~107 ns/call (including the per-call result-segment allocation through the caller's allocator, as in the generated wrappers).
…ibration Freezes the phase 8.2 WP2 record (no code change): - req/plans/phase8-2-reference-findings.md: the frozen contract is amended in place — section 13.1 gains the optional job instrumentation field, section 13.3 gains the worker instrumentation echo, the iff-instrumented identity.host block, the binding_signatures field, the per-request binding_counts/tracing fields, and the WP2 memory rule (replacing the WP1 "stays null" sentence), and section 14.3 records the session orchestration / native_cache fields and the full-cell-identifier fork_order. Section 13.4 item 20 lists the amendments. The new section 15 records the WP2 work: the agent generalization, the WP1 boundary-bug fix (with the caveat that the section 14.6 session 'd' numbers ran under the old Java timer semantics), the instrumentation and host-observation implementation, the FFM calibration design with its carrier-fidelity argument and the local measurement baseline, and the test evidence. - jmlx-benchmarks/README.md: the decode-baseline section documents the spec orchestration / per-worker instrumentation fields, the -<instrumentation> cell suffix, the native-cache directory and the diagnostic report fields; a new section documents the opt-in benchmarkFfmCalibration task and the FfmTransitionEstimate main class.
…ibration fixes Review round 1 of PR #90 (findings 15.6): 1. BaselineAdapters records explicit prefill-release, per-step release and request-release windows (idempotent scope close in finally). 2. FfmTransitionEstimate aggregates per request: one sample per (measured request, position), aggregates restricted to kind == measured, new measured_requests field, refusal of non-controlled adapters and of reports without measured requests. 3. Duplicate cell keys now include the worker instrumentation. 4. Worker request bookkeeping takes wall/CPU clocks before the RSS sample. 5. windowObserver's binding-count boundaries attribute start/work/stop failures correctly (best-effort agent stop, original error rethrown). 6. FfmCalWorker warms constant MethodHandle holders, passes carriers by value (matching the pinned header), and the loop floor does real work. 7. All launcher subprocess waits are bounded; the binding agent jar path is validated to exist before worker launch. 8. asInt replaced by asLong for the 64-bit downcall counters. 9. Native-cache isolation documented as a guard, not a redirect. The FFM calibration was re-measured with the corrected worker; findings 15.4/15.6 record the new numbers and superseded source SHA, the estimate tests gained the two-measured-request and same-cell-different-instrumentation cases, and the protocol test pins the 33-window layout.
…tes, adapter exception paths
…ted descriptors, failure paths Review round 3 of PR #90 (findings 15.8): 1. FfmCalWorker: every sample gets a fresh CalScope (created before its timed batch, closed after the sample, outside the timed region), mirroring the production step-scope lifecycle, instead of one fork-lifetime set/arena holding 1.6M segments. 2. BaselineAdapters.releaseWindow catches Throwable and calls scope.close() directly when the window dies before the close runs (the close is idempotent); a fallback-close failure is attached to the release failure, never masking the in-flight original. 3. FfmTransitionEstimate accumulates every budget in double and rounds once per output value (a per-term whole-ns rounding biased a ~50-binding step by up to +/-25 ns). 4. BaselineWorker restricts the nativeBindingDescriptors lookups to the binding names with a nonzero count in some measured window's binding_counts.total (each lookup links the binding's downcall); a run that failed before any window completed falls back to every instrumented name, keeping the nonempty-descriptor invariant. 5. ffm_cal.c carrier-fidelity comment: the mlx_optional_* wrappers are value/has-value pairs padded to 8 bytes, not single pointer members (comment-only; the rebuilt dylib is bit-identical to the round-2 pin). 6. Direct tests: BaselineAdaptersTest +4 (the releaseWindow failure paths with a throwing PhaseObserver around a real scope close) and BaselineWorkerProtocolTest +3 (fake-agent fallback binding_descriptors, the counted-only restriction, a window work OOME propagating unwrapped with the recording stopped). Calibration re-measured: source SHA cb30e80b..., binary bit-identical to the round-2 pin (f1e2e7d7...); struct shapes now 52.2-55.1 ns/call (the ~10x drop from round 2 is the expected correction of finding 1's artifact); end-to-end per-step estimate 76.65/106.68 us, about 0.53% of the measured step wall time (was about 3.5%), and the counted-only descriptors reproduce the full-descriptor estimates exactly. Findings 15.3/15.4/15.5 amended in place; 15.8 records the round.
…property hygiene, Validated factory
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Phase 8.2 WP2 — the plan's §5.1/§6 work package: per-binding downcall counting, host observations and the FFM transition calibration for the decode-baseline workers (WP1 is PR #89; the frozen contract is findings §13, amended in place by this WP). Full record in
req/plans/phase8-2-reference-findings.mdsection 15.What lands
jmlx-coretestFixtures (never a published API):BindingCounterAgent— generalized out of the jmlx-models tests:jmlx.binding.*property prefix,signatures()(binding name → the generated static wrapper's JVM descriptor, recorded at class-load time),failureReason()(an instrumented run fails rather than counting silently), windowedstart()/stop().:jmlx-core:bindingAgent(builds the-javaagentjar) + opt-in:jmlx-core:bindingCoverage(assertsmlx_array_new_data >= 2,mlx_array_free >= 2,mlx_matmul == 1in one window). The generated bindings stay byte-identical — the agent transforms at class-load and edits nothing.jmlx-modelstest sources (benchmark-only, never published API):t0; Java startedt0beforefreshCaches, so the oldttft_sincluded prefill cache creation and the per-stepsincluded per-step cache creation.step_region_sandstreaming_decode_tpsare unchanged under both definitions. The frozen §13.3 wording was already correct: this was a Java bug, not a recorded deviation.BaselineAdapters.PhaseObserver— strictly additive phase windows (they never change the work, scope lifetimes or timer boundaries): controlled/fixed-continuation exposeprefill-construction/prefill-sync/prefill-selectionat position 0 plus per-stepconstruction/sync/selectionatprompt_length + k + 1; stock exposes onerequestwindow (a documented controlled-first limitation).BaselineWorkerinstrumentation: jobinstrumentation∈ none|binding-count|tracing, gated before model load (agent installed / thread CPU time API); per-workeridentity.hostblock (present iff instrumented); per-requestbinding_counts(with the non-empty-total gate),tracing(span/windows/residual) andmemory(instrumented requests only; the pre-requestresetPeak()excludes model load,rss_peak_bytesa documented lower bound on the true peak); a mid-run agent failure fails the run.jmlx-benchmarks:orchestration(spec-order|alternating) and per-workerinstrumentation, the-<instrumentation>cell-directory suffix, the per-session isolatednative-cache(retained on success, deleted on any failure), full-cell-identifierfork_order, theinstrumentedexclusion reason, and the validator rules (the memory rule, the cross-mode absence).ffm_cal.c(one non-elidable no-op per dominant generated-wrapper shape; the carrier-fidelity argument documented in the source — every opaque mlx-c struct is a single 8-byte pointer, register-identical tovoid*on AArch64),FfmCalWorker(fresh JVM, the downcall invoked in the generated wrappers' owninvokeExactform),FfmCalibration(builds underbuild/withcc -shared -O2— no production native artifact;build-info.jsonwith compiler/flags/source+binary SHA-256s; per-shape p10/median/p90 with every raw sample retained, plus the empty-loop timing floor) andFfmTransitionEstimate(pure: a binding-count worker report + a calibration summary → the per-decode-step budget with its p10/p90 range, coverage, uncovered signatures and the explicit assumptions — a calibrated estimate, not a causal partition).:jmlx-benchmarks:benchmarkFfmCalibrationtask only.benchmarkDecode,benchmarkBatchDecode,benchmarkTf32andbenchmarkBaselineare preserved; ordinarycheckgains no downloads, Python execution or C-compiler prerequisite (the calibration test is skipped whenccis absent).Frozen-contract amendments (findings §13.4 item 20)
The frozen §13.1/§13.3/§14.3 records are amended in place: the job/worker
instrumentationfield, the iff-instrumentedidentity.host,binding_signatures, the per-requestbinding_counts/tracing, the WP2 memory rule (replacing the WP1 "stays null" sentence), the sessionorchestration/native_cache, and the full-cellfork_order. No schema bump — every new field is additive, and the headline job-file shape is preserved (the launcher writesinstrumentationonly when it is notnone).Evidence
:jmlx-benchmarks:checkgreen with 89 tests (49 launcher / 28 validator / 6 statistics / 1 calibration / 3 estimate / 2 TF32); the:jmlx-modelsbaseline suites green (5/9/22 — up from the 20 protocol tests of §14.6); rootcheckgreen on this head; the opt-in:jmlx-core:bindingCoverageexecuted against the staged native runtime.(M)I~31.8 ns/call, the 4-5 carrier shapes ~40-47 ns/call,(M)M~41 ns/call,(AM)M~107 ns/call (the per-call result-segment allocation through the caller's allocator — real generated-wrapper behavior for struct returns, kept and named in assumption 1)..github/workflows/ci.ymlis unchanged.