sglang 0.3.1__py3-none-any.whl → 0.3.1.post2__py3-none-any.whl

sglang/srt/lora/lora_manager.py

@@ -21,12 +21,15 @@ import re
 from dataclasses import dataclass
 
 import torch
-from flashinfer import SegmentGEMMWrapper
 
 from sglang.srt.lora.lora import LoRAAdapter, get_lora_layer
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.utils import replace_submodule
+from sglang.srt.utils import is_hip, replace_submodule
+
+# ROCm: flashinfer available later
+if not is_hip():
+    from flashinfer import SegmentGEMMWrapper
 
 
 def get_stacked_name(name):
@@ -96,10 +99,10 @@ class LoRAManager:
         # get configs and target modules
         self.configs = {}
         self.origin_target_modules = set()
-        for path in self.lora_paths:
-            self.configs[path] = LoRAConfig(path)
+        for name, path in self.lora_paths.items():
+            self.configs[name] = LoRAConfig(path)
             self.origin_target_modules = set(self.origin_target_modules) | set(
-                self.configs[path].target_modules
+                self.configs[name].target_modules
             )
         self.target_modules = set(
             [
@@ -114,11 +117,11 @@ class LoRAManager:
         # load all weights to cpu
         self.loras = []
         self.lora_id = {}
-        for path in self.lora_paths:
-            self.lora_id[path] = len(self.loras)
+        for name in self.lora_paths.keys():
+            self.lora_id[name] = len(self.loras)
             self.loras.append(
                 LoRAAdapter(
-                    path, self.configs[path], self.base_hf_config, self.load_config
+                    name, self.configs[name], self.base_hf_config, self.load_config
                 )
             )
             self.loras[-1].initialize_weights()

sglang/srt/managers/io_struct.py

@@ -133,6 +133,9 @@ class GenerateReqInput:
                 self.image_data = [None] * num
             elif not isinstance(self.image_data, list):
                 self.image_data = [self.image_data] * num
+            elif isinstance(self.image_data, list):
+                # multi-image with n > 1
+                self.image_data = self.image_data * num
 
             if self.sampling_params is None:
                 self.sampling_params = [{}] * num

sglang/srt/managers/policy_scheduler.py

@@ -119,19 +119,32 @@ class PrefillAdder:
         self.running_batch = running_batch
         self.new_token_ratio = new_token_ratio
         self.rem_total_tokens = rem_total_tokens - mixed_with_decode_tokens
-        self.rem_total_tokens_ = self.rem_total_tokens
-        self.total_tokens = rem_total_tokens
         self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= mixed_with_decode_tokens
 
+        self.cur_rem_tokens = rem_total_tokens - mixed_with_decode_tokens
+
         self.req_states = None
         self.can_run_list = []
         self.new_inflight_req = None
         self.log_hit_tokens = 0
         self.log_input_tokens = 0
 
+        if running_batch is not None:
+            # Pre-remove the tokens which will be occupied by the running requests
+            self.rem_total_tokens -= sum(
+                [
+                    min(
+                        (r.sampling_params.max_new_tokens - len(r.output_ids)),
+                        CLIP_MAX_NEW_TOKENS,
+                    )
+                    * self.new_token_ratio
+                    for r in running_batch.reqs
+                ]
+            )
+
     def no_remaining_tokens(self):
         return (
             self.rem_total_tokens <= 0
@@ -141,31 +154,14 @@ class PrefillAdder:
                 if self.rem_chunk_tokens is not None
                 else False
             )
-        )
-
-    def remove_running_tokens(self, running_batch: ScheduleBatch):
-        self.rem_total_tokens -= sum(
-            [
-                min(
-                    (r.sampling_params.max_new_tokens - len(r.output_ids)),
-                    CLIP_MAX_NEW_TOKENS,
-                )
-                * self.new_token_ratio
-                for r in running_batch.reqs
-            ]
-        )
-        self.rem_total_tokens_ -= sum(
-            [
-                r.sampling_params.max_new_tokens - len(r.output_ids)
-                for r in running_batch.reqs
-            ]
+            or self.cur_rem_tokens <= 0
         )
 
     def _prefill_one_req(
         self, prefix_len: int, extend_input_len: int, max_new_tokens: int
     ):
         self.rem_total_tokens -= extend_input_len + max_new_tokens
-        self.rem_total_tokens_ -= extend_input_len + max_new_tokens
+        self.cur_rem_tokens -= extend_input_len
         self.rem_input_tokens -= extend_input_len
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= extend_input_len
@@ -173,29 +169,7 @@ class PrefillAdder:
         self.log_hit_tokens += prefix_len
         self.log_input_tokens += extend_input_len
 
-    def add_inflight_req_ignore_eos(self, req: Req):
-        truncated = req.extend_input_len > self.rem_chunk_tokens
-        req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
-        req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
-        self.can_run_list.append(req)
-
-        self._prefill_one_req(
-            0,
-            req.extend_input_len,
-            (
-                min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
-                if not truncated
-                else 0
-            ),
-        )
-
-        # Return if chunked prefill not finished
-        return req if truncated else None
-
     def add_inflight_req(self, req: Req):
-        if req.sampling_params.ignore_eos:
-            return self.add_inflight_req_ignore_eos(req)
-
         truncated = req.extend_input_len > self.rem_chunk_tokens
         req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
@@ -225,7 +199,7 @@ class PrefillAdder:
             self.rem_total_tokens += delta
 
     def add_one_req_ignore_eos(self, req: Req):
-        def get_req_state(r):
+        def add_req_state(r, insert_sort=False):
             new_token_ratio = (
                 1.0 if r.sampling_params.ignore_eos else self.new_token_ratio
             )
@@ -235,56 +209,38 @@ class PrefillAdder:
             tokens_occupied = len(r.origin_input_ids) + len(r.output_ids)
 
             if tokens_left > 0:
-                return (tokens_left, tokens_occupied)
-
-            return None
-
-        # Quick Check
-        can_run = False
-        if (
-            req.extend_input_len + req.sampling_params.max_new_tokens
-            <= self.rem_total_tokens
-        ):
-            can_run = True
-
-        if not can_run:
-            if self.req_states is None:
-                self.req_states = []
-                if self.running_batch is not None:
-                    for r in self.running_batch.reqs:
-                        state = get_req_state(r)
-                        if state is not None:
-                            self.req_states.append(state)
-                for r in self.can_run_list:
-                    state = get_req_state(r)
-                    if state is not None:
-                        self.req_states.append(state)
-                state = get_req_state(req)
-                if state is not None:
-                    self.req_states.append(state)
-
-                self.req_states.sort(key=lambda x: x[0])
-            else:
-                state = get_req_state(req)
-                if state is not None:
-                    for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
-                        if tokens_left >= state[0]:
-                            self.req_states.insert(i, state)
+                if not insert_sort:
+                    self.req_states.append((tokens_left, tokens_occupied))
+                else:
+                    for i in range(len(self.req_states)):
+                        if tokens_left <= self.req_states[i][0]:
                             break
-                    else:
-                        self.req_states.append(state)
-
-            tokens_freed = 0
-            for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
-                decode_steps = (
-                    self.req_states[i + 1][0]
-                    if i + 1 < len(self.req_states)
-                    else tokens_left
-                )
-                bs = len(self.req_states) - i
-                if self.total_tokens + tokens_freed - decode_steps * bs <= 0:
-                    return False
-                tokens_freed += tokens_occupied
+                    self.req_states.insert(i, (tokens_left, tokens_occupied))
+
+        if self.req_states is None:
+            self.req_states = []
+            add_req_state(req)
+            if self.running_batch is not None:
+                for r in self.running_batch.reqs:
+                    add_req_state(r)
+            for r in self.can_run_list:
+                add_req_state(r)
+            self.req_states.sort(key=lambda x: x[0])
+        else:
+            add_req_state(req, insert_sort=True)
+
+        cur_rem_tokens = self.cur_rem_tokens - len(req.origin_input_ids)
+        tokens_freed = 0
+        for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
+            decode_steps = (
+                self.req_states[i + 1][0]
+                if i + 1 < len(self.req_states)
+                else tokens_left
+            )
+            bs = len(self.req_states) - i
+            if cur_rem_tokens + tokens_freed - decode_steps * bs <= 0:
+                return False
+            tokens_freed += tokens_occupied
 
         if req.extend_input_len <= self.rem_chunk_tokens:
             self.can_run_list.append(req)

sglang/srt/managers/schedule_batch.py

@@ -40,7 +40,7 @@ global_server_args_dict = {
     "attention_backend": ServerArgs.attention_backend,
     "sampling_backend": ServerArgs.sampling_backend,
     "triton_attention_reduce_in_fp32": ServerArgs.triton_attention_reduce_in_fp32,
-    "enable_mla": ServerArgs.enable_mla,
+    "disable_mla": ServerArgs.disable_mla,
     "torchao_config": ServerArgs.torchao_config,
 }
 
@@ -360,6 +360,7 @@ class ScheduleBatch:
     tree_cache: BasePrefixCache
 
     forward_mode: ForwardMode = None
+    sampling_info: SamplingBatchInfo = None
 
     # Batched arguments to model runner
     input_ids: torch.Tensor = None

sglang/srt/managers/tp_worker.py

@@ -198,6 +198,7 @@ class ModelTpServer:
                     "trust_remote_code": server_args.trust_remote_code,
                 },
                 skip_tokenizer_init=server_args.skip_tokenizer_init,
