sglang 0.5.0rc0__py3-none-any.whl → 0.5.0rc2__py3-none-any.whl

sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py

@@ -5,9 +5,9 @@ import logging
 import os
 import signal
 import threading
-from collections import OrderedDict
+from abc import ABC, abstractmethod
 from functools import wraps
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import torch
 
@@ -17,6 +17,75 @@ from sglang.srt.mem_cache.storage.hf3fs.client_hf3fs import Hf3fsClient
 logger = logging.getLogger(__name__)
 
 
+class Hf3fsMetadataInterface(ABC):
+    """Interface for HF3FS metadata operations."""
+
+    @abstractmethod
+    def initialize(self, rank: int, num_pages: int) -> None:
+        """Initialize the metadata service with specified number of pages."""
+        pass
+
+    @abstractmethod
+    def reserve_and_allocate_page_indices(
+        self,
+        rank: int,
+        keys: List[Tuple[str, str]],
+    ) -> List[Tuple[bool, int]]:
+        """
+        Reserve and allocate page indices for the specified keys.
+        Args:
+            rank: The rank of the process.
+            keys: The keys to reserve and allocate page indices for. Each tuple contains a key and the key of its prefix block.
+        Returns:
+            List[Tuple[bool, int]]: A list of tuples, where each tuple contains a boolean indicating whether the key has existed and an integer indicating the allocated page index.
+        """
+        pass
+
+    @abstractmethod
+    def confirm_write(
+        self,
+        rank: int,
+        written_keys_to_confirm: List[Tuple[str, int]],
+        pages_to_release: List[int],
+    ) -> None:
+        """
+        Confirm that key-value pairs have been successfully written to storage.
+        Args:
+            rank: The rank of the process.
+            written_keys_to_confirm: A list of tuples, where each tuple contains a key and its corresponding page index.
+            pages_to_release: A list of page indices to be released.
+        """
+        pass
+
+    @abstractmethod
+    def get_page_indices(self, rank: int, keys: List[str]) -> List[Optional[int]]:
+        """
+        Get page indices for the specified keys.
+        Args:
+            rank: The rank of the process.
+            keys: A list of keys.
+        Returns:
+            List[Optional[int]]: A list of integers representing the page indices for the specified keys.
+                                 If a key is not found, the corresponding index will be None.
+        """
+        pass
+
+    @abstractmethod
+    def delete_keys(self, rank: int, keys: List[str]) -> None:
+        """Delete specified keys and their associated pages."""
+        pass
+
+    @abstractmethod
+    def exists(self, rank: int, keys: List[str]) -> List[bool]:
+        """Check if the specified keys exist."""
+        pass
+
+    @abstractmethod
+    def clear(self, rank: int) -> None:
+        """Clear all key-value pairs and page allocations for the specified rank."""
+        pass
+
+
 class AtomicCounter:
     def __init__(self, n: int):
         assert n > 0
@@ -48,32 +117,32 @@ class HiCacheHF3FS(HiCacheStorage):
 
     def __init__(
         self,
+        rank: int,
         file_path: str,
         file_size: int,
         numjobs: int,
         bytes_per_page: int,
         entries: int,
         dtype: torch.dtype,
+        metadata_client: Hf3fsMetadataInterface,
     ):
+        self.rank = rank
         self.file_path = file_path
         self.file_size = file_size
         self.numjobs = numjobs
         self.bytes_per_page = bytes_per_page
         self.entries = entries
         self.dtype = dtype
+        self.metadata_client = metadata_client
 
         self.numel = self.bytes_per_page // self.dtype.itemsize
-
         self.num_pages = self.file_size // self.bytes_per_page
 
         logger.info(
-            "HiCacheHF3FS "
-            f"file_path = {self.file_path}, "
-            f"file_size = {self.file_size/(2**30):.2f} GB, "
-            f"numjobs = {self.numjobs}, "
-            f"bytes_per_page = {self.bytes_per_page/(2**20):.2f} MB, "
-            f"entries = {self.entries}, "
-            f"num_pages = {self.num_pages}"
+            f"[Rank {self.rank}] HiCacheHF3FS Client Initializing: "
+            f"file_path={self.file_path}, "
+            f"file_size={self.file_size / (2 ** 30):.2f} GB, "
+            f"num_pages={self.num_pages}"
         )
 
         self.ac = AtomicCounter(self.numjobs)
