Skip to content

WIP: PD-DFlash Task 7 — BM2 + batched prefetch - #160

Draft
drunkcoding wants to merge 7 commits into
devfrom
feat/pd-dflash-task7-batched-prefetch
Draft

WIP: PD-DFlash Task 7 — BM2 + batched prefetch#160
drunkcoding wants to merge 7 commits into
devfrom
feat/pd-dflash-task7-batched-prefetch

Conversation

@drunkcoding

Copy link
Copy Markdown
Contributor

WIP: PD-DFlash Task 7 — BM2 issuance micro-bench + batched prefetch_tensors

Draft. 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 diff
against dev also shows #157's Phase-A commits. Does not implement the later
priority-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:

mode p50 p90 p99
python-per-expert 1132.1 µs 1891.6 2238.7
batched-pybind 101.5 µs 105.3 111.0
cpp-internal null (no native timer)

Batching collapses issuance 1132 → 101.5 µs (~11.2×, ~1.03 ms/block saved) by
replacing 2·E_l·L pybind crossings with one C++-internal loop. The
justified-window interval [batched_p50, per_expert_p50) = [101.5, 1132) µs
covers any realistic route-ahead window t_draft+t_router, so per-expert
issuance is exposed and the batched call hides it → ship_batched=true.

Provenance / caveats (reported honestly):

  • Target substitution: the plan's Qwen/Qwen3-Coder-30B-A3B is not in the HF
    cache; 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.
  • Window: no Phase-A t_draft+t_router JSON exists yet (Task 3 not run on
    hardware); 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

  • BM2 (benchmarks/dflash/bench_prefetch_issuance.py): import-safe
    bm2_decision / percentiles_us / build_bm2_report at module scope; CLI
    runner synthesises the real saturated tensor-id list and times with
    perf_counter_ns; unavailable modes reported null, never zero.
  • Batched C++ (core/prefetch/archer_prefetch_handle.{h,cpp},
    core/python/py_archer_prefetch.cpp): EnqueuePrefetchTensors(tensor_ids, priority=1) enqueues one Task/tensor entirely in C++ (order- and
    default-device-preserving); prefetch_tensors rebound to it, the old no-op
    binding retired, enqueue_prefetch retained.
  • Python seam (moe_infinity/memory/expert_prefetcher.py):
    prefetch_experts_list issues one batched call when the engine exposes it,
    else the byte-for-byte per-expert fallback; empty input is a no-op.

Verification

  • CPU: test_prefetch_perf_reports.py (20) + test_speculative_prefetch.py
    (batch/fallback/empty) + test_route_ahead_wire.py50 passed
    (wire assertions made mechanism-agnostic: a batched call carries the same
    ordered ids the per-expert fallback would).
  • Native _store rebuilt on SM120; prefetch_tensors(tensor_ids, priority=1)
    signature verified; opt-in test_prefetch_native_gpu.py -k batched3
    passed
    on a real offloaded handle (single module-scoped load).
  • Offload no-regression: tests/test_gpt_oss_offload_topology.py +
    tests/python/unit/test_gpt_oss_mxfp4_dispatch.py9 passed.
  • ruff check + ruff format clean; lsp diagnostics 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.

drunkcoding 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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant