sglang 0.4.10.post1__py3-none-any.whl → 0.4.10.post2__py3-none-any.whl

sglang/srt/mem_cache/cpp_radix_tree/radix_tree.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import torch
+from torch.utils.cpp_extension import load
+
+_abs_path = os.path.dirname(os.path.abspath(__file__))
+radix_tree_cpp = load(
+    name="radix_tree_cpp",
+    sources=[
+        f"{_abs_path}/tree_v2_binding.cpp",
+        f"{_abs_path}/tree_v2_debug.cpp",
+        f"{_abs_path}/tree_v2.cpp",
+    ],
+    extra_cflags=["-O3", "-std=c++20"],
+)
+
+if TYPE_CHECKING:
+
+    class TreeNodeCpp:
+        """
+        A placeholder for the TreeNode class. Cannot be constructed elsewhere.
+        """
+
+    class IOHandle:
+        """
+        A placeholder for the IOHandle class. Cannot be constructed elsewhere.
+        """
+
+    class RadixTreeCpp:
+        def __init__(
+            self,
+            disabled: bool,
+            host_size: Optional[int],
+            page_size: int,
+            write_through_threshold: int,
+        ):
+            """
+            Initializes the RadixTreeCpp instance.
+            Args:
+                disabled (bool): If True, the radix tree is disabled.
+                host_size (Optional[int]): Size of the radix tree on the CPU. None means no CPU tree.
+                page_size (int): Size of the page for the radix tree.
+                write_through_threshold (int): Threshold for writing through from GPU to CPU.
+            """
+            self.tree = radix_tree_cpp.RadixTree(  # type: ignore
+                disabled, host_size, page_size, write_through_threshold
+            )
+
+        def match_prefix(
+            self, prefix: List[int]
+        ) -> Tuple[List[torch.Tensor], int, TreeNodeCpp, TreeNodeCpp]:
+            """
+            Matches a prefix in the radix tree.
+            Args:
+                prefix (List[int]): The prefix to match.
+            Returns:
+                Tuple[List[torch.Tensor], TreeNodeCpp, TreeNodeCpp]:
+                    0. A list of indices that is matched by the prefix on the GPU.
+                    1. Sum length of the indices matched on the CPU.
+                    2. The last node of the prefix matched on the GPU.
+                    3. The last node of the prefix matched on the CPU.
+            """
+            return self.tree.match_prefix(prefix)
+
+        def evict(self, num_tokens: int) -> List[torch.Tensor]:
+            """
+            Evicts a number of tokens from the radix tree.
+            Args:
+                num_tokens (int): The number of tokens to evict.
+            Returns:
+                List[torch.Tensor]: A list of indices that were evicted.
+            """
+            return self.tree.evict(num_tokens)
+
+        def lock_ref(self, handle: TreeNodeCpp, lock: bool) -> None:
+            """
+            Locks or unlocks a reference to a tree node.
+            After locking, the node will not be evicted from the radix tree.
+            Args:
+                handle (TreeNodeCpp): The tree node to lock or unlock.
+                lock (bool): If True, locks the node; if False, unlocks it.
+            """
+            return self.tree.lock_ref(handle, lock)
+
+        def writing_through(
+            self, key: List[int], indices: torch.Tensor
+        ) -> Tuple[List[Tuple[IOHandle, torch.Tensor, torch.Tensor]], int]:
+            """
+            Inserts a key-value pair into the radix tree and perform write-through check.
+            Args:
+                key (List[int]): The key to insert.
+                indices (torch.Tensor): The value associated with the key.
+            Returns:
+                Tuple[List[Tuple[IOHandle, torch.Tensor, torch.Tensor]], int]:
+                    0. A list of (IOHandle, device indices, host indices) tuples.
+                       These IOhandles require write-through to the CPU in python side.
+                    1. The number of indices that are matched on device.
+            """
+            return self.tree.writing_through(key, indices)
+
+        def loading_onboard(
+            self,
+            host_node: TreeNodeCpp,
+            new_device_indices: torch.Tensor,
+        ) -> Tuple[IOHandle, List[torch.Tensor]]:
+            """
+            Updates the device indices of tree nodes within a range on the tree.
+            Args:
+                host_node (TreeNodeCpp): The tree node on the host, must be descendant of device_node.
+                new_device_indices (torch.Tensor): The new device indices to set.
+                    The length of this tensor must be exactly host indices length.
+            Returns:
+                Tuple[IOHandle, List[torch.Tensor]]:
+                    0. An IOHandle that requires loading to the CPU in python side.
+                    1. A list of host indices corresponding to the new device indices.
+            """
+            return self.tree.loading_onboard(host_node, new_device_indices)
+
+        def commit_writing_through(self, handle: IOHandle, success: bool) -> None:
+            """
+            Commits the write-through process for a tree node.
+            Args:
+                handle (IOHandle): The IOHandle to commit.
+                success (bool): If True, commits the write-through; if False, just indicates failure.
+            """
+            return self.tree.commit_writing_through(handle, success)
+
+        def commit_loading_onboard(self, handle: IOHandle, success: bool) -> None:
+            """
+            Commits the load onboard process for tree nodes within a range on the tree.
+            Args:
+                handle (IOHandle): The IOHandle to commit.
+                success (bool): If True, commits the load-onboard; if False, just indicates failure.
+            """
+            return self.tree.commit_loading_onboard(handle, success)
+
+        def evictable_size(self) -> int:
+            """
+            Returns the size of the evictable part of the radix tree.
+            This is the size of the part that can be evicted from the GPU (ref_count = 0).
+            Returns:
+                int: The size of the evictable part.
+            """
+            return self.tree.evictable_size()
+
+        def protected_size(self) -> int:
+            """
+            Returns the size of the protected part of the radix tree.
+            This is the size of the part that cannot be evicted from the GPU (ref_count > 0).
+            Returns:
+                int: The size of the protected part.
+            """
+            return self.tree.protected_size()
+
+        def total_size(self) -> int:
+            """
+            Returns the total size of the radix tree (including CPU nodes).
+            Returns:
+                int: The total size of the radix tree.
+            """
+            return self.tree.total_size()
+
+        def reset(self) -> None:
+            """
+            Resets the radix tree, clearing all nodes and indices.
+            """
+            return self.tree.reset()
+
+        def debug_print(self) -> None:
+            """
+            Prints the internal state of the radix tree for debugging purposes.
+            """
+            return self.tree.debug_print()
+
+else:
+    # Real implementation of the classes for runtime
+    RadixTreeCpp = radix_tree_cpp.RadixTree
+    TreeNodeCpp = object
+    IOHandle = object

