Module dflash_head 

2026-09-25: DFlash block-diffusion draft head (arXiv 2602.06036), implementing DraftProposer.

A small Qwen3-style transformer drafter that emits a block of draft tokens in one forward pass with bidirectional in-block attention. It is conditioned on target hidden states captured at target_layer_ids, projected by one fc layer at the drafter’s input.

Owner: model-arch (DFlash drafter). Invariants:

• block_g() stays in 2..=gamma.max(2) (set_block_g clamps).

Modules

levers

2026-09-25: DFlash drafter levers, resolved from the environment once per head and carried on it.

Structs

BlockDiffusionDraftHead

2026-09-25: Block-diffusion draft head; its public API is the DraftProposer trait. It uses the target’s token embedding and LM head (embed_tokens_shared, lm_head_shared / lm_head_nvfp4) alongside the drafter checkpoint’s own fc, hidden_norm, norm and layers. The geometry fields come from the drafter’s config.

DflashKernels

2026-09-25: Kernel handles for the DFlash drafter, all resolved once in BlockDiffusionDraftHead::from_weights. A handle resolved there with try_kernel is KernelHandle(0) when the target’s build lacks the kernel.

DflashLayer

2026-09-25: One drafter layer’s weights: the bf16 norms (with per-head Q/K RMSNorm), projections and MLP; FP8 copies of the seven GEMM weights when the head is Fp8Weights (None otherwise); and the DFlash2 conv weights (attention_conv.*, mlp_conv.*), None when the checkpoint has none.

DflashProposerState

2026-09-25: Per-sequence DFlash drafter state. One block table serves every drafter layer: the paged cache has a pool per layer, addressed by the same slots.

DflashScratch

2026-09-25: Device scratch for the γ-block forward, allocated once per head in from_weights.rs. nb below is the head’s max_batch.

ProposeGraphs

2026-09-25: The propose graphs per block width (see BlockDiffusionDraftHead::propose_graphs). warmup counts the eager passes run at a width; its capture starts once the count reaches DFlashLevers::propose_warmup_n (METRALE_DFLASH_PROPOSE_WARMUP_N, default 2).

Enums

DflashQuantization

2026-09-25: Drafter weight precision. from_weights builds the head as Bf16 and switches it to Fp8Weights unless METRALE_DFLASH_DRAFTER_FP8=0 or the fp8_gemm_n128_row_scaled kernels are missing.

Functions

dflash_ctx_cap

2026-09-25: The DFlash context-window bound, in tokens: the most recent target positions the drafter may accumulate. Both the per-sequence ctx accumulator (proposer.rs) and the model’s whole-prompt hidden capture (metrale-model-engine’s impl_a1.rs) are sized from it.