"tools/dataset_converters/kitti_data_utils.py" did not exist on "53488dad6e63236bd31aa4f6414c2fb12ecdc6d8"
scheduler.py 108 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
# 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.
# ==============================================================================
14
15
"""A scheduler that manages a tensor parallel GPU worker."""

16
import faulthandler
17
import logging
18
import os
19
import signal
20
import sys
Lianmin Zheng's avatar
Lianmin Zheng committed
21
import threading
22
import time
23
from collections import deque
Lianmin Zheng's avatar
Lianmin Zheng committed
24
from concurrent import futures
25
from dataclasses import dataclass
26
from http import HTTPStatus
27
from types import SimpleNamespace
28
from typing import Dict, List, Optional, Tuple, Union
29

30
import psutil
31
import setproctitle
32
import torch
33
import zmq
34
from torch.distributed import barrier
35

36
from sglang.global_config import global_config
Lianmin Zheng's avatar
Lianmin Zheng committed
37
from sglang.srt.configs.model_config import ModelConfig
38
39
40
41
from sglang.srt.constrained.base_grammar_backend import (
    INVALID_GRAMMAR_OBJ,
    create_grammar_backend,
)
Byron Hsu's avatar
Byron Hsu committed
42
43
44
45
46
47
48
49
50
51
52
from sglang.srt.disaggregation.decode import (
    DecodePreallocQueue,
    DecodeTransferQueue,
    SchedulerDisaggregationDecodeMixin,
)
from sglang.srt.disaggregation.prefill import (
    PrefillBootstrapQueue,
    SchedulerDisaggregationPrefillMixin,
)
from sglang.srt.disaggregation.utils import (
    DisaggregationMode,
53
    MetadataBuffers,
Byron Hsu's avatar
Byron Hsu committed
54
    ReqToMetadataIdxAllocator,
55
    TransferBackend,
56
    prepare_abort,
Byron Hsu's avatar
Byron Hsu committed
57
)
58
from sglang.srt.distributed import get_pp_group, get_world_group
fzyzcjy's avatar
fzyzcjy committed
59
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
xm:D's avatar
xm:D committed
60
61
62
63
64
from sglang.srt.hf_transformers_utils import (
    get_processor,
    get_tokenizer,
    get_tokenizer_from_processor,
)
65
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
66
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
67
from sglang.srt.layers.moe import initialize_moe_config
68
69
from sglang.srt.managers.io_struct import (
    AbortReq,
70
71
    BatchTokenizedEmbeddingReqInput,
    BatchTokenizedGenerateReqInput,
72
73
    ClearHiCacheReqInput,
    ClearHiCacheReqOutput,
74
    CloseSessionReqInput,
75
    ExpertDistributionReq,
76
    ExpertDistributionReqOutput,
77
78
    FlushCacheReqInput,
    FlushCacheReqOutput,
79
    FreezeGCReq,
80
81
    GetInternalStateReq,
    GetInternalStateReqOutput,
82
    GetWeightsByNameReqInput,
83
    HealthCheckOutput,
84
85
    InitWeightsSendGroupForRemoteInstanceReqInput,
    InitWeightsSendGroupForRemoteInstanceReqOutput,
86
    InitWeightsUpdateGroupReqInput,
87
88
    LoadLoRAAdapterReqInput,
    LoadLoRAAdapterReqOutput,
89
    MultiTokenizerRegisterReq,
90
    MultiTokenizerWrapper,
91
92
    OpenSessionReqInput,
    OpenSessionReqOutput,
93
    ProfileReq,
94
95
    ReleaseMemoryOccupationReqInput,
    ResumeMemoryOccupationReqInput,
96
97
    RpcReqInput,
    RpcReqOutput,
98
99
    SendWeightsToRemoteInstanceReqInput,
    SendWeightsToRemoteInstanceReqOutput,
100
101
    SetInternalStateReq,
    SetInternalStateReqOutput,
102
103
    SlowDownReqInput,
    SlowDownReqOutput,
104
105
    TokenizedEmbeddingReqInput,
    TokenizedGenerateReqInput,
106
107
    UnloadLoRAAdapterReqInput,
    UnloadLoRAAdapterReqOutput,
Chayenne's avatar
Chayenne committed
108
    UpdateWeightFromDiskReqInput,
109
    UpdateWeightsFromDistributedReqInput,
110
    UpdateWeightsFromTensorReqInput,
111
)
112
from sglang.srt.managers.mm_utils import init_embedding_cache
113
114
from sglang.srt.managers.schedule_batch import (
    FINISH_ABORT,
Mick's avatar
Mick committed
115
    MultimodalInputs,
116
117
    Req,
    ScheduleBatch,
118
    global_server_args_dict,
119
)
120
121
122
123
124
from sglang.srt.managers.schedule_policy import (
    AddReqResult,
    PrefillAdder,
    SchedulePolicy,
)
fzyzcjy's avatar
fzyzcjy committed
125
from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
126
127
128
129
from sglang.srt.managers.scheduler_metrics_mixin import (
    RECORD_STEP_TIME,
    SchedulerMetricsMixin,
)
130
131
132
from sglang.srt.managers.scheduler_output_processor_mixin import (
    SchedulerOutputProcessorMixin,
)
133
from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin
134
from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper
135
136
137
from sglang.srt.managers.scheduler_update_weights_mixin import (
    SchedulerUpdateWeightsMixin,
)
138
from sglang.srt.managers.session_controller import Session
139
from sglang.srt.managers.tp_worker import TpModelWorker
140
from sglang.srt.managers.tp_worker_overlap_thread import TpModelWorkerClient
141
from sglang.srt.managers.utils import DPBalanceMeta, validate_input_length
tarinkk's avatar
tarinkk committed
142
from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
143
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
144
from sglang.srt.mem_cache.lora_radix_cache import LoRARadixCache
145
from sglang.srt.mem_cache.radix_cache import RadixCache
Hanming Lu's avatar
Hanming Lu committed
146
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
Lianmin Zheng's avatar
Lianmin Zheng committed
147
from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
148
from sglang.srt.parser.reasoning_parser import ReasoningParser
149
from sglang.srt.server_args import PortArgs, ServerArgs
150
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
151
from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
152
from sglang.srt.two_batch_overlap import TboDPAttentionPreparer
153
from sglang.srt.utils import (
154
    DynamicGradMode,
155
    broadcast_pyobj,
fzyzcjy's avatar
fzyzcjy committed
156
    configure_gc_logger,
157
    configure_logger,
Lianmin Zheng's avatar
Lianmin Zheng committed
158
    disable_request_logging,
159
    freeze_gc,
160
    get_available_gpu_memory,
161
    get_bool_env_var,
162
    get_zmq_socket,
163
    is_cpu,
Lianmin Zheng's avatar
Lianmin Zheng committed
164
    kill_itself_when_parent_died,
165
    numa_bind_to_node,
166
    point_to_point_pyobj,
167
    pyspy_dump_schedulers,
168
169
    require_mlp_sync,
    require_mlp_tp_gather,
170
    set_gpu_proc_affinity,
171
172
173
    set_random_seed,
    suppress_other_loggers,
)
174
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
175
176
177

logger = logging.getLogger(__name__)

178
# Test retract decode for debugging purposes
179
TEST_RETRACT = get_bool_env_var("SGLANG_TEST_RETRACT")
180
GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
181

182
183
_is_cpu = is_cpu()

184

185
186
@dataclass
class GenerationBatchResult:
187
188
189
    logits_output: Optional[LogitsProcessorOutput]
    pp_hidden_states_proxy_tensors: Optional[torch.Tensor]
    next_token_ids: Optional[List[int]]
190
191
    extend_input_len_per_req: List[int]
    extend_logprob_start_len_per_req: List[int]
192
    bid: int
193
    can_run_cuda_graph: bool
194
195
196
197
198
199
200
201


@dataclass
class EmbeddingBatchResult:
    embeddings: torch.Tensor
    bid: int


Byron Hsu's avatar
Byron Hsu committed
202
203
class Scheduler(
    SchedulerOutputProcessorMixin,
204
205
206
    SchedulerUpdateWeightsMixin,
    SchedulerProfilerMixin,
    SchedulerMetricsMixin,
Byron Hsu's avatar
Byron Hsu committed
207
208
209
    SchedulerDisaggregationDecodeMixin,
    SchedulerDisaggregationPrefillMixin,
):
210
211
212
213
214
215
216
217
    """A scheduler that manages a tensor parallel GPU worker."""

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        gpu_id: int,
        tp_rank: int,
Cheng Wan's avatar
Cheng Wan committed
218
        moe_ep_rank: int,
219
        pp_rank: int,
220
        dp_rank: Optional[int],
221
        dp_balance_meta: Optional[DPBalanceMeta] = None,
222
223
    ):
        # Parse args
224
        self.server_args = server_args
225
        self.tp_rank = tp_rank
Cheng Wan's avatar
Cheng Wan committed
226
        self.moe_ep_rank = moe_ep_rank
227
        self.pp_rank = pp_rank
228
        self.dp_rank = dp_rank
229
        self.tp_size = server_args.tp_size
Cheng Wan's avatar
Cheng Wan committed
230
        self.moe_ep_size = server_args.ep_size
231
232
        self.pp_size = server_args.pp_size
        self.dp_size = server_args.dp_size
233
        self.schedule_policy = server_args.schedule_policy
234
        self.enable_lora = server_args.enable_lora
235
        self.max_loras_per_batch = server_args.max_loras_per_batch
236
        self.enable_overlap = not server_args.disable_overlap_schedule
237
        self.skip_tokenizer_init = server_args.skip_tokenizer_init
238
        self.enable_metrics = server_args.enable_metrics
239
240
241
        self.enable_metrics_for_all_schedulers = (
            server_args.enable_metrics_for_all_schedulers
        )
242
        self.enable_kv_cache_events = server_args.kv_events_config is not None
243
        self.stream_interval = server_args.stream_interval
244
245
246
        self.spec_algorithm = SpeculativeAlgorithm.from_string(
            server_args.speculative_algorithm
        )
247
248
        self.gpu_id = gpu_id
        self.enable_hierarchical_cache = server_args.enable_hierarchical_cache
249
        self.enable_hicache_storage = server_args.hicache_storage_backend is not None
Lianmin Zheng's avatar
Lianmin Zheng committed
250
        self.page_size = server_args.page_size
251

252
        self.attn_tp_rank, self.attn_tp_size, self.attn_dp_rank = (
253
254
255
256
257
258
259
260
            compute_dp_attention_world_info(
                server_args.enable_dp_attention,
                self.tp_rank,
                self.tp_size,
                self.dp_size,
            )
        )

261
262
263
        # Init model config
        self.model_config = ModelConfig.from_server_args(server_args)

264
265
        # Init inter-process communication
        context = zmq.Context(2)
266
        self.idle_sleeper = None
267
        if self.pp_rank == 0 and self.attn_tp_rank == 0:
268
            self.recv_from_tokenizer = get_zmq_socket(
Lianmin Zheng's avatar
Lianmin Zheng committed
269
                context, zmq.PULL, port_args.scheduler_input_ipc_name, False
270
            )
271
272
273
274
            self.recv_from_rpc = get_zmq_socket(
                context, zmq.DEALER, port_args.rpc_ipc_name, False
            )

275
            self.send_to_tokenizer = get_zmq_socket(
Lianmin Zheng's avatar
Lianmin Zheng committed
276
                context, zmq.PUSH, port_args.tokenizer_ipc_name, False
277
            )
278
            if server_args.skip_tokenizer_init:
279
                # Directly send to the TokenizerManager
280
                self.send_to_detokenizer = get_zmq_socket(
Lianmin Zheng's avatar
Lianmin Zheng committed
281
                    context, zmq.PUSH, port_args.tokenizer_ipc_name, False
282
283
                )
            else:
284
                # Send to the DetokenizerManager
285
                self.send_to_detokenizer = get_zmq_socket(
Lianmin Zheng's avatar
Lianmin Zheng committed
286
                    context, zmq.PUSH, port_args.detokenizer_ipc_name, False
287
                )
288

289
290
291
292
293
294
295
            if self.server_args.sleep_on_idle:
                self.idle_sleeper = IdleSleeper(
                    [
                        self.recv_from_tokenizer,
                        self.recv_from_rpc,
                    ]
                )
296
        else:
297
            self.recv_from_tokenizer = None
298
            self.recv_from_rpc = None
299
300
            self.send_to_tokenizer = SimpleNamespace(send_pyobj=lambda x: None)
            self.send_to_detokenizer = SimpleNamespace(send_pyobj=lambda x: None)
301

302
303
304
305
306
        if self.current_scheduler_metrics_enabled():
            self.send_metrics_from_scheduler = get_zmq_socket(
                context, zmq.PUSH, port_args.metrics_ipc_name, False
            )

307
        # Init tokenizer
308
        self.init_tokenizer()
309

310
311
312
        # Init moe config
        self.init_moe_config()

313
314
315
316
317
318
319
320
321
        # Set reasoning_parser and think_end_id if --reasoning_parser is enabled
        if self.server_args.reasoning_parser and self.tokenizer:
            reasoning_parser = ReasoningParser(
                model_type=self.server_args.reasoning_parser, stream_reasoning=False
            )
            self.tokenizer.think_end_id = self.tokenizer.encode(
                reasoning_parser.detector.think_end_token, add_special_tokens=False
            )[0]

322
323
324
325
        # Check whether overlap can be enabled
        if not self.is_generation:
            self.enable_overlap = False
            logger.info("Overlap scheduler is disabled for embedding models.")
326

327
        # Launch a tensor parallel worker
328
        if self.enable_overlap:
329
            TpWorkerClass = TpModelWorkerClient
330
331
        else:
            TpWorkerClass = TpModelWorker
332

333
        self.tp_worker = TpWorkerClass(
334
            server_args=server_args,
335
336
            gpu_id=gpu_id,
            tp_rank=tp_rank,
Cheng Wan's avatar
Cheng Wan committed
337
            moe_ep_rank=moe_ep_rank,
338
            pp_rank=pp_rank,
339
            dp_rank=dp_rank,
340
            nccl_port=port_args.nccl_port,
341
        )
342

343
        # Launch a draft worker for speculative decoding
344
345
346
347
348
349
        if self.spec_algorithm.is_eagle():
            from sglang.srt.speculative.eagle_worker import EAGLEWorker

            self.draft_worker = EAGLEWorker(
                gpu_id=gpu_id,
                tp_rank=tp_rank,
Cheng Wan's avatar
Cheng Wan committed
350
                moe_ep_rank=moe_ep_rank,
351
352
353
354
355
                server_args=server_args,
                nccl_port=port_args.nccl_port,
                target_worker=self.tp_worker,
                dp_rank=dp_rank,
            )
