feat(turbomind): EAGLE3 and Qwen3.5 MTP speculative decoding

lmdeploy/metrics/stats.py

@@ -285,10 +285,25 @@ def update_from_output(self, outputs: EngineOutput):
         """Update from engine output."""
         spec_info = getattr(outputs.req_metrics, 'spec_info', None)
         if spec_info:
-            self.num_drafts += 1
-            self.num_draft_tokens += spec_info['num_draft_tokens']
-            self.num_accepted_tokens += spec_info['num_accepted_tokens']
-            self.num_accepted_tokens_per_pos[:spec_info['num_accepted_tokens']] += 1
+            if 'num_drafts' in spec_info:
+                self.num_drafts += \
+                    spec_info['num_drafts']
+                self.num_draft_tokens += \
+                    spec_info['num_draft_tokens']
+                self.num_accepted_tokens += \
+                    spec_info['num_accepted_tokens']
+                self.num_accepted_tokens_per_pos += \
+                    np.asarray(
+                        spec_info[
+                            'num_accepted_tokens_per_pos'])
+            else:
+                self.num_drafts += 1
+                self.num_draft_tokens += \
+                    spec_info['num_draft_tokens']
+                self.num_accepted_tokens += \
+                    spec_info['num_accepted_tokens']
+                self.num_accepted_tokens_per_pos[
+                    :spec_info['num_accepted_tokens']] += 1
 
     def update_per_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
         """Update with per draft stats."""

lmdeploy/serve/core/async_engine.py

