sglang 0.3.0__py3-none-any.whl → 0.3.1.post1__py3-none-any.whl

sglang/srt/model_executor/model_runner.py

@@ -25,12 +25,6 @@ from typing import Optional, Tuple, Type
 
 import torch
 import torch.nn as nn
-from flashinfer import (
-    BatchDecodeWithPagedKVCacheWrapper,
-    BatchPrefillWithPagedKVCacheWrapper,
-    BatchPrefillWithRaggedKVCacheWrapper,
-)
-from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 from vllm.config import DeviceConfig, LoadConfig
 from vllm.config import ModelConfig as VllmModelConfig
 from vllm.distributed import (
@@ -43,32 +37,34 @@ from vllm.distributed.parallel_state import in_the_same_node_as
 from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import ModelRegistry
 
-from sglang.global_config import global_config
+from sglang.srt.configs.model_config import AttentionArch, ModelConfig
+from sglang.srt.layers.attention_backend import FlashInferAttnBackend, TritonAttnBackend
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.layers.sampler import SampleOutput
+from sglang.srt.layers.sampler import Sampler
+from sglang.srt.lora.lora_manager import LoRAManager
 from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
     MLATokenToKVPool,
     ReqToTokenPool,
 )
-from sglang.srt.model_config import AttentionArch, ModelConfig
-from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
+from sglang.srt.model_executor.forward_batch_info import InputMetadata
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     get_available_gpu_memory,
     is_generation_model,
-    is_llama3_405b_fp8_head_16,
     is_multimodal_model,
     monkey_patch_vllm_dummy_weight_loader,
     monkey_patch_vllm_p2p_access_check,
-    monkey_patch_vllm_qvk_linear_loader,
 )
 
 logger = logging.getLogger(__name__)
 
 
 class ModelRunner:
+    """ModelRunner runs the forward passes of the models."""
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -92,13 +88,15 @@ class ModelRunner:
         )
         global_server_args_dict.update(
             {
-                "disable_flashinfer": server_args.disable_flashinfer,
-                "disable_flashinfer_sampling": server_args.disable_flashinfer_sampling,
+                "attention_backend": server_args.attention_backend,
+                "sampling_backend": server_args.sampling_backend,
                 "triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32,
                 "enable_mla": server_args.enable_mla,
+                "torchao_config": server_args.torchao_config,
             }
         )
 