356
357
358
359
360
361
362
363
364
365
366
367
        elif self.spec_algorithm.is_standalone():
            from sglang.srt.speculative.standalone_worker import StandaloneWorker

            self.draft_worker = StandaloneWorker(
                gpu_id=gpu_id,
                tp_rank=tp_rank,
                moe_ep_rank=moe_ep_rank,
                server_args=server_args,
                nccl_port=port_args.nccl_port,
                target_worker=self.tp_worker,
                dp_rank=dp_rank,
            )
368
369
370
        else:
            self.draft_worker = None

371
        # Get token and memory info from the model worker
372
373
374
375
        (
            self.max_total_num_tokens,
            self.max_prefill_tokens,
            self.max_running_requests,
376
            self.max_queued_requests,
377
            self.max_req_len,
378
379
            self.max_req_input_len,
            self.random_seed,
380
            self.device,
381
382
383
384
385
            worker_global_server_args_dict,
            _,
            _,
            _,
        ) = self.tp_worker.get_worker_info()
386
387
388
389
390
391
392
393
        if global_server_args_dict["max_micro_batch_size"] is None:
            global_server_args_dict["max_micro_batch_size"] = max(
                self.max_running_requests // server_args.pp_size, 1
            )

        self.tp_group = self.tp_worker.get_tp_group()
        self.tp_cpu_group = self.tp_group.cpu_group
        self.attn_tp_group = self.tp_worker.get_attention_tp_group()
394
        self.attn_tp_cpu_group = self.tp_worker.get_attention_tp_cpu_group()
395
396
397
        self.pp_group = get_pp_group()
        self.world_group = get_world_group()

398
        self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
399
        global_server_args_dict.update(worker_global_server_args_dict)
400
        set_random_seed(self.random_seed)
401

402
        # Hybrid memory pool
Hanming Lu's avatar
Hanming Lu committed
403
404
405
406
407
408
409
        self.is_hybrid = self.tp_worker.is_hybrid
        if self.is_hybrid:
            self.sliding_window_size = self.tp_worker.sliding_window_size
            self.full_tokens_per_layer, self.swa_tokens_per_layer = (
                self.tp_worker.get_tokens_per_layer_info()
            )

410
        # Print debug info
411
        if tp_rank == 0:
412
413
414
            avail_mem = get_available_gpu_memory(
                self.device, self.gpu_id, empty_cache=False
            )
415
416
417
418
419
            logger.info(
                f"max_total_num_tokens={self.max_total_num_tokens}, "
                f"chunked_prefill_size={server_args.chunked_prefill_size}, "
                f"max_prefill_tokens={self.max_prefill_tokens}, "
                f"max_running_requests={self.max_running_requests}, "
420
                f"context_len={self.model_config.context_len}, "
421
                f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB"
422
            )
423

Lianmin Zheng's avatar
Lianmin Zheng committed
424
        # Init memory pool and cache
425
        self.init_memory_pool_and_cache()
426
427
428

        # Init running status
        self.waiting_queue: List[Req] = []
429
        # The running decoding batch for continuous batching
Lianmin Zheng's avatar
Lianmin Zheng committed
430
        self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False)
431
        # The current forward batch
Lianmin Zheng's avatar
Lianmin Zheng committed
432
        self.cur_batch: Optional[ScheduleBatch] = None
Lianmin Zheng's avatar
Lianmin Zheng committed
433
        # The last forward batch
434
        self.last_batch: Optional[ScheduleBatch] = None
Lianmin Zheng's avatar
Lianmin Zheng committed
435
436
        self.forward_ct = 0
        self.forward_ct_decode = 0
437
        self.num_generated_tokens = 0
Liangsheng Yin's avatar
Liangsheng Yin committed
438
        self.last_prefill_tokens = 0
439
440
        self.last_decode_stats_tic = time.perf_counter()
        self.last_prefill_stats_tic = time.perf_counter()
441
        self.return_health_check_ct = 0
442
443
444
445
446
        self.num_retracted_reqs: int = 0
        self.num_paused_reqs: int = 0
        self.kv_transfer_speed_gb_s: float = 0.0
        self.kv_transfer_latency_ms: float = 0.0
        self.sessions: Dict[str, Session] = {}
447
        self.current_stream = torch.get_device_module(self.device).current_stream()
448
449
        if self.device == "cpu":
            self.current_stream.synchronize = lambda: None  # No-op for CPU
450
        self.forward_sleep_time = None
451

452
453
        # Init chunked prefill
        self.chunked_prefill_size = server_args.chunked_prefill_size
454
455
        if self.chunked_prefill_size <= 0:  # -1 means disable
            self.chunked_prefill_size = None
456
        self.chunked_req = None
457
458
459
460
        self.is_mixed_chunk = (
            self.chunked_prefill_size is not None and server_args.enable_mixed_chunk
        )

Lianmin Zheng's avatar
Lianmin Zheng committed
461
        # Init the grammar backend for constrained generation
462
        self.grammar_queue: List[Req] = []
463
        if not server_args.skip_tokenizer_init:
464
            self.grammar_backend = create_grammar_backend(
465
466
467
468
                server_args,
                self.tokenizer,
                self.model_config.vocab_size,
                self.model_config.hf_eos_token_id,
469
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
470
471
        else:
            self.grammar_backend = None
472

473
        # Init schedule policy and new token estimation
474
        self.policy = SchedulePolicy(
Lianmin Zheng's avatar
Lianmin Zheng committed
475
476
477
            self.schedule_policy,
            self.tree_cache,
            self.enable_hierarchical_cache,
478
        )
479
480
481
        assert (
            server_args.schedule_conservativeness >= 0
        ), "Invalid schedule_conservativeness"
482
483
        self.init_new_token_ratio = min(
            global_config.default_init_new_token_ratio
484
485
            * server_args.schedule_conservativeness,
            1.0,
486
        )
487
488
489
490
491
492
493
494
495
496
        self.min_new_token_ratio = min(
            self.init_new_token_ratio
            * global_config.default_min_new_token_ratio_factor,
            1.0,
        )
        self.new_token_ratio_decay = (
            self.init_new_token_ratio - self.min_new_token_ratio
        ) / global_config.default_new_token_ratio_decay_steps
        self.new_token_ratio = self.init_new_token_ratio

Lianmin Zheng's avatar
Lianmin Zheng committed
497
498
499
500
        # Init watchdog thread
        self.watchdog_timeout = server_args.watchdog_timeout
        t = threading.Thread(target=self.watchdog_thread, daemon=True)
        t.start()
501
        self.parent_process = psutil.Process().parent()
502
503

        # Init memory saver, profiler and metric stats
504
505
506
        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
            enable=server_args.enable_memory_saver
        )
507
        self.offload_tags = set()
limingshu's avatar
limingshu committed
508
        self.init_profiler()
509

510
        self.recv_skipper = SchedulerRecvSkipper.maybe_create(server_args)
fzyzcjy's avatar
fzyzcjy committed
511
512
513
514
515
516
        self.input_blocker = (
            SchedulerInputBlocker(noop=self.attn_tp_rank != 0)
            if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN")
            else None
        )

517
        # Init metrics stats
518
        self.init_metrics(tp_rank, pp_rank, dp_rank)
519
        self.init_kv_events(server_args.kv_events_config)
520
        self.init_dp_balance(dp_balance_meta)
521

522
523
524
525
526
527
528
529
530
        # Init disaggregation
        self.disaggregation_mode = DisaggregationMode(
            self.server_args.disaggregation_mode
        )
        self.init_disaggregation()

        if get_bool_env_var("SGLANG_GC_LOG"):
            configure_gc_logger()

531
532
        # Init request dispatcher
        self._request_dispatcher = TypeBasedDispatcher(
533
534
535
            [
                (TokenizedGenerateReqInput, self.handle_generate_request),
                (TokenizedEmbeddingReqInput, self.handle_embedding_request),
536
537
                (BatchTokenizedGenerateReqInput, self.handle_batch_generate_request),
                (BatchTokenizedEmbeddingReqInput, self.handle_batch_embedding_request),
538
                (FlushCacheReqInput, self.flush_cache_wrapped),
539
                (ClearHiCacheReqInput, self.clear_hicache_storage_wrapped),
540
                (AbortReq, self.abort_request),
541
542
                (OpenSessionReqInput, self.open_session),
                (CloseSessionReqInput, self.close_session),
543
544
                (UpdateWeightFromDiskReqInput, self.update_weights_from_disk),
                (InitWeightsUpdateGroupReqInput, self.init_weights_update_group),
545
546
547
548
549
550
551
552
                (
                    InitWeightsSendGroupForRemoteInstanceReqInput,
                    self.init_weights_send_group_for_remote_instance,
                ),
                (
                    SendWeightsToRemoteInstanceReqInput,
                    self.send_weights_to_remote_instance,
                ),
553
554
555
556
557
558
                (
                    UpdateWeightsFromDistributedReqInput,
                    self.update_weights_from_distributed,
                ),
                (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
                (GetWeightsByNameReqInput, self.get_weights_by_name),
559
560
                (ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
                (ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
561
                (SlowDownReqInput, self.slow_down),
562
                (ProfileReq, self.profile),
563
                (FreezeGCReq, self.handle_freeze_gc),
564
                (GetInternalStateReq, self.get_internal_state),
565
                (SetInternalStateReq, self.set_internal_state),
566
                (RpcReqInput, self.handle_rpc_request),
567
                (ExpertDistributionReq, self.expert_distribution_handle),
568
569
                (LoadLoRAAdapterReqInput, self.load_lora_adapter),
                (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
570
                (MultiTokenizerRegisterReq, self.register_multi_tokenizer),
571
572
573
            ]
        )

574
575
576
    def init_tokenizer(self):
        server_args = self.server_args
        self.is_generation = self.model_config.is_generation
577

578
579
580
581
582
583
584
585
586
        if server_args.skip_tokenizer_init:
            self.tokenizer = self.processor = None
        else:
            if self.model_config.is_multimodal:
                self.processor = get_processor(
                    server_args.tokenizer_path,
                    tokenizer_mode=server_args.tokenizer_mode,
                    trust_remote_code=server_args.trust_remote_code,
                    revision=server_args.revision,
587
                    use_fast=not server_args.disable_fast_image_processor,
588
                )
xm:D's avatar
xm:D committed
589
                self.tokenizer = get_tokenizer_from_processor(self.processor)
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
            else:
                self.tokenizer = get_tokenizer(
                    server_args.tokenizer_path,
                    tokenizer_mode=server_args.tokenizer_mode,
                    trust_remote_code=server_args.trust_remote_code,
                    revision=server_args.revision,
                )

    def init_memory_pool_and_cache(self):
        server_args = self.server_args

        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
            self.tp_worker.get_memory_pool()
        )

        if (
            server_args.chunked_prefill_size is not None
            and server_args.disable_radix_cache
        ):
Hanming Lu's avatar
Hanming Lu committed
609
            if self.is_hybrid:
tarinkk's avatar
tarinkk committed
610
611
612
613
                ChunkCacheClass = SWAChunkCache
            else:
                ChunkCacheClass = ChunkCache
            self.tree_cache = ChunkCacheClass(
614
615
                req_to_token_pool=self.req_to_token_pool,
                token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
616
                page_size=self.page_size,
617
618
            )
        else:
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
            if os.environ.get("SGLANG_EXPERIMENTAL_CPP_RADIX_TREE") == "1":
                # lazy import to avoid JIT overhead
                from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp

                self.tree_cache = RadixCacheCpp(
                    disable=False,
                    use_hicache=self.enable_hierarchical_cache,
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool=self.token_to_kv_pool_allocator,
                    tp_cache_group=self.tp_cpu_group,
                    page_size=self.page_size,
                    hicache_ratio=server_args.hicache_ratio,
                    hicache_size=server_args.hicache_size,
                    hicache_write_policy=server_args.hicache_write_policy,
                    enable_kv_cache_events=self.enable_kv_cache_events,
                )
            elif self.enable_hierarchical_cache:
636
637
638
                self.tree_cache = HiRadixCache(
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
639
640
641
642
643
                    tp_cache_group=(
                        self.attn_tp_cpu_group
                        if self.server_args.enable_dp_attention
                        else self.tp_cpu_group
                    ),
644
                    page_size=self.page_size,
645
                    hicache_ratio=server_args.hicache_ratio,
Zhiqiang Xie's avatar
Zhiqiang Xie committed
646
647
                    hicache_size=server_args.hicache_size,
                    hicache_write_policy=server_args.hicache_write_policy,
648
                    hicache_io_backend=server_args.hicache_io_backend,
649
                    hicache_mem_layout=server_args.hicache_mem_layout,
650
                    enable_metrics=self.enable_metrics,
651
                    hicache_storage_backend=server_args.hicache_storage_backend,
pansicheng's avatar
pansicheng committed
652
                    hicache_storage_prefetch_policy=server_args.hicache_storage_prefetch_policy,
653
654
                    model_name=server_args.served_model_name,
                    storage_backend_extra_config=server_args.hicache_storage_backend_extra_config,
655
                )
656
657
658
                self.tp_worker.register_hicache_layer_transfer_counter(
                    self.tree_cache.cache_controller.layer_done_counter
                )
Hanming Lu's avatar
Hanming Lu committed
659
660
661
662
663
664
665
666
667
668
669
            elif self.is_hybrid:
                assert (
                    self.server_args.disaggregation_mode == "null"
                ), "Hybrid mode does not support disaggregation yet"
                self.tree_cache = SWARadixCache(
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
                    sliding_window_size=self.sliding_window_size,
                    page_size=self.page_size,
                    disable=server_args.disable_radix_cache,
                )
670
671
672
673
674
675
676
677
678
679
680
681
682
            elif self.enable_lora:
                assert (
                    not self.enable_hierarchical_cache
                ), "LoRA radix cache doesn't support hierarchical cache"
                assert (
                    self.schedule_policy == "fcfs"
                ), "LoRA radix cache only supports FCFS policy"
                self.tree_cache = LoRARadixCache(
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
                    page_size=self.page_size,
                    disable=server_args.disable_radix_cache,
                )
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
            elif server_args.enable_lmcache:
                from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
                    LMCRadixCache,
                )

                self.tree_cache = LMCRadixCache(
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
                    page_size=self.page_size,
                    disable=server_args.disable_radix_cache,
                    model_config=self.model_config,
                    tp_size=self.tp_size,
                    rank=self.tp_rank,
                    tp_group=self.tp_group,
                )
698
699
700
701
            else:
                self.tree_cache = RadixCache(
                    req_to_token_pool=self.req_to_token_pool,
                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
Lianmin Zheng's avatar
Lianmin Zheng committed
702
                    page_size=self.page_size,
703
                    disable=server_args.disable_radix_cache,
704
                    enable_kv_cache_events=self.enable_kv_cache_events,
705
706
707
708
709
710
711
712
713
714
715
716
                )

        self.decode_mem_cache_buf_multiplier = (
            1
            if self.spec_algorithm.is_none()
            else (
                server_args.speculative_num_draft_tokens
                + (
                    server_args.speculative_eagle_topk
                    * server_args.speculative_num_steps
                )
            )
717
        )
718

719
720
721
        embedding_cache_size = int(os.environ.get("SGLANG_VLM_CACHE_SIZE_MB", "100"))
        init_embedding_cache(embedding_cache_size * 1024 * 1024)

Byron Hsu's avatar
Byron Hsu committed
722
    def init_disaggregation(self):
723
724
725
726
        self.transfer_backend = TransferBackend(
            self.server_args.disaggregation_transfer_backend
        )

Byron Hsu's avatar
Byron Hsu committed
727
728
729
730
        if (
            self.disaggregation_mode == DisaggregationMode.DECODE
        ):  # *2 for the headroom.
            buffer_size = (self.req_to_token_pool.size) * 2
Byron Hsu's avatar
Byron Hsu committed
731
            self.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
Byron Hsu's avatar
Byron Hsu committed
732
733
                buffer_size
            )
734
735
            self.disagg_metadata_buffers = MetadataBuffers(
                buffer_size,
736
737
                hidden_size=self.model_config.hf_text_config.hidden_size,
                dtype=self.model_config.dtype,
738
739
                custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
            )
Byron Hsu's avatar
Byron Hsu committed
740
741
742

            # The decode requests polling kv cache
            self.disagg_decode_transfer_queue = DecodeTransferQueue(
743
                gloo_group=self.attn_tp_cpu_group,
Byron Hsu's avatar
Byron Hsu committed
744
                req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator,
745
                tp_rank=self.tp_rank,
746
                metadata_buffers=self.disagg_metadata_buffers,
Byron Hsu's avatar
Byron Hsu committed
747
748
                scheduler=self,
                tree_cache=self.tree_cache,
Byron Hsu's avatar
Byron Hsu committed
749
750
751
752
753
754
            )

            # The decode requests pending for pre-allocation
            self.disagg_decode_prealloc_queue = DecodePreallocQueue(
                req_to_token_pool=self.req_to_token_pool,
                token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
Byron Hsu's avatar
Byron Hsu committed
755
756
757
758
759
                draft_token_to_kv_pool=(
                    None
                    if self.draft_worker is None
                    else self.draft_worker.model_runner.token_to_kv_pool
                ),
Byron Hsu's avatar
Byron Hsu committed
760
                req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator,
761
                metadata_buffers=self.disagg_metadata_buffers,
Byron Hsu's avatar
Byron Hsu committed
762
763
764
                scheduler=self,
                transfer_queue=self.disagg_decode_transfer_queue,
                tree_cache=self.tree_cache,
765
                gloo_group=self.attn_tp_cpu_group,
Byron Hsu's avatar
Byron Hsu committed
766
767
                tp_rank=self.tp_rank,
                tp_size=self.tp_size,
768
769
                dp_size=self.server_args.dp_size,
                gpu_id=self.gpu_id,
Byron Hsu's avatar
Byron Hsu committed
770
                bootstrap_port=self.server_args.disaggregation_bootstrap_port,
771
772
                max_total_num_tokens=self.max_total_num_tokens,
                prefill_pp_size=self.server_args.disaggregation_prefill_pp,
773
                num_reserved_decode_tokens=self.server_args.num_reserved_decode_tokens,
774
                transfer_backend=self.transfer_backend,
Byron Hsu's avatar
Byron Hsu committed
775
            )
Liangsheng Yin's avatar
Liangsheng Yin committed
776

Byron Hsu's avatar
Byron Hsu committed
777
778
779
        elif self.disaggregation_mode == DisaggregationMode.PREFILL:
            # *2 for the headroom.
            buffer_size = self.max_running_requests * 2
Byron Hsu's avatar
Byron Hsu committed
780
            self.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
Byron Hsu's avatar
Byron Hsu committed
781
782
                buffer_size
            )
