Skip to content

feat: overlap-aware expert prefetch - #180

Open
drunkcoding wants to merge 10 commits into
mainfrom
plan/overlap-aware-expert-prefetch
Open

feat: overlap-aware expert prefetch#180
drunkcoding wants to merge 10 commits into
mainfrom
plan/overlap-aware-expert-prefetch

Conversation

@drunkcoding

@drunkcoding drunkcoding commented Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Description

Implements the Momus-approved overlap-aware expert prefetch plan
(docs/superpowers/plans/2026-08-21-overlap-aware-expert-prefetch.md). Admits
speculative expert transfers only when their exact stored bytes fit a
conservatively calibrated transfer window derived from measured bandwidth and
the current layer's measured compute time. The native router, expert weights,
and dispatched expert set remain authoritative — budgeting applies only to early
cache warming.

What landed (9 atomic commits, strict TDD)

  • Deterministic controller (moe_infinity/memory/overlap_budget.py): pure
    EWMA calibration, whole-expert greedy byte admission, cold-start cap, and
    exact-route correction accounting (coverage/waste/late/uncovered/cancel).
  • Config surface (overlap_prefetch_policy=off|observe|enforce + knobs and
    the shared gpu_only_expert_routing field) with fail-closed validation.
  • Native bounded queues: header-only PrefetchQueueAccounting value helper
    (GPU-free, unit-tested), RemovalReason/RunningOutcome enums, a
    RunPrefetchTaskNoThrow worker boundary, and thread-safe
    admit/cancel/drain/inflight pool + handle + pybind APIs. Legacy
    prefetch_tensors unchanged.
  • Native compute intervals: testable ExpertTimingLifecycle
    (free/pending/quarantined over a CUDA adapter) plus dispatcher-owned
    timing-enabled epoch/start/stop events bracketing ForwardHelper only,
    nonzero invocation ids, and set_inputs_with_invocation /
    set_overlap_compute_timing_enabled / drain_compute_samples bindings.
  • Prefetcher + executor wiring: controller ownership, capability detection,
    exact byte costing, exact-route correction before every enqueue, one-layer
    admitted-only pinning, and (invocation_id, layer_id)-exact calibration with
    error/finally cleanup that never masks a wait_expert error.
  • Profiling + metrics: IOProfiler.record() instant events and overlap byte
    stages in ExpertTracer.
  • Benchmark + report gates: paired off|observe|enforce
    bench_overlap_prefetch.py with a pure build_report, run_all --scenario overlap, and run_decision_profile --overlap-prefetch-policy.
  • Docs: configuration knobs, formula, cross-plan compatibility table,
    staged rollout/rollback, and the risk/mitigation table.

Verification

  • CPU deterministic gates: 109 passed (controller, config, profiler,
    speculative-prefetch, route-ahead wire/metrics, report, output-equivalence,
    qwen3.5 rollback).
  • Native GPU-free gates: 10 passed — PrefetchQueueAccounting (5) and
    ExpertTimingLifecycle (5) via the opt-in MOE_BUILD_TESTS root build.
  • Native extension (prefetch_op, --whole-archive over archer_core) builds
    and links cleanly with all new bindings.
  • Opt-in native compute-sample smoke test collects and SKIPs cleanly without
    MOE_DFLASH_SERVING_GPU.

Dependencies / limitations

Independently implementable on eager routing. First-release co-enablement with
GPU-only routing is rejected until route-ready ordering is reconciled;
ArcherConfig raises when gpu_only_expert_routing=True is combined with an
active overlap policy.

Type of Change

  • Production implementation
  • Documentation update

Performance / support evidence

  • Not applicable — benchmark methodology only; results are labeled MEASURED
    with no speedup claim.

Checklist

  • Source-aligned implementation, strict TDD
  • Default-off fallback preserves legacy path byte-for-byte
  • Native router semantics unchanged

drunkcoding @sisyphus-dev-ai
docs(plan): add overlap-aware expert prefetch plan
89e1ac6 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding

drunkcoding commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Momus plan review: [OKAY] on the first review round.