+        # Model-specific adjustment
         if self.is_multimodal_model:
             logger.info(
                 "Automatically turn off --chunked-prefill-size and adjust --mem-fraction-static for multimodal models."
@@ -106,15 +104,19 @@ class ModelRunner:
             server_args.chunked_prefill_size = None
             server_args.mem_fraction_static *= 0.95
 
+        # Init componnets
         min_per_gpu_memory = self.init_torch_distributed()
+        self.sampler = Sampler()
         self.load_model()
+        if server_args.lora_paths is not None:
+            self.init_lora_manager()
         self.init_memory_pool(
             min_per_gpu_memory,
-            server_args.max_num_reqs,
+            server_args.max_running_requests,
             server_args.max_total_tokens,
         )
         self.init_cublas()
-        self.init_flashinfer()
+        self.init_attention_backend()
         self.init_cuda_graphs()
 
     def init_torch_distributed(self):
@@ -162,10 +164,13 @@ class ModelRunner:
         return min_per_gpu_memory
 
     def load_model(self):
-        torch.set_num_threads(1)
         logger.info(
             f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
+
+        # This can reduce thread conflicts and speed up weight loading.
+        torch.set_num_threads(1)
+
         if torch.cuda.get_device_capability()[0] < 8:
             logger.info(
                 "Compute capability below sm80. Use float16 due to lack of bfloat16 support."
@@ -174,6 +179,7 @@ class ModelRunner:
             if torch.cuda.get_device_capability()[1] < 5:
                 raise RuntimeError("SGLang only supports sm75 and above.")
 
+        # Prepare the vllm model config
         monkey_patch_vllm_dummy_weight_loader()
         self.device_config = DeviceConfig()
         self.load_config = LoadConfig(load_format=self.server_args.load_format)
@@ -184,23 +190,16 @@ class ModelRunner:
             tokenizer_mode=None,
             trust_remote_code=self.server_args.trust_remote_code,
             dtype=self.server_args.dtype,
-            seed=42,
+            seed=self.server_args.random_seed,
             skip_tokenizer_init=True,
         )
-
-        # A temporary hack to fix the num_heads for meta-llama/Meta-Llama-3.1-405B-FP8 checkpoints
-        # Drop this after Sept, 2024.
-        if is_llama3_405b_fp8_head_16(self.model_config) and self.tp_size <= 8:
-            self.model_config.hf_config.num_key_value_heads = 8
-            self.vllm_model_config.hf_config.num_key_value_heads = 8
-            monkey_patch_vllm_qvk_linear_loader()
-
-        self.dtype = self.vllm_model_config.dtype
         if self.model_config.model_override_args is not None:
             self.vllm_model_config.hf_config.update(
                 self.model_config.model_override_args
             )
+        self.dtype = self.vllm_model_config.dtype
 
+        # Load the model
         self.model = get_model(
             model_config=self.vllm_model_config,
             load_config=self.load_config,
@@ -251,20 +250,20 @@ class ModelRunner:
                 tokenizer_mode=None,
                 trust_remote_code=self.server_args.trust_remote_code,
                 dtype=self.server_args.dtype,
-                seed=42,
+                seed=self.server_args.random_seed,
                 skip_tokenizer_init=True,
             )
         except Exception as e:
-            logger.error(f"Failed to load model config: {e}")
-            return False, "Failed to update model weights"
+            message = f"Failed to load model config: {e}."
+            return False, message
 
         load_config = LoadConfig(load_format=load_format)
 
         # Only support vllm DefaultModelLoader for now
         loader = get_model_loader(load_config)
         if not isinstance(loader, DefaultModelLoader):
-            logger.error("Failed to get weights iterator: Unsupported loader")
-            return False, "Failed to update model weights"
+            message = f"Failed to get model loader: {loader}."
+            return False, message
 
         def get_weight_iter(config):
             iter = loader._get_weights_iterator(
@@ -289,14 +288,14 @@ class ModelRunner:
             try:
                 iter = get_weight_iter(vllm_model_config)
             except Exception as e:
-                message = f"Failed to get weights iterator: {e}"
-                logger.error(message)
+                message = f"Failed to get weights iterator: {e}."
                 return False, message
             try:
                 model = model_load_weights(self.model, iter)
             except Exception as e:
-                message = f"Failed to update weights: {e}. \n Rolling back to original weights"
-                logger.error(message)
+                message = (
+                    f"Failed to update weights: {e}.\nRolling back to original weights."
+                )
                 del iter
                 gc.collect()
                 iter = get_weight_iter(self.vllm_model_config)
@@ -311,7 +310,18 @@ class ModelRunner:
         self.model_config.path = model_path
 
         logger.info("Update weights end.")
-        return True, "Succeeded to update model weights"
+        return True, "Succeeded to update model weights."
+
+    def init_lora_manager(self):
+        self.lora_manager = LoRAManager(
+            base_model=self.model,
+            lora_paths=self.server_args.lora_paths,
+            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,
+        )
+        logger.info("LoRA manager ready.")
 
     def profile_max_num_token(self, total_gpu_memory: int):
         available_gpu_memory = get_available_gpu_memory(
@@ -343,8 +353,8 @@ class ModelRunner:
     def init_memory_pool(
         self,
         total_gpu_memory: int,
-        max_num_reqs: int = None,
-        max_total_tokens: int = None,
+        max_num_reqs: Optional[int] = None,
+        max_total_tokens: Optional[int] = None,
     ):
         if self.server_args.kv_cache_dtype == "auto":
             self.kv_cache_dtype = self.dtype
@@ -378,7 +388,7 @@ class ModelRunner:
                     ),
                     2048,
                 ),
-                5120,
+                4096,
             )
 
         self.req_to_token_pool = ReqToTokenPool(
@@ -396,9 +406,6 @@ class ModelRunner:
                 qk_rope_head_dim=self.model_config.qk_rope_head_dim,
                 layer_num=self.model_config.num_hidden_layers,
             )
-            logger.info("using MLA Triton implementaion, flashinfer is disabled")
-            # FIXME: temporarily only Triton MLA is supported
-            self.server_args.disable_flashinfer = True
         else:
             self.token_to_kv_pool = MHATokenToKVPool(
                 self.max_total_num_tokens,
@@ -421,118 +428,46 @@ class ModelRunner:
         c = a @ b
         return c
 
-    def init_flashinfer(self):
-        """Init flashinfer attention kernel wrappers."""
-        if self.server_args.disable_flashinfer:
-            assert (
-                self.sliding_window_size is None
-            ), "turn on flashinfer to support window attention"
-            self.flashinfer_prefill_wrapper_ragged = None
-            self.flashinfer_prefill_wrapper_paged = None
-            self.flashinfer_decode_wrapper = None
-            return
-
-        if not _grouped_size_compiled_for_decode_kernels(
-            self.model_config.num_attention_heads // self.tp_size,
-            self.model_config.get_num_kv_heads(self.tp_size),
-        ):
-            use_tensor_cores = True
-        else:
-            use_tensor_cores = False
-
-        if self.sliding_window_size is None:
-            self.flashinfer_workspace_buffer = torch.empty(
-                global_config.flashinfer_workspace_size,
-                dtype=torch.uint8,
-                device="cuda",
-            )
-            self.flashinfer_prefill_wrapper_ragged = (
-                BatchPrefillWithRaggedKVCacheWrapper(
-                    self.flashinfer_workspace_buffer, "NHD"
-                )
-            )
-            self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
-                self.flashinfer_workspace_buffer, "NHD"
-            )
-            self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-                self.flashinfer_workspace_buffer,
-                "NHD",
-                use_tensor_cores=use_tensor_cores,
+    def init_attention_backend(self):
+        """Init attention kernel backend."""
+        if self.server_args.attention_backend == "flashinfer":
+            self.attn_backend = FlashInferAttnBackend(self)
+        elif self.server_args.attention_backend == "triton":
+            assert self.sliding_window_size is None, (
+                "Window attention is not supported in the triton attention backend. "
+                "Please use `--attention-backend flashinfer`."
             )
+            self.attn_backend = TritonAttnBackend(self)
         else:
-            self.flashinfer_workspace_buffer = torch.empty(
-                global_config.flashinfer_workspace_size,
-                dtype=torch.uint8,
-                device="cuda",
+            raise ValueError(
+                f"Invalid attention backend: {self.server_args.attention_backend}"
             )
-            self.flashinfer_prefill_wrapper_ragged = None
-            self.flashinfer_prefill_wrapper_paged = []
-            self.flashinfer_decode_wrapper = []
-            for i in range(2):
-                self.flashinfer_prefill_wrapper_paged.append(
-                    BatchPrefillWithPagedKVCacheWrapper(
-                        self.flashinfer_workspace_buffer, "NHD"
-                    )
-                )
-                self.flashinfer_decode_wrapper.append(
-                    BatchDecodeWithPagedKVCacheWrapper(
-                        self.flashinfer_workspace_buffer,
-                        "NHD",
-                        use_tensor_cores=use_tensor_cores,
-                    )
-                )
 
     def init_cuda_graphs(self):
         """Capture cuda graphs."""
+        from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
+
+        self.cuda_graph_runner = None
+
         if not self.is_generation:
             # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
             return
 
-        from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
-
-        if self.server_args.disable_cuda_graph or self.server_args.disable_flashinfer:
-            self.cuda_graph_runner = None
+        if self.server_args.disable_cuda_graph:
             return
 
         logger.info("Capture cuda graph begin. This can take up to several minutes.")
-
-        if self.server_args.disable_cuda_graph_padding:
-            batch_size_list = list(range(1, 32)) + [64, 128]
-        else:
-            batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 21)]
-
-        self.cuda_graph_runner = CudaGraphRunner(
-            self,
-            max_batch_size_to_capture=max(batch_size_list),
-            use_torch_compile=self.server_args.enable_torch_compile,
-            disable_padding=self.server_args.disable_cuda_graph_padding,
-        )
-        try:
-            self.cuda_graph_runner.capture(batch_size_list)
-        except RuntimeError as e:
-            raise Exception(
-                f"Capture cuda graph failed: {e}\n"
-                "Possible solutions:\n"
-                "1. disable cuda graph by --disable-cuda-graph\n"
-                "2. set --mem-fraction-static to a smaller value\n"
-                "3. disable torch compile by not using --enable-torch-compile\n"
-                "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
-            )
+        self.cuda_graph_runner = CudaGraphRunner(self)
 
     @torch.inference_mode()
     def forward_decode(self, batch: ScheduleBatch):
