From ceeaec207887206818e94b378139be0756738097 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:07:58 -0700 Subject: [PATCH] [PD] Support --enable-unified-memory with PD disaggregation (kimi-linear MLA hybrid-Mamba) (#33362) Co-authored-by: Claude Fable 5 --- python/sglang/srt/disaggregation/decode.py | 32 +++- .../srt/disaggregation/mooncake/conn.py | 118 +++++++++--- python/sglang/srt/disaggregation/prefill.py | 33 +++- python/sglang/srt/disaggregation/utils.py | 44 +++++ python/sglang/srt/managers/scheduler.py | 14 +- python/sglang/srt/mem_cache/allocator/base.py | 10 + .../srt/mem_cache/kv_cache_configurator.py | 27 ++- .../srt/mem_cache/multi_ended_allocator.py | 37 +++- .../srt/mem_cache/unified_memory_pool.py | 64 ++++++- python/sglang/srt/server_args.py | 26 ++- .../test_disaggregation_unified_memory.py | 152 +++++++++++++++ .../test_unified_memory_move_gate.py | 179 ++++++++++++++++++ ...test_priority_scheduling_disaggregation.py | 3 + .../unit/mem_cache/test_hisparse_allocator.py | 21 +- .../test_pd_envelope_transfer_layout.py | 126 ++++++++++++ 15 files changed, 835 insertions(+), 51 deletions(-) create mode 100644 test/registered/disaggregation/test_disaggregation_unified_memory.py create mode 100644 test/registered/unit/disaggregation/test_unified_memory_move_gate.py create mode 100644 test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 0a2d8b36a..232e266a3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 5780fdf7e..255ff045f 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index d50c635db..09ba12c24 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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 diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index e6f35fac7..16e303377 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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 ######################### diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index dd035c148..f79f0108e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 532cee46a..103324dcd 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -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() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index b8f72b146..8aa875482 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -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 diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 6f70d23a4..98a7bb190 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -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) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 6c9aa877b..0d8acc7b1 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 477a648b9..a11b04a34 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 " diff --git a/test/registered/disaggregation/test_disaggregation_unified_memory.py b/test/registered/disaggregation/test_disaggregation_unified_memory.py new file mode 100644 index 000000000..02e2cac4c --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_unified_memory.py @@ -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() diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py new file mode 100644 index 000000000..fc93bf7d4 --- /dev/null +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -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() diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index afac4e240..cabe21066 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -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) diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index bbf0428e1..923267af5 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -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}) diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py new file mode 100644 index 000000000..f357c2aa9 --- /dev/null +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -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()