Generated: 2026-05-04 Purpose: Map every DFlash paper concept to the z-lab reference implementation, the Luce C++ GGUF implementation, and the dflash-robot target abstraction.
Sources:
- Paper: "DFlash: Block Diffusion for Flash Speculative Decoding" (arXiv 2602.06036v1)
- Paper notes: /home/am/dflash-gguf/notes/dflash-paper-first-principles.md
- z-lab code: /tmp/zlab-dflash-inspect/dflash/model.py, benchmark.py
- Luce code: /tmp/lucebox-hub-dflash-inspect/dflash/src/.cpp, test/.cpp
| Paper Concept | Paper Section | z-lab Implementation | Luce Implementation | dflash-robot Target Abstraction | Gap / Risk |
|---|---|---|---|---|---|
| Target prefill | Sec 3.1 (step 1-3) | model.py:dflash_generate() L74-100 — runs target(input_ids, output_hidden_states=True), samples first token, extracts hidden states for the prompt. Single batched forward. |
test_dflash.cpp L2114-2253 — token-segmented prefill loop (chunked ubatch) OR L1966-2113 layer-segmented prefill. Both iterate build_target_step() with capture=true to populate cache.target_feat ring buffer and KV/SSM caches. migrate_prefill_cache() promotes prefill-only cache to full decode cache. |
ModelAdapter.prefill(tokens, capture_layers=True) — must populate KV cache + target_feat buffer. Must support chunked prefill for long contexts. Layer-segmented optional. |
z-lab does one big batched forward; Luce chunks at ubatch boundary. dflash-robot needs to support both strategies and expose ubatch size as tunable. Hybrid SSM+attention targets (Qwen3.5) require migrate_prefill_cache for rollback tensor allocation — pure-attention targets won't need this. |
| Target hidden-state layer selection | Sec 3.2, App. A | model.py:build_target_layer_ids() L27-36 — computes num_draft_layers evenly-spaced layer IDs from [1, num_target_layers-3]. Single layer for 1-draft-layer, else linearly spaced. DFlashDraftModel.__init__() L312-314 reads dflash_config.target_layer_ids or calls build_target_layer_ids(). |
gguf_target_loader.cpp L344-349 — computes capture_layer_ids[k] = 1 + k * step where step = (n_layer - 2) / (N - 1), stored in TargetWeights::capture_layer_ids[DFLASH27B_DRAFT_N_TARGET_LAYERS]. qwen35_target_graph.cpp L1141-1177 build_single_layer() — checks w.capture_layer_ids[k] == layer_idx and copies cur into cache.target_feat ring buffer at the matched capture index. |
ModelAdapter.extract_features(layer_ids, token_range) — must expose hidden states from arbitrary selected layers. The layer ID list is draft metadata (stored in draft GGUF as dflash.n_target_layers). Must handle ring-buffer wraparound. |
Layer selection formula differs slightly: z-lab uses start=1, end=n-3; Luce uses 1 + k*((n-2)/(N-1)). These are equivalent for typical configs but dflash-robot must document which formula is canonical and store it in draft metadata. |
| Hidden feature concatenation/fusion | Sec 3.2, Fig 2 | model.py:extract_context_feature() L39-45 — selects hidden_states[layer_id + 1] for each target layer, concatenates along feature dim (torch.cat(..., dim=-1)). DFlashDraftModel.fc L317 — nn.Linear(N*hidden, hidden) projects concatenated features to hidden size. |
qwen3_dflash_graph.cpp L53-60 — target_feat = ggml_mul_mat(ctx, w.fc, in.target_hidden_cat) then ggml_rms_norm + ggml_mul(ctx, target_feat, w.hidden_norm). Input is [5*hidden, ctx_len] concatenated on device; fc projects to [hidden, ctx_len]. |
DraftAdapter.fuse_features(concatenated_hidden_states) — the fc projection and hidden_norm are draft-model weights. dflash-robot must load them from the draft GGUF and apply them in the draft graph. The concatenation layout [N*hidden, seq] must be consistent between target capture and draft input. |
The z-lab path is trivial (PyTorch cat+linear). Luce implements the same logic in ggml ops. dflash-robot must ensure the column-major layout of the target_feat ring buffer matches the fc weight's expected input shape. bf16 storage in Luce's ring buffer requires f32 widen before fc (done via CUDA kernel dflash27b_launch_bf16_to_f32). |
| Target-feature KV injection | Sec 3.2, Fig 2 | model.py:Qwen3DFlashAttention.forward() L211-255 — k_ctx = self.k_proj(target_hidden), v_ctx = self.v_proj(target_hidden), then k = torch.cat([k_ctx, k_noise], dim=1), v = torch.cat([v_ctx, v_noise], dim=1). Target features are projected through the SAME wk/wv as noise tokens and concatenated along the sequence dimension. |
qwen3_dflash_graph.cpp L79-88 — Kctx = ggml_mul_mat(ctx, L.wk, target_feat), Kn = ggml_mul_mat(ctx, L.wk, hn), Vctx = ggml_mul_mat(ctx, L.wv, target_feat), Vn = ggml_mul_mat(ctx, L.wv, hn), then K = ggml_concat(ctx, Kctx, Kn, 1), V = ggml_concat(ctx, Vctx, Vn, 1). Concatenation along sequence dim (ne[1]). |
DraftAdapter.inject_target_features(features, positions) — target features are injected into every draft layer's K/V via shared wk/wv projections, concatenated with noise K/V along sequence dim. This is a per-layer operation. The DraftGraphInputs.target_hidden_cat tensor drives this. |
This is the core DFlash mechanism. Both implementations agree: project target features through draft's wk/wv, concat with noise K/V. dflash-robot must ensure the concatenation order (target features first, then noise tokens) is consistent with position encoding. The RoPE positions for K must span [0..ctx_len+q_len-1] with target features at positions [0..ctx_len-1]. |
| Mask/noise embeddings | Sec 3.1 (step 8) | model.py:dflash_generate() L108-121 — block_output_ids = output_ids[:, start:start+block_size].clone() where position 0 = last_tok, positions 1..B-1 = mask_token_id. target.model.embed_tokens(block_output_ids) produces noise embeddings. Mask token ID from config. |
test_dflash.cpp L2319-2322 — noise_ids[0] = last_tok; for (int i=1; i<q_len; i++) noise_ids[i] = mask_tok; then w.embedder.embed(noise_ids, q_len, noise_embed_buf). Mask token = DFLASH27B_DRAFT_MASK_TOKEN_ID (constant). qwen3_dflash_graph.cpp L63 — ggml_tensor * h = in.noise_embed starts the draft from noise embeddings. |
DraftAdapter.propose(anchor_token, block_size, cache_state) — must construct [anchor, MASK, MASK, ...] token IDs, embed them via target's embedding table, and pass as input to draft graph. The mask token ID is draft metadata. |
z-lab embeds via target's embed_tokens (GPU). Luce embeds via CPU CpuEmbedder::embed() (mmap'd token_embd stays on CPU). dflash-robot must decide: GPU embedding (requires uploading token_embd to GPU) or CPU embedding (saves VRAM). For large vocab models, CPU embedding saves ~1GB VRAM but adds latency. |
| Non-causal block draft | Sec 3.1 (step 9), Sec 3.3 | model.py:DFlashDraftModel.forward() L323-347 — passes is_causal=False to attention. All positions in the block attend to each other bidirectionally. dflash_generate() L118 — is_causal=False in the draft call. |
qwen3_dflash_graph.cpp L120-124 — ggml_flash_attn_ext(ctx, Q, K, V, /*mask=*/nullptr, scale, 0.0f, 0.0f). No mask = non-causal. The draft graph explicitly passes nullptr for the attention mask. |
Draft graph must use non-causal (bidirectional) attention within the draft block. The is_causal=False / mask=nullptr is the key semantic. dflash-robot's draft graph builder must NOT apply a causal mask to the draft attention. |
Both implementations agree: no mask = non-causal. This is correct for block diffusion where all masked positions are denoised simultaneously. dflash-robot must document this clearly — it's the fundamental difference from autoregressive drafters like EAGLE. |
| Candidate token selection | Sec 3.1 (step 10) | model.py:dflash_generate() L121 — block_output_ids[:, 1:] = sample(draft_logits). Samples/argmaxs from draft output logits at positions 1..B-1. Position 0 stays as last_tok. sample() L48-54 — argmax for temperature<1e-5, else multinomial. |
test_dflash.cpp L2429-2442 — GPU-side ggml_argmax on draft logits (L1022-1026 in build_draft_step), or CPU argmax for split-GPU path. draft_tok[0] = last_tok pinned back. DDTree mode: extract_draft_topk() L165-229 extracts top-K per position for tree construction. |
DraftAdapter.propose() returns candidate token block. Must support greedy (argmax) and optionally top-K for tree-structured verify. GPU-side argmax preferred to avoid vocab-size CPU transfer. |
z-lab always uses argmax or multinomial sampling. Luce adds DDTree top-K extraction (up to 8 candidates per position) for tree-structured verification. dflash-robot must support both chain mode (argmax only) and tree mode (top-K), with tree mode as optional enhancement. |
| Target verify | Sec 3.1 (step 11) | model.py:dflash_generate() L126-132 — target(block_output_ids, position_ids, past_key_values, use_cache=True, output_hidden_states=True). Single batched forward over the entire candidate block. |
test_dflash.cpp L2838-2915 — batched verify: build_target_step(sg, w, cache, backend, committed, q_len, with_mask=true, capture=true), embeds draft tokens, runs one forward with causal mask. Sequential verify: L2883-2915 — q_len independent single-token decodes (for debugging). DDTree verify: L2528-2596 — build_target_step_tree() with ancestor-only mask + parent_ids for tree-mode DeltaNet. |
ModelAdapter.verify(candidate_tokens, cache_state, capture_layers=True) — runs target forward on candidate block with causal mask. Must update KV/SSM cache and capture target_feat for the verified positions. Returns logits for all positions. |
z-lab runs one batched forward. Luce supports three modes: batched (fast), sequential (correct reference), and DDTree (tree-structured). The batched vs sequential discrepancy (z-lab issue #57) is a numerical concern. dflash-robot should default to batched and offer sequential for debugging. |
| Acceptance prefix calculation | Sec 3.1 (step 12-13) | model.py:dflash_generate() L135 — acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1). Compares draft candidates[1:] to target posterior[:-1], cumprod finds longest prefix match. Standard speculative decoding comparison. |
test_dflash.cpp L2921-2954 — for (int i=0; i<q_len-1; i++) { if (draft_tok[i+1] == target_tok[i]) accept_n++; else break; }. Same logic: compare draft[i+1] vs target[i], accept longest matching prefix. DDTree: follow_verified_tree() L418-437 walks the tree from root following target.argmax at each node, collecting accepted path. |
RuntimeOrchestrator owns acceptance logic. Chain mode: linear prefix match. Tree mode: tree walk. Must report acceptance_length for metrics. | Both implementations agree on the comparison: draft[i+1] vs target[i]. The cumprod trick in z-lab is elegant but equivalent to the loop in Luce. DDTree adds tree-walk acceptance which can accept more tokens via sibling branches. |
| Bonus token handling | Sec 3.1 (step 14) | model.py:dflash_generate() L137 — output_ids[:, start + acceptance_length + 1] = posterior[:, acceptance_length]. Inserts the first mismatching target token as the bonus. Also handles all-accepted case: posterior[:, q_len-1] when acceptance == q_len-1. |
test_dflash.cpp L2936-2954 — two strategies: (1) Legacy: bonus_tok = target_tok[accept_n - 1], committed = accept_n + 1. (2) Fast-rollback: no explicit bonus, commit_n = accept_n, next last_tok = target_tok[commit_n - 1], bonus becomes next iter's draft_tok[0]. |
RuntimeOrchestrator. Must handle both explicit bonus (legacy path with replay) and implicit bonus (fast-rollback path). The bonus token is the target's correction at the first mismatch. | z-lab always commits accept_n + 1 (including bonus). Luce's fast-rollback path commits only accept_n and uses the target's next-token prediction as the implicit bonus for the next iteration. Both produce identical output streams. dflash-robot should support both paths — fast-rollback is preferred for performance. |
| Cache crop/rollback | Sec 3.1 (step 15) | model.py:dflash_generate() L120 — past_key_values_draft.crop(start) after draft forward. L139 — past_key_values_target.crop(start) after accepting. Uses HuggingFace DynamicCache.crop(). |
test_dflash.cpp L2487-2489 snapshot_ssm_state(cache) before verify (legacy path). L3080 restore_ssm_state(cache) to rollback. Fast-rollback: L2980-3057 uses per-step SSM intermediates (dflash27b_launch_f16_to_f32) to restore state at commit_n-1 without replay. Conv state rolled back via cudaMemcpy2DAsync. KV cache: L2677-2823 compaction via cudaMemcpyAsync for DDTree mode. qwen35_target_graph.cpp L328-340 snapshot_ssm_state() / restore_ssm_state() — device-side ggml_backend_tensor_copy for SSM+conv state. |
ModelAdapter.snapshot_cache(), rollback_cache(accepted_length), crop_cache(length). Must handle KV cache cropping (trivial for attention-only), SSM state rollback (requires snapshot/restore for hybrid models), and conv state rollback. The rollback mechanism depends on target architecture type. |
z-lab's crop is trivial (DynamicCache). Luce has three rollback strategies: (1) snapshot+restore (legacy), (2) per-step intermediate capture (fast-rollback), (3) DDTree tree-aware rollback with parent-chain conv reconstruction. dflash-robot must abstract rollback as architecture-dependent. Attention-only targets: simple KV crop. Hybrid SSM targets: intermediate-state capture. This is the biggest architecture-specific complexity. |
| Block size metadata | Sec 4.3, App. A | model.py:DFlashDraftModel.__init__() L319 — self.block_size = config.block_size. Stored in draft config. dflash_generate() L76 — block_size = model.block_size if block_size is None else block_size. Overridable at inference time. |
gguf_draft_loader.cpp L158 — block_sz = read_u32("dflash.block_size", 0). Loaded from draft GGUF metadata. L177-188 — validates against DFLASH27B_DRAFT_BLOCK_SIZE compile-time constant; rejects mismatched GGUFs. Used throughout as q_len = DFLASH27B_DRAFT_BLOCK_SIZE. |
DraftAdapter.training_block_size, DraftAdapter.max_supported_block_size. Stored in draft GGUF metadata (dflash.block_size). Runtime should allow smaller inference block sizes than training block size (paper finding: larger→smaller generalizes well). |
Luce currently hardcodes block_size as a compile-time constant and rejects drafts that disagree. dflash-robot must read block_size from draft metadata AND support runtime override to a smaller value. This requires the draft graph builder to handle variable q_len. |
| Draft depth metadata | Sec 4.2, App. A | model.py:DFlashDraftModel.__init__() L309-311 — self.layers = nn.ModuleList([Qwen3DFlashDecoderLayer(...) for layer_idx in range(config.num_hidden_layers)]). Depth = config.num_hidden_layers of the draft model. Paper ablates 1/3/5/8 layers. |
gguf_draft_loader.cpp L153 — n_layer = read_u32("block_count", 0). Draft layer count from GGUF metadata. Graph builder iterates for (int il=0; il<w.n_layer; il++) in qwen3_dflash_graph.cpp L65. |
DraftAdapter.depth or DraftAdapter.n_layer. Stored in draft GGUF metadata (qwen35-dflash-draft.block_count). Must be exposed as metadata for benchmarking draft-depth/acceptance trade-off. |
Both implementations read depth from config. dflash-robot must expose it as tunable metadata and benchmark different depths. Paper finding: 5 layers was best average speedup even when 8 had longer acceptance. |
| Position-decayed loss | Sec 3.4, App. B | Not directly in inference code. Paper describes training loss with exp(-decay * (position - anchor)) weighting. Decay hyperparameter: 7 for block_size 16, 5 for block_size 10, 4 for block_size 8. |
Not in Luce code (inference-only). | Not applicable for inference runtime. Should be documented in draft metadata schema for future training support. dflash.loss_decay in GGUF metadata. |
Training-only concern. dflash-robot should store the decay value in draft metadata for reproducibility but doesn't need it at inference time. No gap for inference; gap only if training is implemented later. |
| Random anchor training blocks | Sec 3.4 | Not in inference code. Paper: randomly sample anchor positions from responses, construct blocks with clean anchor + masked future positions. 512 anchor positions per sequence. | Not in Luce code (inference-only). | Not applicable for inference. Relevant only for future draft training pipeline. | Training-only concern. dflash-robot should design the target_feat cache format and draft metadata schema to be compatible with future training, but this doesn't affect inference. |
| Adaptive block-size opportunity | Sec 4.3 | Not implemented. Paper mentions: "adaptive block-size scheduling is a promising future optimization." Models trained with larger blocks generalize to smaller inference blocks. | Not implemented in Luce. Block size is compile-time constant. | RuntimeOrchestrator block-size policy. Should allow runtime block_size < training block_size. Future: adaptive scheduling based on recent acceptance lengths (e.g., shrink block when acceptance drops, grow when high). | Neither implementation supports adaptive block size. dflash-robot should expose the hook (block_size_policy(acceptance_history) → int) but start with a fixed block size. This is a low-risk future enhancement. The paper explicitly recommends benchmarking block 8, 10, 16 on local hardware. |
| File | Key Functions | Paper Concepts |
|---|---|---|
dflash/model.py |
build_target_layer_ids(), extract_context_feature(), sample(), dflash_generate(), Qwen3DFlashAttention, Qwen3DFlashDecoderLayer, DFlashDraftModel |
All inference concepts |
dflash/benchmark.py |
_run_transformers(), _run_mlx(), _run_server(), _print_decode_summary() |
Benchmarking, acceptance histogram |
| File | Key Functions | Paper Concepts |
|---|---|---|
src/qwen35_target_graph.cpp |
create_target_cache(), free_target_cache(), reset_target_cache(), migrate_prefill_cache(), snapshot_ssm_state(), restore_ssm_state(), snapshot_target_cache(), restore_target_cache(), snapshot_target_cache_thin(), restore_target_cache_chain(), build_full_attn_block(), build_delta_net_block(), build_single_layer(), build_qwen35_graph(), build_qwen35_layer() |
Target prefill, KV cache management, hidden-state capture, SSM/conv rollback, feature ring buffer |
src/qwen3_dflash_graph.cpp |
build_draft_graph() |
Feature fusion (fc + hidden_norm), non-causal attention, KV injection (wk/wv on target_feat + noise concat), RoPE, SwiGLU FFN |
src/gguf_target_loader.cpp |
load_target_gguf(), CpuEmbedder::embed() |
Target weight loading, capture_layer_ids computation, CPU embedding |
src/gguf_draft_loader.cpp |
load_draft_gguf() |
Draft weight loading, block_size/n_target_layers validation |
src/dflash_graph.h |
DraftGraphInputs, DraftGraphOutputs, build_draft_graph() |
Draft graph interface: noise_embed, target_hidden_cat, positions_q, positions_k, lm_head |
test/test_dflash.cpp |
main(), build_target_step(), build_draft_step(), build_target_step_tree(), build_lm_head_projection_step(), extract_draft_topk(), build_ddtree(), follow_verified_tree(), build_causal_mask(), build_tree_mask(), draft_feature_mirror_*() |
Full spec decode loop: prefill, noise block, draft forward, verify, accept, rollback, DDTree tree-structured verify |
test/test_generate.cpp |
main(), build_step_step(), run_step() |
Baseline autoregressive generation (no draft) |
Draft layer attention:
Q = wq @ noise_hidden → [q_dim, q_len]
K_ctx = wk @ target_feat → [kv_dim, ctx_len] (from target hidden states)
K_noise = wk @ noise_hidden → [kv_dim, q_len] (from draft hidden states)
K = concat[K_ctx, K_noise] → [kv_dim, ctx_len + q_len]
V = concat[V_ctx, V_noise] → [kv_dim, ctx_len + q_len]
RoPE(Q, positions_q); RoPE(K, positions_k)
attn = flash_attn(Q, K, V, mask=null) ← NON-CAUSAL
| Strategy | Used By | Mechanism | Overhead |
|---|---|---|---|
| Cache crop (DynamicCache.crop) | z-lab | Tensor slice | Trivial |
| Snapshot + restore | Luce legacy | Device-side tensor copy before verify; restore on partial accept | O(n_layers × state_size) |
| Per-step intermediate capture | Luce fast-rollback | Patched gated_delta_net kernel writes per-step states; rollback reads single slot | O(n_layers × state_size) but no pre-verify snapshot |
| DDTree tree-aware rollback | Luce DDTree | Intermediate capture + parent-chain conv reconstruction + KV/feat compaction | Most complex; enables tree-structured acceptance |
| Aspect | z-lab | Luce |
|---|---|---|
| Storage | In-memory PyTorch tensor | bf16 ring buffer in TargetCache::target_feat |
| Layout | [seq_len, N*hidden] (PyTorch row-major) |
[N*hidden, target_feat_cap] (ggml column-major) |
| Capacity | Full sequence | Ring buffer, default cap=4096 |
| Transfer to draft | Direct tensor reference | bf16→f32 widen via CUDA kernel or peer copy |
- Sampling/speculative sampling: Both use greedy argmax. Paper mentions sampling support requires proper rejection logic.
- Adaptive block size: Neither implements it. Paper recommends it as future work.
- Training pipeline: Neither z-lab inference code nor Luce implements training. Paper's training algorithm (random anchors, position-decayed loss) is documented but not coded.
- Multi-model generic adapter: Both are Qwen3/Qwen3.5-specific. dflash-robot's adapter pattern is the key differentiator.
-
ModelAdapter must support hybrid SSM+attention targets — rollback is fundamentally different for SSM layers vs attention layers. The
migrate_prefill_cache,snapshot_ssm_state,restore_ssm_statepattern from Luce must be generalized. -
Draft GGUF must carry metadata — block_size, n_target_layers, mask_token_id, loss_decay, training_block_size. Current Luce GGUF format stores these; dflash-robot should formalize the schema.
-
Feature mirror for multi-GPU — Luce's
DraftFeatureMirrorhandles the case where target and draft are on different GPUs. dflash-robot should support this via the adapter pattern. -
DDTree is an enhancement, not a requirement — The core DFlash algorithm is chain-mode (argmax draft → causal verify → prefix accept). DDTree adds tree-structured verify on top. dflash-robot should implement chain mode first, DDTree as optional.
-
Rollback abstraction is the hardest part — Attention-only targets (Llama, Qwen3) need only KV crop. Hybrid targets (Qwen3.5) need SSM/conv intermediate-state capture. This must be a first-class abstraction in the ModelAdapter.