sglang 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

sglang/srt/model_executor/model_runner.py

@@ -15,13 +15,13 @@ limitations under the License.
 
 """ModelRunner runs the forward passes of the models."""
 
+import gc
 import importlib
 import importlib.resources
 import logging
 import pkgutil
-import warnings
 from functools import lru_cache
-from typing import Optional, Type
+from typing import Optional, Tuple, Type
 
 import torch
 import torch.nn as nn
@@ -37,10 +37,15 @@ from vllm.distributed import (
     get_tp_group,
     init_distributed_environment,
     initialize_model_parallel,
+    set_custom_all_reduce,
 )
+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.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.sampler import SampleOutput
 from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
@@ -53,7 +58,7 @@ from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     get_available_gpu_memory,
     is_generation_model,
-    is_llama3_405b_fp8,
+    is_llama3_405b_fp8_head_16,
     is_multimodal_model,
     monkey_patch_vllm_dummy_weight_loader,
     monkey_patch_vllm_p2p_access_check,
@@ -87,22 +92,35 @@ class ModelRunner:
             {
                 "disable_flashinfer": server_args.disable_flashinfer,
                 "disable_flashinfer_sampling": server_args.disable_flashinfer_sampling,
-                "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
+                "triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32,
                 "enable_mla": server_args.enable_mla,
             }
         )
 
+        min_per_gpu_memory = self.init_torch_distributed()
+        self.load_model()
+        self.init_memory_pool(
+            min_per_gpu_memory,
+            server_args.max_num_reqs,
+            server_args.max_total_tokens,
+        )
+        self.init_cublas()
+        self.init_flashinfer()
+        self.init_cuda_graphs()
+
+    def init_torch_distributed(self):
         # Init torch distributed
         torch.cuda.set_device(self.gpu_id)
-        logger.info(f"[gpu={self.gpu_id}] Init nccl begin.")
+        logger.info("Init nccl begin.")
 
-        if not server_args.enable_p2p_check:
+        if not self.server_args.enable_p2p_check:
             monkey_patch_vllm_p2p_access_check(self.gpu_id)
 
-        if server_args.nccl_init_addr:
-            nccl_init_method = f"tcp://{server_args.nccl_init_addr}"
+        if self.server_args.nccl_init_addr:
+            nccl_init_method = f"tcp://{self.server_args.nccl_init_addr}"
         else:
             nccl_init_method = f"tcp://127.0.0.1:{self.nccl_port}"
