sglang 0.4.7.post1__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/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

@@ -29,10 +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
@@ -52,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"):
@@ -184,8 +202,8 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
         return nn.Identity()
 
 
-if not _is_cuda:
+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,
             )

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()
 

sglang/srt/layers/attention/cutlass_mla_backend.py

@@ -122,6 +122,7 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
     def init_cuda_graph_state(
         self,
         max_bs: int,
+        max_num_tokens: int,
         block_kv_indices: Optional[torch.Tensor] = None,
     ):
         if block_kv_indices is None:

sglang/srt/layers/attention/flashattention_backend.py

@@ -1120,7 +1120,7 @@ class FlashAttentionBackend(AttentionBackend):
 
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
-    def init_cuda_graph_state(self, max_bs: int):
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         """Initialize CUDA graph state for the attention backend.
 
         Args:
@@ -1704,14 +1704,15 @@ class FlashAttentionBackend(AttentionBackend):
 
                 # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
                 metadata_expand = self.target_verify_metadata_topk_expand[bs]
+
                 # metadata_expand.max_seq_len_q = 1, already set in capture
                 # metadata_expand.cu_seqlens_q already set in capture
-
                 offsets = torch.arange(
                     self.speculative_num_draft_tokens, device=device
                 ).unsqueeze(
                     0
                 )  # shape: (1, self.speculative_num_draft_tokens)
+
                 cols = offsets.expand(seq_lens.numel(), -1) + seq_lens.unsqueeze(1)
                 cum_len = torch.nn.functional.pad(
                     torch.cumsum(
@@ -1728,17 +1729,20 @@ class FlashAttentionBackend(AttentionBackend):
                 ).view(1, -1)
                 # avoid extracting padded seq indices which will be out of boundary
                 mask_extraction_indices[
-                    :, spec_info.positions.numel() * self.speculative_num_draft_tokens :
+                    :,
+                    spec_info.positions.numel() * self.speculative_num_draft_tokens :,
                 ].fill_(0)
-
                 mask = spec_info.custom_mask[mask_extraction_indices].view(
                     -1, self.speculative_num_draft_tokens
                 )  # (bsz * draft_num, draft_num)
+
                 col_indices = offsets.expand(
                     mask.shape[0], self.speculative_num_draft_tokens
                 )
                 keys = torch.where(
-                    mask, col_indices, col_indices + self.speculative_num_draft_tokens
+                    mask,
+                    col_indices,
+                    col_indices + self.speculative_num_draft_tokens,
                 )
                 _, sort_order = torch.sort(keys, dim=1)
 
@@ -1747,6 +1751,7 @@ class FlashAttentionBackend(AttentionBackend):
                     .gather(1, cols)
                     .repeat_interleave(self.speculative_num_draft_tokens, dim=0)
                 )  # (bsz, draft_num)
+
                 metadata_expand.page_table.copy_(
                     non_masked_page_table.gather(1, sort_order)
                 )
@@ -1758,6 +1763,7 @@ class FlashAttentionBackend(AttentionBackend):
                         dtype=torch.int32,
                     )
                 )
+
         elif forward_mode.is_draft_extend():
             metadata = self.draft_extend_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens)
@@ -1767,7 +1773,11 @@ class FlashAttentionBackend(AttentionBackend):
                 torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
             )
             accept_length = spec_info.accept_length[:bs]
-            metadata.max_seq_len_q = max(spec_info.accept_length_cpu) + 1
+            if spec_info.accept_length_cpu:
+                metadata.max_seq_len_q = max(spec_info.accept_length_cpu) + 1
+            else:
+                metadata.max_seq_len_q = 1
+
             metadata.cu_seqlens_q[1:].copy_(
                 torch.cumsum(accept_length, dim=0, dtype=torch.int32)
             )
@@ -1807,7 +1817,7 @@ class FlashAttentionBackend(AttentionBackend):
 
     def get_cuda_graph_seq_len_fill_value(self):
         """Get the fill value for sequence length in CUDA graph."""
-        return 0
+        return 1
 
     def _init_local_attn_metadata(self, metadata: FlashAttentionMetadata, device):
         """Centralized utility to initialize local_attn_metadata if chunked attention is enabled."""
@@ -1999,9 +2009,9 @@ class FlashAttentionMultiStepBackend:
         for i in range(self.speculative_num_steps - 1):
             self.attn_backends[i].init_forward_metadata(forward_batch)
 
-    def init_cuda_graph_state(self, max_bs: int):
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         for i in range(self.speculative_num_steps):
-            self.attn_backends[i].init_cuda_graph_state(max_bs)
+            self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens)
 
     def init_forward_metadata_capture_cuda_graph(
         self,