[RFC] [MLX] Write down the per-step contract that a replacement event loop must satisfy
apple-silicon
`Scheduler.run_batch` runs the torch forward and performs the per step bookkeeping backend neutral readers depend on. `event_loop_overlap_mlx` replaces the first and inherits the second, with nothing enforcing the inheritance. Four bugs across July hit four of the five clauses below: `TypeError` at startup, silent cache loss, silent wrong output, a segfault behind a flag pair. The fifth returns the right answer by accident.
The pattern is older than those. #29217, merged in June, states in its own root cause section: the MLX loop bypasses `run_batch`, where `forward_ct` is incremented and the profiler predicate is called. Its fix mirrored the missing write inside the loop. Clause 3 then broke again in July and was fixed the same way by #32447.
Line numbers anchored to `50029f0`, which predates #32447; the `scheduler_mixin.py` anchors have drifted by a couple of lines since.
## The contract
A step that emits tokens must leave the following true when it ends.
1. Slots recorded. `alloc_for_decode` writes `req_to_token[row, position]` (`allocation.py:578`), read back as a pool slot id by `_sync_decode_kv_to_pool`.
2. Counters advanced. `seq_lens`, `kv_committed_len`, `kv.kv_allocated_len` (`schedule_batch.py:2884-2892`), read by radix insertion as the visibility bound.
3. Step stamped. `forward_ct`, `forward_iter`, `launch_ts` (`scheduler.py:3371-3373`), read by the watchdog, the profiler and the SWA eviction cadence.
4. Relays published or provably unread. `future_map.output_tokens_buf` is written only from the torch branch of `run_batch` (`scheduler.py:3460`).
5. In flight predicates test any overlap loop. `enable_overlap` and `enable_overlap_mlx` are mutually exclusive by construction (`scheduler.py:382`).
| Clause | On the MLX loop | Status |
|---|---|---|
| 1 | Skipped on the chained path | Sync side clamped by #30147; open allocation |
| 2 | Skipped on the chained path | Open, #30093 |
| 3 | Fixed | #29217, #32447 |
| 4 | Never published; inert for pure decode, live for a mixed batch | Path disabled on MLX by #32591 |
| 5 | Wrong flag, masked | #32934 |
## Clause 2
The `seq_lens` undercount never self corrects, since the next fresh `prepare_for_decode` adds one regardless of chain length and `get_next_batch_to_run` returns `running_batch` by identity (`scheduler.py:2883`), so the stale value sits on the live object for the request's lifetime. #30093 documents the per chain skip; the permanence is what makes it a contract violation rather than a transient. #30147 clamps the sync to the committed bound and does not advance it, so clause 2 stays open after the PR lands.
The retraction budget reads these counters, but chained steps never allocate, so `kv_committed_len` tracks real pool occupancy and `check_decode_mem` stays self consistent. The measured cost is radix truncation, not a memory pressure miscalculation.
## Clause 5
`is_fully_idle` (`scheduler.py:3814`) checks `not self.enable_overlap or len(self.result_queue) == 0`, so MLX never inspects the queue. It returns the right answer anyway because `running_batch` is non empty whenever the queue is. Nothing enforces coupling and `release_memory_occupation` asserts on this predicate.
## Scope
#32321 closes reaching through the worker, which produces absence bugs; this is state never written, which produces staleness. Nothing asks chained decode to change.
## Proposed enforcement
One registered unit test driving a short prefill and decode workload through both loops and asserting all five clauses at each step boundary, failing with the clause number and the field name rather than a downstream symptom. #32934 is a worked example of the shape at one clause: five lines of production change, three cases, red on the parent and green on the branch. Adding the other four is incremental, one clause per commit, and none of it requires a model.
2 条评论