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

Init TBO with dp_padded batch (#11423)

Browse Source 
Co-authored-by: Cheng Wan <wan4ch@gmail.com>
Co-authored-by: Yuhao Yao <37280700+yuhyao@users.noreply.github.com>
This commit is contained in:
Quanfeng Li
2025-12-02 10:34:26 -08:00
committed by GitHub
co-authored by Cheng Wan Yuhao Yao
parent 0141ca370f
commit 427b08e24d
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/forward_batch_info.py
+9 -9
View File
@@ -360,8 +360,6 @@ class ForwardBatch:
batch: ModelWorkerBatch, batch: ModelWorkerBatch,
model_runner: ModelRunner, model_runner: ModelRunner,
): ):
from sglang.srt.two_batch_overlap import TboForwardBatchPreparer
ret = cls( ret = cls(
forward_mode=batch.forward_mode, forward_mode=batch.forward_mode,
batch_size=len(batch.seq_lens), batch_size=len(batch.seq_lens),
@@ -436,9 +434,6 @@ class ForwardBatch:
if ret.forward_mode.is_idle(): if ret.forward_mode.is_idle():
ret.positions = torch.empty((0,), dtype=torch.int64, device=device) ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
TboForwardBatchPreparer.prepare(
ret, is_draft_worker=model_runner.is_draft_worker
)
return ret return ret
# Override the positions with diffusion LLM or spec_info # Override the positions with diffusion LLM or spec_info
@@ -496,10 +491,6 @@ class ForwardBatch:
if model_runner.server_args.enable_lora: if model_runner.server_args.enable_lora:
model_runner.lora_manager.prepare_lora_batch(ret) model_runner.lora_manager.prepare_lora_batch(ret)
TboForwardBatchPreparer.prepare(
ret, is_draft_worker=model_runner.is_draft_worker
)
return ret return ret
def merge_mm_inputs(self) -> Optional[MultimodalInputs]: def merge_mm_inputs(self) -> Optional[MultimodalInputs]:
@@ -728,6 +719,8 @@ class ForwardBatch:
) )
def prepare_mlp_sync_batch(self, model_runner: ModelRunner): def prepare_mlp_sync_batch(self, model_runner: ModelRunner):
from sglang.srt.two_batch_overlap import TboForwardBatchPreparer
assert self.global_num_tokens_cpu is not None assert self.global_num_tokens_cpu is not None
assert self.global_num_tokens_for_logprob_cpu is not None assert self.global_num_tokens_for_logprob_cpu is not None
@@ -792,6 +785,8 @@ class ForwardBatch:
) )
else: else:
bs = self.batch_size = num_tokens bs = self.batch_size = num_tokens
elif self.forward_mode.is_extend():
self.extend_num_tokens = num_tokens
# padding # padding
self._pad_inputs_to_size(model_runner, num_tokens, bs) self._pad_inputs_to_size(model_runner, num_tokens, bs)
@@ -799,10 +794,15 @@ class ForwardBatch:
global_num_tokens_pinned = torch.tensor(global_num_tokens, pin_memory=True) global_num_tokens_pinned = torch.tensor(global_num_tokens, pin_memory=True)
self.global_num_tokens_gpu.copy_(global_num_tokens_pinned, non_blocking=True) self.global_num_tokens_gpu.copy_(global_num_tokens_pinned, non_blocking=True)
TboForwardBatchPreparer.prepare(
batch=self, is_draft_worker=model_runner.is_draft_worker
)
def _pad_inputs_to_size(self, model_runner: ModelRunner, num_tokens, bs): def _pad_inputs_to_size(self, model_runner: ModelRunner, num_tokens, bs):
# padding # padding
self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens) self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens)
self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs) self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs)
self.lora_ids.extend((bs - len(self.lora_ids)) * [None])
seq_len_fill_value = ( seq_len_fill_value = (
model_runner.attn_backend.get_cuda_graph_seq_len_fill_value() model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()