model_runner.py 71.4 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.
# ==============================================================================
Lianmin Zheng's avatar
Lianmin Zheng committed
14
"""ModelRunner runs the forward passes of the models."""
15

16
import datetime
17
import gc
18
import inspect
Shuo Yang's avatar
Shuo Yang committed
19
import json
20
import logging
21
import os
22
import time
23
24
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union
Lianmin Zheng's avatar
Lianmin Zheng committed
25
26

import torch
27
import torch.distributed as dist
28
29
30
31

from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import AttentionArch, ModelConfig
32
from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp
33
from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS
34
from sglang.srt.distributed import (
zhyncs's avatar
zhyncs committed
35
    get_tp_group,
36
    get_world_group,
zhyncs's avatar
zhyncs committed
37
38
    init_distributed_environment,
    initialize_model_parallel,
39
    set_custom_all_reduce,
40
    set_mscclpp_all_reduce,
zhyncs's avatar
zhyncs committed
41
)
42
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
fzyzcjy's avatar
fzyzcjy committed
43
44
45
46
47
48
49
50
51
52
53
54
55
from sglang.srt.eplb.eplb_manager import EPLBManager
from sglang.srt.eplb.expert_distribution import (
    ExpertDistributionRecorder,
    get_global_expert_distribution_recorder,
    set_global_expert_distribution_recorder,
)
from sglang.srt.eplb.expert_location import (
    ExpertLocationMetadata,
    compute_initial_expert_location_metadata,
    get_global_expert_location_metadata,
    set_global_expert_location_metadata,
)
from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
56
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
57
58
from sglang.srt.layers.dp_attention import (
    get_attention_tp_group,
59
    get_attention_tp_size,
60
61
    initialize_dp_attention,
)
Liangsheng Yin's avatar
Liangsheng Yin committed
62
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
63
64
65
from sglang.srt.layers.quantization import (
    deep_gemm_wrapper,
    monkey_patch_isinstance_for_vllm_base_layer,
66
)
67
from sglang.srt.layers.sampler import Sampler
68
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
69
from sglang.srt.layers.utils import is_sm100_supported
70
from sglang.srt.lora.lora_manager import LoRAManager
71
from sglang.srt.lora.lora_registry import LoRARef
72
73
74
75
from sglang.srt.managers.schedule_batch import (
    GLOBAL_SERVER_ARGS_KEYS,
    global_server_args_dict,
)
76
from sglang.srt.mem_cache.allocator import (
77
    AscendPagedTokenToKVPoolAllocator,
78
79
    BaseTokenToKVPoolAllocator,
    PagedTokenToKVPoolAllocator,
tarinkk's avatar
tarinkk committed
80
    SWATokenToKVPoolAllocator,
81
82
    TokenToKVPoolAllocator,
)
83
from sglang.srt.mem_cache.memory_pool import (
84
85
    AscendMLAPagedTokenToKVPool,
    AscendTokenToKVPool,
Shuo Yang's avatar
Shuo Yang committed
86
    DoubleSparseTokenToKVPool,
87
88
89
    MHATokenToKVPool,
    MLATokenToKVPool,
    ReqToTokenPool,
tarinkk's avatar
tarinkk committed
90
    SWAKVPool,
91
)
Yineng Zhang's avatar
Yineng Zhang committed
92
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
93
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
94
from sglang.srt.model_loader import get_model
95
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
Lianmin Zheng's avatar
Lianmin Zheng committed
96
from sglang.srt.model_loader.utils import set_default_torch_dtype
97
from sglang.srt.model_loader.weight_utils import default_weight_loader
98
from sglang.srt.patch_torch import monkey_patch_torch_reductions
99
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
Lianmin Zheng's avatar
Lianmin Zheng committed
100
from sglang.srt.server_args import ServerArgs
101
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
102
from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
103
from sglang.srt.utils import (
104
    MultiprocessingSerializer,
105
    cpu_has_amx_support,
106
    dynamic_import,
107
    enable_show_time_cost,
108
    get_available_gpu_memory,
109
    get_bool_env_var,
110
    get_cpu_ids_by_node,
111
    init_custom_process_group,
112
    is_fa3_default_architecture,
113
    is_flashinfer_available,
HAI's avatar
HAI committed
114
    is_hip,
115
    is_hopper_with_cuda_12_3,
116
    is_no_spec_infer_or_topk_one,
117
    is_npu,
118
    monkey_patch_p2p_access_check,
119
    monkey_patch_vllm_gguf_config,
120
    set_cpu_offload_max_bytes,
121
    set_cuda_arch,
122
)
123

124
_is_hip = is_hip()
125
_is_npu = is_npu()
126
_is_cpu_amx_available = cpu_has_amx_support()
127

Lianmin Zheng's avatar
Lianmin Zheng committed
128
# Use a small KV cache pool size for tests in CI
129
SGLANG_CI_SMALL_KV_SIZE = os.getenv("SGLANG_CI_SMALL_KV_SIZE", None)
Lianmin Zheng's avatar
Lianmin Zheng committed
130
131

# Detect stragger ranks in model loading
132
133
UNBALANCED_MODEL_LOADING_TIMEOUT_S = 300

Lianmin Zheng's avatar
Lianmin Zheng committed
134
135
logger = logging.getLogger(__name__)

136

137
138
139
140
141
142
143
144
145
146
147
148
149
class RankZeroFilter(logging.Filter):
    """Filter that only allows INFO level logs from rank 0, but allows all other levels from any rank."""

    def __init__(self, is_rank_zero):
        super().__init__()
        self.is_rank_zero = is_rank_zero

    def filter(self, record):
        if record.levelno == logging.INFO:
            return self.is_rank_zero
        return True


Lianmin Zheng's avatar
Lianmin Zheng committed
150
class ModelRunner:
151
152
    """ModelRunner runs the forward passes of the models."""

Lianmin Zheng's avatar
Lianmin Zheng committed
153
154
    def __init__(
        self,
155
        model_config: ModelConfig,
156
157
158
159
        mem_fraction_static: float,
        gpu_id: int,
        tp_rank: int,
        tp_size: int,
Cheng Wan's avatar
Cheng Wan committed
160
161
        moe_ep_rank: int,
        moe_ep_size: int,
162
163
        pp_rank: int,
        pp_size: int,
164
        nccl_port: int,
Lianmin Zheng's avatar
Lianmin Zheng committed
165
        server_args: ServerArgs,
166
        is_draft_worker: bool = False,
167
        req_to_token_pool: Optional[ReqToTokenPool] = None,
168
        token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
Lianmin Zheng's avatar
Lianmin Zheng committed
169
    ):
170
        # Parse args
Lianmin Zheng's avatar
Lianmin Zheng committed
171
        self.mem_fraction_static = mem_fraction_static
Zhang, Liangang's avatar
Zhang, Liangang committed
172
        self.device = server_args.device
173
        self.gpu_id = gpu_id
174
175
176
177

        # Apply the rank zero filter to logger
        if not any(isinstance(f, RankZeroFilter) for f in logger.filters):
            logger.addFilter(RankZeroFilter(tp_rank == 0))
Lianmin Zheng's avatar
Lianmin Zheng committed
178
179
        self.tp_rank = tp_rank
        self.tp_size = tp_size
Cheng Wan's avatar
Cheng Wan committed
180
181
        self.moe_ep_rank = moe_ep_rank
        self.moe_ep_size = moe_ep_size
182
        self.dp_size = server_args.dp_size
183
184
        self.pp_rank = pp_rank
        self.pp_size = pp_size
185
        self.model_config = model_config
Zhang, Liangang's avatar
Zhang, Liangang committed
186
        self.dist_port = nccl_port
Lianmin Zheng's avatar
Lianmin Zheng committed
187
        self.server_args = server_args
188
        self.is_draft_worker = is_draft_worker
189
190
        self.is_generation = model_config.is_generation
        self.is_multimodal = model_config.is_multimodal
191
192
193
        self.is_multimodal_chunked_prefill_supported = (
            model_config.is_multimodal_chunked_prefill_supported
        )
194
195
196
        self.spec_algorithm = SpeculativeAlgorithm.from_string(
            server_args.speculative_algorithm
        )
197
        self.page_size = server_args.page_size
198
199
        self.req_to_token_pool = req_to_token_pool
        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
tarinkk's avatar
tarinkk committed
200
        self.is_hybrid = model_config.is_hybrid
Baizhou Zhang's avatar
Baizhou Zhang committed
201
        self.use_mla_backend = self.model_config.attention_arch == AttentionArch.MLA
Chang Su's avatar
Chang Su committed
202
        self.attention_chunk_size = model_config.attention_chunk_size
Ke Bao's avatar
Ke Bao committed
203

204
205
        self.forward_pass_id = 0

206
        # Model-specific adjustment
207
        self.model_specific_adjustment()
Shuo Yang's avatar
Shuo Yang committed
208

209
210
        if server_args.show_time_cost:
            enable_show_time_cost()
211
212

        # Global vars