783
784
            self.disagg_metadata_buffers = MetadataBuffers(
                buffer_size,
785
786
                hidden_size=self.model_config.hf_text_config.hidden_size,
                dtype=self.model_config.dtype,
787
788
                custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
            )
Byron Hsu's avatar
Byron Hsu committed
789

Liangsheng Yin's avatar
Liangsheng Yin committed
790
            self.disagg_prefill_bootstrap_queue = PrefillBootstrapQueue(
Byron Hsu's avatar
Byron Hsu committed
791
                token_to_kv_pool=self.token_to_kv_pool_allocator.get_kvcache(),
Byron Hsu's avatar
Byron Hsu committed
792
793
794
795
796
                draft_token_to_kv_pool=(
                    None
                    if self.draft_worker is None
                    else self.draft_worker.model_runner.token_to_kv_pool
                ),
Byron Hsu's avatar
Byron Hsu committed
797
                req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator,
798
                metadata_buffers=self.disagg_metadata_buffers,
Byron Hsu's avatar
Byron Hsu committed
799
800
                tp_rank=self.tp_rank,
                tp_size=self.tp_size,
Byron Hsu's avatar
Byron Hsu committed
801
                gpu_id=self.gpu_id,
Byron Hsu's avatar
Byron Hsu committed
802
                bootstrap_port=self.server_args.disaggregation_bootstrap_port,
803
                gloo_group=self.attn_tp_cpu_group,
Byron Hsu's avatar
Byron Hsu committed
804
805
806
                max_total_num_tokens=self.max_total_num_tokens,
                decode_tp_size=self.server_args.disaggregation_decode_tp,
                decode_dp_size=self.server_args.disaggregation_decode_dp,
807
                scheduler=self,
Byron Hsu's avatar
Byron Hsu committed
808
809
810
                pp_rank=self.pp_rank,
                pp_size=self.pp_size,
                transfer_backend=self.transfer_backend,
Byron Hsu's avatar
Byron Hsu committed
811
812
            )
            # The prefill requests that are in the middle of kv sending
813
            self.disagg_prefill_inflight_queue: List[Req] = []
Byron Hsu's avatar
Byron Hsu committed
814

815
816
817
818
    def init_moe_config(self):
        if hasattr(self.model_config.hf_config, "num_experts_per_tok"):
            initialize_moe_config(self.server_args)

819
    @DynamicGradMode()
820
    def event_loop_normal(self):
821
        """A normal scheduler loop."""
822
        while True:
Lianmin Zheng's avatar
Lianmin Zheng committed
823
824
            recv_reqs = self.recv_requests()
            self.process_input_requests(recv_reqs)
825

826
            batch = self.get_next_batch_to_run()
Lianmin Zheng's avatar
Lianmin Zheng committed
827
            self.cur_batch = batch
828
829
830
831

            if batch:
                result = self.run_batch(batch)
                self.process_batch_result(batch, result)
Lianmin Zheng's avatar
Lianmin Zheng committed
832
            else:
Lianmin Zheng's avatar
Lianmin Zheng committed
833
                # When the server is idle, do self-check and re-init some states
834
                self.self_check_during_idle()
835
836

            self.last_batch = batch
837

838
    @DynamicGradMode()
Lianmin Zheng's avatar
Lianmin Zheng committed
839
    def event_loop_overlap(self):
840
        """A scheduler loop that overlaps the CPU processing and GPU computation."""
841
        self.result_queue = deque()
Lianmin Zheng's avatar
Lianmin Zheng committed
842
843
844
845
846
847
848

        while True:
            recv_reqs = self.recv_requests()
            self.process_input_requests(recv_reqs)

            batch = self.get_next_batch_to_run()
            self.cur_batch = batch
849

Lianmin Zheng's avatar
Lianmin Zheng committed
850
            if batch:
851
                batch.launch_done = threading.Event()
Lianmin Zheng's avatar
Lianmin Zheng committed
852
                result = self.run_batch(batch)
853
                self.result_queue.append((batch.copy(), result))
Lianmin Zheng's avatar
Lianmin Zheng committed
854

855
                if self.last_batch is None:
856
                    # Create a dummy first batch to start the pipeline for overlap schedule.
857
858
859
860
861
862
                    # It is now used for triggering the sampling_info_done event.
                    tmp_batch = ScheduleBatch(
                        reqs=None,
                        forward_mode=ForwardMode.DUMMY_FIRST,
                        next_batch_sampling_info=self.tp_worker.cur_sampling_info,
                    )
863
                    self.process_batch_result(tmp_batch, None, batch.launch_done)
864

Lianmin Zheng's avatar
Lianmin Zheng committed
865
            if self.last_batch:
866
                # Process the results of the last batch
867
                tmp_batch, tmp_result = self.result_queue.popleft()
868
869
870
                tmp_batch.next_batch_sampling_info = (
                    self.tp_worker.cur_sampling_info if batch else None
                )
871
872
873
874
                # NOTE: we should use current launched batch's launch_done event Instead of the last batch's
                self.process_batch_result(
                    tmp_batch, tmp_result, batch.launch_done if batch else None
                )
Lianmin Zheng's avatar
Lianmin Zheng committed
875
            elif batch is None:
Lianmin Zheng's avatar
Lianmin Zheng committed
876
                # When the server is idle, do self-check and re-init some states
877
                self.self_check_during_idle()
Lianmin Zheng's avatar
Lianmin Zheng committed
878
879
880

            self.last_batch = batch

881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
    @DynamicGradMode()
    def event_loop_pp(self):
        """A non-overlap scheduler loop for pipeline parallelism."""
        mbs = [None] * self.pp_size
        last_mbs = [None] * self.pp_size
        self.running_mbs = [
            ScheduleBatch(reqs=[], batch_is_full=False) for _ in range(self.pp_size)
        ]
        bids = [None] * self.pp_size
        pp_outputs: Optional[PPProxyTensors] = None
        while True:
            server_is_idle = True
            for mb_id in range(self.pp_size):
                self.running_batch = self.running_mbs[mb_id]
                self.last_batch = last_mbs[mb_id]

                recv_reqs = self.recv_requests()
                self.process_input_requests(recv_reqs)
                mbs[mb_id] = self.get_next_batch_to_run()
                self.running_mbs[mb_id] = self.running_batch

                self.cur_batch = mbs[mb_id]
                if self.cur_batch:
                    server_is_idle = False
                    result = self.run_batch(self.cur_batch)

907
                # (last rank) send the outputs to the next step
908
909
910
911
912
913
                if self.pp_group.is_last_rank:
                    if self.cur_batch:
                        next_token_ids, bids[mb_id] = (
                            result.next_token_ids,
                            result.bid,
                        )
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
                        if self.cur_batch.return_logprob:
                            pp_outputs = PPProxyTensors(
                                {
                                    "next_token_ids": next_token_ids,
                                    "extend_input_len_per_req": result.extend_input_len_per_req,
                                    "extend_logprob_start_len_per_req": result.extend_logprob_start_len_per_req,
                                }
                                | (
                                    {
                                        f"logits_output.{k}": v
                                        for k, v in result.logits_output.__dict__.items()
                                    }
                                    if result.logits_output is not None
                                    else {}
                                )
                            )
                        else:
                            pp_outputs = PPProxyTensors(
                                {
                                    "next_token_ids": next_token_ids,
                                }
                            )
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
                        # send the output from the last round to let the next stage worker run post processing
                        self.pp_group.send_tensor_dict(
                            pp_outputs.tensors,
                            all_gather_group=self.attn_tp_group,
                        )

                # receive outputs and post-process (filter finished reqs) the coming microbatch
                next_mb_id = (mb_id + 1) % self.pp_size
                next_pp_outputs = None
                if mbs[next_mb_id] is not None:
                    next_pp_outputs: Optional[PPProxyTensors] = PPProxyTensors(
                        self.pp_group.recv_tensor_dict(
                            all_gather_group=self.attn_tp_group
                        )
                    )
                    mbs[next_mb_id].output_ids = next_pp_outputs["next_token_ids"]
952
953
954
955
956
957
958
959
960
                    logits_output_args = {
                        k[len("logits_output.") :]: v
                        for k, v in next_pp_outputs.tensors.items()
                        if k.startswith("logits_output.")
                    }
                    if len(logits_output_args) > 0:
                        logits_output = LogitsProcessorOutput(**logits_output_args)
                    else:
                        logits_output = None
961
                    output_result = GenerationBatchResult(
962
                        logits_output=logits_output,
963
964
                        pp_hidden_states_proxy_tensors=None,
                        next_token_ids=next_pp_outputs["next_token_ids"],
965
966
967
968
969
970
                        extend_input_len_per_req=next_pp_outputs.tensors.get(
                            "extend_input_len_per_req", None
                        ),
                        extend_logprob_start_len_per_req=next_pp_outputs.tensors.get(
                            "extend_logprob_start_len_per_req", None
                        ),
971
                        bid=bids[next_mb_id],
972
                        can_run_cuda_graph=result.can_run_cuda_graph,
973
974
975
976
                    )
                    self.process_batch_result(mbs[next_mb_id], output_result)
                    last_mbs[next_mb_id] = mbs[next_mb_id]

977
                # (not last rank)
978
979
980
                if not self.pp_group.is_last_rank:
                    if self.cur_batch:
                        bids[mb_id] = result.bid
981
982
                    # carry the outputs to the next stage
                    # send the outputs from the last round to let the next stage worker run post processing
983
984
985
986
987
988
989
                    if pp_outputs:
                        self.pp_group.send_tensor_dict(
                            pp_outputs.tensors,
                            all_gather_group=self.attn_tp_group,
                        )

                    # send out reqs to the next stage
990
                    dp_offset = self.attn_dp_rank * self.attn_tp_size
991
992
993
994
                    if self.attn_tp_rank == 0:
                        point_to_point_pyobj(
                            recv_reqs,
                            self.pp_rank * self.tp_size + dp_offset,
995
                            self.world_group.device_group,
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
                            self.pp_rank * self.tp_size + dp_offset,
                            (self.pp_rank + 1) * self.tp_size + dp_offset,
                        )

                    # send out proxy tensors to the next stage
                    if self.cur_batch:
                        self.pp_group.send_tensor_dict(
                            result.pp_hidden_states_proxy_tensors,
                            all_gather_group=self.attn_tp_group,
                        )

                pp_outputs = next_pp_outputs

            # When the server is idle, self-check and re-init some states
            if server_is_idle:
1011
1012
                # When the server is idle, do self-check and re-init some states
                self.self_check_during_idle()
1013

1014
1015
    def recv_requests(self) -> List[Req]:
        """Receive results at tp_rank = 0 and broadcast it to all other TP ranks."""
1016
1017
1018
1019
1020
1021
1022
1023

        if self.recv_skipper is not None:
            last_forward_mode = (
                self.last_batch.forward_mode if self.last_batch is not None else None
            )
            if not self.recv_skipper.handle(last_forward_mode):
                return []

1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
        if self.pp_rank == 0:
            if self.attn_tp_rank == 0:
                recv_reqs = []

                while True:
                    try:
                        recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
                    except zmq.ZMQError:
                        break
                    recv_reqs.append(recv_req)

                while True:
                    try:
                        recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK)
                    except zmq.ZMQError:
                        break
                    recv_reqs.append(recv_rpc)
            else:
                recv_reqs = None
