[Feature]: Support speculative decoding #3945

RunningLeon · 2025-09-09T06:26:14Z

Motivation

Support speculative decoding

Examples

pipeline

from lmdeploy import pipeline, PytorchEngineConfig
from lmdeploy.messages import SpeculativeConfig


if __name__ == '__main__':

    model_path = 'meta-llama/Llama-3.1-8B-Instruct'
    spec_cfg = SpeculativeConfig(method='eagle3', 
                                    num_speculative_tokens=3,
                                    model='yuhuili/EAGLE3-LLaMA3.1-Instruct-8B',
                                    )
    pipe = pipeline(model_path, 
                    backend_config=PytorchEngineConfig(max_batch_size=128),
                    speculative_config=spec_cfg)
    response = pipe(['Hi, pls intro yourself', 'Shanghai is'])
    print(response)

serving

HF_HOME=/nvme1/shared/huggingface_hub \
TRANSFORMERS_OFFLINE=1 \
CUDA_VISIBLE_DEVICES=7 \
lmdeploy serve api_server \
meta-llama/Llama-3.1-8B-Instruct \
--backend pytorch \
--server-port 24545 \
--speculative-draft-model yuhuili/EAGLE3-LLaMA3.1-Instruct-8B \
--speculative-algorithm eagle3 \
--speculative-num-draft-tokens 3 \
--max-batch-size 128 \
--enable-metrics

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

Pre-commit or other linting tools are used to fix the potential lint issues.
The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
The documentation has been modified accordingly, like docstring or example tutorials.

grimoire · 2025-10-27T03:42:37Z

lmdeploy/pytorch/engine/engine.py

+        self.speculative_config = speculative_config
+
+        if speculative_config is not None:
+            engine_config.prefill_interval = 16


this change was used for debugging and can be removed

grimoire · 2025-10-27T03:45:22Z

lmdeploy/pytorch/engine/engine.py

+            device_type=engine_config.device_type,
+            distributed_executor_backend=engine_config.distributed_executor_backend,
+            dtype=engine_config.dtype,
+            speculative_config=speculative_config,


Speculative_config is define outside pytorch engine. It is not a good design to use it in any module besides Engine.

grimoire · 2025-10-27T03:46:14Z

lmdeploy/pytorch/engine/engine.py

        # input ids
        token_ids = [msg.token_ids for msg in messages]
+
+        # spec decode


remove the comment.

grimoire · 2025-10-27T03:47:02Z

lmdeploy/pytorch/engine/engine.py


+    def _debug_spec_stats(self, batched_outputs: BatchedOutputs, is_decoding: bool = False):
+        """Debugging spec stats."""
+        is_debugging = True


Do we need this after release the feature?

It's just for debugging.

grimoire · 2025-10-27T04:21:17Z

lmdeploy/pytorch/config.py

            logger.warning(f'Overriding HF config with {hf_overrides}')
            override_hf_config(model_config.hf_config, hf_overrides)

+        # for serialization of transformers modules


This might not work

It works with tp case on one node, but not teste on dp case on multiple nodes

grimoire · 2025-10-27T04:33:38Z

lmdeploy/pytorch/engine/model_agent.py

    inputs: ModelInputs,
    cache_engine: CacheEngine,
    stream: torch.cuda.Stream = None,
+    output_position_ids: bool = False,


output position_ids is cheap, we can always output it.

grimoire · 2025-10-27T04:43:04Z

lmdeploy/pytorch/engine/model_agent.py

