sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/disaggregation/prefill.py

@@ -20,7 +20,7 @@ Life cycle of a request in the prefill server
 from __future__ import annotations
 
 import logging
-import threading
+import time
 from collections import deque
 from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Type
@@ -42,14 +42,18 @@ from sglang.srt.disaggregation.utils import (
     poll_and_all_reduce,
     prepare_abort,
 )
-from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
-from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
-from sglang.srt.utils import (
-    DynamicGradMode,
-    broadcast_pyobj,
-    point_to_point_pyobj,
-    require_mlp_sync,
+from sglang.srt.managers.schedule_batch import (
+    FINISH_LENGTH,
+    Req,
+    RequestStage,
+    ScheduleBatch,
 )
+from sglang.srt.mem_cache.memory_pool import (
+    HybridLinearKVPool,
+    NSATokenToKVPool,
+    SWAKVPool,
+)
+from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj, require_mlp_sync
 
 if TYPE_CHECKING:
     from torch.distributed import ProcessGroup
@@ -140,6 +144,28 @@ class PrefillBootstrapQueue:
         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
         kv_args.gpu_id = self.scheduler.gpu_id
 
+        if hasattr(self.token_to_kv_pool, "get_state_buf_infos"):
+            state_data_ptrs, state_data_lens, state_item_lens = (
+                self.token_to_kv_pool.get_state_buf_infos()
+            )
+            kv_args.state_data_ptrs = state_data_ptrs
+            kv_args.state_data_lens = state_data_lens
+            kv_args.state_item_lens = state_item_lens
+
+            if isinstance(self.token_to_kv_pool, SWAKVPool):
+                kv_args.state_type = "swa"
+            elif isinstance(self.token_to_kv_pool, HybridLinearKVPool):
+                kv_args.state_type = "mamba"
+            elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
+                kv_args.state_type = "nsa"
+            else:
+                kv_args.state_type = "none"
+        else:
+            kv_args.state_data_ptrs = []
+            kv_args.state_data_lens = []
+            kv_args.state_item_lens = []
+            kv_args.state_type = "none"
+
         kv_manager_class: Type[BaseKVManager] = get_kv_class(
             self.transfer_backend, KVClassType.MANAGER
         )
@@ -170,6 +196,7 @@ class PrefillBootstrapQueue:
             pp_rank=self.pp_rank,
         )
         self._process_req(req)
+        req.add_latency(RequestStage.PREFILL_PREPARE)
         self.queue.append(req)
 
     def extend(self, reqs: List[Req], num_kv_heads: int) -> None:
@@ -256,8 +283,11 @@ class PrefillBootstrapQueue:
 
             num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size)
             req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
+
             bootstrapped_reqs.append(req)
             indices_to_remove.add(i)
+            req.time_stats.wait_queue_entry_time = time.perf_counter()
+            req.add_latency(RequestStage.PREFILL_BOOTSTRAP)
 
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
@@ -322,30 +352,21 @@ class SchedulerDisaggregationPrefillMixin:
             if require_mlp_sync(self.server_args):
                 batch = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
+
+            batch_result = None
             if batch:
-                result = self.run_batch(batch)
-                self.result_queue.append((batch.copy(), result))
-
-                if self.last_batch is None:
-                    # Create a dummy first batch to start the pipeline for overlap schedule.
-                    # It is now used for triggering the sampling_info_done event.
-                    tmp_batch = ScheduleBatch(
-                        reqs=None,
-                        forward_mode=ForwardMode.DUMMY_FIRST,
-                        next_batch_sampling_info=self.tp_worker.cur_sampling_info,
-                    )
-                    self.set_next_batch_sampling_info_done(tmp_batch)
+                batch_result = self.run_batch(batch)
+                self.result_queue.append((batch.copy(), batch_result))
 
             if self.last_batch:
                 tmp_batch, tmp_result = self.result_queue.popleft()
-                tmp_batch.next_batch_sampling_info = (
-                    self.tp_worker.cur_sampling_info if batch else None
-                )
                 self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
 
             if len(self.disagg_prefill_inflight_queue) > 0:
                 self.process_disagg_prefill_inflight_queue()
 
