Skip to content

Adjust InputeMetadata and ScheduleBatch #981

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

Merged
hnyls2002 merged 13 commits into from
Aug 8, 2024
Merged

Adjust InputeMetadata and ScheduleBatch #981

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
557492c
clean
hnyls2002 Aug 8, 2024
80b067f
move init_triton
hnyls2002 Aug 8, 2024
efe19de
simplify the code
hnyls2002 Aug 8, 2024
c09f580
position -> int64
hnyls2002 Aug 8, 2024
05014c1
update
hnyls2002 Aug 8, 2024
8f85613
from_batch
hnyls2002 Aug 8, 2024
74f3b80
simplify the code
hnyls2002 Aug 8, 2024
1e6c735
optimize position computation
hnyls2002 Aug 8, 2024
7f07c73
simplify code
hnyls2002 Aug 8, 2024
366636a
reduce cpu/gpu sync
hnyls2002 Aug 8, 2024
beb5a94
remove some vars in schedule batch
hnyls2002 Aug 8, 2024
8cfb80b
move multimultal init
hnyls2002 Aug 8, 2024
0f9f284
clean
hnyls2002 Aug 8, 2024
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
60 changes: 16 additions & 44 deletions python/sglang/srt/managers/schedule_batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,7 +307,6 @@ class ScheduleBatch:
input_ids: torch.Tensor = None
req_pool_indices: torch.Tensor = None
seq_lens: torch.Tensor = None
prefix_lens: torch.Tensor = None
position_ids_offsets: torch.Tensor = None
out_cache_loc: torch.Tensor = None
extend_num_tokens: int = None
Expand All @@ -316,11 +315,6 @@ class ScheduleBatch:
return_logprob: bool = False
top_logprobs_nums: List[int] = None

# For multimodal
pixel_values: List[torch.Tensor] = None
image_sizes: List[List[int]] = None
image_offsets: List[int] = None

# Batched sampling params
temperatures: torch.Tensor = None
top_ps: torch.Tensor = None
Expand Down Expand Up @@ -412,59 +406,40 @@ def batch_sampling_params(self, vocab_size, int_token_logit_bias):
self.logit_bias[i][: len(int_token_logit_bias)] = int_token_logit_bias

def prepare_for_extend(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
device = "cuda"
bs = self.batch_size()
reqs = self.reqs
input_ids = [r.input_ids[len(r.prefix_indices) :] for r in reqs]
prefix_indices = [r.prefix_indices for r in reqs]

# Handle prefix
extend_lens = []
prefix_lens = []
extend_num_tokens = sum(len(ids) for ids in input_ids)
seq_lens = []

# Allocate memory
req_pool_indices_cpu = self.alloc_req_slots(bs)
out_cache_loc = self.alloc_token_slots(extend_num_tokens)

pt = 0
for i, req in enumerate(reqs):
req.req_pool_idx = req_pool_indices_cpu[i]
extend_lens.append(len(input_ids[i]))
pre_len, seq_len = len(req.prefix_indices), len(req.input_ids)
ext_len = seq_len - pre_len
seq_lens.append(seq_len)

if len(prefix_indices[i]) == 0:
prefix_lens.append(0)
else:
prefix_lens.append(len(prefix_indices[i]))
if pre_len > 0:
self.req_to_token_pool.req_to_token[req.req_pool_idx][
: len(prefix_indices[i])
] = prefix_indices[i]

seq_lens.append(prefix_lens[-1] + extend_lens[-1])
:pre_len
] = req.prefix_indices

# Allocate memory
seq_lens, prefix_lens = np.array(seq_lens), np.array(prefix_lens)
extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
out_cache_loc = self.alloc_token_slots(extend_num_tokens)

pt = 0
for i, req in enumerate(reqs):
self.req_to_token_pool.req_to_token[req.req_pool_idx][
prefix_lens[i] : prefix_lens[i] + extend_lens[i]
] = out_cache_loc[pt : pt + extend_lens[i]]
pt += extend_lens[i]
self.req_to_token_pool.req_to_token[req.req_pool_idx][pre_len:seq_len] = (
out_cache_loc[pt : pt + ext_len]
)
pt += ext_len

