feat: DFlash2 block-diffusion drafter for Qwen 3.8 (opt-in, byte-identical) - #22
Merged
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 76b711b25b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
NeelM0906
added a commit
that referenced
this pull request
Aug 26, 2026
The config guards accepted a checkpoint with the expected vocabulary and target-layer count but a different hidden_size (Codex P2 on PR #22). The drafter shares the target's embedding/LM head and consumes tap captures of the target's hidden rows, so the mismatch would mis-stride those kernels or read out of bounds on the GPU instead of failing cleanly. Validate hidden_size against the target before any buffer is built. Config-only test fixtures prove the guards fire before model.safetensors is opened; without the guard the mismatched checkpoint sails through to the weights load.
NeelM0906
force-pushed
the
claude/qwen38-dflash2-drafter
branch
from
August 26, 2026 17:36
76b711b to
db38166
Compare
NeelM0906
added a commit
that referenced
this pull request
Aug 26, 2026
The config guards accepted a checkpoint with the expected vocabulary and target-layer count but a different hidden_size (Codex P2 on PR #22). The drafter shares the target's embedding/LM head and consumes tap captures of the target's hidden rows, so the mismatch would mis-stride those kernels or read out of bounds on the GPU instead of failing cleanly. Validate hidden_size against the target before any buffer is built. Config-only test fixtures prove the guards fire before model.safetensors is opened; without the guard the mismatched checkpoint sails through to the weights load.
NeelM0906
force-pushed
the
claude/qwen38-dflash2-drafter
branch
from
August 26, 2026 18:15
db38166 to
e919d79
Compare
Ports z-lab model_mlx.py semantics to Metal: BF16 multi-x GEMV, block attention over [sliding ctx | bidirectional block], grouped dynamic conv (prepare/finish), top-16 + CPU path selector, tap-gather staging, and an mmap safetensors loader. FP32 residual stream (o/down projections overflow FP16 on outlier channels — the Inkling residual lesson). Every kernel is CPU-reference tested; the assembled core and selector match the fp32 MLX reference (fixtures under scratch/, env-gated) — the bf16 reference disagrees with itself-in-fp32 by more than this port does. Not wired into decode yet.
Verify stages tap rows for the drafter; accept shrinks to the GDN restore + tap commit; continuity breaks (plain-decode gaps, rewinds) reset the drafter with positions realigned — quality-only, bytes unchanged. Draft capacity defaults to 6 when the drafter is active.
… taps The BF16 drafter (3.85 GB) both thrashed the 24 GB host against the 15 GB target and ran hand-rolled scalar GEMVs at ~8.6 GB/s — decode collapsed to 0.6 tok/s. Projections now quantize at load into the production multi-x INT4 format (disk-cached beside the checkpoint), small tensors copy to an anonymous GPU buffer, and no command buffer touches file-backed memory. Draft rounds: 2.13 s -> 32 ms. Prefill now stages the prompt's trailing context window of tap rows so the first rounds draft warm, and spec stats gain a per-draft-position accept profile. Measured (M5 24 GB, greedy 96-128 tok, same-session A/B vs MTP k=3): bench-protocol chat prompt 11.81 vs 11.38 tok/s (+3.7%), accept profile 75/53/42/33/25/25 vs 75/34/18; raw code prompt 13.8 vs 15.6 (-11%). Bytes identical everywhere. The limiter is the verify pass's linear per-row cost (~15 ms/row of sequential attention+GDN dispatches), which taxes k=6 harder than MTP's k=3 — batching verify rows is the unlock.
The config guards accepted a checkpoint with the expected vocabulary and target-layer count but a different hidden_size (Codex P2 on PR #22). The drafter shares the target's embedding/LM head and consumes tap captures of the target's hidden rows, so the mismatch would mis-stride those kernels or read out of bounds on the GPU instead of failing cleanly. Validate hidden_size against the target before any buffer is built. Config-only test fixtures prove the guards fire before model.safetensors is opened; without the guard the mismatched checkpoint sails through to the weights load.
NeelM0906
force-pushed
the
claude/qwen38-dflash2-drafter
branch
from
August 26, 2026 20:13
e919d79 to
505f83c
Compare
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.
Stacked on #21 (retarget as the chain merges). Adopts the DFlash2 drafter (z-lab/Inco, ICML 2026) as an alternative draft source for the Qwen 3.8 speculative rounds — the verify/accept/rollback machinery and therefore every emitted byte is unchanged for any draft quality.
What's here
Qwen38DFlash2Drafter,Metal/DFlash2/): 5-layer block-diffusion forward with KV-injected target taps, bidirectional block attention, grouped dynamic conv, top-16 + CPU path selector, mmap safetensors loader. FP32 residual stream (o/down projections overflow FP16 on outlier channels).DFlash2Int4Slab): projections quantize into the production multi-x format, disk-cached beside the checkpoint. The BF16 drafter (3.85 GB) thrashed the 24 GB host and reproduced the DSV4 file-backed wiring collapse — draft rounds went 2.13 s → 32 ms with the slab.MFERENCE_DFLASH2_DIRenables; per-draft-position accept telemetry added toMFERENCE_PHASES=1.Gates
scratch/).Measured (M5 24 GB, same-session A/B vs MTP k=3, greedy)
Tokens/round 3.53 vs 2.27 on the bench prompt — the drafts are genuinely deeper (accept profile 75/53/42/33/25/25 vs 75/34/18), consistent with llama.cpp's independent +3%-on-prose finding. The limiter is the verify pass's linear per-row cost (~15 ms/row of sequential attention+GDN dispatches), which taxes k=6 harder than k=3.
Follow-ups (docs/QWEN38_DFLASH2.md)
🤖 Generated with Claude Code