-        if (
-            self.cuda_graph_runner
-            and self.cuda_graph_runner.can_run(len(batch.reqs))
-            and batch.sampling_info.can_run_in_cuda_graph()
-        ):
+        if self.server_args.lora_paths is not None:
+            self.lora_manager.prepare_lora_batch(batch)
+
+        if self.cuda_graph_runner and self.cuda_graph_runner.can_run(len(batch.reqs)):
             return self.cuda_graph_runner.replay(batch)
 
-        input_metadata = InputMetadata.from_schedule_batch(
-            self,
-            batch,
-            ForwardMode.DECODE,
-        )
+        input_metadata = InputMetadata.from_schedule_batch(self, batch)
 
         return self.model.forward(
             batch.input_ids, input_metadata.positions, input_metadata
@@ -540,11 +475,10 @@ class ModelRunner:
 
     @torch.inference_mode()
     def forward_extend(self, batch: ScheduleBatch):
-        input_metadata = InputMetadata.from_schedule_batch(
-            self,
-            batch,
-            forward_mode=ForwardMode.EXTEND,
-        )
+        input_metadata = InputMetadata.from_schedule_batch(self, batch)
+        if self.server_args.lora_paths is not None:
+            self.lora_manager.prepare_lora_batch(batch, input_metadata.extend_seq_lens)
+
         if self.is_generation:
             return self.model.forward(
                 batch.input_ids, input_metadata.positions, input_metadata
@@ -560,11 +494,7 @@ class ModelRunner:
 
     @torch.inference_mode()
     def forward_extend_multi_modal(self, batch: ScheduleBatch):
-        input_metadata = InputMetadata.from_schedule_batch(
-            self,
-            batch,
-            forward_mode=ForwardMode.EXTEND,
-        )
+        input_metadata = InputMetadata.from_schedule_batch(self, batch)
         return self.model.forward(
             batch.input_ids,
             input_metadata.positions,
@@ -574,17 +504,56 @@ class ModelRunner:
             input_metadata.image_offsets,
         )
 
-    def forward(
-        self, batch: ScheduleBatch, forward_mode: ForwardMode
-    ) -> Tuple[SampleOutput, LogitsProcessorOutput]:
-        if self.is_multimodal_model and forward_mode == ForwardMode.EXTEND:
+    def forward(self, batch: ScheduleBatch) -> Tuple[LogitsProcessorOutput]:
+        assert batch.forward_mode is not None
+
+        if self.is_multimodal_model and batch.forward_mode.is_extend():
             return self.forward_extend_multi_modal(batch)
-        elif forward_mode == ForwardMode.DECODE:
+        elif batch.forward_mode.is_decode():
             return self.forward_decode(batch)
-        elif forward_mode == ForwardMode.EXTEND:
+        elif batch.forward_mode.is_extend():
             return self.forward_extend(batch)
         else:
-            raise ValueError(f"Invaid forward mode: {forward_mode}")
+            raise ValueError(f"Invaid forward mode: {batch.forward_mode}")
+
+    def _apply_logits_bias(
+        self, logits: torch.Tensor, sampling_info: SamplingBatchInfo
+    ):
+        # Apply logit_bias
+        if sampling_info.logit_bias is not None:
+            logits.add_(sampling_info.logit_bias)
+
+        # min-token, presence, frequency
+        if sampling_info.linear_penalties is not None:
+            logits += sampling_info.linear_penalties
+
+        # repetition
+        if sampling_info.scaling_penalties is not None:
+            logits = torch.where(
+                logits > 0,
+                logits / sampling_info.scaling_penalties,
+                logits * sampling_info.scaling_penalties,
+            )
+
+        # Apply regex vocab_mask
+        if sampling_info.vocab_mask is not None:
+            logits = logits.masked_fill(sampling_info.vocab_mask, float("-inf"))
+
+        return logits
+
+    def sample(
+        self, logits_output: LogitsProcessorOutput, batch: ScheduleBatch
+    ) -> torch.Tensor:
+        # Put CPU-heavy tasks here. They will be overlapped with the forward pass.
+        batch.sampling_info.update_regex_vocab_mask(batch)
+        batch.sampling_info.update_penalties()
+        logits = self._apply_logits_bias(
+            logits_output.next_token_logits, batch.sampling_info
+        )
+
+        # Sample the next tokens.
+        next_token_ids = self.sampler(logits, batch.sampling_info)
+        return next_token_ids
 
 
 @lru_cache()