sglang/python/sglang/srt/disaggregation/utils.py

from __future__ import annotations

import os
import random
import logging
from collections import deque
from contextlib import nullcontext
from enum import Enum
from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Type, overload

import numpy as np
import torch
import torch.distributed as dist

from sglang.srt.environ import envs
from sglang.srt.utils import is_npu

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from sglang.srt.disaggregation.base.conn import KVArgs
    from sglang.srt.disaggregation.common.conn import (
        CommonKVBootstrapServer,
        CommonKVManager,
        CommonKVReceiver,
        CommonKVSender,
    )
    from sglang.srt.managers.schedule_batch import Req

#########################
# Constants & Enums
#########################
FAKE_BOOTSTRAP_HOST = "2.2.2.2"


class DisaggregationMode(Enum):
    NULL = "null"
    PREFILL = "prefill"
    DECODE = "decode"


#########################
# Synchronization
#########################

# env var for testing failure, convert to float explicitly
FAILURE_PROB = float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", 0))


def poll_and_all_reduce(pollers, gloo_group: dist.ProcessGroup):
    # at a certain prob, the poll is failed to simulate failure
    if FAILURE_PROB > 0:
        from sglang.srt.disaggregation.base import KVPoll

        polls = [
            int(KVPoll.Failed) if random.random() < FAILURE_PROB else int(poller.poll())
            for poller in pollers
        ]
    else:
        polls = [int(poller.poll()) for poller in pollers]
    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
    dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group)
    return tensor_to_reduce.tolist()


def _cp_shared_kv_poll_debug_enabled() -> bool:
    return envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get()


def _poll_queue_debug_hash(debug_ids: Optional[list[str]]) -> int:
    if not debug_ids:
        return 0

    # Stable bounded FNV-1a style digest.  Keep it in signed-int64 range because
    # the debug consensus uses torch.int64 CPU collectives.
    value = 1469598103934665603
    mask = (1 << 63) - 1
    for item in debug_ids:
        for byte in str(item).encode("utf-8", errors="replace"):
            value ^= byte
            value = (value * 1099511628211) & mask
        value ^= 0xFF
        value = (value * 1099511628211) & mask
    return int(value)


def _validate_poll_queue_consensus(
    *,
    label: str,
    scope: str,
    local_len: int,
    debug_ids: Optional[list[str]],
    group: dist.ProcessGroup,
) -> None:
    local_hash = _poll_queue_debug_hash(debug_ids)
    signature = torch.tensor([local_len, local_hash], dtype=torch.int64, device="cpu")
    min_signature = signature.clone()
    max_signature = signature.clone()

    dist.all_reduce(min_signature, op=dist.ReduceOp.MIN, group=group)
    dist.all_reduce(max_signature, op=dist.ReduceOp.MAX, group=group)

    if torch.equal(min_signature, max_signature):
        return

    ids_head = list(debug_ids[:8]) if debug_ids is not None else None
    message = (
        "[CP_SHARED_KV_FAIL_FAST][poll_queue] "
        f"label={label} scope={scope} local_len={local_len} "
        f"min_len={int(min_signature[0].item())} "
        f"max_len={int(max_signature[0].item())} "
        f"local_hash={local_hash} min_hash={int(min_signature[1].item())} "
        f"max_hash={int(max_signature[1].item())} ids_head={ids_head}"
    )
    logger.error(message)
    raise RuntimeError(message)


