From 6f96a1c269fd9d2b80c0623416dbec04dbe9de35 Mon Sep 17 00:00:00 2001
From: wxiwnd <wxiwnd@outlook.com>
Date: Tue, 30 Jun 2026 00:05:16 +0800
Subject: [PATCH] Enable flashinfer TRTLLM for FP8 EAGLE MoE

---
 .../sglang/srt/model_loader/weight_utils.py   | 19 +++++++++
 python/sglang/srt/server_args.py              | 41 +++++++++++--------
 2 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py
index d1681aaa8..d179734f0 100644
--- a/python/sglang/srt/model_loader/weight_utils.py
+++ b/python/sglang/srt/model_loader/weight_utils.py
@@ -622,6 +622,25 @@ def nextn_moe_is_fp4_quantized_in_checkpoint(
     )
 
 
+def nextn_moe_is_fp8_quantized_in_checkpoint(
+    model_dir: Optional[str],
+    revision: Optional[str],
+    num_hidden_layers: Optional[int],
+) -> Optional[bool]:
+    """True/False if the checkpoint's NextN/MTP MoE experts ship fp8 scales."""
+    if not model_dir or num_hidden_layers is None:
+        return None
+    weight_map = load_safetensors_index_weight_map(model_dir, revision)
+    if weight_map is None:
+        return None
+    expert_prefix = f"model.layers.{num_hidden_layers}.mlp.experts."
+    return any(
+        name.startswith(expert_prefix)
+        and (name.endswith(".weight_scale_inv") or name.endswith(".weight_scale"))
+        for name in weight_map
+    )
+
+
 # For models like Mistral-7B-v0.3, there are both sharded
 # safetensors files and a consolidated safetensors file.
 # Passing both of these to the weight loader functionality breaks.
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index ab8583604..6598a4b45 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -3303,17 +3303,12 @@ class ServerArgs:
         from sglang.srt.layers.moe.utils import MoeRunnerBackend
         from sglang.srt.model_loader.weight_utils import (
             nextn_moe_is_fp4_quantized_in_checkpoint,
+            nextn_moe_is_fp8_quantized_in_checkpoint,
         )
 
-        # An fp4 NextN/MTP layer (e.g. GLM-5.2-NVFP4) is a genuine DeepSeek-V3 fp4 MoE,
-        # identical in kind to the main model's MoE — and the fp4 trtllm kernel supports
-        # the DeepSeek routing it carries. So the EAGLE draft can run the SAME
-        # flashinfer_trtllm path as the main model (consistent, well-tested, correct EP
-        # scale handling). The trtllm restriction below is real only for a *bf16* MTP
-        # (DeepSeek-R1/V3-FP4): the bf16 trtllm MoE requires RenormalizeNaive routing,
-        # and the draft routing is hard to guarantee there. We distinguish the two from
-        # the checkpoint's safetensors index (ground truth: do the MTP experts ship fp4
-        # scales), the same signal used to keep/drop modelopt_fp4 for the NextN module.
+        # Quantized NextN/MTP MoE can run the same non-routed flashinfer_trtllm
+        # path as the main model. A bf16 MTP still cannot: bf16 trtllm MoE needs
+        # RenormalizeNaive routing, which the draft path does not guarantee.
         _main_is_trtllm = self.moe_runner_backend in [
             "flashinfer_trtllm",
             "flashinfer_trtllm_routed",
@@ -3324,20 +3319,32 @@ class ServerArgs:
             if self.speculative_draft_model_path
             else self.revision
         )
+        _num_hidden_layers = getattr(
+            self.get_model_config().hf_config, "num_hidden_layers", None
+        )
         _nextn_is_fp4 = _main_is_trtllm and (
             nextn_moe_is_fp4_quantized_in_checkpoint(
                 _spec_model_dir,
                 _spec_revision,
-                getattr(self.get_model_config().hf_config, "num_hidden_layers", None),
+                _num_hidden_layers,
             )
             is True
         )
+        _nextn_is_fp8 = _main_is_trtllm and (
+            nextn_moe_is_fp8_quantized_in_checkpoint(
+                _spec_model_dir,
+                _spec_revision,
+                _num_hidden_layers,
+            )
+            is True
+        )
+        _nextn_can_use_trtllm = _nextn_is_fp4 or _nextn_is_fp8
 
         if self.speculative_moe_runner_backend is None:
-            if _nextn_is_fp4:
-                # fp4 NextN -> the draft runs the same fp4 trtllm MoE path as main.
-                # Always the NON-routed kernel: the draft's FlashInferFP4MoE only
-                # implements trtllm_fp4_block_scale_moe, not the _routed variant.
+            if _nextn_can_use_trtllm:
+                # Quantized NextN -> the draft runs the same trtllm MoE path as main.
+                # Always the NON-routed kernel; speculative MoE does not support the
+                # _routed variant.
                 self.speculative_moe_runner_backend = "flashinfer_trtllm"
             elif _main_is_trtllm:
                 self.speculative_moe_runner_backend = "auto"
@@ -3346,10 +3353,10 @@ class ServerArgs:
         elif MoeRunnerBackend(
             self.speculative_moe_runner_backend
         ).is_flashinfer_trtllm():
-            assert _nextn_is_fp4, (
+            assert _nextn_can_use_trtllm, (
                 "Speculative MoE runner backend flashinfer_trtllm is only supported for "
-                "an fp4 NextN/MTP layer (the bf16 trtllm MoE requires RenormalizeNaive "
-                "routing); use triton or auto for a bf16 MTP."
+                "a quantized fp4/fp8 NextN/MTP layer (the bf16 trtllm MoE requires "
+                "RenormalizeNaive routing); use triton or auto for a bf16 MTP."
             )
         elif MoeRunnerBackend(
             self.speculative_moe_runner_backend