sglang/srt/mem_cache/hicache_storage.py

@@ -33,8 +33,7 @@ class HiCacheStorage(ABC):
     It abstracts the underlying storage mechanism, allowing different implementations to be used.
     """
 
-    # todo, translate tensor object access for different TP ranks
-    # potentially pass model and TP configs into storage backend
+    # todo, potentially pass model and TP configs into storage backend
     # todo, the page size of storage backend does not have to be the same as the same as host memory pool
 
     @abstractmethod
@@ -117,35 +116,28 @@ class HiCacheFile(HiCacheStorage):
     def get(
         self,
         key: str,
-        target_location: Optional[Any] = None,
+        target_location: torch.Tensor,
         target_sizes: Optional[Any] = None,
     ) -> torch.Tensor | None:
         key = self._get_suffixed_key(key)
         tensor_path = os.path.join(self.file_path, f"{key}.bin")
         try:
-            if target_location is not None:
-                # Load directly into target_location's memory buffer
-                with open(tensor_path, "rb") as f:
-                    target_location.set_(
-                        torch.frombuffer(f.read(), dtype=target_location.dtype)
-                        .reshape(target_location.shape)
-                        .storage()
-                    )
-                return target_location
-            else:
-                loaded_tensor = torch.load(tensor_path)
-                if isinstance(loaded_tensor, torch.Tensor):
-                    return loaded_tensor
-                else:
-                    logger.error(f"Loaded data for key {key} is not a tensor.")
-                    return None
+            # Load directly into target_location's memory buffer
+            with open(tensor_path, "rb") as f:
+                target_location.set_(
+                    torch.frombuffer(f.read(), dtype=target_location.dtype)
+                    .reshape(target_location.shape)
+                    .untyped_storage()
+                )
+            return target_location
         except FileNotFoundError:
+            logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
             return None
 
     def batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[Any] = None,
+        target_locations: List[torch.Tensor],
         target_sizes: Optional[Any] = None,
     ) -> List[torch.Tensor | None]:
         return [
@@ -168,7 +160,7 @@ class HiCacheFile(HiCacheStorage):
             logger.debug(f"Key {key} already exists. Skipped.")
             return True
         try:
-            torch.save(value, tensor_path)
+            value.contiguous().view(dtype=torch.uint8).numpy().tofile(tensor_path)
             return True
         except Exception as e:
             logger.error(f"Failed to save tensor {key}: {e}")

sglang/srt/mem_cache/radix_cache_cpp.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Set
+
+import torch
+
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
+from sglang.srt.mem_cache.cpp_radix_tree.radix_tree import (
+    IOHandle,
+    RadixTreeCpp,
+    TreeNodeCpp,
+)
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+logger = logging.getLogger(__name__)
+
+
+class RadixCacheCpp(BasePrefixCache):
+    def _merge_tensor(self, l: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Merge a list of tensors into a single tensor.
+        Args:
+            l (List[torch.Tensor]): List of tensors to merge.
+        Returns:
+            torch.Tensor: Merged tensor.
+        """
+        if len(l) == 0:
+            return torch.empty(0, dtype=torch.int64, device=self.device)
+        elif len(l) == 1:
+            return l[0]
+        else:
+            return torch.cat(l)
+
+    def __init__(
+        self,
+        disable: bool,
+        use_hicache: bool,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool: BaseTokenToKVPoolAllocator,
+        tp_cache_group: torch.distributed.ProcessGroup,
+        page_size: int,
+        hicache_ratio: float,
+        hicache_size: int,
+        hicache_write_policy: str,
+        enable_kv_cache_events: bool = False,
+        hicache_oracle: bool = False,
+        enable_write_cancel: bool = False,
+    ):
+        self.disable = disable
+        self.enable_write_cancel = enable_write_cancel
+
+        assert (
+            enable_kv_cache_events is False
+        ), "HiRadixCache does not support kv cache events yet"
+        self.kv_cache = token_to_kv_pool.get_kvcache()
+
+        # record the nodes with ongoing write through
+        self.ongoing_write_through: Set[IOHandle] = set()
+        # record the node segments with ongoing load back
+        self.ongoing_load_back: Set[IOHandle] = set()
+        # todo: dynamically adjust the threshold
+        self.write_through_threshold = (
+            1 if hicache_write_policy == "write_through" else 2
+        )
+        self.device = token_to_kv_pool.device
+        self.token_to_kv_pool = token_to_kv_pool
+        self.req_to_token_pool = req_to_token_pool
+        self.page_size = page_size
+
+        self.tp_group = tp_cache_group
+
+        if not use_hicache:
+            self.tree = RadixTreeCpp(
+                disabled=self.disable,
+                page_size=page_size,
+                host_size=None,  # no host cache, this should be removed in the future
+                write_through_threshold=self.write_through_threshold,
+            )
+            self.cache_controller = None
+            return  # early return if hicache is not used
+
+        raise NotImplementedError("Host cache is not supported yet")
+
+    def reset(self):
+        if self.cache_controller is not None:
+            # need to clear the acks before resetting the cache controller
+            raise NotImplementedError("Host cache is not supported yet")
+        self.tree.reset()
+
+    def match_prefix(self, key: List[int], **kwargs) -> MatchResult:
+        device_indices_vec, host_indices_length, node_gpu, node_cpu = (
+            self.tree.match_prefix(key)
+        )
+        return MatchResult(
+            device_indices=self._merge_tensor(device_indices_vec),
+            last_device_node=node_gpu,
+            last_host_node=node_cpu,
+            host_hit_length=host_indices_length,
+        )
+
+    def _insert(self, key: List[int], value: torch.Tensor) -> int:
+        """
+        Insert a key-value pair into the radix tree.
+        Args:
+            key (List[int]): The key to insert, represented as a list of integers.
+            value (torch.Tensor): The value to associate with the key.
+        Returns:
+            int: Number of device indices that were already present in the tree before the insertion.
+        """
+        ongoing_write, length = self.tree.writing_through(key, value)
+        if self.cache_controller is None:
+            assert len(ongoing_write) == 0, "Implementation error"
+            return length
+
+        raise NotImplementedError("Host cache is not supported yet")
+
+    def dec_lock_ref(self, node: TreeNodeCpp):
+        """
+        Decrement the reference count of a node to root of the radix tree.
+        Args:
+            node (TreeNodeCpp): The handle of the node to decrement the reference count for.
+        """
+        self.tree.lock_ref(node, False)  # do not increment
+
+    def inc_lock_ref(self, node: TreeNodeCpp):
+        """
+        Increment the reference count of from a node to root of the radix tree.
+        Args:
+            node (TreeNodeCpp): The handle of the node to increment the reference count for.
+        """
+        self.tree.lock_ref(node, True)
+
+    def evict(self, num_tokens: int):
+        evicted_device_indices = self.tree.evict(num_tokens)
+        for indice in evicted_device_indices:
+            self.token_to_kv_pool.free(indice)
+
+    def evictable_size(self):
+        return self.tree.evictable_size()
+
+    def protected_size(self):
+        return self.tree.protected_size()
+
+    def total_size(self):
+        return self.tree.total_size()
+
+    def cache_finished_req(self, req: Req):
+        """Cache request when it finishes."""
+        assert req.req_pool_idx is not None
+        token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+        overall_len = len(token_ids)  # prefill + decode
+        kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, :overall_len]
+
+        # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
+        # it will automatically align them, but length of them should be equal
+        old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
+        new_prefix_len = self._insert(token_ids, kv_indices)
+
+        # NOTE: kv_indices[:old_prefix_len] == req.prefix_indices
+        assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
+
+        # KVCache between old & new is newly generated, but already exists in the pool
+        # we need to free this newly generated kv indices
+        if old_prefix_len < new_prefix_len:
+            self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
+
+        # need to free the unaligned part, since it cannot be inserted into the radix tree
+        if self.page_size != 1 and (  # unaligned tail only exists when page_size > 1
+            (unaligned_len := overall_len % self.page_size) > 0
+        ):
+            # NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
+            self.token_to_kv_pool.free(kv_indices[overall_len - unaligned_len :])
+
+        # Remove req slot release the cache lock
+        self.dec_lock_ref(req.last_node)
+        self.req_to_token_pool.free(req.req_pool_idx)
+
+    def cache_unfinished_req(self, req: Req):
+        """Cache request when it is unfinished."""
+        assert req.req_pool_idx is not None
+        token_ids = req.fill_ids
+        prefill_len = len(token_ids)  # prefill only (maybe chunked)
+        kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, :prefill_len]
+
+        # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
+        # it will automatically align them, but length of them should be equal
+        old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
+        new_prefix_len = self._insert(token_ids, kv_indices)
+
+        # NOTE: kv_indices[:old_prefix_len] == req.prefix_indices
+        assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
+
+        # TODO(dark): optimize the `insert` and `match` (e.g. merge into 1 function)
+        # The prefix indices need to updated to reuse the kv indices in the pool
+        new_indices_vec, _, new_last_node, _ = self.tree.match_prefix(token_ids)
+        new_indices = self._merge_tensor(new_indices_vec)
+        assert new_prefix_len <= len(new_indices)
+
+        # KVCache between old & new is newly generated, but already exists in the pool
+        # we need to free this newly generated kv indices and reuse the indices in the pool
+        if old_prefix_len < new_prefix_len:
+            self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
+            reused_indices = new_indices[old_prefix_len:new_prefix_len]
+            self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, old_prefix_len:new_prefix_len
+            ] = reused_indices
+
+        if req.last_node != new_last_node:
+            self.dec_lock_ref(req.last_node)
+            self.inc_lock_ref(new_last_node)
+
+        # NOTE: there might be unaligned tail, so we may need to append it
+        assert len(new_indices) <= prefill_len < len(new_indices) + self.page_size
+        if self.page_size != 1 and len(new_indices) < prefill_len:
+            req.prefix_indices = torch.cat(
+                [new_indices, kv_indices[len(new_indices) :]]
+            )
+        else:
+            req.prefix_indices = new_indices
+        req.last_node = new_last_node
+
+    def pretty_print(self):
+        return self.tree.debug_print()