drunkcoding and others added 9 commits August 28, 2026 22:21
drunkcoding @sisyphus-dev-ai
feat: add overlap prefetch budget model
5152888 
Pure-Python deterministic overlap-budget controller: EWMA calibration, whole-expert greedy byte admission, cold-start cap, and exact-route correction accounting (coverage/waste/late/uncovered/cancel). No Torch/CUDA imports; fully CPU-testable.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: configure overlap-aware prefetch policy
ebe2b78 
Add default-off off|observe|enforce policy plus validated EWMA/safety/cold-start/window/inflight knobs and the shared gpu_only_expert_routing field. Reject gpu_only_expert_routing=True together with an active overlap policy in the first release. Document the formula, cold-start, fail-closed behavior, rollout, and cross-plan compatibility table.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: report overlap prefetch effectiveness
03099af 
Add IOProfiler.record() public instant-event API (dur_ns=0) that merges caller fields and rejects keys overwriting reserved event columns. Register overlap byte stages (prefetch_budget/admit/complete/cancel/late/waste) in ExpertTracer for consistent byte aggregation.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: bound and observe native prefetch queues
60af9f9 
Add header-only PrefetchQueueAccounting value helper (GPU-free, unit-tested) with RemovalReason/RunningOutcome enums and a RunPrefetchTaskNoThrow noexcept worker boundary that retires bytes exactly once. Add Task generation/layer/bytes metadata, PrefetchAdmission/PrefetchSample structs, and thread-safe AdmitPrefetchTasks/CancelQueuedPrefetch/DrainPrefetchSamples/GetInflightPrefetchBytes pool APIs backed by the accounting helper. Bridge through ArcherPrefetchHandle with tensor-to-node dedupe and bind schedule_prefetch_tensors/cancel_prefetch_generation/drain_prefetch_samples/get_inflight_prefetch_bytes plus read-only value structs. Legacy prefetch_tensors behavior unchanged. Add opt-in MOE_BUILD_TESTS root build and a GPU-free GoogleTest for the accounting invariants.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: expose native expert compute intervals
ff87468 
Add testable ExpertTimingLifecycle (free/pending/quarantined state machine over a CudaTimingAdapter, GPU-free unit-tested) and wire dispatcher-owned, timing-enabled epoch/start/stop events around ForwardHelper only. MoEMLP::forward records kernel start/stop immediately around ForwardHelper (excluding input copy and syncs). GPUExecFunc lazily creates per-worker timing objects after cudaSetDevice only when the atomic enable flag is set, brackets forward, captures separate host forward-return/output-complete steady-clock delays, polls pending stops between work items, and quarantines exception pairs behind a same-stream fence until a query or stream sync proves completion. Every ExpertComputeSample carries a nonzero invocation id threaded through CallArgs/ExecArgs via set_inputs_with_invocation. Bind set_inputs_with_invocation/set_overlap_compute_timing_enabled/drain_compute_samples and the read-only expert_compute_sample struct; legacy set_inputs unchanged. Add the GPU-free lifecycle GoogleTest and an opt-in native compute-sample smoke test.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: budget expert candidates by overlap bytes
3824aba 
Give ExpertPrefetcher an overlap controller, a NativeOverlapCapabilities adapter, exact byte costing, and native admission/cancellation/telemetry bridges (plan_candidates, correct_to_native_route, abort_prefetch_generations, drain_native_prefetch_samples, observe_compute_samples, record_stale_compute_samples, overlap_prefetch_stats). configure_overlap_policy leaves off byte-for-byte, enables dispatcher timing only when the full capability set is present, records a capability miss under observe on an older extension while keeping legacy issuance, and fails closed under enforce. Add ExpertPredictor.ranked_candidates. Configure the policy in model_offload right after the exact byte map is registered.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: reconcile prefetch with native expert routes
4e4b325 
Apply exact-route correction before every expert enqueue and budget exact-route warming through plan_candidates on the supported eager path, pinning and reporting only admitted experts while the full native union is still dispatched. Thread a per-dispatch invocation id via set_inputs_with_invocation when the dispatcher timing capability is complete, else legacy set_inputs. Extend _pending_prefetch to carry owned generations and the invocation id; _maybe_route_ahead_prefetch returns (fired, generations). Restructure wait_dispatch_local so pending state is detached before waiting and all policy cleanup runs in finalize_policy: cancel owned generations on wait error, accept only the exact (invocation_id, layer_id) compute samples (counting the rest as stale), correct/re-trigger on success, and always drain native samples in finally without masking the original wait error.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
bench: measure overlap-aware expert prefetch
47f9298 
Add bench_overlap_prefetch.py: a pure build_report that derives coverage/waste/late ratios and p50/p95 latency, enforces the byte invariants (completed<=admitted, covered<=route, late<=route, wasted<=completed), and labels the result MEASURED with no speedup key; plus a paired off|observe|enforce CUDA driver that runs each policy in a separate process, resets the cache between arms, and compares output ids exactly against the off oracle. Extend run_decision_profile.py with --overlap-prefetch-policy and knob flags and append overlap_prefetch_stats; add run_all.py --scenario overlap that shells out without altering other scenarios. Document commands, interpretation, and the no-speedup promise; add the CPU-only report-gate test.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
docs: define overlap prefetch rollout and rollback
0ba8420 
Document the staged off->observe->enforce rollout per model/hardware pair, config-only rollback (set overlap_prefetch_policy=off and restart; deploy the prior package if the extension is suspect), rollback triggers, and the full risk/detection/mitigation table in docs/configuration.md.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding

drunkcoding commented Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

Independent implementation verification:

  • Focused CPU suites: 79 passed.
  • Python LSP scan: 0 errors across 50 files.
  • Branch clean and synchronized.

PR remains draft for native/CUDA maintainer review.

@drunkcoding
drunkcoding marked this pull request as ready for review September 2, 2026 21:53
@drunkcoding drunkcoding changed the title WIP: plan overlap-aware expert prefetch feat: overlap-aware expert prefetch Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@drunkcoding