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

Remove overlap thread (#11210)

Browse Source 
Co-authored-by: Lianmin Zheng <15100009+merrymercy@users.noreply.github.com>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2025-10-07 20:12:12 +08:00
committed by GitHub
co-authored by Lianmin Zheng Hanming Lu
parent 24bc3fb0f9
commit 1519a89cfd
14 changed files with 280 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tp_worker.py
+33 -28
View File
@@ -15,14 +15,12 @@
from __future__ import annotations
import logging
import threading
from typing import TYPE_CHECKING, Optional, Tuple, Union
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.io_struct import (
DestroyWeightsUpdateGroupReqInput,
GetWeightsByNameReqInput,
@@ -36,13 +34,10 @@ from sglang.srt.managers.io_struct import (
UpdateWeightsFromTensorReqInput,
)
from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_args_dict
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardBatchOutput,
PPProxyTensors,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
@@ -236,9 +231,8 @@ class TpModelWorker:
def forward_batch_generation(
self,
model_worker_batch: ModelWorkerBatch,
launch_done: Optional[threading.Event] = None,
is_verify: bool = False,
) -> ForwardBatchOutput:
) -> GenerationBatchResult:
# update the consumer index of hicache to the running batch
self.set_hicache_consumer(model_worker_batch.hicache_consumer_index)
@@ -256,32 +250,43 @@ class TpModelWorker:
logits_output, can_run_cuda_graph = self.model_runner.forward(
forward_batch, pp_proxy_tensors=pp_proxy_tensors
)
if launch_done is not None:
launch_done.set()
skip_sample = is_verify or model_worker_batch.is_prefill_only
next_token_ids = None
if not skip_sample:
next_token_ids = self.model_runner.sample(logits_output, forward_batch)
elif model_worker_batch.return_logprob and not is_verify:
# NOTE: Compute logprobs without full sampling
self.model_runner.compute_logprobs_only(
logits_output, model_worker_batch
)
return ForwardBatchOutput(
batch_result = GenerationBatchResult(
logits_output=logits_output,
next_token_ids=next_token_ids,
can_run_cuda_graph=can_run_cuda_graph,
)
if is_verify:
# Skip sampling and return logits for target forward
return batch_result
if model_worker_batch.delay_sample_launch:
batch_result.delay_sample_launch = True
batch_result.forward_batch = forward_batch
return batch_result
if model_worker_batch.is_prefill_only:
# For prefill-only requests, create dummy token IDs on CPU
batch_result.next_token_ids = torch.zeros_like(
model_worker_batch.input_ids, dtype=torch.long
)
if model_worker_batch.return_logprob:
# NOTE: Compute logprobs without full sampling
self.model_runner.compute_logprobs_only(
logits_output, model_worker_batch
)
else:
batch_result.next_token_ids = self.model_runner.sample(
logits_output, forward_batch
)
return batch_result
else:
pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
forward_batch,
pp_proxy_tensors=pp_proxy_tensors,
)
return ForwardBatchOutput(
pp_proxy_tensors=pp_proxy_tensors,
return GenerationBatchResult(
pp_hidden_states_proxy_tensors=pp_proxy_tensors,
can_run_cuda_graph=can_run_cuda_graph,
)