WIP: PD-DFlash Task 7 — BM2 + batched prefetch - #160
Draft
drunkcoding wants to merge 7 commits into
Draft
Conversation
added 7 commits
August 14, 2026 22:22
Record exact stored FP4/FP8 expert payload bytes through the route-ahead observer so the measure-first gate can report wasted_prefetch_bytes as bytes, not an expert count. RouteAheadStats.observe_layer gains an optional expert_nbytes map and RouteAheadStepSummary/as_dict expose predicted/kept/ wasted byte fields (None when unavailable -- never a fabricated average). ExpertPrefetcher.expert_nbytes_map is populated at registration time in model_offload from live params (before offload placeholders erase shape); the executor seam forwards it None-safely (isinstance-dict guard keeps mocks and resident runs at None). Strictly additive and off by default.
Opt-in RTX PRO 6000 B0-B3 route-ahead serving runner. The CLI module is import-safe (torch/moe_infinity imported only inside the GPU path), emits one JSON row per (model,baseline,B,concurrency,repeat) matching REQUIRED_METRICS, validates the device is an RTX PRO 6000 (12,0), refuses resident B0/B1/B2, and blocks B2 as BLOCKED_UNTIL_2D_SCHEDULER until the 2-D scheduler lands. Wraps draft/router/issue/verify/H2D in the frozen NVTX ranges the BM4 parser keys on. GPU test is gated on MOE_DFLASH_SERVING_GPU (1 skipped, side-effect free collection); CPU contract test locks the pure matrix/schema/writer logic.
Add BM1 summarise_row (pass iff t_router < t_verify, ratio + raw terms retained) and a CPU-safe aggregation CLI: group raw rows into §8 matrices keyed by (model,block,concurrency), permit blocked B2 via --allow-blocked, attach BM1, and emit result_matrix.json/CSV/Markdown for validate_result_matrix. Refactors the per-baseline metric check out of validate_result_matrix (behavior-preserving).
run_phase_a.sh drives the full §8 matrix for both required MoE targets on one RTX PRO 6000 with the documented env (HF_HOME, MOE_ENABLE_SM120, device-memory- ratio<0.9 to force offload) and aggregates into result_matrix.json. All inputs are documented env vars in the script header.
Task 7 of the PD-DFlash serving plan (design §10 BM2). Adds an import-safe issuance micro-benchmark that times enqueuing a saturated E_l×L block of offloaded expert tensors three ways -- python-per-expert, batched-pybind, and (reserved) cpp-internal -- reporting p50/p90/p99 microseconds and the ship gate. - benchmarks/dflash/bench_prefetch_issuance.py: pure bm2_decision / percentiles_us / build_bm2_report at module scope (torch + moe_infinity lazily imported in the GPU runner), plus the CLI runner that synthesises the real saturated tensor-id list from a loaded target and times warmup=20 / iterations=200 with perf_counter_ns. Unavailable candidate modes are reported null, never zero. - tests/python/dflash/test_prefetch_perf_reports.py: CPU-only decision-rule and report-schema tests (the plan's four exact assertions plus boundary cases). BM2 alone gates the batched-issuance C++ hop; committed independently of it.
Task 8 of the PD-DFlash serving plan (candidate hop 1), retained because BM2 passed on offloaded gpt-oss-20b: python-per-expert issuance p50=1132 µs vs batched-pybind p50=101 µs over a 768-tensor saturated block (~11× fewer pybind crossings), so per-expert issuance is exposed and the batched call hides it. - core/prefetch: add ArcherPrefetchHandle::EnqueuePrefetchTensors(tensor_ids, priority=1), which constructs and enqueues one Task per tensor entirely in C++, preserving input order and node default devices (mirrors EnqueuePrefetch). - core/python: bind prefetch_tensors -> EnqueuePrefetchTensors and retire the old no-op PrefetchTensors(request_id, buffer) binding; enqueue_prefetch retained. - expert_prefetcher: prefetch_experts_list issues one batched prefetch_tensors call when the engine exposes it, else the byte-for-byte per-expert fallback; empty input is a no-op. - tests: batch/fallback/empty coverage in test_speculative_prefetch.py; wire assertions made mechanism-agnostic (batched call carries the same ordered ids); opt-in native GPU smoke test (single module-scoped offload load).
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.
WIP: PD-DFlash Task 7 — BM2 issuance micro-bench + batched
prefetch_tensorsDraft. Do not merge. Implements Phase C's first C++ hop from
docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md(Task 7 = BM2 issuance micro-bench, Task 8 = the batched-issuance C++
candidate). Enforces the plan's hard rule — no C++ ships without its paired
benchmark passing — by measuring BM2 on hardware before keeping the C++.
Stacks on #157 (
feat/pd-dflash-phaseA-runner). Until #157 merges, the diffagainst
devalso shows #157's Phase-A commits. Does not implement the laterpriority-band (BM3) / overlap (BM4) / e2e (BM5) hops.
BM2 gate decision — JUSTIFIED ✅
Measured on a real offloaded
openai/gpt-oss-20b(RTX PRO 6000, SM120,CUDA_VISIBLE_DEVICES=0,device_memory_ratio=0.2), 768-tensor saturated block(
E_l×L), warmup=20, iterations=200:python-per-expertbatched-pybindcpp-internalnull(no native timer)Batching collapses issuance 1132 → 101.5 µs (~11.2×, ~1.03 ms/block saved) by
replacing
2·E_l·Lpybind crossings with one C++-internal loop. Thejustified-window interval
[batched_p50, per_expert_p50) = [101.5, 1132) µscovers any realistic route-ahead window
t_draft+t_router, so per-expertissuance is exposed and the batched call hides it →
ship_batched=true.Provenance / caveats (reported honestly):
Qwen/Qwen3-Coder-30B-A3Bis not in the HFcache;
gpt-oss-20b(cached, offload-validated) is used. Qwen3-Coder-30B has~8× more expert tensors, so per-expert issuance (and the batching win) would be
strictly larger — the substitution is conservative.
t_draft+t_routerJSON exists yet (Task 3 not run onhardware); the gate is reported as a window-independent crossover plus a
1000 µs reference. Both yield
candidate_required=true,ship_batched=true.What's here
benchmarks/dflash/bench_prefetch_issuance.py): import-safebm2_decision/percentiles_us/build_bm2_reportat module scope; CLIrunner synthesises the real saturated tensor-id list and times with
perf_counter_ns; unavailable modes reportednull, never zero.core/prefetch/archer_prefetch_handle.{h,cpp},core/python/py_archer_prefetch.cpp):EnqueuePrefetchTensors(tensor_ids, priority=1)enqueues oneTask/tensor entirely in C++ (order- anddefault-device-preserving);
prefetch_tensorsrebound to it, the old no-opbinding retired,
enqueue_prefetchretained.moe_infinity/memory/expert_prefetcher.py):prefetch_experts_listissues one batched call when the engine exposes it,else the byte-for-byte per-expert fallback; empty input is a no-op.
Verification
test_prefetch_perf_reports.py(20) +test_speculative_prefetch.py(batch/fallback/empty) +
test_route_ahead_wire.py— 50 passed(wire assertions made mechanism-agnostic: a batched call carries the same
ordered ids the per-expert fallback would).
_storerebuilt on SM120;prefetch_tensors(tensor_ids, priority=1)signature verified; opt-in
test_prefetch_native_gpu.py -k batched— 3passed on a real offloaded handle (single module-scoped load).
tests/test_gpt_oss_offload_topology.py+tests/python/unit/test_gpt_oss_mxfp4_dispatch.py— 9 passed.ruff check+ruff formatclean;lspdiagnostics clean on changed files.Explicitly NOT in scope
No BM3 priority band, no BM4 overlap, no BM5 e2e; no scheduler work; no GLM /
feat/dflash-spec-decode/ #156 / #158 changes. GPU work is opt-in(
MOE_DFLASH_SERVING_GPU=1) and collects/skips cleanly when ungated.