sglang/srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp

@@ -0,0 +1,35 @@
+#include <torch/extension.h>
+
+#include <cstring>
+#include <vector>
+
+void read_shm(const torch::Tensor &shm, std::vector<torch::Tensor> dst) {
+  py::gil_scoped_release release;
+  char *src_ptr = static_cast<char *>(shm.data_ptr());
+  size_t current = 0;
+  for (size_t i = 0; i < dst.size(); ++i) {
+    auto &t = dst[i];
+    size_t t_bytes = t.numel() * t.element_size();
+    char *dst_ptr = static_cast<char *>(t.data_ptr());
+    std::memcpy(dst_ptr, src_ptr + current, t_bytes);
+    current += t_bytes;
+  }
+}
+
+void write_shm(const std::vector<torch::Tensor> src, torch::Tensor &shm) {
+  py::gil_scoped_release release;
+  char *dst_ptr = static_cast<char *>(shm.data_ptr());
+  size_t current = 0;
+  for (size_t i = 0; i < src.size(); ++i) {
+    auto &t = src[i];
+    size_t t_bytes = t.numel() * t.element_size();
+    char *src_ptr = static_cast<char *>(t.data_ptr());
+    std::memcpy(dst_ptr + current, src_ptr, t_bytes);
+    current += t_bytes;
+  }
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("read_shm", &read_shm, "Read tensors from shared memory");
+  m.def("write_shm", &write_shm, "Write tensors to shared memory");
+}