Lianmin Zheng's avatar
Lianmin Zheng committed
1043
        else:
1044
            if self.attn_tp_rank == 0:
1045
                dp_offset = self.attn_dp_rank * self.attn_tp_size
1046
1047
1048
                recv_reqs = point_to_point_pyobj(
                    [],
                    self.pp_rank * self.tp_size + dp_offset,
1049
                    self.world_group.device_group,
1050
1051
1052
1053
1054
                    (self.pp_rank - 1) * self.tp_size + dp_offset,
                    self.pp_rank * self.tp_size + dp_offset,
                )
            else:
                recv_reqs = None
1055

fzyzcjy's avatar
fzyzcjy committed
1056
1057
1058
        if self.input_blocker is not None:
            recv_reqs = self.input_blocker.handle(recv_reqs)

1059
1060
1061
1062
1063
1064
        if self.server_args.enable_dp_attention:
            if self.attn_tp_rank == 0:
                work_reqs = [
                    req
                    for req in recv_reqs
                    if isinstance(
1065
1066
1067
1068
1069
1070
1071
                        req,
                        (
                            TokenizedGenerateReqInput,
                            TokenizedEmbeddingReqInput,
                            BatchTokenizedGenerateReqInput,
                            BatchTokenizedEmbeddingReqInput,
                        ),
1072
1073
1074
1075
1076
1077
                    )
                ]
                control_reqs = [
                    req
                    for req in recv_reqs
                    if not isinstance(
1078
1079
1080
1081
1082
1083
1084
                        req,
                        (
                            TokenizedGenerateReqInput,
                            TokenizedEmbeddingReqInput,
                            BatchTokenizedGenerateReqInput,
                            BatchTokenizedEmbeddingReqInput,
                        ),
1085
1086
1087
1088
1089
1090
1091
1092
1093
                    )
                ]
            else:
                work_reqs = None
                control_reqs = None

            if self.attn_tp_size != 1:
                work_reqs = broadcast_pyobj(
                    work_reqs,
1094
                    self.attn_tp_group.rank,
1095
                    self.attn_tp_cpu_group,
1096
                    src=self.attn_tp_group.ranks[0],
1097
1098
1099
                )
            if self.tp_size != 1:
                control_reqs = broadcast_pyobj(
1100
1101
1102
1103
                    control_reqs,
                    self.tp_group.rank,
                    self.tp_cpu_group,
                    src=self.tp_group.ranks[0],
1104
1105
1106
                )
            recv_reqs = work_reqs + control_reqs
        elif self.tp_size != 1:
1107
1108
1109
1110
1111
1112
            recv_reqs = broadcast_pyobj(
                recv_reqs,
                self.tp_group.rank,
                self.tp_cpu_group,
                src=self.tp_group.ranks[0],
            )
1113
1114
        return recv_reqs

Lianmin Zheng's avatar
Lianmin Zheng committed
1115
    def process_input_requests(self, recv_reqs: List):
1116
        for recv_req in recv_reqs:
1117
1118
            # If it is a health check generation request and there are running requests, ignore it.
            if is_health_check_generate_req(recv_req) and (
1119
1120
1121
                self.chunked_req is not None
                or not self.running_batch.is_empty()
                or len(self.offload_tags) > 0
1122
1123
1124
1125
            ):
                self.return_health_check_ct += 1
                continue

1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
            # If it is a work request, accept or reject the request based on the request queue size.
            if is_work_request(recv_req):
                if len(self.waiting_queue) + 1 > self.max_queued_requests:
                    abort_req = AbortReq(
                        recv_req.rid,
                        finished_reason={
                            "type": "abort",
                            "status_code": HTTPStatus.SERVICE_UNAVAILABLE,
                            "message": "The request queue is full.",
                        },
                    )
                    self.send_to_tokenizer.send_pyobj(abort_req)
                    continue
1139

1140
1141
            # If it is a MultiTokenizerWrapper, unwrap it and handle the inner request.
            if isinstance(recv_req, MultiTokenizerWrapper):
1142
1143
1144
1145
                worker_id = recv_req.worker_id
                recv_req = recv_req.obj
                output = self._request_dispatcher(recv_req)
                if output is not None:
1146
                    output = MultiTokenizerWrapper(worker_id, output)
1147
1148
1149
                    self.send_to_tokenizer.send_pyobj(output)
                continue

1150
            output = self._request_dispatcher(recv_req)
1151
            if output is not None:
1152
1153
1154
1155
1156
                if isinstance(output, RpcReqOutput):
                    if self.recv_from_rpc is not None:
                        self.recv_from_rpc.send_pyobj(output)
                else:
                    self.send_to_tokenizer.send_pyobj(output)
1157
1158
1159
1160
1161

    def handle_generate_request(
        self,
        recv_req: TokenizedGenerateReqInput,
    ):
1162
        self.maybe_update_dp_balance_data(recv_req)
1163

1164
        # Create a new request
1165
1166
1167
1168
1169
        if (
            recv_req.session_params is None
            or recv_req.session_params.id is None
            or recv_req.session_params.id not in self.sessions
        ):
Rin Intachuen's avatar
Rin Intachuen committed
1170
1171
1172
1173
1174
1175
            if recv_req.input_embeds is not None:
                # Generate fake input_ids based on the length of input_embeds
                seq_length = len(recv_req.input_embeds)
                fake_input_ids = [1] * seq_length
                recv_req.input_ids = fake_input_ids

1176
1177
1178
1179
            if recv_req.bootstrap_port is None:
                # Use default bootstrap port
                recv_req.bootstrap_port = self.server_args.disaggregation_bootstrap_port

1180
1181
1182
1183
1184
            req = Req(
                recv_req.rid,
                recv_req.input_text,
                recv_req.input_ids,
                recv_req.sampling_params,
Lianmin Zheng's avatar
Lianmin Zheng committed
1185
1186
                return_logprob=recv_req.return_logprob,
                top_logprobs_num=recv_req.top_logprobs_num,
1187
                token_ids_logprob=recv_req.token_ids_logprob,
Lianmin Zheng's avatar
Lianmin Zheng committed
1188
                stream=recv_req.stream,
1189
                lora_id=recv_req.lora_id,
Rin Intachuen's avatar
Rin Intachuen committed
1190
                input_embeds=recv_req.input_embeds,
Lianmin Zheng's avatar
Lianmin Zheng committed
1191
                custom_logit_processor=recv_req.custom_logit_processor,
1192
                return_hidden_states=recv_req.return_hidden_states,
1193
                eos_token_ids=self.model_config.hf_eos_token_id,
1194
                bootstrap_host=recv_req.bootstrap_host,
1195
                bootstrap_port=recv_req.bootstrap_port,
1196
                bootstrap_room=recv_req.bootstrap_room,
1197
                data_parallel_rank=recv_req.data_parallel_rank,
1198
                vocab_size=self.model_config.vocab_size,
1199
1200
            )
            req.tokenizer = self.tokenizer
Lianmin Zheng's avatar
Lianmin Zheng committed
1201

1202
1203
1204
            if self.disaggregation_mode != DisaggregationMode.NULL:
                # Invalid request for disaggregated mode
                if recv_req.bootstrap_room is None:
1205
                    error_msg = (
1206
1207
1208
                        f"Invalid request: Disaggregated request received without "
                        f"boostrap room id. {req.rid=}"
                    )
1209
                    logger.error(error_msg)
1210
                    prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST)
1211
1212
1213
                    self.stream_output([req], req.return_logprob)
                    return

1214
1215
1216
1217
            if (
                recv_req.session_params is not None
                and recv_req.session_params.id is not None
            ):
1218
                req.set_finish_with_abort(
1219
                    f"Invalid request: session id {recv_req.session_params.id} does not exist"
1220
                )
1221
                self._add_request_to_queue(req)
1222
1223
                return
        else:
1224
1225
            # Create a new request from a previous session
            session = self.sessions[recv_req.session_params.id]
1226
            req = session.create_req(recv_req, self.tokenizer)
1227
            if isinstance(req.finished_reason, FINISH_ABORT):
1228
                self._add_request_to_queue(req)
1229
                return
1230

1231
        # Handle multimodal inputs
Mick's avatar
Mick committed
1232
1233
        if recv_req.mm_inputs is not None:
            image_inputs = MultimodalInputs.from_dict(recv_req.mm_inputs)
1234
            # Expand a single image token into multiple dummy tokens for receiving image embeddings
1235
            req.origin_input_ids = self.pad_input_ids_func(
1236
                req.origin_input_ids, image_inputs
1237
            )
1238
            req.extend_image_inputs(image_inputs)
1239

1240
            if len(req.origin_input_ids) >= self.max_req_input_len:
1241
1242
1243
1244
1245
                req.set_finish_with_abort(
                    error_msg=(
                        "Multimodal prompt is too long after expanding multimodal tokens. "
                        f"After expanding {len(req.origin_input_ids_unpadded)=} => {len(req.origin_input_ids)} >= {self.max_req_input_len}."
                    )
1246
                )
1247
                self._add_request_to_queue(req)
1248
1249
                return

1250
        # Validate prompt length
1251
1252
1253
1254
1255
1256
        error_msg = validate_input_length(
            req,
            self.max_req_input_len,
            self.server_args.allow_auto_truncate,
        )
        if error_msg:
1257
            req.set_finish_with_abort(error_msg)
1258
            self._add_request_to_queue(req)
1259
            return
1260

1261
        # Copy more attributes
1262
        if recv_req.logprob_start_len == -1 or not recv_req.return_logprob:
1263
            # By default, only return the logprobs for output tokens
1264
1265
1266
1267
1268
1269
1270
            # For prefill-only requests with logprob_start_len == -1, set logprob_start_len beyond input sequence
            # to skip input logprob computation entirely
            if req.is_prefill_only:
                req.logprob_start_len = len(req.origin_input_ids)
            else:
                # TODO: For text generation, evaluate setting logprob_start_len to len(req.origin_input_ids) as well
                req.logprob_start_len = len(req.origin_input_ids) - 1
1271
1272
1273
        else:
            req.logprob_start_len = recv_req.logprob_start_len

1274
1275
1276
        if not req.is_prefill_only and req.logprob_start_len >= len(
            req.origin_input_ids
        ):
1277
            error_msg = f"{req.logprob_start_len=} is higher than the number of input tokens {len(req.origin_input_ids)=}. Please use a smaller logprob_start_len."
1278
            req.logprob_start_len = len(req.origin_input_ids) - 1
1279
            req.set_finish_with_abort(error_msg)
1280
1281
1282
            self._add_request_to_queue(req)
            return

1283
1284
1285
1286
1287
1288
        req.sampling_params.max_new_tokens = min(
            (
                req.sampling_params.max_new_tokens
                if req.sampling_params.max_new_tokens is not None
                else 1 << 30
            ),
1289
            self.max_req_len - len(req.origin_input_ids) - 1,
1290
1291
        )

1292
1293
1294
1295
1296
        # Init grammar cache for this request
        add_to_grammar_queue = False
        if (
            req.sampling_params.json_schema is not None
            or req.sampling_params.regex is not None
1297
            or req.sampling_params.ebnf is not None
1298
            or req.sampling_params.structural_tag is not None
1299
1300
1301
1302
1303
1304
        ):
            assert self.grammar_backend is not None
            if req.sampling_params.json_schema is not None:
                key = ("json", req.sampling_params.json_schema)
            elif req.sampling_params.regex is not None:
                key = ("regex", req.sampling_params.regex)
1305
1306
            elif req.sampling_params.ebnf is not None:
                key = ("ebnf", req.sampling_params.ebnf)
1307
1308
            elif req.sampling_params.structural_tag:
                key = ("structural_tag", req.sampling_params.structural_tag)
1309

1310
1311
1312
1313
1314
            value, cache_hit = self.grammar_backend.get_cached_or_future_value(key)
            req.grammar = value

            if not cache_hit:
                req.grammar_key = key
1315
                add_to_grammar_queue = True
1316
1317
1318
1319
            else:
                if value is INVALID_GRAMMAR_OBJ:  # We hit a cached invalid grammar.
                    error_msg = f"Invalid grammar request with cache hit: {key=}"
                    req.set_finish_with_abort(error_msg)
1320
1321

        if add_to_grammar_queue:
1322
            req.queue_time_start = time.perf_counter()
1323
1324
            self.grammar_queue.append(req)
        else:
1325
1326
            self._add_request_to_queue(req)

1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
    def handle_batch_generate_request(
        self,
        recv_req: BatchTokenizedGenerateReqInput,
    ):
        """Handle optimized batch generate request."""
        logger.debug(f"Processing batch generate request with {len(recv_req)} requests")

        # Process each request in the batch
        for tokenized_req in recv_req:
            self.handle_generate_request(tokenized_req)

1338
    def _add_request_to_queue(self, req: Req):
1339
        req.queue_time_start = time.perf_counter()
Byron Hsu's avatar
Byron Hsu committed
1340
        if self.disaggregation_mode == DisaggregationMode.PREFILL:
1341
            self._prefetch_kvcache(req)
Byron Hsu's avatar
Byron Hsu committed
1342
1343
1344
            self.disagg_prefill_bootstrap_queue.add(
                req, self.model_config.num_key_value_heads
            )
Byron Hsu's avatar
Byron Hsu committed
1345
1346
1347
        elif self.disaggregation_mode == DisaggregationMode.DECODE:
            self.disagg_decode_prealloc_queue.add(req)
        else:
1348
            self._prefetch_kvcache(req)
Byron Hsu's avatar
Byron Hsu committed
1349
1350
            self.waiting_queue.append(req)

1351
1352
1353
    def _prefetch_kvcache(self, req: Req):
        if self.enable_hicache_storage:
            req.init_next_round_input(self.tree_cache)
1354
1355
1356
1357
1358
            if req.last_node.backuped:
                # only to initiate the prefetch if the last node is backuped
                # otherwise, the allocated GPU memory must be locked for integrity
                last_hash = req.last_host_node.get_last_hash_value()
                matched_len = len(req.prefix_indices) + req.host_hit_length
1359
1360
1361
1362
1363
                new_input_tokens = req.fill_ids[matched_len:]
                self.tree_cache.prefetch_from_storage(
                    req.rid, req.last_host_node, new_input_tokens, last_hash
                )

1364
    def _extend_requests_to_queue(self, reqs: List[Req], is_retracted: bool = False):
1365
        if self.disaggregation_mode == DisaggregationMode.PREFILL:
Byron Hsu's avatar
Byron Hsu committed
1366
1367
1368
            self.disagg_prefill_bootstrap_queue.extend(
                reqs, self.model_config.num_key_value_heads
            )
1369
1370
        elif self.disaggregation_mode == DisaggregationMode.DECODE:
            # If this is a decode server, we put the request to the decode pending prealloc queue
1371
            self.disagg_decode_prealloc_queue.extend(reqs, is_retracted)
Byron Hsu's avatar
Byron Hsu committed
1372
1373
        else:
            self.waiting_queue.extend(reqs)
