[aqilmarwan]

Closes #156.

Makes dflash_config.target_layer_ids required instead of guessing it, and deletes build_target_layer_ids. Rebased onto 07ebd93. Full evidence in #156.

self.target_layer_ids = _draft_value(config, "target_layer_ids")
if self.target_layer_ids is None:
    raise ValueError("Draft config must define dflash_config.target_layer_ids.")

Why:

• The guess never runs. All 20 checkpoints I could download set the key.
• If it did run it would usually be wrong. It passes the draft's layer count where the formula needs the number of target layers read, which matches only 6 of 20 checkpoints (18 of 20 with the right count). A wrong guess sets fc's input width, so it surfaces as a weight shape mismatch rather than a missing config key.
• Because dict.get evaluates its default eagerly, the guess ran on every load, making the unread num_target_layers mandatory in every config. That's why the MLX DFlashConfig field goes too.

MLX already required the key with no fallback; torch now matches.

Verified by building the draft model from all 20 reachable configs, including both DFlash 2 ones, and checking target_layer_ids, fc.in_features, block_size and mask_token_id against each. All load, same as main. No released checkpoint changes behaviour.

[harshal-96]

I hit this same helper while porting DFlash to LMDeploy and reviewing the TensorRT-LLM and LMDeploy implementations, so I read this change closely.

Verified against the evidence in #156 and my own ports: every published checkpoint sets dflash_config.target_layer_ids, the taps are a training-time choice baked into fc's input width, and DFlash 2 (Qwen3.8-27B) proves matching counts no longer imply matching ids, so no load-time formula can be correct. Requiring the key and turning a shape mismatch into a named config error is the right call, and dropping the unused num_target_layers field from the MLX config matches.

One data point on ecosystem impact: this helper was copied into downstream engines. SGLang's srt/speculative/dflash_utils.py still carries it as a live fallback in resolve_target_layer_ids, where the count-collision case (right tap count, wrong layers) would load cleanly and silently read the wrong target layers. I am filing that separately against SGLang, referencing this issue. LMDeploy's in-flight DFlash PR carries the same fallback and I will flag it there too. Landing this in the reference repo gives those fixes something to point at.

LGTM.

[aqilmarwan]

gentle ping - @jianc99