213
        global_server_args_dict.update(
214
215
216
            {k: getattr(server_args, k) for k in GLOBAL_SERVER_ARGS_KEYS}
            | {
                # TODO it is indeed not a "server args"
217
                "use_mla_backend": self.use_mla_backend,
218
                "speculative_algorithm": self.spec_algorithm,
219
220
            }
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
221

222
        # CPU offload
223
224
        set_cpu_offload_max_bytes(int(server_args.cpu_offload_gb * 1024**3))

225
226
227
228
        # Init OpenMP threads binding for CPU
        if self.device == "cpu":
            self.init_threads_binding()

229
        # Get memory before model loading
230
        min_per_gpu_memory = self.init_torch_distributed()
231

232
        # Update deep gemm configure
233
234
        if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
            deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args)
235

Lianmin Zheng's avatar
Lianmin Zheng committed
236
        # If it is a draft model, tp_group can be different
237
238
        self.initialize(min_per_gpu_memory)

239
240
241
242
        # temporary cached values
        self.support_pp = (
            "pp_proxy_tensors" in inspect.signature(self.model.forward).parameters
        )
243
        self._model_update_group = {}
244

245
246
    def initialize(self, min_per_gpu_memory: float):
        server_args = self.server_args
247

248
249
250
251
        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
            enable=self.server_args.enable_memory_saver
        )

252
253
254
255
256
257
258
259
        if not self.is_draft_worker:
            set_global_expert_location_metadata(
                compute_initial_expert_location_metadata(server_args, self.model_config)
            )
            if self.tp_rank == 0 and get_bool_env_var(
                "SGLANG_LOG_EXPERT_LOCATION_METADATA"
            ):
                logger.info(
260
                    f"Initial expert_location_metadata: {get_global_expert_location_metadata()}"
261
262
263
264
265
266
267
268
269
270
                )

            set_global_expert_distribution_recorder(
                ExpertDistributionRecorder.init_new(
                    server_args,
                    get_global_expert_location_metadata(),
                    rank=self.tp_rank,
                )
            )

271
272
273
274
275
        self.eplb_manager = (
            EPLBManager(self)
            if self.server_args.enable_eplb and (not self.is_draft_worker)
            else None
        )
276
        self.expert_location_updater = ExpertLocationUpdater()
277

278
        # Load the model
279
        self.sampler = Sampler()
280
        self.load_model()
281

282
        # Check if the model is using hybrid SWA
Hanming Lu's avatar
Hanming Lu committed
283
284
285
286
287
288
289
290
291
        if (
            not self.server_args.disable_hybrid_swa_memory
            and self.sliding_window_size is not None
            and self.sliding_window_size > 0
        ):
            architectures = self.model_config.hf_config.architectures
            if architectures and not any("Llama4" in arch for arch in architectures):
                self.is_hybrid = self.model_config.is_hybrid = True

292
293
294
295
296
297
298
299
        # For MTP models like DeepSeek-V3 or GLM-4.5, the MTP layer(s) are used separately as draft
        # models for speculative decoding. In those cases, `num_nextn_predict_layers` is used to
        # determine the number of layers.
        model_has_mtp_layers = self.model_config.num_nextn_predict_layers is not None
        model_num_layers = (
            self.model_config.num_nextn_predict_layers
            if self.is_draft_worker and model_has_mtp_layers
            else self.model_config.num_hidden_layers
300
        )
301
302
        self.start_layer = getattr(self.model, "start_layer", 0)
        self.end_layer = getattr(self.model, "end_layer", model_num_layers)
303
        self.num_effective_layers = self.end_layer - self.start_layer
304
305
306
        assert (not model_has_mtp_layers) or (
            self.num_effective_layers == model_num_layers
        ), "PP is not compatible with MTP models."
307

308
        # Apply torchao quantization
309
310
311
312
313
314
        torchao_applied = getattr(self.model, "torchao_applied", False)
        # In layered loading, torchao may have been applied
        if not torchao_applied:
            apply_torchao_config_to_model(
                self.model, global_server_args_dict["torchao_config"]
            )
315

316
        # Apply torch TP if the model supports it
317
318
319
320
        supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
        if self.tp_size > 1 and supports_torch_tp:
            self.apply_torch_tp()

321
        # Init lora
322
        if server_args.enable_lora:
323
            self.init_lora_manager()
324
325

        # Init memory pool and attention backends
326
327
        self.init_memory_pool(
            min_per_gpu_memory,
328
            server_args.max_running_requests,
329
330
            server_args.max_total_tokens,
        )
Zhang, Liangang's avatar
Zhang, Liangang committed
331
332
333
334
335
        if self.device == "cuda":
            self.init_cublas()
            self.init_attention_backend()
            self.init_cuda_graphs()
        else:
336
            self.cuda_graph_runner = None
337
            self.cuda_graph_mem_usage = 0
Zhang, Liangang's avatar
Zhang, Liangang committed
338
            self.init_attention_backend()
339

James Liu's avatar
James Liu committed
340
341
        # auxiliary hidden capture mode. TODO: expose this to server args?
        if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
lukec's avatar
lukec committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
            # load draft config
            draft_model_config = ModelConfig.from_server_args(
                server_args,
                model_path=(server_args.speculative_draft_model_path),
                is_draft_model=True,
            )

            try:
                # get the aux layer from draft model config
                eagle_config = getattr(
                    draft_model_config.hf_config, "eagle_config", None
                )
                eagle_aux_hidden_state_layer_ids = eagle_config[
                    "eagle_aux_hidden_state_layer_ids"
                ]
            except:
                # if there is no aux layer, set to None
                eagle_aux_hidden_state_layer_ids = None

            self.model.set_eagle3_layers_to_capture(eagle_aux_hidden_state_layer_ids)
James Liu's avatar
James Liu committed
362

363
364
365
    def model_specific_adjustment(self):
        server_args = self.server_args

366
367
368
        if (
            server_args.attention_backend == "intel_amx"
            and server_args.device == "cpu"
369
            and not _is_cpu_amx_available
370
371
372
373
374
375
        ):
            logger.info(
                "The current platform does not support Intel AMX, will fallback to torch_native backend."
            )
            server_args.attention_backend = "torch_native"

376
        if server_args.attention_backend is None:
377
            """
Lianmin Zheng's avatar
Lianmin Zheng committed
378
379
            Auto select the fastest attention backend.

380
381
382
383
384
            1. Models with MHA Architecture (e.g: Llama, QWen)
                1.1 We will turn on FA3 on hopper unless user use spec decode with topk > 1 or page_size > 1.
                1.2 In other cases, we will use flashinfer if available, otherwise use triton.
            2. Models with MLA Architecture and using FA3
                2.1 We will use FA3 backend on hopper.
385
386
                2.2 We will use Flashinfer backend on blackwell.
                2.3 Otherwise, we will use triton backend.
387
388
            """

389
            if not self.use_mla_backend:
Lianmin Zheng's avatar
Lianmin Zheng committed
390
                # MHA architecture
391
                if (
392
                    is_hopper_with_cuda_12_3()
393
394
                    and is_no_spec_infer_or_topk_one(server_args)
                    and is_fa3_default_architecture(self.model_config.hf_config)
Zhiqiang Xie's avatar
Zhiqiang Xie committed
395
                    and (not server_args.enable_hierarchical_cache)
396
397
                ):
                    server_args.attention_backend = "fa3"
398
399
                elif _is_hip:
                    server_args.attention_backend = "aiter"
400
401
                elif _is_npu:
                    server_args.attention_backend = "ascend"
402
403
404
405
                else:
                    server_args.attention_backend = (
                        "flashinfer" if is_flashinfer_available() else "triton"
                    )
406
            else:
Lianmin Zheng's avatar
Lianmin Zheng committed
407
                # MLA architecture
Zhiqiang Xie's avatar
Zhiqiang Xie committed
408
409
410
                if is_hopper_with_cuda_12_3() and (
                    not server_args.enable_hierarchical_cache
                ):
411
                    server_args.attention_backend = "fa3"
412
413
                elif is_sm100_supported():
                    server_args.attention_backend = "flashinfer"
414
415
416
417
418
419
420
421
422
                elif _is_hip:
                    head_num = self.model_config.get_num_kv_heads(self.tp_size)
                    # TODO current aiter only support head number 16 or 128 head number
                    if (
                        head_num == 128 or head_num == 16
                    ) and self.spec_algorithm.is_none():
                        server_args.attention_backend = "aiter"
                    else:
                        server_args.attention_backend = "triton"
423
424
                elif _is_npu:
                    server_args.attention_backend = "ascend"
425
426
                else:
                    server_args.attention_backend = "triton"
427
            logger.info(
428
                f"Attention backend not explicitly specified. Use {server_args.attention_backend} backend by default."
429
            )
430
        elif self.use_mla_backend:
431
            if server_args.device != "cpu":
432
                if server_args.attention_backend in [
433
                    "aiter",
434
435
436
437
                    "flashinfer",
                    "fa3",
                    "triton",
                    "flashmla",
438
                    "cutlass_mla",
439
                    "trtllm_mla",
440
                    "ascend",
441
                ]:
442
443
444
                    logger.info(
                        f"MLA optimization is turned on. Use {server_args.attention_backend} backend."
                    )
445
                else:
446
447
448
449
                    raise ValueError(
                        f"Invalid attention backend for MLA: {server_args.attention_backend}"
                    )
            else:
450
451
452
453
                if server_args.attention_backend != "intel_amx":
                    raise ValueError(
                        "MLA optimization not supported on CPU except for intel_amx backend."
                    )
454

455
456
457
458
459
460
461
462
463
464
        if (
            server_args.attention_backend == "fa3"
            and server_args.kv_cache_dtype == "fp8_e5m2"
        ):
            logger.warning(
                "FlashAttention3 only supports fp8_e4m3 if using FP8; "
                "Setting attention backend to triton."
            )
            server_args.attention_backend = "triton"

465
        if server_args.enable_double_sparsity:
466
467
468
            logger.info(
                "Double sparsity optimization is turned on. Use triton backend without CUDA graph."
            )
469
470
471
472
473
474
475
476
477
            server_args.attention_backend = "triton"
            server_args.disable_cuda_graph = True
            if server_args.ds_heavy_channel_type is None:
                raise ValueError(
                    "Please specify the heavy channel type for double sparsity optimization."
                )
            self.init_double_sparsity_channel_config(server_args.ds_heavy_channel_type)

        if self.is_multimodal:
478
479
480
            if not self.is_multimodal_chunked_prefill_supported:
                server_args.chunked_prefill_size = -1
                logger.info(
481
                    f"Automatically turn off --chunked-prefill-size as it is not supported for "
482
483
                    f"{self.model_config.hf_config.model_type}"
                )
484

485
486
487
        if not self.use_mla_backend:
            server_args.disable_chunked_prefix_cache = True
        elif self.page_size > 1:
488
            logger.info("Disable chunked prefix cache when page size > 1.")
489
490
491
            server_args.disable_chunked_prefix_cache = True

        if not server_args.disable_chunked_prefix_cache:
492
            logger.info("Chunked prefix cache is turned on.")
493

kk's avatar
kk committed
494
495
496
497
        if server_args.attention_backend == "aiter":
            if self.model_config.context_len > 8192:
                self.mem_fraction_static *= 0.85

498
    def init_torch_distributed(self):
499
        logger.info("Init torch distributed begin.")
500

501
502
503
504
505
506
507
508
        try:
            torch.get_device_module(self.device).set_device(self.gpu_id)
        except Exception:
            logger.warning(
                f"Context: {self.device=} {self.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {self.tp_rank=} {self.tp_size=}"
            )
            raise

Zhang, Liangang's avatar
Zhang, Liangang committed
509
510
        if self.device == "cuda":
            backend = "nccl"
511
        elif self.device == "xpu":
512
            backend = "xccl"
513
514
        elif self.device == "hpu":
            backend = "hccl"
515
516
        elif self.device == "cpu":
            backend = "gloo"
517
518
        elif self.device == "npu":
            backend = "hccl"
519

520
        before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
521
        if not self.server_args.enable_p2p_check:
522
523
            monkey_patch_p2p_access_check()

524
        if self.server_args.dist_init_addr:
Zhang, Liangang's avatar
Zhang, Liangang committed
525
            dist_init_method = f"tcp://{self.server_args.dist_init_addr}"
526
        else:
Zhang, Liangang's avatar
Zhang, Liangang committed
527
            dist_init_method = f"tcp://127.0.0.1:{self.dist_port}"
528
        set_custom_all_reduce(not self.server_args.disable_custom_all_reduce)
529
        set_mscclpp_all_reduce(self.server_args.enable_mscclpp)
530
531

        if not self.is_draft_worker:
532
533
534
535
            if self.device == "cpu":
                if _is_cpu_amx_available:
                    # Bind OpenMP threads to CPU cores
                    torch.ops.sgl_kernel.init_cpu_threads_env(self.local_omp_cpuid)
536
537
538
539

                    # Set local size to hint SGLang to use shared memory based AllReduce
                    os.environ["LOCAL_SIZE"] = str(self.tp_size)
                    torch.ops.sgl_kernel.initialize(self.tp_size, self.tp_rank)
540
541
                else:
                    logger.warning(
542
                        "init_cpu_threads_env and shared memory based AllReduce is disabled since intel amx backend is not available"
543
544
                    )

Mick's avatar
Mick committed
545
            # Only initialize the distributed environment on the target model worker.
546
547
            init_distributed_environment(
                backend=backend,
548
549
                world_size=self.tp_size * self.pp_size,
                rank=self.tp_size * self.pp_rank + self.tp_rank,
550
551
                local_rank=self.gpu_id,
                distributed_init_method=dist_init_method,
552
                timeout=self.server_args.dist_timeout,
553
            )
554
555
556
            initialize_model_parallel(
                tensor_model_parallel_size=self.tp_size,
                pipeline_model_parallel_size=self.pp_size,
Cheng Wan's avatar
Cheng Wan committed
557
                expert_model_parallel_size=self.moe_ep_size,
558
                duplicate_tp_group=self.server_args.enable_pdmux,
559
            )
560
561
562
563
564
            initialize_dp_attention(
                enable_dp_attention=self.server_args.enable_dp_attention,
                tp_rank=self.tp_rank,
                tp_size=self.tp_size,
                dp_size=self.server_args.dp_size,
565
                moe_dense_tp_size=self.server_args.moe_dense_tp_size,
566
                pp_size=self.server_args.pp_size,
567
            )
568

569
        min_per_gpu_memory = get_available_gpu_memory(
570
571
572
573
            self.device,
            self.gpu_id,
            distributed=get_world_group().world_size > 1,
            cpu_group=get_world_group().cpu_group,
574
        )
575
        self.tp_group = get_tp_group()
576
        self.attention_tp_group = get_attention_tp_group()
577

578
        # Check memory for tensor parallelism
579
        local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id)
580
        if self.tp_size > 1 and not self.is_draft_worker:
581
            if min_per_gpu_memory < local_gpu_memory * 0.9:
582
583
584
585
586
587
588
589
590
591
                if get_bool_env_var("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK"):
                    logger.warning(
                        "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. "
                        f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
                    )
                else:
                    raise ValueError(
                        "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. "
                        f"{min_per_gpu_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
                    )
Lianmin Zheng's avatar
Lianmin Zheng committed
592

593
594
595
        logger.info(
            f"Init torch distributed ends. mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
        )
596
        return min_per_gpu_memory
597

Lianmin Zheng's avatar
Lianmin Zheng committed
598
    def load_model(self):
599
        before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
600
        logger.info(
Zhang, Liangang's avatar
Zhang, Liangang committed
601
            f"Load weight begin. avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
602
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
603
604

        # This can reduce thread conflicts and speed up weight loading.
605
606
        if self.device != "cpu":
            torch.set_num_threads(1)
Zhang, Liangang's avatar
Zhang, Liangang committed
607
608
        if self.device == "cuda":
            if torch.cuda.get_device_capability()[0] < 8:
609
610
611
                logger.info(
                    "Compute capability below sm80. Use float16 due to lack of bfloat16 support."
                )
Zhang, Liangang's avatar
Zhang, Liangang committed
612
                self.server_args.dtype = "float16"
613
                self.model_config.dtype = torch.float16
Zhang, Liangang's avatar
Zhang, Liangang committed
614
615
                if torch.cuda.get_device_capability()[1] < 5:
                    raise RuntimeError("SGLang only supports sm75 and above.")
Lianmin Zheng's avatar
Lianmin Zheng committed
616

617
618
        set_cuda_arch()

619
        # Prepare the model config
620
621
622
        self.load_config = LoadConfig(
            load_format=self.server_args.load_format,
            download_dir=self.server_args.download_dir,
623
            model_loader_extra_config=self.server_args.model_loader_extra_config,
624
        )
625
626
627
628
        if self.device == "cpu":
            self.model_config = adjust_config_with_unaligned_cpu_tp(
                self.model_config, self.load_config, self.tp_size
            )
629
630
        if self.server_args.load_format == "gguf":
            monkey_patch_vllm_gguf_config()
631
632

        # Load the model
633
634
        # Remove monkey_patch when linear.py quant remove dependencies with vllm
        monkey_patch_vllm_parallel_state()
635
636
        monkey_patch_isinstance_for_vllm_base_layer()

637
        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_WEIGHTS):
638
639
640
641
642
            self.model = get_model(
                model_config=self.model_config,
                load_config=self.load_config,
                device_config=DeviceConfig(self.device),
            )
643
        monkey_patch_vllm_parallel_state(reverse=True)
644
        monkey_patch_isinstance_for_vllm_base_layer(reverse=True)
645

bjmsong's avatar
bjmsong committed
646
647
648
649
650
651
        if self.server_args.kv_cache_dtype == "fp8_e4m3":
            if self.server_args.quantization_param_path is not None:
                if callable(getattr(self.model, "load_kv_cache_scales", None)):
                    self.model.load_kv_cache_scales(
                        self.server_args.quantization_param_path
                    )
652
653
654
655
                    logger.info(
                        "Loaded KV cache scaling factors from %s",
                        self.server_args.quantization_param_path,
                    )