+        set_custom_all_reduce(not self.server_args.disable_custom_all_reduce)
         init_distributed_environment(
             backend="nccl",
             world_size=self.tp_size,
@@ -111,43 +129,43 @@ class ModelRunner:
             distributed_init_method=nccl_init_method,
         )
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
-        self.tp_group = get_tp_group()
-        total_gpu_memory = get_available_gpu_memory(
+        min_per_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
+        self.tp_group = get_tp_group()
 
+        # Currently, there is a bug with mulit-node tensor parallelsim + padded cuda graph,
+        # so we disable padding in cuda graph.
+        if not all(in_the_same_node_as(self.tp_group.cpu_group, source_rank=0)):
+            self.server_args.disable_cuda_graph_padding = True
+            logger.info(
+                "Setting disable_cuda_graph_padding to True because of multi-node tensor parallelism."
+            )
+
+        # Check memory for tensor parallelism
         if self.tp_size > 1:
-            total_local_gpu_memory = get_available_gpu_memory(self.gpu_id)
-            if total_local_gpu_memory < total_gpu_memory * 0.9:
+            local_gpu_memory = get_available_gpu_memory(self.gpu_id)
+            if min_per_gpu_memory < local_gpu_memory * 0.9:
                 raise ValueError(
                     "The memory capacity is unbalanced. Some GPUs may be occupied by other processes."
                 )
 
-        # Load the model and create memory pool
-        self.load_model()
-        self.init_memory_pool(
-            total_gpu_memory,
-            server_args.max_num_reqs,
-            server_args.max_total_tokens,
-        )
-        self.init_cublas()
-        self.init_flashinfer()
-
-        if self.is_generation:
-            # FIXME Currently, cuda graph only capture decode steps, which only exists in causal models
-            # Capture cuda graphs
-            self.init_cuda_graphs()
+        return min_per_gpu_memory
 
     def load_model(self):
         logger.info(
-            f"[gpu={self.gpu_id}] Load weight begin. "
-            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+            f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
+        if torch.cuda.get_device_capability()[0] < 8:
+            logger.info(
+                "Compute capability below sm80. Use float16 due to lack of bfloat16 support."
+            )
+            self.server_args.dtype = "float16"
 
         monkey_patch_vllm_dummy_weight_loader()
-        device_config = DeviceConfig()
-        load_config = LoadConfig(load_format=self.server_args.load_format)
-        vllm_model_config = VllmModelConfig(
+        self.device_config = DeviceConfig()
+        self.load_config = LoadConfig(load_format=self.server_args.load_format)
+        self.vllm_model_config = VllmModelConfig(
             model=self.server_args.model_path,
             quantization=self.server_args.quantization,
             tokenizer=None,
@@ -158,47 +176,132 @@ class ModelRunner:
             skip_tokenizer_init=True,
         )
 
-        if is_llama3_405b_fp8(self.model_config) and self.tp_size <= 8:
-            # A temporary hack to fix the num_heads for meta-llama/Meta-Llama-3.1-405B-FP8 checkpoints
+        # 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
-            vllm_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 = vllm_model_config.dtype
+        self.dtype = self.vllm_model_config.dtype
         if self.model_config.model_overide_args is not None:
-            vllm_model_config.hf_config.update(self.model_config.model_overide_args)
-
-        if (
-            self.server_args.efficient_weight_load
-            and "llama" in self.server_args.model_path.lower()
-            and self.server_args.quantization == "fp8"
-        ):
-            from sglang.srt.model_loader.model_loader import get_model
-        else:
-            from vllm.model_executor.model_loader import get_model
+            self.vllm_model_config.hf_config.update(
+                self.model_config.model_overide_args
+            )
 
         self.model = get_model(
-            model_config=vllm_model_config,
-            device_config=device_config,
-            load_config=load_config,
-            lora_config=None,
-            multimodal_config=None,
+            model_config=self.vllm_model_config,
+            load_config=self.load_config,
+            device_config=self.device_config,
             parallel_config=None,
             scheduler_config=None,
+            lora_config=None,
             cache_config=None,
         )
+        self.sliding_window_size = (
+            self.model.get_attention_sliding_window_size()
+            if hasattr(self.model, "get_attention_sliding_window_size")
+            else None
+        )
         self.is_generation = is_generation_model(
-            self.model_config.hf_config.architectures
+            self.model_config.hf_config.architectures, self.server_args.is_embedding
         )
 
         logger.info(
-            f"[gpu={self.gpu_id}] Load weight end. "
+            f"Load weight end. "
             f"type={type(self.model).__name__}, "
             f"dtype={self.dtype}, "
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
 
-    def profile_max_num_token(self, total_gpu_memory):
+    def update_weights(self, model_path: str, load_format: str):
+        """Update weights in-place."""
+        from vllm.model_executor.model_loader.loader import (
+            DefaultModelLoader,
+            device_loading_context,
+            get_model_loader,
+        )
+        from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+
+        logger.info(
+            f"Update weights begin. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+        target_device = torch.device(self.device_config.device)
+
+        try:
+            # TODO: Use a better method to check this
+            vllm_model_config = VllmModelConfig(
+                model=model_path,
+                quantization=self.server_args.quantization,
+                tokenizer=None,
+                tokenizer_mode=None,
+                trust_remote_code=self.server_args.trust_remote_code,
+                dtype=self.server_args.dtype,
+                seed=42,
+                skip_tokenizer_init=True,
+            )
+        except Exception as e:
+            logger.error(f"Failed to load model config: {e}")
+            return False, "Failed to update model weights"
+
+        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"
+
+        def get_weight_iter(config):
+            iter = loader._get_weights_iterator(
+                config.model,
+                config.revision,
+                fall_back_to_pt=getattr(
+                    self.model, "fall_back_to_pt_during_load", True
+                ),
+            )
+            return iter
+
+        def model_load_weights(model, iter):
+            model.load_weights(iter)
+            for _, module in self.model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is not None:
+                    with device_loading_context(module, target_device):
+                        quant_method.process_weights_after_loading(module)
+            return model
+
+        with set_default_torch_dtype(vllm_model_config.dtype):
+            try:
+                iter = get_weight_iter(vllm_model_config)
+            except Exception as e:
+                message = f"Failed to get weights iterator: {e}"
+                logger.error(message)
+                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)
+                del iter
+                gc.collect()
+                iter = get_weight_iter(self.vllm_model_config)
+                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.vllm_model_config = vllm_model_config
+        self.load_config = load_config
+        self.model_config.path = model_path
+
+        logger.info("Update weights end.")
+        return True, "Succeeded to update model weights"
+
+    def profile_max_num_token(self, total_gpu_memory: int):
         available_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
@@ -209,7 +312,7 @@ class ModelRunner:
             cell_size = (
                 (self.model_config.kv_lora_rank + self.model_config.qk_rope_head_dim)
                 * self.model_config.num_hidden_layers
-                * torch._utils._element_size(self.dtype)
+                * torch._utils._element_size(self.kv_cache_dtype)
             )
         else:
             cell_size = (
@@ -217,7 +320,7 @@ class ModelRunner:
                 * self.model_config.head_dim
                 * self.model_config.num_hidden_layers
                 * 2
-                * torch._utils._element_size(self.dtype)
+                * torch._utils._element_size(self.kv_cache_dtype)
             )
         rest_memory = available_gpu_memory - total_gpu_memory * (
             1 - self.mem_fraction_static
@@ -226,12 +329,30 @@ class ModelRunner:
         return max_num_token
 
     def init_memory_pool(
-        self, total_gpu_memory, max_num_reqs=None, max_total_tokens=None
+        self,
+        total_gpu_memory: int,
+        max_num_reqs: int = None,
+        max_total_tokens: int = None,
     ):
+        if self.server_args.kv_cache_dtype == "auto":
+            self.kv_cache_dtype = self.dtype
+        elif self.server_args.kv_cache_dtype == "fp8_e5m2":
+            if self.server_args.disable_flashinfer or self.server_args.enable_mla:
+                logger.warning(
+                    "FP8 KV cache is not supported for Triton kernel now, using auto kv cache dtype"
+                )
+                self.kv_cache_dtype = self.dtype
+            else:
+                self.kv_cache_dtype = torch.float8_e5m2
+        else:
+            raise ValueError(
+                f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
+            )
+
         self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
         if max_total_tokens is not None:
             if max_total_tokens > self.max_total_num_tokens:
-                warnings.warn(
+                logging.warning(
                     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."
@@ -264,7 +385,7 @@ class ModelRunner:
         ):
             self.token_to_kv_pool = MLATokenToKVPool(
                 self.max_total_num_tokens,
-                dtype=self.dtype,
+                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,
                 layer_num=self.model_config.num_hidden_layers,
@@ -275,13 +396,13 @@ class ModelRunner:
         else:
             self.token_to_kv_pool = MHATokenToKVPool(
                 self.max_total_num_tokens,
-                dtype=self.dtype,
+                dtype=self.kv_cache_dtype,
                 head_num=self.model_config.get_num_kv_heads(self.tp_size),
                 head_dim=self.model_config.head_dim,
                 layer_num=self.model_config.num_hidden_layers,
             )
         logger.info(
-            f"[gpu={self.gpu_id}] Memory pool end. "
+            f"Memory pool end. "
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
 
@@ -295,7 +416,11 @@ class ModelRunner:
         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
@@ -309,36 +434,72 @@ class ModelRunner:
         else:
             use_tensor_cores = False
 
-        self.flashinfer_workspace_buffers = torch.empty(
-            2, global_config.flashinfer_workspace_size, dtype=torch.uint8, device="cuda"
-        )
-        self.flashinfer_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0], "NHD"
-        )
-        self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[1], "NHD"
-        )
-        self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0],
-            "NHD",
-            use_tensor_cores=use_tensor_cores,
-        )
+        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,
+            )
+        else:
+            self.flashinfer_workspace_buffer = torch.empty(
+                global_config.flashinfer_workspace_size,
+                dtype=torch.uint8,
+                device="cuda",
+            )
+            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."""
+        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
             return
 
-        logger.info(
-            f"[gpu={self.gpu_id}] Capture cuda graph begin. This can take up to several minutes."
-        )
-        batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 17)]
+        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)
@@ -354,11 +515,17 @@ class ModelRunner:
 
     @torch.inference_mode()
     def forward_decode(self, batch: ScheduleBatch):
-        if self.cuda_graph_runner and self.cuda_graph_runner.can_run(len(batch.reqs)):
+        if (
+            self.cuda_graph_runner
+            and self.cuda_graph_runner.can_run(len(batch.reqs))
+            and not batch.sampling_info.has_bias()
+        ):
             return self.cuda_graph_runner.replay(batch)
 
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, ForwardMode.DECODE
+            self,
+            batch,
+            ForwardMode.DECODE,
         )
 
         return self.model.forward(
@@ -368,16 +535,29 @@ class ModelRunner:
     @torch.inference_mode()
     def forward_extend(self, batch: ScheduleBatch):
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, forward_mode=ForwardMode.EXTEND
-        )
-        return self.model.forward(
-            batch.input_ids, input_metadata.positions, input_metadata
+            self,
+            batch,
+            forward_mode=ForwardMode.EXTEND,
         )
+        if self.is_generation:
+            return self.model.forward(
+                batch.input_ids, input_metadata.positions, input_metadata
+            )
+        else:
+            # Only embedding models have get_embedding parameter
+            return self.model.forward(
+                batch.input_ids,
+                input_metadata.positions,
+                input_metadata,
+                get_embedding=True,
+            )
 
     @torch.inference_mode()
     def forward_extend_multi_modal(self, batch: ScheduleBatch):
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, forward_mode=ForwardMode.EXTEND
+            self,
+            batch,
+            forward_mode=ForwardMode.EXTEND,
         )
         return self.model.forward(
             batch.input_ids,
@@ -388,7 +568,9 @@ class ModelRunner:
             input_metadata.image_offsets,
         )
 
-    def forward(self, batch: ScheduleBatch, forward_mode: ForwardMode):
+    def forward(
+        self, batch: ScheduleBatch, forward_mode: ForwardMode
+    ) -> Tuple[SampleOutput, LogitsProcessorOutput]:
         if self.is_multimodal_model and forward_mode == ForwardMode.EXTEND:
             return self.forward_extend_multi_modal(batch)
         elif forward_mode == ForwardMode.DECODE:
@@ -444,4 +626,4 @@ def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
 
 
 # Monkey patch model loader
-setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)
+setattr(ModelRegistry, "_try_load_model_cls", load_model_cls_srt)

sglang/srt/models/chatglm.py

@@ -24,8 +24,6 @@ from torch import nn
 from torch.nn import LayerNorm
 from vllm.config import CacheConfig
 from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
@@ -33,18 +31,18 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.sampler import Sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.sampling_metadata import SamplingMetadata
-from vllm.sequence import SamplerOutput
 from vllm.transformers_utils.configs import ChatGLMConfig
 
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 LoraConfig = None
@@ -383,17 +381,11 @@ class ChatGLMForCausalLM(nn.Module):
         input_metadata: InputMetadata,
     ) -> torch.Tensor:
         hidden_states = self.transformer(input_ids, positions, input_metadata)
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
-
-    def sample(
-        self,
-        logits: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> Optional[SamplerOutput]:
-        next_tokens = self.sampler(logits, sampling_metadata)
-        return next_tokens
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         params_dict = dict(self.named_parameters(remove_duplicate=False))

sglang/srt/models/commandr.py

@@ -50,7 +50,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
-from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
@@ -62,8 +61,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
 
+from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -326,6 +327,7 @@ class CohereForCausalLM(nn.Module):
         self.config = config
         self.quant_config = quant_config
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
         self.model = CohereModel(config, quant_config)
 
     @torch.no_grad()
@@ -340,9 +342,11 @@ class CohereForCausalLM(nn.Module):
             positions,
             input_metadata,
         )
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.model.embed_tokens.weight, input_metadata
         )
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/dbrx.py

@@ -45,6 +45,7 @@ from vllm.transformers_utils.configs.dbrx import DbrxConfig
 
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.sampler import Sampler
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
@@ -382,6 +383,7 @@ class DbrxForCausalLM(nn.Module):
             padding_size=DEFAULT_VOCAB_PADDING_SIZE,
         )
         self.logits_processor = LogitsProcessor(config)
+        self.sampler = Sampler()
 
     @torch.no_grad()
     def forward(
@@ -391,9 +393,11 @@ class DbrxForCausalLM(nn.Module):
         input_metadata: InputMetadata,
     ) -> torch.Tensor:
         hidden_states = self.transformer(input_ids, positions, input_metadata)
-        return self.logits_processor(
+        logits_output = self.logits_processor(
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
+        sample_output = self.sampler(logits_output, input_metadata.sampling_info)
+        return sample_output, logits_output
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         expert_params_mapping = [