sglang 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

sglang/srt/entrypoints/openai/serving_score.py

@@ -0,0 +1,61 @@
+import logging
+from typing import Union
+
+from fastapi import Request
+
+from sglang.srt.entrypoints.openai.protocol import (
+    ErrorResponse,
+    ScoringRequest,
+    ScoringResponse,
+)
+from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIServingScore(OpenAIServingBase):
+    """Handler for /v1/score requests"""
+
+    # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
+    # to another module in the future.
+
+    def _request_id_prefix(self) -> str:
+        return "score-"
+
+    def _convert_to_internal_request(
+        self,
+        request: ScoringRequest,
+    ) -> tuple[ScoringRequest, ScoringRequest]:
+        """Convert OpenAI scoring request to internal format"""
+        # For scoring, we pass the request directly as the tokenizer_manager
+        # has a specialized score_request method that doesn't use GenerateReqInput
+
+        return request, request
+
+    async def _handle_non_streaming_request(
+        self,
+        adapted_request: ScoringRequest,
+        request: ScoringRequest,
+        raw_request: Request,
+    ) -> Union[ScoringResponse, ErrorResponse]:
+        """Handle the scoring request"""
+        try:
+            # Use tokenizer_manager's score_request method directly
+            scores = await self.tokenizer_manager.score_request(
+                query=request.query,
+                items=request.items,
+                label_token_ids=request.label_token_ids,
+                apply_softmax=request.apply_softmax,
+                item_first=request.item_first,
+                request=raw_request,
+            )
+
+            # Create response with just the scores, without usage info
+            response = ScoringResponse(
+                scores=scores,
+                model=request.model,
+            )
+            return response
+
+        except ValueError as e:
+            return self.create_error_response(str(e))

sglang/srt/entrypoints/openai/usage_processor.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Mapping, Optional, final
+
+from sglang.srt.entrypoints.openai.protocol import UsageInfo
+
+
+@final
+class UsageProcessor:
+    """Stateless helpers that turn raw token counts into a UsageInfo."""
+
+    @staticmethod
+    def _details_if_cached(count: int) -> Optional[Dict[str, int]]:
+        """Return {"cached_tokens": N} only when N > 0 (keeps JSON slim)."""
+        return {"cached_tokens": count} if count > 0 else None
+
+    @staticmethod
+    def calculate_response_usage(
+        responses: List[Dict[str, Any]],
+        n_choices: int = 1,
+        enable_cache_report: bool = False,
+    ) -> UsageInfo:
+        completion_tokens = sum(r["meta_info"]["completion_tokens"] for r in responses)
+
+        prompt_tokens = sum(
+            responses[i]["meta_info"]["prompt_tokens"]
+            for i in range(0, len(responses), n_choices)
+        )
+
+        cached_details = None
+        if enable_cache_report:
+            cached_total = sum(
+                r["meta_info"].get("cached_tokens", 0) for r in responses
+            )
+            cached_details = UsageProcessor._details_if_cached(cached_total)
+
+        return UsageProcessor.calculate_token_usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            cached_tokens=cached_details,
+        )
+
+    @staticmethod
+    def calculate_streaming_usage(
+        prompt_tokens: Mapping[int, int],
+        completion_tokens: Mapping[int, int],
+        cached_tokens: Mapping[int, int],
+        n_choices: int,
+        enable_cache_report: bool = False,
+    ) -> UsageInfo:
+        # index % n_choices == 0 marks the first choice of a prompt
+        total_prompt_tokens = sum(
+            tok for idx, tok in prompt_tokens.items() if idx % n_choices == 0
+        )
+        total_completion_tokens = sum(completion_tokens.values())
+
+        cached_details = (
+            UsageProcessor._details_if_cached(sum(cached_tokens.values()))
+            if enable_cache_report
+            else None
+        )
+
+        return UsageProcessor.calculate_token_usage(
+            prompt_tokens=total_prompt_tokens,
+            completion_tokens=total_completion_tokens,
+            cached_tokens=cached_details,
+        )
+
+    @staticmethod
+    def calculate_token_usage(
+        prompt_tokens: int,
+        completion_tokens: int,
+        cached_tokens: Optional[Dict[str, int]] = None,
+    ) -> UsageInfo:
+        """Calculate token usage information"""
+        return UsageInfo(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+            prompt_tokens_details=cached_tokens,
+        )

sglang/srt/entrypoints/openai/utils.py

