sglang 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

sglang/srt/disaggregation/decode.py

@@ -21,20 +21,19 @@ Life cycle of a request in the decode server
 from __future__ import annotations
 
 import logging
-import os
 from collections import deque
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
-import numpy as np
 import torch
 from torch.distributed import ProcessGroup
 
-from sglang.srt.disaggregation.base import BaseKVManager, BaseKVReceiver, KVArgs, KVPoll
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.disaggregation.base import BaseKVManager, BaseKVReceiver, KVPoll
 from sglang.srt.disaggregation.utils import (
+    FAKE_BOOTSTRAP_HOST,
     DisaggregationMode,
-    FakeBootstrapHost,
     KVClassType,
     MetadataBuffers,
     ReqToMetadataIdxAllocator,
@@ -46,10 +45,12 @@ from sglang.srt.disaggregation.utils import (
     prepare_abort,
 )
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
+from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
+from sglang.srt.utils import require_mlp_sync
 
 logger = logging.getLogger(__name__)
 
@@ -86,7 +87,7 @@ class DecodeReqToTokenPool:
         self.max_context_len = max_context_len
         self.device = device
         self.pre_alloc_size = pre_alloc_size
-        with memory_saver_adapter.region():
+        with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
             self.req_to_token = torch.zeros(
                 (size + pre_alloc_size, max_context_len),
                 dtype=torch.int32,
@@ -135,7 +136,7 @@ class DecodePreallocQueue:
     def __init__(
         self,
         req_to_token_pool: ReqToTokenPool,
-        token_to_kv_pool_allocator: TokenToKVPoolAllocator,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
         draft_token_to_kv_pool: Optional[KVCache],
         req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
         metadata_buffers: MetadataBuffers,
@@ -145,7 +146,12 @@ class DecodePreallocQueue:
         gloo_group: ProcessGroup,
         tp_rank: int,
         tp_size: int,
+        dp_size: int,
+        gpu_id: int,
         bootstrap_port: int,
+        max_total_num_tokens: int,
+        prefill_pp_size: int,
+        num_reserved_decode_tokens: int,
         transfer_backend: TransferBackend,
     ):
         self.req_to_token_pool = req_to_token_pool
@@ -161,25 +167,33 @@ class DecodePreallocQueue:
         self.gloo_group = gloo_group
         self.tp_rank = tp_rank
         self.tp_size = tp_size
+        self.dp_size = dp_size
+        self.gpu_id = gpu_id
         self.bootstrap_port = bootstrap_port
-
-        self.num_reserved_decode_tokens = int(
-            os.environ.get("SGLANG_NUM_RESERVED_DECODE_TOKENS", "512")
-        )
-
+        self.max_total_num_tokens = max_total_num_tokens
+        self.prefill_pp_size = prefill_pp_size
+        self.num_reserved_decode_tokens = num_reserved_decode_tokens
+        self.transfer_backend = transfer_backend
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
-        self.transfer_backend = transfer_backend
+        self.retracted_queue: List[Req] = []
+        self.prefill_pp_size = prefill_pp_size
         self.kv_manager = self._init_kv_manager()
 
     def _init_kv_manager(self) -> BaseKVManager:
-        kv_args = KVArgs()
-        kv_args.engine_rank = self.tp_rank
+        kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
+        kv_args = kv_args_class()
+
+        attn_tp_size = self.tp_size // self.dp_size
+        kv_args.engine_rank = self.tp_rank % (attn_tp_size)
+        kv_args.decode_tp_size = attn_tp_size
+        kv_args.prefill_pp_size = self.prefill_pp_size
         kv_data_ptrs, kv_data_lens, kv_item_lens = (
             self.token_to_kv_pool.get_contiguous_buf_infos()
         )
-
         if self.draft_token_to_kv_pool is not None:
+            # We should also transfer draft model kv cache. The indices are
+            # always shared with a target model.
             draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
                 self.draft_token_to_kv_pool.get_contiguous_buf_infos()
             )
@@ -194,6 +208,7 @@ class DecodePreallocQueue:
         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
             self.metadata_buffers.get_buf_infos()
         )
