DFlashDraftModel¶
Model type |
|
Class |
|
Task |
|
Source |
|
Description¶
DFlash drafter — a stack of cross-attending Qwen3-style blocks.
Inputs (graph-level, set up by :class:~mobius.tasks.DFlashDraftTask):
noise_embedding: [batch, q_len, hidden] — the target’s
embed_tokens(block_output_ids) where
block_output_ids = [prev_token, mask, ..., mask].
target_hidden: [batch, ctx_len, num_target_layers * hidden] —
concatenated post-residual hidden states of selected target
decoder layers, as configured by
:attr:DFlashConfig.target_layer_ids. Projected down to
hidden by self.fc and RMSNorm’d by self.hidden_norm.
position_ids: [batch, ctx_len + q_len] — absolute positions
for the K side (covers context tokens followed by noise tokens).
q_position_ids: [batch, q_len] — absolute positions for the
Q side (noise tokens only).
past_key_values: per-layer (key, value) cache pairs from prior
speculative steps.
Outputs:
draft_output: Final hidden states [batch, q_len, hidden] for the
target LM head, or logits [batch, q_len, draft_vocab] when the
checkpoint owns output.weight.
present_key_values: updated per-layer cache pairs.
Usage¶
mobius build --model <MODEL_ID> output_dir/
from mobius import build
model = build("<MODEL_ID>")