sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/layers/attention/flashinfer_backend.py

@@ -7,6 +7,7 @@ FlashInfer is faster and Triton is easier to customize.
 Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
 """
 
+import logging
 import os
 from dataclasses import dataclass
 from enum import Enum, auto
@@ -15,20 +16,13 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
 
 import torch
 
-if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1":
-    import logging
-
-    torch._logging.set_logs(dynamo=logging.ERROR)
-    torch._dynamo.config.suppress_errors = True
-
-from sglang.global_config import global_config
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import (
     get_int_env_var,
@@ -41,6 +35,12 @@ if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
 
+logger = logging.getLogger(__name__)
+
+if envs.SGLANG_ENABLE_TORCH_COMPILE.get():
+    torch._logging.set_logs(dynamo=logging.ERROR)
+    torch._dynamo.config.suppress_errors = True
+
 
 if is_flashinfer_available():
     from flashinfer import (
@@ -50,7 +50,6 @@ if is_flashinfer_available():
         fast_decode_plan,
     )
     from flashinfer.cascade import merge_state
-    from flashinfer.decode import _get_range_buf, get_seq_lens
 
 
 class WrapperDispatch(Enum):
@@ -58,6 +57,36 @@ class WrapperDispatch(Enum):
     CROSS_ATTENTION = auto()
 
 
+@dataclass
+class MultiItemScoringParams:
+    """Parameters for multi-item scoring in attention computation.
+
+    Used when processing sequences with multiple items separated by delimiters,
+    where each item needs specific attention patterns that respect item boundaries.
+
+    Attributes:
+        prefix_len_ptr: A uint32 1D tensor indicating the prefix length of each prompt.
+                       The tensor size is equal to the batch size.
+        token_pos_in_items_ptr: A uint16 1D tensor indicating the token position of each item
+                               starting from 0 (delimiter) for each item. For batch size > 1,
+                               sequences are concatenated with zero padding to ensure same length.
+        token_pos_in_items_len: Zero padding length for token_pos_in_items_ptr to handle
+                               batch_size > 1 case. Defines the padded length for each sequence.
+        max_item_len_ptr: A uint16 tensor containing the max token length of all items
+                         for each prompt in the batch.
+
+    """
+
+    prefix_len_ptr: Optional[torch.Tensor] = None
+    token_pos_in_items_ptr: Optional[torch.Tensor] = None
+    token_pos_in_items_len: int = 0
+    max_item_len_ptr: Optional[torch.Tensor] = None
+
+    def is_enabled(self) -> bool:
+        """Check if multi-item scoring is enabled."""
+        return self.prefix_len_ptr is not None
+
+
 @dataclass
 class DecodeMetadata:
     decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper]
@@ -68,6 +97,7 @@ class PrefillMetadata:
     prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper]
     use_ragged: bool
     extend_no_prefix: bool
+    multi_item_params: Optional[MultiItemScoringParams] = None
 
 
 # Reuse this workspace buffer across all flashinfer wrappers
@@ -87,9 +117,15 @@ class FlashInferAttnBackend(AttentionBackend):
         skip_prefill: bool = False,
         kv_indptr_buf: Optional[torch.Tensor] = None,
         kv_last_page_len_buf: Optional[torch.Tensor] = None,
+        init_new_workspace: bool = False,
     ):
         super().__init__()
 
+        # Store multi-item scoring delimiter for efficient access
+        self.multi_item_scoring_delimiter = (
+            model_runner.server_args.multi_item_scoring_delimiter
+        )
+
         # Parse constants
         self.decode_use_tensor_cores = should_use_tensor_core(
             kv_cache_dtype=model_runner.kv_cache_dtype,
@@ -124,7 +160,7 @@ class FlashInferAttnBackend(AttentionBackend):
             or "Qwen3ForCausalLM" in model_runner.model_config.hf_config.architectures
             or "MiMoForCausalLM" in model_runner.model_config.hf_config.architectures
         ):
-            global_config.flashinfer_workspace_size = 512 * 1024 * 1024
+            envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(512 * 1024 * 1024)
 
         # When deterministic inference is enabled, tensor cores should be used for decode
         # Also set split tile sizes for prefill and decode from environment variables, and disable kv split for cuda graph
@@ -144,19 +180,26 @@ class FlashInferAttnBackend(AttentionBackend):
                 "SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE", 2048
             )
             self.disable_cuda_graph_kv_split = True
-            global_config.flashinfer_workspace_size = 2048 * 1024 * 1024
+            envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(2048 * 1024 * 1024)
 
         # Allocate buffers
         global global_workspace_buffer
         if global_workspace_buffer is None:
             # different from flashinfer zero_init_global_workspace_buffer
-            global_workspace_size = global_config.flashinfer_workspace_size
+            global_workspace_size = envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get()
             global_workspace_buffer = torch.empty(
                 global_workspace_size,
                 dtype=torch.uint8,
                 device=model_runner.device,
             )
-        self.workspace_buffer = global_workspace_buffer
+        if init_new_workspace:
+            self.workspace_buffer = torch.empty(
+                envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(),
+                dtype=torch.uint8,
+                device=model_runner.device,
+            )
+        else:
+            self.workspace_buffer = global_workspace_buffer
         max_bs = model_runner.req_to_token_pool.size
         if kv_indptr_buf is None:
             self.kv_indptr = [
@@ -187,7 +230,16 @@ class FlashInferAttnBackend(AttentionBackend):
 
         fmha_backend = "auto"
         if is_sm100_supported():
-            fmha_backend = "cutlass"
+            # Disable CUTLASS backend when piecewise cuda graph is enabled
+            # due to TMA descriptor initialization issues on B200
+            if model_runner.server_args.enable_piecewise_cuda_graph:
+                logger.warning(
+                    "CUTLASS backend is disabled when piecewise cuda graph is enabled "
+                    "due to TMA descriptor initialization issues on B200. "
+                    "Using auto backend instead for stability."
+                )
+            else:
+                fmha_backend = "cutlass"
         self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
             self.workspace_buffer, "NHD", backend=fmha_backend
         )
@@ -229,10 +281,133 @@ class FlashInferAttnBackend(AttentionBackend):
 
         # Other metadata
         self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
+
         self.decode_cuda_graph_metadata = {}
         self.prefill_cuda_graph_metadata = {}  # For verify
         self.draft_extend_cuda_graph_metadata = {}  # For draft extend
 
+    def _process_multi_item_scoring(
+        self, forward_batch: ForwardBatch
+    ) -> MultiItemScoringParams:
+        """Process multi-item scoring tensors for FlashInfer attention.
+
+        This method handles sequences containing multiple "items" separated by delimiter tokens,
+        where each item needs specific attention patterns that respect item boundaries.
+
+        The method produces four key tensors for FlashInfer:
+        - prefix_len_ptr: uint32 tensor with prefix length for each prompt in batch
+        - token_pos_in_items_ptr: uint16 tensor with token positions starting from 0 at delimiters
+        - token_pos_in_items_len: padding length for batch processing
+        - max_item_len_ptr: uint16 tensor with max item length for each prompt
+
+        Args:
+            forward_batch: The forward batch containing input sequences and delimiter info
+
+        Returns:
+            MultiItemScoringParams: The processed multi-item scoring parameters
+
+        Examples:
+            Following FlashInfer definition: for 3 items of length 3, 2, 4 respectively:
+            token_pos_in_items_ptr = [0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0]
+
+            Case 1: Single sequence
+            Text: "What is the capital of France? <delim> London <delim> Paris <delim> Berlin <delim>"
+            Tokens: [What, is, the, capital, of, France, ?, <delim>, London, <delim>, Paris, <delim>, Berlin, <delim>]
+            Indices: [ 0,   1,  2,   3,      4,  5,     6,   7,     8,      9,     10,    11,    12,     13]
+            - prefix_len_ptr: [7] (query length before first delimiter)
+            - token_pos_in_items_ptr: [0, 1, 0, 1, 0, 1, 0] (delim=0, London=1, delim=0, Paris=1, delim=0, Berlin=1, delim=0)
+            - token_pos_in_items_len: 7 (actual length)
+            - max_item_len_ptr: [1] (max item length is 1 token - all options are single tokens)
+
+            Case 2: Batch processing (batch_size=2)
+            Sequence 1: 2 items of length 2, 1 → [0, 1, 2, 0, 1, 0] (6 elements)
+            Sequence 2: 3 items of length 1, 3, 2 → [0, 1, 0, 1, 2, 3, 0, 1, 2, 0] (10 elements)
+            After padding both to length 10:
+            - token_pos_in_items_ptr: [0, 1, 2, 0, 1, 0, 0, 0, 0, 0,    0, 1, 0, 1, 2, 3, 0, 1, 2, 0]
+            - token_pos_in_items_len: 10 (padded length for batch processing)
+            - max_item_len_ptr: [2, 3] (max lengths per sequence)
+        """
+
+        delimiter = self.multi_item_scoring_delimiter
+        if delimiter is None or forward_batch.forward_mode == ForwardMode.DECODE:
+            return MultiItemScoringParams()
+
+        delimiter_mask = forward_batch.input_ids == delimiter
+        prefix_cache_lens = getattr(forward_batch, "extend_prefix_lens", None)
+        extend_seq_lens = getattr(forward_batch, "extend_seq_lens", None)
+        prefix_len_ptr, token_pos_in_items_ptr = [], []
+        token_pos_in_items_len = 0
+
+        # If no extend_seq_lens, treat whole batch as one sequence
+        if extend_seq_lens is None or len(extend_seq_lens) <= 1:
+            extend_seq_lens = [forward_batch.input_ids.size(0)]
+
+        seq_start = 0
+        for i, seq_len in enumerate(extend_seq_lens):
+            seq_end = seq_start + seq_len
+            mask = delimiter_mask[seq_start:seq_end]
+            pos = forward_batch.positions[seq_start:seq_end]
+            delimiter_indices = torch.nonzero(mask, as_tuple=True)[0]
+
+            if len(delimiter_indices) > 0:
+                first_delim = delimiter_indices[0]
+                # Prefix length: store as scalar
+                prefix_len = first_delim + (
+                    prefix_cache_lens[i] if prefix_cache_lens is not None else 0
+                )
+                prefix_len_ptr.append(
+                    prefix_len.item() if torch.is_tensor(prefix_len) else prefix_len
+                )
+
+                # Compute relative positions within items after delimiters
+                diff = pos[first_delim:] - torch.cummax(mask[first_delim:], 0)[1]
+                token_pos = (diff - pos[first_delim]).to(torch.uint16)
+                token_pos_in_items_ptr.append(token_pos)
+
+                # Update forward_batch positions in-place
+                pos[first_delim:] = diff - 1
+                forward_batch.positions[seq_start:seq_end] = pos
+
+            seq_start = seq_end
+
+        # Pad token_pos_in_items_ptr for batch processing
+        if token_pos_in_items_ptr:
+            token_pos_in_items_len = max(t.numel() for t in token_pos_in_items_ptr)
+            device = forward_batch.input_ids.device
+            token_pos_in_items_ptr = [
+                torch.cat(
+                    [
+                        t,
+                        torch.zeros(
+                            token_pos_in_items_len - t.numel(),
+                            dtype=torch.uint16,
+                            device=device,
+                        ),
+                    ]
+                )
+                for t in token_pos_in_items_ptr
+            ]
+
+        if not prefix_len_ptr or not token_pos_in_items_ptr:
+            return MultiItemScoringParams()
+
+        # Build final params
+        device = forward_batch.input_ids.device
+        return MultiItemScoringParams(
+            prefix_len_ptr=torch.tensor(
+                prefix_len_ptr, dtype=torch.uint32, device=device
+            ),
+            token_pos_in_items_ptr=torch.cat(token_pos_in_items_ptr, dim=0),
+            token_pos_in_items_len=token_pos_in_items_len & 0xFFFFFFFF,
+            max_item_len_ptr=torch.stack(
+                [
+                    t.to(torch.int32).max().to(torch.uint16)
+                    for t in token_pos_in_items_ptr
+                ],
+                dim=0,
+            ),
+        )
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         if forward_batch.forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
@@ -280,13 +455,26 @@ class FlashInferAttnBackend(AttentionBackend):
         else:
             prefix_lens = forward_batch.extend_prefix_lens
 
-            if self.is_multimodal:
+            # Disable ragged wrapper and ensure prefix handling for multimodal and multi-item scoring
+            if self.is_multimodal or self.multi_item_scoring_delimiter is not None:
+                # use_ragged = False: Multi-item scoring requires the paged wrapper because:
+                # 1. Ragged wrapper doesn't support the specialized multi-item parameters
+                #    (prefix_len_ptr, token_pos_in_items_ptr, etc.)
+                # 2. Paged wrapper provides better control over attention masking needed
+                #    for respecting item boundaries in multi-item sequences
+                # 3. Custom masking logic conflicts with ragged wrapper's assumptions
                 use_ragged = False
                 extend_no_prefix = False
             else:
                 use_ragged = not self.enable_deterministic
                 extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
 
+            # Process multi-item scoring in attention backend instead of ForwardBatch
+            multi_item_params = MultiItemScoringParams()
+            if self.multi_item_scoring_delimiter is not None:
+                # Use new backend-specific implementation
+                multi_item_params = self._process_multi_item_scoring(forward_batch)
+
             self.indices_updater_prefill.update(
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
@@ -298,9 +486,13 @@ class FlashInferAttnBackend(AttentionBackend):
                 encoder_lens=forward_batch.encoder_lens,
                 spec_info=None,
                 fixed_split_size=self.prefill_split_tile_size,
+                multi_item_params=multi_item_params,
             )
             self.forward_metadata = PrefillMetadata(
-                self.prefill_wrappers_paged, use_ragged, extend_no_prefix
+                self.prefill_wrappers_paged,
+                use_ragged,
+                extend_no_prefix,
+                multi_item_params,
             )
 
     def init_cuda_graph_state(
@@ -531,7 +723,20 @@ class FlashInferAttnBackend(AttentionBackend):
                 forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
                 causal=not layer.is_cross_attention,
                 sm_scale=layer.scaling,
-                window_left=layer.sliding_window_size,
+                # Disable sliding window attention for multi-item scoring:
+                # - Sliding window could cut across item boundaries, breaking semantic coherence
+                # - Multi-item sequences need full attention to properly handle delimiter tokens
+                # - Specialized multi-item parameters (prefix_len_ptr, token_pos_in_items_ptr)
+                #   provide more precise attention control than simple sliding windows
+                # - Item-aware masking takes precedence over window-based masking
+                window_left=(
+                    layer.sliding_window_size
+                    if not (
+                        self.forward_metadata.multi_item_params
+                        and self.forward_metadata.multi_item_params.is_enabled()
+                    )
+                    else -1
+                ),
                 logits_soft_cap=logits_soft_cap,
                 # Must use _float to avoid device-to-host copy that breaks cuda graph capture.
                 k_scale=layer.k_scale_float,
@@ -539,9 +744,13 @@ class FlashInferAttnBackend(AttentionBackend):
             )
         else:
             causal = True
-            if layer.attn_type == AttentionType.ENCODER_ONLY:
-                save_kv_cache = False
+            if (
+                layer.is_cross_attention
+                or layer.attn_type == AttentionType.ENCODER_ONLY
+            ):
                 causal = False
+            if save_kv_cache and layer.attn_type == AttentionType.ENCODER_ONLY:
+                save_kv_cache = False
 
             if self.forward_metadata.extend_no_prefix:
                 # NOTE: FlashInfer currently has limitations with head_dim = 32 or other dimensions
@@ -952,6 +1161,7 @@ class FlashInferIndicesUpdaterPrefill:
         encoder_lens: Optional[torch.Tensor],
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
+        multi_item_params: Optional[MultiItemScoringParams] = None,
     ):
         if use_ragged:
             # TODO: remove this device sync, we can use forward_batch.extend_prefix_lens_cpu
@@ -976,6 +1186,7 @@ class FlashInferIndicesUpdaterPrefill:
             use_ragged,
             spec_info,
             fixed_split_size=fixed_split_size,
+            multi_item_params=multi_item_params,
         )
 
     def update_sliding_window(
@@ -990,6 +1201,7 @@ class FlashInferIndicesUpdaterPrefill:
         encoder_lens: Optional[torch.Tensor],
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
+        multi_item_params: Optional[MultiItemScoringParams] = None,
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -1023,6 +1235,7 @@ class FlashInferIndicesUpdaterPrefill:
                 use_ragged,
                 spec_info,
                 use_sliding_window_kv_pool=use_sliding_window_kv_pool,
+                multi_item_params=multi_item_params,
             )
 
     def update_cross_attention(
@@ -1037,6 +1250,7 @@ class FlashInferIndicesUpdaterPrefill:
         encoder_lens: Optional[torch.Tensor],
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
+        multi_item_params: Optional[MultiItemScoringParams] = None,
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -1063,6 +1277,7 @@ class FlashInferIndicesUpdaterPrefill:
                 self.qo_indptr[wrapper_id],
                 use_ragged,
                 spec_info,
+                multi_item_params=multi_item_params,
             )
 
     def call_begin_forward(
@@ -1081,6 +1296,7 @@ class FlashInferIndicesUpdaterPrefill:
         spec_info: Optional[SpecInput],
         use_sliding_window_kv_pool: bool = False,
         fixed_split_size: Optional[int] = None,
+        multi_item_params: Optional[MultiItemScoringParams] = None,
     ):
         bs = len(seq_lens)
         if spec_info is None:
@@ -1136,6 +1352,22 @@ class FlashInferIndicesUpdaterPrefill:
             )
 
         # cached part
+        # Conditionally set multi-item parameters
+        if multi_item_params is not None and multi_item_params.is_enabled():
+            # Multi-item scoring is active - use specialized parameters and disable generic custom_mask
+            use_custom_mask = None
+            prefix_len_ptr = multi_item_params.prefix_len_ptr
+            token_pos_in_items_ptr = multi_item_params.token_pos_in_items_ptr
+            token_pos_in_items_len = multi_item_params.token_pos_in_items_len
+            max_item_len_ptr = multi_item_params.max_item_len_ptr
+        else:
+            # No multi-item scoring - use standard parameters
+            use_custom_mask = custom_mask
+            prefix_len_ptr = None
+            token_pos_in_items_ptr = None
+            token_pos_in_items_len = 0
+            max_item_len_ptr = None
+
         wrapper_paged.begin_forward(
             qo_indptr,
             kv_indptr,
@@ -1147,9 +1379,13 @@ class FlashInferIndicesUpdaterPrefill:
             1,
             q_data_type=self.q_data_type,
             kv_data_type=self.data_type,
-            custom_mask=custom_mask,
+            custom_mask=use_custom_mask,
             non_blocking=True,
             fixed_split_size=fixed_split_size,
+            prefix_len_ptr=prefix_len_ptr,
+            token_pos_in_items_ptr=token_pos_in_items_ptr,
+            token_pos_in_items_len=token_pos_in_items_len,
+            max_item_len_ptr=max_item_len_ptr,
         )
 
 
@@ -1185,7 +1421,7 @@ class FlashInferMultiStepDraftBackend:
             (max_bs,), dtype=torch.int32, device=model_runner.device
         )
         self.attn_backends: List[FlashInferAttnBackend] = []
-        for i in range(self.speculative_num_steps):
+        for i in range(self.speculative_num_steps - 1):
             self.attn_backends.append(
                 FlashInferAttnBackend(
                     model_runner,
@@ -1273,7 +1509,7 @@ class FlashInferMultiStepDraftBackend:
             device="cuda",
         )
 
-        for i in range(self.speculative_num_steps):
+        for i in range(self.speculative_num_steps - 1):
             self.attn_backends[i].init_cuda_graph_state(
                 max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i]
             )

sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -9,27 +9,20 @@ and uses BatchMLAPaged wrapper for decoding.
 More details can be found in https://docs.flashinfer.ai/api/mla.html
 """
 