+
         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
         kv_args.gpu_id = self.scheduler.gpu_id
         kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
@@ -205,27 +220,84 @@ class DecodePreallocQueue:
         )
         return kv_manager
 
-    def add(self, req: Req) -> None:
+    def add(self, req: Req, is_retracted: bool = False) -> None:
         """Add a request to the pending queue."""
-        if req.bootstrap_host == FakeBootstrapHost:
-            # Fake transfer for warmup reqs
-            kv_receiver_class = get_kv_class(TransferBackend.FAKE, KVClassType.RECEIVER)
+        if self._check_if_req_exceed_kv_capacity(req):
+            return
+
+        if is_retracted:
+            self.retracted_queue.append(req)
         else:
-            kv_receiver_class = get_kv_class(
-                self.transfer_backend, KVClassType.RECEIVER
+            if req.bootstrap_host == FAKE_BOOTSTRAP_HOST:
+                kv_receiver_class = get_kv_class(
+                    TransferBackend.FAKE, KVClassType.RECEIVER
+                )
+            else:
+                kv_receiver_class = get_kv_class(
+                    self.transfer_backend, KVClassType.RECEIVER
+                )
+
+            kv_receiver = kv_receiver_class(
+                mgr=self.kv_manager,
+                bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
+                bootstrap_room=req.bootstrap_room,
+                data_parallel_rank=req.data_parallel_rank,
             )
-        kv_receiver = kv_receiver_class(
-            mgr=self.kv_manager,
-            bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
-            bootstrap_room=req.bootstrap_room,
-            data_parallel_rank=req.data_parallel_rank,
-        )
-        self.queue.append(DecodeRequest(req=req, kv_receiver=kv_receiver))
 
-    def extend(self, reqs: List[Req]) -> None:
+            self.queue.append(
+                DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
+            )
+
+    def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool:
+        if len(req.origin_input_ids) > self.max_total_num_tokens:
+            message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}"
+            logger.error(message)
+            prepare_abort(req, message)
+            self.scheduler.stream_output([req], req.return_logprob)
+            return True
+        return False
+
+    def extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
         """Add a request to the pending queue."""
         for req in reqs:
-            self.add(req)
+            self.add(req, is_retracted=is_retracted)
+
+    def resume_retracted_reqs(self) -> List[Req]:
+        # TODO refactor the scheduling part, reuse with the unified engine logic as much as possible
+
+        # allocate memory
+        resumed_reqs = []
+        indices_to_remove = set()
+        allocatable_tokens = self._allocatable_tokens(count_retracted=False)
+
+        for i, req in enumerate(self.retracted_queue):
+            if self.req_to_token_pool.available_size() <= 0:
+                break
+
+            required_tokens_for_request = (
+                len(req.origin_input_ids)
+                + len(req.output_ids)
+                + self.num_reserved_decode_tokens
+            )
+            if required_tokens_for_request > allocatable_tokens:
+                break
+
+            resumed_reqs.append(req)
+            indices_to_remove.add(i)
+            req.is_retracted = False
+            self._pre_alloc(req)
+            allocatable_tokens -= required_tokens_for_request
+
+            # load from cpu, release the cpu copy
+            req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator)
+
+        self.retracted_queue = [
+            entry
+            for i, entry in enumerate(self.retracted_queue)
+            if i not in indices_to_remove
+        ]
+
+        return resumed_reqs
 
     def _update_handshake_waiters(self) -> None:
         if not self.queue:
@@ -255,6 +327,8 @@ class DecodePreallocQueue:
                     error_message,
                     status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 )
+            else:
+                raise ValueError(f"Unexpected poll case: {poll}")
 
     def pop_preallocated(self) -> List[DecodeRequest]:
         """Pop the preallocated requests from the pending queue (FIFO)."""
