diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py index 928d10156..8fe3b4827 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py @@ -275,6 +275,7 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool): self.auxiliary_state_pool.free(track_buffer) req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None + req.mamba_last_track_idx = None def free_auxiliary_state_cache(self, req, track_buffer_to_keep=None): self.free_mamba_cache( @@ -385,6 +386,7 @@ class MlxAuxiliaryStateComponent(MambaComponent): self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer) req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None + req.mamba_last_track_idx = None req.mamba_last_track_seqlen = None return @@ -409,5 +411,6 @@ class MlxAuxiliaryStateComponent(MambaComponent): self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer) req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None + req.mamba_last_track_idx = None req.mamba_pool_idx = None req.mamba_last_track_seqlen = None diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 3c91ab47a..c69c0b073 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -399,6 +399,7 @@ class MlxModelRunner: return req.mamba_ping_pong_track_buffer = track_buffer req.mamba_next_track_idx = 0 + req.mamba_last_track_idx = 0 pool.store_cache( track_buffer[0], diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c273912b0..3b9152f4c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -939,6 +939,7 @@ class Req(ReqDllmMixin): self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1) self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2) self.mamba_next_track_idx: Optional[int] = None # 0 or 1 + self.mamba_last_track_idx: Optional[int] = None # 0 or 1 self.mamba_last_track_seqlen: Optional[int] = ( None # seq len of the last cached mamba state ) @@ -1693,6 +1694,7 @@ class Req(ReqDllmMixin): self.mamba_pool_idx = None self.mamba_ping_pong_track_buffer = None self.mamba_next_track_idx = None + self.mamba_last_track_idx = None self.mamba_last_track_seqlen = None self.mamba_branching_seqlen = None self.mamba_cow_src_index = None @@ -1883,6 +1885,7 @@ def set_mamba_track_indices_from_reqs( req.mamba_next_track_idx if req.mamba_next_track_idx is not None else 0 for req in batch.reqs ] + batch.mamba_track_buffer_indices = list(track_positions) idx = ( torch.tensor( track_positions, @@ -2079,6 +2082,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For hybrid GDN prefix cache mamba_track_indices: torch.Tensor = None # shape: [b], int64 + # Per-batch snapshot of the logical ping-pong positions selected for this + # forward (normally req.mamba_next_track_idx; spec may override it). Result + # processing uses it to update req.mamba_last_track_idx, since both req-level + # indices may advance under overlap. + mamba_track_buffer_indices: Optional[List[int]] = None # shape: [b], 0 or 1 mamba_track_mask: torch.Tensor = None # shape: [b], bool mamba_track_seqlens: torch.Tensor = None # shape: [b], int64 mamba_track_mask_cpu: Optional[List[bool]] = None # shape: [b] @@ -2666,6 +2674,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # In lazy mode, skip the swap — the second ping-pong slot is not # allocated yet; it will be allocated on demand at the track boundary # in mamba_lazy_prealloc_at_boundary during prepare_for_decode. + req.mamba_last_track_idx = req.mamba_next_track_idx if not mamba_extra_buffer_lazy_enabled(): req.mamba_next_track_idx = ( self.req_to_token_pool.get_mamba_ping_pong_other_idx( @@ -3067,6 +3076,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.mamba_track_indices = torch.empty( (0,), dtype=torch.int64, device=self.device ) + self.mamba_track_buffer_indices = [] else: if mamba_extra_buffer_lazy_enabled(): self.mamba_lazy_prealloc_at_boundary(mamba_track_interval) @@ -3148,6 +3158,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.seq_lens_cpu = self.seq_lens_cpu[keep_indices] self.mamba_track_indices = None + self.mamba_track_buffer_indices = None self.mamba_track_mask = None self.mamba_track_seqlens = None self.mamba_track_mask_cpu = None @@ -3211,6 +3222,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): else: self.seq_lens_cpu = torch.cat([self.seq_lens_cpu, other.seq_lens_cpu]) self.mamba_track_indices = None + self.mamba_track_buffer_indices = None self.mamba_track_mask = None self.mamba_track_seqlens = None self.mamba_track_mask_cpu = None @@ -3272,6 +3284,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): seq_lens_cpu=self.seq_lens_cpu, enable_overlap=self.enable_overlap, mamba_track_indices=self.mamba_track_indices, + mamba_track_buffer_indices=self.mamba_track_buffer_indices, mamba_track_mask=self.mamba_track_mask, mamba_track_seqlens=self.mamba_track_seqlens, mamba_track_mask_cpu=self.mamba_track_mask_cpu, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index c1e17365e..c98becd59 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1016,8 +1016,12 @@ class SchedulerBatchResultProcessor: i: int, logits_output: LogitsProcessorOutput, ): + lazy = mamba_extra_buffer_lazy_enabled() known_mamba_boundary = None + completed_mamba_boundary = None + lookahead = 0 if batch.mamba_track_mask_cpu is not None: + completed_mamba_boundary = bool(batch.mamba_track_mask_cpu[i]) lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i] assert lookahead in (0, 1), ( f"mamba result lookahead={lookahead} for req {req.rid}; " @@ -1028,15 +1032,28 @@ class SchedulerBatchResultProcessor: else: known_mamba_boundary = bool(batch.mamba_track_mask_next_cpu[i]) + if completed_mamba_boundary and not lazy: + req.mamba_last_track_idx = batch.mamba_track_buffer_indices[i] + req.mamba_last_track_seqlen = req.kv_committed_len - lookahead + elif ( + req.finished() + and lazy + and lookahead == 1 + and known_mamba_boundary + and req.mamba_next_track_idx == req.mamba_last_track_idx + ): + req.mamba_lazy_is_insert = False + # Called here (after update_finish_state) so req.finished() is valid # for mamba_lazy_post_decode_at_boundary inside. - if known_mamba_boundary is None or known_mamba_boundary: + should_update = completed_mamba_boundary if lazy else known_mamba_boundary + if should_update is None or should_update: self._mamba_prefix_cache_update( req, batch, result, i, - known_boundary=known_mamba_boundary is True, + known_boundary=not lazy and known_mamba_boundary is True, ) if ( @@ -1129,14 +1146,18 @@ class SchedulerBatchResultProcessor: if not at_boundary: return - req.mamba_last_track_seqlen = track_seqlen + track_idx = req.mamba_next_track_idx + if not known_boundary and batch.mamba_track_buffer_indices is not None: + track_idx = batch.mamba_track_buffer_indices[i] + if not known_boundary: + req.mamba_last_track_seqlen = track_seqlen if lazy: - self.mamba_lazy_post_decode_at_boundary(req, batch) + self.mamba_lazy_post_decode_at_boundary(req, batch, track_idx) else: + if not known_boundary: + req.mamba_last_track_idx = track_idx req.mamba_next_track_idx = ( - batch.req_to_token_pool.get_mamba_ping_pong_other_idx( - req.mamba_next_track_idx - ) + batch.req_to_token_pool.get_mamba_ping_pong_other_idx(track_idx) ) def _mamba_lazy_spec_update( @@ -1193,6 +1214,7 @@ class SchedulerBatchResultProcessor: req.mamba_next_track_idx = planned_pos # else: in-place fallback, or promoted by an earlier confirmation — # keep holds the track_seqlen state either way. + req.mamba_last_track_idx = planned_pos req.mamba_last_track_seqlen = track_seqlen @staticmethod @@ -1216,9 +1238,8 @@ class SchedulerBatchResultProcessor: matches what the forward's tracking mask used: ``prepare_for_decode`` increments both ``seq_lens_cpu`` and ``kv_committed_len`` by 1, then checks - ``seq_lens_cpu % interval == 0``. Using ``kv_committed_len`` - here reproduces that check exactly, and the value is always a - multiple of ``interval`` (hence page-aligned). + ``seq_lens_cpu % interval == 0``. Subtracting the overlap + lookahead from ``kv_committed_len`` reproduces that check. For spec decode, the boundary is detected by comparing the accepted seq_len range against interval boundaries. @@ -1226,9 +1247,10 @@ class SchedulerBatchResultProcessor: interval = get_exec().mamba.mamba_track_interval if batch.spec_algorithm.is_none(): - self._mamba_assert_committed_len_lookahead(req) - if req.kv_committed_len % interval == 0: - return True, req.kv_committed_len + lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i] + committed_len = req.kv_committed_len - lookahead + if committed_len % interval == 0: + return True, committed_len elif result.num_correct_drafts_per_req_cpu is not None: cur = req.seqlen - 1 prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1 @@ -1237,19 +1259,13 @@ class SchedulerBatchResultProcessor: return False, 0 - def mamba_lazy_post_decode_at_boundary(self, req: Req, batch: ScheduleBatch): - """Post-decode cleanup at a lazy-mode track boundary. - - Finished reqs: if prealloc failed (other slot is -1), the forward - overwrote the only slot with corrupted state, so mark - is_insert=False to skip the cache insert. If the other slot is - occupied (stale prealloc from an overlap extra forward), free it - so the prealloc assert in the next prepare_for_decode holds. - - Running reqs: free the old ping-pong slot so we go back to - holding only 1 slot until the next boundary. - """ - other_idx = 1 - req.mamba_next_track_idx + def mamba_lazy_post_decode_at_boundary( + self, req: Req, batch: ScheduleBatch, track_idx: int + ): + """Commit a completed lazy-mode boundary and free its old slot.""" + req.mamba_last_track_idx = track_idx + req.mamba_next_track_idx = track_idx + other_idx = 1 - track_idx other_val = req.mamba_ping_pong_track_buffer[other_idx].item() if other_val != -1: pool = batch.req_to_token_pool @@ -1257,5 +1273,3 @@ class SchedulerBatchResultProcessor: req.mamba_ping_pong_track_buffer[other_idx].unsqueeze(0) ) pool.set_mamba_ping_pong_slot(req, other_idx, -1) - elif req.finished(): - req.mamba_lazy_is_insert = False diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 3da4022d9..5f3728ce8 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1388,14 +1388,8 @@ class HybridReqToTokenPool(ReqToTokenPool): return mamba_next_track_idx def get_mamba_ping_pong_keep_idx(self, req: Req) -> int: - """Return the ping-pong index holding the most recent tracked state. - - In lazy mode the valid state stays at next_track_idx (no eager swap). - In normal mode it is at the "other" index (swapped after each track). - """ - if self.enable_mamba_extra_buffer_lazy: - return req.mamba_next_track_idx - return self.get_mamba_ping_pong_other_idx(req.mamba_next_track_idx) + """Return the ping-pong index holding the most recent tracked state.""" + return req.mamba_last_track_idx def _alloc_ping_pong_buffer(self, req: Req): """Allocate the ping-pong track buffer for a new request. @@ -1422,6 +1416,11 @@ class HybridReqToTokenPool(ReqToTokenPool): buf[:n] = slots req.mamba_ping_pong_track_buffer = buf req.mamba_next_track_idx = 0 + req.mamba_last_track_idx = ( + 0 + if self.enable_mamba_extra_buffer_lazy + else self.get_mamba_ping_pong_other_idx(0) + ) def set_mamba_ping_pong_slot(self, req: Req, idx: int, value): """Update a ping-pong slot value and sync the device-side mapping. @@ -1442,8 +1441,6 @@ class HybridReqToTokenPool(ReqToTokenPool): Returns the old slot index (shape [1]) for cache insertion and replaces it with new_slot so the request can continue tracking. - In lazy mode the valid state is at next_track_idx; in normal mode - it is at the "other" index. """ donate_idx = self.get_mamba_ping_pong_keep_idx(req) mamba_value_donated = ( @@ -1512,6 +1509,7 @@ class HybridReqToTokenPool(ReqToTokenPool): # tensor on the req side while the new pool slot leaks). req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None + req.mamba_last_track_idx = None def clear(self): logger.info("Reset HybridReqToTokenPool") diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 95f45db47..a29180c77 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -60,6 +60,7 @@ class SessionSlot: mamba_pool_idx: Any = None mamba_ping_pong_track_buffer: Any = None mamba_next_track_idx: Any = None + mamba_last_track_idx: Any = None mamba_last_track_seqlen: Any = None mamba_branching_seqlen: Any = None @@ -83,6 +84,7 @@ class SessionSlot: self.mamba_pool_idx = req.mamba_pool_idx self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer self.mamba_next_track_idx = req.mamba_next_track_idx + self.mamba_last_track_idx = req.mamba_last_track_idx self.mamba_last_track_seqlen = req.mamba_last_track_seqlen self.mamba_branching_seqlen = req.mamba_branching_seqlen @@ -99,6 +101,7 @@ class SessionSlot: req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None + req.mamba_last_track_idx = None req.mamba_last_track_seqlen = None req.mamba_branching_seqlen = None @@ -113,6 +116,7 @@ class SessionSlot: req.mamba_pool_idx = self.mamba_pool_idx req.mamba_ping_pong_track_buffer = self.mamba_ping_pong_track_buffer req.mamba_next_track_idx = self.mamba_next_track_idx + req.mamba_last_track_idx = self.mamba_last_track_idx req.mamba_last_track_seqlen = self.mamba_last_track_seqlen req.mamba_branching_seqlen = self.mamba_branching_seqlen diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py index cc6b27dcc..0feeb976c 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py @@ -110,9 +110,9 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): "--hicache-write-policy", "write_through", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--max-total-tokens", "12000", "--max-mamba-cache-size", @@ -171,9 +171,9 @@ class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): "--hicache-storage-prefetch-policy", "wait_complete", "--hicache-io-backend", - "direct", + "kernel", "--hicache-mem-layout", - "page_first_direct", + "page_first", "--hicache-storage-backend", "file", "--max-mamba-cache-size",