Conversation
stikves
marked this pull request as draft
September 21, 2026 21:41
stikves
force-pushed
the
sukru/dflash-fused-export
branch
from
September 21, 2026 21:44
6a02190 to
a806a06
Compare
…+ structural-metadata validation)
stikves
force-pushed
the
sukru/dflash-fused-export
branch
from
September 21, 2026 22:21
a806a06 to
eee61c2
Compare
…d kwargs, namedtuple-safe logits cast
Rename _speculative_metadata to _speculative_decoding_metadata in bundle.py and update all call sites, plus the comment references in pipeline.py and the tests. Also trim the future-fallback paragraph from the add_pytorch_module rationale comment in macos.py.
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
2 participants
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
Adds a
--fused-targetexport option for the Muse Glimmer DFlash drafter, plus the--dflash-draftertwo-function drafter export and structural-metadata validation.--fused-target): exports the 2-outputMuseGlimmerForCausalLMWithDraftergraph emitting(logits, drafter_features)in one model. Adds tuple/list-aware bf16 to fp16 logits casting inBaseForCausalLM.--dflash-drafter): exports the two-entrypoint drafternn.Module:inject_kv(writes target features into a shared ring-KV cache) anddraft(reads them back and produces draft logits) — with per-layer mmap offload during export. Requires--fused-target; mutually exclusive with the ring--with-drafter.metadata.jsonis sourced explicitly from it (raises on divergence) rather than inferred.Why
The DFlash drafter has a custom architecture: the inputs to 5 of the target model's layers are shared with the drafter. Exporting the target and drafter as separate models means copying that shared data between them by hand, or wiring their States together, which the current API makes awkward.
--fused-targetavoids the separate data paths by fusing the target and drafter into one graph with two heads,logitsanddrafter_features, so the shared features flow through directly.Scope
This PR covers the export path only. The new CLI flags default off, so nothing changes unless you opt in. The Swift
DFlashDrafterengine, the speculative decoder cycle, and target integration come in separate later PRs, and this change makes no throughput claim on its own. It sits on top of the merged muse_glimmer drafter (#228), keeps that PR's bidirectional / no-embed-norm draft semantics, and does not regress it.Tests
test_dflash.py+test_stage2_fused_export.py, 33 passing. They cover the two-graph export contract, torch.export tracing of both entrypoints, the tuple-aware logits cast, speculative-metadata merge/validation, and a fused-target check thatdrafter_featuresequals the target's hidden state at the configured layers. All run on tiny configs without the 55GB model. ruff clean.