@@ -262,8 +336,16 @@ class DecodePreallocQueue:
 
         preallocated_reqs = []
         indices_to_remove = set()
-        allocatable_tokens = self._allocatable_tokens()
 
+        # We need to make sure that the sum of inflight tokens and allocatable tokens is greater than maximum input+output length of each inflight request
+        # Otherwise it is possible for one request running decode out of memory, while all other requests are in the transfer queue that cannot be retracted.
+        retractable_tokens = sum(
+            len(r.origin_input_ids) + len(r.output_ids)
+            for r in self.scheduler.running_batch.reqs
+        )
+        allocatable_tokens = self._allocatable_tokens(
+            retractable_tokens=retractable_tokens, count_retracted=True
+        )
         # First, remove all failed requests from the queue
         for i, decode_req in enumerate(self.queue):
             if isinstance(decode_req.req.finished_reason, FINISH_ABORT):
@@ -272,6 +354,7 @@ class DecodePreallocQueue:
                 )
                 indices_to_remove.add(i)
 
+        # Then, preallocate the remaining requests if possible
         for i, decode_req in enumerate(self.queue):
             if i in indices_to_remove:
                 continue
@@ -285,10 +368,23 @@ class DecodePreallocQueue:
             if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
                 break
 
+            # Memory estimation: don't add if the projected memory cannot be met
+            # TODO: add new_token ratio
+            origin_input_len = len(decode_req.req.origin_input_ids)
             required_tokens_for_request = (
-                len(decode_req.req.origin_input_ids) + self.num_reserved_decode_tokens
+                origin_input_len + self.num_reserved_decode_tokens
             )
 
+            if (
+                max(
+                    required_tokens_for_request,
+                    origin_input_len
+                    + decode_req.req.sampling_params.max_new_tokens
+                    - retractable_tokens,
+                )
+                > allocatable_tokens
+            ):
+                break
             if required_tokens_for_request > allocatable_tokens:
                 break
 
@@ -301,7 +397,6 @@ class DecodePreallocQueue:
                 ]
                 .cpu()
                 .numpy()
-                .astype(np.int64)
             )
 
             decode_req.metadata_buffer_index = (
@@ -321,15 +416,35 @@ class DecodePreallocQueue:
 
         return preallocated_reqs
 
-    def _allocatable_tokens(self) -> int:
-        allocatable_tokens = (
-            self.token_to_kv_pool_allocator.available_size()
-            - self.num_reserved_decode_tokens
+    def _allocatable_tokens(
+        self, retractable_tokens: Optional[int] = None, count_retracted: bool = True
+    ) -> int:
+        need_space_for_single_req = (
+            max(
+                [
+                    x.sampling_params.max_new_tokens
+                    + len(x.origin_input_ids)
+                    - retractable_tokens
+                    for x in self.scheduler.running_batch.reqs
+                ]
+            )
+            if retractable_tokens is not None
+            and len(self.scheduler.running_batch.reqs) > 0
+            else 0
+        )
+
+        available_size = self.token_to_kv_pool_allocator.available_size()
+
+        allocatable_tokens = available_size - max(
+            # preserve some space for future decode
+            self.num_reserved_decode_tokens
             * (
                 len(self.scheduler.running_batch.reqs)
                 + len(self.transfer_queue.queue)
                 + len(self.scheduler.waiting_queue)
-            )
+            ),
+            # make sure each request can finish if reach max_tokens with all other requests retracted
+            need_space_for_single_req,
         )
 
         # Note: if the last fake extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
@@ -342,15 +457,27 @@ class DecodePreallocQueue:
                 self.scheduler.last_batch.reqs
             )
 
+        if count_retracted:
+            allocatable_tokens -= sum(
+                [
+                    len(req.origin_input_ids)
+                    + len(req.output_ids)
+                    + self.num_reserved_decode_tokens
+                    for req in self.retracted_queue
+                ]
+            )
         return allocatable_tokens
 
     def _pre_alloc(self, req: Req) -> torch.Tensor:
         """Pre-allocate the memory for req_to_token and token_kv_pool"""
         req_pool_indices = self.req_to_token_pool.alloc(1)
 