-                     f'batch_size={inputs.seq_length.size(0)} '
-                     f'num_tokens={inputs.input_ids.size(-1)} '
-                     f'is_decoding={inputs.is_decoding}')
+        logger.info(f'<ForwardTask> rank[{rank}]: '


info this would be too verbose

grimoire · 2025-10-27T04:46:17Z

lmdeploy/pytorch/engine/mp_engine/zmq_engine.py

+    def __init__(self,
+                 model_path: str,
+                 engine_config: PytorchEngineConfig = None,
+                 speculative_config: SpeculativeConfig = None,


Does RayEngine support speculative_config?

grimoire · 2025-10-27T04:48:18Z

lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py

    return k_states, v_states
+
+
+@triton.testing.perf_report(


remove bench

grimoire · 2025-10-27T04:56:03Z

lmdeploy/pytorch/models/utils/cudagraph.py

        input_buffers['position_ids'] = torch.zeros((1, max_tokens), dtype=torch.int64, device=device)
-        if getattr(self.config, 'use_flash_mla', False) is True:
-            import flash_mla
+        seqlens_dtype = torch.int64


when would we need int64?

the default is int64. while mla, fa3 needs int32.

grimoire · 2025-10-27T07:31:14Z

lmdeploy/pytorch/strategies/ar_spec/cudagraph.py

+        """Get max tokens."""
+        num_tokens = input_ids.size(1)
+        orig_batch = q_seqlens.size(0)
+        if num_tokens == orig_batch:


I do not think sending tensor here is a good idea.

grimoire · 2025-10-27T07:31:51Z

lmdeploy/pytorch/paging/scheduler.py

        self.scheduler_config = scheduler_config
        self.cache_config = cache_config
-
+        self.num_spec_tokens = num_spec_tokens


Has this value been used?

grimoire · 2025-10-27T07:56:03Z

lmdeploy/pytorch/spec_decode/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+from .deepseek_mtp import DeepseekMTP  # noqa F401


I don't want model w/o spec decoding loading these modules.

grimoire · 2025-10-27T08:32:58Z

lmdeploy/pytorch/spec_decode/proposers/base.py

+    def get_logits(self, hidden_states: torch.Tensor):
+        """Get logits of model output."""
+        draft_model = self.model
+        if not isinstance(draft_model, torch.nn.Module):


graph_runner has expose get_logits of model.

Yes. but eagle do not have get_logits while eagle3 has . Base on graph_runner's get_logits method, we cannot differ these two. That's why here check if original model has get_logits

grimoire · 2025-10-27T09:09:35Z

lmdeploy/pytorch/spec_decode/reject_sampler.py

+    expected_output_token_ids = torch.tensor([[0, 1, 2], [0, -1, -1], [1, -1, -1]], dtype=torch.long).cuda()
+
+    draft_probs = None
+    target_probs, draft_token_ids, bonus_token_ids, max_spec_len = torch.load('tmp.pt')


This test requires tmp.pt, do not place it here.
Could we add it in unit test?

RunningLeon requested a review from grimoire September 9, 2025 06:26

RunningLeon added 2 commits October 13, 2025 14:04

support spec with tp and cudagraph

465b533

fa3 cudagraph

2c18726

RunningLeon force-pushed the spec-decode branch from 28ddcc1 to 2c18726 Compare October 17, 2025 04:22

RunningLeon added 4 commits October 22, 2025 14:24

only use fa3 decode for spec decoding model

693998a

update req

68d4705

add simple doc

af5efb1

fix conficts with main

0aeb651

RunningLeon changed the title ~~[WIP]: Support speculative decoding~~ [Feature]: Support speculative decoding Oct 23, 2025

RunningLeon marked this pull request as ready for review October 23, 2025 09:04

grimoire reviewed Oct 27, 2025

View reviewed changes

lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py Outdated

return k_states, v_states

@triton.testing.perf_report(

Copy link

Collaborator

grimoire Oct 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove bench

grimoire reviewed Oct 27, 2025

View reviewed changes

resolve comments

f3f3634

		@@ -0,0 +1,9 @@
		# Copyright (c) OpenMMLab. All rights reserved.

		from .deepseek_mtp import DeepseekMTP # noqa F401

Uh oh!

[Feature]: Support speculative decoding #3945

Are you sure you want to change the base?

[Feature]: Support speculative decoding #3945

Uh oh!

Conversation

RunningLeon commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Examples

pipeline

serving

BC-breaking (Optional)

Use cases (Optional)

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

grimoire Oct 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

RunningLeon commented Sep 9, 2025 •

edited

Loading

grimoire Oct 27, 2025 •

edited

Loading