Module dflash_head 

DFlash block-diffusion draft head implementing DraftProposer.

Block-diffusion drafter (Z Lab, arXiv 2602.06036): a small Qwen3-architecture transformer (8 layers, hidden=2048, GQA 32:4, head_dim=128) that emits γ=16 tokens in a single forward pass via bidirectional in-block attention. Conditioned on five intermediate hidden states captured from the target model at target_layer_ids (e.g., [1, 10, 19, 28, 37] for Qwen3.6-35B-A3B-DFlash), projected through a single fc layer at model entry — NOT per-layer KV injection (early plan was wrong; cf. vLLM qwen3_dflash.py).

Phase 1 deliverable: type + trait wiring. The actual γ-block forward kernel (inferspark_dflash_block_attn_fp8) lands in Phase 2; until then propose() returns the bonus token repeated num_drafts times so the verify path degenerates to single-token decode (acceptance ~100% but no speedup).

Structs

BlockDiffusionDraftHead

Block-diffusion draft head. Public API is the DraftProposer trait.

DflashKernels

Kernel handles for the DFlash γ-block forward chain. All resolved once at BlockDiffusionDraftHead::from_weights against the active GPU backend (which compiles target-specific PTX at startup); subsequent propose() calls just KernelLaunch::new(...).launch(stream).

DflashLayer

Per-drafter-layer Qwen3-style weights. Phase 1 is BF16-only; Phase G (2026-05-28) adds optional FP8 weight fields populated at model load when ATLAS_DFLASH_DRAFTER_FP8=1. The BF16 fields are always present (Fp8 path falls back to them for any GEMM whose Fp8 weight is None).

DflashProposerState

Per-sequence DFlash drafter state. One paged KV cache per drafter layer (8 typical), shared block table across layers since attention shape is identical layer-to-layer for a vanilla Qwen3 architecture. Mirrors MtpProposerState in spirit; the multi-layer cache keeps it distinct.

DflashScratch

Per-step scratch buffers for the γ-block forward.

Enums

DflashQuantization

Drafter-side weight precision. Defaults to BF16. Phase G (2026-05-28) adds Fp8Weights, gated by env var ATLAS_DFLASH_DRAFTER_FP8. The historical SM12.x acceptance collapse note applied to drafter FP8 KV cache (different concern — bidirectional attention math); Phase G targets weight FP8 only, so the risk surface is dynamic-range loss in MLP intermediate activations, which per-row scales mitigate. --mtp-quantization fp8 is still not honored for the DFlash drafter.

Functions

dflash_ctx_cap

The DFlash context-window bound, in tokens: the most recent target positions the drafter is allowed to accumulate and attend to.