feat(dflash): converge DFlash spec-decode stack → dev (supersedes #140) - #143
Merged
Conversation
added 30 commits
July 30, 2026 19:31
Add moe_infinity/spec_decode with DFlashSpeculator, DFlash config reader, and target/draft pairing validation for z-lab GPT-OSS DFlash drafters.
…e=5, transformers>=5.12 guarded)
…ompat re-export preserved)
…nce, first_k_dense, eager attn, MTP-safe)
…decar, SetScales interface)
…nding (bit-exact vs python; dispatcher H2D integration deferred)
… patch, GLM use_native_engine=False, e_score_correction_bias device guard)
…duces native expert-dispatch SIGSEGV
…ense offset) — fixes native dispatch SIGSEGV; tiny GLM now generates end-to-end
…ess-by-construction; validated on tiny model)
…ep stats, expert-fetch hook) — parity preserved
…ash adapter, process-isolated)
… on non-paged HF path; +tokenizer save +gated test)
…erver, process-isolated)
…quant, gated) — Option A==B parity on tiny FP8 GLM; completes T15/T16
…de roofline (scaffold, cross-plan option b)
…ff + predicted roofline -> CSV)
added 25 commits
August 5, 2026 17:58
…reconstructible sliding-window KV rollback
GPT-OSS is excluded from the C++ expert dispatcher and runs a resident
Python expert loop in SyncGptOssMLP.forward, but the loader never
materialized the _PackedExperts params: expert weights stayed zeros,
biases NaN, attention sinks and router garbage, producing incoherent
output. Add _load_resident_gpt_oss to load the real MXFP4 blocks/scales
(uint8, output-major [E,N,K//2]/[E,N,K//32]), biases, router, and
self_attn.sinks into the live params and drop them from name_id_map.
Fix _expert_forward_mxfp4 to feed packed weights without the erroneous
.t() (checkpoint is already output-major; verified rel-err 0.0 vs
reference dequant on gate_up and down proj).
Real 120B: base output now coherent ('The capital of France is Paris.');
DFlash agreement 0.95->1.00, mean acceptance 1.0->6.57, ~21x decode
speedup. 103 dflash tests pass.
Per-token union of the verify block's routed experts, pinned/prefetched via the dispatch_local seam during DFlash verify. Adds union/coverage/waste pure helpers (_prefetch_route), a contextvars route-ahead context (_route_ahead_ctx), opt-in coverage/waste metrics (_route_ahead_stats), an explicit-set mode on ExpertPrefetcher.speculative_prefetch (legacy path byte-identical), and the executor observe seam. Non-spec decode unchanged.
Distribution-preserving block-diffusion accept rule for temperature/top-k/top-p, keeping greedy byte-identical.
Adds a spec branch to ContinuousBatchingEngine.step() for greedy batch==1 sequences: paged-KV rollback primitive (kv_cache.truncate_tokens), per-seq cached-vs-emitted accounting (SpecDecodeState), verify-commit orchestration (spec_verify.apply_verify_step), and variable-commit scheduler accounting (committed_counts, default byte-identical). Spec-off path unchanged.
… hooks, drafter-driven contract Generalizes the native draft->verify->rollback loop to batch>1 (batched build_block/acceptance/committed_tokens with per-seq accept + rollback), wires the sampled accept + route-ahead context/stats into generate(), and makes validate_pairing drafter-driven (reads block_size/target_layer_ids from the drafter config instead of hardcoding the gpt-oss-120b contract) so any z-lab DFlash pair loads.
MOE_DFLASH_GPU-gated: gpt-oss-20b native DFlash == plain greedy (token-identical), and serving-path DFlash == sync DFlash.
…ang-format) Pre-existing formatting drift on the feat/dflash-spec-decode base that CI enforces via 'pre-commit run --all-files' (ruff --fix, ruff-format, clang-format). Mechanical only; no logic changes.
Add null checks in archer_prefetch_handle ReleaseTensor for the first-call (last_node_ null) and missing-node cases to prevent a native segfault during offloaded expert release.
Keep the Qwen3.5-MoE text backbone (incl. GatedDeltaNet linear-attention layers and shared expert) resident on GPU while offloading only the routed experts, and generalize routed-expert key matching via _is_routed_expert_key.
Make DFlash speculative decoding lossless/correct on offloaded, executor-backed Qwen3.5 (rollback via clone instead of in-place copy_; @torch.no_grad over @torch.inference_mode so worker-thread expert index_add_ works), add resident route-ahead observation for the non-executor path, and harden native-engine plumbing in MoE.generate. Includes hybrid-rollback and route-ahead wiring tests.
Resolve model_offload.py get_topology conflict: adopt _is_routed_expert_key (regex) as the routed-expert detector, dropping the parse_expert_id-based is_routed_expert closure from #138. Keeps MTP-layer experts offloaded (dflash intent) but narrows detection to '.mlp.experts.<id>.'; Mixtral '.block_sparse_moe.experts.' is no longer matched -- tracked in #142.
2 tasks
Contributor
Author
|
Merging to converge the DFlash stack into |
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
Converges the DFlash speculative-decoding stack into
dev. This is the up-to-date successor of #140 (feat/dflash-spec-decode) and supersedes it —feat/dflash-tracks-abcdfully contains #140's tree plus 12 newer commits (offloaded Qwen3.5 spec-decode + route-ahead, serving integration, sampled/batched decoding, GLM perf-model, prefetch null-deref fix, ruff/clang-format convergence).origin/devwas merged in and the one conflict resolved (see below), so this is a clean merge.Conflict resolution (
moe_infinity/runtime/model_offload.py)Only
get_topology()'s routed-expert detector conflicted — two parallel implementations:is_routed_expertclosure →parse_expert_idfor qwen3_5, else"expert" in name._is_routed_expert_key→ regex\.mlp\.experts\.\d+\..Resolved by adopting
_is_routed_expert_key(drops the now-dead closure). Rationale: it is the intended detector for the offload/dflash path and, by design, keeps the MTP-layer experts offloaded (the closure'sparse_expert_iddropped them via thelayer_id ≥ num_layersguard).Known limitation — tracked in #142
The adopted regex matches
.mlp.experts.<id>.but not Mixtral's.block_sparse_moe.experts.<id>.naming (see_remap_v5_batched_experts). Mixtral routed experts may therefore be misclassified as non-routed. Documented inline asNOTE(#142)and tracked in #142 (proposed fix: widen regex to\.(mlp|block_sparse_moe)\.experts\.\d+\.+ Mixtral smoke test).Verification
model_offload.py: no conflict markers,py_compileOK, LSP clean.qwen.py,gpt_oss.py(auto-merged)py_compileOK.Notes
devafter this lands.