-        assert req_pool_indices is not None
+        assert (
+            req_pool_indices is not None
+        ), "req_pool_indices is full! There is a bug in memory estimation."
 
         req.req_pool_idx = req_pool_indices[0]
+
         if self.token_to_kv_pool_allocator.page_size == 1:
             kv_loc = self.token_to_kv_pool_allocator.alloc(
                 len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
@@ -375,7 +502,10 @@ class DecodePreallocQueue:
                 ),
                 extend_num_tokens=num_tokens,
             )
-        assert kv_loc is not None
+
+        assert (
+            kv_loc is not None
+        ), "KV cache is full! There is a bug in memory estimation."
 
         self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
 
@@ -395,6 +525,7 @@ class DecodeTransferQueue:
         self,
         gloo_group: ProcessGroup,
         req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
+        tp_rank: int,
         metadata_buffers: MetadataBuffers,
         scheduler: Scheduler,
         tree_cache: BasePrefixCache,
@@ -402,9 +533,11 @@ class DecodeTransferQueue:
         self.queue: List[DecodeRequest] = []
         self.gloo_group = gloo_group
         self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
+        self.tp_rank = tp_rank
         self.metadata_buffers = metadata_buffers
         self.scheduler = scheduler
         self.tree_cache = tree_cache
+        self.spec_algorithm = scheduler.spec_algorithm
 
     def add(self, decode_req: DecodeRequest) -> None:
         self.queue.append(decode_req)
@@ -412,10 +545,9 @@ class DecodeTransferQueue:
     def extend(self, decode_reqs: List[DecodeRequest]) -> None:
         self.queue.extend(decode_reqs)
 
-    def pop_transferred(self) -> List[DecodeRequest]:
+    def pop_transferred(self) -> List[Req]:
         if not self.queue:
             return []
