schedule_batch.py 32.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
"""
Copyright 2023-2024 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

16
17
18
19
20
21
22
23
24
25
26
27
28
"""
Store information about requests and batches.

The following is the flow of data structures for a batch:

ScheduleBatch -> ModelWorkerBatch -> ForwardBatch

- ScheduleBatch is managed by `scheduler.py::Scheduler`.
  It contains high-level scheduling data. Most of the data is on the CPU.
- ModelWorkerBatch is managed by `tp_worker.py::TpModelWorker`.
- ForwardBatch is managed by `model_runner.py::ModelRunner`.
  It contains low-level tensor data. Most of the data consists of GPU tensors.
"""
Lianmin Zheng's avatar
Lianmin Zheng committed
29

Ying Sheng's avatar
Ying Sheng committed
30
import logging
31
from dataclasses import dataclass
32
from typing import List, Optional, Tuple, Union
Lianmin Zheng's avatar
Lianmin Zheng committed
33
34

import torch
35

Liangsheng Yin's avatar
Liangsheng Yin committed
36
from sglang.global_config import global_config
37
38
from sglang.srt.constrained import RegexGuide
from sglang.srt.constrained.jump_forward import JumpForwardMap
39
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
40
from sglang.srt.mem_cache.chunk_cache import ChunkCache
41
from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
42
from sglang.srt.model_executor.forward_batch_info import ForwardMode
43
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
44
from sglang.srt.sampling.sampling_params import SamplingParams
45
from sglang.srt.server_args import ServerArgs
Liangsheng Yin's avatar
Liangsheng Yin committed
46
47

INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
Lianmin Zheng's avatar
Lianmin Zheng committed
48

49
50
# Put some global args for easy access
global_server_args_dict = {
51
52
53
    "attention_backend": ServerArgs.attention_backend,
    "sampling_backend": ServerArgs.sampling_backend,
    "triton_attention_reduce_in_fp32": ServerArgs.triton_attention_reduce_in_fp32,
Ke Bao's avatar
Ke Bao committed
54
    "disable_mla": ServerArgs.disable_mla,
55
    "torchao_config": ServerArgs.torchao_config,
56
57
}

Lianmin Zheng's avatar
Lianmin Zheng committed
58

Ying Sheng's avatar
Ying Sheng committed
59
60
61
logger = logging.getLogger(__name__)


62
63
64
class BaseFinishReason:
    def __init__(self, is_error: bool = False):
        self.is_error = is_error
Lianmin Zheng's avatar
Lianmin Zheng committed
65

66
    def to_json(self):
67
        raise NotImplementedError()
68
69
70


class FINISH_MATCHED_TOKEN(BaseFinishReason):
Mingyi's avatar
Mingyi committed
71
    def __init__(self, matched: Union[int, List[int]]):
72
73
74
        super().__init__()
        self.matched = matched

75
76
77
78
79
    def to_json(self):
        return {
            "type": "stop",  # to match OpenAI API's return value
            "matched": self.matched,
        }
80
81


82
83
class FINISH_MATCHED_STR(BaseFinishReason):
    def __init__(self, matched: str):
84
        super().__init__()
85
        self.matched = matched
86

87
88
89
90
91
    def to_json(self):
        return {
            "type": "stop",  # to match OpenAI API's return value
            "matched": self.matched,
        }
92
93


94
95
class FINISH_LENGTH(BaseFinishReason):
    def __init__(self, length: int):
96
        super().__init__()
97
        self.length = length
98

99
100
101
102
103
    def to_json(self):
        return {
            "type": "length",  # to match OpenAI API's return value
            "length": self.length,
        }
104
105
106
107
108
109


class FINISH_ABORT(BaseFinishReason):
    def __init__(self):
        super().__init__(is_error=True)

110
111
112
113
    def to_json(self):
        return {
            "type": "abort",
        }
114

Lianmin Zheng's avatar
Lianmin Zheng committed
115

Liangsheng Yin's avatar
Liangsheng Yin committed
116
117
@dataclass
class ImageInputs:
118
119
    """The image related inputs."""

Liangsheng Yin's avatar
Liangsheng Yin committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
    pixel_values: torch.Tensor
    image_hash: int
    image_sizes: Optional[list] = None
    image_offsets: Optional[list] = None
    pad_values: Optional[list] = None
    modalities: Optional[list] = None

    image_embeds: Optional[List[torch.Tensor]] = None
    aspect_ratio_ids: Optional[List[torch.Tensor]] = None
    aspect_ratio_mask: Optional[List[torch.Tensor]] = None

    @staticmethod
    def from_dict(obj, vocab_size):
        # Use image hash as fake token_ids, which is then used for prefix matching
        ret = ImageInputs(
            pixel_values=obj["pixel_values"],
            image_hash=hash(tuple(obj["image_hashes"])),
        )
        image_hash = ret.image_hash
        ret.pad_values = [
            (image_hash) % vocab_size,
            (image_hash >> 16) % vocab_size,
            (image_hash >> 32) % vocab_size,
            (image_hash >> 64) % vocab_size,
        ]
        ret.image_sizes = obj["image_sizes"]
        # Only when pixel values is not None we have modalities
147
        ret.modalities = obj["modalities"] or ["image"]
Liangsheng Yin's avatar
Liangsheng Yin committed
148
149
150
        return ret


Lianmin Zheng's avatar
Lianmin Zheng committed
151
class Req:
152
    """The input and output status of a request."""
153

154
155
156
157
158
    def __init__(
        self,
        rid: str,
        origin_input_text: str,
        origin_input_ids: Tuple[int],
159
        sampling_params: SamplingParams,
160
161
        lora_path: Optional[str] = None,
    ):
162
        # Input and output info
Lianmin Zheng's avatar
Lianmin Zheng committed
163
        self.rid = rid
Liangsheng Yin's avatar
Liangsheng Yin committed
164
        self.origin_input_text = origin_input_text
Liangsheng Yin's avatar
Liangsheng Yin committed
165
        self.origin_input_ids_unpadded = origin_input_ids  # Before image padding
Liangsheng Yin's avatar
Liangsheng Yin committed
166
        self.origin_input_ids = origin_input_ids
Liangsheng Yin's avatar
Liangsheng Yin committed
167
        self.output_ids = []  # Each decode stage's output ids
168
        self.fill_ids = None  # fill_ids = origin_input_ids + output_ids
169
170

        self.sampling_params = sampling_params
171
        self.lora_path = lora_path
Liangsheng Yin's avatar
Liangsheng Yin committed
172

173
174
175
        # Memory info
        self.req_pool_idx = None

176
177
178
        # Check finish
        self.tokenizer = None
        self.finished_reason = None
179
        self.stream = False
180

181
        # For incremental decoding
182
183
184
185
186
187
188
189
        # ----- | --------- read_ids -------|
        # ----- |   surr_ids  |
        # xxxxx | xxxxxxxxxxx | xxxxxxxxxxx |
        # ----- ^ ----------- ^ ----------- ^
        # ----- 1 ----------- 2 ----------- 3
        # 1: surr_offset
        # 2: read_offset
        # 3: last token
190
        self.vid = 0  # version id to sync decode status with in detokenizer_manager
Liangsheng Yin's avatar
Liangsheng Yin committed
191
192
193
        self.decoded_text = ""
        self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
        self.read_offset = None
194

195
196
197
        # The number of decoded tokens for token usage report. Note that
        # this does not include the jump forward tokens.
        self.completion_tokens_wo_jump_forward = 0
198

199
        # For vision inputs
Liangsheng Yin's avatar
Liangsheng Yin committed
200
        self.image_inputs: Optional[ImageInputs] = None
201

202
203
        # Prefix info
        self.prefix_indices = []
204
        self.extend_input_len = 0
205
        self.last_node = None
Lianmin Zheng's avatar
Lianmin Zheng committed
206
        self.is_inflight_req = 0
207

208
        # Logprobs (arguments)
209
210
211
        self.return_logprob = False
        self.logprob_start_len = 0
        self.top_logprobs_num = 0
212
213

        # Logprobs (return value)
214
        self.normalized_prompt_logprob = None
215
216
217
218
        self.input_token_logprobs = None
        self.input_top_logprobs = None
        self.output_token_logprobs = []
        self.output_top_logprobs = []
219
220

        # Logprobs (internal values)
Liangsheng Yin's avatar
Liangsheng Yin committed
221
222
223
        # The tokens is prefilled but need to be considered as decode tokens
        # and should be updated for the decode logprobs
        self.last_update_decode_tokens = 0
224
225
226
227
228
        # The relative logprob_start_len in an extend batch
        self.extend_logprob_start_len = 0

        # Embedding
        self.embedding = None
Lianmin Zheng's avatar
Lianmin Zheng committed
229

230
        # Constrained decoding
Liangsheng Yin's avatar
Liangsheng Yin committed
231
232
233
        self.regex_fsm: RegexGuide = None
        self.regex_fsm_state: int = 0
        self.jump_forward_map: JumpForwardMap = None
Liangsheng Yin's avatar
Liangsheng Yin committed
234

235
236
237
238
    # whether request reached finished condition
    def finished(self) -> bool:
        return self.finished_reason is not None

239
    def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None):
240
        self.fill_ids = self.origin_input_ids + self.output_ids
241
242
243
244
        if tree_cache is not None:
            self.prefix_indices, self.last_node = tree_cache.match_prefix(
                rid=self.rid, key=self.adjust_max_prefix_ids()
            )
245
        self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices)
246

247
    def adjust_max_prefix_ids(self):
248
249
        self.fill_ids = self.origin_input_ids + self.output_ids
        input_len = len(self.fill_ids)
250
251
252
253

        # FIXME: To work around some bugs in logprob computation, we need to ensure each
        # request has at least one token. Later, we can relax this requirement and use `input_len`.
        max_prefix_len = input_len - 1
Liangsheng Yin's avatar
Liangsheng Yin committed
254
255
256
257
258

        if self.sampling_params.max_new_tokens > 0:
            # Need at least one token to compute logits
            max_prefix_len = min(max_prefix_len, input_len - 1)

259
        if self.return_logprob:
Liangsheng Yin's avatar
Liangsheng Yin committed
260
261
262
            if self.normalized_prompt_logprob is None:
                # Need at least two tokens to compute normalized logprob
                max_prefix_len = min(max_prefix_len, input_len - 2)
263
            max_prefix_len = min(max_prefix_len, self.logprob_start_len)
264

265
        max_prefix_len = max(max_prefix_len, 0)
266
        return self.fill_ids[:max_prefix_len]
267

Liangsheng Yin's avatar
Liangsheng Yin committed
268
    # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
269
    def init_incremental_detokenize(self):
Liangsheng Yin's avatar
Liangsheng Yin committed
270
271
272
273
274
275
276
277
278
        first_iter = self.surr_offset is None or self.read_offset is None

        if first_iter:
            self.read_offset = len(self.origin_input_ids_unpadded)
            self.surr_offset = max(
                self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0
            )

        all_ids = self.origin_input_ids_unpadded + self.output_ids
279
        return all_ids[self.surr_offset :], self.read_offset - self.surr_offset
Liangsheng Yin's avatar
Liangsheng Yin committed
280

281
    def get_next_inc_detokenization(self):
282
283
        if self.tokenizer is None:
            return False, ""
284
285
        read_ids, read_offset = self.init_incremental_detokenize()
        surr_ids = read_ids[:read_offset]
Liangsheng Yin's avatar
Liangsheng Yin committed
286
287
288
289
290

        surr_text = self.tokenizer.decode(
            surr_ids,
            skip_special_tokens=self.sampling_params.skip_special_tokens,
            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
Liangsheng Yin's avatar
Liangsheng Yin committed
291
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
292
293
294
295
296
297
298
        new_text = self.tokenizer.decode(
            read_ids,
            skip_special_tokens=self.sampling_params.skip_special_tokens,
            spaces_between_special_tokens=self.sampling_params.spaces_between_special_tokens,
        )

        if len(new_text) > len(surr_text) and not new_text.endswith("�"):
299
            return True, new_text[len(surr_text) :]
Liangsheng Yin's avatar
Liangsheng Yin committed
300
301

        return False, ""
Lianmin Zheng's avatar
Lianmin Zheng committed
302

303
    def check_finished(self):
304
        if self.finished():
305
306
            return

Liangsheng Yin's avatar
Liangsheng Yin committed
307
        if len(self.output_ids) >= self.sampling_params.max_new_tokens:
308
309
310
            self.finished_reason = FINISH_LENGTH(
                length=self.sampling_params.max_new_tokens
            )
311
312
            return

313
        last_token_id = self.output_ids[-1]
314
315
316
317
318
319

        matched_eos = last_token_id in self.sampling_params.stop_token_ids

        if self.tokenizer is not None:
            matched_eos |= last_token_id == self.tokenizer.eos_token_id

320
        if matched_eos and not self.sampling_params.ignore_eos:
321
322
323
            self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id)
            return

324
325
326
327
328
329
        if len(self.sampling_params.stop_strs) > 0:
            tail_str = self.tokenizer.decode(
                self.output_ids[-(self.sampling_params.stop_str_max_len + 1) :]
            )

            for stop_str in self.sampling_params.stop_strs:
Liangsheng Yin's avatar
Liangsheng Yin committed
330
                if stop_str in tail_str or stop_str in self.decoded_text:
331
                    self.finished_reason = FINISH_MATCHED_STR(matched=stop_str)
332
333
                    return

Liangsheng Yin's avatar
Liangsheng Yin committed
334
    def jump_forward_and_retokenize(self, jump_forward_str, next_state):
Liangsheng Yin's avatar
Liangsheng Yin committed
335
336
337
338
339
340
        if self.origin_input_text is None:
            # Recovering text can only use unpadded ids
            self.origin_input_text = self.tokenizer.decode(
                self.origin_input_ids_unpadded
            )

Liangsheng Yin's avatar
Liangsheng Yin committed
341
        all_text = self.origin_input_text + self.decoded_text + jump_forward_str
Liangsheng Yin's avatar
Liangsheng Yin committed
342
        all_ids = self.tokenizer.encode(all_text)
343
        if not all_ids:
havetc's avatar
havetc committed
344
            logger.warning("Encoded all_text resulted in empty all_ids")
345
346
            return False

Liangsheng Yin's avatar
Liangsheng Yin committed
347
        prompt_tokens = len(self.origin_input_ids_unpadded)
348
        if prompt_tokens > len(all_ids):
havetc's avatar
havetc committed
349
            logger.warning("prompt_tokens is larger than encoded all_ids")
350
            return False
Liangsheng Yin's avatar
Liangsheng Yin committed
351
352
353

        if all_ids[prompt_tokens - 1] != self.origin_input_ids_unpadded[-1]:
            # TODO(lsyin): fix token fusion
354
            logger.warning(
Liangsheng Yin's avatar
Liangsheng Yin committed
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
                "Token fusion between input and output, try to avoid this by removing the space at the end of the input."
            )
            return False

        old_output_ids = self.output_ids
        self.output_ids = all_ids[prompt_tokens:]
        self.decoded_text = self.decoded_text + jump_forward_str
        self.surr_offset = prompt_tokens
        self.read_offset = len(all_ids)

        # NOTE: A trick to reduce the surrouding tokens decoding overhead
        for i in range(0, INIT_INCREMENTAL_DETOKENIZATION_OFFSET):
            surr_text_ = self.tokenizer.decode(
                all_ids[self.read_offset - i : self.read_offset]
            )
            if not surr_text_.endswith("�"):
                self.surr_offset = self.read_offset - i
                break
Liangsheng Yin's avatar
Liangsheng Yin committed
373
374
375
376
377
378

        self.regex_fsm_state = next_state

        if self.return_logprob:
            # For fast-forward part's logprobs
            k = 0
Liangsheng Yin's avatar
Liangsheng Yin committed
379
380
            for i, old_id in enumerate(old_output_ids):
                if old_id == self.output_ids[i]:
Liangsheng Yin's avatar
Liangsheng Yin committed
381
382
383
                    k = k + 1
                else:
                    break
384
385
            self.output_token_logprobs = self.output_token_logprobs[:k]
            self.output_top_logprobs = self.output_top_logprobs[:k]
Liangsheng Yin's avatar
Liangsheng Yin committed
386
            self.logprob_start_len = prompt_tokens + k
Liangsheng Yin's avatar
Liangsheng Yin committed
387
            self.last_update_decode_tokens = len(self.output_ids) - k
388

Liangsheng Yin's avatar
Liangsheng Yin committed
389
        return True
Liangsheng Yin's avatar
Liangsheng Yin committed
390

Lianmin Zheng's avatar
Lianmin Zheng committed
391
    def __repr__(self):
Liangsheng Yin's avatar
Liangsheng Yin committed
392
        return f"rid(n={self.rid}, " f"input_ids={self.origin_input_ids}, "
Lianmin Zheng's avatar
Lianmin Zheng committed
393
394


395
396
397
bid = 0


398
@dataclass
399
class ScheduleBatch:
400
401
    """Store all inforamtion of a batch."""

402
    # Request, memory pool, and cache
403
404
    reqs: List[Req]
    req_to_token_pool: ReqToTokenPool
405
    token_to_kv_pool: BaseTokenToKVPool
406
    tree_cache: BasePrefixCache
407

Liangsheng Yin's avatar
Liangsheng Yin committed
408
    forward_mode: ForwardMode = None
409
    sampling_info: SamplingBatchInfo = None
Liangsheng Yin's avatar
Liangsheng Yin committed
410

411
    # Batched arguments to model runner
412
413
414
    input_ids: torch.Tensor = None
    req_pool_indices: torch.Tensor = None
    seq_lens: torch.Tensor = None
415
    out_cache_loc: torch.Tensor = None
416

417
418
    output_ids: torch.Tensor = None

419
    # For processing logprobs
420
    return_logprob: bool = False
421
422
423
424
425
426
427
    top_logprobs_nums: Optional[List[int]] = None

    # For extend and mixed chunekd prefill
    prefix_lens: List[int] = None
    extend_lens: List[int] = None
    extend_num_tokens: int = None
    running_bs: int = None
428
    decoding_reqs: List[Req] = None
429

430
431
432
    # Stream
    has_stream: bool = False

Zhang, Liangang's avatar
Zhang, Liangang committed
433
434
435
    # device
    device: str = "cuda"

436
437
438
    # Has regex
    has_regex: bool = False

439
440
    @classmethod
    def init_new(cls, reqs, req_to_token_pool, token_to_kv_pool, tree_cache):
441
        return_logprob = any(req.return_logprob for req in reqs)
442
        has_stream = any(req.stream for req in reqs)
443
        has_regex = any(req.regex_fsm for req in reqs)
444
445
446
447
448
449

        return cls(
            reqs=reqs,
            req_to_token_pool=req_to_token_pool,
            token_to_kv_pool=token_to_kv_pool,
            tree_cache=tree_cache,
450
            return_logprob=return_logprob,
451
            has_stream=has_stream,
Zhang, Liangang's avatar
Zhang, Liangang committed
452
            device=req_to_token_pool.device,
453
            has_regex=has_regex,
Lianmin Zheng's avatar
Lianmin Zheng committed
454
455
        )

456
    def batch_size(self):
457
        return len(self.reqs)
458

Lianmin Zheng's avatar
Lianmin Zheng committed
459
460
461
    def is_empty(self):
        return len(self.reqs) == 0

462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
    def alloc_req_slots(self, num_reqs):
        req_pool_indices = self.req_to_token_pool.alloc(num_reqs)
        if req_pool_indices is None:
            raise RuntimeError(
                "Out of memory. "
                "Please set a smaller number for `--max-running-requests`."
            )
        return req_pool_indices

    def alloc_token_slots(self, num_tokens: int):
        out_cache_loc = self.token_to_kv_pool.alloc(num_tokens)

        if out_cache_loc is None:
            if self.tree_cache is not None:
                self.tree_cache.evict(num_tokens, self.token_to_kv_pool.free)
                out_cache_loc = self.token_to_kv_pool.alloc(num_tokens)

            if out_cache_loc is None:
                logger.error("Prefill out of memory. Try to lower your batch size.")
                if self.tree_cache is not None:
                    self.tree_cache.pretty_print()
                exit(1)

        return out_cache_loc

487
    def prepare_for_extend(self, vocab_size: int):
Liangsheng Yin's avatar
Liangsheng Yin committed
488
489
        self.forward_mode = ForwardMode.EXTEND

490
        bs = len(self.reqs)
Lianmin Zheng's avatar
Lianmin Zheng committed
491
        reqs = self.reqs
492
        input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
493
        extend_num_tokens = sum(len(ids) for ids in input_ids)
Lianmin Zheng's avatar
Lianmin Zheng committed
494
495
        seq_lens = []

496
        # Allocate memory
497
        req_pool_indices = self.alloc_req_slots(bs)
498
        out_cache_loc = self.alloc_token_slots(extend_num_tokens)
499

500
        pt = 0
501
        for i, req in enumerate(reqs):
502
            req.req_pool_idx = req_pool_indices[i]
503
            pre_len, seq_len = len(req.prefix_indices), len(req.fill_ids)
504
            seq_lens.append(seq_len)
505
            assert seq_len - pre_len == req.extend_input_len
Lianmin Zheng's avatar
Lianmin Zheng committed
506

507
            if pre_len > 0:
508
                self.req_to_token_pool.req_to_token[req.req_pool_idx][
509
510
                    :pre_len
                ] = req.prefix_indices
Lianmin Zheng's avatar
Lianmin Zheng committed
511

512
            self.req_to_token_pool.req_to_token[req.req_pool_idx][pre_len:seq_len] = (
513
                out_cache_loc[pt : pt + req.extend_input_len]
514
            )
515
516
517
518
519
520
521
522
523
524
525

            # Compute the relative logprob_start_len in an extend batch
            if req.logprob_start_len >= pre_len:
                extend_logprob_start_len = min(
                    req.logprob_start_len - pre_len, req.extend_input_len - 1
                )
            else:
                extend_logprob_start_len = req.extend_input_len - 1

            req.extend_logprob_start_len = extend_logprob_start_len
            pt += req.extend_input_len
Lianmin Zheng's avatar
Lianmin Zheng committed
526
527

        # Set fields
528
529
530
531
        with out_cache_loc.device:
            self.input_ids = torch.tensor(sum(input_ids, []), dtype=torch.int32)
            self.req_pool_indices = torch.tensor(req_pool_indices)
            self.seq_lens = torch.tensor(seq_lens)
532

Lianmin Zheng's avatar
Lianmin Zheng committed
533
534
        self.extend_num_tokens = extend_num_tokens
        self.out_cache_loc = out_cache_loc
535
536
537
538
539
        if self.return_logprob:
            self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
        self.prefix_lens = [len(r.prefix_indices) for r in reqs]
        self.extend_lens = [r.extend_input_len for r in reqs]
        self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
Lianmin Zheng's avatar
Lianmin Zheng committed
540

541
542
543
        self.sampling_info = SamplingBatchInfo.from_schedule_batch(
            self, vocab_size, global_server_args_dict["disable_penalizer"]
        )
544

545
    def mix_with_running(self, running_batch: "ScheduleBatch"):
546
        self.forward_mode = ForwardMode.MIXED
547
        running_bs = running_batch.batch_size()
548
549
550
551
552

        for req in running_batch.reqs:
            req.fill_ids = req.origin_input_ids + req.output_ids
            req.extend_input_len = 1

553
        input_ids = torch.cat([self.input_ids, running_batch.input_ids])
554
        out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
555
556
        extend_num_tokens = self.extend_num_tokens + running_bs

557
        self.merge_batch(running_batch)
558
559
560
        self.input_ids = input_ids
        self.out_cache_loc = out_cache_loc
        self.extend_num_tokens = extend_num_tokens
561
562

        # NOTE: prefix_indices is what has been cached, but we don't cache each decode step
563
        self.prefix_lens.extend(
564
565
566
567
568
            [
                len(r.origin_input_ids) + len(r.output_ids) - 1
                for r in running_batch.reqs
            ]
        )
569
570
        self.extend_lens.extend([1] * running_bs)
        self.extend_logprob_start_lens.extend([0] * running_bs)
571

572
    def check_decode_mem(self):
573
        bs = len(self.reqs)
Ying Sheng's avatar
Ying Sheng committed
574
        if self.token_to_kv_pool.available_size() >= bs:
575
576
            return True

Mingyi's avatar
Mingyi committed
577
        self.tree_cache.evict(bs, self.token_to_kv_pool.free)
578

579
580
581
582
583
584
585
        if self.token_to_kv_pool.available_size() >= bs:
            return True

        return False

    def retract_decode(self):
        sorted_indices = [i for i in range(len(self.reqs))]
Liangsheng Yin's avatar
Liangsheng Yin committed
586
587

        # TODO(lsyin): improve retraction policy for radix cache
588
        sorted_indices.sort(
Liangsheng Yin's avatar
Liangsheng Yin committed
589
590
591
592
            key=lambda i: (
                len(self.reqs[i].output_ids),
                -len(self.reqs[i].origin_input_ids),
            ),
593
594
595
596
            reverse=True,
        )

        retracted_reqs = []
597
        seq_lens_cpu = self.seq_lens.cpu().numpy()
598
        first_iter = True
Liangsheng Yin's avatar
Liangsheng Yin committed
599
600
601
        while (
            self.token_to_kv_pool.available_size()
            < len(sorted_indices) * global_config.retract_decode_steps
602
            or first_iter
Liangsheng Yin's avatar
Liangsheng Yin committed
603
604
605
606
607
608
609
610
        ):
            if len(sorted_indices) == 1:
                # Corner case: only one request left
                assert (
                    self.token_to_kv_pool.available_size() > 0
                ), "No space left for only one request"
                break

611
            first_iter = False
612
613
614
615
            idx = sorted_indices.pop()
            req = self.reqs[idx]
            retracted_reqs.append(req)

616
617
            if isinstance(self.tree_cache, ChunkCache):
                # ChunkCache does not have eviction
618
619
620
                token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx][
                    : seq_lens_cpu[idx]
                ]
621
                self.token_to_kv_pool.free(token_indices)
622
                self.req_to_token_pool.free(req.req_pool_idx)
623
624
625
626
                del self.tree_cache.entries[req.rid]
            else:
                # TODO: apply more fine-grained retraction
                last_uncached_pos = len(req.prefix_indices)
627
628
629
                token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx][
                    last_uncached_pos : seq_lens_cpu[idx]
                ]
630
                self.token_to_kv_pool.free(token_indices)
631
                self.req_to_token_pool.free(req.req_pool_idx)
632
633
634
635
636
637
638
639
640
641
642

                # release the last node
                self.tree_cache.dec_lock_ref(req.last_node)

                # NOTE(lsyin): we should use the newly evictable memory instantly.
                residual_size = (
                    len(sorted_indices) * global_config.retract_decode_steps
                    - self.token_to_kv_pool.available_size()
                )
                residual_size = max(0, residual_size)
                self.tree_cache.evict(residual_size, self.token_to_kv_pool.free)
Liangsheng Yin's avatar
Liangsheng Yin committed
643

644
            req.prefix_indices = []
645
            req.last_node = None
646
            req.extend_input_len = 0
Liangsheng Yin's avatar
Liangsheng Yin committed
647
648
649
650

            # For incremental logprobs
            req.last_update_decode_tokens = 0
            req.logprob_start_len = 10**9
Liangsheng Yin's avatar
Liangsheng Yin committed
651

652
653
        self.filter_batch(sorted_indices)

Liangsheng Yin's avatar
Liangsheng Yin committed
654
655
656
657
658
659
660
661
662
663
        # Reqs in batch are filtered
        total_decoded_tokens = sum(len(r.output_ids) for r in self.reqs)
        total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in self.reqs)

        new_estimate_ratio = (
            total_decoded_tokens + global_config.retract_decode_steps * len(self.reqs)
        ) / total_max_new_tokens
        new_estimate_ratio = min(1.0, new_estimate_ratio)

        return retracted_reqs, new_estimate_ratio
664

665
    def check_for_jump_forward(self, pad_input_ids_func):
Liangsheng Yin's avatar
Liangsheng Yin committed
666
        jump_forward_reqs = []
667
        keep_indices = set(i for i in range(len(self.reqs)))
Liangsheng Yin's avatar
Liangsheng Yin committed
668
669

        for i, req in enumerate(self.reqs):
Liangsheng Yin's avatar
Liangsheng Yin committed
670
            if req.jump_forward_map is not None:
Liangsheng Yin's avatar
Liangsheng Yin committed
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
                jump_forward_bytes = req.jump_forward_map.jump_forward_byte(
                    req.regex_fsm_state
                )
                if jump_forward_bytes is not None and len(jump_forward_bytes) > 1:
                    suffix_bytes = []
                    continuation_range = range(0x80, 0xC0)
                    cur_state = req.regex_fsm_state
                    while (
                        len(jump_forward_bytes)
                        and jump_forward_bytes[0][0] in continuation_range
                    ):
                        # continuation bytes
                        byte_edge = jump_forward_bytes.pop(0)
                        suffix_bytes.append(byte_edge[0])
                        cur_state = byte_edge[1]

                    suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes]
                    suffix_ids = req.tokenizer.convert_tokens_to_ids(suffix_tokens)

                    # Current ids, for cache and revert
                    cur_all_ids = tuple(req.origin_input_ids + req.output_ids)[:-1]
                    cur_output_ids = req.output_ids

                    req.output_ids.extend(suffix_ids)
695
                    decode_res, new_text = req.get_next_inc_detokenization()
Liangsheng Yin's avatar
Liangsheng Yin committed
696
697
                    if not decode_res:
                        req.output_ids = cur_output_ids
Liangsheng Yin's avatar
Liangsheng Yin committed
698
699
                        continue

sglang's avatar
sglang committed
700
701
702
703
                    (
                        jump_forward_str,
                        next_state,
                    ) = req.jump_forward_map.jump_forward_symbol(cur_state)
Liangsheng Yin's avatar
Liangsheng Yin committed
704
705
706
707
708
709
710
711
712

                    # Make the incrementally decoded text part of jump_forward_str
                    # so that the UTF-8 will not corrupt
                    jump_forward_str = new_text + jump_forward_str
                    if not req.jump_forward_and_retokenize(
                        jump_forward_str, next_state
                    ):
                        req.output_ids = cur_output_ids
                        continue
Liangsheng Yin's avatar
Liangsheng Yin committed
713

714
715
716
                    # The decode status has diverged from detokenizer_manager
                    req.vid += 1

Liangsheng Yin's avatar
Liangsheng Yin committed
717
                    # insert the old request into tree_cache
718
                    self.tree_cache.cache_finished_req(req, cur_all_ids)
Liangsheng Yin's avatar
Liangsheng Yin committed
719

Liangsheng Yin's avatar
Liangsheng Yin committed
720
                    # re-applying image padding
Liangsheng Yin's avatar
Liangsheng Yin committed
721
                    if req.image_inputs is not None:
722
                        req.origin_input_ids = pad_input_ids_func(
Liangsheng Yin's avatar
Liangsheng Yin committed
723
                            req.origin_input_ids_unpadded, req.image_inputs
Liangsheng Yin's avatar
Liangsheng Yin committed
724
725
                        )

Liangsheng Yin's avatar
Liangsheng Yin committed
726
                    jump_forward_reqs.append(req)
727
                    keep_indices.remove(i)
Liangsheng Yin's avatar
Liangsheng Yin committed
728

729
        self.filter_batch(keep_indices=list(keep_indices))
Liangsheng Yin's avatar
Liangsheng Yin committed
730

Liangsheng Yin's avatar
Liangsheng Yin committed
731
        return jump_forward_reqs
Liangsheng Yin's avatar
Liangsheng Yin committed
732

733
    def prepare_for_decode(self):
Liangsheng Yin's avatar
Liangsheng Yin committed
734
735
        self.forward_mode = ForwardMode.DECODE

736
        self.input_ids = self.output_ids
Lianmin Zheng's avatar
Lianmin Zheng committed
737
        self.seq_lens.add_(1)
738
        self.output_ids = None
Lianmin Zheng's avatar
Lianmin Zheng committed
739
740

        # Alloc mem
741
        bs = len(self.reqs)
742
        self.out_cache_loc = self.alloc_token_slots(bs)
Lianmin Zheng's avatar
Lianmin Zheng committed
743
744
745
746
747

        self.req_to_token_pool.req_to_token[
            self.req_pool_indices, self.seq_lens - 1
        ] = self.out_cache_loc

748
749
750
751
752
753
754
755
756
757
758
759
760
761
    def filter_batch(
        self,
        current_inflight_req: Optional[Req] = None,
        keep_indices: Optional[List[int]] = None,
    ):
        if keep_indices is None:
            keep_indices = [
                i
                for i in range(len(self.reqs))
                if not self.reqs[i].finished()
                and self.reqs[i] is not current_inflight_req
            ]

        if keep_indices is None or len(keep_indices) == 0:
762
763
764
765
            # Filter out all requests
            self.reqs = []
            return

766
        if len(keep_indices) == len(self.reqs):
767
768
769
            # No need to filter
            return

770
        self.reqs = [self.reqs[i] for i in keep_indices]
771
        new_indices = torch.tensor(
772
            keep_indices, dtype=torch.int32, device=self.seq_lens.device
773
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
774
        self.req_pool_indices = self.req_pool_indices[new_indices]
775
        self.seq_lens = self.seq_lens[new_indices]
776
        self.out_cache_loc = None
777
        self.output_ids = self.output_ids[new_indices]
778
        self.return_logprob = any(req.return_logprob for req in self.reqs)
779
        if self.return_logprob:
780
            self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in keep_indices]
781
782
        else:
            self.top_logprobs_nums = None
783

784
        self.has_stream = any(req.stream for req in self.reqs)
785
        self.has_regex = any(req.regex_fsm for req in self.reqs)
Lianmin Zheng's avatar
Lianmin Zheng committed
786

787
        self.sampling_info.filter_batch(keep_indices, new_indices)
Lianmin Zheng's avatar
Lianmin Zheng committed
788

789
    def merge_batch(self, other: "ScheduleBatch"):
790
791
792
        # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because
        # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it
        # needs to be called with pre-merged Batch.reqs.
793
        self.sampling_info.merge_batch(other.sampling_info)
794

Lianmin Zheng's avatar
Lianmin Zheng committed
795
796
797
798
        self.req_pool_indices = torch.concat(
            [self.req_pool_indices, other.req_pool_indices]
        )
        self.seq_lens = torch.concat([self.seq_lens, other.seq_lens])
799
        self.out_cache_loc = None
800
801
        if self.output_ids is not None:
            self.output_ids = torch.concat([self.output_ids, other.output_ids])
802
803
804
805
806
807
        if self.return_logprob and other.return_logprob:
            self.top_logprobs_nums.extend(other.top_logprobs_nums)
        elif self.return_logprob:
            self.top_logprobs_nums.extend([0] * len(other.reqs))
        elif other.return_logprob:
            self.top_logprobs_nums = [0] * len(self.reqs) + other.top_logprobs_nums
808
        self.reqs.extend(other.reqs)
809

810
        self.return_logprob = self.return_logprob or other.return_logprob
811
812
        self.has_stream = self.has_stream or other.has_stream
        self.has_regex = self.has_regex or other.has_regex
813
814
815
816
817
818
819
820
821
822
823
824
825

    def get_model_worker_batch(self):
        if self.forward_mode.is_decode():
            extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = (
                image_inputs
            ) = None
        else:
            extend_seq_lens = self.extend_lens
            extend_prefix_lens = self.prefix_lens
            extend_logprob_start_lens = self.extend_logprob_start_lens
            image_inputs = [r.image_inputs for r in self.reqs]

        lora_paths = [req.lora_path for req in self.reqs]
826
827
828
829
830
        if self.has_regex:
            self.sampling_info.regex_fsms = [req.regex_fsm for req in self.reqs]
            self.sampling_info.regex_fsm_states = [
                req.regex_fsm_state for req in self.reqs
            ]
831
832
        else:
            self.sampling_info.regex_fsms = None
833

834
835
836
        global bid
        bid += 1

837
        return ModelWorkerBatch(
838
            bid=bid,
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
            forward_mode=self.forward_mode,
            input_ids=self.input_ids,
            req_pool_indices=self.req_pool_indices,
            seq_lens=self.seq_lens,
            out_cache_loc=self.out_cache_loc,
            return_logprob=self.return_logprob,
            top_logprobs_nums=self.top_logprobs_nums,
            extend_seq_lens=extend_seq_lens,
            extend_prefix_lens=extend_prefix_lens,
            extend_logprob_start_lens=extend_logprob_start_lens,
            image_inputs=image_inputs,
            lora_paths=lora_paths,
            sampling_info=self.sampling_info,
        )

854
855
856
857
858
859
860
    def copy(self):
        return ScheduleBatch(
            reqs=self.reqs,
            req_to_token_pool=self.req_to_token_pool,
            token_to_kv_pool=self.token_to_kv_pool,
            tree_cache=self.tree_cache,
            forward_mode=self.forward_mode,
861
862
863
            output_ids=self.output_ids,
            sampling_info=self.sampling_info,
            decoding_reqs=self.decoding_reqs,
864
865
866
867
868
869
870
871
        )

    def __str__(self):
        return (
            f"ScheduleBatch(forward_mode={self.forward_mode.name}, "
            f"#req={(len(self.reqs))})"
        )

872
873
874

@dataclass
class ModelWorkerBatch:
875
876
    # The batch id
    bid: int
877
878
879
    # The forward mode
    forward_mode: ForwardMode
    # The input ids
880
    input_ids: torch.Tensor
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
    # The indices of requests in the req_to_token_pool
    req_pool_indices: torch.Tensor
    # The sequence length
    seq_lens: torch.Tensor
    # The indices of output tokens in the token_to_kv_pool
    out_cache_loc: torch.Tensor

    # For logprob
    return_logprob: bool
    top_logprobs_nums: Optional[List[int]]

    # For extend
    extend_seq_lens: Optional[List[int]]
    extend_prefix_lens: Optional[List[int]]
    extend_logprob_start_lens: Optional[List[int]]

    # For multimodal
    image_inputs: Optional[List[ImageInputs]]

    # For LoRA
    lora_paths: Optional[List[str]]

    # Sampling info
    sampling_info: SamplingBatchInfo
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922

    def copy(self):
        return ModelWorkerBatch(
            bid=self.bid,
            forward_mode=self.forward_mode,
            input_ids=self.input_ids.clone(),
            req_pool_indices=self.req_pool_indices,
            seq_lens=self.seq_lens,
            out_cache_loc=self.out_cache_loc,
            return_logprob=self.return_logprob,
            top_logprobs_nums=self.top_logprobs_nums,
            extend_seq_lens=self.extend_seq_lens,
            extend_prefix_lens=self.extend_prefix_lens,
            extend_logprob_start_lens=self.extend_logprob_start_lens,
            image_inputs=self.image_inputs,
            lora_paths=self.lora_paths,
            sampling_info=self.sampling_info.copy(),
        )