feat(dflash): converge spec-decode stack → dev (gpt-oss DFlash + route-ahead prototype + GLM-5.2 + perf-model) - #140
Closed
drunkcoding wants to merge 52 commits into
Closed
feat(dflash): converge spec-decode stack → dev (gpt-oss DFlash + route-ahead prototype + GLM-5.2 + perf-model)#140drunkcoding wants to merge 52 commits into
drunkcoding wants to merge 52 commits into
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)
… plots (PDF+PNG) + markdown report
Uncommitted tail from the glm-future-work session (green on tiny GLM tests): - deprecate/ignore glm_fp8_in_store flag; GLM-5.2-FP8 routed experts are always kept FP8 in the host store, resident weights (MLA/DSA/dense MLP/shared expert) dequantized to BF16 on load - rebuild FP8 block scales from checkpoint on reload-from-store - README GLM memory note; _glm_tiny quantize_shared option; test updates + reload parity test Baseline commit to isolate the upcoming resident-weight mis-map fix.
…ernorm mis-map) get_topology built each node's tensor_ids in named_parameters() order, but the archer store lays tensors out in checkpoint/offload (monotonic id) order. When a dense node spans an offload partition boundary, the C++ view loader (model_topology.cpp) packs/creates per-slot views in tensor_ids order and rebinds the wrong slot -> q_a_layernorm.weight[2048] materialized as the [256] router bias, crashing the first forward of the full GLM-5.2-FP8 (RuntimeError 256 vs 2048). Fix: sort each node's tensor_ids by id (== store order) in get_topology so node order matches the physical store layout. Minimal/Python-only; store, ReadTensor and fp8-in-store paths unchanged. Verified on full 753B zai-org/GLM-5.2-FP8: model now LOADED with zero 256-vs-2048 crashes (forward passes the previously-failing q_a_layernorm resident stage). Adds medium-scale reproducer (tests/python/integration/_glm_medium.py + test_glm_medium_forward.py, gated MOE_GLM_MEDIUM=1) preserving the collision dims q_lora_rank=2048 / n_routed_experts=256.
… only The sort must NOT touch expert nodes: their [gate_proj, up_proj, down_proj] tensors are consumed positionally by the fused MoE kernel, and sorting by id can reorder them at full scale -> 'fused_moe_ffn_into: gate/up proj K-dim mismatch'. Dense nodes (single id sublist) still get sorted to fix the q_a_layernorm mis-map; expert nodes (multiple sublists) keep their order. Verified on full 753B GLM-5.2-FP8 (sm_120 build): LOADED + generate(32) completes with no mis-map / no K-dim mismatch / no kernel-image errors.
…resident-weight fetch mis-map) The runtime resolved tensor identity via GetTensorId(data_ptr()); torch's caching allocator reuses freed addresses across the per-forward begin()/end() placeholder cycle, so distinct resident params transiently collided on a data_ptr -> wrong id -> begin() fetched the wrong stored tensor (a [256] router weight into q_b_proj [*,2048] -> 'size mismatch mat(1x2048) vec(256)'; also caused incoherent output). Fix: AcquireTensor/ReleaseTensor take an explicit tensor_id (default UINT32_MAX = old data_ptr fallback); the Python begin/end hooks pass the stable param.ar_id / buffer.ar_id. Backward-compatible (sentinel preserves prior behavior). Verified on full 753B GLM-5.2-FP8 (sm_120): the previously-crashing chat prompt now generates coherent, on-topic output; no q_b_proj shape crash. Offload/medium/tiny/ dsa tests pass.
parse_expert_id() returns None for the MTP layer (layer_id >= num_layers), so on a fresh fp8-in-store offload the MTP layer's routed experts were dequantized to BF16 instead of being kept FP8 for the dispatcher. Classify routed experts by the .mlp.experts.<id>. name pattern instead, so ALL routed experts (main + MTP layers) stay FP8 and record their block scales; shared experts, MLA attention, DSA indexer and dense MLPs still dequantize to BF16. Verified on a fresh real zai-org/GLM-5.2-FP8 store: 58368 routed FP8 scales recorded (57600 main + 768 MTP), coherent generation.
… (spec-off identical)
…pture for spec-decode
…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.
… integration + drafter-driven contract (#139) * feat(dflash): route-ahead expert prefetch (A1-A5) 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. * feat(dflash): sampled (non-greedy) speculative decoding accept ops Distribution-preserving block-diffusion accept rule for temperature/top-k/top-p, keeping greedy byte-identical. * feat(serving): DFlash speculative decoding in continuous-batching engine 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. * feat(dflash): speculator core — batched decoding, sampled/route-ahead 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. * test(dflash): GPU-gated 20b losslessness + serving-vs-sync harnesses MOE_DFLASH_GPU-gated: gpt-oss-20b native DFlash == plain greedy (token-identical), and serving-path DFlash == sync DFlash. * style(dflash): apply ruff-format to new spec-decode/serving files * chore(format): pre-commit run --all-files (repo-wide ruff-format + clang-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. * chore(format): pre-commit convergence (ruff-format idempotency) * fix(prefetch): guard null derefs in ReleaseTensor 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. * feat(qwen3.5): enable offloaded text-only serving 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. * feat(dflash): offloaded Qwen3.5 speculative decoding + route-ahead 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. * chore(format): ruff check --fix + ruff-format convergence (0.6.9) --------- Co-authored-by: drunkcoding <leyang.xue@ed.ac.uk>
drunkcoding
added a commit
that referenced
this pull request
Aug 10, 2026
… (#143) * feat(dflash): add DFlash speculative decoding module for GPT-OSS Add moe_infinity/spec_decode with DFlashSpeculator, DFlash config reader, and target/draft pairing validation for z-lab GPT-OSS DFlash drafters. * docs(dflash): add runnable GPT-OSS speculative decoding example * test(dflash): add greedy-agreement eval harness * feat(glm): register glmmoedsa architecture (GlmMoeDsaForCausalLM, type=5, transformers>=5.12 guarded) * refactor(fp8): relocate dequant_fp8_blockwise to utils/fp8.py (back-compat re-export preserved) * feat(glm): hf_config parse_moe_param/parse_expert_id GLM branch (MTP-78 skip guard) * feat(glm): SyncGlmMoeDsaMoEBlock wrapper (HF-gate routing delegation) * feat(glm): FP8 block dequant-on-load for GLM experts (honors modules_to_not_convert) * feat(glm): DSA indexer classification + IndexShare ownership map * feat(glm): wire GlmMoeDsa into offload engine (patch/unpatch, isinstance, first_k_dense, eager attn, MTP-safe) * docs(glm): add GLM-5.2 to supported models + FP8 offload usage * test(glm): verify MLA/DSA attention tensors stay resident (not misrouted as experts) * test(glm): base contract chain + RAM/GPU-gated e2e smoke (skips until fp8-in-store) * test(glm): verify IndexShare needs no offload change (indexers only on 'full' layers) * feat(glm): fp8-in-store Python layer (glm_fp8_in_store flag, scale sidecar, SetScales interface) * test(glm): budget-split coherence + gated 32k long-context prefill smoke * test(glm): consolidated DSA/IndexShare/MLA integration test (real-config consistency) * feat(glm): DFlash adapter hook (availability + pairing validation, reuses dflash.py) * feat(glm): native FP8 e4m3 block-scale dequant kernel + set_scales binding (bit-exact vs python; dispatcher H2D integration deferred) * fix(glm): model-load fixes exposed by real load (DeepseekV2Moe rename patch, GLM use_native_engine=False, e_score_correction_bias device guard) * test(glm): tiny synthetic GLM reproducer (gated MOE_GLM_TINY) — reproduces native expert-dispatch SIGSEGV * fix(glm): register routed experts (add experts ModuleList + first_k_dense offset) — fixes native dispatch SIGSEGV; tiny GLM now generates end-to-end * feat(glm): MTP speculative decoding (built-in layer 78 drafter, lossless-by-construction; validated on tiny model) * feat(glm): MTP spec-decode instrumentation (accept-length tau, per-step stats, expert-fetch hook) — parity preserved * test(glm): forked spec-decode test runner (MTP lossless + stats + DFlash adapter, process-isolated) * feat(glm): validate GLM serving via OpenAI api_server_v2 (works as-is on non-paged HF path; +tokenizer save +gated test) * test(glm): serving smoke covers completions + chat + streaming (one server, process-isolated) * feat(glm): native fp8-in-store dequant-on-copy (SetScales + MoEMLP dequant, gated) — Option A==B parity on tiny FP8 GLM; completes T15/T16 * feat(perf-model): minimal roofline package + GLM-5.2 ModelParams/decode roofline (scaffold, cross-plan option b) * feat(perf-model): GLM benchmark harness (tiny-model decode + MTP on/off + predicted roofline -> CSV) * feat(perf-model): GLM perf validation/summary + roofline & throughput plots (PDF+PNG) + markdown report * chore(glm): finalize fp8-in-store always-on tail Uncommitted tail from the glm-future-work session (green on tiny GLM tests): - deprecate/ignore glm_fp8_in_store flag; GLM-5.2-FP8 routed experts are always kept FP8 in the host store, resident weights (MLA/DSA/dense MLP/shared expert) dequantized to BF16 on load - rebuild FP8 block scales from checkpoint on reload-from-store - README GLM memory note; _glm_tiny quantize_shared option; test updates + reload parity test Baseline commit to isolate the upcoming resident-weight mis-map fix. * fix(offload): bind resident params to correct offload id (GLM q_a_layernorm mis-map) get_topology built each node's tensor_ids in named_parameters() order, but the archer store lays tensors out in checkpoint/offload (monotonic id) order. When a dense node spans an offload partition boundary, the C++ view loader (model_topology.cpp) packs/creates per-slot views in tensor_ids order and rebinds the wrong slot -> q_a_layernorm.weight[2048] materialized as the [256] router bias, crashing the first forward of the full GLM-5.2-FP8 (RuntimeError 256 vs 2048). Fix: sort each node's tensor_ids by id (== store order) in get_topology so node order matches the physical store layout. Minimal/Python-only; store, ReadTensor and fp8-in-store paths unchanged. Verified on full 753B zai-org/GLM-5.2-FP8: model now LOADED with zero 256-vs-2048 crashes (forward passes the previously-failing q_a_layernorm resident stage). Adds medium-scale reproducer (tests/python/integration/_glm_medium.py + test_glm_medium_forward.py, gated MOE_GLM_MEDIUM=1) preserving the collision dims q_lora_rank=2048 / n_routed_experts=256. * refine(offload): scope resident-weight tensor_ids sort to DENSE nodes only The sort must NOT touch expert nodes: their [gate_proj, up_proj, down_proj] tensors are consumed positionally by the fused MoE kernel, and sorting by id can reorder them at full scale -> 'fused_moe_ffn_into: gate/up proj K-dim mismatch'. Dense nodes (single id sublist) still get sorted to fix the q_a_layernorm mis-map; expert nodes (multiple sublists) keep their order. Verified on full 753B GLM-5.2-FP8 (sm_120 build): LOADED + generate(32) completes with no mis-map / no K-dim mismatch / no kernel-image errors. * fix(offload): bind begin/end to stable tensor id (fixes intermittent resident-weight fetch mis-map) The runtime resolved tensor identity via GetTensorId(data_ptr()); torch's caching allocator reuses freed addresses across the per-forward begin()/end() placeholder cycle, so distinct resident params transiently collided on a data_ptr -> wrong id -> begin() fetched the wrong stored tensor (a [256] router weight into q_b_proj [*,2048] -> 'size mismatch mat(1x2048) vec(256)'; also caused incoherent output). Fix: AcquireTensor/ReleaseTensor take an explicit tensor_id (default UINT32_MAX = old data_ptr fallback); the Python begin/end hooks pass the stable param.ar_id / buffer.ar_id. Backward-compatible (sentinel preserves prior behavior). Verified on full 753B GLM-5.2-FP8 (sm_120): the previously-crashing chat prompt now generates coherent, on-topic output; no q_b_proj shape crash. Offload/medium/tiny/ dsa tests pass. * fix(glm): keep MTP-layer routed experts FP8 on fresh offload parse_expert_id() returns None for the MTP layer (layer_id >= num_layers), so on a fresh fp8-in-store offload the MTP layer's routed experts were dequantized to BF16 instead of being kept FP8 for the dispatcher. Classify routed experts by the .mlp.experts.<id>. name pattern instead, so ALL routed experts (main + MTP layers) stay FP8 and record their block scales; shared experts, MLA attention, DSA indexer and dense MLPs still dequantize to BF16. Verified on a fresh real zai-org/GLM-5.2-FP8 store: 58368 routed FP8 scales recorded (57600 main + 768 MTP), coherent generation. * refactor(engine): extract _generate_standard + add spec_strategy seam (spec-off identical) * feat(engine): rich on-device forward helper + 5-layer hidden-state capture for spec-decode * feat(dflash): accept-rule + block-build pure fns with hand-checked unit tests * feat(dflash): harden drafter loader (trust_remote_code + dim/vocab/mask asserts) * test(dflash): tiny synthetic gpt-oss target + drafter fixtures (CPU determinism) * feat(dflash): native draft->verify->rollback state machine (bonus emitted-not-cached) * fix(dflash): tiny-target sliding_window 8->128 (>= block_size+1) for reconstructible sliding-window KV rollback * feat(dflash): edge cases (accept 0/9, EOS mid-block, max_new_tokens, batch==1 guard) * feat(dflash): wire native speculator as GenerationEngine spec_strategy * test(dflash): tiny-model E2E losslessness parity (native == plain greedy) * test(engine): spec-off byte-identity regression + gpt-oss suite green * test(dflash): GPU-gated 120B validation harness (agreement-rate/accept-len/tok-s) * docs(dflash): native-path example + usage notes (trust_remote_code, TP=2 SM120) * fix(dflash): gather 5-layer feature + posterior to drafter device (multi-GPU TP>1) * fix(gpt-oss): resident-load MXFP4 experts + sinks + router 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. * feat(dflash): route-ahead expert prefetch (A1-A5) 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. * feat(dflash): sampled (non-greedy) speculative decoding accept ops Distribution-preserving block-diffusion accept rule for temperature/top-k/top-p, keeping greedy byte-identical. * feat(serving): DFlash speculative decoding in continuous-batching engine 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. * feat(dflash): speculator core — batched decoding, sampled/route-ahead 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. * test(dflash): GPU-gated 20b losslessness + serving-vs-sync harnesses MOE_DFLASH_GPU-gated: gpt-oss-20b native DFlash == plain greedy (token-identical), and serving-path DFlash == sync DFlash. * style(dflash): apply ruff-format to new spec-decode/serving files * chore(format): pre-commit run --all-files (repo-wide ruff-format + clang-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. * chore(format): pre-commit convergence (ruff-format idempotency) * fix(prefetch): guard null derefs in ReleaseTensor 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. * feat(qwen3.5): enable offloaded text-only serving 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. * feat(dflash): offloaded Qwen3.5 speculative decoding + route-ahead 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. * chore(format): ruff check --fix + ruff-format convergence (0.6.9) --------- Co-authored-by: drunkcoding <leyang.xue@ed.ac.uk>
Contributor
Author
|
Superseded by #143. |
This was referenced Aug 14, 2026
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.
Summary
Converges the long-lived
feat/dflash-spec-decodeintegration branch intodev— the first step of thespec-decode → dev → mainlanding flow. 52 commits ahead ofdev, 0 behind (clean merge), formatting green (ruff 0.6.9).This aggregates the DFlash speculative-decoding work (PRs #139 merged + #135 closed into this branch) plus GLM-5.2-FP8, the perf-model harness, and offload correctness fixes.
What lands
DFlash speculative decoding (gpt-oss) — implements design doc #131
GenerationEngine.spec_strategyRoute-ahead prefetch + serving (tracks A–D, #139) — prototypes design doc #132 §9
_prefetch_route.py,_route_ahead_ctx.py,_route_ahead_stats.py)_dflash_sample_ops.py)GLM-5.2-FP8 (glm_moe_dsa)
glm_moe_dsa+glm_dsaarchitecture; FP8-in-store dequant-on-copy; DSA indexer / IndexShare; MTP speculative decoding; native FP8 block-scale dequant kernel; offload wiringOffload fixes
Not included / deferred
Known follow-up
feat/glm-5.2-fp8-support): this branch and feat(glm): support zai-org/GLM-5.2-FP8 (glm_moe_dsa) with FP8-in-store expert offload #130 carry divergentglm_moe_dsaimplementations. feat(glm): support zai-org/GLM-5.2-FP8 (glm_moe_dsa) with FP8-in-store expert offload #130 has unique work not here (on-GPU FP8 dequant-after-fetch, partition-read coalescing, NVTX3 auto-detect,test_glm_moe_dsa.py). To be reconciled separately after this converges — feat(glm): support zai-org/GLM-5.2-FP8 (glm_moe_dsa) with FP8-in-store expert offload #130 will need rebase/cherry-pick or closure.Flow
spec-decode → dev(this PR) → letdevconverge (with #136/#137/#138) → singledev → main.