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()