add kv copy logic and offline tests

Signed-off-by: morgendave <morgendave@gmail.com>

Author: morgendave

examples/offline_inference/spec_decode.py

@@ -68,8 +68,19 @@ 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)")
+    parser.add_argument(
+        "--no-prefill-token-shift",
+        dest="prefill_token_shift",
+        action="store_false",
+        help="Disable prefill token shift (default: enabled)",
+    )
+    parser.add_argument("--target_kv_layer_copy_from", type=int, default=-1)
+    parser.add_argument(
+        "--draft_kv_layer_copy_to",
+        type=str,
+        default="",
+        help="comma separated list of layer indices to copy to",
+    )
     return parser.parse_args()
 
 
@@ -101,11 +112,24 @@ def main():
 
         elif args.method == "eagle3" and eagle_dir is None:
             eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
+        target_kv_layer_copy_from = args.target_kv_layer_copy_from
+        draft_kv_layers_copy_to = (
+            [int(layer) for layer in args.draft_kv_layer_copy_to.split(",")]
+            if args.draft_kv_layer_copy_to
+            else None
+        )
+        kv_sharing_mapping = None
+        if args.target_kv_layer_copy_from >= 0 and draft_kv_layers_copy_to:
+            kv_sharing_mapping = {
+                f"{layer}": f"{target_kv_layer_copy_from}"
+                for layer in draft_kv_layers_copy_to
+            }
         speculative_config = {
             "method": args.method,
             "model": eagle_dir,
             "num_speculative_tokens": args.num_spec_tokens,
             "prefill_token_shift": args.prefill_token_shift,
+            "kv_sharing_mapping": kv_sharing_mapping,
         }
     elif args.method == "ngram":
         speculative_config = {

vllm/config.py

@@ -2556,6 +2556,7 @@ class SpeculativeConfig:
     """Shift tokens during draft prefill or not"""
 
     # Config for kv sharing, map from base model layer to draft layer
+    # Key is draft layer, value is base layer
     kv_sharing_mapping: SkipValidation[dict[str, str]] = None
     """KV copy mapping for prefill stage from base to draft"""
 

vllm/envs.py

@@ -138,6 +138,7 @@
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: str = "NONE"
     VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
     VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: Optional[int] = None
+    VLLM_DECODE_ONLY_ATTN: bool = False
 
 
 def get_default_cache_root():
@@ -953,7 +954,9 @@ def get_vllm_port() -> Optional[int]:
     # generations on machines < 100 for compressed-tensors
     # models
     "VLLM_USE_NVFP4_CT_EMULATIONS":
-    lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0")))
+    lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))),
+    "VLLM_DECODE_ONLY_ATTN":
+    lambda: os.environ.get("VLLM_DECODE_ONLY_ATTN", "0") == "1"
 }
 
 # --8<-- [end:env-vars-definition]

vllm/model_executor/models/llama4.py

@@ -183,11 +183,6 @@ 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,