sglang/srt/model_executor/cuda_graph_runner.py

@@ -29,6 +29,9 @@ from torch.profiler import ProfilerActivity, profile
 
 from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import get_tensor_model_parallel_rank
+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.logits_processor import LogitsProcessorOutput
@@ -372,6 +375,11 @@ class CudaGraphRunner:
                 dtype=torch.bool,
                 device="cuda",
             )
+            self.next_token_logits_buffer = torch.zeros(
+                (self.max_num_token, self.model_runner.model_config.vocab_size),
+                dtype=torch.float,
+                device="cuda",
+            )
 
         # Capture
         try:
@@ -517,6 +525,7 @@ class CudaGraphRunner:
         else:
             encoder_lens = None
         mrope_positions = self.mrope_positions[:, :bs]
+        next_token_logits_buffer = self.next_token_logits_buffer[:num_tokens]
         self.num_token_non_padded[...] = num_tokens
 
         # pipeline parallelism
@@ -579,6 +588,7 @@ class CudaGraphRunner:
             input_ids=input_ids,
             req_pool_indices=req_pool_indices,
             seq_lens=seq_lens,
+            next_token_logits_buffer=next_token_logits_buffer,
             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,
@@ -643,11 +653,15 @@ class CudaGraphRunner:
 
             run_once()
 