1374
1375
1376

    def handle_embedding_request(
        self,
1377
        recv_req: TokenizedEmbeddingReqInput,
1378
1379
1380
1381
1382
1383
    ):
        req = Req(
            recv_req.rid,
            recv_req.input_text,
            recv_req.input_ids,
            recv_req.sampling_params,
woodx's avatar
woodx committed
1384
            token_type_ids=recv_req.token_type_ids,
1385
1386
1387
        )
        req.tokenizer = self.tokenizer

1388
1389
        # Handle multimodal inputs
        if recv_req.image_inputs is not None:
Mick's avatar
Mick committed
1390
            image_inputs = MultimodalInputs.from_dict(recv_req.image_inputs)
1391
1392
1393
1394
1395
1396
1397
            # Expand a single image token into multiple dummy tokens for receiving image embeddings
            req.origin_input_ids = self.pad_input_ids_func(
                req.origin_input_ids, image_inputs
            )
            req.extend_image_inputs(image_inputs)

            if len(req.origin_input_ids) >= self.max_req_input_len:
1398
1399
1400
1401
1402
                req.set_finish_with_abort(
                    error_msg=(
                        "Multimodal prompt is too long after expanding multimodal tokens. "
                        f"After expanding {len(req.origin_input_ids_unpadded)=} => {len(req.origin_input_ids)} >= {self.max_req_input_len}."
                    )
1403
                )
1404
                self._add_request_to_queue(req)
1405
1406
                return

1407
        # Validate prompts length
1408
        error_msg = validate_input_length(
1409
1410
1411
1412
            req,
            self.max_req_input_len,
            self.server_args.allow_auto_truncate,
        )
1413
        if error_msg:
1414
            self._add_request_to_queue(req)
1415
            return
1416

1417
1418
        # Copy more attributes
        req.logprob_start_len = len(req.origin_input_ids) - 1
1419
        self._add_request_to_queue(req)
1420

1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
    def handle_batch_embedding_request(
        self,
        recv_req: BatchTokenizedEmbeddingReqInput,
    ):
        """Handle optimized batch embedding request."""
        logger.debug(
            f"Processing batch embedding request with {len(recv_req)} requests"
        )

        # Process each request in the batch
        for tokenized_req in recv_req:
            self.handle_embedding_request(tokenized_req)

1434
1435
1436
1437
1438
    def self_check_during_idle(self):
        self.check_memory()
        self.check_tree_cache()
        self.new_token_ratio = self.init_new_token_ratio
        self.maybe_sleep_on_idle()
1439

Lianmin Zheng's avatar
Lianmin Zheng committed
1440
    def check_memory(self):
Hanming Lu's avatar
Hanming Lu committed
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
        if self.is_hybrid:
            (
                full_num_used,
                swa_num_used,
                _,
                _,
                full_available_size,
                full_evictable_size,
                swa_available_size,
                swa_evictable_size,
            ) = self._get_swa_token_info()
            memory_leak = full_num_used != 0 or swa_num_used != 0
            token_msg = (
                f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {self.tree_cache.full_protected_size()=}\n"
                f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}\n"
            )
tarinkk's avatar
tarinkk committed
1457
        else:
Hanming Lu's avatar
Hanming Lu committed
1458
1459
1460
            _, _, available_size, evictable_size = self._get_token_info()
            protected_size = self.tree_cache.protected_size()
            memory_leak = (available_size + evictable_size) != (
1461
1462
1463
                # self.max_total_num_tokens
                # if not self.enable_hierarchical_cache
                # else self.max_total_num_tokens - protected_size
Hanming Lu's avatar
Hanming Lu committed
1464
                self.max_total_num_tokens
1465
                - protected_size
Lianmin Zheng's avatar
Lianmin Zheng committed
1466
            )
Hanming Lu's avatar
Hanming Lu committed
1467
1468
1469
1470
            token_msg = f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, {protected_size=}\n"

        if memory_leak:
            msg = "token_to_kv_pool_allocator memory leak detected! " f"{token_msg}"
Lianmin Zheng's avatar
Lianmin Zheng committed
1471
            raise ValueError(msg)
Lianmin Zheng's avatar
Lianmin Zheng committed
1472

1473
1474
1475
1476
1477
1478
1479
1480
        if self.disaggregation_mode == DisaggregationMode.DECODE:
            req_total_size = (
                self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
            )
        else:
            req_total_size = self.req_to_token_pool.size

        if len(self.req_to_token_pool.free_slots) != req_total_size:
1481
            msg = (
1482
                "req_to_token_pool memory leak detected!"
1483
1484
                f"available_size={len(self.req_to_token_pool.free_slots)}, "
                f"total_size={self.req_to_token_pool.size}\n"
Lianmin Zheng's avatar
Lianmin Zheng committed
1485
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
1486
            raise ValueError(msg)
Lianmin Zheng's avatar
Lianmin Zheng committed
1487

1488
1489
        if (
            self.enable_metrics
1490
            and self.current_scheduler_metrics_enabled()
1491
            and time.perf_counter() > self.metrics_collector.last_log_time + 30
1492
1493
        ):
            # During idle time, also collect metrics every 30 seconds.
Hanming Lu's avatar
Hanming Lu committed
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
            if self.is_hybrid:
                (
                    full_num_used,
                    swa_num_used,
                    full_token_usage,
                    swa_token_usage,
                    _,
                    _,
                    _,
                    _,
                ) = self._get_swa_token_info()
                num_used = max(full_num_used, swa_num_used)
                token_usage = max(full_token_usage, swa_token_usage)
            else:
                num_used, token_usage, _, _ = self._get_token_info()
Lianmin Zheng's avatar
Lianmin Zheng committed
1509
            num_running_reqs = len(self.running_batch.reqs)
1510
1511
            self.stats.num_running_reqs = num_running_reqs
            self.stats.num_used_tokens = num_used
Hanming Lu's avatar
Hanming Lu committed
1512
            self.stats.token_usage = round(token_usage, 2)
1513
1514
            self.stats.gen_throughput = 0
            self.stats.num_queue_reqs = len(self.waiting_queue)
1515
            self.stats.num_grammar_queue_reqs = len(self.grammar_queue)
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
            if self.disaggregation_mode == DisaggregationMode.PREFILL:
                self.stats.num_prefill_prealloc_queue_reqs = len(
                    self.disagg_prefill_bootstrap_queue.queue
                )
                self.stats.num_prefill_inflight_queue_reqs = len(
                    self.disagg_prefill_inflight_queue
                )
            if self.disaggregation_mode == DisaggregationMode.DECODE:
                self.stats.num_decode_prealloc_queue_reqs = len(
                    self.disagg_decode_prealloc_queue.queue
                )
                self.stats.num_decode_transfer_queue_reqs = len(
                    self.disagg_decode_transfer_queue.queue
                )
1530
            self.metrics_collector.log_stats(self.stats)
1531
        self._publish_kv_events()
1532

Hanming Lu's avatar
Hanming Lu committed
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
    def check_tree_cache(self):
        if self.is_hybrid and isinstance(self.tree_cache, SWARadixCache):
            self.tree_cache.sanity_check()

    def _get_token_info(self):
        available_size = self.token_to_kv_pool_allocator.available_size()
        evictable_size = self.tree_cache.evictable_size()
        num_used = self.max_total_num_tokens - (available_size + evictable_size)
        token_usage = num_used / self.max_total_num_tokens
        return num_used, token_usage, available_size, evictable_size

    def _get_swa_token_info(self):
        full_available_size = self.token_to_kv_pool_allocator.full_available_size()
        full_evictable_size = self.tree_cache.full_evictable_size()
        swa_available_size = self.token_to_kv_pool_allocator.swa_available_size()
        swa_evictable_size = self.tree_cache.swa_evictable_size()
        full_num_used = self.full_tokens_per_layer - (
            full_available_size + full_evictable_size
        )
        swa_num_used = self.swa_tokens_per_layer - (
            swa_available_size + swa_evictable_size
        )
        full_token_usage = full_num_used / self.full_tokens_per_layer
        swa_token_usage = swa_num_used / self.swa_tokens_per_layer
        return (
            full_num_used,
            swa_num_used,
            full_token_usage,
            swa_token_usage,
            full_available_size,
            full_evictable_size,
            swa_available_size,
            swa_evictable_size,
        )

1568
    def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
1569
        # Merge the prefill batch into the running batch
1570
1571
1572
1573
1574
        chunked_req_to_exclude = set()
        if self.chunked_req:
            # Move the chunked request out of the batch so that we can merge
            # only finished requests to running_batch.
            chunked_req_to_exclude.add(self.chunked_req)
1575
            self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
1576
            # chunked request keeps its rid but will get a new req_pool_idx
Yi Zhang's avatar
Yi Zhang committed
1577
1578
1579
1580
1581
1582
            if self.tp_worker.worker.model_runner.is_hybrid_gdn:
                self.req_to_token_pool.free(
                    self.chunked_req.req_pool_idx, free_mamba_cache=False
                )
            else:
                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
Lianmin Zheng's avatar
Lianmin Zheng committed
1583
        if self.last_batch and self.last_batch.forward_mode.is_extend():
1584
1585
1586
1587
            if self.last_batch.chunked_req is not None:
                # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req.
                # We need to discard it.
                chunked_req_to_exclude.add(self.last_batch.chunked_req)
Lianmin Zheng's avatar
Lianmin Zheng committed
1588

1589
            # Filter batch
1590
            last_bs = self.last_batch.batch_size()
1591
1592
1593
            self.last_batch.filter_batch(
                chunked_req_to_exclude=list(chunked_req_to_exclude)
            )
1594
            if self.last_batch.batch_size() < last_bs:
Lianmin Zheng's avatar
Lianmin Zheng committed
1595
                self.running_batch.batch_is_full = False
1596

1597
1598
1599
            # Merge the new batch into the running batch.
            # For prefill-only batch, we can avoid going through decoding step.
            if not self.last_batch.is_empty() and not self.last_batch.is_prefill_only:
Lianmin Zheng's avatar
Lianmin Zheng committed
1600
                if self.running_batch.is_empty():
1601
1602
                    self.running_batch = self.last_batch
                else:
Lianmin Zheng's avatar
Lianmin Zheng committed
1603
                    # Merge running_batch with prefill batch
1604
                    self.running_batch.merge_batch(self.last_batch)
1605

1606
        new_batch = self.get_new_batch_prefill()
1607

1608
1609
1610
1611
1612
        need_dp_attn_preparation = require_mlp_sync(self.server_args)

        if need_dp_attn_preparation and not self.spec_algorithm.is_none():
            # In speculative decoding, prefill batches and decode batches cannot be processed in the same DP attention group.
            # We prepare idle batches in advance to skip preparing decode batches when there are prefill batches in the group.
1613
            new_batch = self.prepare_mlp_sync_batch(new_batch)
1614
1615
1616
            need_dp_attn_preparation = new_batch is None

        if new_batch is not None:
1617
1618
1619
1620
            # Run prefill first if possible
            ret = new_batch
        else:
            # Run decode
Lianmin Zheng's avatar
Lianmin Zheng committed
1621
            if not self.running_batch.is_empty():
1622
                self.running_batch = self.update_running_batch(self.running_batch)
Lianmin Zheng's avatar
Lianmin Zheng committed
1623
1624
1625
                ret = self.running_batch if not self.running_batch.is_empty() else None
            else:
                ret = None
1626

1627
1628
        # Handle DP attention
        if need_dp_attn_preparation:
1629
            self.maybe_handle_dp_balance_data()
1630
            ret = self.prepare_mlp_sync_batch(ret)
1631
1632

        return ret
1633

1634
1635
1636
1637
1638
1639
    def get_num_allocatable_reqs(self, running_bs):
        res = global_server_args_dict["max_micro_batch_size"] - running_bs
        if self.pp_size > 1:
            res = min(res, self.req_to_token_pool.available_size())
        return res

Lianmin Zheng's avatar
Lianmin Zheng committed
1640
    def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
Lianmin Zheng's avatar
Lianmin Zheng committed
1641
        # Check if the grammar is ready in the grammar queue
1642
        if self.grammar_queue:
1643
            self.move_ready_grammar_requests()
1644

Lianmin Zheng's avatar
Lianmin Zheng committed
1645
1646
        # Handle the cases where prefill is not allowed
        if (
Lianmin Zheng's avatar
Lianmin Zheng committed
1647
            self.running_batch.batch_is_full or len(self.waiting_queue) == 0
1648
        ) and self.chunked_req is None:
Lianmin Zheng's avatar
Lianmin Zheng committed
1649
1650
            return None

Lianmin Zheng's avatar
Lianmin Zheng committed
1651
        running_bs = len(self.running_batch.reqs)
1652
        # Ignore the check if self.chunked_req is not None.
1653
1654
1655
1656
1657
        # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0,
        # as the space for the chunked request has just been released.
        # In PP case, a chunked req can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict.
        # Instead, we should always allow chunked request to be added, otherwise, there will be a memory leak.
        if self.get_num_allocatable_reqs(running_bs) <= 0 and not self.chunked_req:
Lianmin Zheng's avatar
Lianmin Zheng committed
1658
            self.running_batch.batch_is_full = True
1659
1660
            return None

1661
        if self.enable_hierarchical_cache:
1662
            self.tree_cache.check_hicache_events()
1663

1664
        # Get priority queue
1665
        self.policy.calc_priority(self.waiting_queue)
1666

Lianmin Zheng's avatar
Lianmin Zheng committed
1667
        # Prefill policy
1668
        adder = PrefillAdder(
1669
            self.page_size,
1670
            self.tree_cache,
1671
            self.token_to_kv_pool_allocator,
1672
1673
1674
1675
            self.running_batch,
            self.new_token_ratio,
            self.max_prefill_tokens,
            self.chunked_prefill_size,
1676
            running_bs if self.is_mixed_chunk else 0,
1677
1678
        )

Lianmin Zheng's avatar
Lianmin Zheng committed
1679
        if self.chunked_req is not None:
1680
1681
            self.chunked_req.init_next_round_input()
            self.chunked_req = adder.add_chunked_req(self.chunked_req)
1682

1683
        if self.enable_lora:
1684
            lora_set = set([req.lora_id for req in self.running_batch.reqs])
Lianmin Zheng's avatar
Lianmin Zheng committed
1685

1686
        # Get requests from the waiting queue to a new prefill batch
1687
        for req in self.waiting_queue:
1688
1689
1690
1691
1692

            if self.enable_lora and not self.tp_worker.can_run_lora_batch(
                lora_set
                | set([req.lora_id for req in adder.can_run_list])
                | set([req.lora_id])
1693
            ):
Lianmin Zheng's avatar
Lianmin Zheng committed
1694
                self.running_batch.batch_is_full = True
1695
1696
                break

1697
            if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
Lianmin Zheng's avatar
Lianmin Zheng committed
1698
                self.running_batch.batch_is_full = True
1699
                break
1700

Byron Hsu's avatar
Byron Hsu committed
1701
1702
1703
1704
1705
1706
1707
            if self.disaggregation_mode == DisaggregationMode.PREFILL:
                # In prefill mode, prealloc queue and transfer queue can also take memory,
                # so we need to check if the available size for the actual available size.
                if len(adder.can_run_list) >= self.req_to_token_pool.available_size():
                    self.running_batch.batch_is_full = True
                    break

1708
            if self.enable_hicache_storage:
pansicheng's avatar
pansicheng committed
1709
1710
1711
1712
                prefetch_done = self.tree_cache.check_prefetch_progress(req.rid)
                if not prefetch_done:
                    # skip staging requests that are ongoing prefetch
                    continue
1713

1714
1715
            req.init_next_round_input(self.tree_cache)
            res = adder.add_one_req(req, has_chunked_req=(self.chunked_req is not None))
1716

1717
1718
            if res != AddReqResult.CONTINUE:
                if res == AddReqResult.NO_TOKEN:
1719
1720
                    if self.enable_hierarchical_cache:
                        # Set batch_is_full after making sure there are requests that can be served
Lianmin Zheng's avatar
Lianmin Zheng committed
1721
1722
                        self.running_batch.batch_is_full = len(
                            adder.can_run_list
1723
                        ) > 0 or (not self.running_batch.is_empty())
1724
                    else:
Lianmin Zheng's avatar
Lianmin Zheng committed
1725
                        self.running_batch.batch_is_full = True
1726
1727
                break

Lianmin Zheng's avatar
Lianmin Zheng committed
1728
        # Update waiting queue
1729
        can_run_list: List[Req] = adder.can_run_list
Lianmin Zheng's avatar
Lianmin Zheng committed
1730
1731
        if len(can_run_list) == 0:
            return None
1732
1733
1734
1735

        if self.enable_metrics:
            # only record queue time when enable_metrics is True to avoid overhead
            for req in can_run_list:
1736
                req.queue_time_end = time.perf_counter()
1737

Lianmin Zheng's avatar
Lianmin Zheng committed
1738
1739
1740
        self.waiting_queue = [
            x for x in self.waiting_queue if x not in set(can_run_list)
        ]
1741

1742
1743
1744
        if adder.new_chunked_req is not None:
            assert self.chunked_req is None
            self.chunked_req = adder.new_chunked_req
1745

1746
1747
        if self.chunked_req:
            self.chunked_req.is_chunked += 1
Lianmin Zheng's avatar
Lianmin Zheng committed
1748

1749
        # Print stats
1750
        if self.current_scheduler_metrics_enabled():
1751
            self.log_prefill_stats(adder, can_run_list, running_bs)
1752

Lianmin Zheng's avatar
Lianmin Zheng committed
1753
        # Create a new batch
1754
1755
1756
        new_batch = ScheduleBatch.init_new(
            can_run_list,
            self.req_to_token_pool,
1757
            self.token_to_kv_pool_allocator,
1758
            self.tree_cache,
1759
            self.model_config,
1760
            self.enable_overlap,
1761
            self.spec_algorithm,
1762
            chunked_req=self.chunked_req,
1763
        )
1764
1765
        if self.enable_hierarchical_cache:
            # todo (zhiqiang): disable cuda graph execution if hicache loading triggered
1766
1767
1768
            new_batch.hicache_consumer_index = (
                self.tree_cache.ready_to_load_host_cache()
            )
1769

1770
        new_batch.prepare_for_extend()
1771

Lianmin Zheng's avatar
Lianmin Zheng committed
1772
        # Mixed-style chunked prefill
1773
1774
        if (
            self.is_mixed_chunk
Lianmin Zheng's avatar
Lianmin Zheng committed
1775
            and not self.running_batch.is_empty()
1776
1777
1778
            and not (new_batch.return_logprob or self.running_batch.return_logprob)
        ):
            # TODO (lianmin): support return_logprob + mixed chunked prefill
1779
1780
            self.running_batch.filter_batch()
            if not self.running_batch.is_empty():
1781
                self.running_batch.prepare_for_decode()
1782
1783
                new_batch.mix_with_running(self.running_batch)
                new_batch.decoding_reqs = self.running_batch.reqs
Lianmin Zheng's avatar
Lianmin Zheng committed
1784
1785
1786
            self.running_batch = ScheduleBatch(
                reqs=[], batch_is_full=self.running_batch.batch_is_full
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
1787
1788
        else:
            new_batch.decoding_reqs = None
Lianmin Zheng's avatar
Lianmin Zheng committed
1789
1790
1791

        return new_batch

Lianmin Zheng's avatar
Lianmin Zheng committed
1792
    def update_running_batch(self, batch: ScheduleBatch) -> Optional[ScheduleBatch]:
1793
        """Update the current running decoding batch."""
Lianmin Zheng's avatar
Lianmin Zheng committed
1794
        initial_bs = batch.batch_size()
Lianmin Zheng's avatar
Lianmin Zheng committed
1795

1796
1797
        batch.filter_batch()
        if batch.is_empty():
Lianmin Zheng's avatar
Lianmin Zheng committed
1798
1799
            batch.batch_is_full = False
            return batch
1800

Lianmin Zheng's avatar
Lianmin Zheng committed
1801
        # Check if decode out of memory
1802
        if not batch.check_decode_mem(self.decode_mem_cache_buf_multiplier) or (
1803
            TEST_RETRACT and batch.batch_size() > 10
1804
        ):
Lianmin Zheng's avatar
Lianmin Zheng committed
1805
1806
            old_ratio = self.new_token_ratio

1807
            retracted_reqs, new_token_ratio = batch.retract_decode(self.server_args)
1808
            num_retracted_reqs = len(retracted_reqs)
Lianmin Zheng's avatar
Lianmin Zheng committed
1809
            self.new_token_ratio = new_token_ratio
1810

Lianmin Zheng's avatar
Lianmin Zheng committed
1811
            logger.info(
1812
                "KV cache pool is full. Retract requests. "
1813
                f"#retracted_reqs: {num_retracted_reqs}, "
Lianmin Zheng's avatar
Lianmin Zheng committed
1814
1815
                f"#new_token_ratio: {old_ratio:.4f} -> {self.new_token_ratio:.4f}"
            )
1816

1817
            self._extend_requests_to_queue(retracted_reqs, is_retracted=True)
1818
            self.total_retracted_reqs += num_retracted_reqs
Lianmin Zheng's avatar
Lianmin Zheng committed
1819
1820
        else:
            self.new_token_ratio = max(
1821
                self.new_token_ratio - self.new_token_ratio_decay,
Lianmin Zheng's avatar
Lianmin Zheng committed
1822
1823
1824
                self.min_new_token_ratio,
            )

Lianmin Zheng's avatar
Lianmin Zheng committed
1825
        if batch.batch_size() < initial_bs:
Lianmin Zheng's avatar
Lianmin Zheng committed
1826
            batch.batch_is_full = False
Lianmin Zheng's avatar
Lianmin Zheng committed
1827
1828

        # Update batch tensors
1829
        batch.prepare_for_decode()
Lianmin Zheng's avatar
Lianmin Zheng committed
1830
        return batch
Lianmin Zheng's avatar
Lianmin Zheng committed
1831

1832
1833
1834
    def run_batch(
        self, batch: ScheduleBatch
    ) -> Union[GenerationBatchResult, EmbeddingBatchResult]:
1835
        """Run a batch."""
Lianmin Zheng's avatar
Lianmin Zheng committed
1836
1837
        self.forward_ct += 1

1838
1839
        # Whether to run the profiler
        self._profile_batch_predicate(batch)
1840
1841
1842
1843
        if self.forward_sleep_time is not None:
            logger.info(f"Scheduler.run_batch sleep {self.forward_sleep_time}s")
            time.sleep(self.forward_sleep_time)

1844
        # Run forward
1845
        if self.is_generation:
1846
1847
            if self.spec_algorithm.is_none():
                model_worker_batch = batch.get_model_worker_batch()
1848

1849
                if self.pp_group.is_last_rank:
1850
                    logits_output, next_token_ids, can_run_cuda_graph = (
1851
1852
1853
                        self.tp_worker.forward_batch_generation(model_worker_batch)
                    )
                else:
1854
                    pp_hidden_states_proxy_tensors, _, can_run_cuda_graph = (
1855
1856
                        self.tp_worker.forward_batch_generation(model_worker_batch)
                    )
1857
                bid = model_worker_batch.bid
Lianmin Zheng's avatar
Lianmin Zheng committed
1858
            else:
1859
1860
1861
                (
                    logits_output,
                    next_token_ids,
1862
                    bid,
1863
                    num_accepted_tokens,
1864
                    can_run_cuda_graph,
1865
                ) = self.draft_worker.forward_batch_speculative_generation(batch)
1866
1867
1868
                bs = batch.batch_size()
                self.spec_num_total_accepted_tokens += num_accepted_tokens + bs
                self.spec_num_total_forward_ct += bs
1869
                self.num_generated_tokens += num_accepted_tokens
1870
1871
1872

            if self.pp_group.is_last_rank:
                batch.output_ids = next_token_ids
1873

1874
1875
1876
            # These 2 values are needed for processing the output, but the values can be
            # modified by overlap schedule. So we have to copy them here so that
            # we can use the correct values in output processing.
1877
            if batch.return_logprob or self.spec_algorithm.is_eagle():
1878
                extend_input_len_per_req = [req.extend_input_len for req in batch.reqs]
1879
1880
1881
            else:
                extend_input_len_per_req = None
            if batch.return_logprob:
1882
1883
1884
1885
1886
1887
                extend_logprob_start_len_per_req = [
                    req.extend_logprob_start_len for req in batch.reqs
                ]
            else:
                extend_logprob_start_len_per_req = None

1888
            ret = GenerationBatchResult(
1889
1890
1891
1892
1893
1894
1895
                logits_output=logits_output if self.pp_group.is_last_rank else None,
                pp_hidden_states_proxy_tensors=(
                    pp_hidden_states_proxy_tensors
                    if not self.pp_group.is_last_rank
                    else None
                ),
                next_token_ids=next_token_ids if self.pp_group.is_last_rank else None,
1896
1897
                extend_input_len_per_req=extend_input_len_per_req,
                extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
1898
                bid=bid,
1899
                can_run_cuda_graph=can_run_cuda_graph,
1900
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
1901
1902
1903
        else:  # embedding or reward model
            model_worker_batch = batch.get_model_worker_batch()
            embeddings = self.tp_worker.forward_batch_embedding(model_worker_batch)
1904
1905
1906
            ret = EmbeddingBatchResult(
                embeddings=embeddings, bid=model_worker_batch.bid
            )
1907
        return ret
Chayenne's avatar
Chayenne committed
1908

1909
1910
1911
1912
    def process_batch_result(
        self,
        batch: ScheduleBatch,
        result: Union[GenerationBatchResult, EmbeddingBatchResult],
1913
        launch_done: Optional[threading.Event] = None,
1914
    ):
Lianmin Zheng's avatar
Lianmin Zheng committed
1915
        if batch.forward_mode.is_decode():
1916
            self.process_batch_result_decode(batch, result, launch_done)
1917
        elif batch.forward_mode.is_extend():
1918
            self.process_batch_result_prefill(batch, result, launch_done)
1919
1920
        elif batch.forward_mode.is_idle():
            if self.enable_overlap:
1921
                self.tp_worker.resolve_last_batch_result(launch_done)
1922
                self.set_next_batch_sampling_info_done(batch)
1923
        elif batch.forward_mode.is_dummy_first():
1924
            self.set_next_batch_sampling_info_done(batch)
Lianmin Zheng's avatar
Lianmin Zheng committed
1925

1926
1927
1928
        self.maybe_send_health_check_signal()

    def maybe_send_health_check_signal(self):
1929
1930
1931
1932
1933
1934
1935
        if self.return_health_check_ct:
            # Return some signal for the health check.
            # This is used to prevent the health check signal being blocked by long context prefill.
            # However, one minor issue is that this code path does not check the status of detokenizer manager.
            self.return_health_check_ct -= 1
            self.send_to_tokenizer.send_pyobj(HealthCheckOutput())

1936
1937
    def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch):
        return self.prepare_mlp_sync_batch_raw(
1938
1939
1940
            local_batch,
            dp_size=self.server_args.dp_size,
            attn_tp_size=self.attn_tp_size,
1941
            tp_group=self.tp_group,
1942
1943
1944
1945
            get_idle_batch=self.get_idle_batch,
            disable_cuda_graph=self.server_args.disable_cuda_graph,
            spec_algorithm=self.spec_algorithm,
            speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
1946
            require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
1947
            disable_overlap_schedule=self.server_args.disable_overlap_schedule,
1948
1949
1950
        )

    @staticmethod
1951
    def prepare_mlp_sync_batch_raw(
1952
1953
1954
        local_batch: ScheduleBatch,
        dp_size,
        attn_tp_size: int,
1955
        tp_group,
1956
1957
1958
1959
        get_idle_batch,
        disable_cuda_graph: bool,
        spec_algorithm,
        speculative_num_draft_tokens,
1960
        require_mlp_tp_gather: bool,
1961
        disable_overlap_schedule: bool,
1962
    ):
1963
1964
1965
        # Check if other DP workers have running batches
        if local_batch is None:
            num_tokens = 0
1966
            num_tokens_for_logprob = 0
1967
1968
        elif local_batch.forward_mode.is_decode():
            num_tokens = local_batch.batch_size()
1969
            num_tokens_for_logprob = num_tokens
1970
1971
        else:
            num_tokens = local_batch.extend_num_tokens
1972
            num_tokens_for_logprob = sum(
Lianmin Zheng's avatar
Lianmin Zheng committed
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
                [
                    # We should have at least 1 token for sample in every case.
                    max(extend_len - logprob_start_len, 1)
                    for logprob_start_len, extend_len in zip(
                        local_batch.extend_logprob_start_lens, local_batch.extend_lens
                    )
                ]
            )

        if local_batch is None or local_batch.forward_mode.is_decode_or_idle():
            can_cuda_graph = 1
        else:
            can_cuda_graph = 0

        is_extend_in_batch = (
            local_batch.forward_mode.is_extend() if local_batch else False
        )
1990
1991

        tbo_preparer = TboDPAttentionPreparer()
1992
1993
1994
1995
1996
1997
        if disable_overlap_schedule:
            group = tp_group.device_group
            device = tp_group.device
        else:
            group = tp_group.cpu_group
            device = "cpu"
1998

Lianmin Zheng's avatar
Lianmin Zheng committed
1999
2000
2001
2002
        local_info = torch.tensor(
            [
                num_tokens,
                can_cuda_graph,
2003
                num_tokens_for_logprob,
Lianmin Zheng's avatar
Lianmin Zheng committed
2004
                is_extend_in_batch,
2005
2006
2007
                *tbo_preparer.prepare_all_gather(
                    local_batch,
                ),
Lianmin Zheng's avatar
Lianmin Zheng committed
2008
2009
            ],
            dtype=torch.int64,
2010
            device=device,
Lianmin Zheng's avatar
Lianmin Zheng committed
2011
2012
        )
        global_info = torch.empty(
2013
            (dp_size, attn_tp_size, 6),
Lianmin Zheng's avatar
Lianmin Zheng committed
2014
            dtype=torch.int64,
2015
            device=device,
Lianmin Zheng's avatar
Lianmin Zheng committed
2016
        )
2017
        torch.distributed.all_gather_into_tensor(
Lianmin Zheng's avatar
Lianmin Zheng committed
2018
2019
            global_info.flatten(),
            local_info,
2020
            group=group,
2021
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
2022
2023
2024
2025
        global_num_tokens = global_info[:, 0, 0].tolist()
        can_cuda_graph = min(global_info[:, 0, 1].tolist())
        global_num_tokens_for_logprob = global_info[:, 0, 2].tolist()
        is_extend_in_batch = global_info[:, 0, 3].tolist()
2026

2027
2028
2029
2030
        tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output(
            global_info[:, :, 4:6]
        )

Lianmin Zheng's avatar
Lianmin Zheng committed
2031
        if local_batch is None and max(global_num_tokens) > 0:
2032
            local_batch = get_idle_batch()
2033
2034

        if local_batch is not None:
2035
            # TODO: handle the case when moe_dense_tp_size != 1
2036
            if not require_mlp_tp_gather:
2037
2038
2039
2040
2041
2042
2043
                local_batch.global_num_tokens = [num_tokens]
                local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob]
            else:
                local_batch.global_num_tokens = global_num_tokens
                local_batch.global_num_tokens_for_logprob = (
                    global_num_tokens_for_logprob
                )
2044
            local_batch.is_extend_in_batch = any(is_extend_in_batch)
2045
2046
            local_batch.tbo_split_seq_index = tbo_split_seq_index
            local_batch.global_forward_mode = global_forward_mode
2047

2048
            # Check forward mode for cuda graph
2049
            if not disable_cuda_graph:
Lianmin Zheng's avatar
Lianmin Zheng committed
2050
                local_batch.can_run_dp_cuda_graph = can_cuda_graph
2051

2052
        return local_batch
2053
2054
2055
2056
2057

    def get_idle_batch(self):
        idle_batch = ScheduleBatch.init_new(
            [],
            self.req_to_token_pool,
2058
            self.token_to_kv_pool_allocator,
2059
2060
2061
            self.tree_cache,
            self.model_config,
            self.enable_overlap,
2062
            self.spec_algorithm,
2063
2064
2065
2066
        )
        idle_batch.prepare_for_idle()
        return idle_batch

2067
2068
    def move_ready_grammar_requests(self):
        """Move requests whose grammar objects are ready from grammar_queue to waiting_queue."""
2069

2070
        num_ready_reqs = 0
2071
        num_timeout_reqs = 0
2072
2073
        for req in self.grammar_queue:
            try:
2074
2075
2076
                if req.finished():  # It is aborted by AbortReq
                    num_ready_reqs += 1
                    continue
2077
                req.grammar = req.grammar.result(timeout=0.03)
2078
2079
2080
2081
2082
                self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
                if req.grammar is INVALID_GRAMMAR_OBJ:
                    req.set_finish_with_abort(
                        f"Invalid grammar request: {req.grammar_key=}"
                    )
2083
2084
                num_ready_reqs += 1
            except futures._base.TimeoutError:
2085
                req.grammar_wait_ct += 1
2086
2087
                # NOTE(lianmin): this timeout is the waiting time of the above line. It is
                # not the waiting time from it enters the grammar queue.
2088
                if req.grammar_wait_ct > GRAMMAR_TIMEOUT / 0.03:
2089
                    num_timeout_reqs = 1
2090
2091
                break

2092
        if self.server_args.enable_dp_attention:
2093
2094
            tp_size = self.attn_tp_size
            tp_group = self.attn_tp_cpu_group
2095
        else:
2096
2097
2098
2099
2100
            tp_size = self.tp_size
            tp_group = self.tp_cpu_group

        if tp_size > 1:
            # Sync across TP ranks to make sure they have the same number of ready requests
2101
            tensor = torch.tensor([num_ready_reqs, num_timeout_reqs], dtype=torch.int32)
2102
2103
2104
            torch.distributed.all_reduce(
                tensor, op=torch.distributed.ReduceOp.MAX, group=tp_group
            )
2105
            num_ready_reqs_max, num_timeout_reqs_max = tensor.tolist()
2106

2107
            for i in range(num_ready_reqs, num_ready_reqs_max):
2108
                req = self.grammar_queue[i]
2109
2110
                if req.finished():  # It is aborted by AbortReq
                    continue
2111
                req.grammar = req.grammar.result()
2112
2113
2114
2115
2116
2117
2118
2119
                self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
                if req.grammar is INVALID_GRAMMAR_OBJ:
                    req.set_finish_with_abort(
                        f"Invalid grammar request: {req.grammar_key=}"
                    )
        else:
            num_ready_reqs_max = num_ready_reqs
            num_timeout_reqs_max = num_timeout_reqs
2120

2121
2122
2123
2124
2125
2126
2127
        for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_max):
            req = self.grammar_queue[i]
            req.grammar.cancel()
            error_msg = f"Grammar preprocessing timed out for {req.grammar_key=}"
            req.set_finish_with_abort(error_msg)
            self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ)
        num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max
