sglang 0.4.10.post2__py3-none-any.whl → 0.5.0rc1__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,11 @@ 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,
+)
 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 +259,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}")
@@ -576,11 +583,11 @@ class CudaGraphRunner:
             )
 
         if self.model_runner.server_args.enable_lora:
-            # It is safe to capture CUDA graph using empty LoRA path, as the LoRA kernels will always be launched whenever
-            # `--enable-lora` is set to True (and return immediately if the LoRA path is empty for perf optimization).
-            lora_paths = [None] * bs
+            # It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever
+            # `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization).
+            lora_ids = [None] * bs
         else:
-            lora_paths = None
+            lora_ids = None
 
         forward_batch = ForwardBatch(
             forward_mode=self.capture_forward_mode,
@@ -589,6 +596,7 @@ class CudaGraphRunner:
             req_pool_indices=req_pool_indices,
             seq_lens=seq_lens,
             next_token_logits_buffer=next_token_logits_buffer,
+            orig_seq_lens=seq_lens,
             req_to_token_pool=self.model_runner.req_to_token_pool,
             token_to_kv_pool=self.model_runner.token_to_kv_pool,
             attn_backend=self.model_runner.attn_backend,
@@ -607,11 +615,11 @@ class CudaGraphRunner:
             capture_hidden_mode=self.capture_hidden_mode,
             num_token_non_padded=self.num_token_non_padded,
             global_forward_mode=self.capture_forward_mode,
-            lora_paths=lora_paths,
+            lora_ids=lora_ids,
         )
         self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
 
-        if lora_paths is not None:
+        if lora_ids is not None:
             self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
 
         # Attention backend
@@ -728,10 +736,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():
@@ -746,7 +756,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,
@@ -765,7 +785,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

@@ -180,6 +180,9 @@ class ForwardBatch:
     # The sum of all sequence lengths
     seq_lens_sum: int
 
+    # The original sequence length without being chunked. Qwen-1M related.
+    orig_seq_lens: Optional[torch.Tensor] = None
+
     # Optional seq_lens on cpu
     seq_lens_cpu: Optional[torch.Tensor] = None
 
@@ -248,7 +251,7 @@ class ForwardBatch:
     encoder_out_cache_loc: Optional[torch.Tensor] = None
 
     # For LoRA
-    lora_paths: Optional[List[str]] = None
+    lora_ids: Optional[List[str]] = None
 
     # For input embeddings
     input_embeds: Optional[torch.Tensor] = None
@@ -321,13 +324,14 @@ class ForwardBatch:
             encoder_out_cache_loc=batch.encoder_out_cache_loc,
             seq_lens_sum=batch.seq_lens_sum,
             seq_lens_cpu=batch.seq_lens_cpu,
+            orig_seq_lens=batch.orig_seq_lens,
             return_logprob=batch.return_logprob,
             top_logprobs_nums=batch.top_logprobs_nums,
             token_ids_logprobs=batch.token_ids_logprobs,
             is_extend_in_batch=batch.is_extend_in_batch,
             can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph,
             global_forward_mode=batch.global_forward_mode,
-            lora_paths=batch.lora_paths,
+            lora_ids=batch.lora_ids,
             sampling_info=batch.sampling_info,
             req_to_token_pool=model_runner.req_to_token_pool,
             token_to_kv_pool=model_runner.token_to_kv_pool,
@@ -420,16 +424,12 @@ class ForwardBatch:
                 batch.extend_prefix_lens, dtype=torch.int32
             ).to(device, non_blocking=True)
             ret.extend_num_tokens = batch.extend_num_tokens
-            if support_triton(model_runner.server_args.attention_backend):
-                positions, ret.extend_start_loc = compute_position_triton(
-                    ret.extend_prefix_lens,
-                    ret.extend_seq_lens,
-                    ret.extend_num_tokens,
-                )
-            else:
-                positions, ret.extend_start_loc = compute_position_torch(
-                    ret.extend_prefix_lens, ret.extend_seq_lens
-                )
+            positions, ret.extend_start_loc = compute_position(
+                model_runner.server_args.attention_backend,
+                ret.extend_prefix_lens,
+                ret.extend_seq_lens,
+                ret.extend_num_tokens,
+            )
             if ret.positions is None:
                 ret.positions = positions
             ret.extend_prefix_lens_cpu = batch.extend_prefix_lens
@@ -632,8 +632,10 @@ class ForwardBatch:
         self.dp_padding_mode = dp_padding_mode
 
         if dp_padding_mode.is_max_len():
-            # when DP gather mode is all gather, we will use all_gather_into_tensor to gather hidden states,
-            # where transferred tokens should be padded to the same length.
+            # when DP gather mode is all gather, we will use
+            # all_gather_into_tensor to gather hidden states, where transferred
+            # tokens should be padded to the same length. We will also use
+            # reduce-scatter instead of all-reduce after MLP.
             max_num_tokens = max(global_num_tokens)
             global_num_tokens = [max_num_tokens] * sync_group_size
             buffer_len = max_num_tokens * sync_group_size
@@ -651,12 +653,30 @@ class ForwardBatch:
         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
-
         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)
@@ -687,6 +707,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)
 
@@ -710,7 +731,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
@@ -882,6 +905,25 @@ class PPProxyTensors:
         return f"PPProxyTensors(tensors={self.tensors})"
 
 
+def compute_position(
+    attn_backend: str,
+    extend_prefix_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    extend_seq_lens_sum: int,
+):
+    if support_triton(attn_backend):
+        positions, extend_start_loc = compute_position_triton(
+            extend_prefix_lens,
+            extend_seq_lens,
+            extend_seq_lens_sum,
+        )
+    else:
+        positions, extend_start_loc = compute_position_torch(
+            extend_prefix_lens, extend_seq_lens
+        )
+    return positions, extend_start_loc
+
+
 def compute_position_triton(
     extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum
 ):