@@ -84,15 +153,11 @@ class HiCacheHF3FS(HiCacheStorage):
             for _ in range(numjobs)
         ]
         self.executor = concurrent.futures.ThreadPoolExecutor(
-            max_workers=self.numjobs, thread_name_prefix="HiCacheHF3FS"
+            max_workers=self.numjobs, thread_name_prefix=f"HiCacheHF3FS-Rank{self.rank}"
         )
 
-        # Implemented a preliminary single-file page_hash -> file_offset index as interim storage.
-        # Future iterations may adopt a global KVCache manager to coordinate external cache instances
-        # through centralized metadata orchestration.
+        self.metadata_client.initialize(self.rank, self.num_pages)
         self.lock = threading.RLock()
-        self.free_pages = list(range(self.num_pages))
-        self.key_to_index = OrderedDict()
 
         atexit.register(self.close)
 
@@ -104,15 +169,22 @@ class HiCacheHF3FS(HiCacheStorage):
     def from_env_config(
         rank: int, bytes_per_page: int, dtype: torch.dtype
     ) -> "HiCacheHF3FS":
+        from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import (
+            Hf3fsGlobalMetadataClient,
+            Hf3fsLocalMetadataClient,
+        )
+
         config_path = os.getenv(HiCacheHF3FS.default_env_var)
         if not config_path:
             return HiCacheHF3FS(
+                rank=rank,
                 file_path=f"/data/hicache.{rank}.bin",
                 file_size=1 << 40,
                 numjobs=16,
                 bytes_per_page=bytes_per_page,
                 entries=8,
                 dtype=dtype,
+                metadata_client=Hf3fsLocalMetadataClient(),
             )
 
         try:
@@ -121,6 +193,7 @@ class HiCacheHF3FS(HiCacheStorage):
         except Exception as e:
             raise RuntimeError(f"Failed to load config from {config_path}: {str(e)}")
 
