[PD] Support --enable-unified-memory with PD disaggregation (kimi-linear MLA hybrid-Mamba) (#33362)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
co-authored by
Claude Fable 5
Claude Fable 5
parent
8a7c8a72d6
commit
ceeaec2078
@@ -327,6 +327,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
self.transfer_queue = transfer_queue
|
||||
self.tree_cache = tree_cache
|
||||
self.gloo_group = gloo_group
|
||||
# Destinations visible to prefill but not yet on the transfer queue.
|
||||
self._num_published_destinations = 0
|
||||
self.tp_rank = tp_rank
|
||||
self.tp_size = tp_size
|
||||
self.dp_size = dp_size
|
||||
@@ -1151,14 +1153,21 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
decode_req.req.req_pool_idx
|
||||
][total_prefix_len:origin_input_len]
|
||||
kv_indices = (
|
||||
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
|
||||
kv_indices
|
||||
)
|
||||
)
|
||||
|
||||
seq_len = origin_input_len
|
||||
|
||||
def _mamba_payload():
|
||||
return [
|
||||
self.req_to_token_pool.req_index_to_mamba_index_mapping[
|
||||
decode_req.req.req_pool_idx
|
||||
]
|
||||
self.req_to_token_pool.translate_mamba_indices(
|
||||
self.req_to_token_pool.req_index_to_mamba_index_mapping[
|
||||
decode_req.req.req_pool_idx
|
||||
]
|
||||
)
|
||||
.cpu()
|
||||
.numpy()
|
||||
]
|
||||
@@ -1306,6 +1315,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
decode_req.kv_receiver,
|
||||
decode_req.req.build_rebootstrap_payload(),
|
||||
)
|
||||
self._num_published_destinations += 1
|
||||
preallocated_reqs.append(decode_req)
|
||||
indices_to_remove.add(i)
|
||||
decode_req.req.time_stats.set_decode_transfer_queue_entry_time()
|
||||
@@ -1316,6 +1326,18 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
|
||||
return preallocated_reqs, failed_reqs
|
||||
|
||||
@property
|
||||
def has_published_destinations(self) -> bool:
|
||||
"""Whether any destination address is visible to prefill but not yet
|
||||
protected by the transfer queue."""
|
||||
return self._num_published_destinations > 0
|
||||
|
||||
def note_destinations_queued(self, count: int) -> None:
|
||||
"""Hand `count` published destinations over to the transfer queue."""
|
||||
self._num_published_destinations = max(
|
||||
0, self._num_published_destinations - count
|
||||
)
|
||||
|
||||
@property
|
||||
def num_tokens_pre_allocated(self):
|
||||
return sum(
|
||||
@@ -1798,6 +1820,10 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
|
||||
|
||||
def extend(self, decode_reqs: List[DecodeRequest]) -> None:
|
||||
self.queue.extend(decode_reqs)
|
||||
# This queue now covers them.
|
||||
prealloc_queue = self.scheduler.disagg_decode_prealloc_queue
|
||||
if prealloc_queue is not None:
|
||||
prealloc_queue.note_destinations_queued(len(decode_reqs))
|
||||
|
||||
def _commit_transfer_to_req(self, decode_req: DecodeRequest):
|
||||
idx = decode_req.metadata_buffer_index
|
||||
|
||||
@@ -8,7 +8,7 @@ import struct
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional, Tuple, Union
|
||||
from typing import List, Optional, Set, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
@@ -59,7 +59,7 @@ from sglang.srt.observability.trace import (
|
||||
TraceReqContext,
|
||||
trace_set_thread_info,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_parallel, get_schedule
|
||||
from sglang.srt.runtime_context import get_memory, get_parallel, get_schedule
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils.network import NetworkAddress
|
||||
|
||||
@@ -282,35 +282,41 @@ class MooncakeKVManager(CommonKVManager):
|
||||
def init_engine(self):
|
||||
self.engine = get_mooncake_transfer_engine()
|
||||
|
||||
def register_buffer_to_engine(self):
|
||||
# Batch register KV data buffers
|
||||
if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
|
||||
self.engine.batch_register(
|
||||
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
|
||||
)
|
||||
def _registerable_regions(self) -> List[Tuple[int, int]]:
|
||||
"""(ptr, len) regions to (de)register, exact duplicates removed.
|
||||
|
||||
# Batch register auxiliary data buffers
|
||||
if self.kv_args.aux_data_ptrs and self.kv_args.aux_data_lens:
|
||||
self.engine.batch_register(
|
||||
self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
|
||||
)
|
||||
Deduped because the unified memory pool reports one raw buffer as both
|
||||
its KV and its mamba state component, and double registration fails in
|
||||
the engine.
|
||||
"""
|
||||
regions: List[Tuple[int, int]] = []
|
||||
seen: Set[Tuple[int, int]] = set()
|
||||
|
||||
def add(ptrs: List[int], lens: List[int]) -> None:
|
||||
for ptr, length in zip(ptrs or [], lens or []):
|
||||
if (ptr, length) not in seen:
|
||||
seen.add((ptr, length))
|
||||
regions.append((ptr, length))
|
||||
|
||||
add(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)
|
||||
add(self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens)
|
||||
for ptrs, lens in zip(
|
||||
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
|
||||
):
|
||||
if ptrs and lens:
|
||||
self.engine.batch_register(ptrs, lens)
|
||||
add(ptrs, lens)
|
||||
return regions
|
||||
|
||||
def register_buffer_to_engine(self):
|
||||
regions = self._registerable_regions()
|
||||
if regions:
|
||||
ptrs, lens = zip(*regions)
|
||||
self.engine.batch_register(list(ptrs), list(lens))
|
||||
|
||||
def deregister_buffer_to_engine(self):
|
||||
if self.kv_args.kv_data_ptrs:
|
||||
self.engine.batch_deregister(self.kv_args.kv_data_ptrs)
|
||||
|
||||
if self.kv_args.aux_data_ptrs:
|
||||
self.engine.batch_deregister(self.kv_args.aux_data_ptrs)
|
||||
|
||||
for ptrs in self.kv_args.state_data_ptrs or []:
|
||||
if ptrs:
|
||||
self.engine.batch_deregister(ptrs)
|
||||
regions = self._registerable_regions()
|
||||
if regions:
|
||||
ptrs, _ = zip(*regions)
|
||||
self.engine.batch_deregister(list(ptrs))
|
||||
|
||||
if hasattr(self, "connection_pool"):
|
||||
with self.connection_lock:
|
||||
@@ -775,6 +781,52 @@ class MooncakeKVManager(CommonKVManager):
|
||||
# compared to using multiple threads
|
||||
return process_layers(layers_params)
|
||||
|
||||
def _validate_envelope_kv_layout(
|
||||
self,
|
||||
dst_kv_ptrs: list[int],
|
||||
dst_kv_item_len: Optional[int],
|
||||
dst_attn_tp_size: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Reject a peer whose KV registration shape differs from ours.
|
||||
|
||||
The unified memory pool registers ONE whole-envelope region and
|
||||
addresses the destination as ``dst_ptr + page_id * item_len`` using OUR
|
||||
``item_len``, so a peer on a different page size / spec, or without
|
||||
unified memory, would take envelope-sized blocks at the wrong offsets.
|
||||
Must run before the first RDMA write.
|
||||
|
||||
Scoped to unified memory by config, not by region count: a non-unified
|
||||
PP stage owning a single full-attention layer also registers one region,
|
||||
and `_send_kvcache_generic` pairs that with the peer by layer id.
|
||||
"""
|
||||
if not get_memory().enable_unified_memory:
|
||||
return
|
||||
if dst_attn_tp_size is not None and self.attn_tp_size != dst_attn_tp_size:
|
||||
# The unified mamba state ships as one whole-slot envelope with no
|
||||
# per-tensor dims, so `_send_mamba_state_slice` cannot reslice it and
|
||||
# silently falls back to an unsliced copy. Reject here, before any KV
|
||||
# is written, rather than in `maybe_send_extra` afterwards.
|
||||
raise RuntimeError(
|
||||
"--enable-unified-memory does not support different prefill / "
|
||||
f"decode attention TP sizes (prefill={self.attn_tp_size}, "
|
||||
f"decode={dst_attn_tp_size}): the whole-envelope state cannot "
|
||||
"be TP-resliced."
|
||||
)
|
||||
src_item_lens = self.kv_args.kv_item_lens
|
||||
if (
|
||||
len(src_item_lens) != 1
|
||||
or len(dst_kv_ptrs) != 1
|
||||
or dst_kv_item_len is None
|
||||
or src_item_lens[0] != dst_kv_item_len
|
||||
):
|
||||
raise RuntimeError(
|
||||
"PD KV layout mismatch on the whole-envelope path: prefill has "
|
||||
f"{len(src_item_lens)} KV region(s) with item_lens="
|
||||
f"{src_item_lens}, decode has {len(dst_kv_ptrs)} with item_len="
|
||||
f"{dst_kv_item_len}. With --enable-unified-memory both sides "
|
||||
"must enable it and use the same page size and model spec."
|
||||
)
|
||||
|
||||
def send_kvcache(
|
||||
self,
|
||||
mooncake_session_id: str,
|
||||
@@ -784,7 +836,12 @@ class MooncakeKVManager(CommonKVManager):
|
||||
executor: concurrent.futures.ThreadPoolExecutor,
|
||||
dst_layer_ids: Optional[List[int]] = None,
|
||||
dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
|
||||
dst_kv_item_len: Optional[int] = None,
|
||||
dst_attn_tp_size: Optional[int] = None,
|
||||
):
|
||||
self._validate_envelope_kv_layout(
|
||||
dst_kv_ptrs, dst_kv_item_len, dst_attn_tp_size
|
||||
)
|
||||
dst_device_kv_ptrs = None
|
||||
if dst_device_kv_indices is not None:
|
||||
compression_ratios = self.kv_args.mla_compression_ratios
|
||||
@@ -1243,6 +1300,17 @@ class MooncakeKVManager(CommonKVManager):
|
||||
)
|
||||
|
||||
if st == StateType.MAMBA:
|
||||
if (not src_dim_per_tensor or not dst_dim_per_tensor) and list(
|
||||
src_item_lens
|
||||
) != list(dst_item_lens):
|
||||
raise RuntimeError(
|
||||
"Mamba state layouts differ between prefill and decode "
|
||||
f"(src item_lens={src_item_lens}, dst item_lens="
|
||||
f"{dst_item_lens}) and no per-tensor dim metadata is "
|
||||
"available to reslice. With --enable-unified-memory, "
|
||||
"prefill and decode must both enable it and use equal "
|
||||
"attention TP sizes."
|
||||
)
|
||||
if (
|
||||
target_rank_registration_info is not None
|
||||
and self.attn_tp_size
|
||||
@@ -1693,6 +1761,8 @@ class MooncakeKVManager(CommonKVManager):
|
||||
executor,
|
||||
dst_layer_ids=target_rank_registration_info.dst_kv_layer_ids,
|
||||
dst_device_kv_indices=chunked_dst_device_kv_indice,
|
||||
dst_kv_item_len=target_rank_registration_info.dst_kv_item_len,
|
||||
dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
|
||||
)
|
||||
elif (
|
||||
self.enable_staging
|
||||
|
||||
@@ -790,6 +790,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
# Optimistic bootstrap can fail while this overlapped chunk is
|
||||
# already running. Drop aborted chunks instead of sending KV.
|
||||
if is_aborted(req):
|
||||
self.clear_pending_chunk_send(req)
|
||||
advance_logprob_pt(i, req)
|
||||
req.time_stats.set_last_chunked_prefill_finish_time()
|
||||
continue
|
||||
@@ -980,7 +981,17 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
|
||||
return transferred_rids
|
||||
|
||||
def clear_pending_chunk_send(self: Scheduler, req: Req) -> None:
|
||||
"""Drop `req` from the sent-but-unconcluded chunk set.
|
||||
|
||||
Every path that retires a request without a `last_chunk=True` send must
|
||||
call this: a stale entry holds the unified-memory compaction gate closed
|
||||
for the process lifetime.
|
||||
"""
|
||||
self.disagg_prefill_pending_chunk_rids.discard(req.rid)
|
||||
|
||||
def handle_bootstrap_failure(self: Scheduler, req: Req) -> None:
|
||||
self.clear_pending_chunk_send(req)
|
||||
error_message = (
|
||||
f"Prefill bootstrap failed for request rank={self.ps.tp_rank} "
|
||||
f"{req.rid=} {req.bootstrap_room=}"
|
||||
@@ -1178,9 +1189,11 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
|
||||
def _mamba_payload():
|
||||
return [
|
||||
self.req_to_token_pool.req_index_to_mamba_index_mapping[
|
||||
req.req_pool_idx
|
||||
]
|
||||
self.req_to_token_pool.translate_mamba_indices(
|
||||
self.req_to_token_pool.req_index_to_mamba_index_mapping[
|
||||
req.req_pool_idx
|
||||
]
|
||||
)
|
||||
.cpu()
|
||||
.numpy()
|
||||
]
|
||||
@@ -1287,6 +1300,13 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, seg_start:seg_end
|
||||
]
|
||||
# Unified memory: req_to_token holds VIRTUAL ids; the transfer needs
|
||||
# physical ones. Per segment, since each is its own gather.
|
||||
kv_indices = (
|
||||
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
|
||||
kv_indices
|
||||
)
|
||||
)
|
||||
page_indices = kv_to_page_indices(kv_indices, page_size)
|
||||
segment_is_last = last_chunk and is_final_segment
|
||||
if not req.disagg_kv_sender.should_send_kv_chunk(
|
||||
@@ -1299,6 +1319,12 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
num_kv_tokens=seg_end - seg_start,
|
||||
)
|
||||
req.start_send_idx = end_idx
|
||||
# A last chunk needs no entry: every `last_chunk=True` call site has
|
||||
# already put the request on `disagg_prefill_inflight_queue`.
|
||||
if last_chunk:
|
||||
self.disagg_prefill_pending_chunk_rids.discard(req.rid)
|
||||
else:
|
||||
self.disagg_prefill_pending_chunk_rids.add(req.rid)
|
||||
|
||||
def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None:
|
||||
"""Release KV cache and requeue an optimistic prefill request."""
|
||||
@@ -1308,6 +1334,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
req.reset_for_retract()
|
||||
req.output_ids = array("q")
|
||||
req.start_send_idx = 0
|
||||
self.clear_pending_chunk_send(req) # re-sends from scratch
|
||||
req.tmp_end_idx = -1
|
||||
req.disagg_decode_prefix_len = 0
|
||||
req.early_send_prefix_end = None
|
||||
|
||||
@@ -111,6 +111,50 @@ class DisaggregationMode(Enum):
|
||||
return "unified"
|
||||
|
||||
|
||||
def unified_memory_disagg_move_gate(scheduler):
|
||||
"""Compaction move gate for a PD node running the unified memory pool.
|
||||
|
||||
Returns a predicate that is True only when no transfer can be in flight, so
|
||||
compaction never relocates a page the RDMA engine is reading or writing.
|
||||
Safe to read this state from here: every mover runs on the scheduler thread.
|
||||
|
||||
A page is exposed from the moment its address reaches the peer until the
|
||||
transfer concludes, and for part of that lifetime the request is in NEITHER
|
||||
end's queue -- so queue emptiness alone is not enough:
|
||||
|
||||
- PREFILL: scheduling the final chunk clears `chunked_req` while earlier
|
||||
chunks may still be draining, and the request only reaches the inflight
|
||||
queue later, in the result path.
|
||||
- DECODE: `pop_preallocated` publishes one request's destinations and keeps
|
||||
allocating for the next, whose allocation can urgently flush the peer
|
||||
sub-allocator; the batch reaches the transfer queue only after the loop.
|
||||
"""
|
||||
if scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
|
||||
def prefill_gate() -> bool:
|
||||
return not (
|
||||
scheduler.disagg_prefill_inflight_queue
|
||||
or scheduler.disagg_prefill_pending_chunk_rids
|
||||
)
|
||||
|
||||
return prefill_gate
|
||||
|
||||
if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
|
||||
def decode_gate() -> bool:
|
||||
return not (
|
||||
scheduler.disagg_decode_transfer_queue.queue
|
||||
or scheduler.disagg_decode_prealloc_queue.has_published_destinations
|
||||
)
|
||||
|
||||
return decode_gate
|
||||
|
||||
raise ValueError(
|
||||
"unified_memory_disagg_move_gate: scheduler is not a PD node "
|
||||
f"(mode={scheduler.disaggregation_mode})"
|
||||
)
|
||||
|
||||
|
||||
#########################
|
||||
# Synchronization
|
||||
#########################
|
||||
|
||||
@@ -25,7 +25,7 @@ from collections import deque
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from functools import partial
|
||||
from http import HTTPStatus
|
||||
from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Tuple, Union
|
||||
from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
from sglang.srt.runtime_context import (
|
||||
get_device,
|
||||
@@ -90,6 +90,7 @@ from sglang.srt.disaggregation.utils import (
|
||||
TransferBackend,
|
||||
get_dsa_seed_metadata_dim,
|
||||
prepare_abort,
|
||||
unified_memory_disagg_move_gate,
|
||||
)
|
||||
from sglang.srt.distributed import get_pp_group, get_world_group
|
||||
from sglang.srt.distributed.parallel_state import get_tp_group
|
||||
@@ -1396,9 +1397,19 @@ class Scheduler(
|
||||
)
|
||||
# The prefill requests that are in the middle of kv sending
|
||||
self.disagg_prefill_inflight_queue: List[Req] = []
|
||||
# Requests with a sent chunk that are not yet on the inflight queue.
|
||||
self.disagg_prefill_pending_chunk_rids: Set[str] = set()
|
||||
|
||||
self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
|
||||
|
||||
if (
|
||||
self.enable_unified_memory
|
||||
and self.disaggregation_mode != DisaggregationMode.NULL
|
||||
):
|
||||
self.token_to_kv_pool_allocator.set_disagg_move_gate(
|
||||
unified_memory_disagg_move_gate(self)
|
||||
)
|
||||
|
||||
# Init mm receiver for EPD disaggregation mode
|
||||
if get_disagg().language_only and get_disagg().encoder_transfer_backend in [
|
||||
"zmq_to_scheduler",
|
||||
@@ -2928,6 +2939,7 @@ class Scheduler(
|
||||
req.time_stats.trace_ctx.abort(abort_info={"reason": "Aborted"})
|
||||
req.to_finish = None
|
||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
self.clear_pending_chunk_send(req)
|
||||
req.disagg_kv_sender.abort()
|
||||
maybe_release_metadata_buffer(
|
||||
req, self.req_to_metadata_buffer_idx_allocator
|
||||
|
||||
@@ -82,6 +82,16 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
|
||||
(0,), dtype=self.release_pages.dtype, device=self.device
|
||||
)
|
||||
|
||||
def translate_kv_indices_for_transfer(
|
||||
self, kv_indices: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""Token ids as the PD-disaggregation transfer engine addresses them.
|
||||
|
||||
Identity here: a static pool's token ids index its registered buffers
|
||||
directly. Virtual-id pools must override.
|
||||
"""
|
||||
return kv_indices
|
||||
|
||||
def get_cpu_copy(self, indices, mamba_indices=None):
|
||||
# FIXME: reuse the get_cpu_copy after paged allocator is implemented
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -353,17 +353,29 @@ class KVCacheConfigurator:
|
||||
# Unified-pool fast path: build req_to_token + token_to_kv pool + allocator
|
||||
# from one byte buffer, then return. Gated to the target worker
|
||||
# (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4).
|
||||
if (
|
||||
get_memory().enable_unified_memory
|
||||
and get_disagg().disaggregation_mode == "null"
|
||||
and req_to_token_pool is None
|
||||
):
|
||||
if get_memory().enable_unified_memory and req_to_token_pool is None:
|
||||
pd_enabled = get_disagg().disaggregation_mode != "null"
|
||||
if self.mambaish_config is not None:
|
||||
if pd_enabled and not self.use_mla_backend:
|
||||
raise ValueError(
|
||||
"--enable-unified-memory with PD disaggregation "
|
||||
"currently supports only MLA hybrid-Mamba models "
|
||||
"(e.g. kimi-linear); this model uses the MHA full-"
|
||||
"attention pool. Drop --enable-unified-memory or run "
|
||||
"without PD disaggregation."
|
||||
)
|
||||
bundle = self._init_unified_mamba_pools(
|
||||
max_num_reqs=sizes.max_running_requests,
|
||||
max_total_num_tokens=sizes.max_total_num_tokens,
|
||||
)
|
||||
elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config):
|
||||
if pd_enabled:
|
||||
raise ValueError(
|
||||
"--enable-unified-memory with PD disaggregation does "
|
||||
"not support hybrid-SWA models yet (no whole-envelope "
|
||||
"transfer scheme for the SWA sub-pool). Drop "
|
||||
"--enable-unified-memory or run without PD."
|
||||
)
|
||||
bundle = self._init_unified_swa_pools(
|
||||
max_num_reqs=sizes.max_running_requests,
|
||||
full_max_total_num_tokens=sizes.full_max_total_num_tokens,
|
||||
@@ -564,6 +576,11 @@ class KVCacheConfigurator:
|
||||
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
|
||||
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
|
||||
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
|
||||
decode_pre_alloc_size=(
|
||||
get_disagg().disaggregation_decode_extra_slots
|
||||
if get_disagg().disaggregation_mode == "decode"
|
||||
else 0
|
||||
),
|
||||
mamba_full_memory_ratio=get_schedule().mamba_full_memory_ratio,
|
||||
# Overlap mode: the allocator's `free` drops a wait_stream(forward_stream)
|
||||
# barrier so eager compaction serializes after the in-flight forward's
|
||||
|
||||
@@ -25,7 +25,7 @@ from __future__ import annotations
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
from typing import Callable, Dict, List, Optional, Set, Tuple
|
||||
|
||||
import torch
|
||||
from torch.profiler import record_function
|
||||
@@ -218,6 +218,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
|
||||
_STATS_INSTANCES.add(self)
|
||||
_install_signal_handlers_once()
|
||||
self.live_page_count = 0
|
||||
# While this returns False, `_flush` must not relocate any page.
|
||||
self.disagg_move_gate: Optional[Callable[[], bool]] = None
|
||||
self._latest_forward_done_event: Optional[torch.cuda.Event] = None
|
||||
# Most-recent forward's (done_event, out_cache_loc_virtual) for `_flush`'s
|
||||
# write-race check. Single slot: at most ONE forward in flight per call site.
|
||||
@@ -393,6 +395,12 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
|
||||
peer = self._peer
|
||||
if peer is None or not peer.lazy_compaction:
|
||||
return 0
|
||||
if peer.disagg_move_gate is not None and not peer.disagg_move_gate():
|
||||
# The peer cannot compact while a PD transfer is in flight, so these
|
||||
# holes are not realizable. Crediting them would let the scheduler
|
||||
# admit work that `_flush_peer_for_alloc` then cannot satisfy, and
|
||||
# the caller treats a failed alloc as a memory-estimation bug.
|
||||
return 0
|
||||
return len(peer._free_phys_pages) * peer.entry_bytes_per_page
|
||||
|
||||
def schedulable_available_size(self) -> int:
|
||||
@@ -1058,6 +1066,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
|
||||
self._compact_pending_impl(freed_physical_pages)
|
||||
|
||||
def _compact_pending_impl(self, freed_physical_pages: torch.Tensor) -> None:
|
||||
assert self.disagg_move_gate is None, (
|
||||
f"_compact_pending({self.sub_pool_name!r}): eager compaction ran with "
|
||||
"a PD-disaggregation move gate installed; PD requires lazy_compaction."
|
||||
)
|
||||
freed_set = set(int(x) for x in freed_physical_pages.tolist())
|
||||
if not freed_set:
|
||||
return
|
||||
@@ -1440,6 +1452,9 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
|
||||
"""
|
||||
if not self.lazy_compaction:
|
||||
return 0
|
||||
if self.disagg_move_gate is not None and not self.disagg_move_gate():
|
||||
# Holes stay in the free list; the next flush picks them up.
|
||||
return 0
|
||||
self._stats_n_flush_calls += 1
|
||||
with record_function("MultiEndedAlloc._flush"):
|
||||
self._drain_pending_reuse(urgent=urgent)
|
||||
@@ -1936,6 +1951,26 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
|
||||
to the physical translate when `kernel_page_multiplier == 1` (MHA)."""
|
||||
return self.full_attn_allocator.translate_kv_loc_dense(loc, out=out)
|
||||
|
||||
def translate_kv_indices_for_transfer(
|
||||
self, kv_indices: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""Virtual TOKEN ids -> PHYSICAL token ids for the PD transfer engine.
|
||||
|
||||
PHYSICAL, not dense: the transfer registers page ENVELOPES (see
|
||||
`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`).
|
||||
"""
|
||||
return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
|
||||
|
||||
def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None:
|
||||
"""Install the PD-disaggregation move gate on both sub-allocators."""
|
||||
assert self.lazy_compaction, (
|
||||
"PD disaggregation with the unified memory pool requires lazy "
|
||||
"compaction (eager free-path compaction moves pages under "
|
||||
"in-flight transfers)."
|
||||
)
|
||||
self.full_attn_allocator.disagg_move_gate = gate
|
||||
self.mamba_allocator.disagg_move_gate = gate
|
||||
|
||||
def is_slot_allocated(self, slot: int) -> bool:
|
||||
return self.full_attn_allocator.is_slot_allocated(slot)
|
||||
|
||||
|
||||
@@ -42,6 +42,7 @@ from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_page_major_mha_views,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
HybridLinearKVPool,
|
||||
HybridReqToTokenPool,
|
||||
MambaPool,
|
||||
MHATokenToKVPool,
|
||||
@@ -663,6 +664,20 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
|
||||
def get_kv_size_bytes(self):
|
||||
return 0 # UnifiedKVPool logs the total; per-sub-pool would double-count
|
||||
|
||||
def get_contiguous_buf_infos(self):
|
||||
"""PD-transfer registration: ONE entry, the raw buffer, addressed as
|
||||
``raw_ptr + physical_page_id * page_envelope_bytes``.
|
||||
|
||||
The transfer item is the whole page envelope (all layers of one page)
|
||||
rather than a per-layer region, because the per-layer dense views
|
||||
overlap and index in dense ids. Both sides must therefore build the
|
||||
pool with identical specs.
|
||||
"""
|
||||
# The address formula omits the anchor; a nonzero one would mis-address.
|
||||
assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
|
||||
raw = self._unified_buffer._raw
|
||||
return [raw.data_ptr()], [raw.numel()], [self._page_bytes]
|
||||
|
||||
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
|
||||
"""Relocate whole page envelopes.
|
||||
|
||||
@@ -812,6 +827,31 @@ class UnifiedMambaPool(MambaPool):
|
||||
# Physical-slot copy used by the allocator's `_compact_pending`.
|
||||
MambaPool.copy_from(self, src_index, dst_index)
|
||||
|
||||
# -- PD state transfer (StateType.MAMBA) --
|
||||
# The transfer item is the whole per-slot envelope, addressed as
|
||||
# `raw_ptr + physical_slot * entry_bytes`. An envelope cannot be TP-resliced
|
||||
# or PP-subset, so the per-tensor metadata below stays empty and both sides
|
||||
# must build identical mamba specs (equal attn TP, pp=1).
|
||||
|
||||
def get_contiguous_buf_infos(self):
|
||||
# The address formula omits the anchor; a nonzero one would mis-address.
|
||||
assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
|
||||
spec = self._unified_buffer.mamba_spec(self._sub_pool_name)
|
||||
raw = self._unified_buffer._raw
|
||||
return [raw.data_ptr()], [raw.numel()], [spec.entry_bytes()]
|
||||
|
||||
def get_state_dim_per_tensor(self):
|
||||
return []
|
||||
|
||||
def get_state_layer_ids(self):
|
||||
return []
|
||||
|
||||
def get_state_slice_outer_counts(self):
|
||||
return []
|
||||
|
||||
def get_state_conv_shard_groups(self):
|
||||
return []
|
||||
|
||||
|
||||
class UnifiedMambaSlotAllocator:
|
||||
"""Mamba slot allocator (PHYSICAL view) for the unified memory pool.
|
||||
@@ -929,6 +969,7 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
|
||||
speculative_num_draft_tokens: Optional[int] = None,
|
||||
enable_overlap_schedule: bool = True,
|
||||
start_layer: Optional[int] = None,
|
||||
pre_alloc_size: int = 0,
|
||||
):
|
||||
self._unified_buffer = unified_buffer
|
||||
self._mamba_sub_pool_name = mamba_sub_pool_name
|
||||
@@ -936,7 +977,10 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
|
||||
unified_buffer.max_slots(mamba_sub_pool_name) - 1
|
||||
) # reserve slot 0
|
||||
super().__init__(
|
||||
size=size,
|
||||
# `DecodeReqToTokenPool` semantics: rows cover the preallocated
|
||||
# requests too, while `self.size` (rebound below) stays the
|
||||
# running-request cap the scheduler and leak invariant expect.
|
||||
size=size + pre_alloc_size,
|
||||
mamba_size=self._shared_mamba_size,
|
||||
mamba_spec_state_size=mamba_spec_state_size,
|
||||
max_context_len=max_context_len,
|
||||
@@ -949,6 +993,8 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
|
||||
enable_overlap_schedule=enable_overlap_schedule,
|
||||
start_layer=start_layer,
|
||||
)
|
||||
self.size = size
|
||||
self.pre_alloc_size = pre_alloc_size
|
||||
|
||||
def _init_mamba_pool(
|
||||
self,
|
||||
@@ -1013,6 +1059,17 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
|
||||
return self.mamba_allocator.translate(virtual_ids).to(torch.int32)
|
||||
|
||||
|
||||
class UnifiedHybridLinearKVPool(HybridLinearKVPool):
|
||||
"""`HybridLinearKVPool` over unified sub-pools (full = Unified{MLA,MHA},
|
||||
mamba = UnifiedMambaPool)."""
|
||||
|
||||
def get_kv_layer_ids(self):
|
||||
# Empty: the KV component is one whole-envelope entry, so there are no
|
||||
# per-layer entries to pair by layer id (the sender falls back to
|
||||
# positional pairing).
|
||||
return []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Factory
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1054,9 +1111,9 @@ def init_unified_mamba_pools(
|
||||
mamba_full_memory_ratio: Optional[float] = None, # informational only
|
||||
forward_stream: Optional[torch.cuda.Stream] = None,
|
||||
lazy_compaction: bool = False,
|
||||
decode_pre_alloc_size: int = 0,
|
||||
) -> UnifiedPoolBundle:
|
||||
"""Build the Mamba-hybrid unified-memory-pool stack."""
|
||||
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
|
||||
from sglang.srt.mem_cache.multi_ended_allocator import (
|
||||
UnifiedMambaTokenToKVPoolAllocator,
|
||||
)
|
||||
@@ -1132,6 +1189,7 @@ def init_unified_mamba_pools(
|
||||
speculative_num_draft_tokens=speculative_num_draft_tokens,
|
||||
enable_overlap_schedule=not disable_overlap_schedule,
|
||||
start_layer=start_layer,
|
||||
pre_alloc_size=decode_pre_alloc_size,
|
||||
)
|
||||
if use_mla_backend:
|
||||
# start_layer stays 0: HybridLinearKVPool patches layer ids to the dense
|
||||
@@ -1153,7 +1211,7 @@ def init_unified_mamba_pools(
|
||||
full_attn_layer_ids_for_pool = (
|
||||
[0] if is_draft_worker else list(full_attention_layer_ids)
|
||||
)
|
||||
token_to_kv_pool = HybridLinearKVPool(
|
||||
token_to_kv_pool = UnifiedHybridLinearKVPool(
|
||||
page_size=page_size,
|
||||
size=max_total_num_tokens,
|
||||
dtype=kv_cache_dtype,
|
||||
|
||||
@@ -8007,9 +8007,29 @@ class ServerArgs:
|
||||
def _handle_unified_memory_pool(self):
|
||||
if not self.enable_unified_memory:
|
||||
return
|
||||
assert self.disaggregation_mode == "null", (
|
||||
"--enable-unified-memory is not yet compatible with PD " "disaggregation."
|
||||
)
|
||||
if self.disaggregation_mode != "null":
|
||||
# Constraints of the whole-envelope transfer; see
|
||||
# UnifiedMLATokenToKVPool.get_contiguous_buf_infos.
|
||||
assert self.disaggregation_transfer_backend == "mooncake", (
|
||||
"--enable-unified-memory with PD disaggregation supports only "
|
||||
"the mooncake transfer backend; got "
|
||||
f"{self.disaggregation_transfer_backend!r}."
|
||||
)
|
||||
assert self.pp_size == 1, (
|
||||
"--enable-unified-memory with PD disaggregation does not support "
|
||||
"pipeline parallelism (whole-envelope transfer has no per-layer "
|
||||
"entries to subset)."
|
||||
)
|
||||
assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), (
|
||||
"--enable-unified-memory with PD disaggregation requires lazy "
|
||||
"compaction; unset SGLANG_DISABLE_LAZY_COMPACTION."
|
||||
)
|
||||
assert not self.enable_hisparse, (
|
||||
"--enable-unified-memory with PD disaggregation is not compatible "
|
||||
"with --enable-hisparse: the decode-side HiSparse prealloc path "
|
||||
"ships host/C4 rows straight from the allocator, bypassing the "
|
||||
"virtual->physical translation the unified pool needs."
|
||||
)
|
||||
assert self.speculative_algorithm in (None, "DSPARK"), (
|
||||
"--enable-unified-memory only supports --speculative-algorithm "
|
||||
"DSPARK (chain draft); other speculative algorithms are not yet "
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
"""PD disaggregation with --enable-unified-memory (MLA hybrid-Mamba).
|
||||
|
||||
Guards the unified-memory PD transfer scheme end to end: whole page-envelope
|
||||
KV registration (`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`), whole
|
||||
slot-envelope KDA/mamba state transfer, virtual->physical index translation at
|
||||
the prefill send / decode prealloc sites, and the compaction move gate. A
|
||||
regression in any of them shifts the decode-side KV/state bytes and breaks
|
||||
logprob parity with the non-PD unified-memory reference.
|
||||
|
||||
`--attention-backend` is deliberately NOT pinned, matching
|
||||
`models_e2e/test_kimi_linear_unified_memory.py`, which documents that pinning
|
||||
hides defects reachable only under the resolved default. The transferred bytes
|
||||
are backend-independent, so the default (fa3 on this suite's H100 runner) covers
|
||||
this file's subject either way. The linear-attn/Mamba backends stay pinned to
|
||||
triton -- the page-major layout requires them.
|
||||
|
||||
`--enable-deterministic-inference` is deliberately NOT set. It would only guard
|
||||
against batch-shape-dependent kernel variation, and the reference and P+D paths
|
||||
run the same shapes: measured, two fresh servers on separate GPUs produce
|
||||
bit-identical logits without it. Setting it would narrow the test to the
|
||||
batch-invariant op set and a non-default sampling backend -- a less
|
||||
representative config -- and couple a PD-transfer test to the deterministic code
|
||||
path, so a defect there would fail this file for an unrelated reason.
|
||||
"""
|
||||
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
assert_process_healthy,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100")
|
||||
|
||||
KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
|
||||
SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
|
||||
SERVER_ARGS = [
|
||||
"--skip-tokenizer-init",
|
||||
"--random-seed",
|
||||
"1",
|
||||
"--enable-unified-memory",
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
"--max-mamba-cache-size",
|
||||
"32",
|
||||
"--max-total-tokens",
|
||||
"4096",
|
||||
"--cuda-graph-backend-decode",
|
||||
"disabled",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryDisaggregation(PDDisaggregationServerBase):
|
||||
"""1 prefill + 1 decode, both with --enable-unified-memory, vs a non-PD
|
||||
unified-memory reference server."""
|
||||
|
||||
prefill_tp_size = 1
|
||||
decode_tp_size = 1
|
||||
decode_base_gpu_id = 1
|
||||
extra_prefill_args = SERVER_ARGS
|
||||
extra_decode_args = SERVER_ARGS
|
||||
extra_prefill_env = SERVER_ENV
|
||||
extra_decode_env = SERVER_ENV
|
||||
baseline_args = SERVER_ARGS
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = KIMI_LINEAR_MODEL
|
||||
|
||||
@staticmethod
|
||||
def generate(base_url):
|
||||
response = requests.post(
|
||||
base_url + "/generate",
|
||||
json={
|
||||
"input_ids": [1] + [100 + i % 1000 for i in range(256)],
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 4,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 5,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["meta_info"]
|
||||
|
||||
def test_logprob_parity(self):
|
||||
baseline = popen_launch_server(
|
||||
self.model,
|
||||
self.lb_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--trust-remote-code"] + self.baseline_args,
|
||||
env=SERVER_ENV,
|
||||
)
|
||||
try:
|
||||
reference = self.generate(self.lb_url)
|
||||
finally:
|
||||
kill_process_tree(baseline.pid, wait_timeout=60)
|
||||
time.sleep(5)
|
||||
|
||||
self.launch_all()
|
||||
disaggregated = self.generate(self.lb_url)
|
||||
|
||||
reference_logprobs = reference["output_token_logprobs"]
|
||||
disaggregated_logprobs = disaggregated["output_token_logprobs"]
|
||||
self.assertEqual(
|
||||
[item[1] for item in reference_logprobs],
|
||||
[item[1] for item in disaggregated_logprobs],
|
||||
)
|
||||
self.assertEqual(len(reference_logprobs), 4)
|
||||
for reference_item, disaggregated_item in zip(
|
||||
reference_logprobs, disaggregated_logprobs
|
||||
):
|
||||
self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
|
||||
|
||||
assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
|
||||
assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
|
||||
assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
|
||||
|
||||
|
||||
class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregation):
|
||||
"""Multi-chunk prefill (257-token prompt, 64-token chunks): each chunk's KV
|
||||
pages are translated to physical ids and shipped while later chunks still
|
||||
run, exercising the chunked send path and the prefill-side move gate
|
||||
(`chunked_req.start_send_idx > 0`). The reference server uses the same
|
||||
chunk size so any parity break isolates to the PD transfer.
|
||||
"""
|
||||
|
||||
_chunked_args = SERVER_ARGS + ["--chunked-prefill-size", "64"]
|
||||
extra_prefill_args = _chunked_args
|
||||
extra_decode_args = _chunked_args
|
||||
baseline_args = _chunked_args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Regression tests for the unified-memory PD compaction move gate.
|
||||
|
||||
The gate decides when lazy compaction may relocate physical pages. A page is
|
||||
exposed to the peer from the moment its address is published until the transfer
|
||||
concludes, and for part of that lifetime the request sits in NEITHER end's
|
||||
queue. Both cases below are exactly those windows: an earlier version of the
|
||||
predicates looked only at `disagg_prefill_inflight_queue` /
|
||||
`disagg_decode_transfer_queue` (plus `scheduler.chunked_req`) and returned True
|
||||
here, letting compaction move pages under in-flight RDMA -- silent KV
|
||||
corruption with no crash.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from typing import List, Optional, Set
|
||||
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
unified_memory_disagg_move_gate,
|
||||
)
|
||||
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _FakeTransferQueue:
|
||||
def __init__(self):
|
||||
self.queue: List[object] = []
|
||||
|
||||
|
||||
class _FakePreallocQueue:
|
||||
"""Mirrors the real queue's published-destination bookkeeping."""
|
||||
|
||||
def __init__(self):
|
||||
self._num_published_destinations = 0
|
||||
|
||||
@property
|
||||
def has_published_destinations(self) -> bool:
|
||||
return self._num_published_destinations > 0
|
||||
|
||||
def note_destinations_published(self) -> None:
|
||||
self._num_published_destinations += 1
|
||||
|
||||
def note_destinations_queued(self, count: int) -> None:
|
||||
self._num_published_destinations = max(
|
||||
0, self._num_published_destinations - count
|
||||
)
|
||||
|
||||
|
||||
class _FakeScheduler:
|
||||
def __init__(self, mode: DisaggregationMode):
|
||||
self.disaggregation_mode = mode
|
||||
self.chunked_req: Optional[object] = None
|
||||
self.disagg_prefill_inflight_queue: List[object] = []
|
||||
self.disagg_prefill_pending_chunk_rids: Set[str] = set()
|
||||
self.disagg_decode_transfer_queue = _FakeTransferQueue()
|
||||
self.disagg_decode_prealloc_queue = _FakePreallocQueue()
|
||||
|
||||
|
||||
class TestDecodeMoveGate(CustomTestCase):
|
||||
def test_closed_while_destination_published_but_not_queued(self):
|
||||
"""`pop_preallocated` publishes request A's destination addresses via
|
||||
`send_metadata`, then keeps allocating for request B in the same loop;
|
||||
the batch only reaches the transfer queue after the loop returns. B's
|
||||
allocation can urgently flush the peer sub-allocator, so the gate must
|
||||
stay closed across that window even though the transfer queue is empty.
|
||||
"""
|
||||
scheduler = _FakeScheduler(DisaggregationMode.DECODE)
|
||||
gate = unified_memory_disagg_move_gate(scheduler)
|
||||
self.assertTrue(gate(), "idle decode node should allow compaction")
|
||||
|
||||
# A's destination is now visible to prefill; transfer queue still empty.
|
||||
scheduler.disagg_decode_prealloc_queue.note_destinations_published()
|
||||
self.assertFalse(scheduler.disagg_decode_transfer_queue.queue)
|
||||
self.assertFalse(gate())
|
||||
|
||||
# Handing the batch to the transfer queue transfers responsibility.
|
||||
scheduler.disagg_decode_transfer_queue.queue.append(object())
|
||||
scheduler.disagg_decode_prealloc_queue.note_destinations_queued(1)
|
||||
self.assertFalse(gate(), "transfer queue still holds it")
|
||||
|
||||
scheduler.disagg_decode_transfer_queue.queue.clear()
|
||||
self.assertTrue(gate())
|
||||
|
||||
|
||||
class TestPrefillMoveGate(CustomTestCase):
|
||||
def test_closed_after_final_chunk_clears_chunked_req(self):
|
||||
"""Scheduling the final chunk clears `scheduler.chunked_req`, but the
|
||||
request only reaches `disagg_prefill_inflight_queue` later in the result
|
||||
path. Earlier middle chunks may still be draining in that window, so the
|
||||
gate must not key off `chunked_req` alone.
|
||||
"""
|
||||
scheduler = _FakeScheduler(DisaggregationMode.PREFILL)
|
||||
gate = unified_memory_disagg_move_gate(scheduler)
|
||||
self.assertTrue(gate(), "idle prefill node should allow compaction")
|
||||
|
||||
# A middle chunk went out for rid "r0".
|
||||
scheduler.chunked_req = object()
|
||||
scheduler.disagg_prefill_pending_chunk_rids.add("r0")
|
||||
self.assertFalse(gate())
|
||||
|
||||
# Final chunk scheduled: chunked_req cleared, not yet inflight-queued.
|
||||
scheduler.chunked_req = None
|
||||
self.assertFalse(scheduler.disagg_prefill_inflight_queue)
|
||||
self.assertFalse(gate())
|
||||
|
||||
# Last chunk sent: the request is on the inflight queue, which covers it.
|
||||
scheduler.disagg_prefill_inflight_queue.append(object())
|
||||
scheduler.disagg_prefill_pending_chunk_rids.discard("r0")
|
||||
self.assertFalse(gate())
|
||||
|
||||
scheduler.disagg_prefill_inflight_queue.clear()
|
||||
self.assertTrue(gate())
|
||||
|
||||
def test_reopens_when_middle_sent_request_is_retired_without_final_chunk(self):
|
||||
"""A request aborted after a middle chunk never reaches a `last_chunk`
|
||||
send, so its rid is only dropped by the abort/release cleanup. Without
|
||||
that discard the gate stays closed for the process lifetime and lazy
|
||||
compaction never packs the free list again -- a liveness leak that ends
|
||||
in allocation failure despite reclaimable space.
|
||||
"""
|
||||
scheduler = _FakeScheduler(DisaggregationMode.PREFILL)
|
||||
gate = unified_memory_disagg_move_gate(scheduler)
|
||||
|
||||
scheduler.chunked_req = object()
|
||||
scheduler.disagg_prefill_pending_chunk_rids.add("r0")
|
||||
self.assertFalse(gate())
|
||||
|
||||
# Aborted mid-chunking: chunked_req dropped, no final send, never queued.
|
||||
scheduler.chunked_req = None
|
||||
scheduler.disagg_prefill_pending_chunk_rids.discard("r0")
|
||||
self.assertTrue(gate(), "abort cleanup must let compaction resume")
|
||||
|
||||
|
||||
class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
|
||||
"""`schedulable_available_size` credits holes a peer urgent-flush would
|
||||
release. While the move gate is closed that flush relocates nothing, so
|
||||
crediting them lets the scheduler admit work `_flush_peer_for_alloc` cannot
|
||||
satisfy; the alloc then returns None and the decode prealloc path treats
|
||||
that as a memory-estimation bug and aborts the scheduler.
|
||||
"""
|
||||
|
||||
class _Peer:
|
||||
def __init__(self, gate):
|
||||
self.lazy_compaction = True
|
||||
self._free_phys_pages = [0, 1, 2, 3] # only len() is read
|
||||
self.entry_bytes_per_page = 512
|
||||
self.disagg_move_gate = gate
|
||||
|
||||
class _Owner:
|
||||
def __init__(self, peer):
|
||||
self._peer = peer
|
||||
|
||||
def _credit(self, gate):
|
||||
peer = self._Peer(gate)
|
||||
owner = self._Owner(peer)
|
||||
return MultiEndedAllocator._peer_drainable_hole_bytes(owner)
|
||||
|
||||
def test_credit_follows_the_gate(self):
|
||||
# No PD gate installed (non-disagg): holes are realizable as before.
|
||||
self.assertEqual(self._credit(gate=None), 4 * 512)
|
||||
# Gate open: peer can compact, so the credit stands.
|
||||
self.assertEqual(self._credit(gate=lambda: True), 4 * 512)
|
||||
# Gate closed: an urgent flush would move nothing, so credit nothing.
|
||||
self.assertEqual(self._credit(gate=lambda: False), 0)
|
||||
|
||||
|
||||
class TestMoveGateRejectsNonPdNode(CustomTestCase):
|
||||
def test_null_mode_is_rejected(self):
|
||||
"""The gate is only meaningful on a PD node; a NULL-mode scheduler is a
|
||||
wiring bug and must not silently produce an always-open predicate."""
|
||||
scheduler = _FakeScheduler(DisaggregationMode.NULL)
|
||||
with self.assertRaises(ValueError):
|
||||
unified_memory_disagg_move_gate(scheduler)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -109,6 +109,9 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
|
||||
queue.pending_reqs = []
|
||||
queue.retracted_queue = []
|
||||
queue.num_reserved_decode_tokens = 0
|
||||
# `pop_preallocated` credits this counter; `__new__` skips the __init__
|
||||
# that seeds it.
|
||||
queue._num_published_destinations = 0
|
||||
queue._resolve_pending_reqs = MagicMock()
|
||||
queue._update_handshake_waiters = MagicMock()
|
||||
queue._allocatable_tokens = MagicMock(return_value=1000)
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
from sglang.srt.mem_cache.allocator.hisparse import (
|
||||
DeepSeekV4HiSparseTokenToKVPoolAllocator,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
@@ -214,14 +215,18 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
|
||||
manager._send_kvcache_generic = MagicMock(return_value=0)
|
||||
executor = MagicMock()
|
||||
|
||||
manager.send_kvcache(
|
||||
"session",
|
||||
np.array([1], dtype=np.int32),
|
||||
[10000, 20000, 30000],
|
||||
np.array([7], dtype=np.int32),
|
||||
executor,
|
||||
dst_device_kv_indices=np.array([21], dtype=np.int32),
|
||||
)
|
||||
# send_kvcache reads the memory bag (the unified-memory envelope-layout
|
||||
# check), so the context has to be published. This is the non-unified
|
||||
# path -- pin that explicitly rather than leaning on the default.
|
||||
with get_context().override_server_args(enable_unified_memory=False):
|
||||
manager.send_kvcache(
|
||||
"session",
|
||||
np.array([1], dtype=np.int32),
|
||||
[10000, 20000, 30000],
|
||||
np.array([7], dtype=np.int32),
|
||||
executor,
|
||||
dst_device_kv_indices=np.array([21], dtype=np.int32),
|
||||
)
|
||||
|
||||
kwargs = manager._send_kvcache_generic.call_args.kwargs
|
||||
self.assertEqual(kwargs["dst_device_data_ptrs"], {20000, 30000})
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Derived-property tests for the PD whole-envelope transfer addressing.
|
||||
|
||||
PD disaggregation transfers the unified memory pool as whole envelopes with
|
||||
``addr = raw_ptr + physical_index * item_len`` (see
|
||||
``UnifiedMLATokenToKVPool.get_contiguous_buf_infos`` /
|
||||
``UnifiedMambaPool.get_contiguous_buf_infos`` and mooncake's
|
||||
``_send_kvcache_generic`` / ``_send_mamba_state``). That contract only holds if
|
||||
the page-major view builders keep (a) one page's data for ALL layers inside one
|
||||
contiguous ``page_envelope_bytes`` block, and (b) one mamba slot's conv+temporal
|
||||
state for all layers inside one contiguous ``entry_bytes`` block. A
|
||||
"looks equivalent" reordering of the view layout (e.g. layer-major across
|
||||
pages) would silently corrupt every PD transfer while all kernels keep working,
|
||||
because kernels read through the strided views, not through raw offsets.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_dense_mla_views,
|
||||
build_page_major_mamba_views,
|
||||
mamba_entry_bytes,
|
||||
mla_entry_bytes,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=60, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestMLAEnvelopeTransferAddressing(CustomTestCase):
|
||||
def test_page_envelope_matches_dense_views(self):
|
||||
"""Every (page, layer, slot) row written through the dense MLA views
|
||||
must land at raw_ptr + page * page_envelope_bytes + layer-block offset,
|
||||
i.e. inside the page's transfer envelope."""
|
||||
layer_num, page_size, kv_dim, num_pages = 3, 4, 8, 6
|
||||
store_dtype = torch.bfloat16
|
||||
row_bytes = kv_dim * store_dtype.itemsize
|
||||
page_bytes = page_size * layer_num * row_bytes
|
||||
self.assertEqual(
|
||||
page_bytes,
|
||||
page_size
|
||||
* mla_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
kv_cache_dim=kv_dim,
|
||||
itemsize=store_dtype.itemsize,
|
||||
),
|
||||
)
|
||||
# +1 page envelope of tail pad, as UnifiedKVPool allocates for MLA.
|
||||
raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8)
|
||||
views = build_dense_mla_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
kv_cache_dim=kv_dim,
|
||||
store_dtype=store_dtype,
|
||||
page_size=page_size,
|
||||
num_pages=num_pages,
|
||||
anchor_bytes=0,
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
for page in range(num_pages):
|
||||
for layer in range(layer_num):
|
||||
for off in range(page_size):
|
||||
dense_id = page * layer_num * page_size + off
|
||||
val = torch.randn(kv_dim, dtype=store_dtype)
|
||||
views[layer][dense_id, 0] = val
|
||||
start = (
|
||||
page * page_bytes
|
||||
+ layer * page_size * row_bytes
|
||||
+ off * row_bytes
|
||||
)
|
||||
got = raw[start : start + row_bytes].view(store_dtype)
|
||||
self.assertTrue(torch.equal(got, val), (page, layer, off))
|
||||
|
||||
|
||||
class TestMambaEnvelopeTransferAddressing(CustomTestCase):
|
||||
def test_slot_envelope_is_self_contained(self):
|
||||
"""A slot's conv+temporal state for all layers must live exactly in
|
||||
raw[slot * entry_bytes : (slot+1) * entry_bytes]: no byte outside the
|
||||
envelope may change, and the payload byte count must fill it."""
|
||||
layer_num, max_slots = 2, 5
|
||||
conv_shapes = ((3, 4), (2, 6))
|
||||
temporal_shape = (2, 3, 4)
|
||||
conv_dtype = torch.bfloat16
|
||||
temporal_dtype = torch.float32
|
||||
entry = mamba_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
conv_state_shapes=conv_shapes,
|
||||
conv_dtype=conv_dtype,
|
||||
temporal_state_shape=temporal_shape,
|
||||
temporal_dtype=temporal_dtype,
|
||||
)
|
||||
raw = torch.zeros(max_slots * entry, dtype=torch.uint8)
|
||||
conv_views, temporal_view = build_page_major_mamba_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
conv_state_shapes=conv_shapes,
|
||||
conv_dtype=conv_dtype,
|
||||
temporal_state_shape=temporal_shape,
|
||||
temporal_dtype=temporal_dtype,
|
||||
max_slots=max_slots,
|
||||
anchor_bytes=0,
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
for slot in range(max_slots):
|
||||
raw.zero_()
|
||||
n_payload = 0
|
||||
for i, conv_view in enumerate(conv_views):
|
||||
val = torch.randn((layer_num,) + conv_shapes[i], dtype=conv_dtype)
|
||||
conv_view[:, slot] = val
|
||||
n_payload += val.numel() * val.element_size()
|
||||
val = torch.randn((layer_num,) + temporal_shape, dtype=temporal_dtype)
|
||||
temporal_view[:, slot] = val
|
||||
n_payload += val.numel() * val.element_size()
|
||||
|
||||
outside = torch.cat([raw[: slot * entry], raw[(slot + 1) * entry :]])
|
||||
self.assertTrue(
|
||||
bool(outside.eq(0).all()),
|
||||
f"slot {slot} state bled outside its transfer envelope",
|
||||
)
|
||||
self.assertEqual(n_payload, entry)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user