-import os
 from dataclasses import dataclass
 from functools import partial
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import torch
 
-if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1":
-    import logging
-
-    torch._logging.set_logs(dynamo=logging.ERROR)
-    torch._dynamo.config.suppress_errors = True
-
-from sglang.global_config import global_config
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.dp_attention import get_attention_tp_size
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import (
     is_flashinfer_available,
@@ -38,10 +31,19 @@ from sglang.srt.utils import (
 )
 
 if TYPE_CHECKING:
+    from sglang.srt.layers.attention.flashinfer_mla_backend import (
+        FlashInferMlaAttnBackend,
+    )
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
     from sglang.srt.speculative.spec_info import SpecInput
 
+if envs.SGLANG_ENABLE_TORCH_COMPILE.get():
+    import logging
+
+    torch._logging.set_logs(dynamo=logging.ERROR)
+    torch._dynamo.config.suppress_errors = True
+
 if is_flashinfer_available():
     from flashinfer import (
         BatchMLAPagedAttentionWrapper,
@@ -66,7 +68,7 @@ global_workspace_buffer = None
 
 class FlashInferMhaChunkKVRunner:
     def __init__(
-        self, model_runner: ModelRunner, attn_backend: "FlashInferMlaAttnBackend"
+        self, model_runner: ModelRunner, attn_backend: FlashInferMlaAttnBackend
     ):
         # Parse Constants
         self.num_local_heads = (
@@ -80,6 +82,7 @@ class FlashInferMhaChunkKVRunner:
 
         # Buffers and wrappers
         self.qo_indptr = attn_backend.qo_indptr
+        self.kv_indptr = attn_backend.kv_indptr
         self.workspace_buffer = attn_backend.workspace_buffer
         self.fmha_backend = attn_backend.fmha_backend
 
@@ -130,9 +133,14 @@ class FlashInferMhaChunkKVRunner:
             )
         # ragged prefill
         if not disable_flashinfer_ragged:
+            kv_indptr = (
+                qo_indptr
+                if not forward_batch.mha_one_shot
+                else self.kv_indptr[: bs + 1]
+            )
             self.ragged_wrapper.begin_forward(
                 qo_indptr=qo_indptr,
-                kv_indptr=qo_indptr,
+                kv_indptr=kv_indptr,
                 num_qo_heads=self.num_local_heads,
                 num_kv_heads=self.num_local_heads,
                 head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
@@ -154,7 +162,7 @@ class FlashInferMhaChunkKVRunner:
             chunk_idx = forward_batch.prefix_chunk_idx
             assert chunk_idx >= 0
             wrapper = self.chunk_ragged_wrappers[chunk_idx]
-            o1, s1 = wrapper.forward_return_lse(
+            o = wrapper.forward_return_lse(
                 q.view(-1, layer.tp_q_head_num, layer.head_dim),
                 k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype),
                 v.view(-1, layer.tp_v_head_num, layer.v_head_dim).to(q.dtype),
@@ -163,7 +171,12 @@ class FlashInferMhaChunkKVRunner:
                 logits_soft_cap=logits_soft_cap,
             )
         else:
-            o1, s1 = self.ragged_wrapper.forward_return_lse(
+            forward = (
+                self.ragged_wrapper.forward_return_lse
+                if forward_batch.mha_return_lse
+                else self.ragged_wrapper.forward
+            )
+            o = forward(
                 q.view(-1, layer.tp_q_head_num, layer.head_dim),
                 k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype),
                 v.view(-1, layer.tp_v_head_num, layer.v_head_dim).to(q.dtype),
@@ -171,8 +184,7 @@ class FlashInferMhaChunkKVRunner:
                 sm_scale=layer.scaling,
                 logits_soft_cap=logits_soft_cap,
             )
-
-        return o1, s1
+        return o
 
 
 class FlashInferMLAAttnBackend(AttentionBackend):
@@ -193,9 +205,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         self.skip_prefill = skip_prefill
         self.enable_chunk_kv = (
             not skip_prefill
-            and global_server_args_dict["disaggregation_mode"] != "decode"
-            and not global_server_args_dict["disable_chunked_prefix_cache"]
-            and not global_server_args_dict["flashinfer_mla_disable_ragged"]
+            and get_global_server_args().disaggregation_mode != "decode"
+            and not get_global_server_args().disable_chunked_prefix_cache
+            and not get_global_server_args().flashinfer_mla_disable_ragged
         )
         self.page_size = model_runner.page_size
 
@@ -204,7 +216,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         if global_workspace_buffer is None:
             # different from flashinfer zero_init_global_workspace_buffer
             global_workspace_buffer = torch.empty(
-                global_config.flashinfer_workspace_size,
+                envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(),
                 dtype=torch.uint8,
                 device=model_runner.device,
             )
@@ -306,7 +318,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             prefix_lens = forward_batch.extend_prefix_lens
             extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
             use_ragged = (
-                not global_server_args_dict["flashinfer_mla_disable_ragged"]
+                not get_global_server_args().flashinfer_mla_disable_ragged
                 and extend_no_prefix
             )
 
@@ -510,15 +522,13 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         q_rope: Optional[torch.Tensor] = None,
         k_rope: Optional[torch.Tensor] = None,
     ):
-        if (
-            forward_batch.attn_attend_prefix_cache is not None
-            and forward_batch.mha_return_lse
+        if forward_batch.attn_attend_prefix_cache is not None and any(
+            forward_batch.extend_prefix_lens_cpu
         ):  # MHA Chunk
             assert self.enable_chunk_kv
             assert q_rope is None
             assert k_rope is None
-            o1, s1 = self.mha_chunk_kv_cache.forward(q, k, v, layer, forward_batch)
-            return o1, s1
+            return self.mha_chunk_kv_cache.forward(q, k, v, layer, forward_batch)
 
         cache_loc = forward_batch.out_cache_loc
         logits_soft_cap = layer.logit_cap
@@ -916,7 +926,7 @@ class FlashInferMLAMultiStepDraftBackend:
         )
 
         self.attn_backends = []
-        for i in range(self.speculative_num_steps):
+        for i in range(self.speculative_num_steps - 1):
             self.attn_backends.append(
                 FlashInferMLAAttnBackend(
                     model_runner,
@@ -998,7 +1008,7 @@ class FlashInferMLAMultiStepDraftBackend:
             device="cuda",
         )
 
-        for i in range(self.speculative_num_steps):
+        for i in range(self.speculative_num_steps - 1):
             self.attn_backends[i].init_cuda_graph_state(
                 max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i]
             )
@@ -1060,7 +1070,7 @@ def fast_mla_decode_plan(
 
     try:
         # Standard version with just the required arguments (no use_profiler)
-        self._cached_module.plan.default(
+        self._cached_module.plan(
             self._float_workspace_buffer,
             self._int_workspace_buffer,
             self._pin_memory_int_workspace_buffer,