Module dflash_loader 

DFlash drafter weight loader.

Loads z-lab/Qwen3.6-{27B,35B-A3B}-DFlash-style drafter checkpoints into the typed DflashWeights structure consumed by crate::layers::BlockDiffusionDraftHead. The drafter is a small Qwen3-architecture transformer (8 layers, hidden=2048, GQA 32:4) with these distinctive parts vs. a vanilla Qwen3:

• model.fc — [len(target_layer_ids) * target_hidden, draft_hidden] BF16 projection that maps the stack of captured target hidden states into the drafter’s input space.
• model.hidden_norm — RMSNorm applied to the projected target context before mixing with token embeddings.
• lm_head — drafter ships its own (NOT tied to target’s), so tie_word_embeddings=false.
• Optional d2t — draft-vocab → target-vocab id remap (absent when drafter shares vocab with target, as in Qwen3.6-35B-A3B-DFlash where both = 248320).
• Special mask_token_id (248070 for Qwen3.6-DFlash) used for the γ “to-be-predicted” positions in block diffusion.

Under TP the drafter is not sharded — it’s small (~1–2 GB BF16), every rank loads the full set. Mirrors the existing MTP-under-TP pattern (MTP loads ALL experts on every rank — no EP all_reduce needed).

Structs

DflashConfig

Drafter HF config.json (subset Atlas consumes). Mirrors z-lab/Qwen3.6-35B-A3B-DFlash/config.json field names verbatim so serde_json::from_str works directly on the raw file.

DflashLayerWeights

Per-drafter-layer raw weights (BF16). Same shape across all 8 layers.

DflashRopeScaling

Subset of HF rope_scaling block consumed by Atlas. Mirrors the field names in transformers’ Qwen3 config so serde_json::from_str works directly on the drafter’s config.json.

DflashSubConfig

Nested dflash_config block in the drafter’s config.json.

DflashWeights

Raw weight bundle for the DFlash drafter, post-load.

Functions

load_dflash_weights

Load DFlash drafter weights from a separate WeightStore pointing at the drafter checkpoint.

parse_dflash_config

Parse a DFlash drafter’s config.json into a DflashConfig. Used by main.rs after fetching the drafter’s HF metadata to size the runtime BlockDiffusionDraftHead (layer count, head_dim, vocab_size, the target_layer_ids capture indices).

store_has_dflash_weights

Probe a WeightStore for the presence of DFlash drafter weights. Returns true if the store contains the unique fc.weight tensor that DFlash drafters ship — a lightweight detection that doesn’t load any data. Both bare-key and model.-prefixed layouts are accepted; the canonical z-lab/Qwen3.6-{27B,35B-A3B}-DFlash checkpoints ship the bare layout (verified against commit 42d3b34, May 2026).