2128

2129
        self._extend_requests_to_queue(self.grammar_queue[:num_ready_reqs])
2130
2131
        self.grammar_queue = self.grammar_queue[num_ready_reqs:]

2132
2133
2134
2135
2136
2137
2138
    def set_next_batch_sampling_info_done(self, batch: ScheduleBatch):
        if batch.next_batch_sampling_info:
            if batch.next_batch_sampling_info.grammars is not None:
                batch.next_batch_sampling_info.update_regex_vocab_mask()
                self.current_stream.synchronize()
            batch.next_batch_sampling_info.sampling_info_done.set()

2139
2140
2141
    def watchdog_thread(self):
        """A watch dog thread that will try to kill the server itself if one forward batch takes too long."""
        self.watchdog_last_forward_ct = 0
2142
        self.watchdog_last_time = time.perf_counter()
2143
2144

        while True:
2145
            current = time.perf_counter()
2146
2147
2148
2149
2150
2151
2152
2153
2154
            if self.cur_batch is not None:
                if self.watchdog_last_forward_ct == self.forward_ct:
                    if current > self.watchdog_last_time + self.watchdog_timeout:
                        break
                else:
                    self.watchdog_last_forward_ct = self.forward_ct
                    self.watchdog_last_time = current
            time.sleep(self.watchdog_timeout // 2)

Lianmin Zheng's avatar
Lianmin Zheng committed
2155
2156
        if not disable_request_logging():
            # Print batch size and memory pool info to check whether there are de-sync issues.
Hanming Lu's avatar
Hanming Lu committed
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
            if self.is_hybrid:
                (
                    _,
                    _,
                    _,
                    _,
                    full_available_size,
                    full_evictable_size,
                    swa_available_size,
                    swa_evictable_size,
                ) = self._get_swa_token_info()
                info_msg = (
                    f"{full_available_size=}, "
                    f"{full_evictable_size=}, "
                    f"{swa_available_size=}, "
                    f"{swa_evictable_size=}, "
                )
            else:
                _, _, available_size, evictable_size = self._get_token_info()
                info_msg = f"{available_size=}, " f"{evictable_size=}, "
Lianmin Zheng's avatar
Lianmin Zheng committed
2177
2178
2179
            logger.error(
                f"{self.cur_batch.batch_size()=}, "
                f"{self.cur_batch.reqs=}, "
Hanming Lu's avatar
Hanming Lu committed
2180
                f"{info_msg}"
Lianmin Zheng's avatar
Lianmin Zheng committed
2181
2182
            )

2183
        pyspy_dump_schedulers()
Lianmin Zheng's avatar
Lianmin Zheng committed
2184
        logger.error(f"Watchdog timeout ({self.watchdog_timeout=})")
2185
2186
        print(file=sys.stderr, flush=True)
        print(file=sys.stdout, flush=True)
Lianmin Zheng's avatar
Lianmin Zheng committed
2187
2188

        # Wait for some time so that the parent process can print the error.
2189
2190
2191
        time.sleep(5)
        self.parent_process.send_signal(signal.SIGQUIT)

2192
2193
2194
    def flush_cache_wrapped(self, recv_req: FlushCacheReqInput):
        success = self.flush_cache()
        return FlushCacheReqOutput(success=success)
2195

2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
    def clear_hicache_storage_wrapped(self, recv_req: ClearHiCacheReqInput):
        if self.enable_hierarchical_cache:
            self.tree_cache.clear_storage_backend()
            logger.info("Hierarchical cache cleared successfully!")
            if_success = True
        else:
            logging.warning("Hierarchical cache is not enabled.")
            if_success = False
        return ClearHiCacheReqOutput(success=if_success)

2206
    def flush_cache(self):
2207
        """Flush the memory pool and cache."""
2208
2209
2210
2211
2212
        if (
            len(self.waiting_queue) == 0
            and self.running_batch.is_empty()
            and (self.pp_size == 1 or all(x.is_empty() for x in self.running_mbs))
        ):
2213
2214
            self.cur_batch = None
            self.last_batch = None
2215
            self.tree_cache.reset()
2216
            if self.grammar_backend:
Lianmin Zheng's avatar
Lianmin Zheng committed
2217
                self.grammar_backend.reset()
2218
            self.req_to_token_pool.clear()
2219
            self.token_to_kv_pool_allocator.clear()
2220
2221
2222

            if not self.spec_algorithm.is_none():
                self.draft_worker.model_runner.req_to_token_pool.clear()
2223
                self.draft_worker.model_runner.token_to_kv_pool_allocator.clear()
2224
2225
2226
2227
2228

            self.num_generated_tokens = 0
            self.forward_ct_decode = 0
            self.spec_num_total_accepted_tokens = 0
            self.spec_num_total_forward_ct = 0
2229
2230
            self.cum_spec_accept_length = 0
            self.cum_spec_accept_count = 0
2231
2232
2233
2234
2235
2236
2237
            torch.cuda.empty_cache()
            logger.info("Cache flushed successfully!")
            if_success = True
        else:
            logging.warning(
                f"Cache not flushed because there are pending requests. "
                f"#queue-req: {len(self.waiting_queue)}, "
Lianmin Zheng's avatar
Lianmin Zheng committed
2238
                f"#running-req: {len(self.running_batch.reqs)}"
2239
2240
2241
2242
            )
            if_success = False
        return if_success

Liangsheng Yin's avatar
Liangsheng Yin committed
2243
2244
    def get_load(self):
        # TODO(lsyin): use dynamically maintained num_waiting_tokens
Hanming Lu's avatar
Hanming Lu committed
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
        if self.is_hybrid:
            load_full = (
                self.full_tokens_per_layer
                - self.token_to_kv_pool_allocator.full_available_size()
                - self.tree_cache.full_evictable_size()
            )
            load_swa = (
                self.swa_tokens_per_layer
                - self.token_to_kv_pool_allocator.swa_available_size()
                - self.tree_cache.swa_evictable_size()
            )
            load = max(load_full, load_swa)
        else:
            load = (
                self.max_total_num_tokens
                - self.token_to_kv_pool_allocator.available_size()
                - self.tree_cache.evictable_size()
            )
Liangsheng Yin's avatar
Liangsheng Yin committed
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
        load += sum(len(req.origin_input_ids) for req in self.waiting_queue)
        if self.disaggregation_mode == DisaggregationMode.PREFILL:
            load += sum(
                len(req.origin_input_ids)
                for req in self.disagg_prefill_bootstrap_queue.queue
            )
        elif self.disaggregation_mode == DisaggregationMode.DECODE:
            load += sum(
                len(req.req.origin_input_ids)
                for req in self.disagg_decode_prealloc_queue.queue
            )

        return load

2277
2278
2279
    def get_internal_state(self, recv_req: GetInternalStateReq):
        ret = dict(global_server_args_dict)
        ret["last_gen_throughput"] = self.last_gen_throughput
2280
2281
2282
2283
2284
2285
2286
2287
2288
        ret["memory_usage"] = {
            "weight": round(
                self.tp_worker.worker.model_runner.weight_load_mem_usage, 2
            ),
            "kvcache": round(
                self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 2
            ),
            "token_capacity": int(self.max_total_num_tokens),
        }
2289

2290
2291
2292
        ret["memory_usage"]["graph"] = round(
            self.tp_worker.worker.model_runner.graph_mem_usage, 2
        )
2293

2294
2295
2296
2297
2298
2299
        if not self.spec_algorithm.is_none() and self.cum_spec_accept_count > 0:
            ret["avg_spec_accept_length"] = (
                self.cum_spec_accept_length / self.cum_spec_accept_count
            )
        if RECORD_STEP_TIME:
            ret["step_time_dict"] = self.step_time_dict
Liangsheng Yin's avatar
Liangsheng Yin committed
2300
2301
2302
2303

        ret["load"] = self.get_load()

        return GetInternalStateReqOutput(internal_state=ret)
2304
2305
2306
2307
2308

    def set_internal_state(self, recv_req: SetInternalStateReq):
        server_args_dict = recv_req.server_args
        args_allow_update = set(
            [
2309
                "max_micro_batch_size",
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
                "speculative_accept_threshold_single",
                "speculative_accept_threshold_acc",
            ]
        )
        if_success = True
        for k, v in server_args_dict.items():
            if k not in args_allow_update:
                logging.warning(f"Updating {k} is not supported.")
                if_success = False
                break
2320
2321
2322
2323
2324
2325
2326
2327
            elif k == "max_micro_batch_size" and (
                v > self.max_running_requests // self.pp_size or v < 1
            ):
                logging.warning(
                    f"Updating {k} to {v} is rejected because it is out of the valid range [1, {self.max_running_requests // self.pp_size}]."
                )
                if_success = False
                break
2328
2329
2330
2331
2332
2333
2334
2335
2336
        if if_success:
            if not self.spec_algorithm.is_none() and self.cum_spec_accept_count > 0:
                avg_spec_accept_length = (
                    self.cum_spec_accept_length / self.cum_spec_accept_count
                )
                logger.info(f"{avg_spec_accept_length=}")
            self.cum_spec_accept_length = self.cum_spec_accept_count = 0
            for k, v in server_args_dict.items():
                global_server_args_dict[k] = v
2337
            logger.info(f"Global server args updated! {global_server_args_dict=}")
2338
2339
2340
2341
2342
        return SetInternalStateReqOutput(
            updated=True,
            server_args=global_server_args_dict,
        )

2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
    def handle_rpc_request(self, recv_req: RpcReqInput):
        # Handle RPC requests
        logger.info(
            f"handle_rpc_request: {recv_req.method}, param: {recv_req.parameters}"
        )

        success = True
        exec = None
        try:
            func = getattr(self, recv_req.method)
            func(recv_req.parameters)
        except Exception as e:
            success = False
            exec = e
            logger.error(f"Failed to call rpc {recv_req.method}: {str(e)}")

        barrier()
        return RpcReqOutput(success, "" if not exec else str(exec))

2362
2363
    def abort_request(self, recv_req: AbortReq):
        # Delete requests in the waiting queue
Lianmin Zheng's avatar
Lianmin Zheng committed
2364
        to_del = []
2365
        for i, req in enumerate(self.waiting_queue):
2366
            if recv_req.abort_all or req.rid.startswith(recv_req.rid):
Lianmin Zheng's avatar
Lianmin Zheng committed
2367
                to_del.append(i)
2368

Lianmin Zheng's avatar
Lianmin Zheng committed
2369
        # Sort in reverse order to avoid index issues when deleting
Lianmin Zheng's avatar
Lianmin Zheng committed
2370
        for i in reversed(to_del):
2371
2372
2373
            # Abort method 1: directly pop from the queue
            # This only works for requests that have not started anything.
            # We still need to send something back to TokenizerManager to clean up the state.
Lianmin Zheng's avatar
Lianmin Zheng committed
2374
            req = self.waiting_queue.pop(i)
2375
2376
2377
            if self.enable_hicache_storage:
                # to release prefetch events associated with the request
                self.tree_cache.release_aborted_request(req.rid)
Lianmin Zheng's avatar
Lianmin Zheng committed
2378
            self.send_to_tokenizer.send_pyobj(AbortReq(req.rid))
2379
2380
2381
2382
            # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
            if self.disaggregation_mode == DisaggregationMode.DECODE:
                self.tree_cache.cache_finished_req(req)

2383
            logger.debug(f"Abort queued request. {req.rid=}")
2384

2385
2386
2387
2388
2389
        # Delete the requests in the grammar queue
        for req in self.grammar_queue:
            # Abort method 2: call `set_finish_with_abort`
            # The request will still run one prefill forward pass.
            # In this case, we change the input_ids to be only one token to make this prefill cheap.
2390
            if recv_req.abort_all or req.rid.startswith(recv_req.rid):
2391
                logger.debug(f"Abort grammar queue request. {req.rid=}")
2392
2393
                if req.grammar:
                    req.grammar.cancel()
2394
2395
                req.set_finish_with_abort("Aborted by AbortReq.")

2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
        # Delete requests not in the waiting queue when PD disaggregation is enabled
        if self.disaggregation_mode == DisaggregationMode.PREFILL:
            # Abort requests that have not yet been bootstrapped
            for i, req in enumerate(self.disagg_prefill_bootstrap_queue.queue):
                logger.debug(f"Abort bootstrap queue request. {req.rid=}")
                if recv_req.abort_all or req.rid.startswith(recv_req.rid):
                    if hasattr(req.disagg_kv_sender, "abort"):
                        req.disagg_kv_sender.abort()

            # Abort in-flight requests
            for i, req in enumerate(self.disagg_prefill_inflight_queue):
                logger.debug(f"Abort inflight queue request. {req.rid=}")
                if recv_req.abort_all or req.rid.startswith(recv_req.rid):
                    if hasattr(req.disagg_kv_sender, "abort"):
                        req.disagg_kv_sender.abort()

        elif self.disaggregation_mode == DisaggregationMode.DECODE:
            # Abort requests that have not yet finished preallocation
            for i, decode_req in enumerate(self.disagg_decode_prealloc_queue.queue):
                logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
                if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
                    if hasattr(decode_req.kv_receiver, "abort"):
                        decode_req.kv_receiver.abort()

            # Abort requests waiting for kvcache to release tree cache
            for i, decode_req in enumerate(self.disagg_decode_transfer_queue.queue):
                logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}")
                if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
                    if hasattr(decode_req.kv_receiver, "abort"):
                        decode_req.kv_receiver.abort()