-        global global_graph_memory_pool
-        with torch.cuda.graph(graph, pool=global_graph_memory_pool, stream=stream):
+        if get_global_graph_memory_pool() is None:
+            set_global_graph_memory_pool(torch.cuda.graph_pool_handle())
+        # Set graph pool id globally to be able to use symmetric memory
+        set_graph_pool_id(get_global_graph_memory_pool())
+        with torch.cuda.graph(
+            graph, pool=get_global_graph_memory_pool(), stream=stream
+        ):
             out = run_once()
 
-        global_graph_memory_pool = graph.pool()
         return graph, out
 
     def recapture_if_needed(self, forward_batch: ForwardBatch):

sglang/srt/model_executor/forward_batch_info.py

@@ -38,6 +38,7 @@ import torch
 import triton
 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,
     get_attention_dp_rank,
@@ -188,6 +189,7 @@ class ForwardBatch:
     token_ids_logprobs: Optional[List[List[int]]] = None
 
     # For logits and logprobs post processing
+    next_token_logits_buffer: torch.Tensor = None
     temp_scaled_logprobs: bool = False
     temperature: torch.Tensor = None
     top_p_normalized_logprobs: bool = False
@@ -644,12 +646,17 @@ class ForwardBatch:
             device=model_runner.device,
         )
 
-        bs = self.batch_size
         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
+
+        bs = self.batch_size
+
         # 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)
@@ -657,6 +664,9 @@ class ForwardBatch:
         seq_len_fill_value = (
             model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
         )
+        self.seq_lens_sum = self.seq_lens_sum + seq_len_fill_value * (
+            bs - self.seq_lens.shape[0]
+        )
         self.seq_lens = self._pad_tensor_to_size(
             self.seq_lens, bs, value=seq_len_fill_value
         )
@@ -700,7 +710,7 @@ class ForwardBatch:
 
     def post_forward_mlp_sync_batch(self, logits_output: LogitsProcessorOutput):
 
-        bs = self.batch_size
+        bs = getattr(self, "raw_bs", self.batch_size)
 
         if self.spec_info is not None:
             if self.forward_mode.is_decode():  # draft
@@ -839,7 +849,7 @@ class ForwardBatch:
 
 
 def enable_num_token_non_padded(server_args):
-    return server_args.enable_ep_moe or server_args.enable_deepep_moe
+    return get_moe_expert_parallel_world_size() > 1
 
 
 class PPProxyTensors:

sglang/srt/model_executor/model_runner.py

@@ -60,6 +60,7 @@ from sglang.srt.layers.dp_attention import (
     initialize_dp_attention,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.moe.utils import DeepEPMode, MoeA2ABackend
 from sglang.srt.layers.quantization import (
     deep_gemm_wrapper,
     monkey_patch_isinstance_for_vllm_base_layer,
@@ -217,6 +218,10 @@ class ModelRunner:
                 "use_mla_backend": self.use_mla_backend,
                 "speculative_algorithm": self.spec_algorithm,
             }
+            | {
+                "moe_a2a_backend": MoeA2ABackend(server_args.moe_a2a_backend),
+                "deepep_mode": DeepEPMode(server_args.deepep_mode),
+            }
         )
 
         # CPU offload