Function load_dflash_weights 

pub fn load_dflash_weights(
    drafter_store: &WeightStore,
    drafter_config: &DflashConfig,
    _gpu: &dyn GpuBackend,
    _tp_size: usize,
) -> Result<Option<DflashWeights>>

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

The drafter ships its weights at the root of the safetensors file (no model. prefix), in the same naming convention as a vanilla Qwen3 transformer minus embed_tokens and lm_head. Atlas’s runtime fills those two from the target model’s embedding / LM head at construction time — exactly mirroring vLLM’s “absent in checkpoint → skip_substrs → share with parent” flow.

The probed key list (verified against z-lab/Qwen3.6-35B-A3B-DFlash):

  fc.weight                                              [H, 5*H_target]
  hidden_norm.weight                                     [H]
  norm.weight                                            [H]
  layers.{0..L-1}.input_layernorm.weight                 [H]
  layers.{0..L-1}.post_attention_layernorm.weight        [H]
  layers.{0..L-1}.self_attn.q_proj.weight                [Q*Hd, H]
  layers.{0..L-1}.self_attn.k_proj.weight                [Kv*Hd, H]
  layers.{0..L-1}.self_attn.v_proj.weight                [Kv*Hd, H]
  layers.{0..L-1}.self_attn.o_proj.weight                [H, Q*Hd]
  layers.{0..L-1}.self_attn.q_norm.weight                [Hd]
  layers.{0..L-1}.self_attn.k_norm.weight                [Hd]
  layers.{0..L-1}.mlp.gate_proj.weight                   [I, H]
  layers.{0..L-1}.mlp.up_proj.weight                     [I, H]
  layers.{0..L-1}.mlp.down_proj.weight                   [H, I]

where H=2048, H_target=2048, Q=32, Kv=4, Hd=128, I=6144, L=8 for Qwen3.6-35B-A3B-DFlash.

Under TP the drafter is replicated, not sharded — tp_size>1 produces the same per-rank result as tp_size=1. Memory cost: ~948 MB BF16 per rank, trivially below the 119 GB GB10 budget.