# Set fields
with torch.device("cuda"):
self.input_ids = torch.tensor(sum(input_ids, []), dtype=torch.int32)
self.req_pool_indices = torch.tensor(req_pool_indices_cpu)
self.seq_lens = torch.tensor(seq_lens, dtype=torch.int32)
self.position_ids_offsets = torch.zeros((bs,), dtype=torch.int32)

self.pixel_values = [r.pixel_values for r in reqs]
self.image_sizes = [r.image_size for r in reqs]
self.image_offsets = [
(r.image_offset - p_len) if r.image_offset is not None else 0
for r, p_len in zip(reqs, prefix_lens)
]
self.prefix_lens = torch.tensor(prefix_lens, dtype=torch.int32, device=device)
self.position_ids_offsets = torch.zeros((bs,), dtype=torch.int64)

self.extend_num_tokens = extend_num_tokens
self.out_cache_loc = out_cache_loc
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
Expand Down Expand Up @@ -642,7 +617,6 @@ def prepare_for_decode(self, input_ids=None):
]
self.input_ids = torch.tensor(input_ids, dtype=torch.int32, device="cuda")
self.seq_lens.add_(1)
self.prefix_lens = None

# Alloc mem
bs = self.batch_size()
Expand All @@ -667,7 +641,6 @@ def filter_batch(self, unfinished_indices: List[int]):
self.seq_lens = self.seq_lens[new_indices]
self.input_ids = None
self.req_pool_indices = self.req_pool_indices[new_indices]
self.prefix_lens = None
self.position_ids_offsets = self.position_ids_offsets[new_indices]
self.out_cache_loc = None
self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in unfinished_indices]
Expand All @@ -692,7 +665,6 @@ def merge(self, other: "ScheduleBatch"):
[self.req_pool_indices, other.req_pool_indices]
)
self.seq_lens = torch.concat([self.seq_lens, other.seq_lens])
self.prefix_lens = None
self.position_ids_offsets = torch.concat(
[self.position_ids_offsets, other.position_ids_offsets]
)
Expand Down
18 changes: 9 additions & 9 deletions python/sglang/srt/model_executor/cuda_graph_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from sglang.srt.model_executor.forward_batch_info import (
ForwardMode,
InputMetadata,
init_flashinfer_args,
update_flashinfer_indices,
)
from sglang.srt.utils import monkey_patch_vllm_all_gather

Expand Down Expand Up @@ -165,7 +165,7 @@ def capture_one_batch_size(self, bs, forward):
paged_kv_indices_buffer=self.flashinfer_kv_indices,
paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[:bs],
)
init_flashinfer_args(
update_flashinfer_indices(
ForwardMode.DECODE,
self.model_runner,
req_pool_indices,
Expand All @@ -176,19 +176,19 @@ def capture_one_batch_size(self, bs, forward):

# Run and capture
def run_once():
input_metadata = InputMetadata.create(
self.model_runner,
input_metadata = InputMetadata(
forward_mode=ForwardMode.DECODE,
batch_size=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
prefix_lens=None,
position_ids_offsets=position_ids_offsets,
req_to_token_pool=self.model_runner.req_to_token_pool,
token_to_kv_pool=self.model_runner.token_to_kv_pool,
out_cache_loc=out_cache_loc,
return_logprob=False,
top_logprobs_nums=0,
skip_flashinfer_init=True,
positions=(seq_lens - 1).to(torch.int64),
flashinfer_decode_wrapper=flashinfer_decode_wrapper,
)
input_metadata.flashinfer_decode_wrapper = flashinfer_decode_wrapper

return forward(input_ids, input_metadata.positions, input_metadata)

Expand Down Expand Up @@ -222,7 +222,7 @@ def replay(self, batch: ScheduleBatch):
self.out_cache_loc[:raw_bs] = batch.out_cache_loc

# FlashInfer inputs
init_flashinfer_args(
update_flashinfer_indices(
ForwardMode.DECODE,
self.model_runner,
self.req_pool_indices[:bs],
Expand Down
Loading