bjmsong's avatar
bjmsong committed
656
657
658
659
660
661
662
663
664
665
666
667
668
                else:
                    raise RuntimeError(
                        "Using FP8 KV cache and scaling factors provided but "
                        "model %s does not support loading scaling factors.",
                        self.model.__class__,
                    )
            else:
                logger.warning(
                    "Using FP8 KV cache but no scaling factors "
                    "provided. Defaulting to scaling factors of 1.0. "
                    "This may lead to less accurate results!"
                )

669
        # Parse other args
Hanming Lu's avatar
Hanming Lu committed
670
671
672
673
674
        self.sliding_window_size = None
        if hasattr(self.model, "get_attention_sliding_window_size"):
            self.sliding_window_size = self.model.get_attention_sliding_window_size()
        elif self.model_config.attention_chunk_size is not None:
            self.sliding_window_size = self.model_config.attention_chunk_size
675
            logger.info(
Hanming Lu's avatar
Hanming Lu committed
676
677
678
                f"Setting sliding_window_size to be attention_chunk_size: {self.sliding_window_size}"
            )

679
        self.dtype = self.model_config.dtype
680

681
        after_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
682
        self.weight_load_mem_usage = before_avail_memory - after_avail_memory
683
        logger.info(
684
            f"Load weight end. "
685
            f"type={type(self.model).__name__}, "
Lianmin Zheng's avatar
Lianmin Zheng committed
686
            f"dtype={self.dtype}, "
687
            f"avail mem={after_avail_memory:.2f} GB, "
688
            f"mem usage={self.weight_load_mem_usage:.2f} GB."
689
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
690

691
692
693
694
695
696
697
698
699
700
701
702
        # Handle the case where some ranks do not finish loading.
        try:
            dist.monitored_barrier(
                group=get_tp_group().cpu_group,
                timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S),
                wait_all_ranks=True,
            )
        except RuntimeError:
            raise ValueError(
                f"TP rank {self.tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node."
            ) from None

703
    def update_expert_location(
704
705
706
        self,
        new_expert_location_metadata: ExpertLocationMetadata,
        update_layer_ids: List[int],
707
    ):
708
        self.expert_location_updater.update(
709
710
            self.model.routed_experts_weights_of_layer,
            new_expert_location_metadata,
711
            update_layer_ids=update_layer_ids,
712
713
714
715
            nnodes=self.server_args.nnodes,
            rank=self.tp_rank,
        )

716
717
718
719
    def update_weights_from_disk(
        self, model_path: str, load_format: str
    ) -> tuple[bool, str]:
        """Update engine weights in-place from the disk."""
720
        logger.info(
Chayenne's avatar
Chayenne committed
721
            f"Update engine weights online from disk begin. "
Zhang, Liangang's avatar
Zhang, Liangang committed
722
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
723
724
        )

Zhang, Liangang's avatar
Zhang, Liangang committed
725
        target_device = torch.device(self.device)
726
        self.model_config.model_path = model_path
727
728
        load_config = LoadConfig(load_format=load_format)

Lianmin Zheng's avatar
Lianmin Zheng committed
729
        # Only support DefaultModelLoader for now
730
731
        loader = get_model_loader(load_config)
        if not isinstance(loader, DefaultModelLoader):
Lianmin Zheng's avatar
Lianmin Zheng committed
732
733
            message = f"Failed to get model loader: {loader}."
            return False, message
734
735
736

        def get_weight_iter(config):
            iter = loader._get_weights_iterator(
737
                DefaultModelLoader.Source.init_new(config, self.model)
738
739
740
741
            )
            return iter

        def model_load_weights(model, iter):
742
            DefaultModelLoader.load_weights_and_postprocess(model, iter, target_device)
743
744
            return model

745
        with set_default_torch_dtype(self.model_config.dtype):
746
            try:
747
                iter = get_weight_iter(self.model_config)
748
            except Exception as e:
Lianmin Zheng's avatar
Lianmin Zheng committed
749
                message = f"Failed to get weights iterator: {e}."
750
751
752
753
                return False, message
            try:
                model = model_load_weights(self.model, iter)
            except Exception as e:
Lianmin Zheng's avatar
Lianmin Zheng committed
754
755
756
                message = (
                    f"Failed to update weights: {e}.\nRolling back to original weights."
                )
757
758
                del iter
                gc.collect()
759
                iter = get_weight_iter(self.model_config)
760
761
762
763
764
765
766
767
                self.model = model_load_weights(self.model, iter)
                return False, message

        self.model = model
        self.server_args.model_path = model_path
        self.server_args.load_format = load_format
        self.load_config = load_config

768
        logger.info("Update weights end.")
Lianmin Zheng's avatar
Lianmin Zheng committed
769
        return True, "Succeeded to update model weights."
770

771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
    def init_weights_update_group(
        self,
        master_address,
        master_port,
        rank_offset,
        world_size,
        group_name,
        backend="nccl",
    ):
        """Initialize the Torch process group for model parameter updates.

        `_model_update_group` is used in the RLHF workflow, where rank
        0 is the actor model in the training engine, and the other ranks are
        the inference engine, which is used for rollout.

        In the RLHF workflow, the training engine updates the model
        weights/parameters online, and broadcasts them to the inference
        engine through the `_model_update_group` process group.
        """
        assert (
            torch.distributed.is_initialized()
        ), "Default torch process group must be initialized"
        assert group_name != "", "Group name cannot be empty"

        rank = rank_offset + self.tp_rank

        logger.info(
            f"init custom process group: master_address={master_address}, master_port={master_port}, "
799
            f"rank_offset={rank_offset}, rank={rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
800
801
802
        )

        try:
803
            self._model_update_group[group_name] = init_custom_process_group(
804
805
806
807
808
809
810
811
812
813
814
815
                backend=backend,
                init_method=f"tcp://{master_address}:{master_port}",
                world_size=world_size,
                rank=rank,
                group_name=group_name,
            )
            return True, "Succeeded to initialize custom process group."
        except Exception as e:
            message = f"Failed to initialize custom process group: {e}."
            logger.error(message)
            return False, message

816
    def update_weights_from_distributed(self, names, dtypes, shapes, group_name):
817
818
819
820
821
822
823
824
825
826
        """
        Update specific parameter in the model weights online
        through `_model_update_group` process group.

        Args:
            name: the name of the parameter to be updated.
            dtype: the data type of the parameter to be updated.
            shape: the shape of the parameter to be updated.
        """

827
828
829
830
        assert group_name in self._model_update_group, (
            f"Group {group_name} not in {list(self._model_update_group.keys())}. "
            "Please call `init_weights_update_group` first."
        )
831
832

        try:
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
            weights = []
            handles = []
            for name, dtype, shape in zip(names, dtypes, shapes):
                target_dtype = (
                    dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
                )
                weight = torch.empty(shape, dtype=target_dtype, device=self.device)
                handles.append(
                    torch.distributed.broadcast(
                        weight,
                        src=0,
                        group=self._model_update_group[group_name],
                        async_op=True,
                    )
                )
                weights.append((name, weight))
            for handle in handles:
                handle.wait()

            self.model.load_weights(weights)
            return True, f"Succeeded to update parameter online."
854
855
856
857
858
859
860
861
862
863

        except Exception as e:
            error_msg = (
                f"Failed to update parameter online: {e}. "
                f"The full weights of the ModelRunner are partially updated. "
                f"Please discard the whole weights."
            )
            logger.error(error_msg)
            return False, error_msg

864
865
866
867
868
869
870
871
872
873
874
    def update_weights_from_tensor(
        self,
        named_tensors: List[Tuple[str, Union[torch.Tensor, "LocalSerializedTensor"]]],
        load_format: Optional[str] = None,
    ):
        named_tensors = [
            (name, _unwrap_tensor(tensor, tp_rank=self.tp_rank))
            for name, tensor in named_tensors
        ]
        if load_format == "direct":
            _model_load_weights_direct(self.model, named_tensors)
875
876
877
        elif load_format in self.server_args.custom_weight_loader:
            custom_loader = dynamic_import(load_format)
            custom_loader(self.model, named_tensors)
878
879
880
881
        elif load_format is None:
            self.model.load_weights(named_tensors)
        else:
            raise NotImplementedError(f"Unknown load_format={load_format}")
882
        return True, "Success"
883

884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
    def get_weights_by_name(
        self, name: str, truncate_size: int = 100
    ) -> Optional[torch.Tensor]:
        """Get the weights of the parameter by its name. Similar to `get_parameter` in Hugging Face.

        Only used for unit test with an unoptimized performance.
        For optimized performance, please use torch.save and torch.load.
        """
        # TODO: (chenyang) Add support for Qwen models.
        try:
            return self.model.get_weights_by_name(
                name, truncate_size, tp_size=self.tp_size
            )
        except Exception as e:
            logger.error(f"Error when getting parameter {name}: {e}")
            return None

901
902
903
904
905
906
907
    def init_lora_manager(self):
        self.lora_manager = LoRAManager(
            base_model=self.model,
            base_hf_config=self.model_config.hf_config,
            max_loras_per_batch=self.server_args.max_loras_per_batch,
            load_config=self.load_config,
            dtype=self.dtype,
908
            lora_backend=self.server_args.lora_backend,
909
910
            tp_size=self.tp_size,
            tp_rank=self.tp_rank,
911
912
            max_lora_rank=self.server_args.max_lora_rank,
            target_modules=self.server_args.lora_target_modules,
913
            lora_paths=self.server_args.lora_paths,
914
        )