@@ -0,0 +1,72 @@
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from sglang.srt.entrypoints.openai.protocol import (
+    ChatCompletionRequest,
+    CompletionRequest,
+    LogProbs,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def to_openai_style_logprobs(
+    input_token_logprobs=None,
+    output_token_logprobs=None,
+    input_top_logprobs=None,
+    output_top_logprobs=None,
+):
+    ret_logprobs = LogProbs()
+
+    def append_token_logprobs(token_logprobs):
+        for logprob, _, token_text in token_logprobs:
+            ret_logprobs.tokens.append(token_text)
+            ret_logprobs.token_logprobs.append(logprob)
+
+            # Not supported yet
+            ret_logprobs.text_offset.append(-1)
+
+    def append_top_logprobs(top_logprobs):
+        for tokens in top_logprobs:
+            if tokens is not None:
+                ret_logprobs.top_logprobs.append(
+                    {token[2]: token[0] for token in tokens}
+                )
+            else:
+                ret_logprobs.top_logprobs.append(None)
+
+    if input_token_logprobs is not None:
+        append_token_logprobs(input_token_logprobs)
+    if output_token_logprobs is not None:
+        append_token_logprobs(output_token_logprobs)
+    if input_top_logprobs is not None:
+        append_top_logprobs(input_top_logprobs)
+    if output_top_logprobs is not None:
+        append_top_logprobs(output_top_logprobs)
+
+    return ret_logprobs
+
+
+def process_hidden_states_from_ret(
+    ret_item: Dict[str, Any],
+    request: Union[
+        ChatCompletionRequest,
+        CompletionRequest,
+    ],
+) -> Optional[List]:
+    """Process hidden states from a ret item in non-streaming response.
+
+    Args:
+        ret_item: Response item containing meta_info
+        request: The original request object
+
+    Returns:
+        Processed hidden states for the last token, or None
+    """
+    if not request.return_hidden_states:
+        return None
+
+    hidden_states = ret_item["meta_info"].get("hidden_states", None)
+    if hidden_states is not None:
+        hidden_states = hidden_states[-1] if len(hidden_states) > 1 else []
+    return hidden_states

sglang/srt/eplb_simulator/__init__.py

@@ -0,0 +1 @@
+from . import reader

sglang/srt/eplb_simulator/reader.py

@@ -0,0 +1,51 @@
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+from sglang.srt.managers.expert_distribution import (
+    _convert_global_physical_count_to_logical_count,
+)
+
+convert_global_physical_count_to_logical_count = (
+    _convert_global_physical_count_to_logical_count
+)
+
+
+def read_mode_per_pass(dir_data: Path):
+    """Read data from ExpertDistributionRecorder when recorded with mode `per_pass`"""
+
+    # gpc := global_physical_count
+    gpc_of_forward_pass_and_rank = defaultdict(lambda: defaultdict())
+    for path in tqdm(list(dir_data.glob("*.pt"))):
+        data_pack = torch.load(path, weights_only=True)
+        last_physical_to_logical_map = data_pack["last_physical_to_logical_map"]
+        for record in data_pack["records"]:
+            forward_pass_id = record["forward_pass_id"]
+            rank = record["rank"]
+            assert (
+                gpc_of_forward_pass_and_rank[forward_pass_id].get(rank) is None
+            ), f"Duplicated {forward_pass_id=} {rank=}"
+            gpc_of_forward_pass_and_rank[forward_pass_id][rank] = record[
+                "global_physical_count"
+            ]
+
+    forward_pass_ids = sorted(gpc_of_forward_pass_and_rank.keys())
+    print(f"Make {forward_pass_ids=} into array")
+
+    items = []
+    for forward_pass_id, gpc_of_rank in sorted(gpc_of_forward_pass_and_rank.items()):
+        gpc_of_rank_tensor = torch.stack(
+            [gpc for rank, gpc in sorted(gpc_of_rank.items())]
+        ).sum(dim=0)
+        items.append(gpc_of_rank_tensor)
+
+    gpc_of_forward_pass = torch.stack(items)
+    print(f"{gpc_of_forward_pass.shape=}")
+
+    return dict(
+        global_physical_count_of_forward_pass=gpc_of_forward_pass,
+        last_physical_to_logical_map=last_physical_to_logical_map,
+        forward_pass_ids=forward_pass_ids,
+    )

sglang/srt/function_call/base_format_detector.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List
 from partial_json_parser.core.exceptions import MalformedJSON
 from partial_json_parser.core.options import Allow
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
     ToolCallItem,
@@ -16,7 +17,6 @@ from sglang.srt.function_call.utils import (
     _is_complete_json,
     _partial_json_loads,
 )
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 
@@ -111,11 +111,10 @@ class BaseFormatDetector(ABC):
         # The current_text has tool_call if it is the start of a new tool call sequence
         # or it is the start of a new tool call after a tool call separator, when there is a previous tool call
         if not (
-            self.bot_token in current_text
-            or current_text.startswith("{")
+            self.has_tool_call(current_text)
             or (
                 self.current_tool_id > 0
-                and current_text.startswith(self.tool_call_separator + "{")
+                and current_text.startswith(self.tool_call_separator)
             )
         ):
             # Only clear buffer if we're sure no tool call is starting
@@ -143,6 +142,10 @@ class BaseFormatDetector(ABC):
             try:
                 if current_text.startswith(self.bot_token):
                     start_idx = len(self.bot_token)
+                elif self.current_tool_id > 0 and current_text.startswith(
+                    self.tool_call_separator + self.bot_token
+                ):
+                    start_idx = len(self.tool_call_separator + self.bot_token)
                 elif self.current_tool_id > 0 and current_text.startswith(
                     self.tool_call_separator
                 ):

sglang/srt/function_call/deepseekv3_detector.py

@@ -3,6 +3,7 @@ import logging
 import re
 from typing import List
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
@@ -12,7 +13,6 @@ from sglang.srt.function_call.core_types import (
 )
 from sglang.srt.function_call.ebnf_composer import EBNFComposer
 from sglang.srt.function_call.utils import _is_complete_json
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 

sglang/srt/function_call/ebnf_composer.py

@@ -211,20 +211,74 @@ class EBNFComposer:
             properties = params.get("properties", {})
             required_props = set(params.get("required", []))
 
-            # Build argument rules for this tool
-            arg_rules = []
+            # The generated pattern ensures:
+            # 1. Required properties appear first, joined by commas
+            # 2. Optional properties are wrapped with comma included: ( "," ( "prop" : value )? )?
+            # 3. For multiple optional properties, we allow flexible ordering:
+            #    - Each optional can be skipped entirely
+            #    - They can appear in any combination
+            #
+            # Example patterns generated:
+            # - One required, one optional:
+            #   "{" "location" ":" string ( "," ( "unit" ":" enum ) )? "}"
+            #   Allows: {"location": "Paris"} or {"location": "Paris", "unit": "celsius"}
+            #
+            # - Multiple optional properties with flexible ordering:
+            #   "{" "req" ":" string ( "," ( "opt1" ":" value ( "," "opt2" ":" value )? | "opt2" ":" value ) )? "}"
+            #   Allows: {"req": "x"}, {"req": "x", "opt1": "y"}, {"req": "x", "opt2": "z"},
+            #           {"req": "x", "opt1": "y", "opt2": "z"}
+            #
+            # - All optional properties with flexible ordering:
+            #   "{" ( "opt1" ":" value ( "," "opt2" ":" value )? | "opt2" ":" value )? "}"
+            #   Allows: {}, {"opt1": "x"}, {"opt2": "y"}, {"opt1": "x", "opt2": "y"}
+
+            prop_kv_pairs = {}
+            ordered_props = list(properties.keys())
+
             for prop_name, prop_schema in properties.items():
                 value_rule = EBNFComposer.get_value_rule(prop_schema, function_format)
                 # Create key=value pair
                 pair = key_value_template.format(key=prop_name, valrule=value_rule)
-
-                if prop_name not in required_props:
-                    pair = f"[ {pair} ]"
-
-                arg_rules.append(pair)
-
-            # Combine all argument rules
-            combined_args = ' "," '.join(arg_rules) if arg_rules else ""
+                prop_kv_pairs[prop_name] = pair
+
+            # Separate into required and optional while preserving order
+            required = [p for p in ordered_props if p in required_props]
+            optional = [p for p in ordered_props if p not in required_props]
+
+            # Build the combined rule
+            rule_parts = []
+
+            # Add required properties joined by commas
+            if required:
+                rule_parts.append(' "," '.join(prop_kv_pairs[k] for k in required))
+
+            # Add optional properties with flexible ordering
+            if optional:
+                # Build alternatives where any optional property can appear first
+                opt_alternatives = []
+                for i in range(len(optional)):
+                    # Build pattern for optional[i] appearing first
+                    opt_parts = []
+                    for j in range(i, len(optional)):
+                        if j == i:
+                            opt_parts.append(prop_kv_pairs[optional[j]])
+                        else:
+                            opt_parts.append(f' ( "," {prop_kv_pairs[optional[j]]} )?')
+                    opt_alternatives.append("".join(opt_parts))
+
+                # Wrap with appropriate comma handling based on whether we have required properties
+                if required:
+                    # Required properties exist, so optional group needs outer comma
+                    rule_parts.append(' ( "," ( ')
+                    rule_parts.append(" | ".join(opt_alternatives))
+                    rule_parts.append(" ) )?")
+                else:
+                    # All properties are optional
+                    rule_parts.append("( ")
+                    rule_parts.append(" | ".join(opt_alternatives))
+                    rule_parts.append(" )?")
+
+            combined_args = "".join(rule_parts)
             arguments_rule = args_template.format(arg_rules=combined_args)
 
             # Add the function call rule and its arguments rule

sglang/srt/function_call/function_call_parser.py

@@ -1,6 +1,12 @@
 import logging
 from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Type, Union
 
+from sglang.srt.entrypoints.openai.protocol import (
+    StructuralTagResponseFormat,
+    StructuresResponseFormat,
+    Tool,
+    ToolChoice,
+)
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import ToolCallItem
 from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector
@@ -8,12 +14,6 @@ from sglang.srt.function_call.llama32_detector import Llama32Detector
 from sglang.srt.function_call.mistral_detector import MistralDetector
 from sglang.srt.function_call.pythonic_detector import PythonicDetector
 from sglang.srt.function_call.qwen25_detector import Qwen25Detector
-from sglang.srt.openai_api.protocol import (
-    StructuralTagResponseFormat,
-    StructuresResponseFormat,
-    Tool,
-    ToolChoice,
-)
 
 logger = logging.getLogger(__name__)
 

sglang/srt/function_call/llama32_detector.py

@@ -2,6 +2,7 @@ import json
 import logging
 from typing import List
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
@@ -9,7 +10,6 @@ from sglang.srt.function_call.core_types import (
     _GetInfoFunc,
 )
 from sglang.srt.function_call.ebnf_composer import EBNFComposer
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 

sglang/srt/function_call/mistral_detector.py

@@ -3,6 +3,7 @@ import logging
 import re
 from typing import List
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
@@ -10,7 +11,6 @@ from sglang.srt.function_call.core_types import (
     _GetInfoFunc,
 )
 from sglang.srt.function_call.ebnf_composer import EBNFComposer
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 

sglang/srt/function_call/pythonic_detector.py

@@ -4,6 +4,7 @@ import logging
 import re
 from typing import List, Optional
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
@@ -12,7 +13,6 @@ from sglang.srt.function_call.core_types import (
     _GetInfoFunc,
 )
 from sglang.srt.function_call.ebnf_composer import EBNFComposer
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 

sglang/srt/function_call/qwen25_detector.py

@@ -3,6 +3,7 @@ import logging
 import re
 from typing import List
 
+from sglang.srt.entrypoints.openai.protocol import Tool
 from sglang.srt.function_call.base_format_detector import BaseFormatDetector
 from sglang.srt.function_call.core_types import (
     StreamingParseResult,
@@ -10,7 +11,6 @@ from sglang.srt.function_call.core_types import (
     _GetInfoFunc,
 )
 from sglang.srt.function_call.ebnf_composer import EBNFComposer
-from sglang.srt.openai_api.protocol import Tool
 
 logger = logging.getLogger(__name__)
 

sglang/srt/{openai_api/utils.py → jinja_template_utils.py}

@@ -1,11 +1,12 @@
-"""
-Utility functions for OpenAI API adapter.
+"""Template utilities for Jinja template processing.
+
+This module provides utilities for analyzing and processing Jinja chat templates,
+including content format detection and message processing.
 """
 
 import logging
-from typing import Dict, List
 
-import jinja2.nodes
+import jinja2
 import transformers.utils.chat_template_utils as hf_chat_utils
 
 logger = logging.getLogger(__name__)
@@ -75,7 +76,7 @@ def _try_extract_ast(chat_template: str):
         return None
 
 
-def detect_template_content_format(chat_template: str) -> str:
+def detect_jinja_template_content_format(chat_template: str) -> str:
     """
     Detect whether a chat template expects 'string' or 'openai' content format.
 

sglang/srt/layers/activation.py

@@ -20,6 +20,7 @@ from typing import Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PretrainedConfig
 
 from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import (
@@ -28,9 +29,19 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.utils import is_cuda, set_weight_attrs
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+    is_cuda,
+    is_npu,
+    set_weight_attrs,
+)
+from sglang.utils import resolve_obj_by_qualname
 
 _is_cuda = is_cuda()
+_is_npu = is_npu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
 
 if _is_cuda:
     from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
@@ -50,6 +61,15 @@ class SiluAndMul(CustomOp):
         silu_and_mul(x, out)
         return out
 
+    def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
+        if _is_cpu_amx_available:
+            d = x.shape[-1] // 2
+            output_shape = x.shape[:-1] + (d,)
+            out = torch.ops.sgl_kernel.silu_and_mul_cpu(x)
+            return out
+        else:
+            return self.forward_native(x)
+
 
 class GeluAndMul(CustomOp):
     def __init__(self, approximate="tanh"):
@@ -165,8 +185,25 @@ def get_act_fn(
     return act_fn
 
 
-if not _is_cuda:
+def get_cross_encoder_activation_function(config: PretrainedConfig):
+    if (
+        hasattr(config, "sbert_ce_default_activation_function")
+        and config.sbert_ce_default_activation_function is not None
+    ):
+
+        function_name = config.sbert_ce_default_activation_function
+        assert function_name.startswith("torch.nn.modules."), (
+            "Loading of activation functions is restricted to "
+            "torch.nn.modules for security reasons"
+        )
+        return resolve_obj_by_qualname(function_name)()
+    else:
+        # adapt bge-reranker
+        return nn.Identity()
+
+
+if not (_is_cuda or _is_npu or (_is_cpu and _is_cpu_amx_available)):
     logger.info(
-        "sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries."
+        "sgl-kernel is not available on Non-NV platforms or Non-AMX CPUs. Fallback to other kernel libraries."
     )
     from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul

sglang/srt/layers/attention/aiter_backend.py

@@ -324,7 +324,10 @@ class AiterAttnBackend(AttentionBackend):
                 )
 
     def init_cuda_graph_state(
-        self, max_bs: int, kv_indices_buf: Optional[torch.Tensor] = None
+        self,
+        max_bs: int,
+        max_num_tokens: int,
+        kv_indices_buf: Optional[torch.Tensor] = None,
     ):
         self.cuda_graph_kv_last_page_len = torch.ones(max_bs, dtype=torch.int)
         if kv_indices_buf is None:
@@ -338,7 +341,7 @@ class AiterAttnBackend(AttentionBackend):
 
         if not self.skip_prefill:
             self.cuda_graph_custom_mask = torch.zeros(
-                (max_bs * self.max_context_len),
+                (max_num_tokens * self.max_context_len),
                 dtype=torch.uint8,
                 device=self.device,
             )
@@ -717,6 +720,11 @@ class AiterIndicesUpdaterPrefill:
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.update = self.update_single_wrapper
 
+        # get the last index of the pool
+        self.pool_size = (
+            model_runner.token_to_kv_pool.size + model_runner.token_to_kv_pool.page_size
+        ) - 1
+
         self.kv_indices = None
         self.max_q_len = 0
         self.max_kv_len = 0
@@ -754,8 +762,16 @@ class AiterIndicesUpdaterPrefill:
             # Normal extend
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
-            kv_indices = torch.empty(
-                paged_kernel_lens_sum + 256,
+
+            # (TODO: Kk) WA - CI test_moe_eval_accuracy_large.py
+            # mha_batch_prefill reads 128 data to do computatoin
+            # if real data is not long enough then original padding value 0 is used
+            # but the 0 location will be made nan (noqa) in cuda graph capture mode
+            # this will cause the output tensor value becomes nan
+            # WA is to assure that last index of pool not changed
+            kv_indices = torch.full(
+                (paged_kernel_lens_sum + 128,),
+                self.pool_size,
                 dtype=torch.int32,
                 device=req_pool_indices.device,
             )

sglang/srt/layers/attention/base_attn_backend.py

@@ -19,7 +19,7 @@ class AttentionBackend(ABC):
         """Init the metadata for a forward pass."""
         raise NotImplementedError()
 
-    def init_cuda_graph_state(self, max_bs: int):
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         """Init the global shared states for cuda graph."""
         raise NotImplementedError()