From 4a8e76805c27d8fffed4b8759593310a6ebde298 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:21:59 -0700 Subject: [PATCH] feat(mem_cache): unified memory pool for hybrid Mamba / SWA models (#29678) Co-authored-by: lch1475369 --- .../advanced_features/server_arguments.mdx | 6 + python/sglang/srt/environ.py | 9 + .../hardware_backend/npu/memory_pool_npu.py | 4 +- .../attention/hybrid_linear_attn_backend.py | 363 +-- .../layers/attention/mamba/mamba2_metadata.py | 7 + .../mamba/mamba_state_scatter_triton.py | 9 + .../srt/layers/attention/triton_backend.py | 251 +- python/sglang/srt/managers/schedule_policy.py | 112 +- python/sglang/srt/managers/scheduler.py | 31 +- .../scheduler_components/invariant_checker.py | 12 +- .../sglang/srt/mem_cache/allocator/mamba.py | 7 + python/sglang/srt/mem_cache/allocator/swa.py | 9 + python/sglang/srt/mem_cache/common.py | 88 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 21 +- python/sglang/srt/mem_cache/memory_pool.py | 78 +- .../srt/mem_cache/multi_ended_allocator.py | 2474 +++++++++++++++ .../sglang/srt/mem_cache/swa_memory_pool.py | 2 +- .../srt/mem_cache/triton_ops/virtual_slot.py | 96 + .../mamba_component.py | 6 +- .../srt/mem_cache/unified_memory_pool.py | 1369 +++++++++ .../model_runner_kv_cache_mixin.py | 183 ++ .../runner/decode_cuda_graph_runner.py | 21 +- python/sglang/srt/server_args.py | 52 + .../unit/mem_cache/test_full_loc_fast_path.py | 206 ++ .../unit/mem_cache/test_layout_compat.py | 355 +++ .../mem_cache/test_multi_ended_allocator.py | 2729 +++++++++++++++++ .../unit/mem_cache/test_store_cache_4d.py | 138 +- .../mem_cache/test_unified_mamba_views.py | 291 ++ 28 files changed, 8378 insertions(+), 551 deletions(-) create mode 100644 python/sglang/srt/mem_cache/multi_ended_allocator.py create mode 100644 python/sglang/srt/mem_cache/triton_ops/virtual_slot.py create mode 100644 python/sglang/srt/mem_cache/unified_memory_pool.py create mode 100644 test/registered/unit/mem_cache/test_full_loc_fast_path.py create mode 100644 test/registered/unit/mem_cache/test_layout_compat.py create mode 100644 test/registered/unit/mem_cache/test_multi_ended_allocator.py create mode 100644 test/registered/unit/mem_cache/test_unified_mamba_views.py diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index 98c5f5d49..de9bc606a 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -515,6 +515,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s Enable the page-major KV layout: lay out the Mamba state and full/SWA KV caches in a page-granularity envelope (page is the outermost axis, layer-major within a page) instead of the default per-layer (layer-major) layout. Requires the Triton attention / linear-attn / Mamba backends (`--attention-backend triton`, and for hybrid models `--linear-attn-backend triton --mamba-backend triton`). `False` bool flag (set to enable) + + + `--enable-unified-memory` + For hybrid Mamba/GDN and hybrid SWA models, replace the statically-partitioned pools (full-attention KV + SWA/Mamba conv/SSM state) with a single byte buffer split dynamically between the sub-pools, so KV-vs-state capacity flexes with the workload instead of being fixed at startup. Implies `--enable-page-major-kv-layout`. Requires the Triton attention / linear-attn / Mamba backends; monolithic (decode) cuda-graph capture only; not yet compatible with PD disaggregation or speculative decoding. + `False` + bool flag (set to enable) `--swa-full-tokens-ratio` diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7e7dc1bbe..fefa853c4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -301,6 +301,15 @@ class Envs: SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500) SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64) SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False) + # Disable lazy compaction in the unified memory pool allocator and + # fall back to the per-free eager compaction. Used for production + # A/B and quick rollback. Default False (lazy compaction on). + SGLANG_DISABLE_LAZY_COMPACTION = EnvBool(False) + # Sort the multi-ended allocator's free list after a merge (perf A/B knob). + SGLANG_SORT_FREE_LIST_AFTER_MERGE = EnvBool(False) + # Periodically log lazy-compaction stats per sub-pool (observability only). + SGLANG_LOG_LAZY_COMPACTION_STATS = EnvBool(False) + SGLANG_LOG_LAZY_COMPACTION_STATS_INTERVAL_SEC = EnvInt(30) SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(True) SGLANG_TEST_DISAGG_FAILURE_PROB = EnvFloat(0.0) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index f60083929..718b46220 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -179,7 +179,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): layer_id_override: Optional[int] = None, dcp_kv_mask: Optional[torch.Tensor] = None, ): - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) if layer_id_override is not None: layer_id = layer_id_override else: @@ -441,7 +441,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): cache_k: torch.Tensor, cache_v: torch.Tensor, ): - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) layer_id = layer.layer_id if cache_k.dtype != self.dtype: cache_k = cache_k.to(self.dtype) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 57100a9b9..93df15c6d 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -35,15 +35,15 @@ class MambaAttnBackendBase(AttentionBackend): self.is_draft_worker = model_runner.is_draft_worker self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool self.token_to_kv_pool = model_runner.token_to_kv_pool + self.enable_unified_memory = model_runner.server_args.enable_unified_memory self.forward_metadata: ForwardMetadata = None self.state_indices_list = [] - # GDN ReplaySSM (slice 1b): per-bs STATIC per-row write-cursor buffers - # for cuda-graph. Allocated lazily in init_cuda_graph_state only when - # --enable-linear-replayssm is set; stays None otherwise. + # Static (max_bs,) track-dest buffer captured by pointer, refreshed in-place + # each replay; the captured track-save reads this, not the InputBuffer slot. + self.mamba_track_indices_buf = None + # Per-bs static write-cursor / force-flush buffers for cuda-graph; None + # unless --enable-linear-replayssm is set. self.replayssm_write_pos_list = None - # GDN ReplaySSM (slice 2b): per-bs STATIC per-row force-flush buffers - # for cuda-graph, parallel to replayssm_write_pos_list. Same lifetime - # (None unless the flag is on). self.replayssm_force_flush_list = None self.query_start_loc_list = [] self.retrieve_next_token_list = [] @@ -66,8 +66,10 @@ class MambaAttnBackendBase(AttentionBackend): forward_batch.mamba_clear_indices is not None and len(forward_batch.mamba_clear_indices) > 0 ): + # mamba_pool is a pure PHYSICAL store; translate before zeroing or + # clear_slots zeroes the wrong physical slots. self.req_to_token_pool.mamba_pool.clear_slots( - forward_batch.mamba_clear_indices + self._translate_mamba_indices(forward_batch.mamba_clear_indices) ) if ( forward_batch.mamba_cow_src_indices is not None @@ -75,22 +77,28 @@ class MambaAttnBackendBase(AttentionBackend): ): ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None) if ckpt_pool is not None: - # int8 checkpoints: dequantize the cached state (src = int8 ckpt slot) - # into the request's active bf16 slot (dst). + # int8 checkpoints: dequantize src int8 ckpt slot into the active bf16 dst. ckpt_pool.load_to_active( self.req_to_token_pool.mamba_pool, forward_batch.mamba_cow_src_indices, forward_batch.mamba_cow_dst_indices, ) else: + # mamba_pool is a pure PHYSICAL store; translate both COW slot ids. self.req_to_token_pool.mamba_pool.copy_from( - forward_batch.mamba_cow_src_indices, - forward_batch.mamba_cow_dst_indices, + self._translate_mamba_indices(forward_batch.mamba_cow_src_indices), + self._translate_mamba_indices(forward_batch.mamba_cow_dst_indices), ) forward_batch.mamba_clear_indices = None forward_batch.mamba_cow_src_indices = None forward_batch.mamba_cow_dst_indices = None + def _translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor: + """Virtual->physical mamba slot-id translate (identity for the non-unified + pool). Must run everywhere mamba ids feed the SSM/conv kernels or mamba-pool + state ops, incl. the cuda-graph replay-prep copy into ``state_indices_list``.""" + return self.req_to_token_pool.translate_mamba_indices(mamba_indices) + def _forward_metadata(self, forward_batch: ForwardBatch): bs = forward_batch.batch_size @@ -106,7 +114,14 @@ class MambaAttnBackendBase(AttentionBackend): mamba_cache_indices = self.req_to_token_pool.get_mamba_indices( forward_batch.req_pool_indices ) - _real_bs = getattr(forward_batch, "_original_batch_size", None) + # Translate virtual->physical BEFORE the padding sentinel below, so the + # gather reads only real ids; padded rows are then poisoned to -1 (skipped). + mamba_cache_indices = self._translate_mamba_indices(mamba_cache_indices) + if forward_batch.mamba_track_indices is not None: + forward_batch.mamba_track_indices = self._translate_mamba_indices( + forward_batch.mamba_track_indices + ) + _real_bs = forward_batch._original_batch_size if _real_bs is not None and _real_bs < mamba_cache_indices.shape[0]: mamba_cache_indices = mamba_cache_indices.clone() mamba_cache_indices[_real_bs:] = -1 @@ -117,11 +132,8 @@ class MambaAttnBackendBase(AttentionBackend): query_start_loc = torch.arange( 0, bs + 1, dtype=torch.int32, device=self.device ) - # GDN ReplaySSM (slice 1a): the ring cursor is a per-slot - # decode-position counter shared by ALL GDN layers in this forward. - # Manage it exactly ONCE here (not per-layer): snapshot this step's - # value for the batch's slots, hand it to the layers, then advance - # the persistent buffer mod L for the next step. + # The ring cursor is a per-slot decode counter shared by all GDN layers; + # manage it once here (snapshot, hand to layers, advance mod L), not per-layer. mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None) write_pos_buf = ( getattr(mamba_pool, "replayssm_write_pos", None) @@ -130,27 +142,16 @@ class MambaAttnBackendBase(AttentionBackend): ) if write_pos_buf is not None: slots = mamba_cache_indices.to(torch.long) - # Padded rows carry slot == -1; clamp so the per-row gather stays - # in-bounds (the kernel zeroes padded rows via state_idx < 0). + # Padded rows carry slot == -1; clamp the gather in-bounds (kernel + # zeroes padded rows via state_idx < 0). safe_slots = slots.clamp(min=0) replayssm_write_pos = write_pos_buf[safe_slots].clone() L = mamba_pool.linear_replayssm_cache_len - # KDA (per-K gate) ships without radix coordination for now: no - # track-boundary force-flush, so the ring flushes only at the - # natural write_pos == L-1 wrap. GDN keeps the radix-aligned - # force-flush (slice 2b). Gate on the pool's recorded gate type. + # KDA has no radix coordination: flush only on the natural write_pos + # == L-1 wrap. GDN adds the radix-aligned force-flush below. is_kda = getattr(mamba_pool, "replayssm_is_kda", False) - # GDN ReplaySSM (slice 2b): per-row force-flush at the radix - # track boundary. THE alignment: the radix mamba track snapshots - # temporal[slot] when seq_lens_cpu % mamba_track_interval == 0 - # (extra_buffer: schedule_batch.prepare_for_decode builds - # `mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)` - # off the SAME post-increment seq_lens_cpu used here). We source - # the flush from the identical seq_lens + condition so the kernel - # folds the ring into temporal[slot] on EXACTLY the steps the - # snapshot reads it. seq_lens_cpu is the committed length AFTER - # this decode token (incremented in prepare_for_decode before the - # forward), matching the track. int32, one entry per batch row. + # Force-flush on the radix track's seq_lens % mamba_track_interval + # == 0 boundary so the ring folds into temporal[slot] when read. if not is_kda: force_flush_bool = self._replayssm_track_flush_mask( forward_batch.seq_lens_cpu, bs @@ -158,16 +159,11 @@ class MambaAttnBackendBase(AttentionBackend): replayssm_force_flush = force_flush_bool.to( device=self.device, dtype=torch.int32 ) - # Advance only the VALID (non-padded) slots. Scatter over the - # unique valid slots to avoid duplicate-index races (padded rows - # all clamp to slot 0, which a real row may also occupy). A - # forced flush empties the ring -> next write_pos is 0 (same as - # the natural wrap at write_pos == L-1). + # Advance only valid slots, scattered over unique slots (dup-index + # race; padded rows clamp to 0); a forced flush -> next write_pos 0. valid_mask = slots >= 0 valid_slots = slots[valid_mask] if valid_slots.numel() > 0: - # Per-row "did this step flush?": natural wrap OR forced. - # (KDA has no forced flush -> force_flush is None -> pure wrap.) flushed = replayssm_write_pos == (L - 1) if replayssm_force_flush is not None: flushed = flushed | (replayssm_force_flush != 0) @@ -176,9 +172,8 @@ class MambaAttnBackendBase(AttentionBackend): torch.zeros_like(replayssm_write_pos), (replayssm_write_pos + 1) % L, ) - # Dedup valid slots; for duplicates a scatter picks one - # arbitrary row, but all rows of a given slot share the same - # write_pos/flush, so the value is identical regardless. + # Dedup: rows sharing a slot share write_pos/flush, so the + # scattered value is identical regardless of which row wins. uniq_slots, inv = torch.unique(valid_slots, return_inverse=True) next_for_valid = next_pos[valid_mask] new_vals = torch.empty( @@ -190,9 +185,8 @@ class MambaAttnBackendBase(AttentionBackend): write_pos_buf[uniq_slots] = new_vals elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True): if forward_batch.forward_mode.is_draft_extend_v2(): - # HybridLinearAttnBackend.init_forward_metadata calls all sub-backends - # unconditionally, but DRAFT_EXTEND_V2 only runs full-attn layers in - # the draft model, so mamba metadata can be skipped. + # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; + # skip mamba metadata. query_start_loc = None elif forward_batch.forward_mode.is_target_verify(): query_start_loc = torch.arange( @@ -208,7 +202,7 @@ class MambaAttnBackendBase(AttentionBackend): retrieve_next_sibling = ( forward_batch.spec_info.retrieve_next_sibling ) - # retrieve_next_token is None during dummy run so skip tensor creation + # None during dummy run if retrieve_next_token is not None: retrieve_parent_token = torch.empty_like(retrieve_next_token) else: @@ -245,6 +239,9 @@ class MambaAttnBackendBase(AttentionBackend): return ForwardMetadata( query_start_loc=query_start_loc, mamba_cache_indices=mamba_cache_indices, + # Physical track destinations (None when tracking off); cuda-graph + # supplies this via the static backend buffer in _replay_metadata. + mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None), retrieve_next_token=retrieve_next_token, retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, @@ -263,8 +260,6 @@ class MambaAttnBackendBase(AttentionBackend): forward_batch: ForwardBatch, in_capture: bool = False, ): - # seq_lens_cpu is unused by _replay_metadata for the non-target-verify - # case but kept in the contract for compatibility. self.forward_metadata = self._replay_metadata( forward_batch.batch_size, forward_batch.req_pool_indices, @@ -275,6 +270,7 @@ class MambaAttnBackendBase(AttentionBackend): 0 if in_capture else getattr(forward_batch, "num_padding", None) ), in_capture=in_capture, + mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None), ) def init_forward_metadata(self, forward_batch: ForwardBatch): @@ -284,28 +280,10 @@ class MambaAttnBackendBase(AttentionBackend): def _init_track_conv_indices( self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch ): - """ - Compute indices for extracting conv states from the input sequence during extend. - - In Mamba models, the conv layer maintains a sliding window of recent inputs. - After processing a prefill chunk, we need to save the last `conv_state_len` tokens - of the processed region for prefix caching. - - The key insight is that FLA (Flash Linear Attention) and Mamba2 processes sequences in chunks - of the chunk size (FLA_CHUNK_SIZE=64 for FLA, mamba_chunk_size for Mamba2). - We only track the conv state up to the last complete chunk boundary (aligned_len). - - start_indices is the starting token index of the conv state to track in this extend batch. - indices include all pos to track in this extend batch, conv_state_len for each req that - needs to be tracked (i.e. mamba_track_mask is True) - - Returns: - indices: Tensor of shape [num_tracked_requests, conv_state_len] containing - flattened positions into the packed input tensor. - """ + """Flattened input positions of conv states to track during extend (up to + the last complete chunk boundary, mamba_track_mask rows only).""" conv_state_len = self.conv_states_shape[-1] - # Calculate the end position of the last aligned chunk lens_to_track = ( forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens ) @@ -314,7 +292,6 @@ class MambaAttnBackendBase(AttentionBackend): start_indices = query_start_loc[:-1] + aligned_len - conv_state_len start_indices = start_indices[forward_batch.mamba_track_mask] - # Create indices: [batch_size, conv_state_len] indices = start_indices.unsqueeze(-1) + torch.arange( conv_state_len, device=self.device, @@ -326,46 +303,11 @@ class MambaAttnBackendBase(AttentionBackend): def _init_track_ssm_indices( self, mamba_cache_indices: torch.Tensor, forward_batch: ForwardBatch ): - """ - Compute source and destination indices for tracking SSM states for prefix caching. - - After processing a prefill, we need to save the SSM recurrent state for prefix caching. - The kernel outputs intermediate hidden states `h` at each chunk boundary, - plus a `last_recurrent_state` at the end of the chunked prefill size. - - The chunk size varies by model type: - - FLA models: FLA_CHUNK_SIZE (64) - - Mamba2 models: mamba_chunk_size (256) - - The challenge is that sequences may or may not end on a chunk boundary: - - Aligned case (len % chunk_size == 0): The to-cache state is stored in - the last_recurrent_state. - - Unaligned case (len % chunk_size != 0): The last_recurrent_state includes the - unaligned position, but we only want state up to the last chunk boundary. - We must extract from the intermediate `h` tensor at the appropriate chunk index. - - We compute the src and dst indices for all requests that need to be cached - (i.e. mamba_track_mask is True) based on the rule above. - - For example (assuming chunk_size=64): - 1. If chunked prefill length is < chunk_size, then only final state has value. - In this case we cache `final` state. - 2. If chunked prefill length == chunk_size, then only final state has value. - In this case we cache pos chunk_size, from `final` state. - 3. If chunked prefill length > chunk_size and < 2 * chunk_size, then both h and - final state have value. We cache pos chunk_size from `h` state. - 4. If chunked prefill length == 2 * chunk_size, then both h and final state have - value. We cache pos 2 * chunk_size from `final` state. Note `h` doesn't include - the final position. - - Returns: - track_ssm_h_src: Source indices into the packed `h` tensor (for unaligned seqs) - track_ssm_h_dst: Destination cache slot indices (for unaligned seqs) - track_ssm_final_src: Source indices into last_recurrent_state buffer (for aligned seqs) - track_ssm_final_dst: Destination cache slot indices (for aligned seqs) - """ + """src/dst indices to track SSM states for prefix caching: aligned seqs + cache last_recurrent_state, unaligned cache intermediate `h` at the last + chunk boundary.""" mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size - # Move to CPU to avoid kernel launches for masking operations + # CPU to avoid kernel launches for the masking ops mamba_track_mask = forward_batch.mamba_track_mask.cpu() extend_seq_lens = forward_batch.extend_seq_lens.cpu() mamba_track_indices = forward_batch.mamba_track_indices.cpu() @@ -373,38 +315,33 @@ class MambaAttnBackendBase(AttentionBackend): mamba_track_seqlens = forward_batch.mamba_track_seqlens.cpu() prefix_lens = forward_batch.extend_prefix_lens.cpu() - # Calculate the number of hidden states per request if isinstance(self, Mamba2AttnBackend): num_h_states = extend_seq_lens // mamba_cache_chunk_size else: num_h_states = (extend_seq_lens - 1) // mamba_cache_chunk_size + 1 - # Calculate the starting offset for each sequence in the packed batch track_ssm_src_offset = torch.zeros_like(num_h_states) track_ssm_src_offset[1:] = torch.cumsum(num_h_states[:-1], dim=0) - # Filter variables by track mask lens_to_track = mamba_track_seqlens - prefix_lens lens_masked = lens_to_track[mamba_track_mask] offset_masked = track_ssm_src_offset[mamba_track_mask] dst_masked = mamba_track_indices[mamba_track_mask] - # Determine if the sequence ends at a chunk boundary is_aligned = (lens_masked % mamba_cache_chunk_size) == 0 - # Case 1: Aligned. Use last_recurrent_state from ssm_states. + # Aligned: last_recurrent_state from ssm_states. track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned] track_ssm_final_dst = dst_masked[is_aligned] - # Case 2: Unaligned. Use intermediate state from h. - # TODO: if support mamba_cache_chunk_size % page size != 0, then need to modify this + # Unaligned: intermediate state from h. + # TODO: handle mamba_cache_chunk_size % page size != 0 not_aligned = ~is_aligned track_ssm_h_src = offset_masked[not_aligned] + ( lens_masked[not_aligned] // mamba_cache_chunk_size ) track_ssm_h_dst = dst_masked[not_aligned] - # Move back to GPU return ( track_ssm_h_src.to(self.device, non_blocking=True), track_ssm_h_dst.to(self.device, non_blocking=True), @@ -427,12 +364,8 @@ class MambaAttnBackendBase(AttentionBackend): ) def _replayssm_enabled(self) -> bool: - """True iff --enable-linear-replayssm allocated the persistent ring cursor. - - The per-slot ``replayssm_write_pos`` buffer on MambaPool is None unless - the flag is set, so it doubles as the on/off gate (same signal that - ``_forward_metadata`` / ``GDNAttnBackend.forward_decode`` already use). - """ + """True iff --enable-linear-replayssm allocated the ring cursor + (MambaPool.replayssm_write_pos doubles as the on/off gate).""" mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None) if mamba_pool is None: return False @@ -441,21 +374,12 @@ class MambaAttnBackendBase(AttentionBackend): def _replayssm_track_flush_mask( self, seq_lens_cpu: torch.Tensor, bs: int ) -> torch.Tensor: - """Per-row bool flush mask == the radix mamba-track snapshot condition. - - THE alignment (slice 2b): the radix mamba track snapshots temporal[slot] - exactly when ``seq_lens_cpu % mamba_track_interval == 0`` (the same mask - ``schedule_batch.prepare_for_decode`` builds for extra_buffer at - ``mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)``). Both - read the SAME post-increment ``seq_lens_cpu`` (committed length AFTER - this decode token), so the kernel force-flush fires on EXACTLY the steps - the snapshot reads the checkpoint -- no off-by-one. Returns a CPU bool - tensor of length ``bs`` (caller moves it to device as int32). - """ + """Per-row (length bs) bool flush mask = the radix track's seq_lens_cpu % + mamba_track_interval == 0, so force-flush and snapshot fire on the same + steps (no off-by-one).""" interval = get_global_server_args().mamba_track_interval if seq_lens_cpu is None: - # Decode without a CPU seq-len mirror should not happen for the - # supported (no_buffer, radix-on) config, but stay safe: never flush. + # Should not happen for the supported config; stay safe and never flush. return torch.zeros((bs,), dtype=torch.bool) mask = (seq_lens_cpu[:bs].to(torch.int64) % interval) == 0 if mask.shape[0] < bs: @@ -468,16 +392,15 @@ class MambaAttnBackendBase(AttentionBackend): max_num_tokens % max_bs == 0 ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" draft_token_num = max_num_tokens // max_bs - # GDN ReplaySSM (slice 1b): per-batch-size STATIC per-row write-cursor - # buffers the kernel reads. Captured into the graph by pointer, so they - # must be the SAME tensor objects refreshed in-place each replay. Sized - # and indexed like state_indices_list ((i+1,), indexed [bs - 1]). Left - # None when the flag is off so the dispatch falls through unchanged. + # Per-bs static write-cursor / force-flush buffers, captured by pointer + + # refreshed in-place each replay; sized like state_indices_list. None when off. self.replayssm_write_pos_list = [] if self._replayssm_enabled() else None - # GDN ReplaySSM (slice 2b): static per-bs force-flush buffers, captured - # by pointer and refreshed in-place per replay just like the write-pos - # buffers. None when the flag is off. self.replayssm_force_flush_list = [] if self._replayssm_enabled() else None + # int64 to match DecodeInputBuffers.mamba_track_indices + the track-save + # kernel's int64 index load. Refreshed in-place by _replay_metadata. + self.mamba_track_indices_buf = torch.zeros( + (max_bs,), dtype=torch.int64, device=self.device + ) for i in range(max_bs): self.state_indices_list.append( torch.full( @@ -556,32 +479,26 @@ class MambaAttnBackendBase(AttentionBackend): else: raise ValueError(f"Invalid forward mode: {forward_mode=}") mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + # Captured Mamba kernels read state_indices_list as PHYSICAL ids; translate + # before copying (no-op for non-unified pool). + mamba_indices = self._translate_mamba_indices(mamba_indices) self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) - # GDN ReplaySSM (slice 1b): point at the STATIC per-bs write-cursor - # buffer (no advance, no snapshot — capture records the pointer; its - # zeros are overwritten in-place by _replay_metadata before each - # replay). None when the flag is off. Same per-bs tensor object that - # _replay_metadata refreshes, so the captured pointer stays valid. + # Capture records the pointer to the static per-bs buffers; their zeros are + # overwritten in-place by _replay_metadata before each replay. None when off. replayssm_write_pos = ( self.replayssm_write_pos_list[bs - 1] if self.replayssm_write_pos_list is not None else None ) - # GDN ReplaySSM (slice 2b): point at the STATIC per-bs force-flush - # buffer (same capture-by-pointer contract as write_pos; refreshed - # in-place by _replay_metadata before each replay). None when off. replayssm_force_flush = ( self.replayssm_force_flush_list[bs - 1] if self.replayssm_force_flush_list is not None else None ) - # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask if forward_mode.is_target_verify() and self.topk > 1: - # They are None during cuda graph capture so skip the copy_... - # self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrieve_next_token) - # self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrieve_next_sibling) + # retrieve_* are None during capture, so skip the copy. return ForwardMetadata( query_start_loc=self.query_start_loc_list[bs - 1], mamba_cache_indices=self.state_indices_list[bs - 1], @@ -608,6 +525,7 @@ class MambaAttnBackendBase(AttentionBackend): seq_lens_cpu: Optional[torch.Tensor], num_padding: Optional[int] = None, in_capture: bool = False, + mamba_track_indices: Optional[torch.Tensor] = None, ): if num_padding is None: if seq_lens_cpu is None: @@ -619,16 +537,22 @@ class MambaAttnBackendBase(AttentionBackend): # Make sure forward metadata is correctly handled for padding reqs req_pool_indices[bs - num_padding :] = 0 mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + # Translate using the LIVE v2p table BEFORE the padding sentinel below; + # captured Mamba kernels read state_indices_list as PHYSICAL ids. + mamba_indices = self._translate_mamba_indices(mamba_indices) mamba_indices[bs - num_padding :] = -1 self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) - # GDN ReplaySSM (slice 1b): refresh the STATIC per-row write cursor the - # kernel reads, mirroring the eager snapshot-then-advance in - # _forward_metadata but writing in-place into the captured per-bs buffer - # so the graph's recorded pointer stays valid across replays. Done once - # per forward here (out_graph host op), not per layer. Skipped during - # capture (in_capture): capture runs on dummy slots, so advancing the - # persistent counter then would corrupt real per-slot ring positions; - # the captured buffer's contents are irrelevant at capture time anyway. + # Refresh the static track-dest buffer in-place (translated); the captured + # track-save reads it, leaving the handed-in InputBuffer slot read-only. + track_buf = None + if mamba_track_indices is not None: + track_buf = self.mamba_track_indices_buf + track_buf[: len(mamba_track_indices)].copy_( + self._translate_mamba_indices(mamba_track_indices) + ) + # Refresh the static write cursor in-place (mirrors the eager + # snapshot-then-advance). Skip the advance during capture: dummy slots + # would corrupt real ring positions. replayssm_write_pos = None replayssm_force_flush = None if self.replayssm_write_pos_list is not None: @@ -636,29 +560,20 @@ class MambaAttnBackendBase(AttentionBackend): write_pos_buf = mamba_pool.replayssm_write_pos static_wp = self.replayssm_write_pos_list[bs - 1] static_ff = self.replayssm_force_flush_list[bs - 1] - # Hand the full captured per-bs buffers to the kernel, mirroring how - # mamba_cache_indices = self.state_indices_list[bs - 1] is the full - # (bs,) tensor; the kernel indexes them per decode row. + # Hand the full captured per-bs buffers to the kernel; it indexes per row. replayssm_write_pos = static_wp replayssm_force_flush = static_ff if write_pos_buf is not None: - # mamba_indices: this replay's per-row physical slots (padded - # rows == -1, same tensor fed to state_indices_list above). + # this replay's per-row physical slots (padded rows == -1) slots = mamba_indices.to(torch.long) safe_slots = slots.clamp(min=0) - # Snapshot THIS step's per-slot cursor into the captured buffer - # the kernel reads (in-place copy_, never reassign the object). + # Snapshot this step's cursor into the captured buffer in-place + # (copy_, never reassign the object). static_wp[: len(mamba_indices)].copy_(write_pos_buf[safe_slots]) - # GDN ReplaySSM (slice 2b): refresh the captured force-flush - # buffer in-place from THIS step's seq_lens. THE alignment: same - # `seq_lens_cpu % mamba_track_interval == 0` the radix track uses - # (see _replayssm_track_flush_mask / schedule_batch). During - # capture (seq_lens_cpu is None) leave it zeroed: capture content - # is irrelevant and decode replays overwrite it below. + # Refresh the force-flush buffer in-place from this step's seq_lens + # (same condition as the radix track). Zeroed during capture. force_flush_dev = None - # KDA: no radix coordination -> leave static_ff zeroed and - # force_flush_dev None so the advance below is a pure wrap, - # matching the kernel (a zeroed force_flush flushes nothing). + # KDA: no radix coordination -> leave zeroed so the advance is a pure wrap. is_kda = getattr(mamba_pool, "replayssm_is_kda", False) if ( not is_kda @@ -672,10 +587,8 @@ class MambaAttnBackendBase(AttentionBackend): static_ff.zero_() if not in_capture: L = mamba_pool.linear_replayssm_cache_len - # Advance only VALID (non-padded) slots. A forced flush - # empties the ring -> next write_pos is 0 (same as the - # natural wrap at write_pos == L-1). Use this step's snapshot - # cursor (write_pos_buf[safe_slots]) + the flush flag. + # Advance only valid (non-padded) slots; a forced flush empties + # the ring -> next write_pos 0, like the natural L-1 wrap. valid_mask = slots >= 0 valid_slots = slots[valid_mask] if valid_slots.numel() > 0: @@ -688,8 +601,7 @@ class MambaAttnBackendBase(AttentionBackend): torch.zeros_like(cur_pos), (cur_pos + 1) % L, ) - # Dedup; rows sharing a slot share write_pos+flush, so - # the scattered value is identical for either row. + # Dedup; rows sharing a slot share write_pos+flush. uniq_slots, inv = torch.unique(valid_slots, return_inverse=True) next_for_valid = next_pos[valid_mask] new_vals = torch.empty( @@ -726,7 +638,6 @@ class MambaAttnBackendBase(AttentionBackend): else: raise ValueError(f"Invalid forward mode: {forward_mode=}") - # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask if forward_mode.is_target_verify() and self.topk > 1: if ( spec_info is not None @@ -742,6 +653,7 @@ class MambaAttnBackendBase(AttentionBackend): return ForwardMetadata( query_start_loc=self.query_start_loc_list[bs - 1], mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_track_indices=track_buf, retrieve_next_token=self.retrieve_next_token_list[bs - 1], retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1], retrieve_parent_token=self.retrieve_parent_token_list[bs - 1], @@ -752,6 +664,7 @@ class MambaAttnBackendBase(AttentionBackend): return ForwardMetadata( query_start_loc=self.query_start_loc_list[bs - 1], mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_track_indices=track_buf, replayssm_write_pos=replayssm_write_pos, replayssm_force_flush=replayssm_force_flush, ) @@ -769,27 +682,18 @@ class MambaAttnBackendBase(AttentionBackend): ssm_states: torch.Tensor, cache_indices: torch.Tensor, ): - """ - Track and copy Mamba conv/SSM states during decode for prefix caching. - - During decode, each token update modifies conv_states and ssm_states in-place - at positions indexed by cache_indices (the working slots). For prefix caching, - we need to copy these updated states to persistent cache slots (mamba_track_indices) - so they can be prefix cached. - - This delegates to `track_mamba_states_if_needed`, which performs: - conv_states[mamba_track_indices[i]] = conv_states[cache_indices[i]] - ssm_states[mamba_track_indices[i]] = ssm_states[cache_indices[i]] - for all requests where mamba_track_mask[i] is True. - """ + """Copy decode conv/SSM states to track slots for prefix caching. Track + dests come from the metadata (under cuda-graph: the static buffer), so the + InputBuffer registry slot is never mutated.""" if forward_batch.mamba_track_mask is not None: track_mamba_states_if_needed( conv_states, ssm_states, cache_indices, forward_batch.mamba_track_mask, - forward_batch.mamba_track_indices, + self.forward_metadata.mamba_track_indices, forward_batch.batch_size, + check_freed_slots=self.enable_unified_memory, ) def _track_mamba_state_extend( @@ -799,18 +703,8 @@ class MambaAttnBackendBase(AttentionBackend): ssm_states: torch.Tensor, forward_metadata: ForwardMetadata, ): - """ - Track and copy SSM states during extend for prefix caching. - - After the chunked prefill kernel runs, we need to save the SSM recurrent - state at the last chunk boundary so it can be reused for prefix caching. - The source of the state depends on whether the sequence length is aligned - to the chunk size. See `_init_track_ssm_indices` for more details on how - the source and destination indices are computed. - - Note: Conv state tracking for extend is handled separately via gather operations - using indices computed by `_init_track_conv_indices`. - """ + """Copy extend SSM state at the last chunk boundary to track slots (source + depends on chunk alignment; see `_init_track_ssm_indices`).""" if forward_metadata.has_mamba_track_mask: h = h.squeeze(0) @@ -861,6 +755,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): 0 if in_capture else getattr(forward_batch, "num_padding", None) ), in_capture=in_capture, + mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None), ) spec_info = forward_batch.spec_info draft_token_num = spec_info.draft_token_num if spec_info is not None else 1 @@ -891,6 +786,12 @@ class Mamba2AttnBackend(MambaAttnBackendBase): use_triton_causal_conv: bool = False, ): assert isinstance(self.forward_metadata, Mamba2Metadata) + # Page-major stores state strided; only the stride-aware Triton causal-conv + # reads it (CUDA causal_conv1d garbles it). A model may also force Triton. + use_triton_causal_conv = ( + use_triton_causal_conv + or get_global_server_args().enable_page_major_kv_layout + ) layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id) mixer_out, intermediate_states = mixer.forward( hidden_states=hidden_states, @@ -922,8 +823,9 @@ class Mamba2AttnBackend(MambaAttnBackendBase): layer_cache.temporal, self.forward_metadata.mamba_cache_indices[-num_decodes:], forward_batch.mamba_track_mask[-num_decodes:], - forward_batch.mamba_track_indices[-num_decodes:], + self.forward_metadata.mamba_track_indices[-num_decodes:], num_decodes, + check_freed_slots=self.enable_unified_memory, ) return mixer_out @@ -952,7 +854,6 @@ class HybridLinearAttnBackend(AttentionBackend): self.full_attn_backend = full_attn_backend self.linear_attn_backend = linear_attn_backend self.attn_backend_list = [full_attn_backend, linear_attn_backend] - # Dispatcher aliases the full-attn backend's pool refs. self.token_to_kv_pool = full_attn_backend.token_to_kv_pool self.req_to_token_pool = full_attn_backend.req_to_token_pool self.max_context_len = getattr(full_attn_backend, "max_context_len", None) @@ -981,8 +882,8 @@ class HybridLinearAttnBackend(AttentionBackend): def init_forward_metadata(self, forward_batch: ForwardBatch): if forward_batch.forward_mode.is_draft_extend_v2(): - # DRAFT_EXTEND_V2 only runs full-attn layers in the draft model, - # so skip linear/mamba backend metadata which requires query_start_loc. + # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; skip + # linear/mamba metadata (it requires query_start_loc). self.full_attn_backend.init_forward_metadata(forward_batch) return for attn_backend in self.attn_backend_list: @@ -991,9 +892,8 @@ class HybridLinearAttnBackend(AttentionBackend): def init_mha_chunk_metadata( self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False ): - # Hybrid MLA models (Ring/Ling, Kimi-Linear) resolve this via - # get_attn_backend(), which returns this wrapper; delegate to the - # full-attn backend so its chunked/one-shot prefill metadata is planned. + # Hybrid MLA models resolve get_attn_backend() to this wrapper; delegate + # so the full-attn backend plans its chunked-prefill metadata. init = getattr(self.full_attn_backend, "init_mha_chunk_metadata", None) if init is not None: init(forward_batch, disable_flashinfer_ragged) @@ -1050,7 +950,6 @@ class HybridLinearAttnBackend(AttentionBackend): return self.full_attn_backend.forward_decode( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) - # Linear attention backend return self.linear_attn_backend.forward_decode( q=q, k=k, @@ -1081,7 +980,6 @@ class HybridLinearAttnBackend(AttentionBackend): return self.full_attn_backend.forward_extend( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) - # Linear attention backend return self.linear_attn_backend.forward_extend( q=q, k=k, @@ -1150,15 +1048,7 @@ class HybridLinearAttnBackend(AttentionBackend): mamba_steps_to_track: Optional[torch.Tensor], model, ): - """ - Update mamba states after MTP verify using fully fused Triton kernel. - - This replaces the original advanced indexing operations with a single fused - gather-scatter kernel that also handles masking internally, avoiding: - - index_elementwise_kernel from tensor[bool_mask] - - index_select kernel launches - - nonzero kernel launches - """ + """Update mamba states after MTP verify via a fused gather-scatter kernel.""" request_number = last_correct_step_indices.shape[0] state_indices_tensor = ( @@ -1176,16 +1066,14 @@ class HybridLinearAttnBackend(AttentionBackend): intermediate_state_cache = mamba_caches.intermediate_ssm intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0] - # Use fully fused kernel that handles masking internally - # This avoids separate nonzero() and index_select() calls fused_mamba_state_scatter_with_mask( ssm_states, intermediate_state_cache, state_indices_tensor, last_correct_step_indices, ) - # conv intermediate uses the deduplicated sliding-window (overlapping) - # layout, so it needs the strided-read scatter variant. + # conv intermediate uses the deduplicated sliding-window layout, so it + # needs the strided-read scatter variant. fused_conv_window_scatter_with_mask( conv_states, intermediate_conv_window_cache, @@ -1193,10 +1081,9 @@ class HybridLinearAttnBackend(AttentionBackend): last_correct_step_indices, ) - # Track indices used for tracking mamba states for prefix cache + # Track indices for prefix cache if mamba_track_indices is not None: assert mamba_steps_to_track is not None - # Use fully fused kernel for track scatter operations fused_mamba_state_scatter_with_mask( ssm_states, intermediate_state_cache, diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 77cd81ce2..d6dda49a6 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -30,6 +30,11 @@ class ForwardMetadata: query_start_loc: torch.Tensor mamba_cache_indices: torch.Tensor mamba_cache_indices_gdn: Optional[torch.Tensor] = None + # Mamba track DESTINATION slots (PHYSICAL, length == batch). Like + # mamba_cache_indices: a backend-owned static buffer under cuda-graph (translated + # in-place each replay), eager sets the translated decode tensor. The decode + # track-save reads THIS, never forward_batch.mamba_track_indices. + mamba_track_indices: Optional[torch.Tensor] = None # GDN ReplaySSM (slice 1a): per-decode-row snapshot of the ring write # cursor for THIS decode step (gathered from the persistent per-slot # buffer, then advanced once for the next step). int32, length == batch. @@ -179,6 +184,7 @@ class Mamba2Metadata(ForwardMetadata): return Mamba2Metadata( query_start_loc=forward_metadata.query_start_loc, mamba_cache_indices=forward_metadata.mamba_cache_indices, + mamba_track_indices=forward_metadata.mamba_track_indices, retrieve_next_token=forward_metadata.retrieve_next_token, retrieve_next_sibling=forward_metadata.retrieve_next_sibling, retrieve_parent_token=forward_metadata.retrieve_parent_token, @@ -274,6 +280,7 @@ class Mamba2Metadata(ForwardMetadata): return Mamba2Metadata( query_start_loc=query_start_loc, mamba_cache_indices=forward_metadata.mamba_cache_indices, + mamba_track_indices=forward_metadata.mamba_track_indices, retrieve_next_token=forward_metadata.retrieve_next_token, retrieve_next_sibling=forward_metadata.retrieve_next_sibling, retrieve_parent_token=forward_metadata.retrieve_parent_token, diff --git a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py index cc2e18313..506459459 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py @@ -23,6 +23,7 @@ def track_mamba_state_if_needed_kernel( conv_state_numel_per_row: tl.constexpr, # total elements per row ssm_state_numel_per_row: tl.constexpr, # total elements per row BLOCK_SIZE: tl.constexpr, + check_freed_slots: tl.constexpr, # only the shared/unified KV pool emits -1 ): """ Track conv_states and ssm_states rows based on track mask. @@ -51,6 +52,12 @@ def track_mamba_state_if_needed_kernel( src_idx = tl.load(cache_indices_ptr + batch_idx).to(tl.int64) dst_idx = tl.load(mamba_track_indices_ptr + batch_idx).to(tl.int64) + # Skip freed slots (-1): `state_ptr + (-1)*stride` would fault. Only the unified + # pool emits -1 tombstones (from the v2p translate); compiled out for static. + if check_freed_slots: + if src_idx < 0 or dst_idx < 0: + return + # Copy conv_states # Each thread handles BLOCK_SIZE elements for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE): @@ -82,6 +89,7 @@ def track_mamba_states_if_needed( mamba_track_mask: torch.Tensor, mamba_track_indices: torch.Tensor, batch_size: int, + check_freed_slots: bool = False, ): """ Track mamba states using Triton kernel for better performance. @@ -113,6 +121,7 @@ def track_mamba_states_if_needed( conv_state_numel_per_row, ssm_state_numel_per_row, BLOCK_SIZE, + check_freed_slots, ) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index ed8579aeb..4e2ff1cb4 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -98,6 +98,9 @@ class ForwardMetadata: swa_attn_logits: Optional[torch.Tensor] = None # full->SWA translated out_cache_loc (SWA KV-store write target) swa_out_cache_loc: Optional[torch.Tensor] = None + # PHYSICAL full-attn write target for the unified pool (eager: translated tensor; + # cuda-graph: capture-stable buffer view). None for non-unified pools. + out_cache_loc_full_physical: Optional[torch.Tensor] = None class TritonAttnBackend(AttentionBackend): @@ -132,39 +135,31 @@ class TritonAttnBackend(AttentionBackend): extend_attention_fwd_unified ) self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices) - # Split-KV EAGLE-verify kernel (ROCm/Triton). Registered here; enabled - # below once topk is known (the path is only valid at topk == 1). + # Split-KV EAGLE-verify kernel; enabled below once topk is known (valid only at topk == 1). self.verify_splitkv_fwd = torch.compiler.disable(verify_splitkv_fwd) # Parse args self.skip_prefill = skip_prefill max_bs = model_runner.req_to_token_pool.size self.sliding_window_size = model_runner.sliding_window_size - # Pool refs — captured at construction so they survive deletion of the - # corresponding ForwardBatch fields. self.req_to_token_pool = model_runner.req_to_token_pool self.token_to_kv_pool = model_runner.token_to_kv_pool self.req_to_token = model_runner.req_to_token_pool.req_to_token self.token_to_kv_pool_allocator = model_runner.token_to_kv_pool_allocator self.use_sliding_window_kv_pool = isinstance(self.token_to_kv_pool, SWAKVPool) - # Pass-through to the Triton attention wrappers so they can extract the - # KV view strides and specialize on the PAGE_SIZE constexpr. At - # page_size=1 the kernel path matches the slot-based envelope addresses. - # `model_runner.page_size` defaults to 1 when `server_args.page_size` is - # None, avoiding the Optional case here. + # Lets the Triton wrappers specialize on PAGE_SIZE; page_size=1 is + # byte-identical to the slot-based envelope. self.page_size = getattr(model_runner, "page_size", 1) or 1 + # Unified pool v2p hook (None = no-op): req_to_token holds VIRTUAL ids but + # kernels need PHYSICAL. Applied eagerly so the captured graph has no translate. + self._translate_kv_loc = getattr( + self.token_to_kv_pool_allocator, "translate_kv_loc", None + ) self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens self.speculative_num_steps = model_runner.server_args.speculative_num_steps self.topk = model_runner.server_args.speculative_eagle_topk or 0 - # Split-KV verify matches extend_attention_fwd only when the EAGLE tree - # reduces to a pure-causal chain, i.e. topk == 1 (the same condition the - # aiter backend's unified-verify uses). For topk > 1 the tree custom_mask - # is not causal, so leave the path off and fall back to the baseline. - # gfx95-only (MI350X/CDNA4): the kernel uses ROCm/CDNA Triton launch hints - # (waves_per_eu, matrix_instr_nonkdim) and its block config is tuned and - # validated only on gfx950. NVIDIA's Triton rejects those kwargs, and the - # path is unvalidated on NV and on other AMD archs, so restrict it to gfx95 - # and fall back to extend_attention_fwd everywhere else. + # Split-KV verify is bit-equivalent only for a pure-causal chain (topk==1) + # and is gfx95-only; else fall back to extend_attention_fwd. self.use_verify_splitkv = ( is_gfx95_supported() and envs.SGLANG_ENABLE_SPLITKV_VERIFY.get() @@ -179,12 +174,9 @@ class TritonAttnBackend(AttentionBackend): self.num_kv_head = model_runner.model_config.get_num_kv_heads( get_parallel().attn_tp_size ) - # The decode triton kernel derives attn_lse offsets from attn_logits - # strides via integer division by v_head_dim (the "// Lv" trick in - # _fwd_kernel_stage1/stage2), so attn_logits.shape[-1] must exactly - # match the layer's v_head_dim. For hybrid SWA models where SWA and - # full-attention layers use different v_head_dim (e.g. Gemma 4: - # swa=256, full=512), we allocate a second buffer for SWA layers. + # The decode kernel's "// Lv" stride trick requires attn_logits.shape[-1] + # to exactly match the layer's v_head_dim, so hybrid SWA models with + # differing SWA/full v_head_dim need a second buffer for SWA layers. full_v_head_dim = model_runner.model_config.v_head_dim swa_v_head_dim = model_runner.model_config.swa_v_head_dim if self.sliding_window_size is not None and swa_v_head_dim != full_v_head_dim: @@ -217,13 +209,9 @@ class TritonAttnBackend(AttentionBackend): self.max_context_len, ) if _is_gfx942: - # gfx942 (MI300X / MI325X) has 304 CUs, so #20479's next_power_of_2(sm_count) - # rounds up to 512 — twice MI355X's natural cap of 256 — and the persistent - # cuda_graph_attn_logits fp32 buffer hits ~4 GiB on Kimi-K2.6 (v_head_dim=512), - # faulting in ROCm CUDA graph replay - # (https://github.com/sgl-project/sglang/actions/runs/25513282022/job/74877480809). - # Pin the cap at 256 so gfx942 matches the gfx950 (MI355X) behavior that we - # already validated end-to-end. + # gfx942's 304 CUs round up to 512 splits, doubling the persistent + # fp32 attn_logits buffer to ~4 GiB on Kimi-K2.6 and faulting in + # ROCm graph replay; pin to 256 to match validated gfx950 behavior. self.max_kv_splits = min(self.max_kv_splits, 256) if _is_cuda: self.use_pdl = is_arch_support_pdl() @@ -235,18 +223,15 @@ class TritonAttnBackend(AttentionBackend): and model_runner.server_args.chunked_prefill_size == -1 ) - # Decide whether enable deterministic inference with batch-invariant operations self.enable_deterministic = ( model_runner.server_args.enable_deterministic_inference ) - # Configure deterministic inference settings if self.enable_deterministic: - # Use fixed split tile size for batch invariance + # Fixed split tile size for batch invariance self.split_tile_size = get_int_env_var( "SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE", 256 ) - # Set static_kv_splits to False to use deterministic logic instead self.static_kv_splits = False else: self.split_tile_size = ( @@ -258,14 +243,12 @@ class TritonAttnBackend(AttentionBackend): self.max_context_len + self.split_tile_size - 1 ) // self.split_tile_size - # Check arguments assert not ( model_runner.sliding_window_size is not None and model_runner.model_config.is_encoder_decoder ), "Sliding window and cross attention are not supported together" - # Initialize buffers - # TODO(Jianan Ji): Make sure it behaves as expected when kv_indptr_buf is provided and sliding window is enabled + # TODO(Jianan Ji): verify behavior when kv_indptr_buf is provided and sliding window is enabled if kv_indptr_buf is None: self.kv_indptr = torch.zeros( (max_bs + 1,), dtype=torch.int32, device=model_runner.device @@ -273,8 +256,7 @@ class TritonAttnBackend(AttentionBackend): else: self.kv_indptr = kv_indptr_buf - # If sliding window is enabled, we might need two sets of buffers - # because of interleaved attention types (e.g. for Gemma3) + # Sliding window may need a second buffer for interleaved attention types self.window_kv_indptr = None if self.sliding_window_size is not None and self.sliding_window_size > 0: if kv_indptr_buf is None: @@ -282,7 +264,6 @@ class TritonAttnBackend(AttentionBackend): (max_bs + 1,), dtype=torch.int32, device=model_runner.device ) else: - # When provided a buffer, create a clone for the second buffer self.window_kv_indptr = torch.zeros_like(kv_indptr_buf) if not self.skip_prefill: @@ -294,7 +275,6 @@ class TritonAttnBackend(AttentionBackend): (max_bs + 1,), dtype=torch.int64, device=model_runner.device ) - # Initialize forward metadata self.forward_metadata: ForwardMetadata = None self.cuda_graph_custom_mask = None @@ -314,16 +294,13 @@ class TritonAttnBackend(AttentionBackend): num_group * num_seq == num_token ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" - # Legacy dynamic splitting logic (non-deterministic) if ( self.static_kv_splits or self.device_core_count <= 0 ) and not self.enable_deterministic: num_kv_splits.fill_(self.max_kv_splits) return - # deterministic if self.split_tile_size is not None and self.enable_deterministic: - # expand seq_lens to match num_token if num_group > 1: expanded_seq_lens = seq_lens.repeat_interleave(num_group) else: @@ -362,9 +339,8 @@ class TritonAttnBackend(AttentionBackend): kv_indices: Optional[torch.Tensor] = None, kv_start_idx: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # Build per-DCP-rank sharded KV indptr/indices. eager passes - # kv_indices=None (allocate a fresh tensor); the cuda-graph path passes - # a fixed address-stable buffer to fill in place. + # Build per-DCP-rank sharded KV indptr/indices. eager passes kv_indices=None + # (fresh tensor); cuda-graph passes an address-stable buffer to fill in place. dcp_lens = self._dcp_lens(lens, kv_start_idx) kv_indptr[1 : len(req_pool_indices) + 1] = torch.cumsum(dcp_lens, dim=0) kv_indptr = kv_indptr[: len(req_pool_indices) + 1] @@ -422,10 +398,8 @@ class TritonAttnBackend(AttentionBackend): seq_lens = seq_lens[:bs] req_pool_indices = req_pool_indices[:bs] if self.dcp_size > 1: - # DCP: kv_indptr cumsum and kv_indices are per-rank sharded. Write - # them into the same cuda-graph buffers that - # _build_cuda_graph_forward_metadata reads back - # (self.kv_indptr / self.cuda_graph_kv_indices). + # DCP: per-rank sharded; write into the same cuda-graph buffers + # _build_cuda_graph_forward_metadata reads back. _, _, dcp_seq_lens = self._dcp_kv_indices( req_pool_indices, seq_lens, @@ -439,10 +413,15 @@ class TritonAttnBackend(AttentionBackend): kv_indptr = self._fill_kv_indptr_and_indices( bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices ) + # Unified pool: VIRTUAL ids written here are translated to PHYSICAL in + # init_forward_metadata_out_graph (replay-prep) so the captured graph + # carries zero translate nodes. num_kv_splits_lens = seq_lens window_kv_indptr = self.window_kv_indptr window_kv_lens = None if self.sliding_window_size is not None and self.sliding_window_size > 0: + # Unified pool: leave the window VIRTUAL too (translated alongside the + # full kv_indices later); baseline SWA keeps the eager window translate. window_kv_indptr, _, window_kv_lens, _ = update_sliding_window_buffer( self.window_kv_indptr, self.req_to_token, @@ -452,6 +431,7 @@ class TritonAttnBackend(AttentionBackend): bs, token_to_kv_pool=self.token_to_kv_pool, window_kv_indices=self.cuda_graph_window_kv_indices, + skip_full_to_swa_translation=(self._translate_kv_loc is not None), ) return kv_indptr, window_kv_indptr, window_kv_lens, num_kv_splits_lens @@ -462,12 +442,7 @@ class TritonAttnBackend(AttentionBackend): req_pool_indices: torch.Tensor, spec_info, ): - """Fill all cuda-graph buffers for target_verify mode. - - Returns the ForwardMetadata components: - ``(qo_indptr, kv_indptr, custom_mask, mask_indptr, - window_kv_indptr, window_kv_indices, window_num_kv_splits, window_kv_offsets)`` - """ + """Fill all cuda-graph buffers for target_verify mode.""" qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[: bs + 1] = torch.arange( 0, @@ -529,13 +504,10 @@ class TritonAttnBackend(AttentionBackend): forward_mode: ForwardMode, spec_info: Optional[SpecInput], ): - """Fill QO + KV cuda-graph buffers for draft_extend mode. - - Returns (qo_indptr, kv_indptr, num_tokens_per_bs). - """ + """Fill QO + KV cuda-graph buffers for draft_extend mode.""" seq_lens = seq_lens[:bs] - # V2 draft-extend fills num_draft_tokens per req (the cuda-graph runner's - # token layout); num_steps+1 only equals that when topk == 1. + # V2 draft-extend fills num_draft_tokens per req; num_steps+1 only equals + # that when topk == 1. num_tokens_per_bs = ( self.num_draft_tokens if forward_mode.is_draft_extend_v2() @@ -549,17 +521,9 @@ class TritonAttnBackend(AttentionBackend): dtype=torch.int32, device=self.device, ) - # DRAFT_EXTEND_V2: seq_lens = prefix + extend (bumped on the draft-extend path). - # Triton extend kernel receives extend K/V as separate tensors, so - # kv_indptr/kv_indices must cover only the prefix portion. - # extend_seq_lens_tensor is only attached to spec_info at real - # replay (eagle_draft_extend_cuda_graph_runner.replay); during the - # capture-time warmup it's absent, so fall back to zeros (matches - # the pre-unification capture path in #26651). Clamp at 0 because - # padded rows (raw_bs..bs) leave seq_lens at the fill value (1) - # while extend_seq_lens stays at num_tokens_per_bs, which would - # otherwise produce negative kv_lens; padded rows reference - # reserved req-pool slot 0 and their output is discarded. + # DRAFT_EXTEND_V2: kv_indptr/kv_indices cover only the prefix (extend K/V go + # separately). Capture warmup lacks extend_seq_lens_tensor -> fall back to + # zeros; clamp at 0 so padded rows (seq_lens==fill 1) don't go negative. if ( spec_info is not None and getattr(spec_info, "extend_seq_lens_tensor", None) is not None @@ -586,8 +550,8 @@ class TritonAttnBackend(AttentionBackend): if in_capture: assert forward_batch.encoder_lens is None, "Not supported" - # Multi-step speculative decode: kv buffers come from spec_info - # rather than the cuda-graph pool, so replay is not involved. + # Multi-step spec decode: kv buffers come from spec_info, not the + # cuda-graph pool, so replay is not involved. if forward_mode.is_decode_or_idle() and spec_info is not None: self.forward_metadata = ForwardMetadata( attn_logits=self.cuda_graph_attn_logits, @@ -614,9 +578,16 @@ class TritonAttnBackend(AttentionBackend): forward_mode=forward_mode, spec_info=spec_info, ) + out_cache_loc_full_physical = self._translate_cuda_graph_shared_pool_locs( + forward_batch, bs + ) swa_out_cache_loc = self._fill_cuda_graph_swa_out_cache_loc(forward_batch) self.forward_metadata = self._build_cuda_graph_forward_metadata( - bs, forward_mode, spec_info, swa_out_cache_loc + bs, + forward_mode, + spec_info, + swa_out_cache_loc, + out_cache_loc_full_physical, ) else: self._apply_cuda_graph_metadata( @@ -626,15 +597,16 @@ class TritonAttnBackend(AttentionBackend): forward_mode=forward_mode, spec_info=spec_info, ) - # Metadata view is reused from capture; just refill the buffer. + # Metadata view is reused from capture; just refill the buffers. + self._translate_cuda_graph_shared_pool_locs(forward_batch, bs) self._fill_cuda_graph_swa_out_cache_loc(forward_batch) def _fill_cuda_graph_swa_out_cache_loc( self, forward_batch: ForwardBatch ) -> Optional[torch.Tensor]: - """Refill the SWA write-target buffer from the live out_cache_loc and - return the [:n] view (None for non-SWA / multi-step draft), so the - captured store reads fresh slots on replay.""" + """Refill the SWA write-target buffer from live out_cache_loc, returning the + [:n] view (None for non-SWA / multi-step draft) so the captured store reads + fresh slots on replay.""" if not self.use_sliding_window_kv_pool: return None out_cache_loc = forward_batch.out_cache_loc @@ -650,6 +622,47 @@ class TritonAttnBackend(AttentionBackend): ) return self.cuda_graph_swa_out_cache_loc[:n] + def _translate_cuda_graph_shared_pool_locs( + self, forward_batch: ForwardBatch, bs: int + ) -> Optional[torch.Tensor]: + """Unified pool: eager v2p translate of the cuda-graph read+write LOC buffers, + run BEFORE graph.replay() reading the live post-compaction v2p, so the + captured graph carries zero translate nodes. No-op for non-unified pools. + + Read buffers (full kv_indices, SWA window) are translated IN PLACE; the + full-attn WRITE loc is RETURNED as the [:n] view of the backend-owned + out_cache_loc_full_physical buffer. Eager .item() bounds are fine here + (out-of-graph), so no in-graph translate variant is needed. + """ + if self._translate_kv_loc is None: + return None + # Full-attention read path. + n_kv = int(self.kv_indptr[bs].item()) + if n_kv > 0: + self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc( + self.cuda_graph_kv_indices[:n_kv] + ) + # SWA window read path (hybrid-SWA unified pools only). + if self.sliding_window_size is not None and self.sliding_window_size > 0: + n_win = int(self.window_kv_indptr[bs].item()) + if n_win > 0: + self.cuda_graph_window_kv_indices[:n_win] = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + self.cuda_graph_window_kv_indices[:n_win] + ) + ) + # Full-attention write path: translate out_cache_loc -> physical into the + # capture-stable buffer and RETURN the [:n] view. + out_cache_loc = forward_batch.out_cache_loc + n = out_cache_loc.shape[0] + # Zero the padded tail first: a smaller replay batch leaves [n:] holding + # stale ids that the captured store would write; send them to slot 0 (sink). + self.cuda_graph_out_cache_loc_full_physical[n:].zero_() + self.cuda_graph_out_cache_loc_full_physical[:n].copy_( + self._translate_kv_loc(out_cache_loc) + ) + return self.cuda_graph_out_cache_loc_full_physical[:n] + def init_forward_metadata(self, forward_batch: ForwardBatch): """Init auxiliary variables for triton attention backend.""" @@ -663,20 +676,17 @@ class TritonAttnBackend(AttentionBackend): if forward_batch.forward_mode.is_decode_or_idle(): if spec_info is None or spec_info.kv_indptr is None: - # kv_indptr is None for draft-extend's idle batch (no tree - # indices); build plain metadata from seq_lens. + # kv_indptr is None for draft-extend's idle batch; build from seq_lens. if self.dcp_size > 1: - # DCP: per-rank sharded KV indices (shares _dcp_kv_indices - # with the cuda-graph path). Building full contiguous - # indices here would make each rank read the whole KV - # instead of its owner shard. + # DCP: per-rank sharded KV indices, else each rank reads the + # whole KV instead of its owner shard. kv_indptr, kv_indices, _ = self._dcp_kv_indices( forward_batch.req_pool_indices, forward_batch.seq_lens, self.kv_indptr, ) else: - # gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write). + # gpu_only: seq_lens_sum may be None; over-allocate is safe (ragged write). seq_lens_sum = forward_batch.seq_lens_sum if seq_lens_sum is None: seq_lens_sum = bs * self.max_context_len @@ -689,7 +699,8 @@ class TritonAttnBackend(AttentionBackend): forward_batch.req_pool_indices, kv_indices, ) - # Sliding window + if self._translate_kv_loc is not None: + kv_indices = self._translate_kv_loc(kv_indices) if ( self.sliding_window_size is not None and self.sliding_window_size > 0 @@ -755,8 +766,7 @@ class TritonAttnBackend(AttentionBackend): dtype=torch.int32, device=self.device, ) - # Different with flashinfer kv_indptr and kv_indices construction. - # gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write). + # gpu_only: seq_lens_sum may be None; over-allocate is safe (ragged write). seq_lens_sum = forward_batch.seq_lens_sum if seq_lens_sum is None: seq_lens_sum = bs * self.max_context_len @@ -771,7 +781,7 @@ class TritonAttnBackend(AttentionBackend): ) if self.sliding_window_size is not None and self.sliding_window_size > 0: - # window_kv_offsets is used to calculate the start position in custom mask + # window_kv_offsets gives the start position in custom mask ( window_kv_indptr, window_kv_indices, @@ -808,8 +818,7 @@ class TritonAttnBackend(AttentionBackend): self.kv_indptr, ) else: - # gpu_only leaves _cpu unset; ub-allocate is safe (ragged write - # from GPU tensor, extra tail unused). + # gpu_only leaves _cpu unset; over-allocate is safe (ragged write). if forward_batch.extend_prefix_lens_cpu is not None: kv_indices_len = sum(forward_batch.extend_prefix_lens_cpu) else: @@ -825,7 +834,8 @@ class TritonAttnBackend(AttentionBackend): forward_batch.req_pool_indices, kv_indices, ) - # Sliding window + if self._translate_kv_loc is not None: + kv_indices = self._translate_kv_loc(kv_indices) if self.sliding_window_size is not None and self.sliding_window_size > 0: ( window_kv_indptr, @@ -850,8 +860,7 @@ class TritonAttnBackend(AttentionBackend): mask_indptr = None attn_logits = None attn_lse = None - # Caller usually supplies extend_seq_lens_cpu (eagle_info gpu_only - # sets host-constant mirror); defensive GPU-max fallback if not. + # Defensive GPU-max fallback when extend_seq_lens_cpu is absent. if forward_batch.extend_seq_lens_cpu is not None: max_extend_len = max(forward_batch.extend_seq_lens_cpu) else: @@ -864,6 +873,17 @@ class TritonAttnBackend(AttentionBackend): forward_batch.out_cache_loc ) + # Unified pool full-attention WRITE loc (virtual out_cache_loc -> physical), + # carried in the metadata (-> KVWriteLoc.full_loc). None for non-unified pools. + out_cache_loc_full_physical = None + if ( + self._translate_kv_loc is not None + and forward_batch.out_cache_loc is not None + ): + out_cache_loc_full_physical = self._translate_kv_loc( + forward_batch.out_cache_loc + ) + self.forward_metadata = ForwardMetadata( attn_logits, attn_lse, @@ -880,6 +900,7 @@ class TritonAttnBackend(AttentionBackend): window_kv_offsets, swa_attn_logits=swa_attn_logits, swa_out_cache_loc=swa_out_cache_loc, + out_cache_loc_full_physical=out_cache_loc_full_physical, ) def init_cuda_graph_state( @@ -970,12 +991,22 @@ class TritonAttnBackend(AttentionBackend): device=self.device, ) + if self._translate_kv_loc is not None: + # Unified pool full-attention write-target buffer, refilled at replay + # (-> KVWriteLoc.full_loc). Capture-stable, mirrors cuda_graph_swa_out_cache_loc. + self.cuda_graph_out_cache_loc_full_physical = torch.zeros( + (max_num_tokens,), + dtype=torch.int64, + device=self.device, + ) + def _build_cuda_graph_forward_metadata( self, bs: int, forward_mode: ForwardMode, spec_info: Optional[SpecInput], swa_out_cache_loc: Optional[torch.Tensor] = None, + out_cache_loc_full_physical: Optional[torch.Tensor] = None, ) -> ForwardMetadata: """Construct ForwardMetadata from the current cuda-graph buffer state. @@ -1006,6 +1037,7 @@ class TritonAttnBackend(AttentionBackend): window_kv_offsets=None, swa_attn_logits=self.cuda_graph_swa_attn_logits, swa_out_cache_loc=swa_out_cache_loc, + out_cache_loc_full_physical=out_cache_loc_full_physical, ) elif forward_mode.is_target_verify(): custom_mask = ( @@ -1031,6 +1063,7 @@ class TritonAttnBackend(AttentionBackend): ), window_kv_offsets=self.cuda_graph_window_kv_offsets if swa else None, swa_out_cache_loc=swa_out_cache_loc, + out_cache_loc_full_physical=out_cache_loc_full_physical, ) elif forward_mode.is_draft_extend_v2(): return ForwardMetadata( @@ -1056,6 +1089,7 @@ class TritonAttnBackend(AttentionBackend): window_num_kv_splits=None, window_kv_offsets=None, swa_out_cache_loc=swa_out_cache_loc, + out_cache_loc_full_physical=out_cache_loc_full_physical, ) else: raise ValueError(f"Invalid forward mode: {forward_mode=} for CUDA Graph.") @@ -1183,6 +1217,7 @@ class TritonAttnBackend(AttentionBackend): loc_info = KVWriteLoc( forward_batch.out_cache_loc, self.forward_metadata.swa_out_cache_loc, + full_loc=self.forward_metadata.out_cache_loc_full_physical, ) if layer.k_scale is None: self._set_kv_buffer(forward_batch, layer, loc_info, k, v) @@ -1628,6 +1663,7 @@ class TritonAttnBackend(AttentionBackend): KVWriteLoc( forward_batch.out_cache_loc, self.forward_metadata.swa_out_cache_loc, + full_loc=self.forward_metadata.out_cache_loc_full_physical, ), k, v, @@ -1908,12 +1944,21 @@ def update_sliding_window_buffer( device=None, token_to_kv_pool=None, window_kv_indices=None, + skip_full_to_swa_translation=False, ): """Fill window KV buffers for sliding-window attention. - Pass window_kv_indices to write into a pre-allocated buffer (CUDA-graph - path); omit it (or pass None) to allocate a fresh tensor (eager path, - requires device). + Pass ``window_kv_indices`` to write into a pre-allocated buffer (CUDA-graph + path); omit it (or pass ``None``) to allocate a fresh tensor (eager path, + requires ``device``). + + ``skip_full_to_swa_translation=True`` leaves ``window_kv_indices`` as VIRTUAL + full-token ids (no eager full->swa translate). The unified-memory-pool cuda-graph + builder passes this so the window translate is deferred to + ``TritonAttnBackend._translate_cuda_graph_shared_pool_locs`` (run in + ``init_forward_metadata_out_graph``, BEFORE ``graph.replay()``), which reads + the live v2p and rewrites the static window buffer to swa-physical in place; + baseline SWA leaves it False (eager). """ window_kv_lens = torch.minimum( seq_lens, @@ -1935,7 +1980,9 @@ def update_sliding_window_buffer( window_kv_indices, req_to_token.stride(0), ) - if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"): + if not skip_full_to_swa_translation and hasattr( + token_to_kv_pool, "translate_loc_from_full_to_swa" + ): kv_last_index = window_kv_indptr[-1] window_kv_indices[:kv_last_index] = ( token_to_kv_pool.translate_loc_from_full_to_swa( diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index edd3ce869..ef31b32c7 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -52,6 +52,9 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, zero_match_result, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.server_args import ServerArgs, get_global_server_args @@ -495,6 +498,34 @@ class PrefillAdder: self.rem_swa_token_offset = 0 + # Unified-pool joint budget: a new mamba state consumes shared-gap bytes + # that `rem_total_tokens` (full KV) otherwise counts as free, so reserve + # the gap per new mamba slot or admission over-commits. Gate on the + # ALLOCATOR being the unified Mamba composite, NOT on `is_hybrid_ssm_cache` + # (False for `ChunkCache`, which would skip the reservation on the + # chunk-cache path): the gap coupling is a property of the byte buffer. + self._mamba_slot_cost = 0 + if isinstance( + self.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator + ): + self._mamba_slot_cost = ( + self.token_to_kv_pool_allocator.mamba_slot_full_token_cost() + ) + + # `mamba_gap_reserve` is charged to `rem_total_tokens`, which INCLUDES + # `full_evictable_size()` — but `alloc_req_slots` can only recover + # MAMBA-recoverable bytes for a mamba slot (shared gap + peer holes + + # mamba-evictable radix), NOT full-evictable. Gate new mamba slots on + # that mamba-recoverable budget separately or an over-admit hits the + # fail-loud `RuntimeError`. `None` outside the unified Mamba pool. + self.rem_mamba_slots = None + if self._mamba_slot_cost: + self.rem_mamba_slots = ( + self.token_to_kv_pool_allocator.mamba_allocator.schedulable_available_size() + ) + if self.is_hybrid_ssm_cache: + self.rem_mamba_slots += self.tree_cache.mamba_evictable_size() + self.priority_scheduling_preemption_threshold = ( priority_scheduling_preemption_threshold ) @@ -602,6 +633,21 @@ class PrefillAdder: budget += self.ceil_paged_tokens(swa_host_hit_length) return budget + def _mamba_gap_budget_for_req(self, req: Req) -> int: + """Shared-gap reservation (full-token-equivalents) for a request's new + mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`) + and only when the req has no state yet (`mamba_pool_idx is None`, mirroring + `HybridReqToTokenPool.alloc`); 0 keeps baseline / SWA / non-Mamba unchanged. + + Conservative by design (`_mamba_slot_cost` rounds UP). Does NOT reserve + radix COW headroom or locked-but-evictable bytes — that residual is + backstopped by the fail-loud RuntimeError in `alloc_req_slots`. FIXME: if + over-admission crashes under pressure, make this more conservative (e.g. + multiply by `MAMBA_STATE_PER_REQ_PREFIX_CACHE`).""" + if self._mamba_slot_cost and req.mamba_pool_idx is None: + return self._mamba_slot_cost + return 0 + def ceil_paged_tokens(self, tokens: int) -> int: return -(-tokens // self.page_size) * self.page_size @@ -609,6 +655,10 @@ class PrefillAdder: no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 if not no_token and self.is_hybrid_swa: no_token = self.rem_swa_tokens <= 0 + # Gate new mamba slots separately: rem_total_tokens' full_evictable can't + # cover a mamba slot, which needs mamba-recoverable bytes (see __init__). + if not no_token and self.rem_mamba_slots is not None: + no_token = self.rem_mamba_slots <= 0 if no_token: return AddReqResult.NO_TOKEN @@ -630,14 +680,27 @@ class PrefillAdder: extend_input_len: int, max_new_tokens: int, retracted_stain: bool, + mamba_gap_reserve: int = 0, ): # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative extend_input_len = self.ceil_paged_tokens(extend_input_len) # alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit page_overhead = self.page_size - self.rem_total_token_offset += extend_input_len + max_new_tokens + page_overhead - self.cur_rem_token_offset += extend_input_len + page_overhead + # `mamba_gap_reserve` (shared Mamba pool only; 0 otherwise) charges the new + # mamba state's shared-gap cost to BOTH full budgets: the slot is allocated + # immediately (counts against `cur_rem`) and held for the request lifetime + # (counts against `rem_total`). See `_mamba_gap_budget_for_req`. + self.rem_total_token_offset += ( + extend_input_len + max_new_tokens + page_overhead + mamba_gap_reserve + ) + self.cur_rem_token_offset += ( + extend_input_len + page_overhead + mamba_gap_reserve + ) + # The new mamba slot also consumes one mamba-recoverable slot (gated + # separately so full_evictable can't cover it — see __init__). + if mamba_gap_reserve and self.rem_mamba_slots is not None: + self.rem_mamba_slots -= 1 self.rem_input_tokens -= extend_input_len if self.is_hybrid_swa: @@ -681,7 +744,13 @@ class PrefillAdder: self.can_run_list.append(req) - self._update_prefill_budget(prefix_len, trunc_len, 0, req.retracted_stain) + self._update_prefill_budget( + prefix_len, + trunc_len, + 0, + req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + ) def _req_inc_lock_ref(self, req: Req): result = self.tree_cache.inc_lock_ref(req.last_node) @@ -711,7 +780,11 @@ class PrefillAdder: else 0 ) self._update_prefill_budget( - 0, req.extend_range.length, max_new_tokens, req.retracted_stain + 0, + req.extend_range.length, + max_new_tokens, + req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), ) # Return based on remaining token availability @@ -755,6 +828,7 @@ class PrefillAdder: else 0 ), req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), ) # Return if chunked prefill not finished @@ -782,6 +856,9 @@ class PrefillAdder: req.prefix_indices ) paged_input = self.ceil_paged_tokens(cand_extend_input_len) + # Shared Mamba pool: fold the new mamba state's shared-gap cost into the + # budget gate so admission can't over-commit (0 for baseline / non-Mamba). + paged_input += self._mamba_gap_budget_for_req(req) if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): return AddReqResult.NO_TOKEN if self.is_hybrid_swa: @@ -863,6 +940,7 @@ class PrefillAdder: req.extend_range.length, min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS), req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), ) else: if self.rem_chunk_tokens <= 0: @@ -877,7 +955,13 @@ class PrefillAdder: ) self.can_run_list.append(req) self.new_chunked_req = req - self._update_prefill_budget(0, trunc_len, 0, req.retracted_stain) + self._update_prefill_budget( + 0, + trunc_len, + 0, + req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), + ) return self.budget_state() @@ -906,11 +990,9 @@ class PrefillAdder: if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True): return self.add_one_req_ignore_eos(req) - # Reserve page_size for page-alignment overhead. The paged allocator - # may consume up to one extra page per request (see alloc_extend), and - # _update_prefill_budget already accounts for this in the deduction. - # Without this, admission is more optimistic than the actual budget - # deduction, allowing over-admission when the pool is nearly full. + # Reserve page_size for page-alignment overhead: the paged allocator may + # consume one extra page per request (see alloc_extend), which + # _update_prefill_budget also deducts. max_new = min( max(req.sampling_params.max_new_tokens - len(req.output_ids), 0), CLIP_MAX_NEW_TOKENS, @@ -919,6 +1001,9 @@ class PrefillAdder: req.prefix_indices ) total_tokens = cand_extend_input_len + max_new + self.page_size + # Shared Mamba pool: fold the new mamba state's shared-gap cost into + # `total_tokens` so both `rem_total_tokens` gates reflect the joint budget. + total_tokens += self._mamba_gap_budget_for_req(req) # adjusting the input_tokens based on host_hit_length and page_size real_input_tokens = cand_extend_input_len - req.host_hit_length @@ -1009,6 +1094,7 @@ class PrefillAdder: CLIP_MAX_NEW_TOKENS, ), req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), ) else: # Make sure at least one page is available @@ -1045,7 +1131,11 @@ class PrefillAdder: self._req_inc_lock_ref(req) self._update_prefill_budget( - prefix_len, trunc_len, 0, req.retracted_stain + prefix_len, + trunc_len, + 0, + req.retracted_stain, + mamba_gap_reserve=self._mamba_gap_budget_for_req(req), ) return self.budget_state() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a2e72df20..df5a49d0d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1496,7 +1496,7 @@ class Scheduler( if self.device == "cpu": self.schedule_stream.synchronize = lambda: None # No-op for CPU # The global WAR barrier fences the scheduler's next shared-buffer write - # on the previous forward's read of the shared pool. + # on the previous forward's read of the unified memory pool. self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get() with self.device_module.StreamContext(self.schedule_stream): dispatch_event_loop(self) @@ -1580,6 +1580,14 @@ class Scheduler( # we can process the last batch immediately. if disable_overlap_for_batch: pop_and_process() + # Opportunistic flush at the disable_overlap sync boundary: + # forward_stream is idle (prev forward drained, next not launched), + # so `_flush`'s non-urgent guard compacts freely. Sync-free, best-effort. + if self.server_args.enable_unified_memory: + try: + self.token_to_kv_pool_allocator.flush_opportunistic() + except Exception: + pass # Launch the current batch if batch: @@ -2527,7 +2535,6 @@ class Scheduler( release_kv_cache(req, self.tree_cache, is_insert=False) self.chunked_req = None - self._chunked_req_scheduled_last_iter = False self._pending_chunked_abort_req = None self.ipc_channels.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req) logger.debug(f"Abort chunked prefill request. {req.rid=}") @@ -3237,6 +3244,20 @@ class Scheduler( self.batch_record_buf[self.batch_record_ct].extend( batch_result.extra_keep_alive_refs ) + if self.server_args.enable_unified_memory: + # Record a `forward_done` event after the forward (before + # copy_to_cpu); lazy-compaction `_flush` gates src reuse on + # it. Only the unified pool's allocator exposes these hooks. + allocator = self.token_to_kv_pool_allocator + forward_done = self.device_module.Event() + forward_done.record(stream=self.forward_stream) + allocator.set_latest_forward_done_event(forward_done) + # Write-set classification: hand the allocator this + # forward's virtual out_cache_loc as a tensor ref (no GPU work). + allocator.set_inflight_forward( + forward_done, + batch.out_cache_loc, + ) # FIXME(lsyin): maybe move this to forward_batch_generation batch_result.copy_done = self.device_module.Event() if batch_result.delay_sample_func is None: @@ -3497,6 +3518,12 @@ class Scheduler( if not self.is_fully_idle(): return + if self.server_args.enable_unified_memory: + try: + self.token_to_kv_pool_allocator.flush_opportunistic() + except Exception: + pass + # memory leak check (skipped for hisparse — pool counters intentionally # diverge during host-backup, see _get_swa_token_info clamp). if not self.enable_hisparse: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 9e625d404..633baf8ca 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -88,8 +88,16 @@ class SchedulerInvariantChecker: protected = self.tree_cache.full_protected_size() session_held = self.pool_stats_observer.session_held_full_tokens() total = self.full_tokens_per_layer - elif self.is_hybrid_ssm and self.tree_cache.supports_mamba(): - protected = self.tree_cache.full_protected_size() + elif self.is_hybrid_ssm: + # Branch on cache type for the protected accessor (MambaRadixCache + # splits full/mamba; ChunkCache only has the single protected_size). + # Use the allocator's `.size` for `total`: static max_total_num_tokens for + # non-unified pools, the dynamic byte-coordinated cap (matching + # `available_size`) for the unified pool. + if self.tree_cache.supports_mamba(): + protected = self.tree_cache.full_protected_size() + else: + protected = self.tree_cache.protected_size() session_held = self.pool_stats_observer.session_held_tokens() total = self.token_to_kv_pool_allocator.size else: diff --git a/python/sglang/srt/mem_cache/allocator/mamba.py b/python/sglang/srt/mem_cache/allocator/mamba.py index b6ddbdf4a..5ffc99611 100644 --- a/python/sglang/srt/mem_cache/allocator/mamba.py +++ b/python/sglang/srt/mem_cache/allocator/mamba.py @@ -48,6 +48,13 @@ class MambaSlotAllocator: def available_size(self) -> int: return len(self.free_slots) + def schedulable_available_size(self) -> int: + """Planner-facing free count. Identity to ``available_size`` for the + static pool (slot-count and byte-coordinated views coincide); the shared + ``UnifiedMambaSlotAllocator`` overrides it with the byte-coordinated view. + Lets ``alloc_req_slots`` call it uniformly without a getattr fallback.""" + return self.available_size() + def alloc_group_begin(self, num_reqs: int): """Pre-allocate a batch of slots for match_prefix to amortize overhead.""" self._alloc_iter = None diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 3dce815f3..ebfec94b7 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -112,6 +112,15 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def swa_available_size(self): return self.swa_attn_allocator.available_size() + # Slot-conservation views for the leak invariant. On the non-shared allocator + # the static budget IS physical (conserve == physical); the shared composite + # overrides these with the static-cap view. + def _conserve_full_available_size(self): + return self.full_available_size() + + def _conserve_swa_available_size(self): + return self.swa_available_size() + @property def size(self): return min(self._size_full, self._size_swa) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e50fcfc14..1ffbe218b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -87,11 +87,8 @@ def free_swa_out_of_window_slots( req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor) # Subtract an extra page_size so the eviction frontier never reaches the - # radix tree insert boundary (page_floor(seq_len)). This keeps at least one - # page of non-evicted SWA KV for the tree to store as a non-tombstone node, - # preserving cache reuse in multi-turn scenarios. Without this, leaf nodes - # may become tombstoned, causing SWA memory leak. - # See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart). + # radix tree insert boundary, keeping >=1 page of non-evicted SWA KV for the + # tree to store as a non-tombstone node (else leaf nodes tombstone -> SWA leak). if drop_page_margin or envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get(): evict_threshold = pre_len - sliding_window_size else: @@ -220,9 +217,7 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int: if server_args.speculative_algorithm is None: return 1 - # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per - # decode step (draft chain and verify block share the reservation). - + # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step. spec_steps = server_args.speculative_num_steps or 1 spec_topk = server_args.speculative_eagle_topk or 1 spec_tokens = server_args.max_speculative_num_draft_tokens @@ -234,10 +229,9 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int: if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv(): return max(spec_steps * spec_topk, spec_tokens) else: - # page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree - # footprint. Per topk branch needs ceil((last_page_len + num_steps) / page) - # pages; the partial tail page can be up to page_size - 1, and each branch - # gets its own (duplicated) copy -- so reserve for all topk branches. + # spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per + # topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch + # duplicated -- reserve for all topk branches. num_new_pages_per_topk = ( (page_size - 1) + spec_steps + page_size - 1 ) // page_size @@ -256,12 +250,10 @@ def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> in def get_req_to_token_extra_context_len(server_args: ServerArgs) -> int: """req_to_token row headroom beyond the model context length. - Sized to hold the decode over-allocation (kv_committed_len + - get_alloc_reserve_per_decode). The spec v2 page>1 topk>1 holey draft footprint - can outgrow the default num_draft_tokens headroom (PR #26972). + Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey + draft footprint can outgrow the default num_draft_tokens headroom. """ - # FIXME(lsyin): this is the temporary fix for the context length issue when - # using speculative decoding + # FIXME(lsyin): temporary fix for the context length issue under spec decoding extra = 4 + (server_args.max_speculative_num_draft_tokens or 0) if ( server_args.speculative_algorithm is not None @@ -327,13 +319,8 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): def _compute_dsv4_state_lens(batch, *, is_decode: bool): - """Per-req c{4,128}_state pool alloc lens (a ``DSV4StateLens``) for this - alloc step. The DSV4-NPU allocator owns the computation (it also mutates the - per-req cumulative state on each ``Req``); we just trigger it here, right - before the paged alloc that consumes the result. - - None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``) - so the ``alloc_paged_token_slots_*`` forwarding stays a no-op. + """Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step. + None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``). """ allocator = batch.token_to_kv_pool_allocator if not hasattr(allocator, "compute_dsv4_state_lens_extend"): @@ -371,8 +358,7 @@ def alloc_paged_token_slots_extend( extra_alloc_kwargs = {} if is_dsv4: extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Pass the per-req tables in per call for the c-pool / state last_loc - # lookup; the allocator holds no reference to the pool. + # Per-call per-req tables for the c-pool / state last_loc lookup. if batch is not None: extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool if dsv4_state_lens is not None: @@ -415,10 +401,20 @@ def alloc_req_slots( reqs: list[Req], tree_cache: BasePrefixCache | None, ) -> list[int]: - """Allocate request slots from the pool.""" + """Allocate request slots from the pool. + + Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An + alloc failure here means the admission budget (``PrefillAdder``) was wrong + and should surface rather than be masked. + """ num_reqs = len(reqs) if isinstance(req_to_token_pool, HybridReqToTokenPool): - mamba_available_size = req_to_token_pool.mamba_allocator.available_size() + # Byte-coordinated for the shared allocator (accounts for the peer full + # sub-pool's bytes); plain slot free count for the non-shared one. + mamba_available_size = ( + req_to_token_pool.mamba_allocator.schedulable_available_size() + ) + # Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk. if tree_cache.supports_mamba(): factor = ( MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY @@ -433,24 +429,19 @@ def alloc_req_slots( mamba_num = max(0, mamba_state_needed - mamba_available_size) tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) req_pool_indices = req_to_token_pool.alloc(reqs) - if req_pool_indices is None: raise RuntimeError( "alloc_req_slots runs out of memory. " "Please set a smaller number for `--max-running-requests`. " - f"{req_to_token_pool.available_size()=}, " - f"{num_reqs=}, " + f"{req_to_token_pool.available_size()=}, {num_reqs=}, " ) return req_pool_indices def _alloc_page_size(batch: ScheduleBatch) -> int: - # DCP (HIP & CUDA only) swaps in a PagedTokenToKVPoolAllocator whose - # page_size is server_args.page_size * dcp_size, so it can be > 1 even when - # tree_cache.page_size (== server_args.page_size) is 1. Only on the HIP DCP - # path do we branch on the real allocator's page_size so the paged path is - # taken; everywhere else tree_cache.page_size is authoritative and the two - # are equal (dcp_size == 1), so behavior is unchanged. + # DCP swaps in an allocator whose page_size is server_args.page_size * + # dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on + # the real allocator's page_size there. Elsewhere the two are equal. if (_is_hip or _is_cuda) and get_global_server_args().dcp_size > 1: return batch.tree_cache.token_to_kv_pool_allocator.page_size return batch.tree_cache.page_size @@ -462,10 +453,9 @@ def alloc_for_extend( """ Allocate KV cache for extend batch and write to req_to_token_pool. - Returns: - out_cache_loc: allocated cache locations - req_pool_indices_device: request pool indices as a device tensor - req_pool_indices_cpu: request pool indices as a CPU tensor (host mirror) + Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)`` + (the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError`` + if the pool can't satisfy the batch (fail-loud — see its docstring). """ # free out-of-window swa tokens batch.maybe_evict_swa() @@ -478,7 +468,7 @@ def alloc_for_extend( prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True) extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True) - # Allocate req slots + # Allocate req slots (raises RuntimeError if the pool is exhausted) req_pool_indices = alloc_req_slots( batch.req_to_token_pool, batch.reqs, batch.tree_cache ) @@ -489,8 +479,6 @@ def alloc_for_extend( if _alloc_page_size(batch) == 1: out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) else: - # Since tree_cache.page_size is (page_size * dcp_world_size), for dcp - # on cuda platform, always use alloc_paged_token_slots_extend # Paged allocation - build last_loc last_loc = [ (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) @@ -524,8 +512,7 @@ def alloc_for_extend( batch.req_to_token_pool, ) - # DSV4-NPU hook: write c4/c128/swa per-req tables from the stashed bundle. - # No-op on non-DSV4 paths (out_cache_loc_dsv4 stays None there). + # DSV4-NPU hook: no-op on non-DSV4 paths. if _is_npu: maybe_write_dsv4_extend( batch, @@ -559,8 +546,7 @@ def alloc_paged_token_slots_decode( extra_alloc_kwargs = {} if is_dsv4: extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Per-call per-req tables for the last_loc lookup; the allocator holds - # no reference to the pool. + # Per-call per-req tables for the last_loc lookup. if batch is not None: extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool if dsv4_state_lens is not None: @@ -633,8 +619,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: (batch.req_pool_indices, locs), out_cache_loc.to(torch.int32) ) - # DSV4-NPU hook: post-decode write of c4/c128/swa per-req tables from the - # stashed bundle. No-op on non-DSV4 paths (out_cache_loc_dsv4 stays None). + # DSV4-NPU hook: no-op on non-DSV4 paths. if _is_npu: maybe_write_dsv4_decode( batch, @@ -698,8 +683,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr req.mamba_pool_idx is not None ), "mamba state is freed while the tree cache does not manage mamba states" tree_cache.req_to_token_pool.free_mamba_cache(req) - # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the - # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here. + # DSV4-NPU's free() also releases c4/c128 state pages; no-op for others. tree_cache.req_to_token_pool.free(req) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 4c9ab82fd..368781d8d 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -46,6 +46,9 @@ from sglang.srt.mem_cache.base_prefix_cache import ( ) from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.utils import split_node_hash_value from sglang.srt.server_args import get_global_server_args @@ -420,9 +423,15 @@ class LRUList: class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): def __init__(self, params: CacheInitParams): - assert isinstance( - params.token_to_kv_pool_allocator, TokenToKVPoolAllocator - ) or isinstance(params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator) + assert ( + isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator) + or isinstance( + params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator + ) + or isinstance( + params.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator + ) + ) self.req_to_token_pool: HybridReqToTokenPool = params.req_to_token_pool self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator self.mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size @@ -694,8 +703,12 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ) else: mamba_value_donated = self._alloc_mamba_slot() + # mamba_pool is a pure PHYSICAL store; translate both slot ids + # virtual->physical (identity for the non-unified memory pool) before the copy. + translate = self.req_to_token_pool.translate_mamba_indices self.req_to_token_pool.mamba_pool.copy_from( - req.mamba_pool_idx.unsqueeze(0), mamba_value_donated + translate(req.mamba_pool_idx.unsqueeze(0)), + translate(mamba_value_donated), ) result = self.insert( diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 2a710e675..e7d522728 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -970,6 +970,14 @@ class HybridReqToTokenPool(ReqToTokenPool): def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor: return self.req_index_to_mamba_index_mapping[req_indices] + def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor: + """Virtual->physical mamba-slot translate. Identity for a static pool + (slots are physical); UnifiedHybridReqToTokenPool overrides it for the + unified memory pool, where mamba slot ids are virtual. Callers translate + before calling the pool's physical-id state ops (copy_from / clear_slots + / get_cpu_copy / load_cpu_copy).""" + return mamba_indices + def mamba2_layer_cache(self, layer_id: int): assert layer_id in self.mamba_map if self.layer_transfer_counter is not None: @@ -1134,29 +1142,43 @@ class HybridReqToTokenPool(ReqToTokenPool): class KVWriteLoc: """Write target(s) for ``KVCache.set_kv_buffer``. - ``loc`` is the full-pool write location; ``swa_loc`` is the pre-translated - full->SWA location for hybrid SWA pools (``None`` otherwise). Bundling them - lets a backend issue one ``set_kv_buffer`` call regardless of pool type. + All location info lives here (in the attention metadata), NOT in the pool: + - ``loc``: the generic per-token write location (the allocated + ``out_cache_loc``). VIRTUAL under the unified memory pool (it indexes the + virtual slot space); already physical for a non-unified memory pool. + - ``swa_loc``: the pre-translated SWA-sub-pool PHYSICAL location for hybrid + SWA pools (``None`` otherwise). + - ``full_loc``: the pre-translated full-attention-sub-pool PHYSICAL location + for the unified memory pool (``None`` otherwise), computed once per forward in + attention metadata (``ForwardMetadata.out_cache_loc_full_physical``). The + shared full pool writes it directly; the pool never translates (replacing + the former per-layer v2p gather / ``set_full_loc`` pin). + + ``swa_loc`` and ``full_loc`` are the parallel pair (each a pre-resolved + PHYSICAL loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``); + ``loc`` is the generic, possibly-virtual fallback. Bundling them lets a + backend issue one ``set_kv_buffer`` call regardless of pool type. """ loc: torch.Tensor swa_loc: Optional[torch.Tensor] = None + full_loc: Optional[torch.Tensor] = None def __post_init__(self): - # swa_out_cache_loc is computed once at metadata-init time from the - # full (possibly padded) out_cache_loc. Piecewise CUDA graphs later - # narrow out_cache_loc to real_num_tokens per layer, so swa_loc can - # be longer than loc. Slice to match since both are in the same - # per-token order. + # swa_loc / full_loc are resolved once at metadata-init from the full + # (padded) out_cache_loc; piecewise/DP-padded paths later narrow loc per + # layer, so slice these pre-resolved locs to match (same per-token order). if self.swa_loc is not None and self.swa_loc.shape[0] != self.loc.shape[0]: self.swa_loc = self.swa_loc[: self.loc.shape[0]] + if self.full_loc is not None and self.full_loc.shape[0] != self.loc.shape[0]: + self.full_loc = self.full_loc[: self.loc.shape[0]] def unwrap_write_loc(loc_info): - """Return ``(loc, swa_loc)`` from a ``KVWriteLoc`` or a bare loc tensor.""" + """Return ``(loc, swa_loc, full_loc)`` from a ``KVWriteLoc`` or a bare loc.""" if isinstance(loc_info, KVWriteLoc): - return loc_info.loc, loc_info.swa_loc - return loc_info, None + return loc_info.loc, loc_info.swa_loc, loc_info.full_loc + return loc_info, None, None class KVCache(abc.ABC): @@ -1648,7 +1670,7 @@ class MHATokenToKVPool(KVCache): layer_id_override: Optional[int] = None, dcp_kv_mask: Optional[torch.Tensor] = None, ): - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) # Catch stale slot ids here instead of as illegal-addr / silent KV # corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT). maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)") @@ -2123,7 +2145,7 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): v_scale: Optional[float] = None, layer_id_override: Optional[int] = None, ): - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA-FP4)") from sglang.srt.model_executor.runner import get_is_capture_mode @@ -2346,6 +2368,9 @@ class HybridLinearKVPool(KVCache): qk_rope_head_dim: int = None, start_layer: Optional[int] = None, full_kv_pool_class: Optional[type] = None, + # When provided (shared-KV-pool path), use this pool for the + # full-attention layers instead of constructing one internally. + full_kv_pool: Optional[KVCache] = None, ): self.size = size self.dtype = dtype @@ -2357,8 +2382,15 @@ class HybridLinearKVPool(KVCache): self.head_num = head_num self.head_dim = head_dim self.mamba_pool = mamba_pool + # virtual->physical mamba-slot translate for the HiCache offload path; + # identity for a static pool, the allocator's `translate` for the unified pool. + self._mamba_translate = lambda ids: ids self.use_mla = use_mla - if not use_mla: + if full_kv_pool is not None: + # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool + # aliasing the shared byte buffer. + self.full_kv_pool = full_kv_pool + elif not use_mla: TokenToKVPoolClass = MHATokenToKVPool if current_platform.is_out_of_tree(): @@ -2492,11 +2524,16 @@ class HybridLinearKVPool(KVCache): v_scale: float = 1.0, dcp_kv_mask: Optional[torch.Tensor] = None, ): + # Write-location info lives in the metadata (`KVWriteLoc`). `full_loc` is the + # unified pool's pre-translated PHYSICAL loc (None for a static pool, where + # `loc` is already physical) — either way the pool writes a PHYSICAL loc. + loc, _, full_loc = unwrap_write_loc(loc) layer_id = self._transfer_full_attention_id(layer.layer_id) if not self.use_mla: + write_loc = full_loc if full_loc is not None else loc self.full_kv_pool.set_kv_buffer( None, - loc, + write_loc, cache_k, cache_v, k_scale, @@ -2518,8 +2555,9 @@ class HybridLinearKVPool(KVCache): def get_cpu_copy(self, indices, mamba_indices=None): kv_cpu = self.full_kv_pool.get_cpu_copy(indices) + # mamba_pool stores PHYSICAL ids; translate the (unified-pool virtual) ids first. mamba_cpu = ( - self.mamba_pool.get_cpu_copy(mamba_indices) + self.mamba_pool.get_cpu_copy(self._mamba_translate(mamba_indices)) if mamba_indices is not None else None ) @@ -2529,7 +2567,9 @@ class HybridLinearKVPool(KVCache): kv_cpu, mamba_cpu = cache_cpu self.full_kv_pool.load_cpu_copy(kv_cpu, indices) if mamba_cpu is not None and mamba_indices is not None: - self.mamba_pool.load_cpu_copy(mamba_cpu, mamba_indices) + self.mamba_pool.load_cpu_copy( + mamba_cpu, self._mamba_translate(mamba_indices) + ) def get_v_head_dim(self): return self.full_kv_pool.get_value_buffer(0).shape[-1] @@ -2676,7 +2716,7 @@ class MLATokenToKVPool(KVCache): cache_k: torch.Tensor, cache_v: torch.Tensor, ): - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 @@ -2882,7 +2922,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): cache_v: torch.Tensor, ): # loc_info may be a KVWriteLoc; MLA pools have no SWA target. - loc, _ = unwrap_write_loc(loc_info) + loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA-FP4)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py new file mode 100644 index 000000000..921badbaf --- /dev/null +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -0,0 +1,2474 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""MultiEndedAllocator: one allocator per sub-pool over a `UnifiedKVPool`. + +`alloc*` run the upstream kernels ONCE in virtual space using `free_virtual_ids` +as the free-page pointer, then bind consumed virtual pages to physical pages so +`translate_kv_loc` resolves. Public methods take/return TOKEN-granular tensors; +`free_virtual_ids` and the v2p/p2v tables are page-granular. For `page_size == 1` +page math collapses to slot math byte-identically. +""" + +from __future__ import annotations + +import inspect +import logging +import os +from typing import Dict, List, Optional, Set, Tuple + +import torch +from torch.profiler import record_function + +from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.paged import ( + alloc_decode_kernel, + alloc_extend_kernel, +) +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.triton_ops.virtual_slot import alloc_bind_inplace +from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool +from sglang.srt.utils.common import get_num_new_pages, next_power_of_2 + +logger = logging.getLogger(__name__) + + +# OFF (default): cat unsorted, `_flush` sorts once. ON: sort after each cat. +_SORT_FREE_LIST_AFTER_MERGE = envs.SGLANG_SORT_FREE_LIST_AFTER_MERGE.get() + + +import atexit +import signal +import time as _time_mod # local alias so tests can patch +import weakref + +_LAZY_COMPACTION_STATS_ENABLED = envs.SGLANG_LOG_LAZY_COMPACTION_STATS.get() +_LAZY_COMPACTION_STATS_INTERVAL_SEC = float( + envs.SGLANG_LOG_LAZY_COMPACTION_STATS_INTERVAL_SEC.get() +) +# Signal handler emits each instance's final counters (atexit misses signal exits). +_STATS_INSTANCES: weakref.WeakSet[MultiEndedAllocator] = weakref.WeakSet() +_SIGNAL_HANDLERS_INSTALLED = False + + +def _emit_all_final_stats(reason: str) -> None: + for inst in list(_STATS_INSTANCES): + try: + inst._emit_stats_final(reason=reason) + except Exception: + pass + + +def _signal_handler(signum, frame): + try: + sig_name = signal.Signals(signum).name + except (ValueError, AttributeError): + sig_name = str(signum) + _emit_all_final_stats(reason=sig_name) + signal.signal(signum, signal.SIG_DFL) + os.kill(os.getpid(), signum) + + +def _install_signal_handlers_once() -> None: + global _SIGNAL_HANDLERS_INSTALLED + if _SIGNAL_HANDLERS_INSTALLED: + return + _SIGNAL_HANDLERS_INSTALLED = True + # Only override the default handler (the scheduler subprocess installs none). + for sig in (signal.SIGTERM, signal.SIGINT): + try: + prev = signal.getsignal(sig) + if prev in (signal.SIG_DFL, signal.SIG_IGN, None): + signal.signal(sig, _signal_handler) + except (ValueError, OSError): + # Raises off the main thread — skip. + pass + + +class MultiEndedAllocator(BaseTokenToKVPoolAllocator): + """Allocator for one sub-pool over a `UnifiedKVPool`.""" + + def __init__( + self, + *, + kvcache, + unified_buffer: UnifiedKVPool, + sub_pool_name: str, + device: str, + is_id_owner: bool, + page_size: int = 1, + need_sort: bool = False, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, + ): + spec = unified_buffer.spec(sub_pool_name) + max_slots = unified_buffer.max_slots(sub_pool_name) + super().__init__( + size=max_slots, + page_size=page_size, + dtype=spec.get_dtype(), + device=device, + kvcache=kvcache, + need_sort=need_sort, + ) + self.unified_buffer = unified_buffer + self.sub_pool_name = sub_pool_name + self.spec = spec + self.max_slots = max_slots + self.grow_direction = spec.grow_direction + self.entry_bytes = spec.entry_bytes() + self.min_slot_index = unified_buffer.min_slot_index(sub_pool_name) + self.is_id_owner = is_id_owner + # Overlap mode: `free` drops a wait_stream(forward_stream) barrier so its + # v2p writes + move kernel serialize after the in-flight forward. + self.forward_stream = forward_stream + + # --- Page-aware bookkeeping --- + # `min_page_index` = ceil(min_slot_index / page_size), keeping the + # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). + self.page_size = page_size + self.num_pages = max_slots // page_size + self.min_page_index = (self.min_slot_index + page_size - 1) // page_size + self.entry_bytes_per_page = self.entry_bytes * page_size + + # v2p / p2v sized by PAGES. Page 0 is the padding anchor; trailing row is + # the -1 sentinel. + self.virtual_to_physical = torch.full( + (self.num_pages + 1,), + -1, + dtype=torch.int64, + device=device, + ) + self.physical_to_virtual = torch.full( + (self.num_pages + 1,), + -1, + dtype=torch.int64, + device=device, + ) + # Back-compat alias (count of virtual PAGES) consulted by is_slot_allocated. + self.num_virtual_ids = self.num_pages + + self._peer: Optional[MultiEndedAllocator] = None + + # Inverse history of relocations (spec rollback), at PAGE granularity. + self._inverse_history: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = ( + [] + ) + + # --- Lazy compaction state (all unused when lazy_compaction=False) --- + # `_free_phys_pages`: GPU free list of physical PAGE ids, sorted at `_flush`. + # `_pending_reuse`: compaction-src pages whose remap completed but whose + # reader event hasn't fired — can't re-enter the free list until the read + # settles (else a future alloc's WRITE races the READ). + # `live_page_count`: CPU slot-conservation counter, invariant under compaction. + # KV copy and v2p/p2v remap both run on `schedule_stream`, so single-stream + # ordering serializes them — no separate copy-done event needed. + self.lazy_compaction = lazy_compaction + self._free_phys_pages: torch.Tensor = torch.empty( + 0, dtype=torch.int64, device=device + ) + # Keyed by Event, ONE entry per BATCH. `(cpu_list, gpu_tensor)`: cpu_list + # drives the Set update (no sync); gpu_tensor is the SAME tensor + # `_commit_move_batch` remapped, kept alive so drain cats it without an H2D. + self._pending_reuse: Dict[ + torch.cuda.Event, + Tuple[List[int], torch.Tensor], + ] = {} + # CPU mirror of `_pending_reuse` for O(1) membership in the survivor walk. + self._pending_reuse_pages_cpu: Set[int] = set() + # Cumulative observability counters (NOT reset at clear()). + self._stats_n_free_lazy: int = 0 + self._stats_n_release_batch: int = 0 + self._stats_n_drain_calls: int = 0 + self._stats_n_drain_did_work: int = 0 + self._stats_n_drained_pages_total: int = 0 + self._stats_n_flush_calls: int = 0 + self._stats_n_flush_did_work: int = 0 + self._stats_n_flush_moves: int = 0 + self._stats_n_pages_absorbed: int = 0 + self._stats_peak_free_list_len: int = 0 + self._stats_peak_pending_pages: int = 0 + self._stats_n_emits: int = 0 + self._stats_last_emit_ts: float = _time_mod.monotonic() + self._stats_final_emitted: bool = False + if _LAZY_COMPACTION_STATS_ENABLED: + atexit.register(self._emit_stats_final, reason="atexit") + _STATS_INSTANCES.add(self) + _install_signal_handlers_once() + self.live_page_count = 0 + 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. + # Only the tensor reference is stored; `_flush` materializes the write-set + # lazily, avoiding a launch-time sync. + self._inflight_forward: Optional[Tuple[torch.cuda.Event, torch.Tensor]] = None + + # Per-call move cap on NON-urgent `_flush`: bounds work per `on_idle()` so a + # large backlog doesn't block ZMQ IPC; the next flush picks up the rest. + # Urgent (alloc-shortfall retry) is uncapped — must drain everything. + self._lazy_max_moves_per_call = int( + os.environ.get("SGLANG_LAZY_COMPACTION_MAX_MOVES_PER_CALL", "4096") + ) + + self.clear() + + logger.info( + "[unified-memory-pool] MultiEndedAllocator(%r) ready: grow=%s, max_slots=%d, " + "min_slot_index=%d, page_size=%d, num_pages=%d, min_page_index=%d, " + "entry_bytes=%d, entry_bytes_per_page=%d, is_id_owner=%s, " + "initial_watermark_page=%d, allocatable_pages=%d", + self.sub_pool_name, + self.grow_direction, + self.max_slots, + self.min_slot_index, + self.page_size, + self.num_pages, + self.min_page_index, + self.entry_bytes, + self.entry_bytes_per_page, + self.is_id_owner, + self.watermark_physical, + self.num_pages - self.min_page_index, + ) + + # -- peer binding -- + + def bind_peer(self, peer: MultiEndedAllocator) -> None: + self._peer = peer + + @property + def peer(self) -> Optional[MultiEndedAllocator]: + return self._peer + + # -- state -- + + def clear(self) -> None: + """Reset to initial state. Pages in `[0, min_page_index)` are reserved.""" + if self.grow_direction == "up": + self.watermark_physical = self.min_page_index + else: + self.watermark_physical = self.num_pages - 1 + self.virtual_to_physical.fill_(-1) + # Virtual page 0 <-> physical page 0 (padding sink). + self.virtual_to_physical[0] = 0 + self.virtual_to_physical[-1] = -1 # trailing sentinel + self.physical_to_virtual.fill_(-1) + self.physical_to_virtual[0] = 0 + self.physical_to_virtual[-1] = -1 + if self.is_id_owner: + self.free_virtual_ids = torch.arange( + self.min_page_index, + self.num_pages, + dtype=torch.int64, + device=self.device, + ) + else: + self.free_virtual_ids = None + self.is_not_in_free_group = True + self.free_group: List[torch.Tensor] = [] + self._inverse_history.clear() + self._free_phys_pages = torch.empty(0, dtype=torch.int64, device=self.device) + self._pending_reuse.clear() + self._pending_reuse_pages_cpu.clear() + self.live_page_count = 0 + self._inflight_forward = None + self._latest_forward_done_event = None + + def backup_state(self): + # Spec-decode allocates only inside a backup window (no free), so + # `_inverse_history` doesn't grow under correct usage. + return ( + self.watermark_physical, + (len(self.free_virtual_ids) if self.is_id_owner else None), + len(self._inverse_history), + ) + + def restore_state(self, state): + watermark, n_free_virtual, n_inverse = state + self.watermark_physical = watermark + if self.is_id_owner and n_free_virtual is not None: + pass # spec asserted off; no free-list rollback. + new_entries = self._inverse_history[n_inverse:] + if new_entries: + logger.warning( + "MultiEndedAllocator.restore_state: %d relocation(s) recorded inside " + "a backup window (sub_pool=%s). Eager compaction is not fully " + "reversible; SGLang's spec path should not produce a free() inside a " + "backup window.", + len(new_entries), + self.sub_pool_name, + ) + del self._inverse_history[n_inverse:] + return new_entries + + def clear_inverse_history(self) -> None: + self._inverse_history.clear() + + # -- size reporting -- + + def _allocated_pages(self) -> int: + """Number of allocated PAGES (TOKEN callers use `allocated_count()`).""" + if self.grow_direction == "up": + return max(0, self.watermark_physical - self.min_page_index) + return max(0, self.num_pages - 1 - self.watermark_physical) + + def allocated_count(self) -> int: + """LIVE allocated TOKENS (excludes lazy holes / pending). + + TOKENS, not pages — the leak checker's invariant is in tokens. Lazy mode + uses `live_page_count` (invariant under compaction); the watermark span + over-counts because holes/pending sit inside it but aren't live. + """ + if self.lazy_compaction: + return self.live_page_count * self.page_size + return self._allocated_pages() * self.page_size + + def is_slot_allocated(self, slot: int) -> bool: + """Whether the PAGE containing this virtual id is in use.""" + virt_page = slot // self.page_size + if virt_page < 0 or virt_page >= self.num_pages: + return False + return int(self.virtual_to_physical[virt_page].item()) != -1 + + def allocator_state_str(self) -> str: + return ( + f"sub_pool={self.sub_pool_name!r}, grow_direction={self.grow_direction}, " + f"is_id_owner={self.is_id_owner}, page_size={self.page_size}, " + f"min_page_index={self.min_page_index}, " + f"num_pages={self.num_pages}, " + f"watermark_physical={self.watermark_physical}, " + f"allocated_pages={self._allocated_pages()}" + ) + + def _byte_high_frontier(self) -> int: + """Byte just past this side's last-allocated page (grow-up) / buffer top (grow-down).""" + if self.grow_direction == "up": + return self.watermark_physical * self.entry_bytes_per_page + return self.num_pages * self.entry_bytes_per_page + + def _byte_low_frontier(self) -> int: + """Byte starting this side's allocatable range (grow-up) / just below its lowest live page (grow-down).""" + if self.grow_direction == "up": + return self.min_page_index * self.entry_bytes_per_page + return (self.watermark_physical + 1) * self.entry_bytes_per_page + + def _current_gap_bytes(self) -> int: + """Free byte band between this side's frontier and the peer's CURRENT frontier.""" + if self.grow_direction == "up": + my_high = self._byte_high_frontier() + peer_low = ( + self._peer._byte_low_frontier() + if self._peer is not None + else self.unified_buffer.total_bytes + ) + return max(0, peer_low - my_high) + my_low = self._byte_low_frontier() + peer_high = self._peer._byte_high_frontier() if self._peer is not None else 0 + return max(0, my_low - peer_high) + + def _available_tokens(self, extra_gap_bytes: int = 0) -> int: + """Tokens allocatable given `extra_gap_bytes` of ADDED gap room + (0 == current realizable; >0 == post-peer-compaction). + + `pages_by_index_space` is OWN index headroom, unaffected by + `extra_gap_bytes`: peer bytes can't add page indices to our own table. + """ + gap_bytes = self._current_gap_bytes() + extra_gap_bytes + pages_by_bytes = gap_bytes // self.entry_bytes_per_page + pages_by_index_space = ( + self.num_pages - self.min_page_index - self._allocated_pages() + ) + pages_extend = min(pages_by_bytes, pages_by_index_space) + # Lazy: drainable holes don't consume new bytes. + pages_drain = len(self._free_phys_pages) if self.lazy_compaction else 0 + return (pages_extend + pages_drain) * self.page_size + + def available_size(self) -> int: + """Tokens allocatable RIGHT NOW (no peer compaction). + + Alloc shortfall gates consult this to decide whether to peer-flush, so it + MUST NOT fold in peer holes (use `schedulable_available_size()` for that). + """ + return self._available_tokens() + + def _peer_drainable_hole_bytes(self) -> int: + """Gap bytes a peer urgent flush would release. Only `_free_phys_pages` + count — NOT `_pending_reuse` (awaiting an event) — so the credit is realizable. + """ + peer = self._peer + if peer is None or not peer.lazy_compaction: + return 0 + return len(peer._free_phys_pages) * peer.entry_bytes_per_page + + def schedulable_available_size(self) -> int: + """Tokens allocatable AFTER a peer urgent-flush (realizable-with-compaction). + Used by composite views; alloc gates use `available_size()`. + """ + return self._available_tokens(extra_gap_bytes=self._peer_drainable_hole_bytes()) + + def _flush_peer_for_alloc(self, need_tokens: int) -> bool: + """One urgent peer-flush on alloc shortfall; returns whether THIS side now + has enough. Only PEER compaction releases gap bytes (own compaction is net 0). + """ + if not (self.lazy_compaction and self._peer is not None): + return False + self._peer._flush(urgent=True) + return need_tokens <= self.available_size() + + # -- physical-slot / physical-page primitives -- + + def take_physical(self, need_size: int) -> Optional[torch.Tensor]: + """Reserve `need_size` TOKENS (multiple of page_size), returning backing + physical PAGE ids, or `None` on shortfall. + + Eager: pure watermark advance. Lazy: drain `_free_phys_pages` holes first, + then extend the watermark (extend first so state is untouched on failure). + """ + with record_function("MultiEndedAlloc.take_physical"): + if need_size <= 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + assert need_size % self.page_size == 0, ( + f"take_physical: need_size={need_size} must be a multiple of " + f"page_size={self.page_size}" + ) + num_pages = need_size // self.page_size + + if not self.lazy_compaction: + return self._take_physical_eager(num_pages) + + # Lazy: slice the GPU free list (no D2H). sort ON: take deepest-in-band + # per direction (greedy clustering). sort OFF: take from front. + n_drain = min(num_pages, int(self._free_phys_pages.shape[0])) + need_more = num_pages - n_drain + + # Extend first (state untouched on failure), then drain holes. + if need_more > 0: + if not self._extend_watermark(need_more): + return None + + if n_drain > 0: + if _SORT_FREE_LIST_AFTER_MERGE: + if self.grow_direction == "up": + drained_t = self._free_phys_pages[:n_drain] + self._free_phys_pages = self._free_phys_pages[n_drain:] + else: + drained_t = self._free_phys_pages[-n_drain:].flip(0) + self._free_phys_pages = self._free_phys_pages[:-n_drain] + else: + drained_t = self._free_phys_pages[:n_drain] + self._free_phys_pages = self._free_phys_pages[n_drain:] + else: + drained_t = None + + self.live_page_count += num_pages + + if drained_t is None: + return self._take_physical_arange(num_pages) + + # Pure drain — clone off the free-list view so rebindings don't pin it. + if need_more == 0: + return drained_t.clone() + + # Mixed: drained holes ++ extended pages (`bind` is order-agnostic). + if self.grow_direction == "up": + new_wm = self.watermark_physical + extended_t = torch.arange( + new_wm - need_more, + new_wm, + dtype=torch.int64, + device=self.device, + ) + else: + new_wm = self.watermark_physical + extended_t = torch.arange( + new_wm + need_more, + new_wm, + -1, + dtype=torch.int64, + device=self.device, + ) + return torch.cat([drained_t, extended_t]) + + def _take_physical_eager(self, num_pages: int) -> Optional[torch.Tensor]: + """Eager-mode take_physical — contiguous range.""" + if self.grow_direction == "up": + start = self.watermark_physical + end_exclusive = start + num_pages + if end_exclusive > self.num_pages: + return None + phys_pages = torch.arange( + start, end_exclusive, dtype=torch.int64, device=self.device + ) + self.watermark_physical = end_exclusive + return phys_pages + else: + end = self.watermark_physical + start = end - num_pages + 1 + if start < self.min_page_index: + return None + phys_pages = torch.arange( + start, end + 1, dtype=torch.int64, device=self.device + ) + self.watermark_physical -= num_pages + return phys_pages + + def _extend_watermark(self, num_pages: int) -> bool: + """Advance the watermark by `num_pages` (lazy-path helper). Returns False + on index-space overflow OR crossing the PEER's byte frontier. + """ + if self.grow_direction == "up": + new_wm = self.watermark_physical + num_pages + if new_wm > self.num_pages: + return False + # Peer (grow-down) sits ABOVE; don't extend past its low frontier. + if self._peer is not None: + peer_low_pages = ( + self._peer._byte_low_frontier() // self.entry_bytes_per_page + ) + if new_wm > peer_low_pages: + return False + self.watermark_physical = new_wm + else: + new_wm = self.watermark_physical - num_pages + if new_wm < self.min_page_index - 1: + return False + # Peer (grow-up) sits BELOW; `new_wm + 1` (our new lowest live page) + # must stay strictly above the peer's high frontier. + if self._peer is not None: + peer_high_pages = ( + self._peer._byte_high_frontier() // self.entry_bytes_per_page + ) + if new_wm + 1 < peer_high_pages: + return False + self.watermark_physical = new_wm + return True + + def _take_physical_arange(self, num_pages: int) -> torch.Tensor: + """Contiguous arange for an already-applied watermark extension.""" + if self.grow_direction == "up": + return torch.arange( + self.watermark_physical - num_pages, + self.watermark_physical, + dtype=torch.int64, + device=self.device, + ) + return torch.arange( + self.watermark_physical + 1, + self.watermark_physical + num_pages + 1, + dtype=torch.int64, + device=self.device, + ) + + def take_physical_pages(self, num_pages: int) -> Optional[torch.Tensor]: + """Page-granular wrapper around ``take_physical``.""" + with record_function("MultiEndedAlloc.take_physical_pages"): + return self.take_physical(num_pages * self.page_size) + + def bind(self, virtual_ids: torch.Tensor, physical_ids: torch.Tensor) -> None: + """Bind page-granular virtual ids to physical ids.""" + with record_function("MultiEndedAlloc.bind"): + self.virtual_to_physical[virtual_ids] = physical_ids + self.physical_to_virtual[physical_ids] = virtual_ids + + def bind_pages( + self, virtual_pages: torch.Tensor, physical_pages: torch.Tensor + ) -> None: + """Page-granular alias of ``bind``.""" + with record_function("MultiEndedAlloc.bind_pages"): + self.bind(virtual_pages, physical_pages) + + # -- fused take_physical_pages + bind_pages -- + + def _alloc_bind_fast_or_slow( + self, v_pages: torch.Tensor, N: int + ) -> Optional[torch.Tensor]: + """Fuse `take_physical_pages` + `bind` into ONE Triton kernel when no + holes need draining; fall through to the slow path (drains holes first) + when holes exist. Returns physical page ids [N], or None on shortfall. + """ + with record_function("MultiEndedAlloc._alloc_bind_fast_or_slow"): + if N == 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + + # FAST PATH: eager, or lazy with no current holes. + if not self.lazy_compaction or self._free_phys_pages.numel() == 0: + start_wm = self.watermark_physical # kernel's `start_phys` + + # Lazy uses `_extend_watermark` (index + peer checks); eager + # inlines the index-only check to match `_take_physical_eager`. + if self.lazy_compaction: + if not self._extend_watermark(N): + return None + else: + if self.grow_direction == "up": + new_wm = start_wm + N + if new_wm > self.num_pages: + return None + self.watermark_physical = new_wm + else: + new_wm = start_wm - N + if new_wm < self.min_page_index - 1: + return None + self.watermark_physical = new_wm + + # Lowest physical id of the new range (both directions yield + # ascending `[start_phys, start_phys + N)`). + if self.grow_direction == "up": + start_phys = start_wm + else: + start_phys = start_wm - N + 1 + + phys_pages = alloc_bind_inplace( + v_pages, + self.virtual_to_physical, + self.physical_to_virtual, + start_phys, + ) + + if self.lazy_compaction: # live_page_count tracked only in lazy mode + self.live_page_count += N + return phys_pages + + # SLOW PATH: holes exist — drain them first, then bind. + phys_pages = self.take_physical_pages(N) + if phys_pages is None: + return None + self.bind(v_pages, phys_pages) + return phys_pages + + # -- translate (virtual TOKEN ids -> physical TOKEN ids) -- + + def translate_kv_loc( + self, + virt_tokens: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Translate token-granular virtual ids to physical ids. + + ``out=`` writes in-place into a caller-owned buffer — required under + cuda-graph capture for buffer-stability (the captured graph records the + gather against a fixed ``data_ptr``). + """ + if out is not None: + assert out.dtype == torch.int64, ( + f"translate_kv_loc: out= dtype must be int64 (matches v2p), " + f"got {out.dtype}" + ) + assert out.shape == virt_tokens.shape, ( + f"translate_kv_loc: out= shape {tuple(out.shape)} must match " + f"virt_tokens shape {tuple(virt_tokens.shape)}" + ) + with record_function("MultiEndedAlloc.translate_kv_loc"): + return self._translate_kv_loc_impl(virt_tokens, out) + + def _translate_kv_loc_impl( + self, + virt_tokens: torch.Tensor, + out: Optional[torch.Tensor], + ) -> torch.Tensor: + # Tombstone-safety clamp: tombstoned v2p entries (-1) must not reach + # `k_buffer[-1]` (illegal access under captured graph replay). Clamp to 0 + # routes any tombstoned read/write to physical slot 0 — reserved + # padding-sink space by the `min_slot_index` invariant (bytes [0, entry_max) + # across all sub-pools hold no real data). + if self.page_size == 1: + if out is not None: + # `index_select(out=out)` forbids index/out aliasing, but the + # canonical caller does in-place `translate(kv_indices, out=kv_indices)`. + # Route through a transient gather + `copy_` to satisfy that contract. + tmp = torch.index_select(self.virtual_to_physical, 0, virt_tokens) + tmp = torch.clamp_min(tmp, 0) + out.copy_(tmp) + return out + result = torch.index_select(self.virtual_to_physical, 0, virt_tokens) + return torch.clamp_min(result, 0) + # page_size > 1: page math. `virt_pages`/`offsets` are fresh, so they + # cannot alias `out` — `index_select(out=out)` is safe. + virt_pages = virt_tokens // self.page_size + offsets = virt_tokens % self.page_size + if out is not None: + torch.index_select(self.virtual_to_physical, 0, virt_pages, out=out) + out.mul_(self.page_size) + out.add_(offsets) + out.clamp_(min=0) # tombstoned page: -1*ps + offset in [-ps, -1] + return out + phys_pages = self.virtual_to_physical[virt_pages] + result = phys_pages * self.page_size + offsets + return torch.clamp_min(result, 0) + + # -- alloc -- + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + """Allocate `need_size` virtual TOKEN ids (id-owner only). Returns + token-granular, page-structured ids, or None on shortfall. + + `need_size` MUST be a multiple of `page_size`. All allocator GPU ops run + on `schedule_stream`; `alloc` needs no `wait_stream` barrier because its + v2p/p2v writes are picked up by the forward via the existing + `forward_stream.wait_stream(schedule_stream)` at the top of `run_batch`. + """ + with record_function("MultiEndedAlloc.alloc"): + assert self.is_id_owner, ( + f"MultiEndedAllocator({self.sub_pool_name!r}).alloc called on a " + "non-id-owner allocator; use alloc_with_virtual instead" + ) + if need_size <= 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + assert need_size % self.page_size == 0, ( + f"MultiEndedAllocator({self.sub_pool_name!r}).alloc: need_size=" + f"{need_size} must be a multiple of page_size={self.page_size}" + ) + if need_size > self.available_size(): + # Shortfall: flush the PEER, not own. Own compaction is net 0 + # (each move trades 1 hole for +1 gap byte); only peer compaction + # releases bytes into the shared gap that own extension consumes. + if not self._flush_peer_for_alloc(need_size): + return None + num_pages = need_size // self.page_size + v_pages = self.free_virtual_ids[:num_pages] + self.free_virtual_ids = self.free_virtual_ids[num_pages:] + phys_pages = self._alloc_bind_fast_or_slow(v_pages, num_pages) + if phys_pages is None: + self.free_virtual_ids = torch.cat([v_pages, self.free_virtual_ids]) + return None + if self.page_size == 1: + return v_pages # v_pages already IS the token id list + # Expand page ids to token ids: (P, 1) * S + (S,) → (P, S) → (P*S,). + return ( + v_pages[:, None] * self.page_size + + torch.arange(self.page_size, device=self.device) + ).reshape(-1) + + def alloc_with_virtual(self, virtual_pages: torch.Tensor) -> None: + """Take physical PAGES for caller-supplied virtual PAGE ids + (physical-holding non-owner; the SWA `swa` sub-allocator). + + Input is virtual PAGE ids (not token ids): the composite snapshots the + virtual pages before the id-owner consumes them from its free-list. + """ + with record_function("MultiEndedAlloc.alloc_with_virtual"): + if virtual_pages.numel() == 0: + return + phys_pages = self._alloc_bind_fast_or_slow( + virtual_pages, int(virtual_pages.numel()) + ) + assert phys_pages is not None, ( + f"MultiEndedAllocator({self.sub_pool_name!r}).alloc_with_virtual: out of " + "physical room (the composite's byte-budget check should have caught this)" + ) + + # -- paged alloc surface -- + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + num_new_pages: Optional[int] = None, + ) -> Optional[torch.Tensor]: + """Allocate ``extend_num_tokens`` new tokens across ``bs`` requests, + preserving the tail-page-reuse contract. + + Runs the kernel in VIRTUAL space (``free_page_ptr == free_virtual_ids``), + so ``out_indices`` are virtual token ids. Each consumed virtual page is + then bound to a physical page on THIS sub-allocator; without that binding + v2p stays -1 and translation yields negative ids → CUDA OOB. + """ + with record_function("MultiEndedAlloc.alloc_extend"): + assert ( + self.is_id_owner + ), f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" + if num_new_pages is None: + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + if num_new_pages > len(self.free_virtual_ids): + return None + # Lazy: physical-capacity pre-check; on shortfall flush the PEER (own + # compaction is internal — see `alloc`). + need_tokens = num_new_pages * self.page_size + if need_tokens > self.available_size(): + if not self._flush_peer_for_alloc(need_tokens): + return None + bs = len(prefix_lens) + if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len( + self.free_virtual_ids + ): + self.merge_and_sort_free() + + # Snapshot the virtual pages the kernel will consume, to bind them to + # physical pages afterward (else v2p stays -1 → CUDA OOB). + if num_new_pages > 0: + new_virtual_pages = self.free_virtual_ids[:num_new_pages].clone() + else: + new_virtual_pages = None + + out_indices = torch.empty( + (extend_num_tokens,), dtype=torch.int64, device=self.device + ) + # `free_virtual_ids` passed as `free_page_ptr`: the kernel does + # `page_id * page_size + offset` regardless of virtual vs physical. + with record_function("MultiEndedAlloc.alloc_extend.kernel"): + alloc_extend_kernel[(bs,)]( + prefix_lens, + seq_lens, + last_loc, + self.free_virtual_ids, + out_indices, + next_power_of_2(bs), + self.page_size, + ) + + # Bind the consumed virtual pages to fresh physical pages here. The + # peer (swa side) binds the same pages via `alloc_with_virtual`. + if new_virtual_pages is not None: + phys_pages = self._alloc_bind_fast_or_slow( + new_virtual_pages, num_new_pages + ) + if phys_pages is None: + return None # defensive; pre-check should have prevented it + + self.free_virtual_ids = self.free_virtual_ids[num_new_pages:] + return out_indices # virtual token ids + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Allocate one new token per request (decode), preserving the + tail-page-reuse contract. Runs in virtual space; binds each consumed + virtual page on THIS sub-allocator (else v2p stays -1 → CUDA OOB). + """ + with record_function("MultiEndedAlloc.alloc_decode"): + assert ( + self.is_id_owner + ), f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" + bs = len(seq_lens) + # CPU-only count BEFORE the kernel, to snapshot the exact slice the + # kernel will consume. + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True + ) + if num_new_pages > len(self.free_virtual_ids): + return None + # Lazy: physical-capacity pre-check; on shortfall flush PEER. + need_tokens = num_new_pages * self.page_size + if need_tokens > self.available_size(): + if not self._flush_peer_for_alloc(need_tokens): + return None + if self.need_sort and bs > len(self.free_virtual_ids): + self.merge_and_sort_free() + + # Most decode steps reuse the prefix's tail page → num_new_pages == 0. + if num_new_pages > 0: + new_virtual_pages = self.free_virtual_ids[:num_new_pages].clone() + else: + new_virtual_pages = None + + out_indices = torch.empty((bs,), dtype=torch.int64, device=self.device) + with record_function("MultiEndedAlloc.alloc_decode.kernel"): + alloc_decode_kernel[(bs,)]( + seq_lens, + last_loc, + self.free_virtual_ids, + out_indices, + next_power_of_2(bs), + self.page_size, + ) + + if new_virtual_pages is not None: + phys_pages = self._alloc_bind_fast_or_slow( + new_virtual_pages, num_new_pages + ) + if phys_pages is None: + return None + + self.free_virtual_ids = self.free_virtual_ids[num_new_pages:] + return out_indices # virtual token ids + + # -- free with eager compaction -- + + def free(self, free_index: torch.Tensor) -> None: + """Free virtual TOKEN ids: recover virtual PAGE ids, un-map v2p/p2v, + (if id-owner) recycle the page ids, trigger eager compaction. + + `free_index` is token-granular and need not be page-aligned. EAGER mode + drops one `wait_stream(forward_stream)` barrier so v2p/p2v writes and the + compaction move serialize with the in-flight forward. LAZY mode needs no + barrier (a freed `v` has no live reader, so the scatters are + disjoint-element from any forward read, atomic on Ampere+/Hopper) and + defers compaction to `_flush`. + """ + with record_function("MultiEndedAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if not self.is_not_in_free_group: + self.free_group.append(free_index) + return + if self.lazy_compaction: + self._free_lazy(free_index) + return + # --- EAGER path --- + # Near-no-op in normal mode (sampling's CPU sync already drained + # forward_stream); in overlap mode it serializes free+compaction with + # the in-flight forward. + if self.forward_stream is not None: + with record_function("MultiEndedAlloc.free.wait_stream"): + torch.cuda.current_stream().wait_stream(self.forward_stream) + with record_function("MultiEndedAlloc.free.v2p_lookup"): + free_v_pages = torch.unique( + free_index.detach().to(torch.int64) // self.page_size + ) + freed_p_pages = self.virtual_to_physical[free_v_pages] + with record_function("MultiEndedAlloc.free.sync_check"): + # `.item()` forces a CPU/GPU sync — own trace region to measure it. + if bool((freed_p_pages < 0).any().item()): + self._raise_stale_slot_assertion( + free_v=free_v_pages, freed_p=freed_p_pages + ) + self.virtual_to_physical[free_v_pages] = -1 + if self.is_id_owner: + self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages]) + self._compact_pending(freed_p_pages) + + def _free_lazy(self, free_index: torch.Tensor) -> None: + """Lazy free path: disjoint-element scatters + ONE `torch.cat` onto + `_free_phys_pages`. No sort, no boundary absorb, no watermark mutation, + no D2H sync. Boundary absorption is deferred to `_flush`. + + ps==1 skips `torch.unique` (token == page and `free_index` is already + unique per caller contract); ps>1 needs it to dedup same-page tokens. + Callers must not double-free: a tombstone (-1) here would be cat'd onto + the free list. + """ + self._stats_n_free_lazy += 1 + with record_function("MultiEndedAlloc._free_lazy"): + with record_function("MultiEndedAlloc._free_lazy.v2p_lookup"): + free_v_pages_raw = free_index.detach().to(torch.int64) + if self.page_size == 1: + free_v_pages = free_v_pages_raw + else: + free_v_pages = torch.unique(free_v_pages_raw // self.page_size) + freed_p_pages = self.virtual_to_physical[free_v_pages] + # Disjoint-element scatters — no barrier (a freed v has no live reader; + # per-element scatter writes are atomic). + self.virtual_to_physical[free_v_pages] = -1 + self.physical_to_virtual[freed_p_pages] = -1 + if self.is_id_owner: + self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages]) + self._free_phys_pages = torch.cat([self._free_phys_pages, freed_p_pages]) + if _SORT_FREE_LIST_AFTER_MERGE: + self._free_phys_pages, _ = torch.sort(self._free_phys_pages) + self.live_page_count -= int(freed_p_pages.shape[0]) + + def _release_phys_pages_batch(self, pages: torch.Tensor) -> None: + """Cat `pages` onto `_free_phys_pages` (+ optional sort). Called by `_flush` + at END to merge event-fired compaction-srcs (`released_fired`) AFTER the + trailing dst-slice, keeping `_free_phys_pages == holes_cpu` during the walk. + + No watermark / `live_page_count` change — these are vacated src positions + re-entering as PURE storage, not freshly-freed live pages. + """ + if pages.numel() == 0: + return + self._stats_n_release_batch += 1 + with record_function("MultiEndedAlloc._release_phys_pages_batch"): + self._free_phys_pages = torch.cat([self._free_phys_pages, pages]) + if _SORT_FREE_LIST_AFTER_MERGE: + self._free_phys_pages, _ = torch.sort(self._free_phys_pages) + + def _compact_pending(self, freed_physical_pages: torch.Tensor) -> None: + """Eager compaction over the freed PHYSICAL pages: move survivors from the + vacated band (K pages adjacent to the watermark) into the holes in the kept + band, advance the watermark, remap the tables. `src`/`dst` are disjoint by + construction, so the batched copy is order-independent. The caller's + `wait_stream` barrier already serialized us with the in-flight forward. + """ + with record_function("MultiEndedAlloc._compact_pending"): + self._compact_pending_impl(freed_physical_pages) + + def _compact_pending_impl(self, freed_physical_pages: torch.Tensor) -> None: + freed_set = set(int(x) for x in freed_physical_pages.tolist()) + if not freed_set: + return + K = len(freed_set) + if self.grow_direction == "up": + # allocated == [min_page_index, old_wm); after the free == [min_page_index, new_wm) + old_wm = self.watermark_physical + new_wm = old_wm - K + assert new_wm >= self.min_page_index, ( + f"_compact_pending({self.sub_pool_name!r}): freeing {K} pages " + f"would push the watermark below min_page_index " + f"({new_wm} < {self.min_page_index})" + ) + assert all(self.min_page_index <= h < old_wm for h in freed_set), ( + f"_compact_pending({self.sub_pool_name!r}): freed physical pages " + f"{sorted(freed_set)} not all within allocated range " + f"[{self.min_page_index}, {old_wm})" + ) + # vacated band = [new_wm, old_wm); kept band = [min_page_index, new_wm) + src_list = [s for s in range(new_wm, old_wm) if s not in freed_set] + dst_list = sorted(h for h in freed_set if h < new_wm) + self.watermark_physical = new_wm + vacated_lo, vacated_hi = new_wm, old_wm + else: + # allocated == (old_wm, num_pages); after the free == (new_wm, num_pages) + old_wm = self.watermark_physical + new_wm = old_wm + K + assert new_wm <= self.num_pages - 1, ( + f"_compact_pending({self.sub_pool_name!r}): freeing {K} pages " + f"would push the watermark above num_pages " + f"({new_wm} > {self.num_pages - 1})" + ) + assert all(old_wm < h < self.num_pages for h in freed_set), ( + f"_compact_pending({self.sub_pool_name!r}): freed physical pages " + f"{sorted(freed_set)} not all within allocated range " + f"({old_wm}, {self.num_pages})" + ) + # vacated band = (old_wm, new_wm] = [old_wm+1, new_wm+1); kept band = (new_wm, num_pages) + src_list = [s for s in range(old_wm + 1, new_wm + 1) if s not in freed_set] + dst_list = sorted(h for h in freed_set if h > new_wm) + self.watermark_physical = new_wm + vacated_lo, vacated_hi = old_wm + 1, new_wm + 1 + + assert len(src_list) == len(dst_list), ( + f"_compact_pending({self.sub_pool_name!r}): {len(src_list)} survivors vs " + f"{len(dst_list)} holes — corrupt allocator state" + ) + + if src_list: + src_pages = torch.tensor(src_list, dtype=torch.int64, device=self.device) + dst_pages = torch.tensor(dst_list, dtype=torch.int64, device=self.device) + v_moved = self.physical_to_virtual[ + src_pages + ].clone() # read before clearing + + # Expand page ids to token ids for the token-granular move kernel. + if self.page_size == 1: + src_t, dst_t = src_pages, dst_pages + else: + offsets = torch.arange( + self.page_size, dtype=torch.int64, device=self.device + ) + src_t = (src_pages[:, None] * self.page_size + offsets).reshape(-1) + dst_t = (dst_pages[:, None] * self.page_size + offsets).reshape(-1) + + # Un-translated copy: the public copy_from translates virtual ids, + # which we must NOT do here. + move_fn = getattr(self._kvcache, "move_kv_cache", None) + if move_fn is not None: + move_fn(dst_t, src_t) + else: + copy_phys = getattr(self._kvcache, "_copy_from_physical", None) + assert copy_phys is not None, ( + f"sub-pool {self.sub_pool_name!r} supports neither move_kv_cache " + "nor _copy_from_physical" + ) + copy_phys(src_t, dst_t) + # Clear the vacated band, then re-bind the relocated dst pages. + self.physical_to_virtual[vacated_lo:vacated_hi] = -1 + self.virtual_to_physical[v_moved] = dst_pages + self.physical_to_virtual[dst_pages] = v_moved + self._inverse_history.append((src_pages, dst_pages, v_moved)) + else: + self.physical_to_virtual[vacated_lo:vacated_hi] = -1 + + # -- lazy compaction primitives -- + + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + """Stash the most-recent forward's `forward_done` event; `_pending_reuse` + uses it to gate src reuse on read-path settling. None = no in-flight forward. + """ + with record_function("MultiEndedAlloc.set_latest_forward_done_event"): + self._latest_forward_done_event = event + + def set_inflight_forward( + self, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + """Stash the just-launched forward's `forward_done` event + virtual + `out_cache_loc` for `_flush`'s write-race check. + + No GPU work — only references; `_flush` materializes the write-set lazily + on `schedule_stream`, avoiding a launch-time sync. Pass + `out_cache_loc_virtual=None` when the forward doesn't write this pool + (e.g. Mamba state, written by mamba kernels not `set_kv_buffer`). No-op + in eager mode. + """ + with record_function("MultiEndedAlloc.set_inflight_forward"): + if not self.lazy_compaction: + return + if out_cache_loc_virtual is None or out_cache_loc_virtual.numel() == 0: + # No write race on this pool — clear the slot so `_flush` + # short-circuits and the prior tensor reference can be GC'd. + self._inflight_forward = None + return + self._inflight_forward = (forward_done, out_cache_loc_virtual) + + def _materialize_inflight_write_set(self) -> Optional[Set[int]]: + """Materialize the in-flight forward's write-set (physical PAGE ids it is + about to write), or `None` if no in-flight forward / already completed. + Called inside `_flush` on `schedule_stream`. Pays a bs-sized D2H sync, but + only once per call and only when a survivor needs classifying. + """ + inflight = self._inflight_forward + if inflight is None: + return None + event, oclv = inflight + # Forward completed → no write race. Clear so later flushes in the same + # tick don't re-check the fired event. + if event.query(): + self._inflight_forward = None + return None + # `oclv` is non-None here (set_inflight_forward clears the slot otherwise). + with record_function("MultiEndedAlloc._materialize_inflight_write_set"): + phys_tokens = self.translate_kv_loc(oclv) + if self.page_size > 1: + phys_pages = (phys_tokens // self.page_size).unique() + else: + phys_pages = phys_tokens + return set(phys_pages.tolist()) # .tolist() syncs schedule_stream + + def _maybe_emit_stats(self) -> None: + """Env-gated periodic stats emit (at most once per interval) at `_flush` end. + Disabled unless `SGLANG_LOG_LAZY_COMPACTION_STATS=1`. + """ + if not _LAZY_COMPACTION_STATS_ENABLED: + return + now = _time_mod.monotonic() + if now - self._stats_last_emit_ts < _LAZY_COMPACTION_STATS_INTERVAL_SEC: + return + self._stats_last_emit_ts = now + self._stats_n_emits += 1 + cur_holes = int(self._free_phys_pages.shape[0]) + cur_pending = len(self._pending_reuse_pages_cpu) + self._stats_peak_free_list_len = max(self._stats_peak_free_list_len, cur_holes) + self._stats_peak_pending_pages = max( + self._stats_peak_pending_pages, cur_pending + ) + sort_tag = "ON" if _SORT_FREE_LIST_AFTER_MERGE else "OFF" + logger.info( + f"[lazy-stats sub={self.sub_pool_name!r}] " + f"free_lazy={self._stats_n_free_lazy} " + f"flush={self._stats_n_flush_calls} " + f"(work={self._stats_n_flush_did_work} " + f"moves={self._stats_n_flush_moves} " + f"abs={self._stats_n_pages_absorbed}) " + f"drain={self._stats_n_drain_did_work}/{self._stats_n_drain_calls} " + f"sort={sort_tag} " + f"peak_holes={self._stats_peak_free_list_len} " + f"peak_pending={self._stats_peak_pending_pages} " + f"cur_holes={cur_holes} cur_pending={cur_pending} " + f"live={self.live_page_count} wm={self.watermark_physical}" + ) + + def _emit_stats_final(self, reason: str = "exit") -> None: + """Force-emit final counters at shutdown (bypasses the interval gate). + Idempotent (signal handler + atexit may both fire); best-effort. + """ + if not _LAZY_COMPACTION_STATS_ENABLED: + return + if self._stats_final_emitted: + return + try: + cur_holes = int(self._free_phys_pages.shape[0]) + cur_pending = len(self._pending_reuse_pages_cpu) + self._stats_peak_free_list_len = max( + self._stats_peak_free_list_len, cur_holes + ) + self._stats_peak_pending_pages = max( + self._stats_peak_pending_pages, cur_pending + ) + sort_tag = "ON" if _SORT_FREE_LIST_AFTER_MERGE else "OFF" + self._stats_final_emitted = True + logger.info( + f"[lazy-stats FINAL sub={self.sub_pool_name!r} reason={reason}] " + f"free_lazy={self._stats_n_free_lazy} " + f"flush={self._stats_n_flush_calls} " + f"(work={self._stats_n_flush_did_work} " + f"moves={self._stats_n_flush_moves} " + f"abs={self._stats_n_pages_absorbed}) " + f"drain={self._stats_n_drain_did_work}/{self._stats_n_drain_calls} " + f"sort={sort_tag} " + f"peak_holes={self._stats_peak_free_list_len} " + f"peak_pending={self._stats_peak_pending_pages} " + f"cur_holes={cur_holes} cur_pending={cur_pending} " + f"live={self.live_page_count} wm={self.watermark_physical} " + f"n_emits={self._stats_n_emits}" + ) + except Exception: + pass + + def _drain_pending_reuse(self, *, urgent: bool) -> None: + """Move ready `_pending_reuse` entries back into `_free_phys_pages` via + pure-GPU `torch.cat`. + + * non-urgent: release only entries whose event is None or has fired. + * urgent: `stream.wait_event` (stream-side dep, not host block) on + unfired events, then release. + + ONE dict entry per BATCH (keyed by Event); cpu_list drives the Set update, + gpu_tensor is cat'd directly. No watermark / `live_page_count` change. + """ + self._stats_n_drain_calls += 1 + if not self._pending_reuse: + return + with record_function("MultiEndedAlloc._drain_pending_reuse"): + ready_tensors: List[torch.Tensor] = [] + ready_entries: List[Tuple[torch.cuda.Event, List[int]]] = [] + for event, (cpu_list, gpu_tensor) in self._pending_reuse.items(): + if event is None or event.query(): + ready_tensors.append(gpu_tensor) + ready_entries.append((event, cpu_list)) + elif urgent: + torch.cuda.current_stream().wait_event(event) + ready_tensors.append(gpu_tensor) + ready_entries.append((event, cpu_list)) + + for event, cpu_list in ready_entries: + del self._pending_reuse[event] + self._pending_reuse_pages_cpu.difference_update(cpu_list) + + if ready_tensors: + self._free_phys_pages = torch.cat( + [self._free_phys_pages] + ready_tensors + ) + self._stats_n_drain_did_work += 1 + self._stats_n_drained_pages_total += sum( + t.numel() for t in ready_tensors + ) + if _SORT_FREE_LIST_AFTER_MERGE: + self._free_phys_pages, _ = torch.sort(self._free_phys_pages) + + def maybe_drain_pending_reuse(self) -> None: + """Public scheduler hook (once per step): flow fired compaction-src pages + back into `_free_phys_pages` for immediate reuse without waiting for `_flush`. + """ + if not self.lazy_compaction: + return + if not self._pending_reuse: + return + self._drain_pending_reuse(urgent=False) + + def _topmost_survivor( + self, + start_hint: Optional[int] = None, + *, + holes_cpu: Optional[List[int]] = None, + j_in: Optional[int] = None, + ) -> Tuple[Optional[int], Optional[int]]: + """Topmost live PAGE in the allocated band (largest `p < watermark` for + grow-up / smallest `p > watermark` for grow-down), excluding holes + (`holes_cpu`, the sorted-ASCENDING snapshot) and `_pending_reuse_pages_cpu`. + + Two-pointer: `p` is monotonic and `holes_cpu` is sorted, so the hole cursor + `j` (threaded back via the returns) advances alongside for O(1) membership; + no exclude-set needed because uncommitted dsts have p2v=-1 and are correctly + reported by the snapshot. Returns `(p, j)`, or `(None, j)` if none. + + `holes_cpu`/`j_in` are optional only for test fixtures (else a `.tolist()` + sync); `_flush` always passes them. + """ + if holes_cpu is None: + holes_cpu = self._free_phys_pages.tolist() + if self.grow_direction == "up": + if start_hint is None or start_hint >= self.watermark_physical: + p = self.watermark_physical - 1 + else: + p = start_hint + j = j_in if j_in is not None else len(holes_cpu) - 1 + while p >= self.min_page_index: + while j >= 0 and holes_cpu[j] > p: + j -= 1 + is_hole = j >= 0 and holes_cpu[j] == p + if is_hole or p in self._pending_reuse_pages_cpu: + if is_hole: + j -= 1 + p -= 1 + continue + return p, j + return None, j + else: + if start_hint is None or start_hint <= self.watermark_physical: + p = self.watermark_physical + 1 + else: + p = start_hint + j = j_in if j_in is not None else 0 + while p < self.num_pages: + while j < len(holes_cpu) and holes_cpu[j] < p: + j += 1 + is_hole = j < len(holes_cpu) and holes_cpu[j] == p + if is_hole or p in self._pending_reuse_pages_cpu: + if is_hole: + j += 1 + p += 1 + continue + return p, j + return None, j + + def _absorb_boundary_holes(self, all_cpu: List[int]) -> Tuple[int, List[int]]: + """Retreat the watermark past free slots ALREADY contiguous with it, slice + them off `_free_phys_pages`, return ``(new_watermark, interior_holes_cpu)``. + `all_cpu` is the sorted-ascending snapshot; interior holes feed the survivor + walk. + """ + M = len(all_cpu) + wm = self.watermark_physical + n = 0 + if self.grow_direction == "up": + while n < M and all_cpu[M - 1 - n] == wm - 1 - n: + n += 1 + new_wm = wm - n + holes_cpu = all_cpu[: M - n] + self._free_phys_pages = self._free_phys_pages[: M - n] + else: + while n < M and all_cpu[n] == wm + 1 + n: + n += 1 + new_wm = wm + n + holes_cpu = all_cpu[n:] + self._free_phys_pages = self._free_phys_pages[n:] + self.watermark_physical = new_wm + self._stats_n_pages_absorbed += n + return new_wm, holes_cpu + + def _settle_inflight_forward(self) -> None: + """Stream-wait the in-flight forward's done event so freed slots are safe + to MOVE (write settled) and REUSE (read settled). The event is recorded + after the WHOLE forward, so one wait covers both hazards; drop the write-set. + """ + ev = self._latest_forward_done_event + if ev is not None: + torch.cuda.current_stream().wait_event(ev) + self._inflight_forward = None + + def _flush(self, *, urgent: bool) -> int: + """One batched compaction pass; returns the number of survivor moves. + + Pipeline (one D2H total, at step 3): + 1. `_drain_pending_reuse` — return read-settled prior srcs. + 2. sort the free list (or skip via env knob; either way ascending after). + 3. `.tolist()` snapshot → `all_cpu` *(the one sync)*. + 4-5. `_absorb_boundary_holes` — retreat past boundary-contiguous holes; + `holes_cpu` = interior holes. After this `_free_phys_pages==holes_cpu`. + 6. (urgent) `_settle_inflight_forward` — wait once so the walk is race-free. + 7. survivor walk — TWO-POINTER: move topmost live slot into the next hole, + STOPPING when the pointers cross (band packed); batch into one + `move_kv_cache` + one v2p/p2v scatter at `_commit_move_batch`. + 8-9. exit: urgent → FULL-PACK reclaim (retreat past ALL holes, empty list); + non-urgent → slice consumed dsts, merge freed srcs back. + + Two hazards per survivor (both keyed on the single `forward_done` event): + * WRITE race — forward overwrites KV[src]; a compaction read corrupts + KV[dst]. Non-urgent STOPS at such a src; urgent settles up front (step 6). + * READ race — forward READS KV[src]; src REUSE must wait the reader event. + `_commit_move_batch` routes such srcs to `_pending_reuse`; urgent's + settle makes them immediately reusable. + + `_topmost_survivor` excludes all p2v=-1 pages, so a `v_moved < 0` in the + loop is a corrupt-state bug and raises. + """ + if not self.lazy_compaction: + return 0 + self._stats_n_flush_calls += 1 + with record_function("MultiEndedAlloc._flush"): + self._drain_pending_reuse(urgent=urgent) + + # Sort ASCENDING (skip if the env knob keeps the list always-sorted). + if not _SORT_FREE_LIST_AFTER_MERGE and self._free_phys_pages.numel() > 1: + self._free_phys_pages, _ = torch.sort(self._free_phys_pages) + + all_cpu = self._free_phys_pages.tolist() # the ONE D2H sync per flush + + # `holes_cpu` = interior holes; `_free_phys_pages == holes_cpu` after. + new_wm, holes_cpu = self._absorb_boundary_holes(all_cpu) + + latest_event = self._latest_forward_done_event + + # Single-pass FULL-PACK (urgent only): the crossing-checked walk packs + # all live below the frontier so the exit can retreat past every + # interior hole at once — but only if each freed src is reuse-safe. + # `_latest_forward_done_event` is recorded after the WHOLE forward, so + # waiting it once settles BOTH hazards; then every src is event-fired + # and the walk runs race-free (empty write_set, no `_pending_reuse`). + single_pass_absorb = urgent and len(holes_cpu) > 0 + if single_pass_absorb: + self._settle_inflight_forward() + latest_event = None # reads/writes settled → srcs are fired + + # write_set: None = not yet materialized (do it inline on the first + # survivor needing the check); set() = no write race; else materialized. + write_set: Optional[Set[int]] = set() if single_pass_absorb else None + + srcs: List[int] = [] + dsts: List[int] = [] + v_moveds: List[int] = [] + + # Flush-scoped accumulator for event-FIRED srcs. `_commit_move_batch` + # appends here instead of catting onto `_free_phys_pages`; the merge is + # deferred to AFTER the trailing dst-slice, keeping `_free_phys_pages` + # byte-identical to `holes_cpu` for the whole walk. That invariant is + # what makes the directional dst-slice correct in both directions + # (catting srcs mid-flush would chop the wrong end / scramble under + # sort=ON, leaving ghost p2v=-1 pages + double-bound dsts). Event- + # PENDING srcs still route to `_pending_reuse` (read-race gating). + released_fired: List[torch.Tensor] = [] + + cursor: Optional[int] = None + j_cursor: Optional[int] = None + + # Dst cursor reads `holes_cpu` directly (no per-dst sync): grow-up from + # the front, grow-down from the back. Consumed prefix/suffix is sliced + # off in one GPU op at exit. + if self.grow_direction == "up": + dst_cursor = 0 + else: + dst_cursor = len(holes_cpu) - 1 + n_dst_consumed = 0 + + move_cap = self._lazy_max_moves_per_call if not urgent else None + + n_moves = 0 + while n_dst_consumed < len(holes_cpu): + src, j_cursor = self._topmost_survivor( + start_hint=cursor, + holes_cpu=holes_cpu, + j_in=j_cursor, + ) + if src is None: + break + + # Case A: write race. + if write_set is None: + materialized = self._materialize_inflight_write_set() + write_set = materialized if materialized is not None else set() + if write_set and src in write_set: + if urgent: + # Commit accumulated moves, then wait the forward so the + # rest of the walk is race-free. + self._commit_move_batch( + srcs, dsts, v_moveds, latest_event, released_fired + ) + n_moves += len(srcs) + srcs.clear() + dsts.clear() + v_moveds.clear() + inflight = self._inflight_forward + if inflight is not None: + torch.cuda.current_stream().wait_event(inflight[0]) + self._inflight_forward = None + write_set = set() # forward drained → no race + latest_event = None + # DO NOT reset cursor/j_cursor: rewinding would re-pick the + # just-committed srcs (now p2v=-1, not in holes_cpu) and + # trip the p2v=-1 assertion. Preserving cursor resumes at + # the blocker itself, which now passes under empty write_set. + continue + else: + break # non-urgent: top blocker stops the walk + + # Case B/C: no write race. dst from holes_cpu by cursor (no sync). + dst = holes_cpu[dst_cursor] + # Two-pointer crossing check: once src and dst cross, the band is + # packed. Moving further would shuffle a hole back toward the + # frontier and block the watermark retreat, so stop — this is what + # lets one urgent pass reclaim ALL holes (not just a contiguous run). + if (self.grow_direction == "up" and src < dst) or ( + self.grow_direction == "down" and src > dst + ): + break + if self.grow_direction == "up": + dst_cursor += 1 + else: + dst_cursor -= 1 + n_dst_consumed += 1 + + v_moved = int(self.physical_to_virtual[src].item()) + if v_moved < 0: + # `_topmost_survivor` excludes all p2v=-1 pages — corrupt state. + raise AssertionError( + f"MultiEndedAllocator({self.sub_pool_name!r})." + f"_flush: topmost survivor p={src} has p2v=-1; " + "this should be impossible (`_topmost_survivor` " + "excludes `holes_cpu` and `_pending_reuse_pages_cpu`)." + f" State: {self.allocator_state_str()}, " + f"#holes={len(holes_cpu)}, " + f"#pending_reuse={len(self._pending_reuse_pages_cpu)}" + ) + + srcs.append(src) + dsts.append(dst) + v_moveds.append(v_moved) + + # Advance cursor strictly past the picked src. + if self.grow_direction == "up": + cursor = src - 1 + else: + cursor = src + 1 + + if move_cap is not None and len(srcs) >= move_cap: + break + + self._commit_move_batch(srcs, dsts, v_moveds, latest_event, released_fired) + n_moves += len(srcs) + + if single_pass_absorb: + # FULL-PACK reclaim (urgent): all interior holes now sit above the + # frontier, so retreat past the whole lot and EMPTY the free list — + # those pages are beyond-frontier free space (reclaimed by the next + # extension), so `released_fired` is simply dropped too. + n_reclaimed = len(holes_cpu) + if self.grow_direction == "up": + self.watermark_physical = new_wm - n_reclaimed + else: + self.watermark_physical = new_wm + n_reclaimed + self._stats_n_pages_absorbed += n_reclaimed + self._free_phys_pages = self._free_phys_pages[:0] + else: + # Non-urgent partial pass: watermark stays; a later flush absorbs the + # now-top holes. `_free_phys_pages` is still == holes_cpu, so the + # consumed dsts are exactly the front (grow-up) / back (grow-down) + # `n_dst_consumed` entries; slice them, then merge freed srcs in one cat. + if n_dst_consumed > 0: + if self.grow_direction == "up": + self._free_phys_pages = self._free_phys_pages[n_dst_consumed:] + else: + self._free_phys_pages = self._free_phys_pages[:-n_dst_consumed] + if released_fired: + self._release_phys_pages_batch( + released_fired[0] + if len(released_fired) == 1 + else torch.cat(released_fired) + ) + if n_moves > 0: + self._stats_n_flush_did_work += 1 + self._stats_n_flush_moves += n_moves + self._maybe_emit_stats() + return n_moves + + def _commit_move_batch( + self, + srcs: List[int], + dsts: List[int], + v_moveds: List[int], + latest_event: Optional[torch.cuda.Event], + released_fired: List[torch.Tensor], + ) -> None: + """Issue ONE `move_kv_cache` + ONE bulk v2p/p2v remap for the accumulated + `(src, dst, v_moved)` triples. Fired srcs accumulate in `released_fired` + (merged by `_flush` AFTER its dst-slice, keeping the free list == holes_cpu); + event-pending srcs route to `_pending_reuse` (read-race gating). + """ + if not srcs: + return + with record_function("MultiEndedAlloc._commit_move_batch"): + src_pages_t = torch.tensor(srcs, dtype=torch.int64, device=self.device) + dst_pages_t = torch.tensor(dsts, dtype=torch.int64, device=self.device) + v_moveds_t = torch.tensor(v_moveds, dtype=torch.int64, device=self.device) + # Expand to token granularity (the move kernel is token-granular). + if self.page_size == 1: + src_t, dst_t = src_pages_t, dst_pages_t + else: + offsets = torch.arange( + self.page_size, + dtype=torch.int64, + device=self.device, + ) + src_t = (src_pages_t[:, None] * self.page_size + offsets).reshape(-1) + dst_t = (dst_pages_t[:, None] * self.page_size + offsets).reshape(-1) + move_fn = getattr(self._kvcache, "move_kv_cache", None) + if move_fn is not None: + move_fn(dst_t, src_t) + else: + copy_phys = getattr(self._kvcache, "_copy_from_physical", None) + assert copy_phys is not None, ( + f"sub-pool {self.sub_pool_name!r} supports neither " + "move_kv_cache nor _copy_from_physical" + ) + copy_phys(src_t, dst_t) + # ONE bulk remap (single-writer on schedule_stream). + self.virtual_to_physical[v_moveds_t] = dst_pages_t + self.physical_to_virtual[dst_pages_t] = v_moveds_t + self.physical_to_virtual[src_pages_t] = -1 + self._inverse_history.append((src_pages_t, dst_pages_t, v_moveds_t)) + # Src disposition — ONE entry per batch. `src_pages_t` is reused as the + # `_pending_reuse` GPU tensor (no second H2D at drain). + event_fired = latest_event is None or latest_event.query() + if event_fired: + released_fired.append(src_pages_t) + else: + srcs_copy: List[int] = list(srcs) # caller mutates `srcs` + self._pending_reuse[latest_event] = (srcs_copy, src_pages_t) + self._pending_reuse_pages_cpu.update(srcs_copy) + + def flush_opportunistic(self) -> int: + """Public, non-urgent flush at quiescent points; never blocks + `schedule_stream`. No-op if `lazy_compaction=False`. + + Empty-set fast-path: the scheduler triggers this very often and ~99% hit + the empty state. Skip whenever there is no possible work — no holes AND no + pending entries (the in-flight write-set only matters when compacting). + """ + with record_function("MultiEndedAlloc.flush_opportunistic"): + if not self.lazy_compaction: + return 0 + if self._free_phys_pages.numel() == 0 and not self._pending_reuse: + return 0 + return self._flush(urgent=False) + + def _raise_stale_slot_assertion(self, *, free_v, freed_p) -> None: + bad = free_v[freed_p < 0].tolist() + frames = inspect.stack()[1:9] + callers = " <- ".join(f"{f.filename.split('/')[-1]}:{f.lineno}" for f in frames) + raise AssertionError( + f"MultiEndedAllocator({self.sub_pool_name!r}).free: virtual id(s) {bad} have " + f"virtual_to_physical == -1 (double-free or never-allocated). " + f"State: {self.allocator_state_str()}. free_index unique={free_v.tolist()}. " + f"recent _inverse_history (last 3): " + f"{[(s.tolist(), d.tolist()) for s, d, _ in self._inverse_history[-3:]]}. " + f"Caller: {callers}." + ) + + # -- free-group -- + + def free_group_begin(self) -> None: + self.is_not_in_free_group = False + self.free_group = [] + + def free_group_end(self) -> None: + self.is_not_in_free_group = True + if self.free_group: + merged = torch.cat(self.free_group) + self.free_group = [] + self.free(merged) + + +class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): + """Composite allocator for the MHA (full-attn) + Mamba hybrid pair. + + The token-slot surface delegates to the full-attn side (`alloc(N)` → + MHA token slots). The Mamba sub-pool's per-request `alloc(1)` is driven + separately by `UnifiedHybridReqToTokenPool`. Both sub-allocators are id-owners + of their own (independent) virtual-id spaces. + """ + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + kvcache, # HybridLinearKVPool + device: str, + page_size: int = 1, + need_sort: bool = False, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, + ): + full_max = unified_buffer.max_slots("full") + super().__init__( + size=full_max - 1, + page_size=page_size, + dtype=unified_buffer.mha_spec("full").store_dtype, + device=device, + kvcache=kvcache, + need_sort=need_sort, + ) + self.unified_buffer = unified_buffer + self._kvcache = kvcache + self.page_size = page_size + self.lazy_compaction = lazy_compaction + + # FULL is page-aware; MAMBA stays page_size=1 (state is per-request, + # orthogonal to the full side's per-token paging). + self.full_attn_allocator = MultiEndedAllocator( + kvcache=kvcache.full_kv_pool, + unified_buffer=unified_buffer, + sub_pool_name="full", + device=device, + is_id_owner=True, + page_size=page_size, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + self.mamba_allocator = MultiEndedAllocator( + kvcache=kvcache.mamba_pool, + unified_buffer=unified_buffer, + sub_pool_name="mamba", + device=device, + is_id_owner=True, + page_size=1, # Mamba state stays slot-granular (1-per-req) + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + self.full_attn_allocator.bind_peer(self.mamba_allocator) + self.mamba_allocator.bind_peer(self.full_attn_allocator) + + # The mamba slot allocator (PHYSICAL view) is built later by + # `init_unified_mamba_pools`, which wraps `self.mamba_allocator` in a + # `UnifiedMambaSlotAllocator` owning the v2p translate; the mamba pool is a + # pure PHYSICAL store. The full-attn KV pool needs no allocator either — + # write locations are resolved in the attention metadata. + + self.is_not_in_free_group = True + self.free_group: List[torch.Tensor] = [] + # Base init left these None; we use watermark math, not free-lists. + self.free_pages = torch.empty(0, dtype=torch.int64, device=device) + self.release_pages = torch.empty(0, dtype=torch.int64, device=device) + + logger.info( + "[unified-memory-pool] UnifiedMambaTokenToKVPoolAllocator ready: " + "full max_slots=%d (min_slot_index=%d, page_size=%d, " + "num_pages=%d), mamba max_slots=%d (min_slot_index=%d), " + "full_available=%d, mamba_available=%d", + self.full_attn_allocator.max_slots, + self.full_attn_allocator.min_slot_index, + self.full_attn_allocator.page_size, + self.full_attn_allocator.num_pages, + self.mamba_allocator.max_slots, + self.mamba_allocator.min_slot_index, + self.full_attn_allocator.available_size(), + self.mamba_allocator.available_size(), + ) + + # -- size: dynamic -- + @property + def size(self) -> int: + # TOKENS. MUST use the SAME available view as `available_size()` so the + # leak invariant self-cancels (available term cancels → check reduces to + # `evictable + ... == allocated`, independent of peer-hole credit). + return ( + self.full_attn_allocator.schedulable_available_size() + + self.full_attn_allocator.allocated_count() + ) + + @size.setter + def size(self, value) -> None: + pass # base init writes here; computed dynamically + + # -- token-slot surface: MHA side -- + + # Realizable-with-compaction view so the retract gate / evict / schedule_policy + # don't over-retract when the mamba peer holds drainable holes an urgent flush + # would convert into shared-gap room. Per-side alloc gates still use the + # un-credited `available_size()` so they flush before extending. + def available_size(self) -> int: + return self.full_attn_allocator.schedulable_available_size() + + def full_available_size(self) -> int: + return self.full_attn_allocator.schedulable_available_size() + + def mamba_slot_full_token_cost(self) -> int: + """Full-token-equivalents of shared-gap bytes ONE mamba state consumes. + + full and mamba share one byte buffer, so a mamba slot removes that many + full-KV tokens from the gap; the prefill planner reserves this so admission + stays inside the JOINT budget. = mamba bytes/slot ÷ full bytes/token, rounded + UP (conservative). Only on the shared composite (non-shared pools are separate, + so the planner sources this via `getattr(..., None)`). + """ + return -( + -self.mamba_allocator.entry_bytes_per_page + // self.full_attn_allocator.entry_bytes + ) + + @property + def size_full(self) -> int: + return self.full_attn_allocator.max_slots - 1 + + @property + def size_mamba(self) -> int: + return self.mamba_allocator.max_slots - 1 + + def debug_print(self) -> str: + return ( + f"#full-available={self.full_attn_allocator.available_size()}, " + f"#mamba-available={self.mamba_allocator.available_size()}" + ) + + def get_kvcache(self): + return self._kvcache + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + with record_function("UnifiedMambaAlloc.alloc"): + return self.full_attn_allocator.alloc(need_size) + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + num_new_pages: Optional[int] = None, + ) -> Optional[torch.Tensor]: + """Paged extend. Mamba state is per-request (doesn't advance per-token), + so forward only to the full sub-allocator.""" + with record_function("UnifiedMambaAlloc.alloc_extend"): + return self.full_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_new_pages, + ) + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Paged decode. Mamba side stays untouched per-decode.""" + with record_function("UnifiedMambaAlloc.alloc_decode"): + return self.full_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + + def translate_kv_loc( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-pool virtual TOKEN ids -> physical TOKEN ids. Delegates to the + full-side sub-allocator. Supports ``out=`` for cuda-graph buffer stability. + `-1` inputs map to `-1` (treated as padding downstream). + """ + result = self.full_attn_allocator.translate_kv_loc(loc, out=out) + return result + + def is_slot_allocated(self, slot: int) -> bool: + return self.full_attn_allocator.is_slot_allocated(slot) + + def allocator_state_str(self) -> str: + return self.full_attn_allocator.allocator_state_str() + + def free(self, free_index: torch.Tensor) -> None: + with record_function("UnifiedMambaAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if not self.is_not_in_free_group: + self.free_group.append(free_index) + return + self.full_attn_allocator.free(free_index) + self.full_attn_allocator.clear_inverse_history() + self.mamba_allocator.clear_inverse_history() + + def free_group_begin(self) -> None: + self.is_not_in_free_group = False + self.free_group = [] + + def free_group_end(self) -> None: + self.is_not_in_free_group = True + if self.free_group: + merged = torch.cat(self.free_group) + self.free_group = [] + self.full_attn_allocator.free(merged) + self.full_attn_allocator.clear_inverse_history() + self.mamba_allocator.clear_inverse_history() + + def backup_state(self): + return [ + self.full_attn_allocator.backup_state(), + self.mamba_allocator.backup_state(), + ] + + def restore_state(self, state): + assert len(state) == 2 + full_rollback = self.full_attn_allocator.restore_state(state[0]) + mamba_rollback = self.mamba_allocator.restore_state(state[1]) + return full_rollback + mamba_rollback + + def clear(self) -> None: + self.full_attn_allocator.clear() + self.mamba_allocator.clear() + self.is_not_in_free_group = True + self.free_group = [] + + # -- Lazy compaction hooks -- + + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" + with record_function("UnifiedMambaAlloc.set_latest_forward_done_event"): + self.full_attn_allocator.set_latest_forward_done_event(event) + self.mamba_allocator.set_latest_forward_done_event(event) + + def set_inflight_forward( + self, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + """Hand the forward's metadata to BOTH sub-pools. Full derives its write-set + from `out_cache_loc`; the Mamba state pool isn't written via `out_cache_loc` + (mamba kernels, not `set_kv_buffer`), so it gets `None`. + """ + with record_function("UnifiedMambaAlloc.set_inflight_forward"): + self.full_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + self.mamba_allocator.set_inflight_forward(forward_done, None) + + def flush_opportunistic(self) -> int: + """Non-urgent flush of BOTH sub-allocators; sync-free. Composite empty-set + fast-path skips both calls when neither side has work. + """ + with record_function("UnifiedMambaAlloc.flush_opportunistic"): + fa = self.full_attn_allocator + ma = self.mamba_allocator + if ( + fa._free_phys_pages.numel() == 0 + and not fa._pending_reuse + and ma._free_phys_pages.numel() == 0 + and not ma._pending_reuse + ): + return 0 + return fa.flush_opportunistic() + ma.flush_opportunistic() + + +class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): + """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). + + Inherits from `SWATokenToKVPoolAllocator` only for the isinstance contract; + we call grand-parent `BaseTokenToKVPoolAllocator.__init__` directly to skip + the parent's static-partition sub-pool allocation (which unified-memory-pool + replaces). + + Capacity views: + - `available_size()`: joint byte-budget, the only safe `alloc(N)` pre-check + (N slots cost N*(entry_full + entry_swa) shared-gap bytes). + - `_conserve_*`: slot-conservation, for the LEAK invariant only. + - `schedulable_*`: byte-coordinated, realizable-with-compaction. + - `full_available_size()` / `swa_available_size()`: per-side scheduler view + = min(conserve, schedulable). + """ + + # Parent's `size` property has no setter but base init does `self.size = size`; + # override with a no-op setter. Reading returns `min(_size_full, _size_swa)`. + @property + def size(self) -> int: + return min(self._size_full, self._size_swa) + + @size.setter + def size(self, value) -> None: + pass + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + kvcache, # UnifiedSWAKVPool + device: str, + full_max_total_num_tokens: int, + swa_max_total_num_tokens: int, + page_size: int = 1, + need_sort: bool = False, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, + ): + # Set _size_full / _size_swa BEFORE base init (read during it). STATIC + # partition caps — the slot-conservation value the leak invariant expects. + self._size_full = full_max_total_num_tokens + self._size_swa = swa_max_total_num_tokens + self._full_max_total_num_tokens = full_max_total_num_tokens + self._swa_max_total_num_tokens = swa_max_total_num_tokens + self.page_size = page_size + + # Skip SWATokenToKVPoolAllocator.__init__; call grand-parent base init + # directly (its `self.size = size` is absorbed by our no-op setter). + BaseTokenToKVPoolAllocator.__init__( + self, + size=full_max_total_num_tokens, + page_size=page_size, + dtype=unified_buffer.mha_spec("full").store_dtype, + device=device, + kvcache=kvcache, + need_sort=need_sort, + ) + self.unified_buffer = unified_buffer + self._kvcache = kvcache + self.lazy_compaction = lazy_compaction + + self.full_attn_allocator = MultiEndedAllocator( + kvcache=kvcache.full_kv_pool, + unified_buffer=unified_buffer, + sub_pool_name="full", + device=device, + is_id_owner=True, + page_size=page_size, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + self.swa_attn_allocator = MultiEndedAllocator( + kvcache=kvcache.swa_kv_pool, + unified_buffer=unified_buffer, + sub_pool_name="swa", + device=device, + is_id_owner=False, # non-owner; consumes virtuals minted by full + page_size=page_size, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + self.full_attn_allocator.bind_peer(self.swa_attn_allocator) + self.swa_attn_allocator.bind_peer(self.full_attn_allocator) + + # The full/SWA KV pools need no allocator wiring (write locations resolved + # in attention metadata); the composite keeps allocators for read-path translates. + kvcache.attach_allocators( + full_allocator=self.full_attn_allocator, + swa_allocator=self.swa_attn_allocator, + ) + + self.is_not_in_free_group = True + self.free_group: List[torch.Tensor] = [] + # Empty (not None) for the leak checker. + self.free_pages = torch.empty(0, dtype=torch.int64, device=device) + self.release_pages = torch.empty(0, dtype=torch.int64, device=device) + + logger.info( + "[unified-memory-pool] UnifiedSWATokenToKVPoolAllocator ready: " + "full max_slots=%d (min_slot_index=%d, entry_bytes=%d), " + "swa max_slots=%d (min_slot_index=%d, entry_bytes=%d), " + "static caps full=%d swa=%d, joint available=%d", + self.full_attn_allocator.max_slots, + self.full_attn_allocator.min_slot_index, + self.full_attn_allocator.entry_bytes, + self.swa_attn_allocator.max_slots, + self.swa_attn_allocator.min_slot_index, + self.swa_attn_allocator.entry_bytes, + self._full_max_total_num_tokens, + self._swa_max_total_num_tokens, + self.available_size(), + ) + + # -- capacity reporting (three-way split) -- + + def available_size(self) -> int: + """Tokens available for `alloc(N)` / `alloc_extend(N)` (TOKENS). + + Joint byte-budget: each composite alloc(1) consumes one full-side AND one + swa-side page (same virtual id). The 3-phase lazy formula consumes both + sides' holes maximally before extending toward the gap (H_f/H_s = holes, + e_f/e_s = bytes/page, R_f/R_s = extension room, G = byte gap): + Phase 1 (both drain, free): K1 = min(H_f, H_s) + Phase 2 (fewer-holes side extends): K2 limited by remaining holes & G + Phase 3 (both extend): K3 = G // (e_f + e_s) + Total capped by index-space rooms (H_f + R_f, H_s + R_s). ps==1 collapses + to slot math. Eager has no holes → original joint formula. + """ + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + e_f = fa.entry_bytes_per_page + e_s = sa.entry_bytes_per_page + # Direction-agnostic shared gap: the free byte band between the two pools. + if fa.grow_direction == "up": + gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) + else: + gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) + R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() + R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() + + if not self.lazy_compaction: + pages_by_bytes = gap_bytes // (e_f + e_s) + return min(pages_by_bytes, R_f, R_s) * self.page_size + + H_f = len(fa._free_phys_pages) + H_s = len(sa._free_phys_pages) + + K1 = min(H_f, H_s) # Phase 1: both drain + + # Phase 2: fewer-holes side extends; more-holes side keeps draining. + if H_f <= H_s: + e_phase2 = e_f + K_phase2_max = H_s + else: + e_phase2 = e_s + K_phase2_max = H_f + K2_room = K_phase2_max - K1 + K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room + gap_bytes -= K2 * e_phase2 + + K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend + + K_total = K1 + K2 + K3 + K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps + return K_total * self.page_size + + # Slot-conservation views — the ONLY views the leak invariant should see + # (returning the byte-coordinated value would flag spurious leaks). + # `allocated_count()` is in TOKENS (the unit the leak check expects). + def _conserve_full_available_size(self) -> int: + return ( + self._full_max_total_num_tokens - self.full_attn_allocator.allocated_count() + ) + + def _conserve_swa_available_size(self) -> int: + return ( + self._swa_max_total_num_tokens - self.swa_attn_allocator.allocated_count() + ) + + # PHYSICAL per-side views read by scheduling / eviction consumers. The + # `min(...)` is sound under dynamic borrowing: the static-conserve cap bounds + # the lending side, the byte-coordinated `schedulable_*` bounds the side that + # has grown into the shared gap; whichever is tighter wins. + def full_available_size(self) -> int: + return min( + self._conserve_full_available_size(), + self.schedulable_full_available_size(), + ) + + def swa_available_size(self) -> int: + return min( + self._conserve_swa_available_size(), + self.schedulable_swa_available_size(), + ) + + # Byte-coordinated, realizable-with-compaction views (peer drainable holes + # credited — see `MultiEndedAllocator.schedulable_available_size`). + def schedulable_full_available_size(self) -> int: + return self.full_attn_allocator.schedulable_available_size() + + def schedulable_swa_available_size(self) -> int: + return self.swa_attn_allocator.schedulable_available_size() + + def _flush_both_for_alloc(self, need_tokens: int) -> bool: + """SWA analogue of `_flush_peer_for_alloc`. Each composite alloc consumes a + full AND a swa page and either side's compaction opens gap for the other, + so flush BOTH (one urgent pass each). + """ + if not self.lazy_compaction: + return need_tokens <= self.available_size() + self.full_attn_allocator._flush(urgent=True) + self.swa_attn_allocator._flush(urgent=True) + return need_tokens <= self.available_size() + + # `size_full` / `size_swa` are inherited; they read `_size_full`/`_size_swa` + # (set to the static caps). We do NOT report `max_slots - 1`: under unified + # memory pool that ~= full_max + swa_max and would over-promise. + + def debug_print(self) -> str: + return ( + f"#full-available={self.full_attn_allocator.available_size()}, " + f"#swa-available={self.swa_attn_allocator.available_size()}, " + f"#joint-available={self.available_size()}" + ) + + def get_kvcache(self): + return self._kvcache + + def translate_kv_loc( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. + Delegates to the full-side sub-allocator. Supports ``out=`` for cuda-graph. + """ + result = self.full_attn_allocator.translate_kv_loc(loc, out=out) + return result + + def translate_loc_from_full_to_swa( + self, + kv_indices: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """SWA-layer read path: virtual TOKEN ids -> swa-physical TOKEN ids (int32, + matching the non-shared API). Page math against the swa side's v2p table. + Supports ``out=`` (int32, same shape) for cuda-graph buffer stability. + """ + if out is not None: + assert out.dtype == torch.int32, ( + f"translate_loc_from_full_to_swa: out= dtype must be int32 " + f"(matches SWA Triton kernel contract), got {out.dtype}" + ) + assert out.shape == kv_indices.shape, ( + f"translate_loc_from_full_to_swa: out= shape " + f"{tuple(out.shape)} must match kv_indices shape " + f"{tuple(kv_indices.shape)}" + ) + # Tombstone-safety clamp (mirrors the full-side clamp): tombstoned (-1) + # v2p_swa entries must not reach `swa_k_buffer[-1]` (illegal under replay). + # Clamp to 0 routes them to the reserved padding sink (slot 0). + if self.swa_attn_allocator.page_size == 1: + if out is not None: + # Gather into a transient int64, then cast into out (`out.copy_`). + tmp = torch.index_select( + self.swa_attn_allocator.virtual_to_physical, 0, kv_indices + ) + tmp = torch.clamp_min(tmp, 0) + out.copy_(tmp.to(torch.int32)) + return out + result = self.swa_attn_allocator.virtual_to_physical[kv_indices] + result = torch.clamp_min(result, 0) + return result.to(torch.int32) + ps = self.swa_attn_allocator.page_size + virt_pages = kv_indices // ps + offsets = kv_indices % ps + swa_phys_pages = self.swa_attn_allocator.virtual_to_physical[virt_pages] + result = (swa_phys_pages * ps + offsets).to(torch.int32) + result = torch.clamp_min(result, 0) + if out is not None: + out.copy_(result) + return out + return result + + # -- alloc -- + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + with record_function("UnifiedSWAAlloc.alloc"): + # Joint pre-check. Both sides are mutual peers (each side's compaction + # opens gap for the other), so flush BOTH on shortfall. + if need_size > self.available_size(): + if not self._flush_both_for_alloc(need_size): + return None + # Snapshot the virtual PAGES full will consume, to bind them on swa too. + num_pages = need_size // self.page_size + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_pages].clone() + + v_tokens = fa.alloc(need_size) + # Post-pre-check failure can only be internal-state inconsistency. + assert v_tokens is not None, ( + "UnifiedSWA.alloc: full.alloc returned None after joint " + "pre-check passed — internal-state inconsistency" + ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return v_tokens + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + ) -> Optional[torch.Tensor]: + """Paged extend. Runs the kernel ONCE in virtual space, then binds the + consumed virtual PAGES on the swa side via `alloc_with_virtual`. Returns + virtual TOKEN ids respecting the tail-page-reuse contract and the + cross-sub-pool identity (same virtual page maps to full- and swa-physical). + """ + with record_function("UnifiedSWAAlloc.alloc_extend"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + need_tokens = num_new_pages * self.page_size + if need_tokens > self.available_size(): + if not self._flush_both_for_alloc(need_tokens): + return None + + # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps + # its view after the slice is consumed. + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + + out_indices = fa.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_new_pages, + ) + assert out_indices is not None, ( + "UnifiedSWA.alloc_extend: full.alloc_extend returned None " + "after joint pre-check passed — internal-state inconsistency" + ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return out_indices # virtual TOKEN ids + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Paged decode. One new token per request (a page is consumed iff the + decode wraps). Same one-kernel-in-virtual-space discipline as ``alloc_extend``. + """ + with record_function("UnifiedSWAAlloc.alloc_decode"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True + ) + need_tokens = num_new_pages * self.page_size + if need_tokens > self.available_size(): + if not self._flush_both_for_alloc(need_tokens): + return None + + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + + out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc) + assert out_indices is not None, ( + "UnifiedSWA.alloc_decode: full.alloc_decode returned None " + "after joint pre-check passed — internal-state inconsistency" + ) + + if new_virtual_pages.numel() > 0: + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + + return out_indices # virtual TOKEN ids + + def is_slot_allocated(self, slot: int) -> bool: + """Token-slot surface = the full side (which owns the virtual ids).""" + return self.full_attn_allocator.is_slot_allocated(slot) + + def allocator_state_str(self) -> str: + return self.full_attn_allocator.allocator_state_str() + + # -- free -- + + def free(self, free_index: torch.Tensor) -> None: + with record_function("UnifiedSWAAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if not self.is_not_in_free_group: + self.free_group.append(free_index) + return + # Free both peers; the per-sub-pool v2p IS the mapping, so order isn't + # load-bearing. Filter the swa side to skip already-tombstoned virtuals + # (`swa.v2p_page == -1` from an earlier `free_swa`); the full side needs + # no filter (it's the lifecycle owner, so every value is still bound). + v = free_index.detach().to(torch.int64) + v_pages = v // self.page_size + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped. + live_token_mask = swa_v2p_pages > 0 + live_tokens = v[live_token_mask] + if live_tokens.numel() > 0: + self.swa_attn_allocator.free(live_tokens) + self.full_attn_allocator.free(v) + self.full_attn_allocator.clear_inverse_history() + self.swa_attn_allocator.clear_inverse_history() + + def free_swa(self, free_index: torch.Tensor) -> None: + """SWA tombstone path: release swa-physical, leave virtual id and + full-physical live. Called by `SWARadixCache._evict_swa_only` when a node + ages past the sliding-window horizon. `swa.v2p_page[v_page] = -1` IS the + tombstone. + """ + if free_index is None or free_index.numel() == 0: + return + # Keep only tokens whose virtual PAGE is still bound on swa (calling + # `swa.free` on an already-tombstoned one would assert). + v = free_index.detach().to(torch.int64) + v_pages = v // self.page_size + # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + live = v[swa_v2p_pages > 0] + if live.numel() == 0: + return + self.swa_attn_allocator.free(live) + self.swa_attn_allocator.clear_inverse_history() + + def set_full_to_swa_mapping( + self, full_indices: torch.Tensor, swa_indices: torch.Tensor + ) -> None: + """No-op stub for HiCache load-back compatibility. In shared mode there is + no mapping tensor (the swa v2p IS the mapping); HiCache for shared SWA is + out of scope. + """ + return + + # -- free-group -- + + def free_group_begin(self) -> None: + self.is_not_in_free_group = False + self.free_group = [] + + def free_group_end(self) -> None: + self.is_not_in_free_group = True + if self.free_group: + merged = torch.cat(self.free_group) + self.free_group = [] + self.free(merged) + + # -- spec-decode hooks (asserted off; preserved for future use) -- + + def backup_state(self): + return [ + self.full_attn_allocator.backup_state(), + self.swa_attn_allocator.backup_state(), + ] + + def restore_state(self, state): + assert len(state) == 2 + full_rollback = self.full_attn_allocator.restore_state(state[0]) + swa_rollback = self.swa_attn_allocator.restore_state(state[1]) + return full_rollback + swa_rollback + + def clear(self) -> None: + self.full_attn_allocator.clear() + self.swa_attn_allocator.clear() + self.is_not_in_free_group = True + self.free_group = [] + + # -- Lazy compaction hooks -- + + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" + with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): + self.full_attn_allocator.set_latest_forward_done_event(event) + self.swa_attn_allocator.set_latest_forward_done_event(event) + + def set_inflight_forward( + self, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + """Hand the forward's metadata to BOTH sub-pools. Each materializes its + write-set via its OWN v2p; the forward writes both sides per new token, + so both get a non-empty in-flight tensor. + """ + with record_function("UnifiedSWAAlloc.set_inflight_forward"): + self.full_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + self.swa_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + + def flush_opportunistic(self) -> int: + """Non-urgent flush of BOTH sub-allocators; sync-free. Composite empty-set + fast-path skips both calls when neither side has work. + """ + with record_function("UnifiedSWAAlloc.flush_opportunistic"): + fa = self.full_attn_allocator + sa = self.swa_attn_allocator + if ( + fa._free_phys_pages.numel() == 0 + and not fa._pending_reuse + and sa._free_phys_pages.numel() == 0 + and not sa._pending_reuse + ): + return 0 + return fa.flush_opportunistic() + sa.flush_opportunistic() diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index e0f8e0028..c56b4238f 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -160,7 +160,7 @@ class SWAKVPool(BaseSWAKVPool): v_scale: float = 1.0, ): # loc_info bundles the full loc and the pre-translated SWA loc. - loc, swa_loc = unwrap_write_loc(loc_info) + loc, swa_loc, _ = unwrap_write_loc(loc_info) layer_id = layer.layer_id layer_id_pool, is_swa_layer = self.layers_mapping[layer_id] if is_swa_layer: diff --git a/python/sglang/srt/mem_cache/triton_ops/virtual_slot.py b/python/sglang/srt/mem_cache/triton_ops/virtual_slot.py new file mode 100644 index 000000000..b54c2da97 --- /dev/null +++ b/python/sglang/srt/mem_cache/triton_ops/virtual_slot.py @@ -0,0 +1,96 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Virtual<->physical slot Triton kernels for the unified memory pool.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +# Fused take-physical-pages + bind for the alloc fast path. Invoked ONLY when +# `_hole_count == 0`; otherwise the slow path drains holes first (Invariant B, +# greedy hole reuse). Caller advances `watermark_physical` and checks overflow +# BEFORE launch, passing the PRE-extension watermark. Cuda-graph safe (no +# `.item()`, no tensor branching); runs on the scheduler thread. + + +@triton.jit +def alloc_bind_inplace_kernel( + v_pages_ptr, # in: [N] int64 — virtual page ids + v2p_ptr, # in/out: int64 — virtual_to_physical table + p2v_ptr, # in/out: int64 — physical_to_virtual table + out_phys_ptr, # out: [N] int64 — physical page ids + N, # runtime: number of pages to allocate + start_phys, # runtime: lowest physical page id in the new range + BLOCK: tl.constexpr, +): + """Fused: ascending arange + out_phys/v2p/p2v scatter. + + Caller pre-adjusts `start_phys` per direction so the range is always + ascending (grow-up: start_wm; grow-down: start_wm - N + 1), making the + v->p mapping byte-identical to the `torch.arange` slow path. + """ + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < N + + phys = (start_phys + offs).to(tl.int64) + v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64) + + # Masked stores skip out-of-range lanes, and `other=0` keeps us off the + # v2p[0]/p2v[0] padding-sink slot. + tl.store(out_phys_ptr + offs, phys, mask=mask) + tl.store(v2p_ptr + v, phys, mask=mask) + tl.store(p2v_ptr + phys, v, mask=mask) + + +ALLOC_BIND_BLOCK = 128 + + +def alloc_bind_inplace( + v_pages: torch.Tensor, + v2p: torch.Tensor, + p2v: torch.Tensor, + start_phys: int, +) -> torch.Tensor: + """Allocate N ascending physical pages from `start_phys` and bind to `v_pages`. + + Caller must advance `watermark_physical` by N and verify overflow BEFORE + calling; this launcher does neither. + """ + N = int(v_pages.numel()) + if N == 0: + return torch.empty(0, dtype=torch.int64, device=v_pages.device) + if not v_pages.is_cuda: + # Pure-torch CPU reference for the CUDA-only kernel. + phys_pages = torch.arange( + start_phys, start_phys + N, dtype=torch.int64, device=v_pages.device + ) + v = v_pages.to(torch.int64) + v2p[v] = phys_pages + p2v[phys_pages] = v + return phys_pages + phys_pages = torch.empty(N, dtype=torch.int64, device=v_pages.device) + grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),) + alloc_bind_inplace_kernel[grid]( + v_pages, + v2p, + p2v, + phys_pages, + N, + start_phys, + BLOCK=ALLOC_BIND_BLOCK, + ) + return phys_pages diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 14ee54867..48ec5e17d 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -333,8 +333,12 @@ class MambaComponent(TreeComponent): ) else: mamba_value_donated = self._alloc_mamba_slot() + # mamba_pool is a pure PHYSICAL store; translate both slot ids + # virtual->physical (identity for the non-unified memory pool) first. + translate = self.cache.req_to_token_pool.translate_mamba_indices self.cache.req_to_token_pool.mamba_pool.copy_from( - req.mamba_pool_idx.unsqueeze(0), mamba_value_donated + translate(req.mamba_pool_idx.unsqueeze(0)), + translate(mamba_value_donated), ) insert_params.mamba_value = mamba_value_donated return cache_len diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py new file mode 100644 index 000000000..bb02e9096 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -0,0 +1,1369 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""UnifiedKVPool — one physical `uint8` byte buffer shared by 2 sub-pools. + +Two `MultiEndedAllocator`s grow from opposite ends; eager-compacting `free` +keeps each pool's byte range hole-free. Layout is envelope-major (a slot's data +for all its layers in one contiguous byte envelope) so a freed slot vacates a +region the peer can grow into. Everything above the allocator stores virtual +slot IDs; the allocator owns the per-sub-pool virtual<->physical tables and +compaction only mutates those (no reference rewriting). +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Dict, List, NamedTuple, Optional, Tuple + +import torch +import triton +from torch.profiler import record_function + +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.mem_cache.layout.page_major import ( + build_page_major_mamba_views, + build_page_major_mha_views, +) +from sglang.srt.mem_cache.memory_pool import ( + HybridReqToTokenPool, + MambaPool, + MHATokenToKVPool, + move_kv_cache_native, + unwrap_write_loc, +) +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d_kernel +from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = logging.getLogger(__name__) + +GB = 1024 * 1024 * 1024 + + +def _prod(iterable) -> int: + out = 1 + for x in iterable: + out *= int(x) + return out + + +def _store_dtype_for(kv_cache_dtype: torch.dtype) -> torch.dtype: + if kv_cache_dtype in (torch.float8_e5m2, torch.float8_e4m3fn): + return torch.uint8 + return kv_cache_dtype + + +@dataclass(frozen=True, kw_only=True) +class SubPoolSpec(ABC): + """Abstract per-slot layout of one sub-pool in a `UnifiedKVPool`.""" + + name: str + layer_num: int + grow_direction: str # "up" | "down" + + def __post_init__(self): + assert self.grow_direction in ( + "up", + "down", + ), f"grow_direction must be 'up' or 'down'; got {self.grow_direction!r}" + assert self.layer_num > 0, f"layer_num must be positive; got {self.layer_num}" + + @abstractmethod + def entry_bytes(self) -> int: + """Bytes for one slot across all `layer_num` layers.""" + raise NotImplementedError + + @abstractmethod + def get_dtype(self) -> torch.dtype: + """Storage dtype (informational). Multi-dtype subclasses return the dominant buffer's.""" + raise NotImplementedError + + +@dataclass(frozen=True, kw_only=True) +class MHASubPoolSpec(SubPoolSpec): + """Per-slot layout of one MHA-shaped sub-pool. `v_head_dim` defaults to `head_dim`.""" + + head_num: int + head_dim: int + store_dtype: torch.dtype + v_head_dim: Optional[int] = None + + def __post_init__(self): + super().__post_init__() + assert self.head_num > 0, f"head_num must be positive; got {self.head_num}" + assert self.head_dim > 0, f"head_dim must be positive; got {self.head_dim}" + if self.v_head_dim is None: + object.__setattr__(self, "v_head_dim", self.head_dim) + assert ( + self.v_head_dim > 0 + ), f"v_head_dim must be positive; got {self.v_head_dim}" + + def k_row_bytes(self) -> int: + return self.head_num * self.head_dim * self.store_dtype.itemsize + + def v_row_bytes(self) -> int: + return self.head_num * self.v_head_dim * self.store_dtype.itemsize + + def entry_bytes(self) -> int: + return self.layer_num * (self.k_row_bytes() + self.v_row_bytes()) + + # Page-major byte math: within a page block K/V group per layer + # [L0_K*ps | L0_V*ps | L1_K*ps | ...]; at ps==1 this collapses to the per-slot envelope. + + def page_bytes(self, page_size: int) -> int: + return page_size * self.entry_bytes() + + def layer_k_offset_in_page(self, layer_id: int, page_size: int) -> int: + return layer_id * page_size * (self.k_row_bytes() + self.v_row_bytes()) + + def layer_v_offset_in_page(self, layer_id: int, page_size: int) -> int: + return ( + self.layer_k_offset_in_page(layer_id, page_size) + + page_size * self.k_row_bytes() + ) + + def get_dtype(self) -> torch.dtype: + return self.store_dtype + + +@dataclass(frozen=True, kw_only=True) +class MambaSubPoolSpec(SubPoolSpec): + """Per-slot layout of one Mamba-shaped sub-pool.""" + + conv_state_shapes: Tuple[Tuple[int, ...], ...] # one shape per conv tensor + conv_dtype: torch.dtype + temporal_state_shape: Tuple[int, ...] + temporal_dtype: torch.dtype + + def __post_init__(self): + super().__post_init__() + assert len(self.conv_state_shapes) > 0, "conv_state_shapes must be non-empty" + + def conv_row_bytes(self, idx: int) -> int: + return _prod(self.conv_state_shapes[idx]) * self.conv_dtype.itemsize + + def temporal_row_bytes(self) -> int: + return _prod(self.temporal_state_shape) * self.temporal_dtype.itemsize + + def entry_bytes(self) -> int: + total = 0 + for i in range(len(self.conv_state_shapes)): + total += self.layer_num * self.conv_row_bytes(i) + total += self.layer_num * self.temporal_row_bytes() + return total + + def get_dtype(self) -> torch.dtype: + return self.conv_dtype # representative state dtype; matches MambaPool.dtype + + +# --------------------------------------------------------------------------- +# UnifiedKVPool — the byte buffer + the strided per-sub-pool views +# --------------------------------------------------------------------------- + + +class UnifiedKVPool: + """One physical `uint8` byte buffer shared by 2 sub-pools, each exposing + strided per-layer views. Allocators keep byte ranges disjoint; no usage tracking here. + """ + + def __init__( + self, + *, + total_bytes: int, + sub_pool_specs: List[SubPoolSpec], + device: str, + enable_memory_saver: bool, + page_size: int = 1, + ): + assert page_size >= 1, f"page_size must be >= 1; got {page_size}" + assert len(sub_pool_specs) == 2, ( + f"UnifiedKVPool currently supports exactly 2 sub-pools; got " + f"{len(sub_pool_specs)} (N>2 is not yet implemented)" + ) + names = [s.name for s in sub_pool_specs] + assert len(set(names)) == 2, f"sub-pool names must be unique; got {names}" + directions = sorted(s.grow_direction for s in sub_pool_specs) + assert directions == ["down", "up"], ( + f"UnifiedKVPool needs one grow-up and one grow-down sub-pool; " + f"got {directions}" + ) + + self.device = device + self.total_bytes = total_bytes + self.sub_pool_specs = sub_pool_specs + self._page_size = page_size + self._specs_by_name: Dict[str, SubPoolSpec] = { + s.name: s for s in sub_pool_specs + } + + self.memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + self._raw = torch.empty(total_bytes, dtype=torch.uint8, device=device) + self._raw.zero_() # unset slots must read as zeros (matches non-shared) + + self._max_slots: Dict[str, int] = {} + self._anchor_bytes: Dict[str, int] = {} + self._min_slot_index: Dict[str, int] = {} + # MHA: (k_buffer, v_buffer); Mamba: (conv_state_list, temporal_state) + self._mha_views: Dict[str, Tuple[List[torch.Tensor], List[torch.Tensor]]] = {} + self._mamba_views: Dict[str, Tuple[List[torch.Tensor], torch.Tensor]] = {} + + # Slot-0 dummy writes for both pools land in [0, entry_max); each pool's + # first allocatable slot is chosen so real data starts at >= entry_max. + entry_max = max(s.entry_bytes() for s in sub_pool_specs) + + for spec in sub_pool_specs: + entry_bytes = spec.entry_bytes() + max_slots = total_bytes // entry_bytes + min_slot_index = (entry_max + entry_bytes - 1) // entry_bytes # ceil + if max_slots <= min_slot_index: + raise RuntimeError( + f"UnifiedKVPool: sub-pool {spec.name!r} fits only {max_slots} " + f"slots in {total_bytes} bytes, but min_slot_index={min_slot_index} " + f"leaves no room for real data. Increase total_bytes." + ) + anchor = 0 + self._max_slots[spec.name] = max_slots + self._anchor_bytes[spec.name] = anchor + self._min_slot_index[spec.name] = min_slot_index + if isinstance(spec, MHASubPoolSpec): + self._mha_views[spec.name] = self._build_mha_views( + spec, + anchor, + max_slots, + page_size=page_size, + ) + elif isinstance(spec, MambaSubPoolSpec): + self._mamba_views[spec.name] = self._build_mamba_views( + spec, anchor, max_slots + ) + else: # pragma: no cover + raise TypeError(f"unsupported SubPoolSpec type: {type(spec)}") + + logger.info( + "[unified-memory-pool] UnifiedKVPool allocated: total_bytes=%.2f GB (=%d B), " + "%d sub-pool(s)", + total_bytes / GB, + total_bytes, + len(sub_pool_specs), + ) + for s in sub_pool_specs: + logger.info( + "[unified-memory-pool] sub-pool %r: kind=%s, layer_num=%d, grow=%s, " + "entry_bytes=%d, max_slots=%d, min_slot_index=%d (slots [0,%d) reserved)", + s.name, + type(s).__name__, + s.layer_num, + s.grow_direction, + s.entry_bytes(), + self._max_slots[s.name], + self._min_slot_index[s.name], + self._min_slot_index[s.name], + ) + + # -- introspection -- + + def spec(self, name: str) -> SubPoolSpec: + return self._specs_by_name[name] + + def mha_spec(self, name: str) -> MHASubPoolSpec: + s = self._specs_by_name[name] + assert isinstance( + s, MHASubPoolSpec + ), f"sub-pool {name!r} is {type(s).__name__}, expected MHASubPoolSpec" + return s + + def mamba_spec(self, name: str) -> MambaSubPoolSpec: + s = self._specs_by_name[name] + assert isinstance( + s, MambaSubPoolSpec + ), f"sub-pool {name!r} is {type(s).__name__}, expected MambaSubPoolSpec" + return s + + def max_slots(self, name: str) -> int: + return self._max_slots[name] + + def min_slot_index(self, name: str) -> int: + return self._min_slot_index[name] + + def anchor_bytes(self, name: str) -> int: + anchor = self._anchor_bytes[name] + assert anchor == 0, f"current design assumes all anchors are 0; got {anchor}" + return anchor + + def mha_views_for(self, name: str) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + return self._mha_views[name] + + def mamba_views_for(self, name: str) -> Tuple[List[torch.Tensor], torch.Tensor]: + return self._mamba_views[name] + + def _build_mha_views( + self, + spec: MHASubPoolSpec, + anchor_bytes: int, + max_slots: int, + page_size: int, + ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + return build_page_major_mha_views( + self._raw, + layer_num=spec.layer_num, + head_num=spec.head_num, + head_dim=spec.head_dim, + v_head_dim=spec.v_head_dim, + store_dtype=spec.store_dtype, + page_size=page_size, + num_pages=max_slots // page_size, + anchor_bytes=anchor_bytes, + ) + + def _build_mamba_views( + self, spec: MambaSubPoolSpec, anchor_bytes: int, max_slots: int + ) -> Tuple[List[torch.Tensor], torch.Tensor]: + return build_page_major_mamba_views( + self._raw, + layer_num=spec.layer_num, + conv_state_shapes=spec.conv_state_shapes, + conv_dtype=spec.conv_dtype, + temporal_state_shape=spec.temporal_state_shape, + temporal_dtype=spec.temporal_dtype, + max_slots=max_slots, + anchor_bytes=anchor_bytes, + ) + + +class UnifiedMHATokenToKVPool(MHATokenToKVPool): + """MHA KV pool whose `k_buffer`/`v_buffer` are strided views into a `UnifiedKVPool`. + + Relocation uses the native move (strided views break the tiled Triton kernel that + assumes stride == row bytes). `set_kv_buffer` gets PHYSICAL slot ids; never translates. + """ + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + sub_pool_name: str, + page_size: int = 1, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + enable_alt_stream: bool = True, + ): + spec = unified_buffer.mha_spec(sub_pool_name) + k_buffer, v_buffer = unified_buffer.mha_views_for(sub_pool_name) + max_slots = unified_buffer.max_slots(sub_pool_name) + + self._unified_buffer = unified_buffer + self._sub_pool_name = sub_pool_name + self._k_views = k_buffer + self._v_views = v_buffer + self._page_size = page_size + + super().__init__( + size=max_slots - 1, # -1 for reserved slot 0 + page_size=page_size, + dtype=spec.store_dtype, + head_num=spec.head_num, + head_dim=spec.head_dim, + layer_num=spec.layer_num, + device=unified_buffer.device, + enable_memory_saver=False, # buffer owned by UnifiedKVPool + v_head_dim=spec.v_head_dim, + start_layer=start_layer, + end_layer=end_layer, + enable_alt_stream=enable_alt_stream, + enable_kv_cache_copy=False, # strided views — force native move + ) + + def _create_buffers(self): + self.k_buffer = self._k_views + self.v_buffer = self._v_views + # For external inspectors only; the native move path doesn't consume them. + self.k_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.k_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.v_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.v_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) + self.data_strides = torch.tensor( + [x.stride(0) * x.dtype.itemsize for x in (self.k_buffer + self.v_buffer)], + device=self.device, + ) + + def _clear_buffers(self): + # Lifetime owned by UnifiedKVPool; do not delete the views. + pass + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # tgt_loc/src_loc are PHYSICAL slot ids; native move only (strided views). + if tgt_loc.numel() == 0: + return + with record_function("UnifiedMHA.move_kv_cache"): + move_kv_cache_native( + self.k_buffer, + self.v_buffer, + tgt_loc, + src_loc, + page_size=self._page_size, + ) + + def get_kv_size_bytes(self): + return 0, 0 # UnifiedKVPool logs the total; per-sub-pool would double-count + + def set_kv_buffer( + self, + layer, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale=None, + v_scale=None, + layer_id_override: Optional[int] = None, + dcp_kv_mask: Optional[torch.Tensor] = None, + ): + # Decode context parallel (dcp_kv_mask) unsupported; fail loud. + assert dcp_kv_mask is None, ( + "UnifiedMHATokenToKVPool.set_kv_buffer: decode context parallel " + "(dcp_kv_mask) is not supported with --enable-unified-memory." + ) + # Bypass super().set_kv_buffer: the parent's `k_cache.view(-1, row_dim)` can't + # merge our 4-D layer-major view (stride[0]=page_bytes) at page_size>1. Call + # store_cache_4d_kernel directly. `loc` is PHYSICAL token ids — no v2p translate. + with record_function("UnifiedMHA.set_kv_buffer"): + if cache_k.dtype != self.dtype: + if k_scale is not None: + cache_k.div_(k_scale) + if v_scale is not None: + cache_v.div_(v_scale) + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) + if self.store_dtype != self.dtype: + cache_k = cache_k.view(self.store_dtype) + cache_v = cache_v.view(self.store_dtype) + + layer_id = ( + layer.layer_id if layer_id_override is None else layer_id_override + ) - self.start_layer + k_view = self.k_buffer[layer_id] + v_view = self.v_buffer[layer_id] + ps = self._page_size + N = loc.numel() + if N == 0: + return + head_num = k_view.shape[2] + head_dim = k_view.shape[3] + v_head_dim = v_view.shape[3] + K_ROW_DIM = head_num * head_dim + V_ROW_DIM = head_num * v_head_dim + BLOCK = 128 + row_dim_max = K_ROW_DIM if K_ROW_DIM > V_ROW_DIM else V_ROW_DIM + store_cache_4d_kernel[(N, triton.cdiv(row_dim_max, BLOCK), 2)]( + k_view, + v_view, + cache_k, + cache_v, + loc, + k_view.stride(0), + k_view.stride(1), + v_view.stride(0), + v_view.stride(1), + cache_k.stride(0), + cache_v.stride(0), + K_ROW_DIM=K_ROW_DIM, + V_ROW_DIM=V_ROW_DIM, + PAGE_SIZE=ps, + BLOCK=BLOCK, + num_warps=4, + ) + + +class UnifiedMambaPool(MambaPool): + """Mamba state pool whose conv/temporal state are strided views into a `UnifiedKVPool`. + + Pure PHYSICAL store: slot lifecycle and the v<->p mapping live in the attached + `UnifiedMambaSlotAllocator`. Does NOT call `super().__init__()` — replicates the + minimal `MambaPool` state against the unified buffer so inherited methods work. + """ + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + sub_pool_name: str, + spec_state_size: int, + mamba_layer_ids: List[int], + enable_memory_saver: bool = False, + speculative_num_draft_tokens: Optional[int] = None, + ): + spec = unified_buffer.mamba_spec(sub_pool_name) + assert spec.layer_num == len(mamba_layer_ids) + conv_views, temporal_view = unified_buffer.mamba_views_for(sub_pool_name) + max_slots = unified_buffer.max_slots(sub_pool_name) + + self._unified_buffer = unified_buffer + self._sub_pool_name = sub_pool_name + + # Replicate the state MambaPool.__init__ would have set. + self._max_size = max_slots - 1 # -1 for reserved slot 0 + self.size = self._max_size + self.device = unified_buffer.device + self.memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + self.enable_custom_mem_pool = False + self.custom_mem_pool = None + self.num_mamba_layers = spec.layer_num + # GDN/KDA ReplaySSM unsupported; replicate parent's disabled-state attrs so + # paths guarded by `replayssm_write_pos is not None` don't AttributeError. + self.enable_linear_replayssm = False + self.linear_replayssm_cache_len = 16 + self.replayssm_write_pos = None + self.replayssm_is_kda = False + + assert ( + conv_views[0].shape[0] == self.num_mamba_layers + ), f"conv_views layers={conv_views[0].shape[0]} vs expected {self.num_mamba_layers}" + assert ( + conv_views[0].shape[1] == self._max_size + 1 + ), f"conv_views slots={conv_views[0].shape[1]} vs expected {self._max_size + 1}" + + # Per-draft-token intermediate buffers have a different outer size + # (spec_state_size+1), so they're NOT in the shared buffer; allocate locally. + temporal_state_shape = spec.temporal_state_shape + conv_state_shape = spec.conv_state_shapes + conv_dtype = spec.conv_dtype + ssm_dtype = spec.temporal_dtype + if speculative_num_draft_tokens is not None: + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + intermediate_ssm_state_cache = torch.zeros( + size=( + self.num_mamba_layers, + spec_state_size + 1, + speculative_num_draft_tokens, + temporal_state_shape[0], + temporal_state_shape[1], + temporal_state_shape[2], + ), + dtype=ssm_dtype, + device=unified_buffer.device, + ) + intermediate_conv_window_cache = [ + torch.zeros( + size=( + self.num_mamba_layers, + spec_state_size + 1, + speculative_num_draft_tokens, + cshape[0], + cshape[1], + ), + dtype=conv_dtype, + device=unified_buffer.device, + ) + for cshape in conv_state_shape + ] + self.mamba_cache = self.SpeculativeState( + conv=list(conv_views), + temporal=temporal_view, + intermediate_ssm=intermediate_ssm_state_cache, + intermediate_conv_window=intermediate_conv_window_cache, + ) + else: + self.mamba_cache = self.State(conv=list(conv_views), temporal=temporal_view) + + self.mem_usage = unified_buffer.total_bytes / GB + logger.info( + "[unified-memory-pool] UnifiedMambaPool(%s) wrapped unified buffer: max_slots=%d, " + "num_mamba_layers=%d", + sub_pool_name, + max_slots, + self.num_mamba_layers, + ) + + # Inherited MambaPool state ops (copy_from/clear_slots/get_cpu_copy/load_cpu_copy) + # take PHYSICAL slot ids; callers translate via the slot allocator first. + + def _copy_from_physical(self, src_index: torch.Tensor, dst_index: torch.Tensor): + # Physical-slot copy used by the allocator's `_compact_pending`. + MambaPool.copy_from(self, src_index, dst_index) + + +class UnifiedMambaSlotAllocator: + """Mamba slot allocator (PHYSICAL view) for the unified memory pool. + + Owns slot alloc/free, sizing, and the v<->p mapping (``translate``), presenting the + upstream ``MambaSlotAllocator`` interface. ``alloc()`` returns VIRTUAL ids and does + NOT clear state — clearing is deferred to ``UnifiedMambaPool.clear_slots``. + """ + + def __init__(self, mea, max_size: int, device: str): + self._multi_ended_allocator = mea + self._max_size = max_size # excludes reserved slot 0 + self._device = device + self._alloc_iter = None # active alloc_group batch iterator + + # -- translation (owns the v<->p mapping) -- + + def translate(self, virtual_ids: torch.Tensor) -> torch.Tensor: + # VIRTUAL -> PHYSICAL slot ids; page_size==1, so a direct v2p gather. + return self._multi_ended_allocator.virtual_to_physical[virtual_ids] + + @property + def virtual_to_physical(self) -> torch.Tensor: + return self._multi_ended_allocator.virtual_to_physical + + # -- sizing / free-list -- + + @property + def size(self) -> int: + return self._max_size + + def available_size(self) -> int: + # Slot-conservation count (max - allocated): the leak-check view, NOT the + # planner value (use schedulable_available_size for that). + return self._max_size - self._multi_ended_allocator.allocated_count() + + def schedulable_available_size(self) -> int: + # Byte-coordinated count (>= N => alloc(N) succeeds); credits the peer's + # drainable holes since alloc flushes the peer before extending. + return self._multi_ended_allocator.schedulable_available_size() + + @property + def free_slots(self) -> torch.Tensor: + # Watermark-derived physical free-list for the invariant checker. + a = self._multi_ended_allocator + assert a.page_size == 1, ( + "UnifiedMambaSlotAllocator.free_slots assumes page_size==1; got " + f"{a.page_size}. Mamba state is per-request, orthogonal to paging." + ) + if a.grow_direction == "up": + start, end = a.watermark_physical, a.num_pages + else: + start, end = a.min_page_index, a.watermark_physical + 1 + if start >= end: + return torch.empty((0,), dtype=torch.int64, device=self._device) + return torch.arange(start, end, dtype=torch.int64, device=self._device) + + # -- slot management (delegates to the MultiEndedAllocator) -- + + def alloc(self, need_size: int): + # alloc_group fast path: single-slot draws from the prefetched batch. + if self._alloc_iter is not None and need_size == 1: + slot = next(self._alloc_iter, None) + if slot is not None: + return slot + return self._multi_ended_allocator.alloc(need_size) # VIRTUAL ids + + def free(self, free_index: torch.Tensor): + return self._multi_ended_allocator.free(free_index) + + def clear(self): + self._alloc_iter = None + return self._multi_ended_allocator.clear() + + def alloc_group_begin(self, num_reqs: int): + """Pre-allocate a batch that ``alloc(1)`` then draws from.""" + self._alloc_iter = None + if num_reqs > 0: + result = self._multi_ended_allocator.alloc(num_reqs) + if result is not None: + self._alloc_iter = iter(result.split(1)) + + def alloc_group_end(self): + """Return any unused pre-allocated slots from the current group.""" + if self._alloc_iter is not None: + remaining = list(self._alloc_iter) + if remaining: + self._multi_ended_allocator.free(torch.cat(remaining)) + self._alloc_iter = None + + def is_slot_allocated(self, slot) -> bool: + return self._multi_ended_allocator.is_slot_allocated(int(slot)) + + def allocator_state_str(self) -> str: + return self._multi_ended_allocator.allocator_state_str() + + +class UnifiedHybridReqToTokenPool(HybridReqToTokenPool): + """`HybridReqToTokenPool` whose `mamba_pool` is a `UnifiedMambaPool`. The inherited + mamba-id state now holds VIRTUAL ids; adds `translate_mamba_indices` for v->p.""" + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + mamba_sub_pool_name: str, + size: int, + mamba_spec_state_size: int, + max_context_len: int, + device: str, + enable_memory_saver: bool, + cache_params, + mamba_layer_ids: List[int], + enable_mamba_extra_buffer: bool, + speculative_num_draft_tokens: Optional[int] = None, + enable_overlap_schedule: bool = True, + start_layer: Optional[int] = None, + ): + self._unified_buffer = unified_buffer + self._mamba_sub_pool_name = mamba_sub_pool_name + self._shared_mamba_size = ( + unified_buffer.max_slots(mamba_sub_pool_name) - 1 + ) # reserve slot 0 + super().__init__( + size=size, + mamba_size=self._shared_mamba_size, + mamba_spec_state_size=mamba_spec_state_size, + max_context_len=max_context_len, + device=device, + enable_memory_saver=enable_memory_saver, + cache_params=cache_params, + mamba_layer_ids=mamba_layer_ids, + enable_mamba_extra_buffer=enable_mamba_extra_buffer, + speculative_num_draft_tokens=speculative_num_draft_tokens, + enable_overlap_schedule=enable_overlap_schedule, + start_layer=start_layer, + ) + + def _init_mamba_pool( + self, + mamba_size: int, + mamba_spec_state_size: int, + cache_params, + mamba_layer_ids: List[int], + device: str, + enable_mamba_extra_buffer: bool, + speculative_num_draft_tokens: Optional[int] = None, + speculative_eagle_topk: Optional[int] = None, + mamba_envelope_layout: bool = False, + enable_linear_replayssm: bool = False, + linear_replayssm_cache_len: int = 16, + ): + # mamba_envelope_layout / speculative_eagle_topk / enable_linear_replayssm / + # linear_replayssm_cache_len: accepted to match the parent signature but NOT + # forwarded — the shared pool's conv/temporal state are fixed-shape views. + assert mamba_size == self._shared_mamba_size, ( + f"UnifiedHybridReqToTokenPool._init_mamba_pool: mamba_size={mamba_size} " + f"!= unified_buffer.max_slots({self._mamba_sub_pool_name!r}) - 1 " + f"= {self._shared_mamba_size}" + ) + assert len(cache_params.layers) >= len(mamba_layer_ids), ( + f"cache_params.layers ({len(cache_params.layers)}) cannot supply " + f"{len(mamba_layer_ids)} mamba layer ids" + ) + self.mamba_pool = UnifiedMambaPool( + unified_buffer=self._unified_buffer, + sub_pool_name=self._mamba_sub_pool_name, + spec_state_size=mamba_spec_state_size, + mamba_layer_ids=mamba_layer_ids, + enable_memory_saver=self.enable_memory_saver, + speculative_num_draft_tokens=speculative_num_draft_tokens, + ) + # Wired in by init_unified_mamba_pools once the mamba allocator exists. + self.mamba_allocator = None + self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)} + self.mamba_ckpt_pool = None # int8 ckpt pool unused; None = feature off + self.device = device + # Sized by req_to_token's first dim (size + 1; row 0 is padding); self.size + # would under-size by one row. + req_pool_size = self.req_to_token.shape[0] + self.req_index_to_mamba_index_mapping: torch.Tensor = torch.zeros( + req_pool_size, dtype=torch.int32, device=self.device + ) + if enable_mamba_extra_buffer: + self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = ( + torch.zeros( + (req_pool_size, self.mamba_ping_pong_track_buffer_size), + # int64 to match the parent's uncast index_put source (int32 dest + # would dtype-mismatch on the first radix prefill). + dtype=torch.int64, + device=self.device, + ) + ) + + def translate_mamba_indices(self, virtual_ids: torch.Tensor) -> torch.Tensor: + """Virtual mamba ids -> physical slot ids.""" + return self.mamba_allocator.translate(virtual_ids).to(torch.int32) + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + + +class UnifiedPoolBundle(NamedTuple): + unified_memory_pool: UnifiedKVPool + token_to_kv_pool: object # HybridLinearKVPool + token_to_kv_pool_allocator: object # UnifiedMambaTokenToKVPoolAllocator + req_to_token_pool: object # UnifiedHybridReqToTokenPool + + +def init_unified_mamba_pools( + *, + device: str, + kv_cache_dtype: torch.dtype, + head_num: int, + head_dim: int, + page_size: int, + start_layer: int, + end_layer: int, + is_draft_worker: bool, + use_mla_backend: bool, + mamba_layer_ids: List[int], + full_attention_layer_ids: List[int], + mamba2_cache_params, + model_context_len: int, + extra_max_context_len: int, + max_total_num_tokens: int, + max_mamba_cache_size: int, + max_num_reqs: int, + enable_memory_saver: bool, + enable_mamba_extra_buffer: bool, + speculative_num_draft_tokens: Optional[int], + disable_overlap_schedule: bool, + need_sort: bool, + mamba_full_memory_ratio: Optional[float] = None, # informational only + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, +) -> 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, + ) + + assert ( + not use_mla_backend + ), "unified memory pool does not support MLA-hybrid-Mamba yet" + # Full sub-pool is page-aware; mamba stays page=1 (state is per-request). + assert page_size >= 1, f"page_size must be >= 1, got {page_size}" + + store_dtype = _store_dtype_for(kv_cache_dtype) + # full-attn at the high-byte end (grow-down), mamba at the low-byte end (grow-up). + full_spec = MHASubPoolSpec( + name="full", + layer_num=len(full_attention_layer_ids), + head_num=head_num, + head_dim=head_dim, + store_dtype=store_dtype, + grow_direction="down", + ) + cp = mamba2_cache_params + mamba_spec = MambaSubPoolSpec( + name="mamba", + layer_num=len(mamba_layer_ids), + conv_state_shapes=tuple(tuple(int(x) for x in s) for s in cp.shape.conv), + conv_dtype=cp.dtype.conv, + temporal_state_shape=tuple(int(x) for x in cp.shape.temporal), + temporal_dtype=cp.dtype.temporal, + grow_direction="up", + ) + total_bytes = ( + max_total_num_tokens * full_spec.entry_bytes() + + max_mamba_cache_size * mamba_spec.entry_bytes() + ) + shared_pool = UnifiedKVPool( + total_bytes=total_bytes, + sub_pool_specs=[full_spec, mamba_spec], + device=device, + enable_memory_saver=enable_memory_saver, + page_size=page_size, + ) + req_to_token_pool = UnifiedHybridReqToTokenPool( + unified_buffer=shared_pool, + mamba_sub_pool_name="mamba", + size=max_num_reqs, + mamba_spec_state_size=max_num_reqs, # outer dim of spec-decode intermediates + max_context_len=model_context_len + extra_max_context_len, + device=device, + enable_memory_saver=enable_memory_saver, + cache_params=mamba2_cache_params, + mamba_layer_ids=mamba_layer_ids, + enable_mamba_extra_buffer=enable_mamba_extra_buffer, + speculative_num_draft_tokens=speculative_num_draft_tokens, + enable_overlap_schedule=not disable_overlap_schedule, + start_layer=start_layer, + ) + unified_full_kv_pool = UnifiedMHATokenToKVPool( + unified_buffer=shared_pool, + sub_pool_name="full", + page_size=page_size, + start_layer=start_layer, + end_layer=end_layer, + ) + full_attn_layer_ids_for_pool = ( + [0] if is_draft_worker else list(full_attention_layer_ids) + ) + token_to_kv_pool = HybridLinearKVPool( + page_size=page_size, + size=max_total_num_tokens, + dtype=kv_cache_dtype, + head_num=head_num, + head_dim=head_dim, + full_attention_layer_ids=full_attn_layer_ids_for_pool, + device=device, + mamba_pool=req_to_token_pool.mamba_pool, + enable_memory_saver=enable_memory_saver, + use_mla=use_mla_backend, + start_layer=start_layer, + full_kv_pool=unified_full_kv_pool, + ) + allocator = UnifiedMambaTokenToKVPoolAllocator( + unified_buffer=shared_pool, + kvcache=token_to_kv_pool, + device=device, + page_size=page_size, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + + # Wrap the composite's mamba MultiEndedAllocator in a slot allocator (PHYSICAL view). + mamba_slot_allocator = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=req_to_token_pool._shared_mamba_size, + device=device, + ) + # `_mamba_translate` feeds the HiCache offload path, GATED OFF here — wired but inert. + req_to_token_pool.mamba_allocator = mamba_slot_allocator + token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate + + logger.info( + "[unified-memory-pool] ============================================================" + ) + logger.info( + "[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=Mamba hybrid" + ) + logger.info( + "[unified-memory-pool] full_layers=%d, mamba_layers=%d, head_num=%d, head_dim=%d, " + "page_size=%d, is_draft_worker=%s", + len(full_attention_layer_ids), + len(mamba_layer_ids), + head_num, + head_dim, + page_size, + is_draft_worker, + ) + logger.info( + "[unified-memory-pool] total_bytes=%d, max_total_num_tokens=%d, max_mamba_cache_size=%d, " + "max_num_reqs=%d, speculative_num_draft_tokens=%s", + total_bytes, + max_total_num_tokens, + max_mamba_cache_size, + max_num_reqs, + speculative_num_draft_tokens, + ) + if mamba_full_memory_ratio is not None: + logger.info( + "[unified-memory-pool] mamba_full_memory_ratio=%s governs the total budget only, " + "not the runtime split.", + mamba_full_memory_ratio, + ) + logger.info( + "[unified-memory-pool] ============================================================" + ) + return UnifiedPoolBundle( + unified_memory_pool=shared_pool, + token_to_kv_pool=token_to_kv_pool, + token_to_kv_pool_allocator=allocator, + req_to_token_pool=req_to_token_pool, + ) + + +# --------------------------------------------------------------------------- +# UnifiedSWAKVPool — hybrid SWA on the shared byte buffer +# --------------------------------------------------------------------------- + + +class UnifiedSWAKVPool(SWAKVPool): + """Shared-buffer replacement for `SWAKVPool`. + + Composes two `UnifiedMHATokenToKVPool` instances (full + swa) aliasing the same + byte buffer. Inherits from `SWAKVPool` only for `isinstance`; does NOT call the + parent `__init__` (it would build static-partition pools). The per-sub-pool v2p + table IS the full->swa mapping, so `register_mapping` is a no-op. + """ + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + swa_attention_layer_ids: List[int], + full_attention_layer_ids: List[int], + page_size: int = 1, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + enable_memory_saver: bool = False, + ): + # Do NOT call super().__init__ — it would allocate static-partition pools. + self.unified_buffer = unified_buffer + self.swa_layer_nums = len(swa_attention_layer_ids) + self.full_layer_nums = len(full_attention_layer_ids) + self.layer_num = self.full_layer_nums + self.swa_layer_nums + self.start_layer = start_layer if start_layer is not None else 0 + self.page_size = page_size + self.layer_transfer_counter = None + + self.size = unified_buffer.max_slots("full") - 1 + self.size_swa = unified_buffer.max_slots("swa") - 1 + + full_spec = unified_buffer.mha_spec("full") + swa_spec = unified_buffer.mha_spec("swa") + assert full_spec.store_dtype == swa_spec.store_dtype, ( + "UnifiedSWAKVPool: full and swa sub-pools must share store_dtype; got " + f"full={full_spec.store_dtype}, swa={swa_spec.store_dtype}" + ) + self.dtype = full_spec.store_dtype + self.head_num = full_spec.head_num + self.head_dim = full_spec.head_dim + self.device = unified_buffer.device + + self.full_kv_pool = UnifiedMHATokenToKVPool( + unified_buffer=unified_buffer, + sub_pool_name="full", + page_size=page_size, + start_layer=start_layer, + end_layer=end_layer, + ) + self.swa_kv_pool = UnifiedMHATokenToKVPool( + unified_buffer=unified_buffer, + sub_pool_name="swa", + page_size=page_size, + start_layer=start_layer, + end_layer=end_layer, + ) + + # disagg/nvlink disabled; keep attrs present to avoid AttributeError. + self.enable_custom_mem_pool = False + self.custom_mem_pool = None + + # {global_layer_id: (per-pool index, is_swa_layer)} + self.layers_mapping: Dict[int, Tuple[int, bool]] = {} + for idx, gid in enumerate(full_attention_layer_ids): + self.layers_mapping[gid] = (idx, False) + for idx, gid in enumerate(swa_attention_layer_ids): + self.layers_mapping[gid] = (idx, True) + + # None so dispatch routes through our v2p-table overrides, not a registered mapping. + self.full_to_swa_index_mapping: Optional[torch.Tensor] = None + + self.mem_usage = 0.0 # cosmetic; UnifiedKVPool logs the real size + + # Wired in via attach_allocators. + self._full_allocator = None + self._swa_allocator = None + + logger.info( + "[unified-memory-pool] UnifiedSWAKVPool wrapped unified buffer: " + "full_layers=%d (max_slots=%d), swa_layers=%d (max_slots=%d), " + "head_num=%d, head_dim=%d", + self.full_layer_nums, + unified_buffer.max_slots("full"), + self.swa_layer_nums, + unified_buffer.max_slots("swa"), + self.head_num, + self.head_dim, + ) + + # -- allocator wiring -- + + def attach_allocators(self, *, full_allocator, swa_allocator) -> None: + """Wire the two `MultiEndedAllocator`s whose v2p tables translate slot ids.""" + self._full_allocator = full_allocator + self._swa_allocator = swa_allocator + + # -- BaseSWAKVPool ABC surface -- + + def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None: + return # no-op in shared mode (the swa-side v2p IS the mapping) + + def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): + """Virtual token ids -> swa-physical token ids (int32).""" + assert self._swa_allocator is not None, ( + "UnifiedSWAKVPool.translate_loc_from_full_to_swa called before " + "attach_allocators" + ) + ps = self._swa_allocator.page_size + if ps == 1: + return self._swa_allocator.virtual_to_physical[kv_indices].to(torch.int32) + virt_pages = kv_indices // ps + offsets = kv_indices % ps + swa_phys_pages = self._swa_allocator.virtual_to_physical[virt_pages] + return (swa_phys_pages * ps + offsets).to(torch.int32) + + def get_state_buf_infos(self): + return self.swa_kv_pool.get_contiguous_buf_infos() + + # -- size/info -- + + def get_kv_size_bytes(self): + return 0, 0 # UnifiedKVPool logs the total; per-side would double-count + + def get_contiguous_buf_infos(self): + return self.full_kv_pool.get_contiguous_buf_infos() + + # -- buffer accessors -- + + def get_key_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) + pool_layer_id, is_swa = self.layers_mapping[layer_id] + pool = self.swa_kv_pool if is_swa else self.full_kv_pool + return pool.get_key_buffer(pool_layer_id) + + def get_value_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) + pool_layer_id, is_swa = self.layers_mapping[layer_id] + pool = self.swa_kv_pool if is_swa else self.full_kv_pool + return pool.get_value_buffer(pool_layer_id) + + def get_kv_buffer(self, layer_id: int): + self._wait_for_layer(layer_id) + pool_layer_id, is_swa = self.layers_mapping[layer_id] + pool = self.swa_kv_pool if is_swa else self.full_kv_pool + return pool.get_kv_buffer(pool_layer_id) + + # -- kv writing -- + + def set_kv_buffer( + self, + layer, + loc_info, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: float = 1.0, + v_scale: float = 1.0, + ): + """Route to the right sub-pool. Both `swa_loc` and `full_loc` are PHYSICAL + (pre-translated once per forward by the attention backend); never translates here. + """ + _, swa_loc, full_loc = unwrap_write_loc(loc_info) + layer_id = layer.layer_id + pool_layer_id, is_swa = self.layers_mapping[layer_id] + if is_swa: + # swa_loc is ALREADY swa-physical. Routed through the UnifiedMHATokenToKVPool + # override (its 4-D layer-major view can't take the parent's view(-1, row_dim)). + assert swa_loc is not None, ( + "UnifiedSWAKVPool.set_kv_buffer: SWA layer received no swa_loc; the " + "attention backend must bundle forward_metadata.swa_out_cache_loc." + ) + self.swa_kv_pool.set_kv_buffer( + None, + swa_loc, + cache_k, + cache_v, + k_scale, + v_scale, + layer_id_override=pool_layer_id, + ) + return + # Full layer: full_loc is full-physical, always precomputed (eager + cuda-graph). + assert full_loc is not None, ( + "UnifiedSWAKVPool.set_kv_buffer: full layer received no full_loc; " + "ForwardMetadata.out_cache_loc_full_physical must be precomputed for " + "the unified memory pool." + ) + self.full_kv_pool.set_kv_buffer( + None, + full_loc, + cache_k, + cache_v, + k_scale, + v_scale, + layer_id_override=pool_layer_id, + ) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # Never called on the composite — compaction runs per-sub-pool via + # UnifiedMHATokenToKVPool.move_kv_cache. + raise NotImplementedError( + "UnifiedSWAKVPool.move_kv_cache should not be called; compaction " + "operates per-sub-pool via UnifiedMHATokenToKVPool.move_kv_cache." + ) + + # -- HiCache shims (translate virtual->physical, then delegate) -- + + @staticmethod + def _virt_tokens_to_phys_tokens( + virt_tokens: torch.Tensor, allocator + ) -> torch.Tensor: + """Virtual TOKEN ids -> physical TOKEN ids (page-aware). Unbound pages yield + negatives; callers filter via `swa_phys >= 0`.""" + ps = allocator.page_size + if ps == 1: + return allocator.virtual_to_physical[virt_tokens] + virt_pages = virt_tokens // ps + offsets = virt_tokens % ps + phys_pages = allocator.virtual_to_physical[virt_pages] + return phys_pages * ps + offsets + + def get_cpu_copy(self, indices, mamba_indices=None): + assert self._full_allocator is not None + assert self._swa_allocator is not None + # `indices` are virtual TOKEN ids; translate per sub-pool. + full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator) + swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator) + full_cpu = self.full_kv_pool.get_cpu_copy(full_phys) + valid = swa_phys >= 0 + swa_cpu = None + if bool(valid.any().item()): + swa_cpu = self.swa_kv_pool.get_cpu_copy(swa_phys[valid]) + return {"full": full_cpu, "swa": swa_cpu} + + def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + assert self._full_allocator is not None + full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator) + self.full_kv_pool.load_cpu_copy(kv_cache_cpu["full"], full_phys) + if kv_cache_cpu.get("swa") is not None: + assert self._swa_allocator is not None + swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator) + self.swa_kv_pool.load_cpu_copy(kv_cache_cpu["swa"], swa_phys) + + +class UnifiedSWAPoolBundle(NamedTuple): + unified_memory_pool: UnifiedKVPool + token_to_kv_pool: object # UnifiedSWAKVPool + token_to_kv_pool_allocator: object # UnifiedSWATokenToKVPoolAllocator + + +def init_unified_swa_pools( + *, + device: str, + kv_cache_dtype: torch.dtype, + head_num: int, + head_dim: int, + v_head_dim: int, + swa_head_num: int, + swa_head_dim: int, + swa_v_head_dim: int, + page_size: int, + start_layer: int, + end_layer: int, + swa_attention_layer_ids: List[int], + full_attention_layer_ids: List[int], + full_max_total_num_tokens: int, + swa_max_total_num_tokens: int, + enable_memory_saver: bool, + need_sort: bool, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, +) -> UnifiedSWAPoolBundle: + """Build the SWA-hybrid unified-memory-pool stack.""" + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + # Both sub-allocators are page-aware: one virtual ID space at PAGE granularity, + # two physical sub-pools compacting pages independently. + assert page_size >= 1, f"page_size must be >= 1, got {page_size}" + assert ( + len(full_attention_layer_ids) > 0 + ), "SWA-hybrid with zero full-attention layers is degenerate" + assert ( + len(swa_attention_layer_ids) > 0 + ), "SWA-hybrid with zero SWA-attention layers is degenerate" + + store_dtype = _store_dtype_for(kv_cache_dtype) + # full-attn at the high-byte end (grow-down), swa at the low-byte end (grow-up). + full_spec = MHASubPoolSpec( + name="full", + layer_num=len(full_attention_layer_ids), + head_num=head_num, + head_dim=head_dim, + v_head_dim=v_head_dim, + store_dtype=store_dtype, + grow_direction="down", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=len(swa_attention_layer_ids), + head_num=swa_head_num, + head_dim=swa_head_dim, + v_head_dim=swa_v_head_dim, + store_dtype=store_dtype, + grow_direction="up", + ) + total_bytes = ( + full_max_total_num_tokens * full_spec.entry_bytes() + + swa_max_total_num_tokens * swa_spec.entry_bytes() + ) + shared_pool = UnifiedKVPool( + total_bytes=total_bytes, + sub_pool_specs=[full_spec, swa_spec], + device=device, + enable_memory_saver=enable_memory_saver, + page_size=page_size, + ) + token_to_kv_pool = UnifiedSWAKVPool( + unified_buffer=shared_pool, + swa_attention_layer_ids=swa_attention_layer_ids, + full_attention_layer_ids=full_attention_layer_ids, + page_size=page_size, + start_layer=start_layer, + end_layer=end_layer, + enable_memory_saver=enable_memory_saver, + ) + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=shared_pool, + kvcache=token_to_kv_pool, + device=device, + full_max_total_num_tokens=full_max_total_num_tokens, + swa_max_total_num_tokens=swa_max_total_num_tokens, + page_size=page_size, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + + logger.info( + "[unified-memory-pool] ============================================================" + ) + logger.info("[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=SWA hybrid") + logger.info( + "[unified-memory-pool] full_layers=%d, swa_layers=%d, head_num=%d, head_dim=%d, " + "v_head_dim=%d, swa_head_num=%d, swa_head_dim=%d, swa_v_head_dim=%d, " + "page_size=%d", + len(full_attention_layer_ids), + len(swa_attention_layer_ids), + head_num, + head_dim, + v_head_dim, + swa_head_num, + swa_head_dim, + swa_v_head_dim, + page_size, + ) + logger.info( + "[unified-memory-pool] total_bytes=%d (=%.2f GB), full_max_total_num_tokens=%d, " + "swa_max_total_num_tokens=%d, joint_available=%d slots", + total_bytes, + total_bytes / GB, + full_max_total_num_tokens, + swa_max_total_num_tokens, + allocator.available_size(), + ) + logger.info( + "[unified-memory-pool] ============================================================" + ) + return UnifiedSWAPoolBundle( + unified_memory_pool=shared_pool, + token_to_kv_pool=token_to_kv_pool, + token_to_kv_pool_allocator=allocator, + ) diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index a87bf4b5f..ef8ea8558 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -62,6 +62,14 @@ if TYPE_CHECKING: from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig +def _should_enable_lazy_compaction() -> bool: + """Lazy compaction default — ON unless + `SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback). + Centralized here so both unified-memory-pool factory call sites stay in sync. + """ + return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get() + + # the ratio of mamba cache pool size to max_running_requests MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2 @@ -341,10 +349,185 @@ class ModelRunnerKVCacheMixin: "attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1." ) + def _init_unified_mamba_pools(self: ModelRunner, max_num_reqs: int): + """Build the shared-KV-pool stack for a hybrid-Mamba model: + one byte buffer split between the full-attn MHA KV pool and the + per-request Mamba state pool, with virtual slot ids above the + allocator.""" + from sglang.srt.mem_cache.unified_memory_pool import init_unified_mamba_pools + + config = self.mambaish_config + assert config is not None + assert ( + not self.use_mla_backend + ), "unified memory pool does not support MLA-hybrid-Mamba yet" + # The full sub-pool is page-aware (via `MultiEndedAllocator(page_size=...)`); + # the mamba sub-pool stays page=1. + assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" + # Mirror the non-shared path's extra_max_context_len computation. + extra_max_context_len = 4 + if self.server_args.speculative_num_draft_tokens is not None: + extra_max_context_len += self.server_args.speculative_num_draft_tokens + + mamba_layer_ids = [ + i + for i in config.mamba2_cache_params.layers + if self.start_layer <= i < self.end_layer + ] + full_attention_layer_ids = [ + i + for i in config.full_attention_layer_ids + if self.start_layer <= i < self.end_layer + ] + + bundle = init_unified_mamba_pools( + device=self.device, + kv_cache_dtype=self.kv_cache_dtype, + head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()), + head_dim=self.model_config.head_dim, + page_size=self.page_size, + start_layer=self.start_layer, + end_layer=self.end_layer, + is_draft_worker=self.is_draft_worker, + use_mla_backend=self.use_mla_backend, + mamba_layer_ids=mamba_layer_ids, + full_attention_layer_ids=full_attention_layer_ids, + mamba2_cache_params=config.mamba2_cache_params, + model_context_len=self.model_config.context_len, + extra_max_context_len=extra_max_context_len, + max_total_num_tokens=self.max_total_num_tokens, + max_mamba_cache_size=self.server_args.max_mamba_cache_size, + max_num_reqs=max_num_reqs, + enable_memory_saver=self.server_args.enable_memory_saver, + enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), + speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, + disable_overlap_schedule=self.server_args.disable_overlap_schedule, + need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"), + mamba_full_memory_ratio=self.server_args.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 + # v2p/KV reads. Near-no-op in normal mode. + forward_stream=self.forward_stream, + # Lazy compaction: default ON, env-var escape hatch for rollback / A/B. + lazy_compaction=_should_enable_lazy_compaction(), + ) + self.req_to_token_pool = bundle.req_to_token_pool + self.token_to_kv_pool = bundle.token_to_kv_pool + self.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator + # Keep a reference so the shared byte buffer is not GC'd. + self._unified_memory_pool = bundle.unified_memory_pool + + def _init_unified_swa_pools(self: ModelRunner, max_num_reqs: int): + """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte + buffer split between the full-attention and SWA KV pools.""" + from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools + + assert self.is_hybrid_swa, "_init_unified_swa_pools called on a non-SWA model" + # Both sub-pools are page-aware; the SWA composite runs alloc_extend_kernel + # once in virtual space and binds the new pages on both sub-allocators. + assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" + assert ( + not self.use_mla_backend + ), "unified memory pool does not support MLA-SWA hybrid yet" + # Mirror the non-shared path's extra_max_context_len computation. + extra_max_context_len = 4 + if self.server_args.speculative_num_draft_tokens is not None: + extra_max_context_len += self.server_args.speculative_num_draft_tokens + self.req_to_token_pool = ReqToTokenPool( + size=max_num_reqs, + max_context_len=self.model_config.context_len + extra_max_context_len, + device=self.device, + enable_memory_saver=self.server_args.enable_memory_saver, + ) + + head_num = self.model_config.get_num_kv_heads(get_attention_tp_size()) + head_dim = self.model_config.head_dim + if self.is_hybrid_swa_compress: + # Asymmetric head dims between full and SWA (NPU compress path): + # pull SWA-specific dims from the hf text config. + v_head_dim = self.model_config.hf_text_config.v_head_dim + swa_head_num = max( + 1, + self.model_config.hf_text_config.swa_num_key_value_heads + // get_attention_tp_size(), + ) + swa_head_dim = self.model_config.hf_text_config.swa_head_dim + swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim + else: + v_head_dim = head_dim + swa_head_num = head_num + swa_head_dim = head_dim + swa_v_head_dim = head_dim + + # Filter layer ids to this worker's [start_layer, end_layer) range. + swa_attention_layer_ids = [ + i + for i in self.model_config.swa_attention_layer_ids + if self.start_layer <= i < self.end_layer + ] + full_attention_layer_ids = [ + i + for i in self.model_config.full_attention_layer_ids + if self.start_layer <= i < self.end_layer + ] + + bundle = init_unified_swa_pools( + device=self.device, + kv_cache_dtype=self.kv_cache_dtype, + head_num=head_num, + head_dim=head_dim, + v_head_dim=v_head_dim, + swa_head_num=swa_head_num, + swa_head_dim=swa_head_dim, + swa_v_head_dim=swa_v_head_dim, + page_size=self.page_size, + start_layer=self.start_layer, + end_layer=self.end_layer, + swa_attention_layer_ids=swa_attention_layer_ids, + full_attention_layer_ids=full_attention_layer_ids, + full_max_total_num_tokens=self.full_max_total_num_tokens, + swa_max_total_num_tokens=self.swa_max_total_num_tokens, + enable_memory_saver=self.server_args.enable_memory_saver, + need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"), + # Overlap mode: same wait_stream(forward_stream) rationale as + # `_init_unified_mamba_pools`. + forward_stream=self.forward_stream, + # Lazy compaction: default ON, with env var escape hatch for rollback / A/B. + lazy_compaction=_should_enable_lazy_compaction(), + ) + self.token_to_kv_pool = bundle.token_to_kv_pool + self.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator + # Keep a reference so the shared byte buffer is not GC'd. + self._unified_memory_pool = bundle.unified_memory_pool + def _init_pools(self: ModelRunner): """Initialize the memory pools.""" max_num_reqs = self.max_running_requests + # 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 ( + self.server_args.enable_unified_memory + and self.server_args.disaggregation_mode == "null" + and self.req_to_token_pool is None + ): + if self.mambaish_config is not None: + self._init_unified_mamba_pools(max_num_reqs) + return + if self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config): + self._init_unified_swa_pools(max_num_reqs) + return + # Fail loud, not silently fall through to the normal pools (which would + # leave the flag a no-op). The feature replaces the HYBRID pools only. + raise ValueError( + "--enable-unified-memory only supports hybrid Mamba and " + "hybrid sliding-window-attention models (DeepSeek-V4 excluded); " + f"the current model ({self.model_config.hf_config.architectures}) " + "is neither, so the unified memory pool cannot be built. Drop " + "--enable-unified-memory for this model." + ) + # Initialize req_to_token_pool if self.req_to_token_pool is None: max_spec_draft_tokens = self.server_args.max_speculative_num_draft_tokens diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index e34f6826d..3b8e753a9 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -157,6 +157,11 @@ def build_replay_fb_view( encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), + # The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE + # for the backend, which copies the result into its own static buffer and + # reads THAT in the decode track-save — this slot is never mutated. None + # when mamba-track is disabled. + mamba_track_indices=getattr(buffers, "mamba_track_indices", None), spec_info=forward_batch.spec_info, ) @@ -705,6 +710,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if self.enable_profile_cuda_graph: self._post_process_after_profile(prof) + # No pool-side pin to clear: the captured full-physical write loc rides the + # backend's `ForwardMetadata.out_cache_loc_full_physical` (-> KVWriteLoc.full_loc). + def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None: avail_mem = get_available_gpu_memory( self.model_runner.device, @@ -775,10 +783,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) def run_once(): - # Must run inside the capture block: warmup mutations here are - # undone by on_after_cuda_graph_warmup so capture starts clean. + # Graph-recordable metadata-prep hook. The unified memory pool + # records ZERO translate nodes here: all its read/write translates + # run eagerly in `init_forward_metadata_out_graph` (replay-prep), so + # the captured graph reads already-physical locs. Base no-op for triton. attn_backend.init_forward_metadata_in_graph(forward_batch) + # No invalidate_loc_cache() here: the unified pool translates its + # locs in `init_forward_metadata_out_graph`, so no cache to invalidate. + forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = ( None ) @@ -834,6 +847,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if (c := self.model_runner.canary_manager) is not None else contextlib.nullcontext() ) + # Full-physical write loc lives in the attention metadata (the backend's + # `out_cache_loc_full_physical` -> KVWriteLoc.full_loc), so the runner + # wires no buffer here. (SWA write loc rides the `swa_out_cache_loc` rail.) + with canary_ctx: shape_key = self._make_graph_key(bs, stream_idx, variant_label) self.backend.capture_one( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7952c1e63..e6bc49376 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -783,6 +783,14 @@ class ServerArgs: "(layer-major) layout. Requires the Triton attention / linear-attn / " "Mamba backends.", ] = False + enable_unified_memory: A[ + bool, + "Replace the statically-partitioned hybrid-model pools (full-attn KV + " + "SWA/Mamba state) with one byte buffer split dynamically between " + "sub-pools. Requires the Triton attention / linear-attn / Mamba " + "backends; not yet compatible with PD disaggregation or speculative " + "decoding.", + ] = False disable_chunked_prefix_cache: A[ bool, "Disable chunked prefix cache feature for deepseek, which should save overhead for short sequences.", @@ -2717,6 +2725,8 @@ class ServerArgs: self._handle_page_major_kv_layout() + self._handle_unified_memory_pool() + # Handle diffusion LLM inference. self._handle_dllm_inference() @@ -6336,7 +6346,49 @@ class ServerArgs: "NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1." ) + 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." + ) + assert self.speculative_algorithm is None, ( + "--enable-unified-memory is not yet compatible with speculative " + "decoding." + ) + assert not (self.enable_hierarchical_cache or self.enable_lmcache), ( + "--enable-unified-memory is not yet compatible with hierarchical / " + "host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): " + "the unified-memory-pool init wires up no host pools, and its device mamba / " + "full-attention slots are VIRTUAL — the host-offload path does not " + "translate them to physical." + ) + assert self.dcp_size == 1, ( + "--enable-unified-memory is not yet compatible with decode context " + "parallelism (--dcp-size > 1): the pool has no DCP-aware masked write " + "path (UnifiedMHATokenToKVPool.set_kv_buffer asserts dcp_kv_mask is None), " + "so a DCP run would boot and then fail on the first KV write." + ) + # Only monolithic decode cuda-graph capture is wired; piecewise prefill + # capture is not. Guard when the user opts into it. + _cg_cfg = self.cuda_graph_config + if _cg_cfg is not None and _cg_cfg.prefill.backend == Backend.TC_PIECEWISE: + raise ValueError( + "--enable-unified-memory supports monolithic (decode) " + "cuda-graph capture only; disable piecewise prefill capture " + "(e.g. --cuda-graph-backend-prefill=disabled)." + ) + # The strided-layout Triton requirement is enforced via + # --enable-page-major-kv-layout (implied by the unified pool in + # _handle_page_major_kv_layout); the model-family gate is enforced at pool + # construction in model_runner_kv_cache_mixin._init_pools. + def _handle_page_major_kv_layout(self): + # The unified pool stores state in the page-major envelope-strided layout, so + # enabling it implies --enable-page-major-kv-layout — routing it through the + # single page-major path + stride-aware Triton asserts (set before the guard). + if self.enable_unified_memory: + self.enable_page_major_kv_layout = True if not self.enable_page_major_kv_layout: return # Only the Triton attention kernels read the strided 4-D envelope K/V diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py new file mode 100644 index 000000000..9c34f30f8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -0,0 +1,206 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Routing tests for the composite write paths (`UnifiedSWAKVPool`, +`HybridLinearKVPool`). + +All write-location info travels in the attention metadata (`KVWriteLoc`); the +pools hold none and never translate — the write loc reaching `set_kv_buffer` is +always PHYSICAL. Two routing contracts are pinned here: + +1. Full-attention. The full-physical loc is carried in `KVWriteLoc.full_loc` + (from `ForwardBatch.out_cache_loc_full_physical`) and written directly. + `UnifiedSWAKVPool` asserts it's present (the unified memory pool always precomputes + it); `HybridLinearKVPool` falls back to `loc` for a static (non-shared) pool, + where `loc` is itself already physical. +2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` rail + (`KVWriteLoc.swa_loc`) and is written directly. + +Pure dispatch tests: the inner sub-pools are recording stubs, so no GPU / real +buffers are needed (CPU CI). + + python -m pytest test/registered/unit/mem_cache/test_full_loc_fast_path.py -v +""" + +import types +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _loc_info(virtual_loc, swa_phys=None, full_phys=None): + from sglang.srt.mem_cache.memory_pool import KVWriteLoc + + return KVWriteLoc(virtual_loc, swa_phys, full_phys) + + +class _RecordingPool: + """Stub sub-pool that records the `loc` and kwargs passed to `set_kv_buffer`.""" + + def __init__(self): + self.calls = [] + + def set_kv_buffer(self, layer, loc, cache_k, cache_v, *args, **kwargs): + self.calls.append((loc, kwargs)) + + +class TestUnifiedSWARouting(unittest.TestCase): + """`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write the full-physical + `full_loc`; SWA layers write the swa-physical `swa_loc`. Both come from the + write metadata; the pool never translates.""" + + def _make_bare_pool(self): + from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool + + # Bypass the heavy __init__; set only the attributes set_kv_buffer reads. + pool = object.__new__(UnifiedSWAKVPool) + pool.full_kv_pool = _RecordingPool() + pool.swa_kv_pool = _RecordingPool() + # layer 0 -> full attention; layer 1 -> SWA. (pool_layer_id, is_swa) + pool.layers_mapping = {0: (0, False), 1: (0, True)} + return pool + + def test_full_layer_writes_full_loc(self): + pool = self._make_bare_pool() + virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64) + swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64) + full_phys = torch.tensor([3, 4, 5], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=0) # full layer + pool.set_kv_buffer( + layer, + _loc_info(virtual_loc, swa_phys, full_phys), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + self.assertEqual(len(pool.full_kv_pool.calls), 1) + forwarded, kwargs = pool.full_kv_pool.calls[0] + # Forward the full-physical tensor from the write metadata, NOT the + # virtual loc. No `already_physical` — the pool only ever gets physical. + self.assertIs(forwarded, full_phys) + self.assertIsNot(forwarded, virtual_loc) + self.assertNotIn("already_physical", kwargs) + + def test_full_layer_requires_full_loc(self): + pool = self._make_bare_pool() + virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64) + swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=0) + # No full_loc precomputed -> fail loud (the unified memory pool must precompute + # out_cache_loc_full_physical) rather than write a virtual loc as physical. + with self.assertRaises(AssertionError): + pool.set_kv_buffer( + layer, + _loc_info(virtual_loc, swa_phys), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + def test_swa_layer_writes_swa_loc(self): + pool = self._make_bare_pool() + virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64) + swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=1) # SWA layer + pool.set_kv_buffer( + layer, + _loc_info(virtual_loc, swa_phys), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + self.assertEqual(len(pool.swa_kv_pool.calls), 1) + forwarded, kwargs = pool.swa_kv_pool.calls[0] + # SWA write rides the backend rail: forward the swa-physical loc directly. + self.assertIs(forwarded, swa_phys) + self.assertNotIn("already_physical", kwargs) + # Full pool untouched for an SWA layer. + self.assertEqual(len(pool.full_kv_pool.calls), 0) + + def test_swa_layer_requires_swa_loc(self): + pool = self._make_bare_pool() + virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=1) # SWA layer + # No swa_loc bundled -> the rail contract is violated; must assert + # rather than silently writing wrong (un-translated) locations. + with self.assertRaises(AssertionError): + pool.set_kv_buffer( + layer, + _loc_info(virtual_loc, None), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + +class TestHybridLinearFullLocRouting(unittest.TestCase): + """`HybridLinearKVPool.set_kv_buffer` (non-MLA) writes the full-physical + `full_loc` from the write metadata when present (unified memory pool), else the + already-physical `loc` (static pool). No translate, no `already_physical`.""" + + def _make_bare_pool(self): + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + pool = object.__new__(HybridLinearKVPool) + pool.full_kv_pool = _RecordingPool() + pool.use_mla = False + pool.full_attention_layer_id_mapping = {0: 0} + return pool + + def test_writes_full_loc_from_write_loc(self): + pool = self._make_bare_pool() + virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64) + full_phys = torch.tensor([2, 3, 4], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=0) + pool.set_kv_buffer( + layer, + _loc_info(virtual_loc, full_phys=full_phys), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + self.assertEqual(len(pool.full_kv_pool.calls), 1) + forwarded, kwargs = pool.full_kv_pool.calls[0] + self.assertIs(forwarded, full_phys) + self.assertIsNot(forwarded, virtual_loc) + self.assertNotIn("already_physical", kwargs) + + def test_falls_back_to_loc_when_absent(self): + # Static (non-shared) pool: no full_loc bundled; `loc` is already + # physical, so write it directly. + pool = self._make_bare_pool() + phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64) + + layer = types.SimpleNamespace(layer_id=0) + pool.set_kv_buffer( + layer, + _loc_info(phys_loc), + torch.zeros(3, 4, 8), + torch.zeros(3, 4, 8), + ) + + self.assertEqual(len(pool.full_kv_pool.calls), 1) + forwarded, kwargs = pool.full_kv_pool.calls[0] + self.assertIs(forwarded, phys_loc) + self.assertNotIn("already_physical", kwargs) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_layout_compat.py b/test/registered/unit/mem_cache/test_layout_compat.py new file mode 100644 index 000000000..41ddbbe1e --- /dev/null +++ b/test/registered/unit/mem_cache/test_layout_compat.py @@ -0,0 +1,355 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Unit tests for the page-major layer-major byte layout. + +Verifies that: +1. The new 4-D ``_build_mha_views`` output exposes correct byte addresses + for each (layer, page, tok_in_page, head, dim) — under both the + degenerate ``page_size=1`` case (byte-identical to the old per-token + envelope) and the new ``page_size>1`` layer-major case. +2. ``MHASubPoolSpec.layer_k_offset_in_page`` / + ``layer_v_offset_in_page`` math matches the layout intent. +3. ``set_kv_buffer`` round-trips correctly for both page sizes. +4. Compaction (``move_kv_cache_native``) moves the right bytes for both + page sizes via the 4-D advanced indexing path. + +CPU-only — no GPU / Triton needed. + + python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import unittest + +import torch + +from sglang.srt.mem_cache.memory_pool import move_kv_cache_native +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedKVPool, +) + +_DEV = "cpu" + + +def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4): + return MHASubPoolSpec( + name=name, + layer_num=layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction=grow, + ) + + +def _make_mamba_spec(name, grow, layer_num=2): + return MambaSubPoolSpec( + name=name, + layer_num=layer_num, + conv_state_shapes=((4, 3),), + conv_dtype=torch.float32, + temporal_state_shape=(2, 2, 2), + temporal_dtype=torch.float32, + grow_direction=grow, + ) + + +class TestMHASpecLayerOffsets(unittest.TestCase): + """Verify ``layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math.""" + + def test_offsets_at_page_size_1_match_envelope(self): + spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + # At ps=1, layer-major within a 1-token page IS envelope-per-token. + # Layer L's K offset = L * (k_row + v_row); V offset = +k_row. + k_row = spec.k_row_bytes() + v_row = spec.v_row_bytes() + for L in range(spec.layer_num): + self.assertEqual( + spec.layer_k_offset_in_page(L, page_size=1), + L * (k_row + v_row), + ) + self.assertEqual( + spec.layer_v_offset_in_page(L, page_size=1), + L * (k_row + v_row) + k_row, + ) + + def test_offsets_at_page_size_gt_1(self): + spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + ps = 8 + k_row = spec.k_row_bytes() + v_row = spec.v_row_bytes() + # Layer L's K block within the page starts at L * ps * (k_row+v_row). + # V block starts at +ps * k_row. + for L in range(spec.layer_num): + self.assertEqual( + spec.layer_k_offset_in_page(L, page_size=ps), + L * ps * (k_row + v_row), + ) + self.assertEqual( + spec.layer_v_offset_in_page(L, page_size=ps), + L * ps * (k_row + v_row) + ps * k_row, + ) + + def test_page_bytes(self): + spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + # page_bytes = page_size * entry_bytes (preserved invariant) + for ps in [1, 8, 64, 256]: + self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes()) + + +class TestBuildMHAViews(unittest.TestCase): + """Verify the 4-D view shape + strides at both page sizes.""" + + def _build(self, page_size, layer_num=3, head_num=2, head_dim=4, n_full_slots=64): + full = _make_mha_spec( + "full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim + ) + swa = _make_mha_spec( + "swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim + ) + # Pad to ensure max_slots % page_size == 0 in both sub-pools. + # entry_bytes is fixed per spec; size accordingly. + total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa], + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + return pool, full + + def test_view_shape_is_4d(self): + for ps in [1, 8]: + pool, spec = self._build(page_size=ps) + k_views, v_views = pool.mha_views_for("full") + self.assertEqual(len(k_views), spec.layer_num) + max_slots = pool.max_slots("full") + for L in range(spec.layer_num): + self.assertEqual(k_views[L].ndim, 4) + self.assertEqual( + tuple(k_views[L].shape), + (max_slots // ps, ps, spec.head_num, spec.head_dim), + ) + self.assertEqual( + tuple(v_views[L].shape), + (max_slots // ps, ps, spec.head_num, spec.v_head_dim), + ) + + def test_strides_at_page_size_1_match_envelope(self): + """At ps=1, the 4-D view's stride[0] equals what today's 3-D view's + stride[0] would have been (= entry_bytes / itemsize).""" + pool, spec = self._build(page_size=1, layer_num=4, head_num=3, head_dim=8) + k_views, _ = pool.mha_views_for("full") + itemsize = spec.store_dtype.itemsize + for L in range(spec.layer_num): + # stride[0] = page_bytes/itemsize = entry_bytes/itemsize at ps=1 + self.assertEqual(k_views[L].stride(0), spec.entry_bytes() // itemsize) + # stride[1] = k_row/itemsize (within-page token stride) + self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize) + # stride[2] = head_dim (head stride) + self.assertEqual(k_views[L].stride(2), spec.head_dim) + # stride[3] = 1 (innermost) + self.assertEqual(k_views[L].stride(3), 1) + + def test_strides_at_page_size_gt_1(self): + pool, spec = self._build(page_size=8, layer_num=4, head_num=3, head_dim=8) + k_views, _ = pool.mha_views_for("full") + itemsize = spec.store_dtype.itemsize + for L in range(spec.layer_num): + # page_bytes = 8 * 4 * (k_row + v_row); stride[0] = that / itemsize + self.assertEqual(k_views[L].stride(0), spec.page_bytes(8) // itemsize) + # token stride within layer L's K block = k_row/itemsize + self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize) + self.assertEqual(k_views[L].stride(2), spec.head_dim) + self.assertEqual(k_views[L].stride(3), 1) + + def test_distinct_layers_dont_alias_at_page_size_gt_1(self): + """Writes to layer 0 must not affect layer 1's K/V values (under + layer-major within-page layout).""" + pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4) + k_views, v_views = pool.mha_views_for("full") + # Set page 0, token 3, layer 0 K to a distinct pattern. + target_val = 0.5 + k_views[0][0, 3] = target_val + # Layer 1 K at the same (page, tok) should remain at default (0.0). + self.assertFalse(torch.all(k_views[1][0, 3] == target_val)) + self.assertTrue(torch.all(k_views[1][0, 3] == 0.0)) + # And layer 0 V at the same (page, tok) should remain at default. + self.assertFalse(torch.all(v_views[0][0, 3] == target_val)) + self.assertTrue(torch.all(v_views[0][0, 3] == 0.0)) + + def test_distinct_pages_dont_alias_at_page_size_gt_1(self): + """Writes to one page must not affect another page.""" + pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4) + k_views, _ = pool.mha_views_for("full") + # Set page 0, token 3, layer 0 K to a distinct pattern. + k_views[0][0, 3] = 1.25 + # Page 1, token 3, layer 0 K should remain at default. + self.assertTrue(torch.all(k_views[0][1, 3] == 0.0)) + + +class TestMoveKVCacheNative4D(unittest.TestCase): + """Verify ``move_kv_cache_native`` handles 4-D buffers at both + page_size=1 (degenerate envelope) and page_size>1 (layer-major).""" + + def _build_buffer( + self, page_size, layer_num=2, head_num=2, head_dim=4, n_full_slots=64 + ): + full = _make_mha_spec( + "full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim + ) + swa = _make_mha_spec( + "swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim + ) + total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa], + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + return pool + + def test_move_kv_cache_page_size_1(self): + pool = self._build_buffer(page_size=1, layer_num=2, head_num=2, head_dim=4) + k_views, v_views = pool.mha_views_for("full") + # Write distinct markers at source slots 5, 6. + for L in range(2): + k_views[L][5, 0] = float(L + 1) + v_views[L][5, 0] = -float(L + 1) + k_views[L][6, 0] = float(L + 10) + v_views[L][6, 0] = -float(L + 10) + # Move 5 -> 8 and 6 -> 9. + move_kv_cache_native( + k_views, + v_views, + tgt_loc=torch.tensor([8, 9], dtype=torch.int64), + src_loc=torch.tensor([5, 6], dtype=torch.int64), + page_size=1, + ) + for L in range(2): + self.assertTrue(torch.all(k_views[L][8, 0] == float(L + 1))) + self.assertTrue(torch.all(v_views[L][8, 0] == -float(L + 1))) + self.assertTrue(torch.all(k_views[L][9, 0] == float(L + 10))) + self.assertTrue(torch.all(v_views[L][9, 0] == -float(L + 10))) + + def test_move_kv_cache_page_size_gt_1(self): + ps = 8 + pool = self._build_buffer(page_size=ps, layer_num=2, head_num=2, head_dim=4) + k_views, v_views = pool.mha_views_for("full") + # Write markers at token ids 5 and 14 (different pages). + for L in range(2): + # token 5 = (page 0, tok 5) + k_views[L][0, 5] = float(L + 1) + v_views[L][0, 5] = -float(L + 1) + # token 14 = (page 1, tok 6) + k_views[L][1, 6] = float(L + 10) + v_views[L][1, 6] = -float(L + 10) + # Move token 5 -> token 23 (page 2, tok 7) and 14 -> 31 (page 3, tok 7). + move_kv_cache_native( + k_views, + v_views, + tgt_loc=torch.tensor([23, 31], dtype=torch.int64), + src_loc=torch.tensor([5, 14], dtype=torch.int64), + page_size=ps, + ) + for L in range(2): + # 23 = page 2, tok 7 + self.assertTrue(torch.all(k_views[L][2, 7] == float(L + 1))) + self.assertTrue(torch.all(v_views[L][2, 7] == -float(L + 1))) + # 31 = page 3, tok 7 + self.assertTrue(torch.all(k_views[L][3, 7] == float(L + 10))) + self.assertTrue(torch.all(v_views[L][3, 7] == -float(L + 10))) + + def test_move_kv_cache_3d_legacy_path_unchanged(self): + """move_kv_cache_native(3-D, page_size=1) must take the legacy + else-branch and be byte-identical to today.""" + k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] + v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] + for L in range(2): + k[L][5] = float(L + 1) + v[L][5] = -float(L + 1) + move_kv_cache_native( + k, + v, + tgt_loc=torch.tensor([7], dtype=torch.int64), + src_loc=torch.tensor([5], dtype=torch.int64), + page_size=1, + ) + for L in range(2): + self.assertTrue(torch.all(k[L][7] == float(L + 1))) + self.assertTrue(torch.all(v[L][7] == -float(L + 1))) + + +class TestByteIdentityAtPageSize1(unittest.TestCase): + """Verify that at page_size=1 the new 4-D view describes the SAME + physical bytes as the old 3-D view would have. The view + semantics differ (4-D vs 3-D shape) but the underlying byte layout is + identical — confirmed by manually computing expected byte offsets and + matching them against the 4-D view's strides + storage_offset. + """ + + def test_byte_addresses_match_envelope(self): + spec = _make_mha_spec("full", "up", layer_num=4, head_num=2, head_dim=4) + ps = 1 + # Build pool. + total = spec.entry_bytes() * 64 + spec.entry_bytes() * 32 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[ + spec, + _make_mha_spec("swa", "down", layer_num=2), + ], + device=_DEV, + enable_memory_saver=False, + page_size=ps, + ) + k_views, v_views = pool.mha_views_for("full") + # For each (layer, slot), compute the expected byte address under + # the envelope layout and verify the 4-D view's data_ptr + + # advanced indexing agrees. + max_slots = pool.max_slots("full") + itemsize = spec.store_dtype.itemsize + base_addr = pool._raw.data_ptr() + for L in range(spec.layer_num): + for s in range(0, max_slots, max(1, max_slots // 4)): + # Envelope: bytes for slot s, layer L's K start at: + # s * entry_bytes + L * (k_row + v_row) + expected_k_byte_offset = s * spec.entry_bytes() + L * ( + spec.k_row_bytes() + spec.v_row_bytes() + ) + # 4-D view: k_views[L][page=s, tok=0, head=0, dim=0] + # storage_offset of the element [s, 0, 0, 0]: + view_offset_elems = ( + k_views[L].storage_offset() + + s * k_views[L].stride(0) + + 0 * k_views[L].stride(1) + + 0 * k_views[L].stride(2) + + 0 * k_views[L].stride(3) + ) + view_byte_offset = view_offset_elems * itemsize + # 4-D view sits over `_raw.view(spec.store_dtype)`, which + # has data_ptr == _raw.data_ptr() (same backing storage). + self.assertEqual(view_byte_offset, expected_k_byte_offset) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py new file mode 100644 index 000000000..db64ceee3 --- /dev/null +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -0,0 +1,2729 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Unit tests for the shared-KV-pool v2 core: UnifiedKVPool views and +MultiEndedAllocator (virtual<->physical slot ids + eager compaction). + +CPU-only — no GPU / Triton needed (the allocator's data-copy delegates to a +fake kvcache here; the UnifiedKVPool view math is pure torch). + + python -m pytest test/registered/unit/mem_cache/test_multi_ended_allocator.py -v +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=8, suite="base-a-test-cpu") + +import random +import unittest + +import torch + +from sglang.srt.mem_cache.multi_ended_allocator import ( + MultiEndedAllocator, + UnifiedSWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedKVPool, +) + +_DEV = "cpu" + + +def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4): + return MHASubPoolSpec( + name=name, + layer_num=layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction=grow, + ) + + +def _make_mamba_spec(name, grow, layer_num=2): + return MambaSubPoolSpec( + name=name, + layer_num=layer_num, + conv_state_shapes=((4, 3),), + conv_dtype=torch.float32, + temporal_state_shape=(2, 2, 2), + temporal_dtype=torch.float32, + grow_direction=grow, + ) + + +class _FakeKVCache: + """Tracks, per *physical* slot, the virtual id whose data lives there. + `move_kv_cache(dst, src)` copies the marker — so after compaction we can + check that the data followed the relocation. + """ + + def __init__(self, max_slots: int): + # buf[p] == virtual id currently stored at physical slot p (-1 if free). + self.buf = torch.full((max_slots,), -1, dtype=torch.int64) + + def move_kv_cache(self, dst_loc: torch.Tensor, src_loc: torch.Tensor): + self.buf[dst_loc] = self.buf[src_loc].clone() + + +class TestUnifiedKVPoolViews(unittest.TestCase): + def test_min_slot_index_and_disjoint_bytes(self): + full = _make_mha_spec("full", "up", layer_num=4) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + entry_max = max(full.entry_bytes(), mamba.entry_bytes()) + total = full.entry_bytes() * 64 + mamba.entry_bytes() * 16 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + for s in (full, mamba): + min_idx = pool.min_slot_index(s.name) + # real data of every pool begins at bytes >= entry_max + self.assertGreaterEqual(min_idx * s.entry_bytes(), entry_max) + self.assertGreater(pool.max_slots(s.name), min_idx) + + def test_mha_view_roundtrip(self): + full = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + swa = _make_mha_spec("swa", "down", layer_num=2, head_num=2, head_dim=4) + total = full.entry_bytes() * 32 + swa.entry_bytes() * 32 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa], + device=_DEV, + enable_memory_saver=False, + ) + k_full, v_full = pool.mha_views_for("full") + k_swa, v_swa = pool.mha_views_for("swa") + self.assertEqual(len(k_full), 3) + self.assertEqual(len(k_swa), 2) + # Write distinct patterns into a couple of slots/layers of "full" and + # confirm they read back, and that "swa" was not disturbed. "full" and + # "swa" share the buffer from byte 0; "full" grows up (low slots) and + # "swa" grows down, so the allocator places "swa" at the high slots. + # Mirror that here — a low "swa" slot would byte-overlap "full" slot 5 + # (a configuration the byte-frontier coordination never produces). + swa_slot = pool.max_slots("swa") - 1 + for lyr in range(3): + k_full[lyr][5] = float(lyr + 1) + v_full[lyr][5] = float(-(lyr + 1)) + for lyr in range(2): + k_swa[lyr][swa_slot] = 99.0 + for lyr in range(3): + self.assertTrue(torch.all(k_full[lyr][5] == float(lyr + 1))) + self.assertTrue(torch.all(v_full[lyr][5] == float(-(lyr + 1)))) + for lyr in range(2): + self.assertTrue(torch.all(k_swa[lyr][swa_slot] == 99.0)) + # "full" slot 5 layer-0 K must not alias "full" slot 6 layer-0 K + self.assertFalse(torch.all(k_full[0][6] == float(1))) + + def test_mamba_view_shapes(self): + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=3) + total = full.entry_bytes() * 16 + mamba.entry_bytes() * 8 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + conv_views, temporal_view = pool.mamba_views_for("mamba") + max_slots = pool.max_slots("mamba") + self.assertEqual(len(conv_views), 1) + self.assertEqual(tuple(conv_views[0].shape), (3, max_slots, 4, 3)) + self.assertEqual(tuple(temporal_view.shape), (3, max_slots, 2, 2, 2)) + # roundtrip a write at (layer=1, slot=4) + conv_views[0][1, 4] = 3.5 + temporal_view[2, 6] = -1.25 + self.assertTrue(torch.all(conv_views[0][1, 4] == 3.5)) + self.assertTrue(torch.all(temporal_view[2, 6] == -1.25)) + + +class TestMultiEndedAllocator(unittest.TestCase): + def _build_pair(self, n_full_slots=64, n_mamba_slots=16): + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + total = full.entry_bytes() * n_full_slots + mamba.entry_bytes() * n_mamba_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + full_alloc = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + ) + mamba_alloc = MultiEndedAllocator( + kvcache=mamba_kv, + unified_buffer=pool, + sub_pool_name="mamba", + device=_DEV, + is_id_owner=True, + ) + full_alloc.bind_peer(mamba_alloc) + mamba_alloc.bind_peer(full_alloc) + return pool, full_alloc, mamba_alloc, full_kv, mamba_kv + + def _check_invariants(self, alloc: MultiEndedAllocator, kv: _FakeKVCache): + v2p = alloc.virtual_to_physical + p2v = alloc.physical_to_virtual + # live virtual ids = those with v2p != -1, excluding the reserved id 0. + live_v = [ + v for v in range(1, alloc.num_virtual_ids) if int(v2p[v].item()) != -1 + ] + # mutual-inverse on the live set + for v in live_v: + p = int(v2p[v].item()) + self.assertEqual(int(p2v[p].item()), v, f"p2v[{p}] != {v}") + # data followed any relocations + self.assertEqual(int(kv.buf[p].item()), v, f"kv.buf[{p}] != {v}") + # allocated physical range is hole-free + matches live count + if alloc.grow_direction == "up": + alloc_lo, alloc_hi = alloc.min_slot_index, alloc.watermark_physical + else: + alloc_lo, alloc_hi = alloc.watermark_physical + 1, alloc.max_slots + self.assertEqual(alloc_hi - alloc_lo, len(live_v)) + for p in range(alloc_lo, alloc_hi): + self.assertNotEqual(int(p2v[p].item()), -1, f"hole at physical {p}") + # free virtual ids ∪ live = [min_slot_index, max_slots) + free_set = set(int(x) for x in alloc.free_virtual_ids.tolist()) + self.assertEqual( + free_set | set(live_v), + set(range(alloc.min_slot_index, alloc.max_slots)), + ) + self.assertEqual(free_set & set(live_v), set()) + + def _alloc(self, alloc: MultiEndedAllocator, kv: _FakeKVCache, n: int): + avail = alloc.available_size() + v = alloc.alloc(n) + if n > avail: + self.assertIsNone(v) + return None + self.assertIsNotNone(v) + self.assertEqual(int(v.numel()), n) + # stamp the data marker at each new physical slot + p = alloc.virtual_to_physical[v] + kv.buf[p] = v + return v + + def _free(self, alloc: MultiEndedAllocator, kv: _FakeKVCache, v: torch.Tensor): + p = alloc.virtual_to_physical[v] + kv.buf[p] = -1 # the freed virtual id's data is gone + alloc.free(v) + + def test_basic_alloc_free_compaction(self): + _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair() + # alloc three batches on the full side + a = self._alloc(full_alloc, full_kv, 3) + b = self._alloc(full_alloc, full_kv, 5) + c = self._alloc(full_alloc, full_kv, 2) + self._check_invariants(full_alloc, full_kv) + # free the middle batch -> forces eager compaction (boundary slots move in) + self._free(full_alloc, full_kv, b) + self._check_invariants(full_alloc, full_kv) + # `a` and `c` virtual ids unchanged; their physical slots may have moved. + for v in a.tolist() + c.tolist(): + self.assertNotEqual(int(full_alloc.virtual_to_physical[v].item()), -1) + # free the boundary batch (no relocation needed) + self._free(full_alloc, full_kv, c) + self._check_invariants(full_alloc, full_kv) + self._free(full_alloc, full_kv, a) + self._check_invariants(full_alloc, full_kv) + self.assertEqual(full_alloc.allocated_count(), 0) + + def test_grow_down_side(self): + _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair() + a = self._alloc(mamba_alloc, mamba_kv, 2) + b = self._alloc(mamba_alloc, mamba_kv, 3) + c = self._alloc(mamba_alloc, mamba_kv, 1) + self._check_invariants(mamba_alloc, mamba_kv) + self._free(mamba_alloc, mamba_kv, b) # interior -> compaction + self._check_invariants(mamba_alloc, mamba_kv) + self._free(mamba_alloc, mamba_kv, a) + self._free(mamba_alloc, mamba_kv, c) + self._check_invariants(mamba_alloc, mamba_kv) + self.assertEqual(mamba_alloc.allocated_count(), 0) + + def test_byte_frontier_coordination(self): + # full has 8 slots' worth of bytes; mamba's entry is larger, so a few + # mamba allocs should shrink full's available_size below its slot headroom. + _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair( + n_full_slots=8, n_mamba_slots=8 + ) + full_avail0 = full_alloc.available_size() + self._alloc(mamba_alloc, mamba_kv, 3) + self.assertLess(full_alloc.available_size(), full_avail0) + # over-alloc the full side -> None + self.assertIsNone(full_alloc.alloc(full_alloc.available_size() + 1)) + + def test_randomized(self): + rng = random.Random(0xC0FFEE) + _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair( + n_full_slots=48, n_mamba_slots=24 + ) + live_full = [] # list of virtual-id tensors still allocated + live_mamba = [] + for _ in range(400): + side = rng.random() < 0.6 # 60% full + alloc, kv, live = ( + (full_alloc, full_kv, live_full) + if side + else (mamba_alloc, mamba_kv, live_mamba) + ) + if rng.random() < 0.55 or not live: + n = rng.randint(1, 5) + v = self._alloc(alloc, kv, n) + if v is not None: + live.append(v) + else: + idx = rng.randrange(len(live)) + v = live.pop(idx) + self._free(alloc, kv, v) + self._check_invariants(full_alloc, full_kv) + self._check_invariants(mamba_alloc, mamba_kv) + # drain + for live, alloc, kv in ( + (live_full, full_alloc, full_kv), + (live_mamba, mamba_alloc, mamba_kv), + ): + for v in live: + self._free(alloc, kv, v) + self._check_invariants(alloc, kv) + self.assertEqual(alloc.allocated_count(), 0) + + def _build_lazy_full(self, n_full_slots=64, n_mamba_slots=16, move_cap=2): + """A LAZY-compaction 'full' (grow-up) allocator + its peer, with a + small per-call move cap so flushes are partial (the retract-pressure + regime where the ghost bug appears).""" + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + total = full.entry_bytes() * n_full_slots + mamba.entry_bytes() * n_mamba_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + full_alloc = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + lazy_compaction=True, + ) + mamba_alloc = MultiEndedAllocator( + kvcache=mamba_kv, + unified_buffer=pool, + sub_pool_name="mamba", + device=_DEV, + is_id_owner=True, + lazy_compaction=True, + ) + full_alloc.bind_peer(mamba_alloc) + mamba_alloc.bind_peer(full_alloc) + full_alloc._lazy_max_moves_per_call = move_cap + return pool, full_alloc, full_kv + + def test_lazy_retract_churn_no_ghost(self): + """Regression: heavy free/flush churn in LAZY mode must never + leave a 'ghost' page (p2v<0 not registered in _free_phys_pages or + _pending_reuse). This reproduces the retract×lazy-compaction pattern. + The fail-fast ghost invariant is enabled so any ghost trips at the + CREATING op (free_lazy / flush_exit), not at a later survivor walk. + + CPU scope: the GPU write-race *urgent* path (unfired CUDA events -> + _pending_reuse) needs a real device; this covers the no-event + move/absorb/free/released_fired bookkeeping. + """ + rng = random.Random(0x5373) + _, alloc, kv = self._build_lazy_full(n_full_slots=64, move_cap=2) + live = [] + # Bias toward near-capacity occupancy (watermark high, holes pile + # up) then churn free/flush — the conditions that surface the bug. + for _ in range(3000): + avail = alloc.available_size() + if (rng.random() < 0.55 and avail > 0) or not live: + n = rng.randint(1, min(5, max(1, avail))) + v = self._alloc(alloc, kv, n) + if v is not None: + live.append(v) + else: + v = live.pop(rng.randrange(len(live))) + self._free(alloc, kv, v) # often non-boundary -> a hole + if rng.random() < 0.5: + alloc.flush_opportunistic() # partial flush (move_cap) + # Drain to quiescence; must end empty AND ghost-free. + for v in live: + self._free(alloc, kv, v) + for _ in range(64): + alloc.flush_opportunistic() + self.assertEqual(alloc.allocated_count(), 0) + + def test_double_free_raises(self): + _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair() + v = self._alloc(full_alloc, full_kv, 3) + self._free(full_alloc, full_kv, v) + with self.assertRaises(AssertionError): + full_alloc.free(v) + + # -- `out=` parameter regression tests -- + + def test_translate_kv_loc_with_out_writes_inplace(self): + """REGRESSION: `translate_kv_loc(virt, out=buf)` must + modify `buf` in place AND preserve `buf.data_ptr()` — the buffer- + stability invariant for cuda-graph capture.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) + ptr_before = buf.data_ptr() + ret = full_alloc.translate_kv_loc(v, out=buf) + self.assertIs(ret, buf, "must return the `out=` buffer, not a fresh tensor") + self.assertEqual( + buf.data_ptr(), ptr_before, "out= buffer's data_ptr must be stable" + ) + # Result matches v2p directly (page_size == 1 here) + expected = full_alloc.virtual_to_physical[v] + self.assertTrue(bool((buf == expected).all().item())) + + def test_translate_kv_loc_without_out_returns_fresh_tensor(self): + """REGRESSION: without `out=`, behavior returns a fresh tensor.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + ret = full_alloc.translate_kv_loc(v) + # Fresh tensor: different storage from v2p table + self.assertNotEqual(ret.data_ptr(), full_alloc.virtual_to_physical.data_ptr()) + expected = full_alloc.virtual_to_physical[v] + self.assertTrue(bool((ret == expected).all().item())) + + def test_translate_kv_loc_out_matches_no_out(self): + """REGRESSION: result of translate_kv_loc(v, out=buf) byte-equals + translate_kv_loc(v).""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) + with_out = full_alloc.translate_kv_loc(v, out=buf) + no_out = full_alloc.translate_kv_loc(v) + self.assertTrue(bool((with_out == no_out).all().item())) + + def test_translate_kv_loc_dtype_assertion(self): + """REGRESSION: wrong-dtype `out=` (int32 instead of int64) raises + AssertionError. Guards against the copy/paste hazard where someone + might allocate the full-physical buffer with the SWA int32 pattern.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + wrong_dtype = torch.empty(v.shape, dtype=torch.int32, device=_DEV) + with self.assertRaises(AssertionError): + full_alloc.translate_kv_loc(v, out=wrong_dtype) + + def test_translate_kv_loc_shape_assertion(self): + """REGRESSION: mismatched `out=` shape raises AssertionError.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + wrong_shape = torch.empty((v.numel() + 1,), dtype=torch.int64, device=_DEV) + with self.assertRaises(AssertionError): + full_alloc.translate_kv_loc(v, out=wrong_shape) + + # REGRESSION: `translate_kv_loc(buf, out=buf)` — same + # tensor for input and output — is the canonical in-place form used by + # the cuda-graph capture/replay paths in `triton_backend.py`: + # + # self._translate_kv_loc(kv_indices, out=kv_indices) + # + # A naive implementation routes this through + # `torch.index_select(v2p, 0, virt_tokens, out=out)`, which crashes with + # "unsupported operation: some elements of the input tensor and the + # written-to tensor refer to a single memory location" + # because index_select does NOT support aliasing between `index` and + # `out`. Fix: gather into a transient buffer then `out.copy_(tmp)`. + def test_translate_kv_loc_with_out_aliasing_input(self): + """REGRESSION: in-place form `translate_kv_loc(buf, out=buf)` must + succeed and produce identical results to the no-out form.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v_orig = self._alloc(full_alloc, full_kv, 5).clone() + # Save the expected output (no-out form) before mutating `buf`. + expected = full_alloc.translate_kv_loc(v_orig) + # Now exercise the aliasing form: buf serves as BOTH input and out. + buf = v_orig.clone() + ptr_before = buf.data_ptr() + ret = full_alloc.translate_kv_loc(buf, out=buf) + self.assertIs(ret, buf) + self.assertEqual( + buf.data_ptr(), + ptr_before, + "out= buffer's data_ptr must be stable (cuda-graph invariant)", + ) + self.assertTrue( + bool((buf == expected).all().item()), + "in-place result must equal no-out result", + ) + + # Tombstone-safety clamp regression. + # + # The captured cuda-graph paths (full-layer set_kv_buffer elif, + # init_forward_metadata_*_cuda_graph kv_indices translate, init_new + # precompute) all eventually call `translate_kv_loc` against `v2p_full`. + # Padded / stale-tail entries in the cuda-graph input buffers can carry + # virtual ids whose v2p entries got tombstoned (-1) by free/compaction + # between replays. Without a clamp, the captured `k_buffer[result[i]]` + # would index at -1 (illegal memory access). The clamp routes those to + # physical slot 0 (the reserved padding sink under the `min_slot_index` + # invariant — bytes [0, entry_max) hold no real data). + # These tests lock in the clamp contract so a future refactor can't + # quietly remove it and re-introduce the crash. + def test_translate_kv_loc_clamps_tombstoned_v2p(self): + """`translate_kv_loc` must clamp `v2p[v] == -1` entries to 0 (the + padding sink). Required for cuda-graph capture safety.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5) + # Inject a tombstone at one of the live virtual id positions WITHOUT + # going through `free` (which would also touch p2v / compaction). + # This emulates the steady-state where a captured-graph input + # buffer's padded/stale entries reference virtual ids that have + # since been tombstoned by free/compaction. + v_tombstoned = int(v[2].item()) + full_alloc.virtual_to_physical[v_tombstoned] = -1 + # No-out form: result must clamp. + out = full_alloc.translate_kv_loc(v) + self.assertTrue( + bool((out >= 0).all().item()), + f"translate_kv_loc must clamp tombstoned entries to >=0, got {out.tolist()}", + ) + self.assertEqual( + int(out[2].item()), + 0, + "tombstoned virtual id must map to slot 0 (padding sink)", + ) + + def test_translate_kv_loc_with_out_clamps_tombstoned_v2p(self): + """`translate_kv_loc(..., out=buf)` (the captured-graph path) must + clamp tombstoned entries in-place.""" + _, full_alloc, _, full_kv, _ = self._build_pair() + v = self._alloc(full_alloc, full_kv, 5).clone() + full_alloc.virtual_to_physical[int(v[1].item())] = -1 + buf = torch.empty_like(v) + ret = full_alloc.translate_kv_loc(v, out=buf) + self.assertIs(ret, buf) + self.assertTrue( + bool((buf >= 0).all().item()), + "out= path must clamp tombstoned entries", + ) + self.assertEqual(int(buf[1].item()), 0) + + +# --------------------------------------------------------------------------- +# Shared SWA composite — unit tests +# --------------------------------------------------------------------------- + + +class _FakeUnifiedSWAKVPool: + """Minimal stand-in for `UnifiedSWAKVPool` that the composite allocator + needs. Exposes the two sub-pool views (each a `_FakeKVCache` with an + `attach_allocator` no-op) and an `attach_allocators` setter. + + CPU-only — avoids constructing a real `UnifiedMHATokenToKVPool` (which + instantiates `MHATokenToKVPool` and is heavier than these tests need). + """ + + class _SubKV(_FakeKVCache): + def __init__(self, max_slots): + super().__init__(max_slots) + self.allocator = None + + def attach_allocator(self, allocator): + self.allocator = allocator + + def __init__(self, shared_pool: UnifiedKVPool): + self.full_kv_pool = self._SubKV(shared_pool.max_slots("full")) + self.swa_kv_pool = self._SubKV(shared_pool.max_slots("swa")) + self._full_allocator = None + self._swa_allocator = None + + def attach_allocators(self, *, full_allocator, swa_allocator): + self._full_allocator = full_allocator + self._swa_allocator = swa_allocator + + +class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): + """Tests for the SWA composite — joint byte-budget, slot-conservation + leak invariant, tombstone semantics for `free_swa`, divergent compaction + of the two sub-pools, and the alloc-rollback path. + + These tests cover the core invariants: joint byte-budget, + slot-conservation, the `schedulable_*` split, and watermark + rollback.""" + + def _build( + self, + n_full_slots=32, + n_swa_slots=16, + full_layer_num=4, + swa_layer_num=2, + head_num=2, + head_dim=4, + ): + full_spec = MHASubPoolSpec( + name="full", + layer_num=full_layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=swa_layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction="down", + ) + total = ( + n_full_slots * full_spec.entry_bytes() + + n_swa_slots * swa_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + device=_DEV, + full_max_total_num_tokens=n_full_slots, + swa_max_total_num_tokens=n_swa_slots, + need_sort=False, + forward_stream=None, + ) + return pool, allocator, kvcache + + def _alloc(self, allocator, kvcache, n): + """Allocate N virtual ids; stamp the data marker on both sub-pools.""" + v = allocator.alloc(n) + if v is None: + return None + full_phys = allocator.full_attn_allocator.virtual_to_physical[v] + swa_phys = allocator.swa_attn_allocator.virtual_to_physical[v] + kvcache.full_kv_pool.buf[full_phys] = v + kvcache.swa_kv_pool.buf[swa_phys] = v + return v + + def _free(self, allocator, kvcache, v): + """Erase markers on both sub-pools (mirror compaction's no-data-at + -freed-slot invariant), then call the composite's free.""" + full_phys = allocator.full_attn_allocator.virtual_to_physical[v] + swa_phys = allocator.swa_attn_allocator.virtual_to_physical[v] + # erase only the LIVE swa entries (`free_swa` may have already + # tombstoned some of `v`). + valid_swa = swa_phys[swa_phys >= 0] + kvcache.full_kv_pool.buf[full_phys] = -1 + kvcache.swa_kv_pool.buf[valid_swa] = -1 + allocator.free(v) + + def _check_sub_pool_invariants(self, sub, kv): + """Per-sub-pool: v2p ∘ p2v identity on the live set, hole-free + allocated band, data followed relocations.""" + v2p = sub.virtual_to_physical + p2v = sub.physical_to_virtual + live_v = [v for v in range(1, sub.num_virtual_ids) if int(v2p[v].item()) != -1] + for v in live_v: + p = int(v2p[v].item()) + self.assertEqual(int(p2v[p].item()), v) + # data marker followed any relocation + self.assertEqual(int(kv.buf[p].item()), v) + if sub.grow_direction == "up": + lo, hi = sub.min_slot_index, sub.watermark_physical + else: + lo, hi = sub.watermark_physical + 1, sub.max_slots + self.assertEqual(hi - lo, len(live_v)) + for p in range(lo, hi): + self.assertNotEqual(int(p2v[p].item()), -1) + + # 1. Both peers hold a physical slot per virtual after composite alloc. + def test_swa_alloc_both_peers_hold(self): + _, allocator, _ = self._build() + v = allocator.alloc(3) + self.assertIsNotNone(v) + self.assertEqual(int(v.numel()), 3) + full_v2p = allocator.full_attn_allocator.virtual_to_physical + swa_v2p = allocator.swa_attn_allocator.virtual_to_physical + for vi in v.tolist(): + self.assertGreaterEqual(int(full_v2p[vi].item()), 0) + self.assertGreaterEqual(int(swa_v2p[vi].item()), 0) + # Full sub-pool is id-owner -> the minted ids are out of free_virtual_ids. + free_full = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertTrue(set(v.tolist()).isdisjoint(free_full)) + # Swa sub-pool is non-owner -> free_virtual_ids is None. + self.assertIsNone(allocator.swa_attn_allocator.free_virtual_ids) + + # 2. Composite `free` releases both sub-pools' v2p; the virtual goes back + # to the full id-owner's free list. + def test_swa_free_releases_both(self): + _, allocator, kvcache = self._build() + v = self._alloc(allocator, kvcache, 3) + self._free(allocator, kvcache, v) + for vi in v.tolist(): + self.assertEqual( + int(allocator.full_attn_allocator.virtual_to_physical[vi].item()), -1 + ) + self.assertEqual( + int(allocator.swa_attn_allocator.virtual_to_physical[vi].item()), -1 + ) + free_full = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertTrue(set(v.tolist()).issubset(free_full)) + + # 3. `free_swa` tombstones swa side only; virtual + full-physical stay live. + def test_swa_free_swa_keeps_virtual_alive(self): + _, allocator, kvcache = self._build() + v = self._alloc(allocator, kvcache, 3) + # Tombstone the middle one. Erase its swa marker first (compaction + # will run inside `free_swa`). + target = v[1:2] + target_swa = allocator.swa_attn_allocator.virtual_to_physical[target] + kvcache.swa_kv_pool.buf[target_swa] = -1 + allocator.free_swa(target) + tgt = int(target.item()) + # full side still bound: + self.assertGreaterEqual( + int(allocator.full_attn_allocator.virtual_to_physical[tgt].item()), 0 + ) + # swa side tombstoned: + self.assertEqual( + int(allocator.swa_attn_allocator.virtual_to_physical[tgt].item()), -1 + ) + # NOT recycled to the id-owner's free list yet: + free_full = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertNotIn(tgt, free_full) + # composite `free` of the same virtual still works (filters out + # already-tombstoned on the swa side). + full_phys = int(allocator.full_attn_allocator.virtual_to_physical[tgt].item()) + kvcache.full_kv_pool.buf[full_phys] = -1 + allocator.free(target) + # now in free list: + free_full = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertIn(tgt, free_full) + + # 4. Compaction diverges between the two sub-pools (each runs its own). + def test_swa_compaction_diverges_physical_layout(self): + _, allocator, kvcache = self._build() + a = self._alloc(allocator, kvcache, 1) + b = self._alloc(allocator, kvcache, 1) + c = self._alloc(allocator, kvcache, 1) + # Snapshot swa-side physical for c BEFORE we free_swa(b). + c_swa_before = int(allocator.swa_attn_allocator.virtual_to_physical[c].item()) + c_full_before = int(allocator.full_attn_allocator.virtual_to_physical[c].item()) + # Tombstone b on swa only. + b_swa = allocator.swa_attn_allocator.virtual_to_physical[b] + kvcache.swa_kv_pool.buf[b_swa] = -1 + allocator.free_swa(b) + # c's full-physical UNCHANGED (full side did not compact): + self.assertEqual( + int(allocator.full_attn_allocator.virtual_to_physical[c].item()), + c_full_before, + ) + # c's swa-physical MUST have moved (b was interior to swa's + # allocated band on grow-down: a then b then c means b is between + # them; freeing b triggers compaction relocating c into b's slot). + c_swa_after = int(allocator.swa_attn_allocator.virtual_to_physical[c].item()) + self.assertNotEqual(c_swa_after, c_swa_before) + # Per-sub-pool invariants still hold. + self._check_sub_pool_invariants( + allocator.full_attn_allocator, kvcache.full_kv_pool + ) + self._check_sub_pool_invariants( + allocator.swa_attn_allocator, kvcache.swa_kv_pool + ) + + # 5. Byte-frontier coordination — peer-aware available_size shrinks as + # the peer grows. + def test_swa_byte_frontier_coordination(self): + _, allocator, kvcache = self._build(n_full_slots=8, n_swa_slots=8) + avail0 = allocator.available_size() + # Allocate enough that the joint budget visibly tightens. + self._alloc(allocator, kvcache, 3) + self.assertLess(allocator.available_size(), avail0) + # Joint budget enforcement: over-alloc returns None. + self.assertIsNone(allocator.alloc(allocator.available_size() + 1)) + + # 6. Randomized stress — invariants under mixed alloc / free / free_swa. + def test_swa_randomized_alloc_free_freeswa(self): + rng = random.Random(0xBADBEE) + _, allocator, kvcache = self._build( + n_full_slots=48, n_swa_slots=24, full_layer_num=3, swa_layer_num=3 + ) + live = [] # list of (virtual-id tensor) + for _ in range(400): + r = rng.random() + if r < 0.5 or not live: # alloc + n = rng.randint(1, 4) + v = self._alloc(allocator, kvcache, n) + if v is not None: + live.append(("live", v)) + elif r < 0.8: # composite free + idx = rng.randrange(len(live)) + kind, v = live.pop(idx) + self._free(allocator, kvcache, v) + else: # free_swa on some entries + idx = rng.randrange(len(live)) + kind, v = live[idx] + if kind != "live": + continue + # Tombstone all of v on swa only. + swa_phys = allocator.swa_attn_allocator.virtual_to_physical[v] + kvcache.swa_kv_pool.buf[swa_phys] = -1 + allocator.free_swa(v) + live[idx] = ("swa_tomb", v) + # Invariants after every op. + self._check_sub_pool_invariants( + allocator.full_attn_allocator, kvcache.full_kv_pool + ) + self._check_sub_pool_invariants( + allocator.swa_attn_allocator, kvcache.swa_kv_pool + ) + # Slot-conservation invariant balances at all times. NOTE: the + # leak view is now `_conserve_*` — the public `full/swa_available_size()` + # returns `min(conserve, schedulable)` (physical), which can be + # strictly smaller (e.g. the reserved sink page). + self.assertEqual( + allocator._conserve_full_available_size(), + allocator._full_max_total_num_tokens + - allocator.full_attn_allocator.allocated_count(), + ) + self.assertEqual( + allocator._conserve_swa_available_size(), + allocator._swa_max_total_num_tokens + - allocator.swa_attn_allocator.allocated_count(), + ) + # Drain. + for _, v in live: + self._free(allocator, kvcache, v) + self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0) + self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0) + + # 7. Joint byte-budget pre-check. + def test_swa_joint_byte_budget_pre_check(self): + # Pick sizes where the byte gap, not slot-index headroom, is the bind. + full_spec = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="down", + ) + n_full, n_swa = 10, 10 + total = n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes() + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + device=_DEV, + full_max_total_num_tokens=n_full, + swa_max_total_num_tokens=n_swa, + need_sort=False, + forward_stream=None, + ) + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + # Compute the "naive min" against the joint budget — at idle, the + # joint budget is strictly less than min(full.available, swa.available) + # because the joint uses (entry_full + entry_swa) per slot. + naive = min(fa.available_size(), sa.available_size()) + joint = allocator.available_size() + # The joint must be no greater than naive (typically strictly less). + self.assertLessEqual(joint, naive) + # And it must equal `gap_bytes // (entry_full + entry_swa)` clamped + # by slot-room. + gap = sa._byte_low_frontier() - fa._byte_high_frontier() + expected = min( + gap // (fa.entry_bytes + sa.entry_bytes), + fa.max_slots - fa.min_slot_index - fa.allocated_count(), + sa.max_slots - sa.min_slot_index - sa.allocated_count(), + ) + self.assertEqual(joint, expected) + + # 8. Watermark rollback on partial alloc failure. + def test_swa_alloc_swa_failure_is_fail_loud(self): + """The SWA composite runs a tight JOINT pre-check before allocating, so + a swa-side ``alloc_with_virtual`` failure after the full-side alloc can + only mean an internal-state inconsistency. By design (``UnifiedSWA.alloc``: + "assert rather than silently rollback") that surfaces as a loud error, + NOT a silent ``None`` / rollback — masking it would hide the bug. The + real ``alloc_with_virtual`` self-asserts on shortfall, so the production + path is fail-loud too; here we force the failure to prove it propagates. + """ + _, allocator, kvcache = self._build() + sa = allocator.swa_attn_allocator + original = sa.alloc_with_virtual + + def _bomb(virtual_ids): + raise AssertionError("synthetic alloc_with_virtual failure") + + sa.alloc_with_virtual = _bomb + try: + with self.assertRaises(AssertionError): + allocator.alloc(3) + finally: + sa.alloc_with_virtual = original + + # -- `out=` parameter regression tests for the SWA composite -- + + def test_swa_translate_kv_loc_with_out_writes_inplace(self): + """REGRESSION: composite delegates to base-class + translate_kv_loc with `out=` passthrough. Result lands in `buf`.""" + _, allocator, _ = self._build() + v = allocator.alloc(4) + self.assertIsNotNone(v) + buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) + ptr_before = buf.data_ptr() + ret = allocator.translate_kv_loc(v, out=buf) + self.assertIs(ret, buf) + self.assertEqual(buf.data_ptr(), ptr_before) + expected = allocator.translate_kv_loc(v) + self.assertTrue(bool((buf == expected).all().item())) + + def test_swa_translate_loc_from_full_to_swa_with_out_writes_inplace(self): + """REGRESSION: `translate_loc_from_full_to_swa(v, out=buf)` + must modify `buf` in place AND preserve `buf.data_ptr()`. `out=` + buffer MUST be int32 (matches SWA Triton kernel contract).""" + _, allocator, _ = self._build() + v = allocator.alloc(4) + self.assertIsNotNone(v) + buf = torch.empty(v.shape, dtype=torch.int32, device=_DEV) + ptr_before = buf.data_ptr() + ret = allocator.translate_loc_from_full_to_swa(v, out=buf) + self.assertIs(ret, buf) + self.assertEqual(buf.data_ptr(), ptr_before) + # Byte-identical to the no-out form: + no_out = allocator.translate_loc_from_full_to_swa(v) + self.assertEqual(no_out.dtype, torch.int32) + self.assertTrue(bool((buf == no_out).all().item())) + + def test_swa_translate_loc_from_full_to_swa_dtype_assertion(self): + """REGRESSION: wrong-dtype `out=` (int64 instead of int32) + raises AssertionError. Guards against accidentally reusing the int64 + full-physical buffer pattern for the SWA precompute.""" + _, allocator, _ = self._build() + v = allocator.alloc(4) + self.assertIsNotNone(v) + wrong_dtype = torch.empty(v.shape, dtype=torch.int64, device=_DEV) + with self.assertRaises(AssertionError): + allocator.translate_loc_from_full_to_swa(v, out=wrong_dtype) + + # Tombstone-safety clamp for SWA — mirrors the full-side test + # in `TestMultiEndedAllocator`. The captured SWA attention kernel reads + # `swa_k_buffer[result[i]]` at replay; without the clamp, a tombstoned + # `v2p_swa[v] == -1` would index at `swa_k_buffer[-1]` (illegal access). + def test_swa_translate_loc_from_full_to_swa_clamps_tombstoned(self): + _, allocator, _ = self._build() + v = allocator.alloc(4) + self.assertIsNotNone(v) + # Inject a tombstone on the swa side at one of the live virtual ids. + v_tomb = int(v[1].item()) + allocator.swa_attn_allocator.virtual_to_physical[v_tomb] = -1 + # No-out form: result must be int32 AND every entry >= 0. + out = allocator.translate_loc_from_full_to_swa(v) + self.assertEqual(out.dtype, torch.int32) + self.assertTrue( + bool((out >= 0).all().item()), + "translate_loc_from_full_to_swa must clamp tombstoned to >=0", + ) + self.assertEqual(int(out[1].item()), 0) + # out= form (int32 buffer) must also clamp. + buf = torch.empty(v.shape, dtype=torch.int32, device=_DEV) + ret = allocator.translate_loc_from_full_to_swa(v, out=buf) + self.assertIs(ret, buf) + self.assertTrue(bool((buf >= 0).all().item())) + self.assertEqual(int(buf[1].item()), 0) + + +# --------------------------------------------------------------------------- +# page_size > 1 — paged unit tests +# --------------------------------------------------------------------------- + + +class TestPagedMultiEndedAllocator(unittest.TestCase): + """Per-sub-pool paged tests for `MultiEndedAllocator(page_size=...)`. + + All tests use ``page_size = 8`` against a buffer sized for ~16 pages per + sub-pool. Invariants are page-granular: free-list, v2p/p2v tables, and + compaction operate on pages. The external API (alloc → token ids, free + takes token ids) is byte-identical to the page_size == 1 case. + """ + + PAGE_SIZE = 8 + + def _build(self, n_full_pages=16, n_swa_pages=8, full_layer_num=2, swa_layer_num=2): + full_spec = MHASubPoolSpec( + name="full", + layer_num=full_layer_num, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=swa_layer_num, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="down", + ) + # entry_bytes_per_page = layer_num * (k_row + v_row) * page_size + # We size the buffer to fit `n_full_pages` full-pages + `n_swa_pages` + # swa-pages (token-equivalent: n_*_pages * page_size). + total = ( + n_full_pages * self.PAGE_SIZE * full_spec.entry_bytes() + + n_swa_pages * self.PAGE_SIZE * swa_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + swa_kv = _FakeKVCache(pool.max_slots("swa")) + full_alloc = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + page_size=self.PAGE_SIZE, + ) + swa_alloc = MultiEndedAllocator( + kvcache=swa_kv, + unified_buffer=pool, + sub_pool_name="swa", + device=_DEV, + is_id_owner=True, + page_size=self.PAGE_SIZE, + ) + full_alloc.bind_peer(swa_alloc) + swa_alloc.bind_peer(full_alloc) + return pool, full_alloc, swa_alloc, full_kv, swa_kv + + def _stamp_tokens( + self, alloc: MultiEndedAllocator, kv: _FakeKVCache, v_tokens: torch.Tensor + ): + """Mark `kv.buf[phys_token] = some_unique_id` for every returned + token. Uses the alloc's v2p_page table to compute physical tokens.""" + if v_tokens.numel() == 0: + return + ps = alloc.page_size + virt_pages = v_tokens // ps + offsets = v_tokens % ps + phys_pages = alloc.virtual_to_physical[virt_pages] + phys_tokens = phys_pages * ps + offsets + kv.buf[phys_tokens] = v_tokens + + def _check_invariants( + self, alloc: MultiEndedAllocator, kv: _FakeKVCache, stamped_tokens: dict + ): + v2p = alloc.virtual_to_physical + p2v = alloc.physical_to_virtual + ps = alloc.page_size + # Live virtual pages (excluding the reserved padding page 0). + live_v_pages = [ + v for v in range(1, alloc.num_pages) if int(v2p[v].item()) != -1 + ] + # Mutual inverse on the live page set. + for v_page in live_v_pages: + p_page = int(v2p[v_page].item()) + self.assertEqual( + int(p2v[p_page].item()), + v_page, + f"p2v[{p_page}] != {v_page}", + ) + # Allocated physical-page range is hole-free + matches live count. + if alloc.grow_direction == "up": + alloc_lo, alloc_hi = alloc.min_page_index, alloc.watermark_physical + else: + alloc_lo, alloc_hi = ( + alloc.watermark_physical + 1, + alloc.num_pages, + ) + self.assertEqual(alloc_hi - alloc_lo, len(live_v_pages)) + for p_page in range(alloc_lo, alloc_hi): + self.assertNotEqual( + int(p2v[p_page].item()), -1, f"hole at physical page {p_page}" + ) + # Free virtual page ids ∪ live = [min_page_index, num_pages). + free_set = set(int(x) for x in alloc.free_virtual_ids.tolist()) + self.assertEqual( + free_set | set(live_v_pages), + set(range(alloc.min_page_index, alloc.num_pages)), + ) + self.assertEqual(free_set & set(live_v_pages), set()) + # For every token we stamped, verify data followed any relocations. + for v_tok, mark in stamped_tokens.items(): + v_page = v_tok // ps + offset = v_tok % ps + p_page_t = int(v2p[v_page].item()) + if p_page_t == -1: + continue # was freed; don't check + phys_tok = p_page_t * ps + offset + self.assertEqual( + int(kv.buf[phys_tok].item()), + mark, + f"data drift: stamped {mark} at virtual token {v_tok} " + f"(page {v_page}+offset {offset}) — found {int(kv.buf[phys_tok].item())}", + ) + + # 1. alloc(N) returns N TOKEN ids that are page-aligned. + def test_paged_alloc_token_aligned(self): + _, full_alloc, swa_alloc, full_kv, swa_kv = self._build() + v = full_alloc.alloc(16) # 2 pages × 8 tokens + self.assertIsNotNone(v) + self.assertEqual(int(v.numel()), 16) + # The output must consist of exactly 2 contiguous page-ranges. + v_pages = sorted(set((v // self.PAGE_SIZE).tolist())) + self.assertEqual(len(v_pages), 2) + for p in v_pages: + page_tokens = sorted(int(t) for t in v if t // self.PAGE_SIZE == p) + self.assertEqual( + page_tokens, + [p * self.PAGE_SIZE + i for i in range(self.PAGE_SIZE)], + "Page contents should be contiguous token ids", + ) + + # 2. alloc(N) requires N % page_size == 0. + def test_paged_alloc_non_aligned_raises(self): + _, full_alloc, _, _, _ = self._build() + with self.assertRaises(AssertionError): + full_alloc.alloc(5) # not a multiple of 8 + + # 3. v2p / p2v tables are sized by PAGES. + def test_paged_v2p_sized_by_pages(self): + pool, full_alloc, _, _, _ = self._build(n_full_pages=10) + # +1 for the trailing -1 sentinel row. + self.assertEqual( + int(full_alloc.virtual_to_physical.numel()), + full_alloc.num_pages + 1, + ) + self.assertEqual( + int(full_alloc.physical_to_virtual.numel()), + full_alloc.num_pages + 1, + ) + # `num_pages` should be > 1 to be a meaningful test. + self.assertGreater(full_alloc.num_pages, 1) + + # 4. Compaction relocates a whole page at once (data follows). + def test_paged_compaction_relocates_whole_pages(self): + _, full_alloc, _, full_kv, _ = self._build() + stamped = {} + # Alloc 3 pages worth of tokens. + a = full_alloc.alloc(self.PAGE_SIZE) # tokens of page X + b = full_alloc.alloc(self.PAGE_SIZE) # tokens of page Y (middle) + c = full_alloc.alloc(self.PAGE_SIZE) # tokens of page Z + + # Stamp each token with a UNIQUE marker. (alloc returns unique virtuals, + # but we want each token to be distinguishable from its in-page + # siblings, so we use the virtual-token value itself.) + for v in (a, b, c): + self._stamp_tokens(full_alloc, full_kv, v) + for t in v.tolist(): + stamped[t] = t + + # Free the MIDDLE page (token ids of `b`). This forces a compaction + # where page `c` (boundary, grow-up) relocates into page `b`'s slot. + full_alloc.free(b) + for t in b.tolist(): + stamped.pop(t, None) + # Erase markers for the freed page in the fake kv buf so the + # invariant check doesn't see stale data. + # (The compaction kernel moved the survivor's data; we don't manually + # touch full_kv.buf for the freed page — the test below verifies that + # `c`'s data followed the relocation.) + self._check_invariants(full_alloc, full_kv, stamped) + # `a` and `c` pages must still be live. + for t in a.tolist(): + v_page = t // self.PAGE_SIZE + self.assertNotEqual(int(full_alloc.virtual_to_physical[v_page].item()), -1) + for t in c.tolist(): + v_page = t // self.PAGE_SIZE + self.assertNotEqual(int(full_alloc.virtual_to_physical[v_page].item()), -1) + + # 5. free() recovers pages via unique(// page_size) — matches upstream. + def test_paged_free_unique_by_page(self): + _, full_alloc, _, full_kv, _ = self._build() + a = full_alloc.alloc(self.PAGE_SIZE * 2) # 2 pages = 2*PS tokens + allocated_count_before = full_alloc.allocated_count() + # `allocated_count()` returns TOKENS. + self.assertEqual(allocated_count_before, 2 * self.PAGE_SIZE) + # Internal page count. + self.assertEqual(full_alloc._allocated_pages(), 2) + # Free a SUBSET of tokens — but covering all tokens of both pages. + # (Matches the upstream contract: caller passes coherent ranges.) + full_alloc.free(a) + self.assertEqual(full_alloc.allocated_count(), 0) + self.assertEqual(full_alloc._allocated_pages(), 0) + + # 6. take_physical overflow check (grow-up direction). + def test_paged_take_physical_overflow_check(self): + _, full_alloc, _, _, _ = self._build(n_full_pages=4) + # Try to take more pages than the buffer can hold; should return None. + # First, fill normally up to the available_size, then over-alloc by 1. + avail = full_alloc.available_size() + n_pages = avail // self.PAGE_SIZE + result = full_alloc.take_physical(n_pages * self.PAGE_SIZE) + self.assertIsNotNone(result) + # Now one more page would overflow. + overflow = full_alloc.take_physical(self.PAGE_SIZE) + self.assertIsNone(overflow, "Overflow should return None, not crash") + + # 7. SWA composite joint byte-budget in page units. + def test_paged_swa_joint_byte_budget(self): + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + full_spec = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="down", + ) + n_full_pages, n_swa_pages = 8, 8 + total = ( + n_full_pages * self.PAGE_SIZE * full_spec.entry_bytes() + + n_swa_pages * self.PAGE_SIZE * swa_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + device=_DEV, + full_max_total_num_tokens=n_full_pages * self.PAGE_SIZE, + swa_max_total_num_tokens=n_swa_pages * self.PAGE_SIZE, + page_size=self.PAGE_SIZE, + need_sort=False, + forward_stream=None, + ) + # available_size() returns TOKENS. The joint byte-budget at page + # granularity uses `entry_sum_per_page = entry_full_per_page + + # entry_swa_per_page`. Pre-check: + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + entry_sum_pp = fa.entry_bytes_per_page + sa.entry_bytes_per_page + gap = sa._byte_low_frontier() - fa._byte_high_frontier() + expected_pages_by_bytes = gap // entry_sum_pp + expected = ( + min( + expected_pages_by_bytes, + fa.num_pages - fa.min_page_index, + sa.num_pages - sa.min_page_index, + ) + * self.PAGE_SIZE + ) + self.assertEqual(allocator.available_size(), expected) + # And it's strictly less than min(fa.available_size, sa.available_size) + # (since the joint cost is heavier than either single-side cost). + self.assertLessEqual( + allocator.available_size(), + min(fa.available_size(), sa.available_size()), + ) + + # 9. REGRESSION: alloc_extend must bind v2p / p2v on + # this allocator. Without binding, `virtual_to_physical[virt_page]` + # stays -1 and `translate_kv_loc(virt_token)` returns negative token + # ids → CUDA OOB in the Triton attention kernel. + def test_paged_alloc_extend_binds_v2p_p2v(self): + from sglang.srt.mem_cache import multi_ended_allocator as mea_mod + + _, full_alloc, _, _, _ = self._build() + PS = self.PAGE_SIZE + free_before = full_alloc.free_virtual_ids.clone() + watermark_before = full_alloc.watermark_physical + allocated_count_before = full_alloc.allocated_count() + + # Stub the kernel — we only need to verify the BINDING contract. + # (Driving the real Triton kernel needs a GPU; the contract we're + # checking is that the v2p/p2v tables get updated regardless of + # what the kernel writes into out_indices.) + original_kernel = mea_mod.alloc_extend_kernel + + class _NoOpKernelGrid: + def __getitem__(self, _grid): + return self + + def __call__(self, *a, **kw): + pass + + mea_mod.alloc_extend_kernel = _NoOpKernelGrid() + try: + # bs=1, prefix=0, seq=2 pages worth, so num_new_pages=2. + prefix_lens = torch.tensor([0], dtype=torch.int64, device=_DEV) + prefix_lens_cpu = torch.tensor([0], dtype=torch.int64) + seq_lens = torch.tensor([2 * PS], dtype=torch.int64, device=_DEV) + seq_lens_cpu = torch.tensor([2 * PS], dtype=torch.int64) + last_loc = torch.tensor([-1], dtype=torch.int64, device=_DEV) + + out = full_alloc.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + 2 * PS, + num_new_pages=2, + ) + finally: + mea_mod.alloc_extend_kernel = original_kernel + + self.assertIsNotNone(out) + # The two virtual pages consumed from the front of free_virtual_ids + # must now be BOUND in v2p_page (not -1). + consumed_pages = free_before[:2] + v2p_values = full_alloc.virtual_to_physical[consumed_pages] + for v_page, p_page in zip(consumed_pages.tolist(), v2p_values.tolist()): + self.assertNotEqual( + p_page, + -1, + f"REGRESSION: virtual page {v_page} not bound after " + f"alloc_extend (translate_kv_loc would return negative)", + ) + # And p2v_page must round-trip. + for v_page, p_page in zip(consumed_pages.tolist(), v2p_values.tolist()): + self.assertEqual(int(full_alloc.physical_to_virtual[p_page].item()), v_page) + # Watermark must have advanced by 2 pages. + # `allocated_count()` returns TOKENS, so it + # advances by 2 * PAGE_SIZE; `_allocated_pages()` is the page count. + self.assertEqual( + full_alloc.allocated_count(), + allocated_count_before + 2 * PS, + ) + self.assertEqual( + full_alloc._allocated_pages(), + (allocated_count_before // PS) + 2, + ) + if full_alloc.grow_direction == "up": + self.assertEqual(full_alloc.watermark_physical, watermark_before + 2) + else: + self.assertEqual(full_alloc.watermark_physical, watermark_before - 2) + # Free-list must have shrunk by 2. + self.assertEqual( + int(full_alloc.free_virtual_ids.numel()), + int(free_before.numel()) - 2, + ) + + # 10. REGRESSION: alloc_decode must bind v2p / p2v on + # this allocator when num_new_pages > 0. Most decode steps reuse the + # prefix's tail page (num_new_pages == 0), but the page-wrapping case + # must update tables. + def test_paged_alloc_decode_binds_v2p_p2v_on_page_wrap(self): + from sglang.srt.mem_cache import multi_ended_allocator as mea_mod + + _, full_alloc, _, _, _ = self._build() + PS = self.PAGE_SIZE + # Pre-allocate ~1 page so an arbitrary `seq_len % page_size == 1` + # decode step triggers a new-page consumption. + v = full_alloc.alloc(PS) + self.assertIsNotNone(v) + free_before = full_alloc.free_virtual_ids.clone() + watermark_before = full_alloc.watermark_physical + allocated_count_before = full_alloc.allocated_count() + + # Build a decode that wraps to a new page: seq_len % page_size == 1 + # (one req that just stepped past a page boundary). The kernel will + # consume 1 new page from `free_virtual_ids[0]`. + seq_lens = torch.tensor([PS + 1], dtype=torch.int64, device=_DEV) + seq_lens_cpu = torch.tensor([PS + 1], dtype=torch.int64) + last_loc = torch.tensor( + # last token of page-N at offset page_size-1. + [int(v[-1].item())], + dtype=torch.int64, + device=_DEV, + ) + + original_kernel = mea_mod.alloc_decode_kernel + + class _NoOpKernelGrid: + def __getitem__(self, _grid): + return self + + def __call__(self, *a, **kw): + pass + + mea_mod.alloc_decode_kernel = _NoOpKernelGrid() + try: + out = full_alloc.alloc_decode(seq_lens, seq_lens_cpu, last_loc) + finally: + mea_mod.alloc_decode_kernel = original_kernel + + self.assertIsNotNone(out) + # 1 virtual page consumed from the head of free_virtual_ids. + consumed_page = int(free_before[0].item()) + # v2p_page must now map to a valid physical page (not -1). + p_page = int(full_alloc.virtual_to_physical[consumed_page].item()) + self.assertNotEqual( + p_page, + -1, + f"REGRESSION: virtual page {consumed_page} not bound after " + f"alloc_decode (translate_kv_loc would return negative)", + ) + # p2v round-trip. + self.assertEqual( + int(full_alloc.physical_to_virtual[p_page].item()), consumed_page + ) + # Watermark must have advanced by 1 page. + # `allocated_count()` returns TOKENS (advance by PAGE_SIZE); + # `_allocated_pages()` is the page count. + self.assertEqual( + full_alloc.allocated_count(), + allocated_count_before + PS, + ) + self.assertEqual( + full_alloc._allocated_pages(), + (allocated_count_before // PS) + 1, + ) + if full_alloc.grow_direction == "up": + self.assertEqual(full_alloc.watermark_physical, watermark_before + 1) + else: + self.assertEqual(full_alloc.watermark_physical, watermark_before - 1) + # Free-list must have shrunk by 1. + self.assertEqual( + int(full_alloc.free_virtual_ids.numel()), + int(free_before.numel()) - 1, + ) + + # 11. REGRESSION: alloc_decode with num_new_pages == 0 + # (the common case — the decode token reuses the prefix's tail page) + # must NOT advance the watermark and NOT touch v2p / p2v. + def test_paged_alloc_decode_no_op_when_no_new_page(self): + from sglang.srt.mem_cache import multi_ended_allocator as mea_mod + + _, full_alloc, _, _, _ = self._build() + PS = self.PAGE_SIZE + # Pre-allocate 2 pages worth. We'll simulate a decode where seq_len + # advances WITHIN the existing tail page (no new page consumed). + v = full_alloc.alloc(PS) + free_before = full_alloc.free_virtual_ids.clone() + watermark_before = full_alloc.watermark_physical + allocated_count_before = full_alloc.allocated_count() + + # seq_len = PS - 1 (just inside the prefix page), pre-prefix-len = PS - 2. + # `(seq_lens % page_size == 1)` is FALSE here, so num_new_pages == 0. + seq_lens = torch.tensor([PS - 1], dtype=torch.int64, device=_DEV) + seq_lens_cpu = torch.tensor([PS - 1], dtype=torch.int64) + last_loc = torch.tensor( + [int(v[PS - 2].item())], + dtype=torch.int64, + device=_DEV, + ) + + original_kernel = mea_mod.alloc_decode_kernel + + class _NoOpKernelGrid: + def __getitem__(self, _grid): + return self + + def __call__(self, *a, **kw): + pass + + mea_mod.alloc_decode_kernel = _NoOpKernelGrid() + try: + out = full_alloc.alloc_decode(seq_lens, seq_lens_cpu, last_loc) + finally: + mea_mod.alloc_decode_kernel = original_kernel + + self.assertIsNotNone(out) + # Nothing should have moved — no new page consumed. + self.assertEqual(full_alloc.watermark_physical, watermark_before) + self.assertEqual(full_alloc.allocated_count(), allocated_count_before) + self.assertEqual( + int(full_alloc.free_virtual_ids.numel()), + int(free_before.numel()), + ) + + # 12. translate_kv_loc preserves token-level identity end-to-end. + def test_paged_translate_kv_loc_token_round_trip(self): + _, full_alloc, _, _, _ = self._build() + v = full_alloc.alloc(self.PAGE_SIZE * 2) + # Build the composite-style translation manually: virt_page * ps + offset. + ps = self.PAGE_SIZE + virt_pages = v // ps + offsets = v % ps + phys_pages = full_alloc.virtual_to_physical[virt_pages] + phys_tokens = phys_pages * ps + offsets + # `phys_tokens` should be a coherent set of two contiguous PAGES. + phys_pages_unique = sorted(set(phys_pages.tolist())) + self.assertEqual(len(phys_pages_unique), 2) + # Within each page the tokens go through offsets 0..7 in order. + for p in phys_pages_unique: + page_phys = sorted( + int(t) + for i, t in enumerate(phys_tokens.tolist()) + if int(phys_pages[i].item()) == p + ) + self.assertEqual( + page_phys, + [p * ps + i for i in range(ps)], + ) + + # REGRESSION: `translate_kv_loc(virt, out=buf)` must work + # under page_size > 1 — the page-math branch writes via + # `index_select(out=out)` + in-place `mul_` / `add_` and must match the + # no-`out=` form byte-for-byte. Tests the actual page-math path of the + # base-class implementation. + def test_paged_translate_kv_loc_with_out(self): + _, full_alloc, _, _, _ = self._build() + ps = self.PAGE_SIZE + v = full_alloc.alloc(2 * ps) + self.assertIsNotNone(v) + # Compare with-out vs no-out. + buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) + ptr_before = buf.data_ptr() + with_out = full_alloc.translate_kv_loc(v, out=buf) + no_out = full_alloc.translate_kv_loc(v) + self.assertIs(with_out, buf, "must return the out= buffer") + self.assertEqual( + buf.data_ptr(), ptr_before, "out= buffer's data_ptr must be stable" + ) + # Page-math correctness: result equals virt_page * ps + offset + # against the real v2p table. + virt_pages = v // ps + offsets = v % ps + phys_pages = full_alloc.virtual_to_physical[virt_pages] + expected = phys_pages * ps + offsets + self.assertTrue(bool((buf == expected).all().item())) + self.assertTrue(bool((with_out == no_out).all().item())) + + # REGRESSION: the in-place aliasing form + # `translate_kv_loc(buf, out=buf)` must work at page_size > 1 too. The + # page-math branch computes `virt_pages` and `offsets` BEFORE writing + # into `out`, so those fresh tensors capture the pre-mutation values of + # virt_tokens. The final result must equal `phys_page * ps + offset` + # against the original input. + def test_paged_translate_kv_loc_with_out_aliasing_input(self): + _, full_alloc, _, _, _ = self._build() + ps = self.PAGE_SIZE + v_orig = full_alloc.alloc(2 * ps) + self.assertIsNotNone(v_orig) + # Expected (no-out form) computed BEFORE mutating buf. + expected = full_alloc.translate_kv_loc(v_orig) + # In-place form: buf serves as both input and out. + buf = v_orig.clone() + ptr_before = buf.data_ptr() + ret = full_alloc.translate_kv_loc(buf, out=buf) + self.assertIs(ret, buf) + self.assertEqual(buf.data_ptr(), ptr_before) + self.assertTrue( + bool((buf == expected).all().item()), + "page>1 in-place result must equal no-out result", + ) + + # REGRESSION: the stale-tail scenario. + # + # A naive cuda-graph replay path in `triton_backend.py` would call + # `_translate_kv_loc(kv_indices, out=kv_indices)` on the WHOLE + # pre-allocated buffer (`self.cuda_graph_kv_indices`), even though only + # the `kv_indptr[-1]`-length prefix was freshly written by + # `create_flashinfer_kv_indices_triton`. Stale tail data, when fed + # through `v2p` repeatedly across replays, eventually produced negative + # values (via `v2p[unbound] = -1 → -1 * page_size + offset` ∈ [-ps,-1]), + # and the NEXT translation's `// page_size` produced `-1`, which CUDA's + # `index_select` rejects with a scatter-gather OOB device-side assert. + # + # The fix in `triton_backend.py` slices the translate to the valid + # prefix `kv_indices[:kv_indptr[-1]]`. This test confirms slicing is + # transparent to the translate — the in-place result on a contiguous + # slice of a larger buffer matches the standalone-tensor result. + def test_paged_translate_kv_loc_on_buffer_slice(self): + _, full_alloc, _, _, _ = self._build() + ps = self.PAGE_SIZE + v = full_alloc.alloc(2 * ps) + self.assertIsNotNone(v) + # Simulate the cuda-graph buffer pattern: a large pre-allocated + # buffer where only a prefix is freshly written. + N = v.numel() + big_buf = torch.zeros((N * 4,), dtype=torch.int64, device=_DEV) + big_buf[:N] = v + # Translate the valid prefix slice in-place (this is what the + # post-fix triton_backend call does). + slice_view = big_buf[:N] + ptr_before = slice_view.data_ptr() + ret = full_alloc.translate_kv_loc(slice_view, out=slice_view) + self.assertIs(ret, slice_view) + self.assertEqual( + slice_view.data_ptr(), + ptr_before, + "slice in-place write must preserve data_ptr", + ) + # The slice's translation must match a standalone translate of v. + expected = full_alloc.translate_kv_loc(v) + self.assertTrue( + bool((slice_view == expected).all().item()), + "slice in-place translate must equal standalone translate", + ) + # And the tail [N:] must remain UNTOUCHED — zeros, not corrupted + # by the translate. + tail = big_buf[N:] + self.assertTrue( + bool((tail == 0).all().item()), + "translating a slice must NOT touch the buffer tail; if this " + "test fails, the translate is reading/writing past the slice " + "bound — the same regression that caused a scatter-gather OOB " + "after several replays.", + ) + + # REGRESSION: tombstone-safety clamp at page_size > 1. + # + # For ps > 1, `v2p_page[vpage] == -1` produces `-1 * ps + offset` for + # output tokens — a negative value in `[-ps, -1]`. Without clamping, + # the captured `k_buffer[result[i]]` is an illegal access. The clamp + # must produce `>= 0` for every output token. + def test_paged_translate_kv_loc_clamps_tombstoned_v2p(self): + _, full_alloc, _, _, _ = self._build() + ps = self.PAGE_SIZE + v = full_alloc.alloc(2 * ps) + self.assertIsNotNone(v) + # Tombstone one page (any v2p_page entry) -> all ps tokens in that + # page should clamp to 0 in the translate output. + tomb_page = int((v[0] // ps).item()) + full_alloc.virtual_to_physical[tomb_page] = -1 + # No-out form. + out = full_alloc.translate_kv_loc(v) + self.assertTrue( + bool((out >= 0).all().item()), + f"paged translate_kv_loc must clamp tombstoned to >=0; got {out.tolist()}", + ) + # The first `ps` tokens belong to the tombstoned page → all 0. + self.assertTrue( + bool((out[:ps] == 0).all().item()), + "all tokens in a tombstoned page must map to slot 0 (padding sink)", + ) + # The second page is still bound; its outputs must be > 0. + self.assertTrue( + bool((out[ps:] > 0).all().item()), + "non-tombstoned pages must still translate to live physical slots", + ) + + def test_paged_translate_kv_loc_with_out_clamps_tombstoned_v2p(self): + _, full_alloc, _, _, _ = self._build() + ps = self.PAGE_SIZE + v = full_alloc.alloc(2 * ps).clone() + tomb_page = int((v[0] // ps).item()) + full_alloc.virtual_to_physical[tomb_page] = -1 + buf = torch.empty_like(v) + ret = full_alloc.translate_kv_loc(v, out=buf) + self.assertIs(ret, buf) + self.assertTrue( + bool((buf >= 0).all().item()), + "paged out= path must clamp tombstoned entries", + ) + self.assertTrue(bool((buf[:ps] == 0).all().item())) + + # 13. REGRESSION: `allocated_count()` MUST return + # TOKENS, not pages — matching upstream's convention that all external + # capacity methods report tokens. At page_size > 1, returning pages + # here breaks the leak invariant + # (`available + evictable + ... == total`, with all terms in tokens). + def test_paged_allocated_count_returns_tokens(self): + _, full_alloc, _, _, _ = self._build() + PS = self.PAGE_SIZE + # Idle → allocated_count == 0. + self.assertEqual(full_alloc.allocated_count(), 0) + # Alloc 2 pages = 2 * PS tokens. + v = full_alloc.alloc(2 * PS) + self.assertIsNotNone(v) + # allocated_count() must report TOKENS (= 2 * PS), not pages (= 2). + self.assertEqual( + full_alloc.allocated_count(), + 2 * PS, + "REGRESSION: allocated_count() must return TOKENS at page_size > 1", + ) + # _allocated_pages() is the page-granular internal helper. + self.assertEqual(full_alloc._allocated_pages(), 2) + + # 14. REGRESSION: the leak-invariant terms used by the + # scheduler runtime checker must all be in TOKENS. Specifically + # `full_available_size() + allocated_tokens == static_cap` must hold for + # the SWA composite. + def test_paged_swa_full_available_size_in_tokens(self): + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + full_spec = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="down", + ) + PS = self.PAGE_SIZE + n_full_pages, n_swa_pages = 16, 16 + total = ( + n_full_pages * PS * full_spec.entry_bytes() + + n_swa_pages * PS * swa_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + full_max = n_full_pages * PS + swa_max = n_swa_pages * PS + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + device=_DEV, + full_max_total_num_tokens=full_max, + swa_max_total_num_tokens=swa_max, + page_size=PS, + need_sort=False, + forward_stream=None, + ) + # Idle: conserve view == cap (in tokens). (The leak invariant reads + # `_conserve_*`; the public `full/swa_available_size()` is now + # `min(conserve, schedulable)` and may be smaller — reserved sink page.) + self.assertEqual(allocator._conserve_full_available_size(), full_max) + self.assertEqual(allocator._conserve_swa_available_size(), swa_max) + + # Alloc 2 pages = 2*PS tokens. + v = allocator.alloc(2 * PS) + self.assertIsNotNone(v) + + # conserve view must drop by 2*PS TOKENS, not by 2 (pages). + self.assertEqual( + allocator._conserve_full_available_size(), + full_max - 2 * PS, + "REGRESSION: the conserve view must drop by token-count, " + "not page-count. A 'pool memory leak detected' crash is " + "caused by a page-count drop here.", + ) + self.assertEqual( + allocator._conserve_swa_available_size(), + swa_max - 2 * PS, + ) + + # First-principles leak invariant: at this point, allocated tokens + # are all "live" (no eviction yet). So: + # total = conserve + allocated_tokens + # where allocated_tokens = full_max - conserve. + allocated_tokens = full_max - allocator._conserve_full_available_size() + self.assertEqual(allocated_tokens, 2 * PS) + self.assertEqual( + allocated_tokens + allocator._conserve_full_available_size(), + full_max, + ) + + # 15. REGRESSION: UnifiedMambaTokenToKVPoolAllocator.size + # must be TOTAL TOKENS (available + allocated, both in tokens). At + # page_size > 1, the earlier `available + allocated_pages` formula gave + # `tokens + pages` which silently broke the chunk-cache Mamba log lines + # (`#full token`, `full token usage`) and would have crashed Mamba+radix + # if radix weren't auto-downgraded to page=1. + def test_paged_mamba_size_in_tokens(self): + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaTokenToKVPoolAllocator, + ) + + # Build a minimal Mamba composite: one MHA spec for full + one + # Mamba spec. The mamba sub-allocator always uses page_size=1, but + # the full sub-allocator uses self.PAGE_SIZE. + PS = self.PAGE_SIZE + full_spec = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + mamba_spec = MambaSubPoolSpec( + name="mamba", + layer_num=2, + conv_state_shapes=((4, 3),), + conv_dtype=torch.float32, + temporal_state_shape=(2, 2, 2), + temporal_dtype=torch.float32, + grow_direction="down", + ) + n_full_pages, n_mamba_slots = 16, 8 + total = ( + n_full_pages * PS * full_spec.entry_bytes() + + n_mamba_slots * mamba_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, mamba_spec], + device=_DEV, + enable_memory_saver=False, + ) + # Build a fake HybridLinearKVPool-like object with two sub-pool kv + # caches. We only need `.full_kv_pool` and `.mamba_pool` with + # `attach_allocator` / `move_kv_cache` stubs. + full_kv = _FakeKVCache(pool.max_slots("full")) + full_kv.attach_allocator = lambda allocator: None + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + mamba_kv.attach_allocator = lambda allocator: None + # _copy_from_physical for the mamba sub-pool (kept un-translated). + mamba_kv._copy_from_physical = lambda src, dst: None + + class _FakeHybridLinearKVPool: + full_kv_pool = full_kv + mamba_pool = mamba_kv + + allocator = UnifiedMambaTokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=_FakeHybridLinearKVPool(), + device=_DEV, + page_size=PS, + need_sort=False, + forward_stream=None, + ) + + # Idle: size == full_available_size() (entirely in tokens). + full_avail_before = allocator.full_attn_allocator.available_size() + self.assertEqual(allocator.size, full_avail_before) + # available_size == size (no allocations yet). + self.assertEqual(allocator.available_size(), allocator.size) + + # Alloc 2 pages = 2*PS tokens on full side. + v = allocator.alloc(2 * PS) + self.assertIsNotNone(v) + + # size should be CONSERVED in tokens: (available + allocated_tokens) + # stays at the initial total. (For the Mamba composite, `.size` is + # dynamic — it shrinks as the peer consumes bytes — but at this + # point the peer is idle so we should see `size == full_avail_before`.) + self.assertEqual( + allocator.full_attn_allocator.available_size() + + allocator.full_attn_allocator.allocated_count(), + full_avail_before, + "REGRESSION: full.available_size() + full.allocated_count() must " + "be conserved at TOKEN granularity (was `tokens + pages` in the " + "buggy revision).", + ) + # And .size matches this conserved sum. + self.assertEqual(allocator.size, full_avail_before) + + # 16. REGRESSION: the page-math helper used by + # `UnifiedSWAKVPool.translate_loc_from_full_to_swa`, + # `UnifiedSWAKVPool.get_cpu_copy`, and `load_cpu_copy` must do + # `virt_pages = loc // page_size; offsets = loc % page_size; + # phys_tokens = v2p_page[virt_pages] * page_size + offsets`. + # + # A naive implementation does `v2p[loc]` directly — + # indexing a page-granular table with token-granular ids, producing + # wrong physical token ids and (when used as Triton kernel inputs) + # OOB reads (the same bug class as the alloc_extend/alloc_decode + # binding regressions above). + # + # We can't easily construct a real UnifiedSWAKVPool in the CPU test shim + # (it inherits SWAKVPool which builds MHATokenToKVPool sub-pools), so + # we exercise the static helper `_virt_tokens_to_phys_tokens` directly. + # The instance methods in production wrap this helper, so the same + # math is covered. + def test_paged_pool_translate_helper_returns_physical_tokens(self): + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool + + full_spec = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=2, + head_num=2, + head_dim=4, + store_dtype=torch.float16, + grow_direction="down", + ) + PS = self.PAGE_SIZE + n_pages = 8 + total = ( + n_pages * PS * full_spec.entry_bytes() + + n_pages * PS * swa_spec.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full_spec, swa_spec], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + allocator = UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + device=_DEV, + full_max_total_num_tokens=n_pages * PS, + swa_max_total_num_tokens=n_pages * PS, + page_size=PS, + need_sort=False, + forward_stream=None, + ) + + # Alloc 2 pages worth of tokens — the swa allocator's v2p_page table + # now has bindings for the consumed virtual pages. + v_tokens = allocator.alloc(2 * PS) + self.assertIsNotNone(v_tokens) + + # The static helper does the page math: same as the instance methods. + swa_phys = UnifiedSWAKVPool._virt_tokens_to_phys_tokens( + v_tokens, allocator.swa_attn_allocator + ) + + # Output must: + # 1. Be non-negative for every input (none unbound at this point). + # 2. Be distinct (one-to-one mapping). + # 3. Match `swa_phys_page * page_size + offset` reconstructed directly. + self.assertTrue( + bool((swa_phys >= 0).all().item()), + "REGRESSION: _virt_tokens_to_phys_tokens returned negative " + "physical token ids (page-math fix likely reverted).", + ) + self.assertEqual( + int(torch.unique(swa_phys).numel()), + int(swa_phys.numel()), + "Physical token ids must be unique (one-to-one mapping).", + ) + virt_pages_in = v_tokens // PS + offsets_in = v_tokens % PS + swa_phys_pages_direct = allocator.swa_attn_allocator.virtual_to_physical[ + virt_pages_in + ] + expected = swa_phys_pages_direct * PS + offsets_in + self.assertTrue( + bool((swa_phys == expected).all().item()), + "REGRESSION: _virt_tokens_to_phys_tokens output must equal " + "v2p_page[virt_pages] * page_size + offsets.", + ) + + # And the composite allocator's translate method must produce the + # same token-granular result (same page math). + composite_out = allocator.translate_loc_from_full_to_swa(v_tokens) + self.assertTrue( + bool((swa_phys.long() == composite_out.long()).all().item()), + "REGRESSION: the UnifiedSWAKVPool helper and the composite " + "allocator's translate_loc_from_full_to_swa must agree.", + ) + + +class TestLazyCompaction(unittest.TestCase): + """Lazy compaction invariants and lazy-vs-eager + equivalence harness. CPU-only (no GPU events; the conservative Phase A + `_flush` uses `wait_stream(forward_stream)` only when `forward_stream is + not None`, so passing `forward_stream=None` keeps it a no-op). + """ + + def _make_full(self, *, lazy: bool, n_full_slots=64, n_mamba_slots=16): + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + total = full.entry_bytes() * n_full_slots + mamba.entry_bytes() * n_mamba_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + full_alloc = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + lazy_compaction=lazy, + ) + mamba_alloc = MultiEndedAllocator( + kvcache=mamba_kv, + unified_buffer=pool, + sub_pool_name="mamba", + device=_DEV, + is_id_owner=True, + lazy_compaction=lazy, + ) + full_alloc.bind_peer(mamba_alloc) + mamba_alloc.bind_peer(full_alloc) + return pool, full_alloc, full_kv + + def _stamp_kv(self, kv: _FakeKVCache, alloc: MultiEndedAllocator, tokens) -> None: + """Write a marker into KV[phys] for each freshly-alloced virtual + token id, so we can later check the data followed any relocation. + """ + for v in tokens.tolist(): + p = int(alloc.virtual_to_physical[v].item()) + kv.buf[p] = int(v) + + def test_lazy_state_initialized(self): + """Lazy allocator initializes the new state cleanly.""" + _pool, fa, _kv = self._make_full(lazy=True) + self.assertTrue(fa.lazy_compaction) + self.assertEqual(len(fa._free_phys_pages), 0) + self.assertEqual(fa._pending_reuse, {}) + self.assertEqual(fa.live_page_count, 0) + # Watermark + free virtual list start equivalent to eager. + self.assertEqual(fa.watermark_physical, fa.min_page_index) + + def test_lazy_alloc_increments_live_page_count(self): + _pool, fa, _kv = self._make_full(lazy=True) + tokens = fa.alloc(8) + self.assertIsNotNone(tokens) + self.assertEqual(int(tokens.numel()), 8) + self.assertEqual(fa.live_page_count, 8) + self.assertEqual(len(fa._free_phys_pages), 0) + + def test_lazy_free_boundary_shortcut(self): + """Boundary absorption is DEFERRED to `_flush` (the hot + path `_free_lazy` does only a `torch.cat`, no watermark mutation). + After `_flush`, the freed boundary page is absorbed into the + watermark. + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(3) # virtual tokens + before_wm = fa.watermark_physical + # Free the last-alloced virtual id (its physical IS the boundary). + last = a[-1:].clone() + fa.free(last) + # Watermark is NOT shrunk inline; freed page is in the + # free list. + self.assertEqual(fa.watermark_physical, before_wm) + self.assertEqual(len(fa._free_phys_pages), 1) + # `live_page_count` is decremented at free time (CPU-side metadata). + self.assertEqual(fa.live_page_count, 2) + # `_flush` runs the complete boundary absorb → watermark shrinks + # by 1, free list emptied. + fa._flush(urgent=True) + self.assertEqual(fa.watermark_physical, before_wm - 1) + self.assertEqual(len(fa._free_phys_pages), 0) + self.assertEqual(fa.live_page_count, 2) + + def test_lazy_free_non_boundary_pushes_hole(self): + """Freeing a non-boundary page enters _free_phys_pages, watermark + stays put. + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + wm_before = fa.watermark_physical + # Free a middle id (NOT the topmost), boundary-shortcut should + # NOT fire. + mid = a[2:3].clone() + fa.free(mid) + self.assertEqual(fa.watermark_physical, wm_before) + self.assertEqual(len(fa._free_phys_pages), 1) + self.assertEqual(fa.live_page_count, 4) + + def test_lazy_free_inward_walk(self): + """The inward walk (multiple contiguous holes absorbed + into the watermark in one pass) is DEFERRED to `_flush`. After + flush, the watermark shrinks past all contiguous-from-boundary + holes, regardless of the order they were freed. + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + wm_before = fa.watermark_physical + # Free a middle slot first → hole. + fa.free(a[2:3].clone()) + # Free the topmost ids — in eager mode these would absorb inline, + # but in lazy mode they cat onto the free list. + fa.free(a[4:5].clone()) + fa.free(a[3:4].clone()) + # 3 entries in the free list now; watermark still at the + # pre-free position. + self.assertEqual(fa.watermark_physical, wm_before) + self.assertEqual(len(fa._free_phys_pages), 3) + # `_flush` runs the complete CPU-side boundary absorb. All 3 + # contiguous holes near the boundary get absorbed in one pass. + fa._flush(urgent=True) + self.assertEqual(fa.watermark_physical, wm_before - 3) + self.assertEqual(len(fa._free_phys_pages), 0) + self.assertEqual(fa.live_page_count, 2) + + def test_lazy_take_physical_drains_holes_first(self): + """alloc reuses holes before extending the watermark.""" + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + # Free two non-boundary virtuals to populate _free_phys_pages. + fa.free(a[1:2].clone()) + fa.free(a[3:4].clone()) + n_holes_before = len(fa._free_phys_pages) + self.assertEqual(n_holes_before, 2) + wm_before = fa.watermark_physical + # Allocate 2 more — both should come from holes, watermark unchanged. + a2 = fa.alloc(2) + self.assertEqual(fa.watermark_physical, wm_before) + self.assertEqual(len(fa._free_phys_pages), 0) + # Allocate 1 more — must extend (no holes left). + a3 = fa.alloc(1) + self.assertEqual(fa.watermark_physical, wm_before + 1) + # All three new alloc batches are non-None. + self.assertIsNotNone(a2) + self.assertIsNotNone(a3) + + def test_lazy_available_size_includes_holes(self): + """available_size counts drainable holes + extension capacity.""" + _pool, fa, _kv = self._make_full(lazy=True) + avail_initial = fa.available_size() + a = fa.alloc(5) + # 3 non-boundary frees → 3 holes; watermark unchanged. + fa.free(a[0:1].clone()) + fa.free(a[1:2].clone()) + fa.free(a[2:3].clone()) + self.assertEqual(len(fa._free_phys_pages), 3) + avail_after = fa.available_size() + # holes (3) + remaining extension capacity == original capacity + # adjusted by the 2 still-live tokens at the top. + # Concretely: avail_after = avail_initial - 2 (live). + self.assertEqual(avail_after, avail_initial - 2) + + def test_lazy_flush_compacts_holes_into_gap(self): + """_flush(urgent=True) moves a survivor into a hole and shrinks the + watermark, freeing bytes back into the shared gap. + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + # Stamp KV so we can assert the data followed the relocation. + self._stamp_kv(_kv, fa, a) + # Free a low-index hole; keep the topmost live. + fa.free(a[1:2].clone()) + self.assertEqual(len(fa._free_phys_pages), 1) + wm_before = fa.watermark_physical + n_moves = fa._flush(urgent=True) + # At least one move should have happened (topmost survivor → hole). + self.assertGreaterEqual(n_moves, 1) + # Watermark shrunk; hole list is now empty. + self.assertLess(fa.watermark_physical, wm_before) + self.assertEqual(len(fa._free_phys_pages), 0) + # live_page_count invariant under compaction. + self.assertEqual(fa.live_page_count, 4) + + def test_lazy_v2p_p2v_identity_after_flush(self): + """After a flush, v2p ∘ p2v == identity on the live set.""" + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(8) + # Free a scattered set of virtuals (not at boundary). + fa.free(a[1:2].clone()) + fa.free(a[3:4].clone()) + fa.free(a[5:6].clone()) + fa._flush(urgent=True) + # For every still-live virtual token, v2p ∘ p2v == identity. + for v in a.tolist(): + p = int(fa.virtual_to_physical[v].item()) + if p == -1: + continue # freed + self.assertEqual(int(fa.physical_to_virtual[p].item()), v) + + def _replay_sequence(self, ops, lazy: bool): + """Run a given alloc/free op trace under eager OR lazy mode and + return the final (live virtual set, alloc-time KV stamps).""" + _pool, fa, kv = self._make_full(lazy=lazy) + live = set() # set of virtual ids + kv_stamps = {} # v -> stamp (the data we wrote at alloc time) + next_stamp = 100 + for kind, n in ops: + if kind == "alloc": + tokens = fa.alloc(n) + if tokens is None: + continue + for v in tokens.tolist(): + p = int(fa.virtual_to_physical[v].item()) + kv.buf[p] = next_stamp + kv_stamps[v] = next_stamp + live.add(v) + next_stamp += 1 + elif kind == "free": + if not live: + continue + # Take up to n from live, deterministically by id. + victims = sorted(live)[:n] + live.difference_update(victims) + fa.free(torch.tensor(victims, dtype=torch.int64)) + # Force final compaction on lazy so the comparison is at quiescence. + if lazy: + fa._flush(urgent=True) + # Read back the data for each live id. + live_data = {} + for v in live: + p = int(fa.virtual_to_physical[v].item()) + live_data[v] = int(kv.buf[p].item()) + return live, live_data, kv_stamps + + def test_lazy_vs_eager_equivalence(self): + """Same random alloc/free sequence under lazy and eager modes must + yield identical live virtual sets AND identical KV reads (the data + followed any relocation). + """ + rng = random.Random(42) + ops = [] + for _ in range(200): + if rng.random() < 0.6: + ops.append(("alloc", rng.randint(1, 6))) + else: + ops.append(("free", rng.randint(1, 4))) + eager_live, eager_data, eager_stamps = self._replay_sequence(ops, lazy=False) + lazy_live, lazy_data, lazy_stamps = self._replay_sequence(ops, lazy=True) + self.assertEqual(eager_live, lazy_live, "live virtual set diverged") + self.assertEqual(eager_stamps, lazy_stamps, "alloc-time stamps diverged") + # For every live id, the data we read back must match what we wrote. + for v in eager_live: + self.assertEqual( + eager_data[v], eager_stamps[v], f"eager: KV[v={v}] != stamp" + ) + self.assertEqual(lazy_data[v], lazy_stamps[v], f"lazy: KV[v={v}] != stamp") + + def test_lazy_hole_set_directional_pop(self): + """The _HoleSet pops smallest-first for grow-up; alloc must drain + the deepest hole first (the greedy clustering rule keeps near- + boundary holes available for cheap absorption by compaction). + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(6) + # Free middle and lower middles so the holes are NOT at boundary. + fa.free(a[1:2].clone()) # frees physical at index v2p[a[1]] + fa.free(a[3:4].clone()) + # Capture which physical pages are now in the hole set. + # `_free_phys_pages` is a torch.Tensor; `.tolist()` returns + # Python ints so `sorted` produces ints (not 0-dim tensors). + holes_before = sorted(fa._free_phys_pages.tolist()) + self.assertEqual(len(holes_before), 2) + # Alloc 1 — should drain a hole (grow-up). + # With sort-after-merge OFF (default), the drain order + # is FIFO over the free-list tensor — NOT "smallest first". + # We only assert that the bound physical is ONE OF the holes. + a2 = fa.alloc(1) + bound_phys = int(fa.virtual_to_physical[int(a2.item())].item()) + self.assertIn(bound_phys, holes_before) + + def test_lazy_non_urgent_stops_at_write_set_blocker(self): + """Write-race case: when the topmost survivor IS in an + in-flight batch's write-set, non-urgent `_flush` STOPS the + boundary walk (skipping past would shuffle holes without + shrinking the watermark — wasted work).""" + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + + class _FakeEvent: + def __init__(self): + self.fired = False + + def query(self): + return self.fired + + ev = _FakeEvent() + fa.set_latest_forward_done_event(ev) + # Free a non-boundary slot to create a compactable hole. + fa.free(a[1:2].clone()) + self.assertEqual(len(fa._free_phys_pages), 1) + # Register an in-flight forward whose write-set INCLUDES the + # topmost survivor. We pass the virtual `out_cache_loc` TENSOR + # (not a materialized physical set) — `_flush` translates it + # lazily on the scheduler thread when classifying survivors. + topmost_phys = int(fa.virtual_to_physical[int(a[-1].item())].item()) + oclv = a[-1:].clone() # virtual id that translates to topmost_phys + fa.set_inflight_forward(ev, oclv) + # Non-urgent flush → case A blocker at the top → STOP. + n_moves = fa._flush(urgent=False) + self.assertEqual( + n_moves, + 0, + "non-urgent flush must STOP when the topmost survivor is in " + "an in-flight write-set", + ) + # State untouched: hole still present. + self.assertEqual(len(fa._free_phys_pages), 1) + self.assertEqual(len(fa._pending_reuse), 0) + # Fire the event (forward done) so the write-set entry prunes; a + # subsequent flush proceeds and releases src directly (event fired + # → no pending reuse entry needed). + ev.fired = True + n_moves2 = fa._flush(urgent=False) + self.assertGreaterEqual(n_moves2, 1) + self.assertEqual(len(fa._pending_reuse), 0) + + def test_lazy_non_urgent_read_race_uses_pending_reuse(self): + """Read-race case (read race, no write race): when the + topmost survivor is NOT in any in-flight write-set, non-urgent + `_flush` compacts immediately (read+read on KV[src] is safe) and + pushes `(src, latest_event)` to `_pending_reuse` so a future + alloc can't write KV[src] while iter N+1's read is still pending. + """ + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(5) + + class _FakeEvent: + def __init__(self): + self.fired = False + + def query(self): + return self.fired + + ev = _FakeEvent() + fa.set_latest_forward_done_event(ev) + # Free a non-boundary slot → 1 hole. + fa.free(a[1:2].clone()) + # Empty in-flight write-set for the in-flight forward → topmost + # survivor is NOT in the write-set → compact proceeds, src goes + # to pending. Pass `None` (or an empty tensor) for + # `out_cache_loc_virtual` to signal "no write race on this pool" + # — this is the same path Mamba uses (forward writes mamba state + # via its own kernels, not via `out_cache_loc`). + fa.set_inflight_forward(ev, None) + n_moves = fa._flush(urgent=False) + self.assertGreaterEqual(n_moves, 1) + # `_pending_reuse` has ONE entry per BATCH (keyed by + # event), not per src. So len(_pending_reuse) == 1 here. The + # total pages held is tracked in `_pending_reuse_pages_cpu`. + self.assertEqual(len(fa._pending_reuse), 1) + self.assertEqual(len(fa._pending_reuse_pages_cpu), n_moves) + # Fire the event and drain — srcs return to availability. + ev.fired = True + fa._drain_pending_reuse(urgent=False) + self.assertEqual(len(fa._pending_reuse), 0) + self.assertEqual(len(fa._pending_reuse_pages_cpu), 0) + + def test_lazy_pending_reuse_urgent_wait(self): + """Under urgent drain, an unfired event triggers wait_event; we + simulate this by checking that the drain ALSO releases unfired + entries (with a fake event whose `query` is False — `wait_event` is + a no-op in CPU mode since there's no current stream's wait_event for + a FakeEvent, so we test the release path).""" + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(4) + + class _FakeEvent: + def __init__(self): + self.waited = False + + def query(self): + return False # never fires + + # Inject ONE batch entry into _pending_reuse keyed by + # Event. Value is `(cpu_list, gpu_tensor)`. The parallel CPU + # set must also be updated. + # (Simulates a prior compaction whose event hasn't fired.) + p = int(fa.virtual_to_physical[int(a[2].item())].item()) + # Clear v2p/p2v so post-drain reuse is safe. + fa.virtual_to_physical[int(a[2].item())] = -1 + fa.physical_to_virtual[p] = -1 + ev = _FakeEvent() + gpu_t = torch.tensor([p], dtype=torch.int64, device=fa.device) + fa._pending_reuse[ev] = ([p], gpu_t) + fa._pending_reuse_pages_cpu.add(p) + # Urgent drain — should release p despite event.query()=False. + # (CPU shim: torch.cuda.current_stream() may not exist; wrap try.) + try: + fa._drain_pending_reuse(urgent=True) + except Exception: + # CPU: wait_event may not work; this test is GPU-only. + self.skipTest("wait_event requires CUDA") + self.assertEqual(len(fa._pending_reuse), 0) + self.assertEqual(len(fa._pending_reuse_pages_cpu), 0) + + def test_lazy_flush_opportunistic_hook(self): + """The public flush_opportunistic method runs the non-urgent path + and is safe to call when no holes exist.""" + _pool, fa, _kv = self._make_full(lazy=True) + # No holes → returns 0 moves, no-op. + self.assertEqual(fa.flush_opportunistic(), 0) + # Create a hole then call flush_opportunistic; latest_event=None + # means src releases immediately. + a = fa.alloc(3) + fa.free(a[0:1].clone()) + moves = fa.flush_opportunistic() + self.assertGreaterEqual(moves, 1) + + +class TestO3FusedAllocBind(unittest.TestCase): + """Fused take_physical_pages + bind_pages. + + GPU-only tests (Triton kernel requires CUDA). Exercise the helper + `_alloc_bind_fast_or_slow` directly: fast-path correctness, slow- + path fallback when holes exist (Invariant B), overflow handling, + eager vs lazy modes, grow-up vs grow-down, and page_size > 1. + """ + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("O3 fused alloc-bind kernel requires CUDA") + + def _make_full( + self, + *, + lazy: bool = True, + n_full_slots: int = 64, + n_mamba_slots: int = 16, + page_size: int = 1, + ): + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + total = full.entry_bytes() * n_full_slots + mamba.entry_bytes() * n_mamba_slots + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device="cuda", + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + fa = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device="cuda", + is_id_owner=True, + page_size=page_size, + lazy_compaction=lazy, + ) + ma = MultiEndedAllocator( + kvcache=mamba_kv, + unified_buffer=pool, + sub_pool_name="mamba", + device="cuda", + is_id_owner=True, + page_size=1, # mamba is per-request, always page=1 + lazy_compaction=lazy, + ) + fa.bind_peer(ma) + ma.bind_peer(fa) + return pool, fa, full_kv + + def test_helper_exists_and_returns_tensor(self): + """The helper `_alloc_bind_fast_or_slow` is wired and returns a + tensor on success.""" + _pool, fa, _kv = self._make_full(lazy=True) + v_pages = torch.tensor([10, 11, 12], dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, 3) + self.assertIsNotNone(phys) + self.assertEqual(phys.shape, (3,)) + self.assertEqual(phys.dtype, torch.int64) + self.assertEqual(phys.device.type, "cuda") + + def test_fast_path_when_no_holes(self): + """When `_free_phys_pages` is empty, the fast path fires. + Verifies: watermark advanced, v2p and p2v scattered correctly, + return tensor matches the kernel's arange.""" + _pool, fa, _kv = self._make_full(lazy=True) + # Sanity: empty holeset. + self.assertEqual(len(fa._free_phys_pages), 0) + wm_before = fa.watermark_physical + # Pick virtual page ids the kernel will bind. + v_pages = torch.tensor([20, 21, 22, 23], dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, 4) + # Watermark advanced by N. + self.assertEqual(fa.watermark_physical, wm_before + 4) + # Returned phys ids match the grow-up arange [wm_before, wm_before+4). + expected_phys = torch.arange( + wm_before, wm_before + 4, dtype=torch.int64, device="cuda" + ) + self.assertTrue(torch.equal(phys, expected_phys)) + # v2p table: each virtual → its physical. + for v, p in zip(v_pages.tolist(), expected_phys.tolist()): + self.assertEqual(int(fa.virtual_to_physical[v].item()), p) + # p2v table: each physical → its virtual. + for v, p in zip(v_pages.tolist(), expected_phys.tolist()): + self.assertEqual(int(fa.physical_to_virtual[p].item()), v) + # live_page_count updated. + self.assertEqual(fa.live_page_count, 4) + + def test_slow_path_when_holes_exist(self): + """Invariant B (greedy hole reuse): when a hole exists, alloc + drains it BEFORE extending the watermark. The fast path MUST + NOT fire.""" + _pool, fa, _kv = self._make_full(lazy=True) + # Build a hole by alloc-then-free-non-boundary. + a = fa.alloc(3) + fa.free(a[0:1].clone()) # frees a non-boundary slot → enters holeset + self.assertEqual(len(fa._free_phys_pages), 1) + # `_free_phys_pages` is a torch.Tensor; read the single + # hole position via `.tolist()` (`._alive` no longer exists). + hole_pos = int(fa._free_phys_pages.tolist()[0]) + wm_before = fa.watermark_physical + # Alloc 1 page via the helper. Slow path should drain the hole. + v_pages = torch.tensor([42], dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, 1) + # Hole drained, NOT a watermark extension. + self.assertEqual(int(phys[0].item()), hole_pos) + self.assertEqual(fa.watermark_physical, wm_before) + self.assertEqual(len(fa._free_phys_pages), 0) + # v2p/p2v updated. + self.assertEqual(int(fa.virtual_to_physical[42].item()), hole_pos) + self.assertEqual(int(fa.physical_to_virtual[hole_pos].item()), 42) + + def test_fast_path_in_eager_mode(self): + """Eager mode (no lazy compaction) ALWAYS uses the fast path — + no holes ever accumulate.""" + _pool, fa, _kv = self._make_full(lazy=False) + self.assertFalse(fa.lazy_compaction) + wm_before = fa.watermark_physical + v_pages = torch.tensor([30, 31, 32], dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, 3) + self.assertEqual(fa.watermark_physical, wm_before + 3) + expected_phys = torch.arange( + wm_before, wm_before + 3, dtype=torch.int64, device="cuda" + ) + self.assertTrue(torch.equal(phys, expected_phys)) + + def test_index_space_overflow_returns_none(self): + """When the requested allocation would overflow `num_pages`, + the helper returns None and leaves the allocator unchanged.""" + _pool, fa, _kv = self._make_full(lazy=True, n_full_slots=8, n_mamba_slots=2) + # Try to alloc more pages than exist. + N = fa.num_pages + 100 + wm_before = fa.watermark_physical + # Note: we need v_pages of size N; but only its NUMEL matters for + # the helper. We pass a dummy tensor of the right shape. + v_pages = torch.zeros(N, dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, N) + self.assertIsNone(phys) + # Allocator state unchanged. + self.assertEqual(fa.watermark_physical, wm_before) + self.assertEqual(fa.live_page_count, 0) + + def test_empty_alloc_returns_empty_tensor(self): + """N=0 returns an empty tensor (no kernel launch, no state change).""" + _pool, fa, _kv = self._make_full(lazy=True) + wm_before = fa.watermark_physical + v_pages = torch.empty(0, dtype=torch.int64, device="cuda") + phys = fa._alloc_bind_fast_or_slow(v_pages, 0) + self.assertIsNotNone(phys) + self.assertEqual(phys.numel(), 0) + self.assertEqual(fa.watermark_physical, wm_before) + + def test_fast_path_equivalent_to_slow_path(self): + """For the same input on an empty-holeset allocator, the fast + path produces byte-identical v2p / p2v / return-tensor to the + slow path (which is the unfused take_physical_pages + bind + sequence). Verifies the kernel correctness against the + reference implementation.""" + # Two identical allocators; one takes fast path, one takes slow. + _pool_a, fa_a, _kv_a = self._make_full(lazy=True) + _pool_b, fa_b, _kv_b = self._make_full(lazy=True) + v_pages = torch.tensor([50, 51, 52, 53, 54], dtype=torch.int64, device="cuda") + # Fast path on fa_a. + phys_a = fa_a._alloc_bind_fast_or_slow(v_pages, 5) + # Slow path on fa_b: directly call take_physical_pages + bind + # (the unfused reference implementation). + phys_b = fa_b.take_physical_pages(5) + fa_b.bind(v_pages, phys_b) + # take_physical_pages already advances live_page_count (matching the + # fused fast path), so no manual bump here. + # Identical return tensors. + self.assertTrue(torch.equal(phys_a, phys_b)) + # Identical v2p / p2v after the operation. + self.assertTrue(torch.equal(fa_a.virtual_to_physical, fa_b.virtual_to_physical)) + self.assertTrue(torch.equal(fa_a.physical_to_virtual, fa_b.physical_to_virtual)) + # Identical watermark + live_page_count. + self.assertEqual(fa_a.watermark_physical, fa_b.watermark_physical) + self.assertEqual(fa_a.live_page_count, fa_b.live_page_count) + + def test_eager_mode_live_page_count_not_updated(self): + """Match `_take_physical_eager` semantics: in eager mode, + `live_page_count` is NOT maintained (the leak-checker uses + `allocated_count()` based on the watermark span). The helper's + fast path must respect this — updating it would break the + invariant that eager-mode `live_page_count == 0` always.""" + _pool, fa, _kv = self._make_full(lazy=False) + self.assertFalse(fa.lazy_compaction) + self.assertEqual(fa.live_page_count, 0) + v_pages = torch.tensor([10, 11, 12], dtype=torch.int64, device="cuda") + fa._alloc_bind_fast_or_slow(v_pages, 3) + # live_page_count UNCHANGED (eager mode invariant). + self.assertEqual(fa.live_page_count, 0) + + def test_lazy_mode_live_page_count_updated_on_fast_path(self): + """Lazy mode: the fast path advances `live_page_count` by N + (matches `take_physical`'s lazy-path bookkeeping).""" + _pool, fa, _kv = self._make_full(lazy=True) + self.assertEqual(fa.live_page_count, 0) + v_pages = torch.tensor([20, 21, 22], dtype=torch.int64, device="cuda") + fa._alloc_bind_fast_or_slow(v_pages, 3) + self.assertEqual(fa.live_page_count, 3) + # Another fast-path call accumulates. + v_pages2 = torch.tensor([23, 24], dtype=torch.int64, device="cuda") + fa._alloc_bind_fast_or_slow(v_pages2, 2) + self.assertEqual(fa.live_page_count, 5) + + def test_lazy_mode_live_page_count_updated_on_slow_path(self): + """Lazy mode + holes exist: the slow path advances + `live_page_count` via the existing `take_physical_pages` call + (which updates it internally). End state must match the fast + path's accumulation.""" + _pool, fa, _kv = self._make_full(lazy=True) + a = fa.alloc(3) + # alloc(3) used the fast path (no holes at the time). + self.assertEqual(fa.live_page_count, 3) + # Free one non-boundary → creates a hole; subsequent alloc takes + # the slow path. + fa.free(a[0:1].clone()) + self.assertEqual(fa.live_page_count, 2) + self.assertEqual(len(fa._free_phys_pages), 1) + # Now alloc(1) takes the slow path (holes exist). Should still + # update live_page_count back to 3. + b = fa.alloc(1) + self.assertIsNotNone(b) + self.assertEqual(fa.live_page_count, 3) + # Verify slow path actually fired: hole drained, watermark unchanged. + self.assertEqual(len(fa._free_phys_pages), 0) + + def test_page_size_gt_1(self): + """Helper works at page_size > 1: virtual ids and table indices + are page-granular. Verifies kernel scatters one v2p entry per + PAGE (not per token).""" + _pool, fa, _kv = self._make_full( + lazy=True, n_full_slots=64, n_mamba_slots=16, page_size=4 + ) + self.assertEqual(fa.page_size, 4) + v_pages = torch.tensor([3, 4, 5], dtype=torch.int64, device="cuda") + wm_before = fa.watermark_physical + phys = fa._alloc_bind_fast_or_slow(v_pages, 3) + self.assertIsNotNone(phys) + self.assertEqual(phys.shape, (3,)) + # Watermark advances by N PAGES (not N tokens). + self.assertEqual(fa.watermark_physical, wm_before + 3) + # v2p table updated at page granularity. + for v, p in zip(v_pages.tolist(), phys.tolist()): + self.assertEqual(int(fa.virtual_to_physical[v].item()), p) + self.assertEqual(int(fa.physical_to_virtual[p].item()), v) + + def test_grow_down_fast_path(self): + """The mamba sub-pool is grow-down. Verifies fast-path arithmetic + in the descending direction.""" + _pool, _fa, _kv = self._make_full(lazy=True) + # Build a grow-down allocator standalone for the test. + from sglang.srt.mem_cache.unified_memory_pool import ( + UnifiedKVPool, + ) + + full = _make_mha_spec("full", "up", layer_num=2) + swa = _make_mha_spec("swa", "down", layer_num=2) # grow-down + total = (full.entry_bytes() + swa.entry_bytes()) * 32 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa], + device="cuda", + enable_memory_saver=False, + ) + full_kv = _FakeKVCache(pool.max_slots("full")) + swa_kv = _FakeKVCache(pool.max_slots("swa")) + fa = MultiEndedAllocator( + kvcache=full_kv, + unified_buffer=pool, + sub_pool_name="full", + device="cuda", + is_id_owner=True, + lazy_compaction=True, + ) + sa = MultiEndedAllocator( + kvcache=swa_kv, + unified_buffer=pool, + sub_pool_name="swa", + device="cuda", + is_id_owner=False, # non-owner, grow-down + lazy_compaction=True, + ) + fa.bind_peer(sa) + sa.bind_peer(fa) + # Grow-down: watermark starts at num_pages - 1, decreases. + self.assertEqual(sa.grow_direction, "down") + wm_before = sa.watermark_physical + v_pages = torch.tensor([5, 6, 7], dtype=torch.int64, device="cuda") + phys = sa._alloc_bind_fast_or_slow(v_pages, 3) + # Grow-down: kernel emits ASCENDING (matches `_take_physical_eager`'s + # `torch.arange(wm-N+1, wm+1)` output). For wm=wm_before, N=3: + # range is [wm_before-2, wm_before-1, wm_before]. + expected = torch.tensor( + [wm_before - 2, wm_before - 1, wm_before], + dtype=torch.int64, + device="cuda", + ) + self.assertTrue(torch.equal(phys, expected)) + # Watermark decreased by N. + self.assertEqual(sa.watermark_physical, wm_before - 3) + # v2p / p2v consistent with the ascending mapping: + # v_pages[0] → wm_before - 2 (lowest of the new range) + # v_pages[2] → wm_before (highest of the new range) + for v, p in zip(v_pages.tolist(), expected.tolist()): + self.assertEqual(int(sa.virtual_to_physical[v].item()), p) + self.assertEqual(int(sa.physical_to_virtual[p].item()), v) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_store_cache_4d.py b/test/registered/unit/mem_cache/test_store_cache_4d.py index 70128452e..5340d5f2d 100644 --- a/test/registered/unit/mem_cache/test_store_cache_4d.py +++ b/test/registered/unit/mem_cache/test_store_cache_4d.py @@ -24,10 +24,10 @@ import torch from sglang.test.ci.ci_register import register_cuda_ci _HAS_CUDA = torch.cuda.is_available() -# The set_kv_buffer integration test needs SharedMHATokenToKVPool, which only -# exists once the shared-memory-pool feature lands; skip it where absent. +# The set_kv_buffer integration test needs UnifiedMHATokenToKVPool, which only +# exists once the shared-KV-pool feature lands; skip it where absent. _HAS_SHARED_POOL = ( - importlib.util.find_spec("sglang.srt.mem_cache.shared_memory_pool") is not None + importlib.util.find_spec("sglang.srt.mem_cache.unified_memory_pool") is not None ) register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") @@ -72,7 +72,7 @@ class TestStoreCache4D(unittest.TestCase): seed: int = 0xC0FFEE, ): torch.manual_seed(seed) - # The shared pool's views are 4-D `(num_pages, page_size, head_num, + # The unified memory pool's views are 4-D `(num_pages, page_size, head_num, # head_dim)` with the trailing two dims contiguous. We allocate two # independent contiguous buffers (one for the kernel-under-test, # one as the legacy-path target) so we can compare them. @@ -154,7 +154,7 @@ class TestStoreCache4D(unittest.TestCase): def test_store_cache_4d_ps1_byte_identical(self): """At page_size=1 the kernel constexpr-folds to the slot-major envelope view. Output must be byte-identical to advanced indexing. - This protects the Stage 1/2/3 green eval matrix from regression.""" + This protects against byte-layout regression.""" self._check_parity( num_pages=64, page_size=1, @@ -228,8 +228,7 @@ class TestStoreCache4D(unittest.TestCase): def test_store_cache_4d_dtype_fp8_e5m2(self): """fp8_e5m2 is used for KV-cache quantization. Caller is responsible - for the cast (Phase 1); the kernel sees same-dtype source and - destination.""" + for the cast; the kernel sees same-dtype source and destination.""" self._check_parity( num_pages=16, page_size=64, @@ -317,31 +316,25 @@ class TestStoreCache4DAssertions(unittest.TestCase): @unittest.skipUnless( _HAS_CUDA and _HAS_SHARED_POOL, - "Triton kernels require CUDA; SharedMHATokenToKVPool required", + "Triton kernels require CUDA; UnifiedMHATokenToKVPool required", ) class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase): """Integration parity test — exercises the kernel through the FULL - ``SharedMHATokenToKVPool.set_kv_buffer`` path, including the - ``_external_allocator`` v2p translation and the dtype cast. Confirms the - production code path produces bit-identical output to a PyTorch - advanced-indexing reference write. + ``UnifiedMHATokenToKVPool.set_kv_buffer`` path (the direct PHYSICAL write + + the dtype cast; the pool no longer translates). Confirms it produces + bit-identical output to a PyTorch advanced-indexing reference write. """ - def _build_pool_and_stub_alloc(self, page_size: int, v2p=None): - """Build a small SharedMHATokenToKVPool wired to a stub allocator. - - By default `virtual_to_physical` is identity (the kernel-vs-legacy - parity tests don't exercise virtual-id semantics). Pass an explicit - `v2p` tensor (sized `max_slots + 1`) to exercise a NON-identity - translation — used by the `set_full_loc` fast-path parity test, which - needs virtual != physical so the precomputed-physical fast path is - meaningfully different from the per-call gather.""" + def _build_pool(self, page_size: int): + """Build a small UnifiedMHATokenToKVPool. The pool writes PHYSICAL locs + directly (no allocator / v2p translate), so `set_kv_buffer` receives the + already-physical write location.""" import torch as _t - from sglang.srt.mem_cache.shared_memory_pool import ( + from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, - SharedMemoryPool, - SharedMHATokenToKVPool, + UnifiedKVPool, + UnifiedMHATokenToKVPool, ) spec = MHASubPoolSpec( @@ -362,15 +355,15 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase): store_dtype=_t.bfloat16, grow_direction="down", ) - pool = SharedMemoryPool( + pool = UnifiedKVPool( total_bytes=total + peer.entry_bytes() * 16, sub_pool_specs=[spec, peer], device="cuda", enable_memory_saver=False, page_size=page_size, ) - kv_pool = SharedMHATokenToKVPool( - shared_buffer=pool, + kv_pool = UnifiedMHATokenToKVPool( + unified_buffer=pool, sub_pool_name="full", page_size=page_size, start_layer=0, @@ -378,21 +371,12 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase): enable_alt_stream=False, ) - # Stub allocator with an identity (default) or caller-supplied v2p. - max_slots = pool.max_slots("full") - if v2p is None: - v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda") - - class _StubAllocator: - virtual_to_physical = v2p - - kv_pool.attach_allocator(_StubAllocator()) return kv_pool def _run_set_kv_buffer_and_compare(self, page_size: int): import torch as _t - kv_pool = self._build_pool_and_stub_alloc(page_size) + kv_pool = self._build_pool(page_size) # A fake `layer` object with the minimum interface # `set_kv_buffer` reads: `.layer_id`. @@ -415,9 +399,8 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase): k_kernel = kv_pool.k_buffer[0].clone() v_kernel = kv_pool.v_buffer[0].clone() - # Reference: PyTorch advanced-indexing into a fresh view. The stub - # allocator's v2p is identity, so physical loc == virtual loc and no - # dtype cast happens (store_dtype == dtype), making this the exact + # Reference: PyTorch advanced-indexing into a fresh view at the same + # (physical) loc, with no dtype cast (store_dtype == dtype) — the exact # write the kernel performs. kv_pool.k_buffer[0].zero_() kv_pool.v_buffer[0].zero_() @@ -449,81 +432,6 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase): def test_integration_ps64(self): self._run_set_kv_buffer_and_compare(page_size=64) - def _run_full_loc_fast_path_parity(self, page_size: int): - """Stage 3.5 fast-path byte-identity: writing through the precomputed - full-physical loc (`set_loc` fast path) must produce a byte-identical - KV buffer to writing the virtual loc and letting `set_kv_buffer` - translate per call. Uses a NON-identity v2p so the two paths are - genuinely different code (fast path skips the gather).""" - import torch as _t - - # Non-identity v2p: reverse-map the physical slot space so virtual i - # lands on a different physical slot. Keep slot 0 -> 0 (padding sink). - # Build the pool once to learn max_slots, then rebuild with the v2p. - probe = self._build_pool_and_stub_alloc(page_size) - max_slots = probe.k_buffer[0].shape[0] * page_size - v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda") - # Shuffle the interior [1, max_slots) so virtual != physical, leave - # 0 (sink) and the trailing sentinel (max_slots -> itself) alone. - interior = _t.randperm(max_slots - 1, device="cuda") + 1 - v2p[1:max_slots] = interior - - kv_pool = self._build_pool_and_stub_alloc(page_size, v2p=v2p) - - class _FakeLayer: - layer_id = 0 - - layer = _FakeLayer() - head_num, head_dim = 4, 64 - N = 16 - num_pages = kv_pool.k_buffer[0].shape[0] - total = num_pages * page_size - # Draw virtual ids from [1, total) (avoid the padding sink at 0). - loc = (_t.randperm(total - 1, device="cuda")[:N] + 1).to(_t.int64) - cache_k = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda") - cache_v = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda") - - # SLOW path: no precompute pinned -> per-call v2p gather inside - # set_kv_buffer translates virtual -> physical. - kv_pool.set_loc(None) - kv_pool.set_kv_buffer(layer, loc, cache_k.clone(), cache_v.clone()) - k_slow = kv_pool.k_buffer[0].clone() - v_slow = kv_pool.v_buffer[0].clone() - - # FAST path: precompute the full-physical loc exactly as - # `set_kv_buffer`'s page math would, pin it via set_loc, and pass - # it as `loc` so the data-ptr fast path fires (no gather). - if page_size == 1: - phys = _t.clamp_min(v2p[loc], 0) - else: - virt_pages = loc // page_size - offsets = loc % page_size - phys = _t.clamp_min(v2p[virt_pages] * page_size + offsets, 0) - kv_pool.k_buffer[0].zero_() - kv_pool.v_buffer[0].zero_() - kv_pool.set_loc(phys) - try: - kv_pool.set_kv_buffer(layer, phys, cache_k.clone(), cache_v.clone()) - k_fast = kv_pool.k_buffer[0].clone() - v_fast = kv_pool.v_buffer[0].clone() - finally: - kv_pool.set_loc(None) - - self.assertTrue( - _t.equal(k_fast, k_slow), - f"K mismatch: full_loc fast path != per-call translate at ps={page_size}", - ) - self.assertTrue( - _t.equal(v_fast, v_slow), - f"V mismatch: full_loc fast path != per-call translate at ps={page_size}", - ) - - def test_full_loc_fast_path_parity_ps1(self): - self._run_full_loc_fast_path_parity(page_size=1) - - def test_full_loc_fast_path_parity_ps64(self): - self._run_full_loc_fast_path_parity(page_size=64) - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_mamba_views.py b/test/registered/unit/mem_cache/test_unified_mamba_views.py new file mode 100644 index 000000000..15871265e --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_mamba_views.py @@ -0,0 +1,291 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` — the +envelope-strided conv/temporal (SSM) state views that back ``UnifiedMambaPool``. + +This isolates the unified-memory-pool Mamba STATE layout from the full model. It guards +against a class of correctness defect where Falcon-H1 greedy decode is garbled +under the unified memory pool, isolated to the Mamba conv/temporal state path: a +stride/offset/alignment bug in the view construction (analogous to the fixed +`_extract_kv_strides` MHA bug). + +Within one slot's envelope the bytes are +``[conv[0]·L0 | conv[0]·L1 | ... | conv[1]·L0 | ... | temporal·L0 | ...]`` and +across slots the layout is envelope (slot stride == entry_bytes). Each returned +view is ``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B) +and temporal dtype (fp32, 4 B) DIFFER, so the temporal view's byte offset must +be a multiple of the temporal itemsize — an alignment hazard that +``_build_mamba_views`` now asserts. + +These tests prove the views: + - round-trip every (tensor, layer, slot) element with the Falcon-like + bf16-conv / fp32-temporal dtype mix (catches stride/offset/alignment bugs); + - do NOT alias each other (conv[i] vs conv[j] vs temporal) or across + layers/slots (catches envelope-overlap); + - match a contiguous ``(num_layers, max_slots, *inner)`` reference exactly + (the shape `MambaPool.State.conv[i]` / `.temporal` expose); + - reject a deliberately mis-aligned spec via the alignment assert. + +Skipped on CPU — these views back GPU kernels and we mirror the GPU path. + + python -m pytest test/registered/unit/mem_cache/test_shared_mamba_views.py -v +""" + +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +_HAS_CUDA = torch.cuda.is_available() +_DEV = "cuda" if _HAS_CUDA else "cpu" + +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") + + +def _make_pool( + *, + mamba_layer_num, + conv_state_shapes, + conv_dtype, + temporal_state_shape, + temporal_dtype, + want_slots=8, + device=_DEV, +): + """Build a minimal 2-sub-pool ``UnifiedKVPool`` (a small MHA grow-up peer + + the Mamba grow-down pool under test) sized to hold >= ``want_slots`` Mamba + slots, and return ``(pool, mamba_spec)``.""" + from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedKVPool, + ) + + mamba_spec = MambaSubPoolSpec( + name="mamba", + layer_num=mamba_layer_num, + grow_direction="down", + conv_state_shapes=tuple(tuple(s) for s in conv_state_shapes), + conv_dtype=conv_dtype, + temporal_state_shape=tuple(temporal_state_shape), + temporal_dtype=temporal_dtype, + ) + # Tiny full-attention peer (required: exactly one grow-up + one grow-down). + full_spec = MHASubPoolSpec( + name="full", + layer_num=1, + head_num=1, + head_dim=8, + store_dtype=torch.bfloat16, + grow_direction="up", + ) + entry_mamba = mamba_spec.entry_bytes() + entry_full = full_spec.entry_bytes() + entry_max = max(entry_mamba, entry_full) + # Need max_slots("mamba") = total // entry_mamba >= want_slots, and total + # large enough that BOTH pools clear their min_slot_index. Add generous + # headroom, then round up to a multiple of 8 (covers bf16/fp32 .view()). + total_bytes = want_slots * entry_mamba + 8 * entry_max + total_bytes = ((total_bytes + 7) // 8) * 8 + pool = UnifiedKVPool( + total_bytes=total_bytes, + sub_pool_specs=[full_spec, mamba_spec], + device=device, + enable_memory_saver=False, + ) + return pool, mamba_spec + + +@unittest.skipUnless(_HAS_CUDA, "shared Mamba views back GPU kernels") +class TestUnifiedMambaViews(unittest.TestCase): + # Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal, several + # layers. (Mamba2 conv state is (conv_dim, kernel-1); temporal/SSM state is + # (nheads, head_dim, ssm_state_size).) + FALCON_KW = dict( + mamba_layer_num=5, # odd, to stress the temporal-offset alignment + conv_state_shapes=[(48, 3)], # conv_dim=48, kernel-1=3 + conv_dtype=torch.bfloat16, + temporal_state_shape=(6, 8, 16), # nheads, head_dim, ssm_state + temporal_dtype=torch.float32, + ) + + def _fill_and_roundtrip(self, pool, mamba_spec): + """Write a distinct random tensor to each conv view + the temporal view + (in their own dtypes), then read all back and assert exact equality. + Writing ALL views first and reading ALL after means any envelope overlap + (conv[i]/conv[j]/temporal aliasing) corrupts an earlier write → mismatch. + """ + conv_views, temporal_view = pool.mamba_views_for("mamba") + torch.manual_seed(0) + refs = [] + for v in conv_views: + r = torch.randn(v.shape, device=v.device).to(v.dtype) + v.copy_(r) + refs.append(r) + rt = torch.randn(temporal_view.shape, device=temporal_view.device).to( + temporal_view.dtype + ) + temporal_view.copy_(rt) + refs.append(rt) + # Read back AFTER all writes. + for i, v in enumerate(conv_views): + self.assertTrue( + torch.equal(v, refs[i]), + f"conv view[{i}] round-trip mismatch (stride/offset/overlap " + f"bug); shape={tuple(v.shape)} stride={v.stride()}", + ) + self.assertTrue( + torch.equal(temporal_view, refs[-1]), + f"temporal view round-trip mismatch; shape={tuple(temporal_view.shape)} " + f"stride={temporal_view.stride()}", + ) + + def test_roundtrip_falcon_like(self): + pool, spec = _make_pool(**self.FALCON_KW) + self._fill_and_roundtrip(pool, spec) + + def test_roundtrip_single_layer_single_slot_edges(self): + # 1 layer, multiple conv tensors, same-dtype conv/temporal. + pool, spec = _make_pool( + mamba_layer_num=1, + conv_state_shapes=[(16, 3), (8, 3)], + conv_dtype=torch.float32, + temporal_state_shape=(4, 8, 16), + temporal_dtype=torch.float32, + want_slots=4, + ) + self._fill_and_roundtrip(pool, spec) + + def test_roundtrip_multi_conv_tensors(self): + # Two conv tensors + bf16/fp32 mix — exercises the per-conv-tensor offset + # accumulation in _build_mamba_views. + pool, spec = _make_pool( + mamba_layer_num=3, + conv_state_shapes=[(32, 3), (16, 3)], + conv_dtype=torch.bfloat16, + temporal_state_shape=(8, 8, 16), + temporal_dtype=torch.float32, + want_slots=6, + ) + self._fill_and_roundtrip(pool, spec) + + def test_no_cross_region_overlap(self): + """Zero buffer; write a sentinel to ONE view; every OTHER view must read + all-zero. Pinpoints conv[i]/conv[j]/temporal aliasing if present.""" + pool, spec = _make_pool(**self.FALCON_KW) + conv_views, temporal_view = pool.mamba_views_for("mamba") + views = list(conv_views) + [temporal_view] + names = [f"conv[{i}]" for i in range(len(conv_views))] + ["temporal"] + for target in range(len(views)): + pool._raw.zero_() + views[target].fill_(7.0) + for other in range(len(views)): + if other == target: + self.assertTrue( + bool((views[other] == 7.0).all().item()), + f"write to {names[target]} did not fully land", + ) + continue + self.assertTrue( + bool((views[other] == 0).all().item()), + f"writing {names[target]} CORRUPTED {names[other]} " + f"(envelope regions overlap)", + ) + + def test_per_layer_per_slot_addressing(self): + """Distinct value per (layer, slot) on the temporal view; verify exact + addressing (no layer/slot aliasing). Uses small integers exactly + representable in the view dtype. + + NB: ``temporal_view`` is a non-contiguous strided view, so we must NOT + ``.reshape()`` it (that would COPY, breaking the alias) — we + broadcast-assign into the view in place and read back via basic + indexing (which keeps the view).""" + pool, spec = _make_pool(**self.FALCON_KW) + _, temporal_view = pool.mamba_views_for("mamba") + N, S = temporal_view.shape[0], temporal_view.shape[1] + inner_ndim = temporal_view.dim() - 2 + # value = layer*S + slot (< N*S, small → exact in fp32) + base = ( + torch.arange(N, device=temporal_view.device)[:, None] * S + + torch.arange(S, device=temporal_view.device)[None, :] + ).to(temporal_view.dtype) + # Broadcast (N, S) over the inner dims, in place into the strided view. + temporal_view[:] = base.view(N, S, *([1] * inner_ndim)) + # Read back the first inner element of every (layer, slot) via basic + # indexing (stays a view). + readback = temporal_view[(slice(None), slice(None)) + (0,) * inner_ndim] + self.assertTrue( + torch.equal(readback, base), + "temporal (layer, slot) addressing wrong — layer/slot stride bug", + ) + + def test_matches_contiguous_reference(self): + """The shared view must be a faithful relabeling of a contiguous + ``(num_layers, max_slots, *inner)`` tensor: identical data written by the + same logical index reads back identically.""" + pool, spec = _make_pool(**self.FALCON_KW) + conv_views, temporal_view = pool.mamba_views_for("mamba") + for v in conv_views + [temporal_view]: + ref = torch.randn(v.shape, device=v.device).to(v.dtype) + contig = ref.clone().contiguous() + v.copy_(ref) + self.assertEqual(tuple(v.shape), tuple(contig.shape)) + self.assertTrue( + torch.equal(v.contiguous(), contig), + "shared view not equivalent to its contiguous counterpart", + ) + + def test_alignment_guard_fires_on_misaligned_spec(self): + """A spec whose conv region (bf16) is an odd multiple of 2 B makes the + per-slot entry (= conv_region + N*temporal_row = 2 B + 4 B = 6 B) NOT a + multiple of the temporal itemsize (fp32, 4 B). The temporal/SSM-state + view's storage_offset is computed by integer-dividing a byte offset by + the temporal itemsize, so this would silently mis-offset the view. + ``_build_mamba_views`` must reject it with a loud alignment assert. + + NOTE: the ``entry_bytes % itemsize`` guard is what fires here, and it + subsumes the conv-region offset check (see the comment in + ``_build_mamba_views``). We assert on the shared "misaligned" wording + rather than on which specific guard trips.""" + with self.assertRaises(AssertionError) as cm: + _make_pool( + mamba_layer_num=1, # entry = 2 B conv + 4 B temporal = 6 B, not %4 + conv_state_shapes=[(1, 1)], + conv_dtype=torch.bfloat16, + temporal_state_shape=(1,), + temporal_dtype=torch.float32, + want_slots=4, + ) + self.assertIn("misalign", str(cm.exception).lower()) + + def test_alignment_ok_for_aligned_spec(self): + """An aligned spec (conv region a multiple of the temporal itemsize) + must build and round-trip cleanly.""" + # conv region = N * conv_dim*(k-1) * 2 ; with conv_dim=2 -> per-layer 2*3*2=12, + # times N=2 = 24, divisible by 4. Aligned. + pool, spec = _make_pool( + mamba_layer_num=2, + conv_state_shapes=[(2, 3)], + conv_dtype=torch.bfloat16, + temporal_state_shape=(2, 4, 4), + temporal_dtype=torch.float32, + want_slots=4, + ) + self._fill_and_roundtrip(pool, spec) + + +if __name__ == "__main__": + unittest.main()