def poll_and_all_reduce_attn_cp_tp_group(
    pollers,
    attn_cp_cpu_group: dist.ProcessGroup,
    attn_tp_cpu_group: dist.ProcessGroup,
    *,
    debug_label: Optional[str] = None,
    debug_ids: Optional[list[str]] = None,
):
    if _cp_shared_kv_poll_debug_enabled():
        label = debug_label or "unknown"
        local_len = len(pollers)
        _validate_poll_queue_consensus(
            label=label,
            scope="attn_tp",
            local_len=local_len,
            debug_ids=debug_ids,
            group=attn_tp_cpu_group,
        )
        _validate_poll_queue_consensus(
            label=label,
            scope="attn_cp",
            local_len=local_len,
            debug_ids=debug_ids,
            group=attn_cp_cpu_group,
        )

    # First sync across attn-tp ranks so all TP participants for a given (dp, cp)
    # shard observe the same status transitions.
    polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)

    # Then sync across attn-cp ranks, so all TPxCP participants in one DP shard
    # converge to the same global status.
    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
    dist.all_reduce(
        tensor_to_reduce,
        op=dist.ReduceOp.MIN,
        group=attn_cp_cpu_group,
    )
    return tensor_to_reduce.tolist()


#########################
# Metadata Buffers
#########################

_EAGLE_ACCEPT_DEBUG_COUNTERS = {}


def eagle_accept_debug_should_log(
    key: str,
    *,
    first: int = 16,
    every: int = 256,
) -> bool:
    if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get():
        return False
    count = _EAGLE_ACCEPT_DEBUG_COUNTERS.get(key, 0) + 1
    _EAGLE_ACCEPT_DEBUG_COUNTERS[key] = count
    return count <= first or (every > 0 and count % every == 0)


def eagle_accept_debug_tensor_digest(tensor: Any, *, sample: int = 64) -> str:
    """Small deterministic tensor summary for EAGLE handoff debugging.

    This intentionally samples only a prefix.  It is enabled only under
    SGLANG_EAGLE_ACCEPT_DEBUG and is for handoff equality checks, not for full
    numerical validation.
    """

    if tensor is None:
        return "None"
    try:
        t = torch.as_tensor(tensor)
    except Exception as exc:  # pragma: no cover - defensive debug helper
        return f"unavailable({type(exc).__name__})"
    shape = tuple(t.shape)
    dtype = str(t.dtype).replace("torch.", "")
    numel = int(t.numel())
    if numel == 0:
        return f"shape={shape} dtype={dtype} numel=0"

    flat = t.detach().reshape(-1)
    sample_count = min(sample, numel)
    try:
        sample_cpu = flat[:sample_count].cpu()
        head_cpu = sample_cpu[: min(8, sample_count)]
        head = head_cpu.tolist()
        if sample_cpu.is_floating_point():
            checksum = float(sample_cpu.float().sum().item())
            abs_checksum = float(sample_cpu.float().abs().sum().item())
            return (
                f"shape={shape} dtype={dtype} sample={sample_count} "
                f"sum={checksum:.6g} abs={abs_checksum:.6g} head={head}"
            )
        checksum = int(sample_cpu.to(torch.int64).sum().item())
        return (
            f"shape={shape} dtype={dtype} sample={sample_count} "
            f"sum={checksum} head={head}"
        )
    except Exception as exc:  # pragma: no cover - defensive debug helper
        return f"shape={shape} dtype={dtype} numel={numel} digest_error={type(exc).__name__}"


def append_cp_draft_state_buffers(
    kv_args: Any,
    draft_state_type: str,
    draft_state_data_ptrs,
    draft_state_data_lens,
    draft_state_item_lens,
    *,
    role: str,
    cp_rank: int,
) -> bool:
    """Append draft NSA state buffers as payload attached to target state buffers.

    Under CP draft shared KV, draft state is not allowed to silently diverge
    from target state. If draft exposes state buffers, target and draft must both
    be NSA state so the transfer still has one target-driven state stream.
    """

    target_state_type = getattr(kv_args, "state_type", "none")
    draft_state_bufs = len(draft_state_data_ptrs)
    can_append = target_state_type == "nsa" and draft_state_type == "nsa"

    if draft_state_bufs and not can_append:
        if envs.SGLANG_CP_DRAFT_SHARED_KV.get():
            raise RuntimeError(
                "CP draft shared KV state mismatch during "
                f"{role}: cp_rank={cp_rank} target_state_type={target_state_type} "
                f"draft_state_type={draft_state_type} "
                f"draft_state_bufs={draft_state_bufs}"
            )

    kv_args.draft_state_type = draft_state_type
    kv_args.draft_state_buffer_start = len(kv_args.state_data_ptrs)
    kv_args.draft_state_buffer_count = 0

    if not draft_state_bufs or not can_append:
        return False

    kv_args.state_data_ptrs += draft_state_data_ptrs
    kv_args.state_data_lens += draft_state_data_lens
    kv_args.state_item_lens += draft_state_item_lens
    if hasattr(kv_args, "state_layer_ids"):
        kv_args.state_layer_ids += [-(i + 1) for i in range(draft_state_bufs)]
    kv_args.draft_state_buffer_count = draft_state_bufs
    return True


