Upload dflash_mlx/universal.py

dflash_mlx/universal.py

@@ -16,6 +16,23 @@ from .adapters import load_target_model, LoadedTargetModel, detect_model_archite
 from .convert import load_mlx_dflash
 
 
+def _build_target_layer_ids(num_target_layers: int, num_draft_layers: int) -> List[int]:
+    """Select target model layer indices for feature extraction.
+    
+    Uniformly samples from shallow to deep layers for cross-layer
+    feature fusion, matching the DFlash paper.
+    """
+    if num_draft_layers == 1:
+        return [num_target_layers // 2]
+    start = 1
+    end = num_target_layers - 3
+    span = end - start
+    return [
+        int(round(start + (i * span) / (num_draft_layers - 1)))
+        for i in range(num_draft_layers)
+    ]
+
+
 class UniversalDFlashDecoder:
     """Universal DFlash decoder that works with any MLX-converted model.
     
@@ -161,9 +178,7 @@ class UniversalDFlashDecoder:
         intermediate_size = int(draft_hidden_size * 2.75)  # Standard SwiGLU ratio
 
         # Target layer ids for feature extraction
-        target_layer_ids = DFlashDraftModel._build_target_layer_ids(
-            None, num_layers, draft_layers
-        )
+        target_layer_ids = _build_target_layer_ids(num_layers, draft_layers)
 
         drafter = DFlashDraftModel(
             vocab_size=vocab_size,
@@ -248,6 +263,7 @@ class UniversalDFlashDecoder:
         """Save the current drafter model."""
         import json
         from pathlib import Path
+        import numpy as np
 
         path = Path(path)
         path.mkdir(parents=True, exist_ok=True)