@@ -145,13 +145,12 @@ def __init__(self,
         self.session_len = (_get_and_verify_max_len(self.hf_cfg, None)
                             if backend_config.session_len is None else backend_config.session_len)
         backend_config.session_len = self.session_len
-        if speculative_config is not None and backend == 'turbomind':
-            logger.warning('speculative decoding is not supported by turbomind ')
         # build backend engine
         if backend == 'turbomind':
             self.engine = self._build_turbomind(model_path=model_path,
                                                 backend_config=backend_config,
                                                 trust_remote_code=trust_remote_code,
+                                                speculative_config=speculative_config,
                                                 **kwargs)
         elif backend == 'pytorch':
             self.engine = self._build_pytorch(model_path=model_path,
@@ -174,8 +173,9 @@ def __init__(self,
         self.backend = backend
         self.request_logger = RequestLogger(max_log_len)
 
-        self.num_spec_token = 0 if backend == 'turbomind' or speculative_config is None \
-            else speculative_config.num_speculative_tokens
+        self.num_spec_token = (0
+                               if speculative_config is None
+                               else speculative_config.num_speculative_tokens)
 
         self.session_mgr = SessionManager()
         self.session_mgr.build_request_handle_pool(self.engine, self.backend_config.max_batch_size)
@@ -201,6 +201,7 @@ def __exit__(self, exc_type, exc_value, traceback):
     def _build_turbomind(self,
                          model_path: str,
                          backend_config: TurbomindEngineConfig | None = None,
+                         speculative_config: SpeculativeConfig | None = None,
                          trust_remote_code: bool = False,
                          **kwargs):
         """Inner build method for turbomind backend."""
@@ -210,7 +211,11 @@ def _build_turbomind(self,
                 'TurboMind was requested but its native module is unavailable.'
             ) from turbomind._import_error
         return turbomind.TurboMind.from_pretrained(
-            model_path, engine_config=backend_config, trust_remote_code=trust_remote_code, **kwargs
+            model_path,
+            engine_config=backend_config,
+            trust_remote_code=trust_remote_code,
+            speculative_config=speculative_config,
+            **kwargs
         )
 
     def _build_pytorch(self,

lmdeploy/turbomind/__init__.py

@@ -1,5 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
+import sys
+
 import torch  # noqa: F401
 
 _import_error = None
@@ -10,6 +12,9 @@
     _tm = None
     _import_error = error
 else:
+    # Expose the extension under its bare name so submodules that predate the
+    # lazy-import design can `import _turbomind` without a second load.
+    sys.modules.setdefault('_turbomind', _tm)
     from .turbomind import TurboMind as TurboMind
 
 

lmdeploy/turbomind/builders/__init__.py

@@ -7,11 +7,13 @@
 from .attention import AttentionBuilder
 from .decoder_layer import DecoderLayerBuilder, DecoderLayerConfig
 from .deltanet import DeltaNetBuilder
+from .eagle3_weight import Eagle3WeightBuilder, Eagle3WeightConfig
 from .ffn import FfnBuilder
 from .mla import MLABuilder
 from .module_list import ModuleListBuilder, ModuleListConfig
 from .moe import MoeBuilder
 from .norm import LayerNormBuilder, NormBuilder, make_layer_norm_config, make_norm_config
+from .qwen3_5_mtp_weight import Qwen35MtpWeightBuilder, Qwen35MtpWeightConfig
 from .text_model import TextModelBuilder
 from .vision_model import VisionModelBuilder
 
@@ -32,6 +34,8 @@
     'DeltaNetBuilder',
     'MLABuilder',
     'DecoderLayerBuilder',
+    'Eagle3WeightBuilder',
+    'Qwen35MtpWeightBuilder',
     'ModuleListBuilder',
     'NormBuilder',
     'LayerNormBuilder',
@@ -40,6 +44,8 @@
     'make_layer_norm_config',
     # C++ config re-exports
     'DecoderLayerConfig',
+    'Eagle3WeightConfig',
+    'Qwen35MtpWeightConfig',
     'ModuleListConfig',
     # Helper functions
 ]

lmdeploy/turbomind/builders/eagle3_weight.py

@@ -0,0 +1,13 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import _turbomind as _tm
+
+from ._base import Builder
+
+Eagle3WeightConfig = _tm.Eagle3WeightConfig
+
+
+class Eagle3WeightBuilder(Builder):
+    """Builder for the EAGLE3 method weight tree at ModelWeight.spec."""
+
+    def add_target_hidden_proj(self, linear):
+        self._add_linear('target_hidden_proj', linear, split_side=None)

lmdeploy/turbomind/builders/qwen3_5_mtp_weight.py

@@ -0,0 +1,12 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import _turbomind as _tm
+
+from ._base import Builder, SplitSide
+
+Qwen35MtpWeightConfig = _tm.Qwen35MtpWeightConfig
+
+
+class Qwen35MtpWeightBuilder(Builder):
+
+    def add_fc(self, linear):
+        self._add_linear('fc', linear, split_side=SplitSide.OUTPUT)

lmdeploy/turbomind/builders/text_model.py

@@ -9,31 +9,38 @@ class TextModelBuilder(Builder):
 
     Constructs a ModelWeight via ``_tm.create_module(ModelWeightConfig)``
     on each context (inherited Builder machinery), then attaches it to
-    externally-owned ``ModelRoot`` sentinel handles as their
-    ``text_model`` child during ``build()``.
+    externally-owned ``ModelRoot`` sentinel handles as the configured
+    root child during ``build()``.
 
     Owns ``tok_embeddings`` (Tensor param) and ``output`` (LinearWeight
     child) commits on the ModelWeight via ``add_token_embeds`` /
     ``add_lm_head``.
     """
 
-    def __init__(self, config, ctx, *, root_handles,
-                 tp: ParallelGroup, vocab_size):
+    def __init__(self,
+                 config,
+                 ctx,
+                 *,
+                 root_handles,
+                 tp: ParallelGroup,
+                 vocab_size,
+                 root_child: str = 'text_model'):
         super().__init__(config, ctx)
         self.tp = tp
         self.config.tp_size = tp.size
         self._root_handles = root_handles
         self._vocab_size = vocab_size
+        self._root_child = root_child
 
     def build(self) -> BuiltModule:
-        """Create ModelWeight via _tm.create_module (via super), then attach
-        each per-GPU ModelWeight handle to its sentinel root via
-        add_child_raw."""
+        """Build and attach each ModelWeight to the configured root child."""
         built = super().build()
-        for i, (root, text_model) in enumerate(
+        for i, (root, model_weight) in enumerate(
                 zip(self._root_handles, built.handles)):
             with self._ctx.devices[i]:
-                root.add_child_raw('text_model', text_model)
+                root.add_child_raw(
+                    self._root_child,
+                    model_weight)
         return built
 
     def add_token_embeds(self, tensor):

lmdeploy/turbomind/model_loader.py

@@ -15,12 +15,22 @@ class ModelLoader:
     to the C++ runtime.
     """
 
-    def __init__(self, model, model_comm, gpu_count, model_path,
-                 data_type, engine_config):
+    def __init__(self,
+                 model,
+                 model_comm,
+                 gpu_count,
+                 model_path,
+                 data_type,
+                 engine_config,
+                 *,
+                 draft_model=None,
+                 draft_model_path=None):
         self.model = model
+        self.draft_model = draft_model
         self.model_comm = model_comm
         self.gpu_count = gpu_count
         self.model_path = model_path
+        self.draft_model_path = draft_model_path
         self.data_type = data_type
         self.engine_config = engine_config
         self._bind_runtime()
@@ -59,33 +69,52 @@ def _bind_runtime(self):
             [(e * mlp_tp.size + m) % dense_size
              for e, m in zip(ep.ranks, mlp_tp.ranks)])
 
-        self.model.bind_runtime(
-            ctx=ctx,
-            root_handles=[mc.root(g) for g in range(self.gpu_count)],
-            attn_tp=attn_tp,
-            mlp_tp=mlp_tp,
-            ep=ep,
-            model_tp=model_tp,
-            dense_tp=dense_tp,
-        )
+        models = (
+            (self.model,)
+            if self.draft_model is None
+            else (self.model, self.draft_model))
+        for model in models:
+            model.bind_runtime(
+                ctx=ctx,
+                root_handles=[
+                    mc.root(g)
+                    for g in range(self.gpu_count)],
+                attn_tp=attn_tp,
+                mlp_tp=mlp_tp,
+                ep=ep,
+                model_tp=model_tp,
+                dense_tp=dense_tp)
 
-    def export(self):
-        ckpt = create_checkpoint(
-            self.model_path,
-            mappings=getattr(self.model, '_loader_mappings', []))
+    @staticmethod
+    def _export_one(model, model_path):
+        checkpoint = create_checkpoint(
+            model_path,
+            mappings=getattr(
+                model, '_loader_mappings', []))
         try:
-            self.model.model(Prefix(ckpt))
+            root = Prefix(checkpoint)
+            if getattr(model, 'prefix', ''):
+                root = root + model.prefix
+            model.model(root)
         finally:
-            ckpt.close()
+            checkpoint.close()
+
+    def export(self):
+        self._export_one(
+            self.model, self.model_path)
+        if self.draft_model is not None:
+            self._export_one(
+                self.draft_model,
+                self.draft_model_path)
         torch.cuda.empty_cache()
 
     def export_iter(self):
-        ckpt = create_checkpoint(
-            self.model_path,
-            mappings=getattr(self.model, '_loader_mappings', []))
-        try:
-            self.model.model(Prefix(ckpt))
+        self._export_one(
+            self.model, self.model_path)
+        yield -1
+        if self.draft_model is not None:
+            self._export_one(
+                self.draft_model,
+                self.draft_model_path)
             yield -1
-        finally:
-            ckpt.close()
         torch.cuda.empty_cache()

lmdeploy/turbomind/models/__init__.py

@@ -10,3 +10,4 @@
 from .qwen2_vl import Qwen2VLModel  # noqa: F401
 from .qwen3 import Qwen3TextModel  # noqa: F401
 from .qwen3_5 import Qwen3_5Model, Qwen3_5TextModel, Qwen3_5VisionModel  # noqa: F401
+from .qwen3_eagle3 import Qwen3Eagle3TextModel  # noqa: F401

lmdeploy/turbomind/models/qwen3_5_mtp.py

@@ -0,0 +1,66 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""Checkpoint-native Qwen3.5 MTP draft weight tree."""
+
+from ..builders import (
+    DecoderLayerBuilder,
+    DecoderLayerConfig,
+    ModuleListBuilder,
+    ModuleListConfig,
+    Qwen35MtpWeightBuilder,
+    Qwen35MtpWeightConfig,
+    TextModelBuilder,
+)
+from .qwen3_5 import Qwen3_5TextModel
+from .utils import make_model_weight_config
+
+
+class Qwen3_5MtpTextModel(Qwen3_5TextModel):
+
+    def __init__(self, cfg, *, resolver):
+        super().__init__(cfg, resolver=resolver)
+        self.tap_layer_ids = [cfg.num_hidden_layers]
+
+    def model(self, pfx):
+        root_cfg = make_model_weight_config(self.cfg)
+        root_cfg.decoder_only = True
+
+        root = TextModelBuilder(
+            root_cfg,
+            self._ctx,
+            root_handles=self._root_handles,
+            tp=self._model_tp,
+            vocab_size=self.cfg.vocab_size,
+            root_child='draft_model')
+
+        root.norm = self.norm(pfx + 'mtp.norm', zero_centered=True)
+        root.layers = self.mtp_layers(pfx + 'mtp.layers')
+        root.spec = self.mtp_spec(pfx + 'mtp')
+        root.build()
+
+    def mtp_layers(self, pfx):
+        layers = ModuleListBuilder(ModuleListConfig(), self._ctx)
+        for i, layer_pfx in pfx.slices(0, 1):
+            layer = DecoderLayerBuilder(DecoderLayerConfig(), self._ctx)
+            layer.attention = self.attn(layer_pfx + 'self_attn')
+            if self._n_experts > 0:
+                layer.moe_ffn = self.moe(layer_pfx + 'mlp')
+            else:
+                layer.feed_forward = self.ffn(layer_pfx + 'mlp',
+                                              self.cfg.intermediate_size,
+                                              tp=self._mlp_tp)
+            layer.attention_norm = self.norm(
+                layer_pfx + 'input_layernorm', zero_centered=True)
+            layer.ffn_norm = self.norm(
+                layer_pfx + 'post_attention_layernorm', zero_centered=True)
+            layers[i] = layer.build()
+        return layers.build()
+
+    def mtp_spec(self, pfx):
+        spec = Qwen35MtpWeightBuilder(Qwen35MtpWeightConfig(), self._ctx)
+        spec.tp = self._model_tp
+        spec.add_fc(self._linear(pfx + 'fc'))
+        spec.pre_fc_norm_embedding = self.norm(
+            pfx + 'pre_fc_norm_embedding', zero_centered=True)
+        spec.pre_fc_norm_hidden = self.norm(
+            pfx + 'pre_fc_norm_hidden', zero_centered=True)
+        return spec.build()