+                constrained_json_whitespace_pattern=server_args.constrained_json_whitespace_pattern,
             )
         self.jump_forward_cache = JumpForwardCache()
 
@@ -414,7 +415,7 @@ class ModelTpServer:
 
         # Truncate prompts that are too long
         if len(req.origin_input_ids) >= self.max_req_input_len:
-            logger.warn(
+            logger.warning(
                 "Request length is longer than the KV cache pool size or "
                 "the max context length. Truncated!!!"
             )
@@ -444,9 +445,6 @@ class ModelTpServer:
             num_mixed_running,
         )
 
-        if self.running_batch is not None:
-            adder.remove_running_tokens(self.running_batch)
-
         has_inflight = self.current_inflight_req is not None
         if self.current_inflight_req is not None:
             self.current_inflight_req.init_next_round_input(
@@ -464,9 +462,6 @@ class ModelTpServer:
             )
 
         for req in self.waiting_queue:
-            if adder.no_remaining_tokens():
-                break
-            req.init_next_round_input(None if prefix_computed else self.tree_cache)
             if (
                 self.lora_paths is not None
                 and len(
@@ -477,6 +472,10 @@ class ModelTpServer:
                 > self.max_loras_per_batch
             ):
                 break
+
+            if adder.no_remaining_tokens():
+                break
+            req.init_next_round_input(None if prefix_computed else self.tree_cache)
             res = adder.add_one_req(req)
             if (
                 not res
@@ -506,6 +505,11 @@ class ModelTpServer:
             else:
                 tree_cache_hit_rate = 0.0
 
+            num_used = self.max_total_num_tokens - (
+                self.token_to_kv_pool.available_size()
+                + self.tree_cache.evictable_size()
+            )
+
             if num_mixed_running > 0:
                 logger.info(
                     f"Prefill batch"
@@ -514,6 +518,7 @@ class ModelTpServer:
                     f"#new-token: {adder.log_input_tokens}, "
                     f"#cached-token: {adder.log_hit_tokens}, "
                     f"cache hit rate: {100.0 * tree_cache_hit_rate:.2f}%, "
+                    f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
                     f"#queue-req: {len(self.waiting_queue) - len(can_run_list) + has_inflight}"
                 )
             else:
@@ -523,6 +528,7 @@ class ModelTpServer:
                     f"#new-token: {adder.log_input_tokens}, "
                     f"#cached-token: {adder.log_hit_tokens}, "
                     f"cache hit rate: {100.0 * tree_cache_hit_rate:.2f}%, "
+                    f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
                     f"#running-req: {running_bs}, "
                     f"#queue-req: {len(self.waiting_queue) - len(can_run_list) + has_inflight}"
                 )
@@ -807,12 +813,10 @@ class ModelTpServer:
                 unfinished_indices.append(i)
 
             if req.finished() or (
-                (
-                    req.stream
-                    and (
-                        self.decode_forward_ct % self.stream_interval == 0
-                        or len(req.output_ids) == 1
-                    )
+                req.stream
+                and (
+                    self.decode_forward_ct % self.stream_interval == 0
+                    or len(req.output_ids) == 1
                 )
             ):
                 output_rids.append(req.rid)
@@ -937,6 +941,8 @@ class ModelTpServer:
         if success:
             flash_cache_success = self.flush_cache()
             assert flash_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
         return success, message
 
 

sglang/srt/model_executor/cuda_graph_runner.py

@@ -41,6 +41,9 @@ if TYPE_CHECKING:
 def _to_torch(model: torch.nn.Module, reverse: bool = False):
     for sub in model._modules.values():
         if isinstance(sub, CustomOp):
+            # NOTE: FusedMoE torch native implementaiton is not efficient
+            if "FusedMoE" in sub.__class__.__name__:
+                continue
             if reverse:
                 sub._forward_method = sub.forward_cuda
                 setattr(sub, "is_torch_compile", False)
@@ -105,23 +108,22 @@ class CudaGraphRunner:
             self.capture_bs = list(range(1, 32)) + [64, 128]
         else:
             self.capture_bs = [1, 2, 4] + [i * 8 for i in range(1, 21)]
-        self.compile_bs = [1, 2, 4, 8, 16, 24, 32] if self.use_torch_compile else []
 
-        # Common inputs
-        self.max_bs = max(self.capture_bs)
-        self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.req_pool_indices = torch.zeros(
-            (self.max_bs,), dtype=torch.int32, device="cuda"
-        )
-        self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.position_ids_offsets = torch.ones(
-            (self.max_bs,), dtype=torch.int32, device="cuda"
-        )
-        self.out_cache_loc = torch.zeros(
-            (self.max_bs,), dtype=torch.int32, device="cuda"
+        self.capture_bs = [
+            bs for bs in self.capture_bs if bs <= model_runner.req_to_token_pool.size
+        ]
+        self.compile_bs = (
+            [
+                bs
+                for bs in self.capture_bs
+                if bs <= self.model_runner.server_args.max_torch_compile_bs
+            ]
+            if self.use_torch_compile
+            else []
         )
 
         # Attention backend
+        self.max_bs = max(self.capture_bs)
         self.model_runner.attn_backend.init_cuda_graph_state(self.max_bs)
         self.seq_len_fill_value = (
             self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
@@ -130,6 +132,16 @@ class CudaGraphRunner:
         if self.use_torch_compile:
             set_torch_compile_config()
 
+        # Common inputs
+        with torch.device("cuda"):
+            self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32)
+            self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
+            self.seq_lens = torch.full(
+                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
+            )
+            self.position_ids_offsets = torch.ones((self.max_bs,), dtype=torch.int32)
+            self.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32)
+
         # Capture
         try:
             self.capture()

sglang/srt/model_executor/model_runner.py

@@ -40,7 +40,7 @@ from vllm.model_executor.models import ModelRegistry
 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, Sampler
+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 (
@@ -54,11 +54,9 @@ 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__)
@@ -88,12 +86,20 @@ class ModelRunner:
         self.is_multimodal_model = is_multimodal_model(
             self.model_config.hf_config.architectures
         )
+
+        if (
+            self.model_config.attention_arch == AttentionArch.MLA
+            and not self.server_args.disable_mla
+        ):
+            logger.info("MLA optimization is tunred on. Use triton backend.")
+            self.server_args.attention_backend = "triton"
+
         global_server_args_dict.update(
             {
                 "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,
+                "disable_mla": server_args.disable_mla,
                 "torchao_config": server_args.torchao_config,
             }
         )
@@ -166,10 +172,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."
@@ -178,6 +187,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)
@@ -188,23 +198,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,
@@ -255,20 +258,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(
@@ -293,14 +296,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)
@@ -315,7 +318,7 @@ 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(
@@ -334,7 +337,7 @@ class ModelRunner:
         )
         if (
             self.model_config.attention_arch == AttentionArch.MLA
-            and self.server_args.enable_mla
+            and not self.server_args.disable_mla
         ):
             cell_size = (
                 (self.model_config.kv_lora_rank + self.model_config.qk_rope_head_dim)
@@ -397,12 +400,12 @@ class ModelRunner:
             )
 
         self.req_to_token_pool = ReqToTokenPool(
-            max_num_reqs,
-            self.model_config.context_len + 8,
+            max_num_reqs + 1,
+            self.model_config.context_len + 4,
         )
         if (
             self.model_config.attention_arch == AttentionArch.MLA
-            and self.server_args.enable_mla
+            and not self.server_args.disable_mla
         ):
             self.token_to_kv_pool = MLATokenToKVPool(
                 self.max_total_num_tokens,
@@ -521,21 +524,6 @@ class ModelRunner:
         else:
             raise ValueError(f"Invaid forward mode: {batch.forward_mode}")
 
-    def _check_sample_results(self, sample_output: SampleOutput):
-        if not torch.all(sample_output.success):
-            probs = sample_output.probs
-            batch_next_token_ids = sample_output.batch_next_token_ids
-            logging.warning("Sampling failed, fallback to top_k=1 strategy")
-            probs = probs.masked_fill(torch.isnan(probs), 0.0)
-            argmax_ids = torch.argmax(probs, dim=-1)
-            batch_next_token_ids = torch.where(
-                sample_output.success, batch_next_token_ids, argmax_ids
-            )
-            sample_output.probs = probs
-            sample_output.batch_next_token_ids = batch_next_token_ids
-
-        return sample_output.batch_next_token_ids
-
     def _apply_logits_bias(
         self, logits: torch.Tensor, sampling_info: SamplingBatchInfo
     ):
@@ -564,13 +552,16 @@ class ModelRunner:
     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_output = self.sampler(logits, batch.sampling_info)
-        return self._check_sample_results(sample_output)
+
+        # Sample the next tokens.
+        next_token_ids = self.sampler(logits, batch.sampling_info)
+        return next_token_ids
 
 
 @lru_cache()