+        # Check required keys (metadata_server_url is now optional)
         required_keys = {
             "file_path_prefix",
             "file_size",
@@ -131,19 +204,33 @@ class HiCacheHF3FS(HiCacheStorage):
         if missing_keys:
             raise ValueError(f"Missing required keys in config: {missing_keys}")
 
+        # Choose metadata client based on configuration
+        if "metadata_server_url" in config and config["metadata_server_url"]:
+            # Use global metadata client to connect to metadata server
+            metadata_server_url = config["metadata_server_url"]
+            metadata_client = Hf3fsGlobalMetadataClient(metadata_server_url)
+            logger.info(
+                f"Using global metadata client with server url: {metadata_server_url}"
+            )
+        else:
+            # Use local metadata client for single-machine deployment
+            metadata_client = Hf3fsLocalMetadataClient()
+
         return HiCacheHF3FS(
+            rank=rank,
             file_path=f"{config['file_path_prefix']}.{rank}.bin",
             file_size=int(config["file_size"]),
             numjobs=int(config["numjobs"]),
             bytes_per_page=bytes_per_page,
             entries=int(config["entries"]),
             dtype=dtype,
+            metadata_client=metadata_client,
         )
 
     def get(
         self, key: str, target_location: Optional[torch.Tensor] = None
     ) -> torch.Tensor | None:
-        return self.batch_get([key], target_location)[0]
+        return self.batch_get([key], [target_location] if target_location else None)[0]
 
     @synchronized()
     def batch_get(
@@ -151,14 +238,14 @@ class HiCacheHF3FS(HiCacheStorage):
         keys: List[str],
         target_locations: Optional[List[torch.Tensor]] = None,
     ) -> List[torch.Tensor | None]:
+        page_indices = self.metadata_client.get_page_indices(self.rank, keys)
+
         batch_indices, file_offsets = [], []
-        for i, key in enumerate(keys):
-            if key not in self.key_to_index:
-                continue
-            batch_indices.append(i)
-            file_offsets.append(self.key_to_index[key] * self.bytes_per_page)
-            self.key_to_index.move_to_end(key)
-        # TODO: target_locations
+        for i, page_index in enumerate(page_indices):
+            if page_index is not None:
+                batch_indices.append(i)
+                file_offsets.append(page_index * self.bytes_per_page)
+
         file_results = [
             torch.empty(self.numel, dtype=self.dtype) for _ in range(len(batch_indices))
         ]
@@ -180,7 +267,9 @@ class HiCacheHF3FS(HiCacheStorage):
             if read_result == self.bytes_per_page:
                 results[batch_index] = file_result
             else:
-                logger.error(f"HiCacheHF3FS get {keys[batch_index]} failed")
+                logger.error(
+                    f"[Rank {self.rank}] HiCacheHF3FS get {keys[batch_index]} failed"
+                )
 
         return results
 
@@ -188,13 +277,21 @@ class HiCacheHF3FS(HiCacheStorage):
         return self.batch_set([key], [value])
 
     def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
-        indices = self.get_batch_set_indices(keys)
+        # Todo: Add prefix block's hash key
+        key_with_prefix = [(key, "") for key in keys]
+        indices = self.metadata_client.reserve_and_allocate_page_indices(
+            self.rank, key_with_prefix
+        )
+
         batch_indices, file_offsets, file_values = [], [], []
-        for i, (value, (is_written, index)) in enumerate(zip(values, indices)):
-            if is_written or index == -1:
+        pages_to_release = []
+
+        for i, (value, (is_written, page_index)) in enumerate(zip(values, indices)):
+            if is_written or page_index == -1:
                 continue
+
             batch_indices.append(i)
-            file_offsets.append(index * self.bytes_per_page)
+            file_offsets.append(page_index * self.bytes_per_page)
             file_values.append(value.contiguous())
 
         futures = [
@@ -211,62 +308,37 @@ class HiCacheHF3FS(HiCacheStorage):
             for result in future.result()
         ]
 
+        written_keys_to_confirm = []
         results = [index[0] for index in indices]
         for batch_index, write_result in zip(batch_indices, write_results):
             key = keys[batch_index]
-            index = indices[batch_index][1]
+            page_index = indices[batch_index][1]
             if write_result:
-                self.key_to_index[key] = index
-                self.key_to_index.move_to_end(key)
+                written_keys_to_confirm.append((key, page_index))
             else:
-                logger.error(f"HiCacheHF3FS set {key} failed")
-                self.free_pages.append(index)
+                logger.error(f"[Rank {self.rank}] HiCacheHF3FS set {key} failed")
+                pages_to_release.append(page_index)
             results[batch_index] = write_result
-        return all(results)
-
-    @synchronized()
-    def get_batch_set_indices(self, keys: List[str]) -> list:
-        ionum = len(keys)
-        # results: tuples of (is_written: bool, page_idx: int)
-        # - is_written: True = hit (no I/O), False = write (miss)
-        # - page_idx: page storing data
-        results = [None] * min(ionum, self.num_pages)
-        if ionum > self.num_pages:
-            results.extend([(False, -1)] * (ionum - self.num_pages))
-
-        new_keys = []
-        for batch_index, key in enumerate(keys[: self.num_pages]):
-            if key in self.key_to_index:
-                results[batch_index] = (True, self.key_to_index[key])
-                self.key_to_index.move_to_end(key)
-            else:
-                new_keys.append((batch_index, key))
 
-        for batch_index, _ in new_keys:
-            index = (
-                self.free_pages.pop()
-                if len(self.free_pages) > 0
-                else self.key_to_index.popitem(last=False)[1]
+        if len(written_keys_to_confirm) > 0 or len(pages_to_release) > 0:
+            self.metadata_client.confirm_write(
+                self.rank, written_keys_to_confirm, pages_to_release
             )
-            results[batch_index] = (False, index)
 
-        return results
+        return all(results)
 
     @synchronized()
     def delete(self, key: str) -> None:
-        if key not in self.key_to_index:
-            return
-        index = self.key_to_index.pop(key)
-        self.free_pages.append(index)
+        self.metadata_client.delete_keys(self.rank, [key])
 
     @synchronized()
     def exists(self, key: str) -> bool:
-        return key in self.key_to_index
+        result = self.metadata_client.exists(self.rank, [key])
+        return result[0] if result else False
 
     @synchronized()
     def clear(self) -> None:
-        self.free_pages = list(range(self.num_pages))
-        self.key_to_index.clear()
+        self.metadata_client.clear(self.rank)
 
     def close(self) -> None:
         try:

sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -18,13 +18,12 @@ DEFAULT_LOCAL_BUFFER_SIZE = 128 * 1024 * 1024  # 128 MB
 logger = logging.getLogger(__name__)
 
 
-def get_hash_str_mooncake(current_page_ids: List, prefix_block_key: str):
+def get_hash_str_mooncake(token_ids: List[int], prior_hash: str = None):
     local_rank = get_tensor_model_parallel_rank()
     prefix_str = ""
-    if prefix_block_key:
-        if len(prefix_block_key):
-            prefix_str = hashlib.sha256(prefix_block_key.encode()).hexdigest()
-    current_token_ids_bytes = np.array(current_page_ids).tobytes()
+    if prior_hash:
+        prefix_str = hashlib.sha256(prior_hash.encode()).hexdigest()
+    current_token_ids_bytes = np.array(token_ids).tobytes()
     current_hash_object = hashlib.sha256(current_token_ids_bytes)
     current_hash_hex = current_hash_object.hexdigest()
     return f"{prefix_str}_{int(current_hash_hex[:16], 16)}_{local_rank}"
@@ -224,13 +223,11 @@ class MooncakeStore(HiCacheStorage):
 
     def exists(self, keys) -> bool | dict:
         _keys = []
-        local_rank = torch.cuda.current_device()
         for key in keys:
             if key is None:
                 return None
-            # Since mooncake store is stored in layer by layer,
-            # only the first layer is checked here.
-            _keys.append(f"{key}_{local_rank}_k")
+
+            _keys.append(f"{key}_k")
         result = {k: v for k, v in zip(keys, self.store.batch_is_exist(_keys))}
         return result
 

sglang/srt/model_executor/cuda_graph_runner.py

@@ -33,7 +33,12 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
 )
 from sglang.srt.distributed.parallel_state import GroupCoordinator, graph_capture
-from sglang.srt.layers.dp_attention import DPPaddingMode, get_attention_tp_size
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    get_attention_tp_rank,
+    get_attention_tp_size,
+    set_dp_buffer_len,
+)
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.torchao_utils import save_gemlite_cache
 from sglang.srt.model_executor.forward_batch_info import (
@@ -255,6 +260,9 @@ class CudaGraphRunner:
         self.dp_size = model_runner.server_args.dp_size
         self.pp_size = model_runner.server_args.pp_size
 
+        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_rank = get_attention_tp_rank()
+
         # Batch sizes to capture
         self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
         rank0_log(f"Capture cuda graph bs {self.capture_bs}")
@@ -342,30 +350,15 @@ class CudaGraphRunner:
                     self.global_num_tokens_for_logprob_gpu = torch.zeros(
                         (self.dp_size,), dtype=torch.int32
                     )
-                    self.gathered_buffer = torch.zeros(
-                        (
-                            self.max_num_token * self.dp_size,
-                            self.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.model_runner.dtype,
-                    )
                 else:
                     assert self.require_attn_tp_gather
                     self.global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
                     self.global_num_tokens_for_logprob_gpu = torch.zeros(
                         (1,), dtype=torch.int32
                     )
-                    self.gathered_buffer = torch.zeros(
-                        (
-                            self.max_num_token,
-                            self.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.model_runner.dtype,
-                    )
             else:
                 self.global_num_tokens_gpu = None
                 self.global_num_tokens_for_logprob_gpu = None
-                self.gathered_buffer = None
 
             self.custom_mask = torch.ones(
                 (
@@ -549,7 +542,7 @@ class CudaGraphRunner:
                     device=input_ids.device,
                 )
             )
-            gathered_buffer = self.gathered_buffer[: num_tokens * self.dp_size]
+            global_dp_buffer_len = num_tokens * self.dp_size
         elif self.require_attn_tp_gather:
             self.global_num_tokens_gpu.copy_(
                 torch.tensor(
@@ -565,9 +558,9 @@ class CudaGraphRunner:
                     device=input_ids.device,
                 )
             )
-            gathered_buffer = self.gathered_buffer[:num_tokens]
+            global_dp_buffer_len = num_tokens
         else:
-            gathered_buffer = None
+            global_dp_buffer_len = None
 
         spec_info = self.get_spec_info(num_tokens)
         if self.capture_hidden_mode != CaptureHiddenMode.FULL:
@@ -600,8 +593,8 @@ class CudaGraphRunner:
             positions=positions,
             global_num_tokens_gpu=self.global_num_tokens_gpu,
             global_num_tokens_for_logprob_gpu=self.global_num_tokens_for_logprob_gpu,
-            dp_padding_mode=DPPaddingMode.get_default_mode_in_cuda_graph(),
-            gathered_buffer=gathered_buffer,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
             mrope_positions=mrope_positions,
             spec_algorithm=self.model_runner.spec_algorithm,
             spec_info=spec_info,
@@ -630,6 +623,7 @@ class CudaGraphRunner:
         def run_once():
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(global_dp_buffer_len, num_tokens)
 
             kwargs = {}
             if (
@@ -729,10 +723,12 @@ class CudaGraphRunner:
         self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
         self.positions[:raw_num_token].copy_(forward_batch.positions)
 
+        seq_lens_cpu = None
         if forward_batch.seq_lens_cpu is not None:
             if bs != raw_bs:
                 self.seq_lens_cpu.fill_(self.seq_len_fill_value)
             self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+            seq_lens_cpu = self.seq_lens_cpu[:bs]
 
         if pp_proxy_tensors:
             for key in self.pp_proxy_tensors.keys():
@@ -747,7 +743,17 @@ class CudaGraphRunner:
             self.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
             self.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
         if enable_num_token_non_padded(self.model_runner.server_args):
-            self.num_token_non_padded.copy_(forward_batch.num_token_non_padded)
+            num_token_non_padded = forward_batch.num_token_non_padded
+            if self.require_gathered_buffer:
+                tokens_per_rank = bs // self.attn_tp_size * self.num_tokens_per_bs
+                num_local_token_non_padded = torch.clamp(
+                    num_token_non_padded - tokens_per_rank * self.attn_tp_rank,
+                    min=0,
+                    max=tokens_per_rank,
+                )
+                self.num_token_non_padded.copy_(num_local_token_non_padded)
+            else:
+                self.num_token_non_padded.copy_(num_token_non_padded)
         if self.enable_two_batch_overlap:
             self.tbo_plugin.replay_prepare(
                 forward_mode=self.capture_forward_mode,
@@ -766,7 +772,7 @@ class CudaGraphRunner:
             self.encoder_lens[:bs] if self.is_encoder_decoder else None,
             self.capture_forward_mode,
             forward_batch.spec_info,
-            seq_lens_cpu=self.seq_lens_cpu[:bs],
+            seq_lens_cpu=seq_lens_cpu,
         )
 
         # Store fields

sglang/srt/model_executor/forward_batch_info.py

@@ -40,9 +40,10 @@ import triton.language as tl
 
 from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size
 from sglang.srt.layers.dp_attention import (
-    DPPaddingMode,
+    DpPaddingMode,
     get_attention_dp_rank,
     get_attention_tp_size,
+    set_dp_buffer_len,
 )
 from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
 from sglang.srt.utils import (
@@ -274,13 +275,13 @@ class ForwardBatch:
     global_num_tokens_for_logprob_cpu: Optional[List[int]] = None
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
     # The padding mode for DP attention
-    dp_padding_mode: Optional[DPPaddingMode] = None
+    dp_padding_mode: Optional[DpPaddingMode] = None
     # for extend, local start pos and num tokens is different in logits processor
     # this will be computed in get_dp_local_info
     # this will be recomputed in LogitsMetadata.from_forward_batch
     dp_local_start_pos: Optional[torch.Tensor] = None  # cached info at runtime
     dp_local_num_tokens: Optional[torch.Tensor] = None  # cached info at runtime
-    gathered_buffer: Optional[torch.Tensor] = None
+    global_dp_buffer_len: Optional[int] = None
     is_extend_in_batch: bool = False
     can_run_dp_cuda_graph: bool = False
     global_forward_mode: Optional[ForwardMode] = None
@@ -628,7 +629,7 @@ class ForwardBatch:
                 (global_num_tokens[i] - 1) // attn_tp_size + 1
             ) * attn_tp_size
 
-        dp_padding_mode = DPPaddingMode.get_dp_padding_mode(global_num_tokens)
+        dp_padding_mode = DpPaddingMode.get_dp_padding_mode(global_num_tokens)
         self.dp_padding_mode = dp_padding_mode
 
         if dp_padding_mode.is_max_len():
@@ -642,23 +643,38 @@ class ForwardBatch:
         else:
             buffer_len = sum(global_num_tokens)
 
-        self.gathered_buffer = torch.zeros(
-            (buffer_len, model_runner.model_config.hidden_size),
-            dtype=model_runner.dtype,
-            device=model_runner.device,
-        )
-
         if len(global_num_tokens) > 1:
             num_tokens = global_num_tokens[get_attention_dp_rank()]
         else:
             num_tokens = global_num_tokens[0]
 
-        if self.forward_mode.is_decode():
-            setattr(self, "raw_bs", self.batch_size)
-            self.batch_size = num_tokens
+        self.global_dp_buffer_len = buffer_len
+        set_dp_buffer_len(buffer_len, num_tokens)
 
         bs = self.batch_size
 
+        if self.forward_mode.is_decode():
+            if self.is_extend_in_batch and dp_padding_mode.is_max_len():
+                setattr(self, "_original_forward_mode", self.forward_mode)
+                self.forward_mode = ForwardMode.EXTEND
+                self.extend_num_tokens = bs
+                self.extend_seq_lens = torch.full_like(self.seq_lens, 1)
+                self.extend_prefix_lens = self.seq_lens - 1
+                self.extend_start_loc = torch.arange(
+                    bs, dtype=torch.int32, device=self.seq_lens.device
+                )
+                self.extend_prefix_lens_cpu = self.extend_prefix_lens.cpu()
+                self.extend_seq_lens_cpu = self.extend_seq_lens.cpu()
+                self.extend_logprob_start_lens_cpu = self.extend_prefix_lens_cpu
+            else:
+                setattr(self, "_original_batch_size", self.batch_size)
+                if self.spec_info is not None:
+                    bs = self.batch_size = (
+                        num_tokens // self.spec_info.num_tokens_per_batch
+                    )
+                else:
+                    bs = self.batch_size = num_tokens
+
         # padding
         self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens)
         self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs)
@@ -689,6 +705,7 @@ class ForwardBatch:
         if self.mrope_positions is not None:
             self.mrope_positions = self._pad_tensor_to_size(self.mrope_positions, bs)
 
+        # TODO: check if we need to pad other tensors
         if self.extend_seq_lens is not None:
             self.extend_seq_lens = self._pad_tensor_to_size(self.extend_seq_lens, bs)
 
@@ -712,7 +729,9 @@ class ForwardBatch:
 
     def post_forward_mlp_sync_batch(self, logits_output: LogitsProcessorOutput):
 
-        bs = getattr(self, "raw_bs", self.batch_size)
+        self.forward_mode = getattr(self, "_original_forward_mode", self.forward_mode)
+        self.batch_size = getattr(self, "_original_batch_size", self.batch_size)
+        bs = self.batch_size
 
         if self.spec_info is not None:
             if self.forward_mode.is_decode():  # draft