fix: apply --draft-attn-impl when loading with --from-pretrained (vllm-project#755)

### Problem

`--draft-attn-impl` is silently dropped on the `--from-pretrained` load
path: HF configs never serialize `_attn_implementation`, and the load
call doesn't re-apply the CLI choice, so the loaded draft always
resolves the model default (`simple_flex_attention`). On backends
without FlexAttention this crashes during offline eval of trained
checkpoints. Details and reproduction in vllm-project#754.

### Fix

Re-apply the CLI selection on the loaded config before model
construction, mirroring the `from_training_args` mechanism
(`transformer_layer_config._attn_implementation`), on both the
weights-loading path and the config-only path. Applying the default
value is behavior-preserving, so this is a no-op for anyone not using
the flag.

### Validation

- Loading a DSpark checkpoint with `--from-pretrained --val-only
--draft-attn-impl <non-default>` now routes to the selected backend
(previously: flex dispatch error).
- Offline eval of a GLM-5.2 draft over six public benchmarks ran
end-to-end on this path; scoring reproducibility across independent runs
Δ=0.14%.

Fixes vllm-project#754

---------

Signed-off-by: Frozen7771 <133861939+Frozen7771@users.noreply.github.com>

Author: Frozen7771

scripts/train.py

@@ -366,6 +366,7 @@ def _build_from_config_only(
     t2d: torch.Tensor | None,
     d2t: torch.Tensor | None,
     verifier_name_or_path: str | None = None,
+    draft_attn_impl: str | None = None,
 ) -> SpeculatorModel:
     """Initialize a fresh draft from a saved speculator *config* (no weights).
 
@@ -374,6 +375,8 @@ def _build_from_config_only(
     no trained draft weights to restore (decoder weights are randomly initialized).
     """
     config = model_class.config_class.from_pretrained(path)
+    if draft_attn_impl is not None:
+        config.transformer_layer_config._attn_implementation = draft_attn_impl
     speculators_config = getattr(config, "speculators_config", None)
     # Fall back to the CLI --verifier-name-or-path only when the saved config has
     # no verifier path -- either null or blanked to "". A real path in the config
@@ -426,6 +429,23 @@ def build_draft_model(
                 t2d=t2d,
                 d2t=d2t,
                 verifier_name_or_path=args.verifier_name_or_path,
+                draft_attn_impl=(
+                    args.draft_attn_impl if args.speculator_type != "mtp" else None
+                ),
+            )
+        if args.speculator_type != "mtp":
+            # _attn_implementation is never serialized by HF configs, so re-apply
+            # the CLI selection before construction -- mirroring from_training_args.
+            # MTP is skipped: its from_training_args never sets the field and its
+            # __init__ resolves its own default ("eager") when it is absent.
+            config = model_class.config_class.from_pretrained(args.from_pretrained)
+            config.transformer_layer_config._attn_implementation = args.draft_attn_impl
+            return model_class.from_pretrained(
+                args.from_pretrained,
+                config=config,
+                t2d=t2d,
+                d2t=d2t,
+                verifier=args.verifier_name_or_path,
             )
         return model_class.from_pretrained(
             args.from_pretrained,

tests/unit/train/test_draft_config_init.py

@@ -434,6 +434,35 @@ def test_build_from_config_only_preserves_existing_verifier_name(tmp_path):
     assert built.config.speculators_config.verifier.name_or_path == "real-verifier"
 
 
+def test_config_roundtrip_drops_attn_implementation(tmp_path):
+    # Precondition of the --from-pretrained bug: HF configs never serialize
+    # _attn_implementation, so the field does not survive a save/load round-trip
+    # and must be re-applied from the CLI selection.
+    config = _make_eagle3_config()
+    config.transformer_layer_config._attn_implementation = "sdpa"
+    save_dir = tmp_path / "roundtrip"
+    config.save_pretrained(str(save_dir))
+
+    reloaded = Eagle3SpeculatorConfig.from_pretrained(str(save_dir))
+
+    assert reloaded.transformer_layer_config._attn_implementation != "sdpa"
+
+
+def test_build_from_config_only_reapplies_draft_attn_impl(tmp_path):
+    model_dir = _save_config_only_dir(tmp_path)
+
+    with patch.object(Eagle3DraftModel, "load_verifier_weights"):
+        built = _build_from_config_only(
+            Eagle3DraftModel,
+            str(model_dir),
+            None,
+            None,
+            draft_attn_impl="sdpa",
+        )
+
+    assert built.config.transformer_layer_config._attn_implementation == "sdpa"
+
+
 # ---------------------------------------------------------------------------
 # build_draft_model: MTP-from-scratch routing
 # ---------------------------------------------------------------------------