[feat] add small vocab table for eagle's draft model[1].

docs/backend/speculative_decoding.ipynb

@@ -135,6 +135,74 @@
     "print_highlight(f\"Response: {response}\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "terminate_process(server_process)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### EAGLE Decoding via Frequency-Ranked Speculative Sampling\n",
+    "\n",
+    "By employing a truncated high-frequency token vocabulary in the draft model, Eagle speculative decoding reduces lm_head computational overhead while accelerating the pipeline without quality degradation. For more details checkout [this paper](https://arxiv.org/pdf/arXiv:2502.14856)\n",
+    "\n",
+    "Set `--speculative-token-map` to use this optimization. You can get the high-frequency token in FR-Spec from https://huggingface.co/thunlp/LLaMA3-Instruct-8B-FR-Spec. Or you can obtain high-frequency token by yourself (https://github.com/thunlp/FR-Spec/tree/main?tab=readme-ov-file#prepare-fr-spec-vocabulary-subset).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sglang.test.test_utils import is_in_ci\n",
+    "\n",
+    "if is_in_ci():\n",
+    "    from patch import launch_server_cmd\n",
+    "else:\n",
+    "    from sglang.utils import launch_server_cmd\n",
+    "\n",
+    "from sglang.utils import wait_for_server, print_highlight, terminate_process\n",
+    "\n",
+    "server_process, port = launch_server_cmd(\n",
+    "    \"\"\"\n",
+    "python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algo EAGLE \\\n",
+    "    --speculative-draft lmzheng/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \\\n",
+    "    --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --speculative-token-map {hot_token_ids.pt} \n",
+    "\"\"\"\n",
+    ")\n",
+    "\n",
+    "wait_for_server(f\"http://localhost:{port}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import openai\n",
+    "\n",
+    "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n",
+    "\n",
+    "response = client.chat.completions.create(\n",
+    "    model=\"meta-llama/Meta-Llama-3-8B-Instruct\",\n",
+    "    messages=[\n",
+    "        {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n",
+    "    ],\n",
+    "    temperature=0,\n",
+    "    max_tokens=64,\n",
+    ")\n",
+    "\n",
+    "print_highlight(f\"Response: {response}\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

python/sglang/srt/models/llama_eagle.py

@@ -117,9 +117,14 @@ def __init__(
         if self.config.tie_word_embeddings:
             self.lm_head = self.model.embed_tokens
         else:
-            self.lm_head = ParallelLMHead(
-                config.vocab_size, config.hidden_size, quant_config=quant_config
-            )
+            if hasattr(config, "hot_vocab_size"):
+                self.lm_head = ParallelLMHead(
+                    config.hot_vocab_size, config.hidden_size, quant_config=quant_config
+                )
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size, config.hidden_size, quant_config=quant_config
+                )
         self.logits_processor = LogitsProcessor(config)
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

python/sglang/srt/server_args.py

@@ -128,6 +128,7 @@ class ServerArgs:
     speculative_num_steps: int = 5
     speculative_eagle_topk: int = 8
     speculative_num_draft_tokens: int = 64
+    speculative_token_map: Optional[str] = None
 
     # Double Sparsity
     enable_double_sparsity: bool = False
@@ -751,6 +752,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
             help="The number of token sampled from draft model in Speculative Decoding.",
             default=ServerArgs.speculative_num_draft_tokens,
         )
+        parser.add_argument(
+            "--speculative-token-map",
+            type=str,
+            help="The path of the draft model's small vocab table.",
+            default=ServerArgs.speculative_token_map,
+        )
 
         # Double Sparsity
         parser.add_argument(

python/sglang/srt/speculative/eagle_worker.py

@@ -44,6 +44,18 @@ def __init__(
         # We will capture it later
         backup_disable_cuda_graph = server_args.disable_cuda_graph
         server_args.disable_cuda_graph = True
+
+        if server_args.speculative_token_map is not None:
+            try:
+                self.hot_token_id = torch.load(server_args.speculative_token_map)
+            except:
+                raise RuntimeError(
+                    f"there is not hot_token_ids.pt file in {self.server_args.speculative_token_map}"
+                )
+            server_args.json_model_override_args = (
+                f'{{"hot_vocab_size": {len(self.hot_token_id)}}}'
+            )
+
         super().__init__(
             gpu_id=gpu_id,
             tp_rank=tp_rank,
@@ -66,6 +78,14 @@ def __init__(
         # Share the embedding and lm_head
         if not self.speculative_algorithm.is_nextn():
             embed, head = self.target_worker.model_runner.model.get_embed_and_head()
+            if server_args.speculative_token_map is not None:
+                head = head.clone()
+                self.hot_token_id = torch.tensor(
+                    self.hot_token_id, dtype=torch.int32, device=head.device
+                )
+                head.data = head.data[self.hot_token_id]
+            else:
+                self.hot_token_id = None
             self.model_runner.model.set_embed_and_head(embed, head)
         self.model_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph
 
@@ -223,6 +243,11 @@ def draft_forward(self, forward_batch: ForwardBatch):
             spec_info.topk_index,
             spec_info.hidden_states,
         )
+        topk_index = (
+            self.hot_token_id[topk_index]
+            if self.hot_token_id is not None
+            else topk_index
+        )
 
         # Return values
         score_list: List[torch.Tensor] = []
@@ -262,6 +287,11 @@ def draft_forward(self, forward_batch: ForwardBatch):
             )
             probs = torch.softmax(logits_output.next_token_logits, dim=-1)
             topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
+            topk_index = (
+                self.hot_token_id[topk_index]
+                if self.hot_token_id is not None
+                else topk_index
+            )
             hidden_states = logits_output.hidden_states
 
         return score_list, token_list, parents_list