chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

support mtp with deepseek r1 nvfp4 model (#13115)

Browse Source 
Co-authored-by: Trevor Morris <tmorris@nvidia.com>
This commit is contained in:
Rain Jiang
2025-12-06 00:45:54 -08:00
committed by GitHub
parent 42fcf5438f
commit ea177372bd
11 changed files with 123 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/speculative/eagle_info.py
View File

@@ -67,6 +67,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
seq_lens_cpu: torch.Tensor
grammar: BaseGrammarObject = None
# Shape info for padding
num_tokens_per_batch: int = -1
def __post_init__(self):
super().__init__(SpecInputType.EAGLE_VERIFY)

18
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -10,7 +10,10 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i
)
from sglang.srt.layers.dp_attention import get_attention_tp_group
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs
from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -132,7 +135,9 @@ class EAGLEWorker(TpModelWorker):
ctx = draft_tp_context(get_attention_tp_group())
else:
ctx = empty_context()
with ctx, speculative_moe_backend_context():
with (
ctx
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
super().__init__(
server_args=server_args,
gpu_id=gpu_id,
@@ -183,7 +188,7 @@ class EAGLEWorker(TpModelWorker):
)
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.init_attention_backend()
self.init_cuda_graphs()
@@ -276,7 +281,7 @@ class EAGLEWorker(TpModelWorker):
)
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.forward_draft_extend(
batch, logits_output.hidden_states, next_token_ids, seq_lens_cpu
)
@@ -289,7 +294,7 @@ class EAGLEWorker(TpModelWorker):
else:
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
spec_info = self.draft(batch)
logits_output, verify_output, model_worker_batch, can_run_cuda_graph = (
self.verify(batch, spec_info)
@@ -297,7 +302,7 @@ class EAGLEWorker(TpModelWorker):
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
# NOTE: We should use `check_forward_draft_extend_after_decode`
# when DP attention is enabled, but it is slow. Skip it for now.
if (
@@ -665,6 +670,7 @@ class EAGLEWorker(TpModelWorker):
def verify(self, batch: ScheduleBatch, spec_info: EagleVerifyInput):
spec_info.prepare_for_verify(batch, self.page_size)
spec_info.num_tokens_per_batch = self.speculative_num_steps + 1
batch.return_hidden_states = False
batch.forward_mode = (
ForwardMode.TARGET_VERIFY

35
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -12,7 +12,10 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_r
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -112,7 +115,7 @@ class EagleDraftWorker(BaseDraftWorker):
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
with empty_context(), speculative_moe_backend_context():
with empty_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
@@ -140,7 +143,7 @@ class EagleDraftWorker(BaseDraftWorker):
)
with self.draft_tp_context(
self.draft_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.init_attention_backend()
self.init_cuda_graphs()
@@ -611,12 +614,15 @@ class EAGLEWorkerV2(BaseSpecWorker):
# Draft prefill
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST
batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill(
model_worker_batch,
batch_output.logits_output.hidden_states,
batch_output.next_token_ids,
)
return batch_output
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
batch_output.next_draft_input = (
self.draft_worker._draft_extend_for_prefill(
model_worker_batch,
batch_output.logits_output.hidden_states,
batch_output.next_token_ids,
)
)
return batch_output
else:
if model_worker_batch.spec_info is None:
model_worker_batch.spec_info = EagleDraftInput.create_idle_input(
@@ -626,11 +632,17 @@ class EAGLEWorkerV2(BaseSpecWorker):
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch)
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
verify_input: EagleVerifyInput = self.draft_worker.draft(
model_worker_batch
)
assert verify_input.is_verify_input()
model_worker_batch.spec_info = verify_input
batch_output = self.verify(model_worker_batch)
self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output)
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.draft_worker._draft_extend_for_decode(
model_worker_batch, batch_output
)
return batch_output
def verify(self, batch: ModelWorkerBatch):
@@ -643,6 +655,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
# Parse args
verify_input: EagleVerifyInput = batch.spec_info
verify_input.num_tokens_per_batch = self.speculative_num_steps + 1
bs = len(batch.seq_lens)
# Batch 1: Target verify

9
python/sglang/srt/speculative/standalone_worker.py
View File

@@ -3,7 +3,10 @@ from typing import Optional
import torch
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.eagle_worker import EAGLEWorker
@@ -67,7 +70,7 @@ class StandaloneWorker(EAGLEWorker):
self.hot_token_id = None
# Init draft worker
with empty_context(), speculative_moe_backend_context():
with empty_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
TpModelWorker.__init__(
self,
server_args=server_args,
@@ -91,7 +94,7 @@ class StandaloneWorker(EAGLEWorker):
)
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.init_attention_backend()
self.init_cuda_graphs()