915

916
    def load_lora_adapter(self, lora_ref: LoRARef):
917
918
919
        """Load a new lora adapter from disk or huggingface."""

        logger.info(
920
            f"LoRA adapter loading starts: {lora_ref}. "
921
922
923
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
        )

924
        result = self.lora_manager.load_lora_adapter(lora_ref)
925
926

        logger.info(
927
            f"LoRA adapter loading completes: {lora_ref}. "
928
929
930
931
932
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
        )

        return result

933
    def unload_lora_adapter(self, lora_ref: LoRARef):
934
935
936
        """Unload a lora adapter that was previously loaded during initialization or dynamic loading."""

        logger.info(
937
            f"LoRA adapter unloading starts: {lora_ref}. "
938
939
940
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
        )

941
        result = self.lora_manager.unload_lora_adapter(lora_ref)
942
943

        logger.info(
944
            f"LoRA adapter unloading completes: {lora_ref}. "
945
946
947
948
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
        )

        return result
949

950
    def profile_max_num_token(self, total_gpu_memory: int):
951
        available_gpu_memory = get_available_gpu_memory(
952
953
954
955
            self.device,
            self.gpu_id,
            distributed=get_world_group().world_size > 1,
            cpu_group=get_world_group().cpu_group,
956
        )
957
958
959
960
961
        if self.is_draft_worker:
            num_layers = getattr(
                self.model_config.hf_config,
                "num_nextn_predict_layers",
                self.num_effective_layers,
962
            )
963
964
965
        else:
            num_layers = self.num_effective_layers
        if self.use_mla_backend:
966
967
            # FIXME: pipeline parallelism is not compatible with mla backend
            assert self.pp_size == 1
968
969
            cell_size = (
                (self.model_config.kv_lora_rank + self.model_config.qk_rope_head_dim)
970
                * num_layers
971
                * torch._utils._element_size(self.kv_cache_dtype)
972
973
974
            )
        else:
            cell_size = (
975
                self.model_config.get_num_kv_heads(get_attention_tp_size())
976
                * self.model_config.head_dim
977
                * num_layers
978
                * 2
979
                * torch._utils._element_size(self.kv_cache_dtype)
980
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
981
982
983
        rest_memory = available_gpu_memory - total_gpu_memory * (
            1 - self.mem_fraction_static
        )
984
        max_num_token = int(rest_memory * (1 << 30) // cell_size)
Lianmin Zheng's avatar
Lianmin Zheng committed
985
986
        return max_num_token

tarinkk's avatar
tarinkk committed
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
    def set_num_token_hybrid(self):
        if (
            "Llama4ForConditionalGeneration"
            in self.model_config.hf_config.architectures
        ):
            temp_ratio = (
                (1 - self.is_hybrid)
                + self.is_hybrid
                * self.attention_chunk_size
                / self.model_config.context_len
            )
            self.swa_max_total_num_tokens = (
                4 * self.max_total_num_tokens * temp_ratio // (3 * temp_ratio + 1)
            )
            self.full_max_total_num_tokens = (
                4 * self.max_total_num_tokens
                - 12 * self.max_total_num_tokens * temp_ratio // (3 * temp_ratio + 1)
            )
            self.swa_max_total_num_tokens = int(
                self.swa_max_total_num_tokens
                // self.server_args.page_size
                * self.server_args.page_size
            )
            self.full_max_total_num_tokens = int(
                self.full_max_total_num_tokens
                // self.server_args.page_size
                * self.server_args.page_size
            )
            self.max_total_num_tokens = self.full_max_total_num_tokens
        else:
Hanming Lu's avatar
Hanming Lu committed
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
            assert self.sliding_window_size is not None and self.sliding_window_size > 0
            full_attention_layer_ids = []
            swa_attention_layer_ids = []

            try:
                layers = self.model.model.layers
            except:
                try:
                    layers = self.model.language_model.model.layers
                except:
1027
1028
1029
1030
1031
                    try:
                        layers = self.model.language_model.layers
                    except:
                        self.is_hybrid = False
                        return
Hanming Lu's avatar
Hanming Lu committed
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066

            for layer in layers:
                if (
                    layer.self_attn.attn.sliding_window_size is None
                    or layer.self_attn.attn.sliding_window_size == -1
                ):
                    full_attention_layer_ids.append(layer.layer_id)
                else:
                    swa_attention_layer_ids.append(layer.layer_id)
            self.model_config.swa_attention_layer_ids = swa_attention_layer_ids
            self.model_config.full_attention_layer_ids = full_attention_layer_ids

            # Algorithm:
            # Existing max_total_num_tokens is per layer and assume all layers have the same number of tokens.
            # - Find total # of tokens available across layers.
            # - Calculate full_max_total_num_tokens and swa_max_total_num_tokens based on the given swa_full_tokens_ratio.
            total_tokens = (
                self.max_total_num_tokens * self.model_config.num_hidden_layers
            )
            full_layers_num = len(full_attention_layer_ids)
            swa_layers_num = len(swa_attention_layer_ids)
            swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio

            # Solve the equations:
            # 1. swa_max_total_num_tokens * swa_layers_num + full_max_total_num_tokens * full_layers_num == total_tokens
            # 2. full_max_total_num_tokens * swa_full_tokens_ratio == swa_max_total_num_tokens
            denominator = swa_full_tokens_ratio * swa_layers_num + full_layers_num
            self.full_max_total_num_tokens = int(total_tokens / denominator)
            self.swa_max_total_num_tokens = int(
                self.full_max_total_num_tokens * swa_full_tokens_ratio
            )
            self.max_total_num_tokens = self.full_max_total_num_tokens

            logger.info(
                f"Use Sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}"
tarinkk's avatar
tarinkk committed
1067
1068
            )

1069
    def init_memory_pool(
1070
1071
        self,
        total_gpu_memory: int,
1072
1073
        max_num_reqs: Optional[int] = None,
        max_total_tokens: Optional[int] = None,
1074
    ):
1075
1076
1077
        if self.server_args.kv_cache_dtype == "auto":
            self.kv_cache_dtype = self.dtype
        elif self.server_args.kv_cache_dtype == "fp8_e5m2":
1078
            if _is_hip:  # Using natively supported format
HAI's avatar
HAI committed
1079
1080
1081
                self.kv_cache_dtype = torch.float8_e5m2fnuz
            else:
                self.kv_cache_dtype = torch.float8_e5m2
bjmsong's avatar
bjmsong committed
1082
        elif self.server_args.kv_cache_dtype == "fp8_e4m3":
1083
1084
1085
            if _is_hip:  # Using natively supported format
                self.kv_cache_dtype = torch.float8_e4m3fnuz
            else:
bjmsong's avatar
bjmsong committed
1086
                self.kv_cache_dtype = torch.float8_e4m3fn
1087
1088
1089
1090
1091
        else:
            raise ValueError(
                f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
            )

1092
        self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104

        if max_num_reqs is None:
            max_num_reqs = min(
                max(
                    int(
                        self.max_total_num_tokens / self.model_config.context_len * 512
                    ),
                    2048,
                ),
                4096,
            )

1105
1106
1107
        if SGLANG_CI_SMALL_KV_SIZE:
            self.max_total_num_tokens = int(SGLANG_CI_SMALL_KV_SIZE)

1108
1109
1110
        if not self.spec_algorithm.is_none():
            if self.is_draft_worker:
                self.max_total_num_tokens = self.server_args.draft_runner_cache_size
1111
                max_num_reqs = self.server_args.max_num_reqs
1112
            else:
1113
1114
                # We are sharing the `token_to_kv_pool`, and both verify and draft tokens
                # can be concurrently allocated, so we should give a headroom for it.
1115
1116
                self.server_args.draft_runner_cache_size = (
                    self.max_total_num_tokens
1117
1118
1119
1120
1121
1122
1123
                    # draft
                    + max_num_reqs
                    * self.server_args.speculative_num_steps
                    * self.server_args.speculative_eagle_topk
                    # verify
                    + max_num_reqs * self.server_args.speculative_num_draft_tokens
                    # buffer
1124
1125
                    + 100
                )
1126
1127
1128
1129
                # Target worker and draft worker shares the same indices for the
                # token_to_kv_pool, so we should make sure to match max_total_num_tokens.
                self.max_total_num_tokens = self.server_args.draft_runner_cache_size
                self.server_args.max_num_reqs = max_num_reqs
1130

1131
1132
        if max_total_tokens is not None:
            if max_total_tokens > self.max_total_num_tokens:
1133
                logging.warning(
1134
1135
1136
1137
1138
                    f"max_total_tokens={max_total_tokens} is larger than the profiled value "
                    f"{self.max_total_num_tokens}. "
                    f"Use the profiled value instead."
                )
            self.max_total_num_tokens = min(self.max_total_num_tokens, max_total_tokens)
1139

1140
1141
1142
1143
1144
        self.max_total_num_tokens = (
            self.max_total_num_tokens
            // self.server_args.page_size
            * self.server_args.page_size
        )
tarinkk's avatar
tarinkk committed
1145
1146
1147
1148
        # create token size for hybrid cache
        if self.is_hybrid:
            self.set_num_token_hybrid()

1149
        if self.max_total_num_tokens <= 0:
1150
            raise RuntimeError(
1151
                "Not enough memory. Please try to increase --mem-fraction-static."
1152
            )
1153

1154
        if self.req_to_token_pool is None:
Byron Hsu's avatar
Byron Hsu committed
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
            if self.server_args.disaggregation_mode == "decode":
                from sglang.srt.disaggregation.decode import DecodeReqToTokenPool

                # subscribe memory for pre-allocated requests
                # if max_num_reqs <= 32, we pre-allocate 2x requests
                pre_alloc_size = max_num_reqs * 2 if max_num_reqs <= 32 else 0
                self.req_to_token_pool = DecodeReqToTokenPool(
                    size=max_num_reqs,
                    max_context_len=self.model_config.context_len + 4,
                    device=self.device,
                    enable_memory_saver=self.server_args.enable_memory_saver,
                    pre_alloc_size=pre_alloc_size,
                )
            else:
                self.req_to_token_pool = ReqToTokenPool(
                    size=max_num_reqs,
                    max_context_len=self.model_config.context_len + 4,
                    device=self.device,
                    enable_memory_saver=self.server_args.enable_memory_saver,
                )
1175
1176
1177
1178
        else:
            # Draft worker shares req_to_token_pool with the target worker.
            assert self.is_draft_worker

1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
        if self.server_args.attention_backend == "ascend" and not self.use_mla_backend:
            self.token_to_kv_pool = AscendTokenToKVPool(
                self.max_total_num_tokens,
                page_size=self.page_size,
                dtype=self.kv_cache_dtype,
                head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()),
                head_dim=self.model_config.head_dim,
                layer_num=self.model_config.num_hidden_layers,
                device=self.device,
                enable_memory_saver=self.server_args.enable_memory_saver,
            )
        elif self.server_args.attention_backend == "ascend" and self.use_mla_backend:
            self.token_to_kv_pool = AscendMLAPagedTokenToKVPool(
                self.max_total_num_tokens,
                page_size=self.page_size,
                dtype=self.kv_cache_dtype,
                kv_lora_rank=self.model_config.kv_lora_rank,
                qk_rope_head_dim=self.model_config.qk_rope_head_dim,
1197
                layer_num=self.num_effective_layers,
1198
1199
1200
1201
1202
1203
                device=self.device,
                enable_memory_saver=self.server_args.enable_memory_saver,
                start_layer=self.start_layer,
                end_layer=self.end_layer,
            )
        elif self.use_mla_backend:
1204
1205
            self.token_to_kv_pool = MLATokenToKVPool(
                self.max_total_num_tokens,
Lianmin Zheng's avatar
Lianmin Zheng committed
1206
                page_size=self.page_size,
1207
                dtype=self.kv_cache_dtype,
1208
1209
                kv_lora_rank=self.model_config.kv_lora_rank,
                qk_rope_head_dim=self.model_config.qk_rope_head_dim,
1210
                layer_num=self.num_effective_layers,
Zhang, Liangang's avatar
Zhang, Liangang committed
1211
                device=self.device,
1212
                enable_memory_saver=self.server_args.enable_memory_saver,
1213
1214
                start_layer=self.start_layer,
                end_layer=self.end_layer,
1215
            )
Shuo Yang's avatar
Shuo Yang committed
1216
1217
1218
        elif self.server_args.enable_double_sparsity:
            self.token_to_kv_pool = DoubleSparseTokenToKVPool(
                self.max_total_num_tokens,
Lianmin Zheng's avatar
Lianmin Zheng committed
1219
                page_size=self.page_size,
Shuo Yang's avatar
Shuo Yang committed
1220
                dtype=self.kv_cache_dtype,
1221
                head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()),
Shuo Yang's avatar
Shuo Yang committed
1222
                head_dim=self.model_config.head_dim,
1223
                layer_num=self.num_effective_layers,
Shuo Yang's avatar
Shuo Yang committed
1224
1225
                device=self.device,
                heavy_channel_num=self.server_args.ds_heavy_channel_num,
1226
                enable_memory_saver=self.server_args.enable_memory_saver,
1227
1228
                start_layer=self.start_layer,
                end_layer=self.end_layer,
Shuo Yang's avatar
Shuo Yang committed
1229
            )