class ReqToMetadataIdxAllocator:
    """A memory pool that maps a request to its first output token location."""

    def __init__(
        self,
        size: int,
    ):
        self.size = size
        self.free_slots = deque(list(range(size)))

    def available_size(self):
        return len(self.free_slots)

    def alloc(self) -> Optional[int]:
        if len(self.free_slots) == 0:
            return None

        return self.free_slots.popleft()

    def free(self, free_index: int):
        self.free_slots.append(free_index)


class MetadataBuffers:
    def __init__(
        self,
        size: int,
        hidden_size: int,
        hidden_states_dtype: torch.dtype,
        max_top_logprobs_num: int = 128,
        custom_mem_pool: torch.cuda.MemPool = None,
    ):
        self.custom_mem_pool = custom_mem_pool
        bootstrap_room_dtype = torch.uint64
        device = "cpu"
        if is_npu():
            # For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel.
            device = "npu"
            # TODO: Fix me when npu backend supports torch.uint64
            bootstrap_room_dtype = torch.int64
        elif self.custom_mem_pool:
            # TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
            device = "cpu"
        elif envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() == "INTRA_NODE_NVLINK":
            device = "cuda"
        with (
            torch.cuda.use_mem_pool(self.custom_mem_pool)
            if self.custom_mem_pool
            else nullcontext()
        ):
            # TODO: abort top_logprobs_num > 128 in PD

            # 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
            )
            self.output_token_logprobs_idx = torch.zeros(
                (size, 16), dtype=torch.int32, device=device
            )
            self.output_top_logprobs_val = torch.zeros(
                (size, max_top_logprobs_num), dtype=torch.float32, device=device
            )
            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=hidden_states_dtype, device=device
            )
            # Request validation: store bootstrap_room to detect metadata corruption
            self.bootstrap_room = torch.zeros(
                (size, 8), dtype=bootstrap_room_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(),
            self.bootstrap_room.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,
            self.bootstrap_room.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,
            self.bootstrap_room[0].nbytes,
        ]
        return ptrs, data_lens, item_lens

    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],
            self.bootstrap_room[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] = (
                    req.output_token_logprobs_val[0]
                )
            if req.output_token_logprobs_idx:  # not none or empty list
                self.output_token_logprobs_idx[req.metadata_buffer_index][0] = (
                    req.output_token_logprobs_idx[0]
                )

            if req.output_top_logprobs_val:  # not none or empty list
                self.output_top_logprobs_val[req.metadata_buffer_index][
                    : len(req.output_top_logprobs_val[0])
                ] = torch.tensor(
                    req.output_top_logprobs_val[0], dtype=torch.float32, device="cpu"
                )
            if req.output_top_logprobs_idx:  # not none or empty list
                self.output_top_logprobs_idx[req.metadata_buffer_index][
                    : len(req.output_top_logprobs_idx[0])
                ] = torch.tensor(
                    req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu"
                )
        # 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
            )
            if eagle_accept_debug_should_log("metadata_set"):
                logger.warning(
                    "[EAGLE_ACCEPT_DEBUG][metadata_set] rid=%s room=%s idx=%s "
                    "output_id=%s cached_tokens=%s topk=%s topk_p=%s "
                    "topk_index=%s hidden=%s",
                    str(getattr(req, "rid", ""))[:8],
                    getattr(req, "bootstrap_room", None),
                    req.metadata_buffer_index,
                    req.output_ids[0] if req.output_ids else None,
                    req.cached_tokens,
                    topk,
                    eagle_accept_debug_tensor_digest(
                        self.output_topk_p[req.metadata_buffer_index, :topk]
                    ),
                    eagle_accept_debug_tensor_digest(
                        self.output_topk_index[req.metadata_buffer_index, :topk]
                    ),
                    eagle_accept_debug_tensor_digest(
                        self.output_hidden_states[req.metadata_buffer_index]
                    ),
                )
        # Store bootstrap_room for validation on decode side
        self.bootstrap_room[req.metadata_buffer_index, 0] = (
            req.bootstrap_room if req.bootstrap_room is not None else 0
        )