+            self.launch_batch_sample_if_needed(batch_result)
+
             if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
                 self.self_check_during_idle()
 
@@ -358,7 +379,6 @@ class SchedulerDisaggregationPrefillMixin:
         self: Scheduler,
         batch: ScheduleBatch,
         result: GenerationBatchResult,
-        launch_done: Optional[threading.Event] = None,
     ) -> None:
         """
         Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
@@ -369,53 +389,47 @@ class SchedulerDisaggregationPrefillMixin:
             next_token_ids,
             extend_input_len_per_req,
             extend_logprob_start_len_per_req,
+            copy_done,
         ) = (
             result.logits_output,
             result.next_token_ids,
             result.extend_input_len_per_req,
             result.extend_logprob_start_len_per_req,
+            result.copy_done,
         )
 
+        if copy_done is not None:
+            copy_done.synchronize()
+
         logprob_pt = 0
         # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
-        if self.enable_overlap:
-            # wait
-            logits_output, next_token_ids, _ = self.tp_worker.resolve_last_batch_result(
-                launch_done
-            )
-        else:
-            next_token_ids = result.next_token_ids.tolist()
-            if batch.return_logprob:
-                if logits_output.next_token_logprobs is not None:
-                    logits_output.next_token_logprobs = (
-                        logits_output.next_token_logprobs.tolist()
-                    )
-                if logits_output.input_token_logprobs is not None:
-                    logits_output.input_token_logprobs = tuple(
-                        logits_output.input_token_logprobs.tolist()
-                    )
+        next_token_ids = result.next_token_ids.tolist()
+        if batch.return_logprob:
+            if logits_output.next_token_logprobs is not None:
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.tolist()
+                )
+            if logits_output.input_token_logprobs is not None:
+                logits_output.input_token_logprobs = tuple(
+                    logits_output.input_token_logprobs.tolist()
+                )
 
         hidden_state_offset = 0
         for i, (req, next_token_id) in enumerate(
             zip(batch.reqs, next_token_ids, strict=True)
         ):
-            req: Req
             if req.is_chunked <= 0:
                 # There is no output_ids for prefill
                 req.output_ids.append(next_token_id)
                 self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
+                req.add_latency(RequestStage.PREFILL_FORWARD)
                 self.disagg_prefill_inflight_queue.append(req)
-                if (
-                    logits_output is not None
-                    and logits_output.hidden_states is not None
-                ):
-                    last_hidden_index = (
-                        hidden_state_offset + extend_input_len_per_req[i] - 1
-                    )
+                if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
+                    req.output_topk_p = batch.spec_info.topk_p[i]
+                    req.output_topk_index = batch.spec_info.topk_index[i]
                     req.hidden_states_tensor = (
-                        logits_output.hidden_states[last_hidden_index].cpu().clone()
+                        batch.spec_info.hidden_states[i].cpu().clone()
                     )
-                    hidden_state_offset += extend_input_len_per_req[i]
                 else:
                     req.hidden_states_tensor = None
                 if req.return_logprob:
@@ -434,6 +448,7 @@ class SchedulerDisaggregationPrefillMixin:
                     )
                     logprob_pt += num_input_logprobs
                 self.send_kv_chunk(req, last_chunk=True)
+                req.time_stats.prefill_transfer_queue_entry_time = time.perf_counter()
 
                 if req.grammar is not None:
                     # FIXME: this try-except block is for handling unexpected xgrammar issue.
@@ -473,8 +488,6 @@ class SchedulerDisaggregationPrefillMixin:
                 if self.enable_overlap:
                     self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
 
-        # We need to remove the sync in the following function for overlap schedule.
-        self.set_next_batch_sampling_info_done(batch)
         self.maybe_send_health_check_signal()
 
     def process_disagg_prefill_inflight_queue(
@@ -531,6 +544,9 @@ class SchedulerDisaggregationPrefillMixin:
             else:
                 assert False, f"Unexpected polling state {poll=}"
 
+        for req in done_reqs:
+            req.time_stats.completion_time = time.perf_counter()
+
         # Stream requests which have finished transfer
         self.stream_output(
             done_reqs,
@@ -539,6 +555,7 @@ class SchedulerDisaggregationPrefillMixin:
         )
         for req in done_reqs:
             req: Req
+            req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
             self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
             req.metadata_buffer_index = -1
 
@@ -609,232 +626,58 @@ class SchedulerDisaggregationPrefillMixin:
             .numpy()
         )
         req.start_send_idx = end_idx
+        state_indices = None
         if last_chunk:
             self.disagg_metadata_buffers.set_buf(req)
+
+            # Prepare extra pool indices for hybrid models
+            if isinstance(
+                self.token_to_kv_pool_allocator.get_kvcache(), HybridLinearKVPool
+            ):
+                # Mamba hybrid model: send single mamba state index
+                state_indices = [
+                    self.req_to_token_pool.req_index_to_mamba_index_mapping[
+                        req.req_pool_idx
+                    ]
+                    .cpu()
+                    .numpy()
+                ]
+            elif isinstance(self.token_to_kv_pool_allocator.get_kvcache(), SWAKVPool):
+                # SWA hybrid model: send last window KV indices
+                seq_len = len(req.fill_ids)
+                window_size = self.sliding_window_size
+                window_start = max(0, seq_len - window_size)
+                window_start = (window_start // page_size) * page_size
+
+                window_kv_indices_full = self.req_to_token_pool.req_to_token[
+                    req.req_pool_idx, window_start:seq_len
+                ]
+
+                # Translate to SWA pool indices
+                window_kv_indices_swa = (
+                    self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
+                        window_kv_indices_full
+                    )
+                )
+                state_indices = window_kv_indices_swa.cpu().numpy()
+                state_indices = kv_to_page_indices(state_indices, page_size)
+            elif isinstance(
+                self.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
+            ):
+                seq_len = len(req.fill_ids)
+                kv_indices_full = self.req_to_token_pool.req_to_token[
+                    req.req_pool_idx, :seq_len
+                ]
+                state_indices = kv_indices_full.cpu().numpy()
+                state_indices = kv_to_page_indices(state_indices, page_size)
+
         page_indices = kv_to_page_indices(kv_indices, page_size)
         if len(page_indices) == 0:
             logger.info(
                 f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
             )
             return
-        req.disagg_kv_sender.send(page_indices)
-
-    # PP
-    @DynamicGradMode()
-    def event_loop_pp_disagg_prefill(self: Scheduler):
-        """
-        An event loop for the prefill server in pipeline parallelism.
-
-        Rules:
-        1. Each stage runs in the same order and is notified by the previous stage.
-        2. Each send/recv operation is blocking and matched by the neighboring stage.
-
-        Regular Schedule:
-        ====================================================================
-        Stage i                   | Stage i+1
-        send ith req              | recv ith req
-        send ith proxy            | recv ith proxy
-        send prev (i+1)th carry   | recv prev (i+1)th carry
-        ====================================================================
-
-        Prefill Server Schedule:
-        ====================================================================
-        Stage i                        | Stage i+1
-        send ith req                   | recv ith req
-        send ith bootstrap req         | recv ith bootstrap req
-        send ith transferred req       | recv ith transferred req
-        send ith proxy                 | recv ith proxy
-        send prev (i+1)th carry        | recv prev (i+1)th carry
-        send prev (i+1)th release req  | recv prev (i+1)th release req
-        ====================================================================
-
-        There are two additional elements compared to the regular schedule:
-
-        1. Bootstrap Requests:
-            a. Instead of polling the status on the current workers, we should wait for the previous stage to notify to avoid desynchronization.
-            b. The first stage polls the status and propagates the bootstrapped requests down to all other stages.
-            c. If the first stage polls successfully, by nature, other ranks are also successful because they performed a handshake together.
-
-        2. Transferred Requests + Release Requests:
-            a. The first stage polls the transfer finished requests, performs an intersection with the next stage's finished requests, and propagates down to the last stage.
-            b. The last stage receives the requests that have finished transfer on all stages (consensus), then sends them to the first stage to release the memory.
-            c. The first stage receives the release requests, releases the memory, and then propagates the release requests down to the last stage.
-        """
-        from sglang.srt.managers.scheduler import GenerationBatchResult
-
-        mbs = [None] * self.pp_size
-        last_mbs = [None] * self.pp_size
-        self.running_mbs = [
-            ScheduleBatch(reqs=[], batch_is_full=False) for _ in range(self.pp_size)
-        ]
-        bids = [None] * self.pp_size
-        pp_outputs: Optional[PPProxyTensors] = None
-
-        # Either success or failed
-        bootstrapped_rids: List[str] = []
-        transferred_rids: List[str] = []
-        release_rids: Optional[List[str]] = None
-
-        # transferred microbatch
-        tmbs = [None] * self.pp_size
-
-        ENABLE_RELEASE = True  # For debug
-
-        while True:
-            server_is_idle = True
-
-            for mb_id in range(self.pp_size):
-                self.running_batch = self.running_mbs[mb_id]
-                self.last_batch = last_mbs[mb_id]
-
-                recv_reqs = self.recv_requests()
-
-                self.process_input_requests(recv_reqs)
-
-                if self.pp_group.is_first_rank:
-                    # First rank, pop the bootstrap reqs from the bootstrap queue
-                    bootstrapped_reqs, failed_reqs = (
-                        self.disagg_prefill_bootstrap_queue.pop_bootstrapped(
-                            return_failed_reqs=True
-                        )
-                    )
-                    bootstrapped_rids = [req.rid for req in bootstrapped_reqs] + [
-                        req.rid for req in failed_reqs
-                    ]
-                    self.waiting_queue.extend(bootstrapped_reqs)
-                else:
-                    # Other ranks, receive the bootstrap reqs info from the previous rank and ensure the consensus
-                    bootstrapped_rids = self.recv_pyobj_from_prev_stage()
-                    bootstrapped_reqs = (
-                        self.disagg_prefill_bootstrap_queue.pop_bootstrapped(
-                            rids_to_check=bootstrapped_rids
-                        )
-                    )
-                    self.waiting_queue.extend(bootstrapped_reqs)
-
-                if self.pp_group.is_first_rank:
-                    transferred_rids = self.get_transferred_rids()
-                # if other ranks,
-                else:
-                    # 1. recv previous stage's transferred reqs info
-                    prev_transferred_rids = self.recv_pyobj_from_prev_stage()
-                    # 2. get the current stage's transferred reqs info
-                    curr_transferred_rids = self.get_transferred_rids()
-                    # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids)
-                    transferred_rids = list(
-                        set(prev_transferred_rids) & set(curr_transferred_rids)
-                    )
-
-                tmbs[mb_id] = transferred_rids
-
-                self.process_prefill_chunk()
-                mbs[mb_id] = self.get_new_batch_prefill()
-                self.running_mbs[mb_id] = self.running_batch
-
-                self.cur_batch = mbs[mb_id]
-                if self.cur_batch:
-                    server_is_idle = False
-                    result = self.run_batch(self.cur_batch)
-
-                # send the outputs to the next step
-                if self.pp_group.is_last_rank:
-                    if self.cur_batch:
-                        next_token_ids, bids[mb_id] = (
-                            result.next_token_ids,
-                            result.bid,
-                        )
-                        pp_outputs = PPProxyTensors(
-                            {
-                                "next_token_ids": next_token_ids,
-                            }
-                        )
-                        # send the output from the last round to let the next stage worker run post processing
-                        self.pp_group.send_tensor_dict(
-                            pp_outputs.tensors,
-                            all_gather_group=self.attn_tp_group,
-                        )
-
-                if ENABLE_RELEASE:
-                    if self.pp_group.is_last_rank:
-                        # At the last stage, all stages has reached the consensus to release memory for transferred_rids
-                        release_rids = transferred_rids
-                        # send to the first rank
-                        self.send_pyobj_to_next_stage(release_rids)
-
-                # receive outputs and post-process (filter finished reqs) the coming microbatch
-                next_mb_id = (mb_id + 1) % self.pp_size
-                next_pp_outputs = None
-                next_release_rids = None
-
-                if mbs[next_mb_id] is not None:
-                    next_pp_outputs: Optional[PPProxyTensors] = PPProxyTensors(
-                        self.pp_group.recv_tensor_dict(
-                            all_gather_group=self.attn_tp_group
-                        )
-                    )
-                    mbs[next_mb_id].output_ids = next_pp_outputs["next_token_ids"]
-                    output_result = GenerationBatchResult(
-                        logits_output=None,
-                        pp_hidden_states_proxy_tensors=None,
-                        next_token_ids=next_pp_outputs["next_token_ids"],
-                        extend_input_len_per_req=None,
-                        extend_logprob_start_len_per_req=None,
-                        bid=bids[next_mb_id],
-                        can_run_cuda_graph=result.can_run_cuda_graph,
-                    )
-                    self.process_batch_result_disagg_prefill(
-                        mbs[next_mb_id], output_result
-                    )
-
-                    last_mbs[next_mb_id] = mbs[next_mb_id]
-
-                if ENABLE_RELEASE:
-                    if tmbs[next_mb_id] is not None:
-                        # recv consensus rids from the previous rank
-                        next_release_rids = self.recv_pyobj_from_prev_stage()
-                        self.process_disagg_prefill_inflight_queue(next_release_rids)
-
-                # carry the outputs to the next stage
-                if not self.pp_group.is_last_rank:
-                    if self.cur_batch:
-                        bids[mb_id] = result.bid
-                    if pp_outputs:
-                        # send the outputs from the last round to let the next stage worker run post processing
-                        self.pp_group.send_tensor_dict(
-                            pp_outputs.tensors,
-                            all_gather_group=self.attn_tp_group,
-                        )
-                    if ENABLE_RELEASE:
-                        if release_rids is not None:
-                            self.send_pyobj_to_next_stage(release_rids)
-
-                if not self.pp_group.is_last_rank:
-                    # send out reqs to the next stage
-                    self.send_pyobj_to_next_stage(recv_reqs)
-                    self.send_pyobj_to_next_stage(bootstrapped_rids)
-                    self.send_pyobj_to_next_stage(transferred_rids)
-
-                    # send out proxy tensors to the next stage
-                    if self.cur_batch:
-                        self.pp_group.send_tensor_dict(
-                            result.pp_hidden_states_proxy_tensors,
-                            all_gather_group=self.attn_tp_group,
-                        )
-
-                pp_outputs = next_pp_outputs
-                release_rids = next_release_rids
-
-                self.running_batch.batch_is_full = False
-
-            if not ENABLE_RELEASE:
-                if len(self.disagg_prefill_inflight_queue) > 0:
-                    self.process_disagg_prefill_inflight_queue()
-
-            # When the server is idle, self-check and re-init some states
-            if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0:
-                self.check_memory()
-                self.check_tree_cache()
-                self.new_token_ratio = self.init_new_token_ratio
+        req.disagg_kv_sender.send(page_indices, state_indices)
 
     def send_pyobj_to_next_stage(self, data):
         if self.attn_tp_rank == 0:

sglang/srt/disaggregation/utils.py

@@ -5,7 +5,7 @@ import random
 from collections import deque
 from contextlib import nullcontext
 from enum import Enum
-from typing import TYPE_CHECKING, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Optional, Type
 
 import numpy as np
 import torch
@@ -85,7 +85,7 @@ class MetadataBuffers:
         self,
         size: int,
         hidden_size: int,
-        dtype: torch.dtype,
+        hidden_states_dtype: torch.dtype,
         max_top_logprobs_num: int = 128,
         custom_mem_pool: torch.cuda.MemPool = None,
     ):
@@ -107,7 +107,9 @@ class MetadataBuffers:
             # We transfer the metadata of first output token to decode
             # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
             self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device)
-
+            self.cached_tokens = torch.zeros(
+                (size, 16), dtype=torch.int32, device=device
+            )
             self.output_token_logprobs_val = torch.zeros(
                 (size, 16), dtype=torch.float32, device=device
             )
@@ -120,33 +122,49 @@ class MetadataBuffers:
             self.output_top_logprobs_idx = torch.zeros(
                 (size, max_top_logprobs_num), dtype=torch.int32, device=device
             )
+            # For PD + spec decode
+            self.output_topk_p = torch.zeros(
+                (size, 16), dtype=torch.float32, device=device
+            )
+            self.output_topk_index = torch.zeros(
+                (size, 16), dtype=torch.int64, device=device
+            )
             self.output_hidden_states = torch.zeros(
-                (size, hidden_size), dtype=dtype, device=device
+                (size, hidden_size), dtype=hidden_states_dtype, device=device
             )
 
     def get_buf_infos(self):
         ptrs = [
             self.output_ids.data_ptr(),
+            self.cached_tokens.data_ptr(),
             self.output_token_logprobs_val.data_ptr(),
             self.output_token_logprobs_idx.data_ptr(),
             self.output_top_logprobs_val.data_ptr(),
             self.output_top_logprobs_idx.data_ptr(),
+            self.output_topk_p.data_ptr(),
+            self.output_topk_index.data_ptr(),
             self.output_hidden_states.data_ptr(),
         ]
         data_lens = [
             self.output_ids.nbytes,
+            self.cached_tokens.nbytes,
             self.output_token_logprobs_val.nbytes,
             self.output_token_logprobs_idx.nbytes,
             self.output_top_logprobs_val.nbytes,
             self.output_top_logprobs_idx.nbytes,
+            self.output_topk_p.nbytes,
+            self.output_topk_index.nbytes,
             self.output_hidden_states.nbytes,
         ]
         item_lens = [
             self.output_ids[0].nbytes,
+            self.cached_tokens[0].nbytes,
             self.output_token_logprobs_val[0].nbytes,
             self.output_token_logprobs_idx[0].nbytes,
             self.output_top_logprobs_val[0].nbytes,
             self.output_top_logprobs_idx[0].nbytes,
+            self.output_topk_p[0].nbytes,
+            self.output_topk_index[0].nbytes,
             self.output_hidden_states[0].nbytes,
         ]
         return ptrs, data_lens, item_lens
@@ -154,16 +172,20 @@ class MetadataBuffers:
     def get_buf(self, idx: int):
         return (
             self.output_ids[idx],
+            self.cached_tokens[idx],
             self.output_token_logprobs_val[idx],
             self.output_token_logprobs_idx[idx],
             self.output_top_logprobs_val[idx],
             self.output_top_logprobs_idx[idx],
+            self.output_topk_p[idx],
+            self.output_topk_index[idx],
             self.output_hidden_states[idx],
         )
 
     def set_buf(self, req: Req):
 
         self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
+        self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
         if req.return_logprob:
             if req.output_token_logprobs_val:  # not none or empty list
                 self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
@@ -186,8 +208,17 @@ class MetadataBuffers:
                 ] = torch.tensor(
                     req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu"
                 )
-        # for PD + spec decode
+        # For PD + spec decode
         if req.hidden_states_tensor is not None:
+            # speculative_eagle_topk should not be greater than 16 currently
+            topk = req.output_topk_p.size(0)
+
+            self.output_topk_p[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_p
+            )
+            self.output_topk_index[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_index
+            )
             self.output_hidden_states[req.metadata_buffer_index].copy_(
                 req.hidden_states_tensor
             )

sglang/srt/distributed/device_communicators/all_reduce_utils.py

@@ -0,0 +1,16 @@
+MiB = 1024 * 1024
+
+SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
+    9: {
+        2: 64 * MiB,  # 64 MB
+        4: 64 * MiB,  # 64 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+    10: {
+        2: 64 * MiB,  # 64 MB
+        4: 64 * MiB,  # 64 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+}