1230
        else:
tarinkk's avatar
tarinkk committed
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
            if self.is_hybrid:
                self.token_to_kv_pool = SWAKVPool(
                    size=self.full_max_total_num_tokens,
                    size_swa=self.swa_max_total_num_tokens,
                    dtype=self.kv_cache_dtype,
                    head_num=self.model_config.get_num_kv_heads(
                        get_attention_tp_size()
                    ),
                    head_dim=self.model_config.head_dim,
                    swa_attention_layer_ids=self.model_config.swa_attention_layer_ids,
                    full_attention_layer_ids=self.model_config.full_attention_layer_ids,
                    enable_kvcache_transpose=False,
                    device=self.device,
                )
            else:
                self.token_to_kv_pool = MHATokenToKVPool(
Lianmin Zheng's avatar
Lianmin Zheng committed
1247
                    self.max_total_num_tokens,
tarinkk's avatar
tarinkk committed
1248
                    page_size=self.page_size,
Lianmin Zheng's avatar
Lianmin Zheng committed
1249
                    dtype=self.kv_cache_dtype,
tarinkk's avatar
tarinkk committed
1250
1251
1252
1253
1254
                    head_num=self.model_config.get_num_kv_heads(
                        get_attention_tp_size()
                    ),
                    head_dim=self.model_config.head_dim,
                    layer_num=self.num_effective_layers,
Lianmin Zheng's avatar
Lianmin Zheng committed
1255
                    device=self.device,
tarinkk's avatar
tarinkk committed
1256
1257
1258
                    enable_memory_saver=self.server_args.enable_memory_saver,
                    start_layer=self.start_layer,
                    end_layer=self.end_layer,
Lianmin Zheng's avatar
Lianmin Zheng committed
1259
                )
tarinkk's avatar
tarinkk committed
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277

        if self.token_to_kv_pool_allocator is None:
            if self.page_size == 1:
                if self.is_hybrid:
                    self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
                        self.full_max_total_num_tokens,
                        self.swa_max_total_num_tokens,
                        dtype=self.kv_cache_dtype,
                        device=self.device,
                        kvcache=self.token_to_kv_pool,
                    )
                else:
                    self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
                        self.max_total_num_tokens,
                        dtype=self.kv_cache_dtype,
                        device=self.device,
                        kvcache=self.token_to_kv_pool,
                    )
Lianmin Zheng's avatar
Lianmin Zheng committed
1278
            else:
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
                if _is_npu:
                    self.token_to_kv_pool_allocator = AscendPagedTokenToKVPoolAllocator(
                        self.max_total_num_tokens,
                        page_size=self.page_size,
                        dtype=self.kv_cache_dtype,
                        device=self.device,
                        kvcache=self.token_to_kv_pool,
                    )
                else:
                    self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
                        self.max_total_num_tokens,
                        page_size=self.page_size,
                        dtype=self.kv_cache_dtype,
                        device=self.device,
                        kvcache=self.token_to_kv_pool,
                    )
1295
1296
1297
        else:
            assert self.is_draft_worker