-
         polls = poll_and_all_reduce(
             [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
         )
@@ -424,7 +556,7 @@ class DecodeTransferQueue:
         indices_to_remove = set()
         for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
             if poll == KVPoll.Failed:
-                error_message = f"Decode transfer failed for request rank={self.scheduler.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
+                error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
                 try:
                     decode_req.kv_receiver.failure_exception()
                 except Exception as e:
@@ -447,6 +579,7 @@ class DecodeTransferQueue:
                 idx = decode_req.metadata_buffer_index
                 (
                     output_id,
+                    output_hidden_states,
                     output_token_logprobs_val,
                     output_token_logprobs_idx,
                     output_top_logprobs_val,
@@ -454,7 +587,8 @@ class DecodeTransferQueue:
                 ) = self.metadata_buffers.get_buf(idx)
 
                 decode_req.req.output_ids.append(output_id[0].item())
-
+                if not self.spec_algorithm.is_none():
+                    decode_req.req.hidden_states_tensor = output_hidden_states
                 if decode_req.req.return_logprob:
                     decode_req.req.output_token_logprobs_val.append(
                         output_token_logprobs_val[0].item()
@@ -499,15 +633,6 @@ class DecodeTransferQueue:
 
 class SchedulerDisaggregationDecodeMixin:
 
-    def _prepare_idle_batch_and_run(self, batch, delay_process=False):
-        batch, _ = self.prepare_dp_attn_batch(batch)
-        result = None
-        if batch:
-            result = self.run_batch(batch)
-            if not delay_process:
-                self.process_batch_result(batch, result)
-        return batch, result
-
     @torch.no_grad()
     def event_loop_normal_disagg_decode(self: Scheduler):
         """A normal scheduler loop for decode worker in disaggregation mode."""
@@ -520,10 +645,7 @@ class SchedulerDisaggregationDecodeMixin:
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
 
-            prepare_dp_attn_flag = (
-                self.server_args.enable_dp_attention
-                or self.server_args.enable_sp_layernorm
-            )
+            prepare_mlp_sync_flag = require_mlp_sync(self.server_args)
 
             if batch:
                 # Generate fake extend output.
@@ -532,24 +654,26 @@ class SchedulerDisaggregationDecodeMixin:
                     self.stream_output(
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
-                    if prepare_dp_attn_flag:
+                    if prepare_mlp_sync_flag:
                         self._prepare_idle_batch_and_run(None)
                 else:
-                    if prepare_dp_attn_flag:
-                        self.prepare_dp_attn_batch(batch)
+                    if prepare_mlp_sync_flag:
+                        self.prepare_mlp_sync_batch(batch)
                     result = self.run_batch(batch)
                     self.process_batch_result(batch, result)
-            elif prepare_dp_attn_flag:
+            elif prepare_mlp_sync_flag:
                 batch, _ = self._prepare_idle_batch_and_run(None)
 
             if batch is None and (
-                len(self.disagg_decode_transfer_queue.queue)
+                len(self.waiting_queue)
+                + len(self.disagg_decode_transfer_queue.queue)
                 + len(self.disagg_decode_prealloc_queue.queue)
                 == 0
             ):
                 # When the server is idle, do self-check and re-init some states
                 self.check_memory()
                 self.new_token_ratio = self.init_new_token_ratio
+                self.maybe_sleep_on_idle()
 
             self.last_batch = batch
 
@@ -568,10 +692,7 @@ class SchedulerDisaggregationDecodeMixin:
             self.cur_batch = batch
             last_batch_in_queue = False
 
-            prepare_dp_attn_flag = (
-                self.server_args.enable_dp_attention
-                or self.server_args.enable_sp_layernorm
-            )
+            prepare_mlp_sync_flag = require_mlp_sync(self.server_args)
 
             if batch:
                 # Generate fake extend output.
@@ -580,7 +701,7 @@ class SchedulerDisaggregationDecodeMixin:
                     self.stream_output(
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
-                    if prepare_dp_attn_flag:
+                    if prepare_mlp_sync_flag:
                         batch_, result = self._prepare_idle_batch_and_run(
                             None, delay_process=True
                         )
@@ -588,8 +709,8 @@ class SchedulerDisaggregationDecodeMixin:
                             result_queue.append((batch_.copy(), result))
                             last_batch_in_queue = True
                 else:
-                    if prepare_dp_attn_flag:
-                        self.prepare_dp_attn_batch(batch)
+                    if prepare_mlp_sync_flag:
+                        self.prepare_mlp_sync_batch(batch)
                     result = self.run_batch(batch)
                     result_queue.append((batch.copy(), result))
 
@@ -604,7 +725,7 @@ class SchedulerDisaggregationDecodeMixin:
                         self.set_next_batch_sampling_info_done(tmp_batch)
                     last_batch_in_queue = True
 
-            elif prepare_dp_attn_flag:
+            elif prepare_mlp_sync_flag:
                 batch, result = self._prepare_idle_batch_and_run(
                     None, delay_process=True
                 )
@@ -621,17 +742,28 @@ class SchedulerDisaggregationDecodeMixin:
                 self.process_batch_result(tmp_batch, tmp_result)
 
             if batch is None and (
-                len(self.disagg_decode_transfer_queue.queue)
+                len(self.waiting_queue)
+                + len(self.disagg_decode_transfer_queue.queue)
                 + len(self.disagg_decode_prealloc_queue.queue)
                 == 0
             ):
                 # When the server is idle, do self-check and re-init some states
                 self.check_memory()
                 self.new_token_ratio = self.init_new_token_ratio
+                self.maybe_sleep_on_idle()
 
             self.last_batch = batch
             self.last_batch_in_queue = last_batch_in_queue
 
+    def _prepare_idle_batch_and_run(self: Scheduler, batch, delay_process=False):
+        batch, _ = self.prepare_mlp_sync_batch(batch)
+        result = None
+        if batch:
+            result = self.run_batch(batch)
+            if not delay_process:
+                self.process_batch_result(batch, result)
+        return batch, result
+
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler,
     ) -> Optional[Tuple[ScheduleBatch, bool]]:
@@ -714,6 +846,13 @@ class SchedulerDisaggregationDecodeMixin:
         return new_batch
 
     def process_decode_queue(self: Scheduler):
+        # try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps
+        resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs()
+        self.waiting_queue.extend(resumed_reqs)
+        if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0:
+            # if there are still retracted requests, we do not allocate new requests
+            return
+
         req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
         self.disagg_decode_transfer_queue.extend(req_conns)
         alloc_reqs = (

sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -126,15 +126,16 @@ class ScheduleBatchDisaggregationDecodeMixin:
             )
             topk_index = topk_index.reshape(b, server_args.speculative_eagle_topk)
 
+            hidden_states_list = [req.hidden_states_tensor for req in self.reqs]
+            hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device)
+
             # local import to avoid circular import
             from sglang.srt.speculative.eagle_utils import EagleDraftInput
 
             spec_info = EagleDraftInput(
                 topk_p=topk_p,
                 topk_index=topk_index,
-                hidden_states=torch.ones(
-                    (b, model_config.hidden_size), device=self.device
-                ),
+                hidden_states=hidden_states,
                 verified_id=self.output_ids,
             )
             spec_info.prepare_for_extend(self)

sglang/srt/disaggregation/fake/__init__.py

@@ -1 +1 @@
-from .conn import FakeKVReceiver, FakeKVSender
+from sglang.srt.disaggregation.fake.conn import FakeKVReceiver, FakeKVSender

sglang/srt/disaggregation/fake/conn.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -8,7 +8,6 @@ from sglang.srt.disaggregation.base.conn import (
     BaseKVManager,
     BaseKVReceiver,
     BaseKVSender,
-    KVArgs,
     KVPoll,
 )
 
@@ -17,7 +16,14 @@ logger = logging.getLogger(__name__)
 
 # For warmup reqs, we don't kv transfer, we use the fake sender and receiver
 class FakeKVSender(BaseKVSender):
-    def __init__(self, mgr: BaseKVManager, bootstrap_addr: str, bootstrap_room: int):
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: int,
+        dest_tp_ranks: List[int],
+        pp_rank: int,
+    ):
         self.has_sent = False
 
     def poll(self) -> KVPoll:
@@ -26,7 +32,7 @@ class FakeKVSender(BaseKVSender):
             return KVPoll.WaitingForInput
         else:
             # Assume transfer completed instantly
-            logger.info("FakeKVSender poll success")
+            logger.debug("FakeKVSender poll success")
             return KVPoll.Success
 
     def init(
@@ -34,17 +40,17 @@ class FakeKVSender(BaseKVSender):
         kv_indices: list[int],
         aux_index: Optional[int] = None,
     ):
-        logger.info(
+        logger.debug(
             f"FakeKVSender init with kv_indices: {kv_indices}, aux_index: {aux_index}"
         )
         pass
 
     def send(
         self,
-        kv_indices: npt.NDArray[np.int64],
+        kv_indices: npt.NDArray[np.int32],
     ):
         self.has_sent = True
-        logger.info(f"FakeKVSender send with kv_indices: {kv_indices}")
+        logger.debug(f"FakeKVSender send with kv_indices: {kv_indices}")
 
     def failure_exception(self):
         raise Exception("Fake KVSender Exception")
@@ -66,12 +72,12 @@ class FakeKVReceiver(BaseKVReceiver):
             return KVPoll.WaitingForInput
         else:
             # Assume transfer completed instantly
-            logger.info("FakeKVReceiver poll success")
+            logger.debug("FakeKVReceiver poll success")
             return KVPoll.Success
 
     def init(self, kv_indices: list[int], aux_index: Optional[int] = None):
         self.has_init = True
-        logger.info(
+        logger.debug(
             f"FakeKVReceiver init with kv_indices: {kv_indices}, aux_index: {aux_index}"
         )