initial commit for non-shifting prefill in eagle, prepare for kv sharing

rope change

rope change

rebase

Author: morgendave

examples/offline_inference/spec_decode.py

@@ -68,6 +68,8 @@ def parse_args():
     parser.add_argument("--model-dir", type=str, default=None)
     parser.add_argument("--eagle-dir", type=str, default=None)
     parser.add_argument("--custom-mm-prompts", action="store_true")
+    parser.add_argument("--no-prefill-token-shift", dest="prefill_token_shift",
+        action="store_false", help="Disable prefill token shift (default: enabled)")
     return parser.parse_args()
 
 
@@ -103,6 +105,7 @@ def main():
             "method": args.method,
             "model": eagle_dir,
             "num_speculative_tokens": args.num_spec_tokens,
+            "prefill_token_shift": args.prefill_token_shift,
         }
     elif args.method == "ngram":
         speculative_config = {

vllm/config.py

@@ -2551,6 +2551,14 @@ class SpeculativeConfig:
         ParallelConfig] = None  # type: ignore
     """The parallel configuration for the draft model initialized internal."""
 
+    # Shift prefill tokens for draft, only used in eagle
+    prefill_token_shift: bool = True
+    """Shift tokens during draft prefill or not"""
+
+    # Config for kv sharing, map from base model layer to draft layer
+    kv_sharing_mapping: SkipValidation[dict[str, str]] = None
+    """KV copy mapping for prefill stage from base to draft"""
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,
@@ -2937,6 +2945,11 @@ def num_lookahead_slots(self) -> int:
     def use_eagle(self) -> bool:
         return self.method in ("eagle", "eagle3", "deepseek_mtp")
 
+    def eagle_shift_prefill_token(self) -> bool:
+        if self.use_eagle():
+            return self.prefill_token_shift
+        return False
+
     def __repr__(self) -> str:
         method = self.method
         model = None if method == "ngram" else self.draft_model_config.model

vllm/model_executor/models/llama4.py

@@ -183,6 +183,11 @@ def __init__(self,
         is_gguf = quant_config and quant_config.get_name() == "gguf"
         if is_gguf and config.model_type == "llama":
             is_neox_style = False
+        elif config.model_type == "eagle":
+            # EAGLE draft model does not use neox style RoPE
+            is_neox_style = False
+        else:
+            is_neox_style = True
 
         self.rotary_emb = get_rope(
             self.head_dim,