1298
        logger.info(
1299
            f"Memory pool end. "
Zhang, Liangang's avatar
Zhang, Liangang committed
1300
            f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
1301
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
1302

Lianmin Zheng's avatar
Lianmin Zheng committed
1303
1304
1305
1306
1307
1308
1309
1310
1311
    def init_cublas(self):
        """We need to run a small matmul to init cublas. Otherwise, it will raise some errors later."""
        dtype = torch.float16
        device = "cuda"
        a = torch.ones((16, 16), dtype=dtype, device=device)
        b = torch.ones((16, 16), dtype=dtype, device=device)
        c = a @ b
        return c

1312
1313
    def init_attention_backend(self):
        """Init attention kernel backend."""
1314
        if self.server_args.enable_two_batch_overlap and not self.is_draft_worker:
1315
1316
1317
1318
1319
            self.attn_backend = TboAttnBackend.init_new(self._get_attention_backend)
        else:
            self.attn_backend = self._get_attention_backend()

    def _get_attention_backend(self):
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
        """Init attention kernel backend."""
        self.decode_attention_backend_str = (
            self.server_args.decode_attention_backend
            if self.server_args.decode_attention_backend
            else self.server_args.attention_backend
        )
        self.prefill_attention_backend_str = (
            self.server_args.prefill_attention_backend
            if self.server_args.prefill_attention_backend
            else self.server_args.attention_backend
        )
        if self.decode_attention_backend_str != self.prefill_attention_backend_str:
            assert (
                self.server_args.speculative_algorithm is None
            ), "Currently HybridAttentionBackend does not support speculative decoding."
            from sglang.srt.layers.attention.hybrid_attn_backend import (
                HybridAttnBackend,
            )

            attn_backend = HybridAttnBackend(
                decode_backend=self._get_attention_backend_from_str(
                    self.decode_attention_backend_str
                ),
                prefill_backend=self._get_attention_backend_from_str(
                    self.prefill_attention_backend_str
                ),
            )
            logger.info(
                f"Using hybrid attention backend for decode and prefill: "
                f"decode_backend={self.decode_attention_backend_str}, "
                f"prefill_backend={self.prefill_attention_backend_str}."
            )
            logger.warning(
                f"Warning: Attention backend specified by --attention-backend or default backend might be overridden."
                f"The feature of hybrid attention backend is experimental and unstable. Please raise an issue if you encounter any problem."
            )
        else:
            attn_backend = self._get_attention_backend_from_str(
                self.server_args.attention_backend
            )

        global_server_args_dict.update(
            {
                "decode_attention_backend": self.decode_attention_backend_str,
                "prefill_attention_backend": self.prefill_attention_backend_str,
            }
        )
        return attn_backend

    def _get_attention_backend_from_str(self, backend_str: str):
        if backend_str == "flashinfer":
1371
1372
1373
1374
            if not self.use_mla_backend:
                from sglang.srt.layers.attention.flashinfer_backend import (
                    FlashInferAttnBackend,
                )
1375

1376
1377
                # Init streams
                if self.server_args.speculative_algorithm == "EAGLE":
1378
1379
1380
1381
1382
                    if (
                        not hasattr(self, "plan_stream_for_flashinfer")
                        or not self.plan_stream_for_flashinfer
                    ):
                        self.plan_stream_for_flashinfer = torch.cuda.Stream()
1383
                return FlashInferAttnBackend(self)
1384
1385
1386
1387
1388
            else:
                from sglang.srt.layers.attention.flashinfer_mla_backend import (
                    FlashInferMLAAttnBackend,
                )

1389
                return FlashInferMLAAttnBackend(self)
1390
        elif backend_str == "aiter":
1391
1392
            from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend

1393
            return AiterAttnBackend(self)
1394
        elif backend_str == "ascend":
1395
1396
1397
            from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend

            return AscendAttnBackend(self)
1398
        elif backend_str == "triton":
1399
1400
1401
1402
1403
            assert not self.model_config.is_encoder_decoder, (
                "Cross attention is not supported in the triton attention backend. "
                "Please use `--attention-backend flashinfer`."
            )
            if self.server_args.enable_double_sparsity:
1404
1405
1406
1407
                from sglang.srt.layers.attention.double_sparsity_backend import (
                    DoubleSparseAttnBackend,
                )

1408
                return DoubleSparseAttnBackend(self)
1409
            else:
1410
1411
                from sglang.srt.layers.attention.triton_backend import TritonAttnBackend

1412
                return TritonAttnBackend(self)
1413
        elif backend_str == "torch_native":
1414
1415
1416
1417
            from sglang.srt.layers.attention.torch_native_backend import (
                TorchNativeAttnBackend,
            )

1418
            return TorchNativeAttnBackend(self)
1419
        elif backend_str == "flashmla":
lukec's avatar
lukec committed
1420
1421
            from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend

1422
            return FlashMLABackend(self)
1423
        elif backend_str == "fa3":
1424
1425
1426
1427
            assert (
                torch.cuda.get_device_capability()[0] == 8 and not self.use_mla_backend
            ) or torch.cuda.get_device_capability()[0] == 9, (
                "FlashAttention v3 Backend requires SM>=80 and SM<=90. "
1428
1429
1430
1431
1432
1433
                "Please use `--attention-backend flashinfer`."
            )
            from sglang.srt.layers.attention.flashattention_backend import (
                FlashAttentionBackend,
            )

1434
            return FlashAttentionBackend(self)
1435
        elif backend_str == "cutlass_mla":
1436
1437
1438
1439
            from sglang.srt.layers.attention.cutlass_mla_backend import (
                CutlassMLABackend,
            )

1440
            return CutlassMLABackend(self)
1441
1442
1443
1444
1445
1446
        elif self.server_args.attention_backend == "trtllm_mla":
            if not self.use_mla_backend:
                raise ValueError("trtllm_mla backend can only be used with MLA models.")
            from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend

            return TRTLLMMLABackend(self)
1447
1448
1449
1450
1451
1452
1453
        elif self.server_args.attention_backend == "intel_amx":
            from sglang.srt.layers.attention.intel_amx_backend import (
                IntelAMXAttnBackend,
            )

            logger.info(f"Intel AMX attention backend is enabled.")
            return IntelAMXAttnBackend(self)
1454
        else:
1455
            raise ValueError(f"Invalid attention backend: {backend_str}")
1456

Shuo Yang's avatar
Shuo Yang committed
1457
1458
1459
1460
1461
1462
1463
    def init_double_sparsity_channel_config(self, selected_channel):
        selected_channel = "." + selected_channel + "_proj"
        self.sorted_channels = []
        # load channel config
        with open(self.server_args.ds_channel_config_path, "r") as f:
            channel_config = json.load(f)

1464
        for i in range(self.start_layer, self.end_layer):
Shuo Yang's avatar
Shuo Yang committed
1465
1466
1467
1468
1469
1470
1471
1472
1473
            key = "model.layers." + str(i) + ".self_attn" + selected_channel
            self.sorted_channels.append(
                torch.tensor(channel_config[key])[
                    :, : self.server_args.ds_heavy_channel_num
                ]
                .contiguous()
                .cuda()
            )

1474
    def init_cuda_graphs(self):
1475
        """Capture cuda graphs."""
1476
        self.cuda_graph_runner = None
1477
        self.cuda_graph_mem_usage = 0
1478

1479
        if not self.is_generation:
1480
            # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
1481
1482
            return

1483
1484
        if self.server_args.disable_cuda_graph:
            return
1485

1486
        tic = time.perf_counter()
1487
        before_mem = get_available_gpu_memory(self.device, self.gpu_id)
1488
        logger.info(
1489
            f"Capture cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
1490
        )
1491
        self.cuda_graph_runner = CudaGraphRunner(self)
1492
        after_mem = get_available_gpu_memory(self.device, self.gpu_id)
1493
        self.cuda_graph_mem_usage = before_mem - after_mem
1494
        logger.info(
1495
            f"Capture cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
1496
            f"mem usage={self.cuda_graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
1497
        )
1498

1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
    def init_threads_binding(self):
        omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
        if omp_cpuids == "all":
            cpu_ids_by_node = get_cpu_ids_by_node()
            n_numa_node = len(cpu_ids_by_node)

            assert self.tp_size <= n_numa_node, (
                f"SGLANG_CPU_OMP_THREADS_BIND is not set, in this case, "
                f"tp_size {self.tp_size} should be smaller than or equal to number of numa node on the machine {n_numa_node}. "
                f"If you need tp_size to be larger than number of numa node, please set the CPU cores for each tp rank via SGLANG_CPU_OMP_THREADS_BIND explicitly. "
                f"For example, on a machine with 2 numa nodes, where core 0-31 are on numa node 0 and core 32-63 are on numa node 1, "
                f"it is suggested to use -tp 2 and bind tp rank 0 to core 0-31 and tp rank 1 to core 32-63. "
                f"This is the default behavior if SGLANG_CPU_OMP_THREADS_BIND is not set and it is the same as setting SGLANG_CPU_OMP_THREADS_BIND=0-31|32-63. "
                f"If you do need tp_size to be larger than the number of numa nodes, you could set SGLANG_CPU_OMP_THREADS_BIND explicitly for example SGLANG_CPU_OMP_THREADS_BIND=0-15|16-31|32-47|48-63 and run with -tp 4. "
                f"If you don't want each tp rank to use all the cores on one numa node, you could set for example SGLANG_CPU_OMP_THREADS_BIND=0-15|32-47 and run with -tp 2."
            )
            if self.tp_size < n_numa_node:
                logger.warning(
                    f"Detected the current machine has {n_numa_node} numa nodes available, but tp_size is set to {self.tp_size}, so only {self.tp_size} numa nodes are used."
                )
            self.local_omp_cpuid = cpu_ids_by_node[self.tp_rank]
        else:
            self.local_omp_cpuid = omp_cpuids.split("|")[self.tp_rank]

1523
    def apply_torch_tp(self):
1524
        logger.info(f"Enabling torch tensor parallelism on {self.tp_size} devices.")
1525
1526
1527
1528
1529
        from sglang.srt.model_parallel import tensor_parallel

        device_mesh = torch.distributed.init_device_mesh(self.device, (self.tp_size,))
        tensor_parallel(self.model, device_mesh)

1530
    def forward_decode(
Cheng Wan's avatar
Cheng Wan committed
1531
1532
1533
1534
        self,
        forward_batch: ForwardBatch,
        skip_attn_backend_init: bool = False,
        pp_proxy_tensors=None,
1535
    ) -> LogitsProcessorOutput:
Cheng Wan's avatar
Cheng Wan committed
1536
1537
        if not skip_attn_backend_init:
            self.attn_backend.init_forward_metadata(forward_batch)
1538
1539
1540
1541
        # FIXME: add pp_proxy_tensors arg to all models
        kwargs = {}
        if self.support_pp:
            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
1542
        return self.model.forward(
1543
1544
1545
1546
            forward_batch.input_ids,
            forward_batch.positions,
            forward_batch,
            **kwargs,
Lianmin Zheng's avatar
Lianmin Zheng committed
1547
1548
        )

1549
    def forward_extend(
1550
1551
1552
1553
1554
        self,
        forward_batch: ForwardBatch,
        skip_attn_backend_init: bool = False,
        pp_proxy_tensors=None,
    ) -> LogitsProcessorOutput:
1555
1556
1557
        if not skip_attn_backend_init:
            self.attn_backend.init_forward_metadata(forward_batch)

1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
        kwargs = {}
        if self.support_pp:
            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
        if forward_batch.input_embeds is not None:
            kwargs["input_embeds"] = forward_batch.input_embeds.bfloat16()
        if not self.is_generation:
            kwargs["get_embedding"] = True
        return self.model.forward(
            forward_batch.input_ids,
            forward_batch.positions,
            forward_batch,
            **kwargs,
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
1571

1572
1573
1574
1575
1576
1577
    def forward_idle(
        self, forward_batch: ForwardBatch, pp_proxy_tensors=None
    ) -> LogitsProcessorOutput:
        kwargs = {}
        if self.support_pp:
            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
Ke Bao's avatar
Ke Bao committed
1578
        return self.model.forward(
1579
1580
1581
1582
            forward_batch.input_ids,
            forward_batch.positions,
            forward_batch,
            **kwargs,
Ke Bao's avatar
Ke Bao committed
1583
1584
        )

1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
    def forward_split_prefill(
        self,
        forward_batch: ForwardBatch,
        reinit_attn_backend: bool = False,
        forward_count: int = 1,
    ) -> LogitsProcessorOutput:
        if forward_batch.split_index == 0 or reinit_attn_backend:
            self.attn_backend.init_forward_metadata(forward_batch)
        next_split_index = min(
            forward_batch.split_index + forward_count,
            self.model_config.num_hidden_layers,
        )
        ret = self.model.forward_split_prefill(
            forward_batch.input_ids,
            forward_batch.positions,
            forward_batch,
            (forward_batch.split_index, next_split_index),
        )
        forward_batch.split_index = next_split_index
        return ret

1606
    def forward(
1607
1608
1609
1610
        self,
        forward_batch: ForwardBatch,
        skip_attn_backend_init: bool = False,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
1611
1612
        reinit_attn_backend: bool = False,
        split_forward_count: int = 1,
1613
1614
1615
1616
1617
1618
1619
    ) -> Tuple[Union[LogitsProcessorOutput, PPProxyTensors], bool]:
        self.forward_pass_id += 1

        with get_global_expert_distribution_recorder().with_forward_pass(
            self.forward_pass_id,
            forward_batch,
        ):
1620
            output = self._forward_raw(
1621
1622
1623
1624
1625
                forward_batch,
                skip_attn_backend_init,
                pp_proxy_tensors,
                reinit_attn_backend,
                split_forward_count,
1626
1627
            )

1628
        if self.eplb_manager is not None:
1629
            self.eplb_manager.on_forward_pass_end()
1630
1631
1632

        return output

1633
1634
1635
1636
1637
    def _forward_raw(
        self,
        forward_batch: ForwardBatch,
        skip_attn_backend_init: bool,
        pp_proxy_tensors: Optional[PPProxyTensors],
1638
1639
        reinit_attn_backend: bool = False,
        split_forward_count: int = 1,
1640
    ) -> Tuple[Union[LogitsProcessorOutput, PPProxyTensors], bool]:
1641
        can_run_cuda_graph = bool(
1642
1643
1644
            forward_batch.forward_mode.is_cuda_graph()
            and self.cuda_graph_runner
            and self.cuda_graph_runner.can_run(forward_batch)
1645
1646
        )
        if can_run_cuda_graph:
1647
            ret = self.cuda_graph_runner.replay(
1648
1649
1650
                forward_batch,
                skip_attn_backend_init=skip_attn_backend_init,
                pp_proxy_tensors=pp_proxy_tensors,
1651
            )
Cheng Wan's avatar
Cheng Wan committed
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
            return ret, can_run_cuda_graph

        # For MLP sync
        if forward_batch.global_num_tokens_cpu is not None:
            forward_batch.prepare_mlp_sync_batch(self)

        if forward_batch.forward_mode.is_decode():
            ret = self.forward_decode(
                forward_batch,
                skip_attn_backend_init=skip_attn_backend_init,
                pp_proxy_tensors=pp_proxy_tensors,
            )
1664
        elif forward_batch.forward_mode.is_extend():
1665
            ret = self.forward_extend(
1666
1667
1668
                forward_batch,
                skip_attn_backend_init=skip_attn_backend_init,
                pp_proxy_tensors=pp_proxy_tensors,
1669
            )
1670
1671
1672
1673
1674
1675
        elif forward_batch.forward_mode.is_split_prefill():
            ret = self.forward_split_prefill(
                forward_batch,
                reinit_attn_backend=reinit_attn_backend,
                forward_count=split_forward_count,
            )
Ke Bao's avatar
Ke Bao committed
1676
        elif forward_batch.forward_mode.is_idle():
1677
            ret = self.forward_idle(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
Lianmin Zheng's avatar
Lianmin Zheng committed
1678
        else:
1679
            raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode}")
1680

Cheng Wan's avatar
Cheng Wan committed
1681
1682
1683
        if forward_batch.global_num_tokens_cpu is not None:
            forward_batch.post_forward_mlp_sync_batch(ret)

1684
1685
        return ret, can_run_cuda_graph

1686
1687
1688
    def _preprocess_logits(
        self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo
    ):
1689
        # Apply logit bias
1690
1691
1692
1693
1694
1695
1696
1697
        if sampling_info.sampling_info_done:
            # Overlap mode: the function update_regex_vocab_mask was executed
            # in process_batch_result of the last batch.
            if sampling_info.grammars:
                sampling_info.sampling_info_done.wait()
        else:
            # Normal mode: Put CPU-heavy tasks here. They will be overlapped with the forward pass.
            sampling_info.update_regex_vocab_mask()
1698
1699
        sampling_info.apply_logits_bias(logits_output.next_token_logits)

1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
    def sample(
        self,
        logits_output: LogitsProcessorOutput,
        forward_batch: ForwardBatch,
    ) -> torch.Tensor:
        """Sample and compute logprobs and update logits_output.

        Args:
            logits_output: The logits output from the model forward
            forward_batch: The forward batch that generates logits_output

        Returns:
            A list of next_token_ids
        """
        # For duplex models with multiple output streams.
        if isinstance(logits_output, tuple):
            return torch.stack(
                [self.sample(values, forward_batch) for values in logits_output],
                axis=-1,
            )
1720

1721
1722
        self._preprocess_logits(logits_output, forward_batch.sampling_info)

1723
1724
1725
        # Sample the next tokens
        next_token_ids = self.sampler(
            logits_output,
1726
            forward_batch.sampling_info,
1727
1728
            forward_batch.return_logprob,
            forward_batch.top_logprobs_nums,
1729
            forward_batch.token_ids_logprobs,
1730
        )
1731
1732
        return next_token_ids

Yineng Zhang's avatar
Yineng Zhang committed
1733
1734
1735
1736
    @property
    def model_is_mrope(self) -> bool:
        """Detect if the model has "mrope" rope_scaling type.
        mrope requires keep "rope_deltas" between prompt and decoding phases."""
1737
        rope_scaling = getattr(self.model_config.hf_text_config, "rope_scaling", {})
Yineng Zhang's avatar
Yineng Zhang committed
1738
1739
        if rope_scaling is None:
            return False
1740
1741
        is_mrope_enabled = "mrope_section" in rope_scaling
        return is_mrope_enabled
1742

1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
    def save_remote_model(self, url: str):
        from sglang.srt.model_loader.loader import RemoteModelLoader

        logger.info(f"Saving model to {url}")
        RemoteModelLoader.save_model(self.model, self.model_config.model_path, url)

    def save_sharded_model(
        self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None
    ):
        from sglang.srt.model_loader.loader import ShardedStateLoader

        logger.info(
            f"Save sharded model to {path} with pattern {pattern} and max_size {max_size}"
        )
        ShardedStateLoader.save_model(self.model, path, pattern, max_size)

1759
1760
1761
1762
1763
1764
1765
1766
1767

def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]):
    params_dict = dict(model.named_parameters())
    for name, tensor in named_tensors:
        default_weight_loader(params_dict[name], tensor)


def _unwrap_tensor(tensor, tp_rank):
    if isinstance(tensor, LocalSerializedTensor):
1768
1769
1770
        monkey_patch_torch_reductions()
        tensor = tensor.get(tp_rank)
    return tensor.to(torch.cuda.current_device())
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781


@dataclass
class LocalSerializedTensor:
    """torch.Tensor that gets serialized by MultiprocessingSerializer (which only serializes a pointer and not the data).
    The i-th element in the list corresponds to i-th rank's GPU."""

    values: List[bytes]

    def get(self, rank: int):
        return MultiprocessingSerializer.deserialize(self.values[rank])