debug followup (#24058)

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.

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

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

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]

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:
+    is_mla_eagle = is_eagle and any(
-        # Dense EAGLE draft model (e.g. Mistral Small 4 EAGLE).
+        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 (
+    if "mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower:
-        "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)