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

llama model and llama eagle3 model support dp-attn (#15268)

Browse Source
This commit is contained in:
chenxu140
2026-01-12 08:54:56 +08:00
committed by GitHub
parent 503c3d9566
commit 38a88479c6
1 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -13,6 +13,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend
from sglang.srt.layers.dp_attention import get_attention_tp_group
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
@@ -118,7 +119,15 @@ 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(), speculative_moe_a2a_backend_context():
# Init draft worker
if server_args.enable_dp_attention and self.speculative_algorithm.is_eagle3():
ctx = draft_tp_context(get_attention_tp_group())
else:
ctx = empty_context()
with (
ctx
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
@@ -626,7 +635,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
# Draft prefill
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
with self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
batch_output.next_draft_input = (
self.draft_worker._draft_extend_for_prefill(
model_worker_batch,
@@ -644,14 +655,18 @@ class EAGLEWorkerV2(BaseSpecWorker):
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
with self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
), 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)
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
with self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.draft_worker._draft_extend_for_decode(
model_worker_batch, batch_output
)