From 4cf109bbd10b9c91cec50c07a736d3c49ba624d2 Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Wed, 29 Apr 2026 16:03:27 +0100 Subject: [PATCH] debug followup (#24058) --- .../Mistral/Mistral-Medium-3.5.mdx | 69 +++++- .../mistral-medium-3-5-deployment.jsx | 9 + python/sglang/srt/models/llava.py | 9 + python/sglang/srt/models/mistral_eagle.py | 208 ++++++++++++++++++ .../utils/hf_transformers/mistral_utils.py | 34 ++- 5 files changed, 322 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/models/mistral_eagle.py diff --git a/docs_new/cookbook/autoregressive/Mistral/Mistral-Medium-3.5.mdx b/docs_new/cookbook/autoregressive/Mistral/Mistral-Medium-3.5.mdx index a5c2a3104..69af325d0 100644 --- a/docs_new/cookbook/autoregressive/Mistral/Mistral-Medium-3.5.mdx +++ b/docs_new/cookbook/autoregressive/Mistral/Mistral-Medium-3.5.mdx @@ -39,9 +39,16 @@ The HuggingFace repo ships both the mistral native layout (`params.json` + `cons ## 2. SGLang Installation -SGLang offers multiple installation methods. You can choose the most suitable installation method based on your hardware platform and requirements. +Refer to the [official SGLang installation guide](../../../docs/get-started/install). -Please refer to the [official SGLang installation guide](../../../docs/get-started/install) for installation instructions. +**Docker Images by Hardware:** + +| Hardware | Docker Image | +| --- | --- | +| H100 / H200 (Hopper, CUDA 12.9) | `lmsysorg/sglang:dev-mistral-medium-3.5` | +| B200 / B300 (Blackwell, CUDA 13.0) | `lmsysorg/sglang:dev-cu13-mistral-medium-3.5` | + +> Day-0 support for Mistral Medium 3.5 is not yet in `lmsysorg/sglang:latest` — pull one of the tags above (matching your GPU's CUDA driver) until the changes propagate to the next stable release. --- @@ -63,6 +70,30 @@ Please refer to the [official SGLang installation guide](../../../docs/get-start - **Reasoning parser**: Enable `--reasoning-parser mistral` to separate `reasoning_content` from the main response content. - **System prompt**: The model ships with a recommended system prompt in `chat_template.jinja` and `SYSTEM_PROMPT.txt`. If you do not pass a system message yourself, the chat template injects Mistral's default (model identity, current date, tool-use guidelines). For full fidelity with Mistral's reference setup, load `SYSTEM_PROMPT.txt` from the HF repo and substitute `{name}`, `{today}`, `{yesterday}` (see Section 4.6). +### 3.3 Speculative Decoding (EAGLE) + +Mistral ships an EAGLE draft head, [`mistralai/Mistral-Medium-3.5-128B-EAGLE`](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B-EAGLE), that lets you run speculative decoding on top of the dense 128B target. The draft is a 2-layer GQA body sharing the target's vocab/head, FP8-quantized like the target (~4 GB), and is meant for low-concurrency latency-bound serving. + +```bash Command +python -m sglang.launch_server \ + --model-path mistralai/Mistral-Medium-3.5-128B \ + --tp 4 \ + --dtype bfloat16 \ + --tool-call-parser mistral \ + --reasoning-parser mistral \ + --speculative-algorithm EAGLE \ + --speculative-draft-model-path mistralai/Mistral-Medium-3.5-128B-EAGLE \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --port 30000 +``` + +- **`--dtype bfloat16` is required.** The draft `params.json` does not carry a `dtype` field, so `--dtype auto` falls back to fp32 and downcasts to fp16, which conflicts with the bf16 target when the embed/head are shared. Setting bf16 explicitly keeps both sides aligned (this is a no-op for the target — it already loads as bf16). +- The draft uses the same vocab and lm_head as the target. Memory overhead on top of the base model is ~4 GB per TP shard. +- `(num-steps, eagle-topk, num-draft-tokens) = (3, 1, 4)` is the recommended starting point. Tune for your workload — wider trees (higher `eagle-topk` / `num-draft-tokens`) help high-acceptance (templated) outputs, narrower trees keep latency tight on more diverse text. +- EAGLE shines at low concurrency. At high concurrency, throughput is dominated by the target's batched forward pass and the draft's contribution shrinks; consider running without EAGLE for batch-serving workloads. + --- ## 4. Model Invocation @@ -396,3 +427,37 @@ Median TTFT (ms): 152.95 Median TPOT (ms): 42.53 ================================================== ``` + +### 5.3 EAGLE Speculative Decoding (Latency) + +Same 4× H200 setup, EAGLE configuration from [Section 3.3](#3-3-speculative-decoding-eagle). Single-stream latency benchmark (`--max-concurrency 1`). + +```bash Command +python3 -m sglang.bench_serving \ + --backend sglang \ + --dataset-name random \ + --num-prompts 10 \ + --max-concurrency 1 \ + --random-input-len 1024 \ + --random-output-len 512 \ + --port 30000 +``` + +**Results:** + +```text Output +============ Serving Benchmark Result ============ +Backend: sglang +Successful requests: 10 +Benchmark duration (s): 27.64 +Total input tokens: 6101 +Total generated tokens: 2684 +Output token throughput (tok/s): 97.10 +Mean E2E Latency (ms): 2762.99 +Median TTFT (ms): 90.69 +Median TPOT (ms): 9.73 +Accept length: 1.72 +================================================== +``` + +EAGLE delivers **~1.41× output throughput and ~29% lower E2E latency** vs. the baseline in [Section 5.2](#5-2-speed-benchmarks) on the same workload. Acceptance length of 1.72 means each draft cycle averages roughly 1.7 accepted tokens. diff --git a/docs_new/src/snippets/autoregressive/mistral-medium-3-5-deployment.jsx b/docs_new/src/snippets/autoregressive/mistral-medium-3-5-deployment.jsx index 8dbe903eb..0976d5b39 100644 --- a/docs_new/src/snippets/autoregressive/mistral-medium-3-5-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/mistral-medium-3-5-deployment.jsx @@ -30,6 +30,15 @@ export const MistralMedium35Deployment = () => { ], commandRule: (value) => value === 'enabled' ? '--tool-call-parser mistral' : null }, + speculative: { + name: 'speculative', + title: 'Speculative Decoding (EAGLE)', + items: [ + { id: 'disabled', label: 'Disabled', default: false }, + { id: 'enabled', label: 'Enabled', default: true } + ], + commandRule: (value) => value === 'enabled' ? '--dtype bfloat16 \\\n --speculative-algorithm EAGLE \\\n --speculative-draft-model-path mistralai/Mistral-Medium-3.5-128B-EAGLE \\\n --speculative-num-steps 3 \\\n --speculative-eagle-topk 1 \\\n --speculative-num-draft-tokens 4' : null + }, }; // 128B dense FP8 ≈ 130GB, plus KV cache headroom diff --git a/python/sglang/srt/models/llava.py b/python/sglang/srt/models/llava.py index f07dcaa56..712c1f4f8 100644 --- a/python/sglang/srt/models/llava.py +++ b/python/sglang/srt/models/llava.py @@ -453,6 +453,15 @@ class LlavaBaseForCausalLM(nn.Module): elif forward_batch.forward_mode.is_decode(): return self.language_model(input_ids, positions, forward_batch) + def get_embed_and_head(self): + # Spec-decode plumbing: expose the LM's embed/head so the EAGLE draft + # can share them with the target. self.language_model is a Llama-family + # CausalLM that defines this method. + return self.language_model.get_embed_and_head() + + def set_embed_and_head(self, embed, head): + self.language_model.set_embed_and_head(embed, head) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # Load clip vision model by cfg['mm_vision_tower']: # huggingface_name or path_of_clip_relative_to_llava_model_dir diff --git a/python/sglang/srt/models/mistral_eagle.py b/python/sglang/srt/models/mistral_eagle.py new file mode 100644 index 000000000..434856026 --- /dev/null +++ b/python/sglang/srt/models/mistral_eagle.py @@ -0,0 +1,208 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""EAGLE draft model for GQA Mistral targets (e.g. Mistral Medium 3.5). + +Reuses ``LlamaForCausalLMEagle`` for the EAGLE machinery (lm_head/embed_tokens +construction, optional tied embeddings, capture-aux-hidden-states plumbing) but +swaps in a Mistral-specific draft model body that: + +- runs through the standard :class:`LlamaDecoderLayer` (GQA), not the layernorm + -less variant ``llama_eagle.LlamaDecoderLayer`` — Mistral's EAGLE checkpoint + ships ``layers.0.attention_norm.weight``, so layer 0 expects the input + layernorm to be present. +- uses ``RowParallelLinear`` for the EAGLE fc fusion layer with a + ``quant_config``, so the FP8-quantized ``eagle_linear`` weights from the + Mistral native checkpoint load via the standard quant pipeline (``LlamaModel`` + in ``llama_eagle.py`` uses a plain :class:`torch.nn.Linear` which cannot + consume FP8 e4m3 tensors). + +The weight name remapping mirrors :class:`MistralForCausalLMMistralFormat` and +adds the eagle-specific entries for ``eagle_linear`` → ``model.fc``. +""" + +import logging +from collections.abc import Iterable +from typing import Optional, Tuple + +import regex as re +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import RowParallelLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.models.llama import LlamaDecoderLayer, LlamaForCausalLM +from sglang.srt.models.llama_eagle import LlamaForCausalLMEagle +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class MistralEagleModel(nn.Module): + """GQA EAGLE draft body with the input-embed ⊕ target-hidden-state fusion.""" + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + assert ( + get_pp_group().world_size == 1 + ), "MistralForCausalLMEagle currently does not support pipeline parallelism" + self.pp_group = get_pp_group() + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer( + config=config, + layer_id=i, + prefix=add_prefix(f"layers.{i}", prefix), + quant_config=quant_config, + ) + for i in range(config.num_hidden_layers) + ] + ) + self.start_layer = 0 + self.end_layer = config.num_hidden_layers + self.fc = RowParallelLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fc", prefix), + input_is_parallel=False, + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + # EAGLE fusion: concat input embedding with target's previous hidden + # state, project back to hidden_size before going through the draft's + # transformer layers. + hidden_states, _ = self.fc( + torch.cat( + (hidden_states, forward_batch.spec_info.hidden_states), + dim=-1, + ) + ) + + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + return hidden_states + residual + + +class MistralForCausalLMEagle(LlamaForCausalLMEagle): + """EAGLE draft for GQA Mistral targets. + + Inherits LlamaForCausalLMEagle for the lm_head/embed_tokens setup and the + capture-aux-hidden-state hooks, then overrides ``self.model`` with the + quant-aware :class:`MistralEagleModel` and applies Mistral native-format + weight remapping during ``load_weights``. + """ + + # fmt: off + remapping = { + r"layers\.(\d+)\.attention_norm\.weight": r"model.layers.\1.input_layernorm.weight", + r"layers\.(\d+)\.attention\.wq\.(\w+)": r"model.layers.\1.self_attn.q_proj.\2", + r"layers\.(\d+)\.attention\.wk\.(\w+)": r"model.layers.\1.self_attn.k_proj.\2", + r"layers\.(\d+)\.attention\.wv\.(\w+)": r"model.layers.\1.self_attn.v_proj.\2", + r"layers\.(\d+)\.attention\.wo\.(\w+)": r"model.layers.\1.self_attn.o_proj.\2", + r"layers\.(\d+)\.ffn_norm\.weight": r"model.layers.\1.post_attention_layernorm.weight", + r"layers\.(\d+)\.feed_forward\.w1\.(\w+)": r"model.layers.\1.mlp.gate_proj.\2", + r"layers\.(\d+)\.feed_forward\.w2\.(\w+)": r"model.layers.\1.mlp.down_proj.\2", + r"layers\.(\d+)\.feed_forward\.w3\.(\w+)": r"model.layers.\1.mlp.up_proj.\2", + r"norm\.weight": "model.norm.weight", + # Eagle-specific: the fc layer that fuses input embeds and target + # hidden states is named `eagle_linear` in the Mistral checkpoint. + # Its FP8 weights live alongside per-tensor activation/weight scales. + r"eagle_linear\.weight": r"model.fc.weight", + r"eagle_linear\.qscale_act": r"model.fc.input_scale", + r"eagle_linear\.qscale_weight": r"model.fc.weight_scale", + # tok_embeddings and output are intentionally absent — EAGLE shares + # both with the target model and the framework ties them at runtime. + } + # fmt: on + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + # Run LlamaForCausalLMEagle.__init__ to set up lm_head/embed_tokens/etc. + # then replace self.model (which uses a plain torch.nn.Linear for fc and + # cannot consume FP8 weights) with our quant-aware draft body. + super().__init__(config=config, quant_config=quant_config, prefix=prefix) + self.model = MistralEagleModel( + config, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # Bypass LlamaForCausalLMEagle.load_weights' "prepend model." behaviour + # because our remap already emits fully-qualified target names. + return LlamaForCausalLM.load_weights( + self, self._remap_mistral_to_llama(weights) + ) + + def _remap_mistral_to_llama( + self, weights: Iterable[Tuple[str, torch.Tensor]] + ) -> Iterable[Tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + if name.startswith("model.") or name.startswith("lm_head."): + yield name, loaded_weight + continue + for k, v in self.remapping.items(): + match = re.fullmatch(k, name) + if match: + name = match.expand(v) + break + else: + logger.warning(f"Unrecognized weight: {name}. Skipping.") + continue + if name.endswith(".qscale_act"): + name = re.sub(r"\.qscale_act$", ".input_scale", name) + elif name.endswith(".qscale_weight"): + name = re.sub(r"\.qscale_weight$", ".weight_scale", name) + yield name, loaded_weight + + +EntryClass = [MistralForCausalLMEagle] diff --git a/python/sglang/srt/utils/hf_transformers/mistral_utils.py b/python/sglang/srt/utils/hf_transformers/mistral_utils.py index fa8f1c623..a1b998c9b 100644 --- a/python/sglang/srt/utils/hf_transformers/mistral_utils.py +++ b/python/sglang/srt/utils/hf_transformers/mistral_utils.py @@ -27,8 +27,12 @@ def adapt_config_dict( is_moe and (config_dict["moe"].get("num_shared_experts") or 0) > 0 ) is_eagle = "eagle" in model.lower() - if is_eagle and not is_moe: - # Dense EAGLE draft model (e.g. Mistral Small 4 EAGLE). + is_mla_eagle = is_eagle and any( + config_dict.get(k) is not None + for k in ("kv_lora_rank", "q_lora_rank", "v_head_dim") + ) + if is_eagle and not is_moe and is_mla_eagle: + # Dense MLA EAGLE draft model (e.g. Mistral Small 4 EAGLE). # Uses MLA attention like MistralLarge3 but has no MoE layers. # Set model_type to deepseek_v3 for MLA support, and override # MoE fields so all layers are dense. @@ -47,6 +51,21 @@ def adapt_config_dict( config_dict["topk_method"] = None config_dict["scoring_func"] = "softmax" config_dict["routing_method_type"] = 1 + elif is_eagle and not is_moe: + # Dense GQA EAGLE draft model (e.g. Mistral Medium 3.5 EAGLE). + # Routes to a Llama-backbone draft body — no MoE shimming required. + config_dict["architectures"] = ["MistralForCausalLMEagle"] + config_dict["model_type"] = "mistral" + config_dict["rope_is_neox_style"] = False + for mla_key in ( + "q_lora_rank", + "qk_rope_head_dim", + "qk_nope_head_dim", + "kv_lora_rank", + "v_head_dim", + ): + if config_dict.get(mla_key) is None: + config_dict.pop(mla_key, None) elif is_moe: if is_mistral_large_3: config_dict = _remap_moe_args(config_dict) @@ -328,9 +347,14 @@ class MistralConfigParser: def is_mistral_model(name) -> bool: """Return True if *name* refers to a Mistral model needing the custom parser.""" lower = str(name).lower() - return ( - "mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower - ) + if "mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower: + return True + # EAGLE drafts for Mistral targets ship native-format only (params.json + + # consolidated.safetensors, no config.json), so route them through the + # custom parser regardless of the base model name. + if "eagle" in lower and "mistral" in lower: + return True + return False @lru_cache(maxsize=2)