Skip to content

[BugFix] support real_bsz #3027

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lizexu123 wants to merge 5 commits into from
Closed

[BugFix] support real_bsz #3027

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9fa0709
support real_bsz
lizexu-baidu Jul 26, 2025
6aeba3b
add ENABLE_V1_KVCACHE_SCHEDULER situation
lizexu-baidu Jul 27, 2025
a52b7c5
add all workers
lizexu-baidu Jul 28, 2025
d9bc3bc
fix
lizexu-baidu Jul 28, 2025
71f7680
fix xpu_model_runner
lizexu-baidu Jul 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions fastdeploy/worker/gcu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.
"""

import copy
import os
import time
from typing import List, Optional
Expand Down Expand Up @@ -146,7 +147,7 @@ def _init_logits_processor(self, request):
schemata_key,
)

def insert_prefill_inputs(self, req_dicts: List[Request]):
def insert_prefill_inputs(self, req_dicts: List[Request], num_running_requests: int = None):
"""
Process inputs for prefill tasks and insert it to share_inputs buffer
"""
Expand Down Expand Up @@ -187,7 +188,7 @@ def get_attr_from_request(request, attr, default_value=None):
self.share_inputs["prompt_ids"][idx : idx + 1, :length] = np.array(request.prompt_token_ids)
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = 0
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = length
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = 1
self.tmp_seq_lens_this_time[idx : idx + 1] = 1
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = 0
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = length
self.share_inputs["prompt_lens"][idx : idx + 1] = length
Expand All @@ -199,7 +200,7 @@ def get_attr_from_request(request, attr, default_value=None):
request.draft_token_ids[0:num_prefill_send_token],
dtype="int64",
)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = num_prefill_send_token
self.tmp_seq_lens_this_time[idx : idx + 1] = num_prefill_send_token
else:
self.share_inputs["pre_ids"][idx : idx + 1] = -1
self.share_inputs["step_idx"][idx : idx + 1] = 0
Expand All @@ -216,14 +217,14 @@ def get_attr_from_request(request, attr, default_value=None):
)
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = token_chunk_size
self.tmp_seq_lens_this_time[idx : idx + 1] = token_chunk_size
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = token_chunk_size
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = token_chunk_size
self.share_inputs["prompt_lens"][idx : idx + 1] = token_chunk_size
else:
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = length
self.tmp_seq_lens_this_time[idx : idx + 1] = length
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = length
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = length
self.share_inputs["prompt_lens"][idx : idx + 1] = length
Expand Down Expand Up @@ -279,6 +280,7 @@ def get_attr_from_request(request, attr, default_value=None):

if self.speculative_method in ["mtp"]:
self.proposer.insert_prefill_inputs(req_dicts)
self.share_inputs["seq_lens_this_time"] = copy.deepcopy(self.tmp_seq_lens_this_time[:num_running_requests])

def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decode_len: int):
"""Set dummy prefill inputs to share_inputs"""
Expand All @@ -297,7 +299,7 @@ def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decod
self.share_inputs["input_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
self.share_inputs["prompt_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
self.share_inputs["eos_token_id"][:] = np.array([2], dtype="int64").reshape(-1, 1)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = input_length
self.tmp_seq_lens_this_time[idx : idx + 1] = input_length
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = input_length
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = input_length
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = 0
Expand All @@ -315,6 +317,7 @@ def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decod
self.share_inputs["block_tables"][idx : idx + 1, :block_num] = np.arange(
idx * block_num, (idx + 1) * block_num, 1
)
self.share_inputs["seq_lens_this_time"] = self.tmp_seq_lens_this_time

def _init_share_inputs(self, max_num_seqs: int):
"""
Expand Down Expand Up @@ -365,7 +368,7 @@ def _init_share_inputs(self, max_num_seqs: int):
self.share_inputs["max_length"] = paddle.full(
[max_num_seqs, 1], self.model_config.max_model_len, dtype="int64"
)
self.share_inputs["seq_lens_this_time"] = paddle.full(max_num_seqs, 0, dtype="int32")
self.tmp_seq_lens_this_time = paddle.full(max_num_seqs, 0, dtype="int32")
self.share_inputs["seq_lens_encoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
self.share_inputs["seq_lens_decoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
self.share_inputs["step_seq_lens_encoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
Expand Down Expand Up @@ -879,6 +882,7 @@ def _get_skip_idx(self, model_forward_batch: Optional[List[Request]] = None):
def execute_model(
self,
model_forward_batch: Optional[List[Request]] = None,
num_running_requests: int = None,
) -> Optional[ModelRunnerOutput]:
"""
The Entrance of model execute.
Expand Down Expand Up @@ -1011,6 +1015,7 @@ class at the server level, which is too granular for ModelRunner.

self._update_chunked_prefill(model_forward_batch)
self._add_cache(model_forward_batch)
self.tmp_seq_lens_this_time[:num_running_requests] = copy.deepcopy(self.share_inputs["seq_lens_this_time"])
return None

def _add_cache(self, model_forward_batch) -> None:
Expand Down
7 changes: 4 additions & 3 deletions fastdeploy/worker/gcu_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,17 +105,18 @@ def initialize_cache(self, num_gpu_blocks: int) -> None:
def execute_model(
self,
model_forward_batch: Optional[List[Request]] = None,
num_running_requests: int = None,
) -> Optional[ModelRunnerOutput]:
""" """
output = self.model_runner.execute_model(model_forward_batch)
output = self.model_runner.execute_model(model_forward_batch, num_running_requests)
return output

def preprocess_new_task(self, req_dicts: List[Request]) -> None:
def preprocess_new_task(self, req_dicts: List[Request], num_running_requests: int) -> None:
"""Process new requests and then start the decode loop
TODO(gongshaotian):The scheduler should schedule the handling of prefill,
and workers and modelrunners should not perceive it.
"""
self.model_runner.insert_prefill_inputs(req_dicts=req_dicts)
self.model_runner.insert_prefill_inputs(req_dicts=req_dicts, num_running_requests=num_running_requests)

def graph_optimize_and_warm_up_model(self) -> None:
"""
Expand Down
30 changes: 19 additions & 11 deletions fastdeploy/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.
"""

import copy
import os
import time
from typing import List, Optional
Expand Down Expand Up @@ -193,7 +194,7 @@ def _init_logits_processor(self, request):

return self.guided_backend.get_logits_processor(schemata_key=schemata_key), schemata_key

def insert_tasks_v1(self, req_dicts: List[Request]):
def insert_tasks_v1(self, req_dicts: List[Request], num_running_requests: int = None):
"""
Process scheduler output tasks, used when ENABLE_V1_KVCACHE_SCHEDULER=1
"""
Expand Down Expand Up @@ -223,7 +224,7 @@ def insert_tasks_v1(self, req_dicts: List[Request]):
)
self.share_inputs["stop_flags"][idx : idx + 1] = False
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = prefill_start_index
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = length
self.tmp_seq_lens_this_time[idx : idx + 1] = length
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = length
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = 0
self.share_inputs["prompt_lens"][idx : idx + 1] = len(input_ids)
Expand All @@ -245,7 +246,7 @@ def insert_tasks_v1(self, req_dicts: List[Request]):
logger.debug(f"Handle preempted request {request} at idx {idx}")
self.share_inputs["block_tables"][idx : idx + 1, :] = -1
self.share_inputs["stop_flags"][idx : idx + 1] = True
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = 0
self.tmp_seq_lens_this_time[idx : idx + 1] = 0
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = 0
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = 0
self.share_inputs["is_block_step"][idx : idx + 1] = False
Expand Down Expand Up @@ -283,7 +284,9 @@ def insert_tasks_v1(self, req_dicts: List[Request]):
if has_prefill_task:
self.share_inputs["not_need_stop"][0] = True

def insert_prefill_inputs(self, req_dicts: List[Request]):
self.share_inputs["seq_lens_this_time"] = copy.deepcopy(self.tmp_seq_lens_this_time[:num_running_requests])

def insert_prefill_inputs(self, req_dicts: List[Request], num_running_requests: int = None):
"""
Process inputs for prefill tasks and insert it to share_inputs buffer
TODO(gongshaotian): Refactor this func
Expand Down Expand Up @@ -319,7 +322,7 @@ def insert_prefill_inputs(self, req_dicts: List[Request]):
self.share_inputs["prompt_ids"][idx : idx + 1, :length] = np.array(request.prompt_token_ids)
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = 0
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = length
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = 1
self.tmp_seq_lens_this_time[idx : idx + 1] = 1
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = 0
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = length
self.share_inputs["prompt_lens"][idx : idx + 1] = length
Expand All @@ -331,7 +334,7 @@ def insert_prefill_inputs(self, req_dicts: List[Request]):
request.draft_token_ids[0:num_prefill_send_token],
dtype="int64",
)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = num_prefill_send_token
self.tmp_seq_lens_this_time[idx : idx + 1] = num_prefill_send_token
else:
self.share_inputs["pre_ids"][idx : idx + 1] = -1
self.share_inputs["step_idx"][idx : idx + 1] = 0
Expand Down Expand Up @@ -366,7 +369,7 @@ def insert_prefill_inputs(self, req_dicts: List[Request]):
)
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = token_chunk_size
self.tmp_seq_lens_this_time[idx : idx + 1] = token_chunk_size
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = token_chunk_size
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = token_chunk_size
self.share_inputs["prompt_lens"][idx : idx + 1] = token_chunk_size
Expand All @@ -384,7 +387,7 @@ def insert_prefill_inputs(self, req_dicts: List[Request]):
else:
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["step_seq_lens_decoder"][idx : idx + 1] = request.get("seq_lens_decoder", 0)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = length
self.tmp_seq_lens_this_time[idx : idx + 1] = length
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = length
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = length
self.share_inputs["prompt_lens"][idx : idx + 1] = length
Expand Down Expand Up @@ -459,6 +462,8 @@ def get_attr_from_request(request, attr, default_value=None):
if self.speculative_method in ["mtp"]:
self.proposer.insert_prefill_inputs(req_dicts)

self.share_inputs["seq_lens_this_time"] = copy.deepcopy(self.tmp_seq_lens_this_time[:num_running_requests])

def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decode_len: int):
"""Set dummy prefill inputs to share_inputs"""
# NOTE(gongshaotian): The maximum decoding length is equal to the expected decoded tokens plus the eos token
Expand All @@ -477,7 +482,7 @@ def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decod
self.share_inputs["input_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
self.share_inputs["prompt_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
self.share_inputs["eos_token_id"][:] = np.array([2], dtype="int64").reshape(-1, 1)
self.share_inputs["seq_lens_this_time"][idx : idx + 1] = input_length
self.tmp_seq_lens_this_time[idx : idx + 1] = input_length
self.share_inputs["step_seq_lens_encoder"][idx : idx + 1] = input_length
self.share_inputs["seq_lens_encoder"][idx : idx + 1] = input_length
self.share_inputs["seq_lens_decoder"][idx : idx + 1] = 0
Expand All @@ -495,6 +500,7 @@ def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int, expected_decod
self.share_inputs["block_tables"][idx : idx + 1, :block_num] = np.arange(
idx * block_num, (idx + 1) * block_num, 1
)
self.share_inputs["seq_lens_this_time"] = self.tmp_seq_lens_this_time

def _init_share_inputs(self, max_num_seqs: int):
"""
Expand Down Expand Up @@ -545,7 +551,7 @@ def _init_share_inputs(self, max_num_seqs: int):
self.share_inputs["max_length"] = paddle.full(
[max_num_seqs, 1], self.model_config.max_model_len, dtype="int64"
)
self.share_inputs["seq_lens_this_time"] = paddle.full(max_num_seqs, 0, dtype="int32")
self.tmp_seq_lens_this_time = paddle.full(max_num_seqs, 0, dtype="int32")
self.share_inputs["seq_lens_encoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
self.share_inputs["seq_lens_decoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
self.share_inputs["step_seq_lens_encoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
Expand Down Expand Up @@ -1029,6 +1035,7 @@ def _dummy_run(
# 7. Updata 'infer_seed' and step_cuda()
self.share_inputs["infer_seed"].add_(self.infer_seed_increment)
self.share_inputs["infer_seed"][:] %= self.MAX_INFER_SEED

step_cuda(
self.share_inputs,
self.parallel_config.block_size,
Expand Down Expand Up @@ -1162,6 +1169,7 @@ def _get_skip_idx(self, model_forward_batch: Optional[List[Request]] = None):
def execute_model(
self,
model_forward_batch: Optional[List[Request]] = None,
num_running_requests: int = None,
) -> Optional[ModelRunnerOutput]:
"""
The Entrance of model execute.
Expand Down Expand Up @@ -1307,9 +1315,9 @@ class at the server level, which is too granular for ModelRunner.
self.speculative_config,
self.parallel_config.enable_prefix_caching,
)

self._update_chunked_prefill(model_forward_batch)
self._add_cache(model_forward_batch)
self.tmp_seq_lens_this_time[:num_running_requests] = copy.deepcopy(self.share_inputs["seq_lens_this_time"])
return None

def _add_cache(self, model_forward_batch) -> None:
Expand Down
9 changes: 5 additions & 4 deletions fastdeploy/worker/gpu_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,20 +175,21 @@ def initialize_cache(self, num_gpu_blocks: int) -> None:
def execute_model(
self,
model_forward_batch: Optional[List[Request]] = None,
num_running_requests: int = None,
) -> Optional[ModelRunnerOutput]:
""" """
output = self.model_runner.execute_model(model_forward_batch)
output = self.model_runner.execute_model(model_forward_batch, num_running_requests)
return output

def preprocess_new_task(self, req_dicts: List[Request]) -> None:
def preprocess_new_task(self, req_dicts: List[Request], num_running_requests: int) -> None:
"""Process new requests and then start the decode loop
TODO(gongshaotian):The scheduler should schedule the handling of prefill,
and workers and modelrunners should not perceive it.
"""
if envs.ENABLE_V1_KVCACHE_SCHEDULER:
self.model_runner.insert_tasks_v1(req_dicts=req_dicts)
self.model_runner.insert_tasks_v1(req_dicts=req_dicts, num_running_requests=num_running_requests)
else:
self.model_runner.insert_prefill_inputs(req_dicts=req_dicts)
self.model_runner.insert_prefill_inputs(req_dicts=req_dicts, num_running_requests=num_running_requests)

def graph_optimize_and_warm_up_model(self) -> None:
"""
Expand Down
Loading
Loading