2427
        # Delete requests in the running batch
Lianmin Zheng's avatar
Lianmin Zheng committed
2428
2429
2430
2431
2432
2433
        if self.cur_batch is self.running_batch or self.cur_batch is None:
            reqs = self.running_batch.reqs
        else:
            reqs = self.running_batch.reqs + self.cur_batch.reqs

        for req in reqs:
2434
2435
2436
            if not req.finished() and (
                recv_req.abort_all or req.rid.startswith(recv_req.rid)
            ):
2437
2438
2439
                # Abort method 3: set `to_abort=True`
                # The request will still run one decode forward pass.
                # Then we reuse all existing code to clean up the KV cache allocation.
Lianmin Zheng's avatar
Lianmin Zheng committed
2440
2441
                logger.debug(f"Abort running request. {req.rid=}")
                req.to_abort = True
2442

2443
2444
2445
    def _pause_engine(self) -> Tuple[List[Req], int]:
        raise NotImplementedError()

2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
    def load_lora_adapter(
        self, recv_req: LoadLoRAAdapterReqInput
    ) -> LoadLoRAAdapterReqOutput:
        """In-place loading a new lora adapter from disk or huggingface."""

        result = self.tp_worker.load_lora_adapter(recv_req)
        return result

    def unload_lora_adapter(
        self, recv_req: UnloadLoRAAdapterReqInput
    ) -> UnloadLoRAAdapterReqOutput:
        """Unload the lora adapter."""

        result = self.tp_worker.unload_lora_adapter(recv_req)
        return result

2462
2463
2464
2465
    def register_multi_tokenizer(self, recv_req: MultiTokenizerRegisterReq):
        self.send_to_detokenizer.send_pyobj(recv_req)
        return recv_req

2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
    def init_weights_send_group_for_remote_instance(
        self, recv_req: InitWeightsSendGroupForRemoteInstanceReqInput
    ):
        """Init the seed and client instance communication group."""
        success, message = self.tp_worker.init_weights_send_group_for_remote_instance(
            recv_req
        )
        return InitWeightsSendGroupForRemoteInstanceReqOutput(success, message)

    def send_weights_to_remote_instance(
        self, recv_req: SendWeightsToRemoteInstanceReqInput
    ):
        """Send the seed instance weights to the destination instance."""
        success, message = self.tp_worker.send_weights_to_remote_instance(recv_req)
        return SendWeightsToRemoteInstanceReqOutput(success, message)

2482
2483
2484
2485
2486
2487
2488
    def slow_down(self, recv_req: SlowDownReqInput):
        t = recv_req.forward_sleep_time
        if t is not None and t <= 0:
            t = None
        self.forward_sleep_time = t
        return SlowDownReqOutput()

2489
2490
    def expert_distribution_handle(self, recv_req: ExpertDistributionReq):
        if recv_req == ExpertDistributionReq.START_RECORD:
2491
            get_global_expert_distribution_recorder().start_record()
2492
        elif recv_req == ExpertDistributionReq.STOP_RECORD:
2493
            get_global_expert_distribution_recorder().stop_record()
2494
        elif recv_req == ExpertDistributionReq.DUMP_RECORD:
2495
            get_global_expert_distribution_recorder().dump_record()
2496
        else:
2497
            raise ValueError(f"Unrecognized ExpertDistributionReq value: {recv_req=}")
2498
        return ExpertDistributionReqOutput()
2499

2500
    def open_session(self, recv_req: OpenSessionReqInput):
2501
2502
2503
2504
        # handle error
        session_id = recv_req.session_id
        if session_id in self.sessions:
            logger.warning(f"session id {session_id} already exist, cannot open.")
2505
            return OpenSessionReqOutput(session_id, False)
2506
        elif session_id is None:
2507
            logger.warning("session id is None, cannot open.")
2508
            return OpenSessionReqOutput(session_id, False)
2509
2510
2511
2512
        else:
            self.sessions[session_id] = Session(
                recv_req.capacity_of_str_len, session_id
            )
2513
            return OpenSessionReqOutput(session_id, True)
2514
2515
2516
2517
2518
2519
2520
2521
2522

    def close_session(self, recv_req: CloseSessionReqInput):
        # handle error
        session_id = recv_req.session_id
        if session_id not in self.sessions:
            logger.warning(f"session id {session_id} does not exist, cannot delete.")
        else:
            del self.sessions[session_id]

2523
2524
    def get_print_prefix(self):
        prefix = ""
2525
2526
        if self.attn_dp_rank is not None:
            prefix += f" DP{self.attn_dp_rank}"
2527
2528
2529
2530
2531
2532
        if self.server_args.tp_size > 1:
            prefix += f" TP{self.tp_rank}"
        if self.pp_size > 1:
            prefix += f" PP{self.pp_rank}"
        return prefix

2533
2534
    def current_scheduler_metrics_enabled(self):
        return self.attn_tp_rank == 0 or self.enable_metrics_for_all_schedulers
2535

2536
2537
2538
    def maybe_sleep_on_idle(self):
        if self.idle_sleeper is not None:
            self.idle_sleeper.maybe_sleep()
2539

2540
2541
2542
2543
2544
2545
    def handle_freeze_gc(self, recv_req: FreezeGCReq):
        """Handle freeze_gc request: freeze scheduler's GC and forward to detokenizer."""
        freeze_gc("Scheduler")
        self.send_to_detokenizer.send_pyobj(recv_req)
        return None

2546

2547
2548
2549
2550
2551
2552
2553
class IdleSleeper:
    """
    In setups which have long inactivity periods it is desirable to reduce
    system power consumption when sglang does nothing. This would lead not only
    to power savings, but also to more CPU thermal headroom when a request
    eventually comes. This is important in cases when multiple GPUs are connected
    as each GPU would otherwise pin one thread at 100% CPU usage.
2554

2555
2556
2557
    The simplest solution is to use zmq.Poller on all sockets that may receive
    data that needs handling immediately.
    """
2558

2559
2560
    def __init__(self, sockets):
        self.poller = zmq.Poller()
2561
        self.last_empty_time = time.time()
2562
2563
2564
2565
2566
        for s in sockets:
            self.poller.register(s, zmq.POLLIN)

    def maybe_sleep(self):
        self.poller.poll(1000)
2567
2568
2569
2570
2571
2572
2573
        if (
            global_config.torch_empty_cache_interval > 0
            and time.time() - self.last_empty_time
            > global_config.torch_empty_cache_interval
        ):
            self.last_empty_time = time.time()
            torch.cuda.empty_cache()
2574

2575

2576
2577
def is_health_check_generate_req(recv_req):
    return getattr(recv_req, "rid", "").startswith("HEALTH_CHECK")
2578

2579
2580

def is_work_request(recv_req):
2581
2582
2583
2584
2585
2586
2587
2588
2589
    return isinstance(
        recv_req,
        (
            TokenizedGenerateReqInput,
            TokenizedEmbeddingReqInput,
            BatchTokenizedGenerateReqInput,
            BatchTokenizedEmbeddingReqInput,
        ),
    )
2590
2591


2592
2593
2594
2595
2596
def run_scheduler_process(
    server_args: ServerArgs,
    port_args: PortArgs,
    gpu_id: int,
    tp_rank: int,
Cheng Wan's avatar
Cheng Wan committed
2597
    moe_ep_rank: int,
2598
    pp_rank: int,
2599
    dp_rank: Optional[int],
2600
    pipe_writer,
2601
    balance_meta: Optional[DPBalanceMeta] = None,
2602
):
2603
2604
2605
    if (numa_node := server_args.numa_node) is not None:
        numa_bind_to_node(numa_node[gpu_id])

2606
    # Generate the prefix
2607
2608
2609
2610
2611
    prefix = ""
    if dp_rank is not None:
        prefix += f" DP{dp_rank}"
    if server_args.tp_size > 1:
        prefix += f" TP{tp_rank}"
Cheng Wan's avatar
Cheng Wan committed
2612
2613
    if server_args.ep_size > 1:
        prefix += f" EP{moe_ep_rank}"
2614
2615
    if server_args.pp_size > 1:
        prefix += f" PP{pp_rank}"
2616

2617
    # Config the process
2618
    setproctitle.setproctitle(f"sglang::scheduler{prefix.replace(' ', '_')}")
2619
    faulthandler.enable()
2620
    kill_itself_when_parent_died()
2621
    parent_process = psutil.Process().parent()
2622

2623
2624
2625
    # [For Router] if env var "SGLANG_DP_RANK" exist, set dp_rank to the value of the env var
    if dp_rank is None and "SGLANG_DP_RANK" in os.environ:
        dp_rank = int(os.environ["SGLANG_DP_RANK"])
2626

Wang Ran (汪然)'s avatar
Wang Ran (汪然) committed
2627
    # Configure the logger
2628
    configure_logger(server_args, prefix=prefix)
2629
    suppress_other_loggers()
2630

2631
    # Set cpu affinity to this gpu process
2632
2633
2634
    if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
        set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, gpu_id)

2635
    # Create a scheduler and run the event loop
2636
    try:
Cheng Wan's avatar
Cheng Wan committed
2637
        scheduler = Scheduler(
2638
2639
2640
2641
2642
2643
2644
2645
            server_args,
            port_args,
            gpu_id,
            tp_rank,
            moe_ep_rank,
            pp_rank,
            dp_rank,
            dp_balance_meta=balance_meta,
Cheng Wan's avatar
Cheng Wan committed
2646
        )
2647
        pipe_writer.send(
Mick's avatar
Mick committed
2648
2649
2650
2651
2652
            {
                "status": "ready",
                "max_total_num_tokens": scheduler.max_total_num_tokens,
                "max_req_input_len": scheduler.max_req_input_len,
            }
2653
        )
Byron Hsu's avatar
Byron Hsu committed
2654

2655
        disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
Byron Hsu's avatar
Byron Hsu committed
2656
        if disaggregation_mode == DisaggregationMode.NULL:
2657
2658
2659
            if server_args.pp_size > 1:
                scheduler.event_loop_pp()
            elif scheduler.enable_overlap:
Byron Hsu's avatar
Byron Hsu committed
2660
2661
2662
2663
                scheduler.event_loop_overlap()
            else:
                scheduler.event_loop_normal()
        elif disaggregation_mode == DisaggregationMode.PREFILL:
2664
2665
2666
            if scheduler.enable_overlap:
                scheduler.event_loop_overlap_disagg_prefill()
            else:
2667
2668
2669
2670
                if server_args.pp_size > 1:
                    scheduler.event_loop_pp_disagg_prefill()
                else:
                    scheduler.event_loop_normal_disagg_prefill()
2671

Byron Hsu's avatar
Byron Hsu committed
2672
        elif disaggregation_mode == DisaggregationMode.DECODE:
2673
2674
2675
2676
            if scheduler.enable_overlap:
                scheduler.event_loop_overlap_disagg_decode()
            else:
                scheduler.event_loop_normal_disagg_decode()
Byron Hsu's avatar
Byron Hsu committed
2677

2678
    except Exception:
2679
2680
2681
        traceback = get_exception_traceback()
        logger.error(f"Scheduler hit an exception: {traceback}")
        parent_process.send_signal(signal.SIGQUIT)