WIP: PD-DFlash Task 6 Step 5 — SpecSession seam + engine byte-demand VERIFY - #163
Merged
Conversation
added 13 commits
August 14, 2026 21:24
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.
…(DRAFT/VERIFY scheduler) Stacks the two independent additive PD-DFlash features from the same plan: - origin/feat/pd-dflash-phaseA-runner (#157): expert_nbytes / ExpertPrefetcher.expert_nbytes_map - origin/feat/pd-dflash-serving-scheduler (#156): Scheduler DRAFT/VERIFY + set/clear/admit_verify_demands Conflicts in benchmarks/dflash/report.py and tests/python/dflash/test_pd_dflash_report.py resolved KEEP-BOTH: #157's versions are strict supersets (Task-1 contract identical, plus BM1/aggregation/CLI), preserving 100% of both sides' behavior.
…lator seam Honest gate: Steps 1-4 of Task 6 are complete via the #156+#157 stack this branch is based on. Step 5 (engine DRAFT->VERIFY->DRAFT with exact per-round expert-byte demand) genuinely requires a public single-round control seam on DFlashSpeculator (begin_session/draft_round/verify_round + a read-only route projection) so the engine can register the pending verify's EXACT summed expert_nbytes BEFORE the verify runs. That seam does not exist (no draft_round/verify_round anywhere), and Task 6's authorized file list scopes only serving/{batch,scheduler,engine}.py -- not spec_decode/dflash.py or distributed/expert_executor.py. Post-hoc RouteAheadStats can't gate admission, fabricating bytes from expert count is forbidden, and re-implementing the draft/verify loop in the engine would be a forbidden second router. Reporting the exact seam instead of hacking one in.
Externalize the _generate_single loop behind begin_session / draft_round / verify_round (Task 6 Step 5). draft_round returns the pending verify's projected demand (block width + a read-only routed-expert union captured from the SAME union_experts_from_mask gate/top-k path the verify uses, via RouteUnionCollector) whose bytes are the EXACT sum of ExpertPrefetcher.expert_nbytes_map -- never a fabricated expert-count estimate. verify_round runs one _verify_target_block under route_ahead_context and reproduces the accept/commit/rollback exactly. generate() / _generate_single / _generate_batched are left byte-identical; the session runs alongside them. Adds project_expert_bytes + the SpecSession / DraftResult / VerifyResult / RouteUnionCollector public surface. CPU tests pin session==generate() token identity (greedy + lossless-sampled), round-API order, immediate-stop anchor, and the exact-bytes projection.
…Step 5) Wire the engine per-round DRAFT->VERIFY->DRAFT lifecycle behind an opt-in gate: only when verify budgets are configured AND the speculator exposes the SpecSession seam. Per round the engine registers the drafter-projected EXACT token/byte demand (set_verify_demand), runs the verify only for the scheduler's admitted verify_seq_ids, advances the sequence PREFILL->DRAFT->(VERIFY->DRAFT)* ->FINISHED, and clears the demand; unadmitted rounds stay DRAFT and their 2-D deficit carries. Reuses #156 admit_verify_demands/set_verify_demand and #157 expert_nbytes_map; adds Scheduler.verify_scheduling_enabled. Default serving is byte-for-byte unchanged: without budgets or with a non-session speculator the delegated request keeps the whole-request generate() path. CPU tests cover exact-byte registration, carried-deficit liveness, the misconfig raise, and default/fallback preservation.
drunkcoding
marked this pull request as ready for review
August 15, 2026 18:22
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
1 participant
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.
What & why
Completes Task 6 Step 5 of the PD-DFlash serving plan: the serving engine now drives per-round DRAFT→VERIFY→DRAFT and registers each pending verify's EXACT token/byte demand so the 2-D verify scheduler governs when VERIFY runs under concurrency. This unblocks the seam reported in
docs/superpowers/notes/2026-08-14-pd-dflash-task6-step5-seam.md.Stacks on #162 (
feat/pd-dflash-task6-engine) which already contains #156 (DRAFT/VERIFY 2-D scheduler) + #157 (expert_nbytesinstrumentation). Target base:dev.The
SpecSessionseam (spec_decode/dflash.py)Externalizes the
_generate_singleloop as a public single-round control surface:begin_session(...) -> SpecSession— prefill/anchor forward; seeds per-round state (target_kv,draft_kv,context_feature,anchor,start,emitted,step_trace). Handles the immediate-stop anchor case.draft_round(session) -> DraftResult— one_run_drafterpass builds the block; returns the pending verify's projected demand (tokens=block_size,expert_union,expert_bytes).verify_round(session) -> VerifyResult— one_verify_target_blockunderroute_ahead_context, exact accept/commit/rollback + step-trace, then updates the projection for the next round.Read-only route projection (no second router)
The projection reuses the existing
route_ahead_context(..., stats=…)recorder slot: aRouteUnionCollector(duck-typingbegin_step/observe_layer/commit_step) captures the verify's ACTUAL per-layer routed union via the sameunion_experts_from_maskgate/top-k primitive the verify'sdispatch_localuses — read-only (no pin, no prefetch dispatch, no expert exec, no routing/output change). Bytes =project_expert_bytes(union, ExpertPrefetcher.expert_nbytes_map)— the exact summed FP4 payload, never an expert-count estimate. This needs zero changes toexpert_executor.py/_route_ahead_ctx.py(a smaller blast radius than the spec's sketch, same contract).Because exact same-round routing is only knowable during the verify forward (each MoE block computes
router_maskinline per layer), the projection is round-to-round: the prior verify's union projects the next round's demand; the first round is a byte-free warm-up (0 bytes).generate()is byte-identicalgenerate()/_generate_single/_generate_batchedare untouched; the session runs alongside. Proven by: (a) all pre-existing spec-decode tests stay green, and (b) a new test asserting the session loop reproducesgenerate()token-for-token (greedy and lossless-sampled).Engine wiring (
serving/engine.py) — opt-in, default unchanged_step_speculative_sessionruns only when verify budgets are configured and the speculator exposes the seam. Per round:draft_round→set_verify_demand(seq_id, tokens=B, expert_bytes=Σ)→schedule()→ run verify only for admittedverify_seq_ids→clear_verify_demand; sequence advances PREFILL→DRAFT→(VERIFY→DRAFT)*→FINISHED. Unadmitted rounds stay DRAFT and their 2-D deficit carries. Reuses #156'sadmit_verify_demands/set_verify_demandand #157'sexpert_nbytes_map; addsScheduler.verify_scheduling_enabled.Default serving is byte-for-byte unchanged: without budgets, or with a non-session speculator, the delegated request keeps the whole-request
generate()path (_step_speculative).Tests / verification
tests/python/dflash/test_spec_session.py(10) — session==generate() identity (greedy+sampled), round-API order, immediate-stop anchor, exact-bytes projection (count ≠ bytes), collector read-only.tests/python/serving/test_dflash_engine_step5.py(5) — exact-byte demand registration, carried-deficit liveness, misconfig raise, default/fallback preservation.test_gpt_oss_mxfp4_dispatchfailures are pre-existing/environmental — unbuiltmoe_infinity._v4_fp4native ext — identical on the base branch; this PR is pure Python and touches no FP4/gpt-oss code.)ruff check+ruff formatclean; LSP diagnostics clean on all changed files.Scope guardrails
Pure Python (no
_store/C++ rebuild). Does not merge, does not touch GLM /feat/dflash-spec-decode/ #158–#161/#119, and adds no second router/prefetch path.Draft — do not merge.