#########################
# Transfer Backend
#########################


class TransferBackend(Enum):
    MOONCAKE = "mooncake"
    MORI = "mori"
    NIXL = "nixl"
    ASCEND = "ascend"
    FAKE = "fake"


class KVClassType(Enum):
    KVARGS = "kvargs"
    MANAGER = "manager"
    SENDER = "sender"
    RECEIVER = "receiver"
    BOOTSTRAP_SERVER = "bootstrap_server"


@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.KVARGS]
) -> Type[KVArgs]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.MANAGER]
) -> Type[CommonKVManager]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.SENDER]
) -> Type[CommonKVSender]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.RECEIVER]
) -> Type[CommonKVReceiver]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.BOOTSTRAP_SERVER]
) -> Type[CommonKVBootstrapServer]: ...


def get_kv_class(
    transfer_backend: TransferBackend, class_type: KVClassType
) -> Optional[Type]:
    from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender

    if transfer_backend == TransferBackend.MOONCAKE:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.mooncake import (
            MooncakeKVBootstrapServer,
            MooncakeKVManager,
            MooncakeKVReceiver,
            MooncakeKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: MooncakeKVManager,
            KVClassType.SENDER: MooncakeKVSender,
            KVClassType.RECEIVER: (MooncakeKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.MORI:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.mori import (
            MoriKVBootstrapServer,
            MoriKVManager,
            MoriKVReceiver,
            MoriKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: MoriKVManager,
            KVClassType.SENDER: MoriKVSender,
            KVClassType.RECEIVER: (MoriKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: MoriKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.ASCEND:
        from sglang.srt.disaggregation.ascend import (
            AscendKVBootstrapServer,
            AscendKVManager,
            AscendKVReceiver,
            AscendKVSender,
        )
        from sglang.srt.disaggregation.base import KVArgs

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: AscendKVManager,
            KVClassType.SENDER: AscendKVSender,
            KVClassType.RECEIVER: (AscendKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: AscendKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.NIXL:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.nixl import (
            NixlKVBootstrapServer,
            NixlKVManager,
            NixlKVReceiver,
            NixlKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: NixlKVManager,
            KVClassType.SENDER: NixlKVSender,
            KVClassType.RECEIVER: (NixlKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: NixlKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.FAKE:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.fake import (
            FakeKVManager,
            FakeKVReceiver,
            FakeKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: FakeKVManager,
            KVClassType.SENDER: FakeKVSender,
            KVClassType.RECEIVER: (FakeKVReceiver),
        }
        return class_mapping.get(class_type)

    raise ValueError(f"Unsupported transfer backend: {transfer_backend}")


#########################
# KV Pages
#########################


def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
    # 1. The page is guaranteed to be full except the last page.
    # 2. page index = kv_index // page_size
    # The return vector is kv_indices[::page_size] // page_size
    if page_size == 1:  # shortcut
        return kv_indices

    return kv_indices[::page_size] // page_size


def kv_to_page_num(num_kv_indices: int, page_size: int):
    # ceil(num_kv_indices / page_size)
    return (num_kv_indices + page_size - 1) // page_size


def page_indices_to_cp_rank_page_indices(
    page_indices: np.ndarray,
    total_pages: int,
    cp_rank: int,
    cp_size: int,
) -> np.ndarray:
    """
    Filter page_indices (which are *global* page ids in the KV pool) to those
    belonging to the given CP rank for this request.

    For a single request, its pages occupy a contiguous global range
    [first_page, first_page + total_pages). We first compute the local
    split [0, total_pages) across cp_size ranks, then shift that local
    range by first_page back into the global page id space and take
    the intersection with page_indices.

    Returns:
        Subset of page_indices that fall in this rank's global
        [start_page, end_page) slice for the given CP rank.
    """
    if cp_size <= 1:
        return page_indices

    if page_indices.size == 0:
        return np.asarray(page_indices)

    first_page = int(page_indices.min())
    base = total_pages // cp_size
    rem = total_pages % cp_size

    if rem == 0:
        local_start = cp_rank * base
        local_end = local_start + base
    else:
        local_start = cp_rank * base + min(cp_rank, rem)
        n_pages = base + (1 if cp_rank < rem else 0)
        local_end = local_start + n_pages

    # Map back to global page ids.
    start_page = first_page + local_start
    end_page = first_page + local_end

    mask = (page_indices >= start_page) & (page_indices < end_page)
    return np.asarray(page_indices)[mask]


def filter_kv_indices_for_cp_rank(
    kv_mgr: CommonKVManager, kv_indices: np.ndarray, index_slice: slice
) -> Tuple[np.ndarray, slice]:
    """Filters kv_indices and index_slice for the current CP rank."""
    total_pages = len(kv_indices)
    cp_rank = kv_mgr.attn_cp_rank
    cp_size = kv_mgr.attn_cp_size

    if cp_size <= 1:
        return kv_indices, index_slice
    if total_pages == 0:
        return kv_indices[:0], slice(index_slice.start, index_slice.start)

    # This rank owns the global page-id range [start_page, end_page); see
    # page_indices_to_cp_rank_page_indices. Test the range directly: the old
    # np.isin(kv_indices, kv_indices[range_mask]) was identical to range_mask
    # (a value is in the filtered subset iff it passes the same range test)
    # but cost an extra O(N log N) sort per chunk send.
    first_page = int(kv_indices.min())
    base = total_pages // cp_size
    rem = total_pages % cp_size
    if rem == 0:
        local_start = cp_rank * base
        local_end = local_start + base
    else:
        local_start = cp_rank * base + min(cp_rank, rem)
        local_end = local_start + base + (1 if cp_rank < rem else 0)
    start_page = first_page + local_start
    end_page = first_page + local_end

    mask = (kv_indices >= start_page) & (kv_indices < end_page)
    if not mask.any():
        new_kv_indices = kv_indices[:0]
        new_index_slice = slice(index_slice.start, index_slice.start)
    else:
        first_pos = int(mask.argmax())
        last_pos = len(mask) - int(mask[::-1].argmax())

        new_kv_indices = kv_indices[first_pos:last_pos]
        new_index_slice = slice(
            index_slice.start + first_pos,
            index_slice.start + last_pos,
        )
    return new_kv_indices, new_index_slice


def filter_kv_pages_for_cp_shared_kv(
    layout,
    logical_pages: np.ndarray,
    chunk_page_start: int,
) -> Tuple[np.ndarray, np.ndarray]:
    logical_pages = np.asarray(logical_pages, dtype=np.int32)
    if logical_pages.size > 0 and int(logical_pages.min()) < 0:
        bad_pages = logical_pages[logical_pages < 0]
        raise ValueError(
            "CP shared KV transfer got negative logical_pages. "
            f"bad_min={int(bad_pages.min())} bad_max={int(bad_pages.max())} "
            f"chunk_page_start={chunk_page_start} num_pages={len(logical_pages)}"
        )
    request_positions = np.arange(
        chunk_page_start,
        chunk_page_start + len(logical_pages),
        dtype=np.int32,
    )
    return layout.filter_owned_pages_np(logical_pages, request_positions)


def select_pages_by_request_positions(
    pages: np.ndarray,
    request_positions: np.ndarray,
) -> np.ndarray:
    """Select per-request page entries by absolute request page positions."""
    pages = np.asarray(pages, dtype=np.int32)
    request_positions = np.asarray(request_positions, dtype=np.int32)
    if request_positions.size == 0:
        return pages[:0]
    if int(request_positions.max()) >= len(pages):
        raise ValueError(
            "request_positions contains an entry outside pages: "
            f"max_position={int(request_positions.max())} pages={len(pages)}"
        )
    return pages[request_positions]


def _transfer_indices_summary(indices: Any) -> str:
    if indices is None:
        return "None"
    arr = np.asarray(indices)
    if arr.size == 0:
        return f"shape={arr.shape} dtype={arr.dtype} size=0"
    head = arr.reshape(-1)[: min(8, arr.size)].tolist()
    tail = arr.reshape(-1)[-min(8, arr.size) :].tolist()
    return (
        f"shape={arr.shape} dtype={arr.dtype} size={arr.size} "
        f"head={head} tail={tail}"
    )


def _slice_summary(index_slice: Optional[slice]) -> str:
    if index_slice is None:
        return "None"
    return (
        f"slice(start={index_slice.start}, stop={index_slice.stop}, "
        f"step={index_slice.step})"
    )


def validate_transfer_page_count_or_raise(
    *,
    prefill_indices: Any,
    dst_indices: Any,
    room: Optional[int],
    cp_rank: Optional[int],
    logical_page_positions: Optional[Any],
    index_slice: Optional[slice],
    is_cp_shared_kv: bool,
    path: str,
) -> None:
    """Fail fast if a transfer chunk would drop or invent pages.

    Mooncake used to truncate source pages when the destination page selection
    was shorter.  Under CP shared-KV this hides a broken logical-page mapping
    and lets decode continue with incomplete KV.  Keep the check as a small
    shared helper so tests can cover the contract without constructing the
    transfer worker thread.
    """

    prefill_len = len(prefill_indices) if prefill_indices is not None else 0
    dst_len = len(dst_indices) if dst_indices is not None else 0
    if prefill_len == dst_len:
        return

    prefix = (
        "[CP_SHARED_KV_FAIL_FAST]"
        if is_cp_shared_kv or logical_page_positions is not None
        else "[KV_TRANSFER_FAIL_FAST]"
    )
    raise RuntimeError(
        f"{prefix}[mooncake_transfer_page_count_mismatch] "
        "source and destination page counts must match exactly; "
        "continuing would silently drop or mis-map KV pages. "
        f"path={path} room={room} cp_rank={cp_rank} "
        f"prefill_pages={prefill_len} dst_pages={dst_len} "
        f"logical_positions={_transfer_indices_summary(logical_page_positions)} "
        f"index_slice={_slice_summary(index_slice)} "
        f"prefill_indices={_transfer_indices_summary(prefill_indices)} "
        f"dst_indices={_transfer_indices_summary(dst_indices)}"
    )


#########################
# Misc
#########################


def is_mla_backend(target_kv_pool) -> bool:
    from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool

    return isinstance(target_kv_pool, MLATokenToKVPool)


def prepare_abort(req: Req, error_message: str, status_code=None):
    from sglang.srt.managers.schedule_batch import FINISH_ABORT

    # populate finish metadata and stream output
    req.finished_reason = FINISH_ABORT(error_message, status_code)

    if req.return_logprob:
        req.input_token_logprobs_val = []
        req.input_token_logprobs_idx = []
        req.input_top_logprobs_val = []
        req.input_top_logprobs_idx = []
        req.input_token_ids_logprobs_val = []
        req.input_token_ids_logprobs_idx = []