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 30ca32f4c..f4a99c4e7 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -140,6 +140,7 @@ class MambaAttnBackendBase(AttentionBackend): self.pad_slot_id = PAD_SLOT_ID self.device = model_runner.device self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.is_draft_worker = model_runner.is_draft_worker self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool self.forward_metadata: ForwardMetadata = None self.state_indices_list = [] @@ -151,6 +152,28 @@ class MambaAttnBackendBase(AttentionBackend): self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None self.conv_states_shape: tuple[int, int] = None + def _execute_deferred_mamba_cow_and_clear(self, forward_batch: ForwardBatch): + """Run deferred clear/COW ops on the forward stream to avoid races.""" + if not forward_batch.forward_mode.is_extend() or self.is_draft_worker: + return + if ( + forward_batch.mamba_clear_indices is not None + and len(forward_batch.mamba_clear_indices) > 0 + ): + self.req_to_token_pool.mamba_pool.clear_slots( + forward_batch.mamba_clear_indices + ) + if ( + forward_batch.mamba_cow_src_indices is not None + and len(forward_batch.mamba_cow_src_indices) > 0 + ): + self.req_to_token_pool.mamba_pool.copy_from( + forward_batch.mamba_cow_src_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 _forward_metadata(self, forward_batch: ForwardBatch): bs = forward_batch.batch_size @@ -240,6 +263,7 @@ class MambaAttnBackendBase(AttentionBackend): ) def init_forward_metadata(self, forward_batch: ForwardBatch): + self._execute_deferred_mamba_cow_and_clear(forward_batch) self.forward_metadata = self._forward_metadata(forward_batch) def _init_track_conv_indices( @@ -645,6 +669,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): self.mamba_chunk_size = config.mamba_chunk_size def init_forward_metadata(self, forward_batch: ForwardBatch): + self._execute_deferred_mamba_cow_and_clear(forward_batch) metadata = self._forward_metadata(forward_batch) self.forward_metadata = Mamba2Metadata.prepare_mixed( metadata, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ca8215909..b7a89b0b4 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -726,6 +726,10 @@ class Req(ReqDllmMixin): # the branching point seqlen to track mamba state. If set, given by prefix match, # it will be the tracked seqlen in the ping pong buffer for the right prefill pass. self.mamba_branching_seqlen: Optional[int] = None + # Deferred COW: source mamba pool index from radix cache node (copy on forward stream) + self.mamba_cow_src_index: Optional[torch.Tensor] = None + # Deferred clear: newly allocated mamba slot needs zeroing on forward stream + self.mamba_needs_clear: bool = False # Check finish self.tokenizer = None @@ -1282,6 +1286,8 @@ class Req(ReqDllmMixin): self.mamba_next_track_idx = None self.mamba_last_track_seqlen = None self.mamba_branching_seqlen = None + self.mamba_cow_src_index = None + self.mamba_needs_clear = False self.already_computed = 0 self.kv_allocated_len = 0 self.kv_committed_len = 0 @@ -1398,6 +1404,26 @@ class _MambaRadixCacheV2TrackEntry(NamedTuple): track_seqlen: int +def set_mamba_track_indices_from_reqs(batch): + """Build mamba_track_indices from req objects (authoritative source).""" + req_to_token_pool = batch.req_to_token_pool + all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[ + batch.req_pool_indices + ] # (bs, ping_pong_size), int64, on device + idx = ( + torch.tensor( + [req.mamba_next_track_idx for req in batch.reqs], + dtype=torch.int64, + pin_memory=True, + ) + .unsqueeze(1) + .to(device=all_buffers.device, non_blocking=True) + ) + batch.mamba_track_indices = ( + torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64) + ) + + @dataclasses.dataclass class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): """Store all information of a batch on the scheduler.""" @@ -1443,6 +1469,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_indices: torch.Tensor = None # shape: [b], int64 mamba_track_mask: torch.Tensor = None # shape: [b], bool mamba_track_seqlens: torch.Tensor = None # shape: [b], int64 + # Deferred mamba init ops: COW pairs and clear indices (performed on forward stream) + mamba_cow_src_indices: torch.Tensor = None + mamba_cow_dst_indices: torch.Tensor = None + mamba_clear_indices: torch.Tensor = None # For multimodal inputs multimodal_inputs: Optional[List] = None @@ -2003,6 +2033,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): device=self.device, ) + # Collect mamba init info for deferred ops on forward stream + if any(req.mamba_pool_idx is not None for req in reqs): + self._collect_deferred_mamba_cow_and_clear(reqs) + if self.model_config.is_encoder_decoder: self.prepare_encoder_info_extend(input_ids, seq_lens) @@ -2084,12 +2118,40 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen) mamba_track_seqlen_aligned = req.mamba_branching_seqlen req.mamba_last_track_seqlen = mamba_track_seqlen_aligned + return _MambaRadixCacheV2TrackEntry( track_mask=mask, track_index=track_index, track_seqlen=mamba_track_seqlen, ) + def _collect_deferred_mamba_cow_and_clear(self, reqs): + """Collect deferred COW/clear info from requests.""" + cow_src_tensors = [] + cow_dst_tensors = [] + clear_tensors = [] + for req in reqs: + if req.mamba_cow_src_index is not None: + cow_src_tensors.append(req.mamba_cow_src_index) + cow_dst_tensors.append(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_cow_src_index = None + req.mamba_needs_clear = False + elif req.mamba_needs_clear: + clear_tensors.append(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_needs_clear = False + self.mamba_cow_src_indices = ( + torch.cat(cow_src_tensors) if cow_src_tensors else None + ) + self.mamba_cow_dst_indices = ( + torch.cat(cow_dst_tensors) if cow_dst_tensors else None + ) + self.mamba_clear_indices = torch.cat(clear_tensors) if clear_tensors else None + + def prepare_for_split_prefill(self): + self.prepare_for_extend() + # For split prefill, we need to set the forward mode to SPLIT_PREFILL + self.forward_mode = ForwardMode.SPLIT_PREFILL + def mix_with_running(self, running_batch: "ScheduleBatch"): self.forward_mode = ForwardMode.MIXED running_bs = running_batch.batch_size() @@ -2379,22 +2441,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): (0,), dtype=torch.int64, device=self.device ) else: - # already on device - all_buffers = torch.stack( - [req.mamba_ping_pong_track_buffer for req in self.reqs] - ) - idx = ( - torch.tensor( - [req.mamba_next_track_idx for req in self.reqs], - dtype=torch.int64, - pin_memory=True, - ) - .unsqueeze(1) - .to(device=all_buffers.device, non_blocking=True) - ) - self.mamba_track_indices = ( - torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64) - ) + set_mamba_track_indices_from_reqs(self) # async H2D self.mamba_track_mask = ( @@ -2467,6 +2514,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.mamba_track_indices = None self.mamba_track_mask = None self.mamba_track_seqlens = None + self.mamba_cow_src_indices = None + self.mamba_cow_dst_indices = None + self.mamba_clear_indices = None self.return_logprob = any(req.return_logprob for req in self.reqs) if self.return_logprob: self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in keep_indices] diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 0475b3515..76a5426aa 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -1049,13 +1049,9 @@ class HiMambaRadixCache(MambaRadixCache): lock_node=mamba_node, error_message="Can not alloc mamba cache", ) - src_index = mamba_node.mamba_value - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) req.mamba_pool_idx = dst_index[0] - else: - src_index = mamba_node.mamba_value - dst_index = req.mamba_pool_idx.unsqueeze(0) - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + req.mamba_cow_src_index = mamba_node.mamba_value + req.mamba_needs_clear = False value = value[:best_value_len] if value: diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 2462af688..81d421613 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -639,45 +639,44 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): page_aligned_token_ids = token_ids[:page_aligned_len] + # Donate the mamba index to the radix cache instead of copying. + # This avoids a data copy that would race with the forward stream. if self.enable_mamba_extra_buffer: - # copy from the ping pong track buffer mamba_ping_pong_track_buffer_to_keep = ( self.req_to_token_pool.get_mamba_ping_pong_other_idx( req.mamba_next_track_idx ) ) - mamba_value = ( + mamba_value_donated = ( req.mamba_ping_pong_track_buffer[mamba_ping_pong_track_buffer_to_keep] .unsqueeze(-1) .clone() ) - else: - mamba_value = self.req_to_token_pool.get_mamba_indices( - req.req_pool_idx - ).unsqueeze(-1) - # radix tree mamba value is forked from req space - mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from(mamba_value) - - # if alloc mamba cache failed, do evict and alloc again - if mamba_value_forked is None: - self.evict(EvictParams(num_tokens=0, mamba_num=1)) - mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from( - mamba_value + new_slot = self._alloc_mamba_slot() + req.mamba_ping_pong_track_buffer[mamba_ping_pong_track_buffer_to_keep] = ( + new_slot[0] ) - assert mamba_value_forked is not None, "Can not alloc mamba cache" + self.req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[ + req.req_pool_idx + ] = req.mamba_ping_pong_track_buffer + else: + mamba_value_donated = self._alloc_mamba_slot() + self.req_to_token_pool.mamba_pool.copy_from( + req.mamba_pool_idx.unsqueeze(0), mamba_value_donated + ) + result = self.insert( InsertParams( key=RadixKey(page_aligned_token_ids, req.extra_key), value=page_aligned_kv_indices, - mamba_value=mamba_value_forked, + mamba_value=mamba_value_donated, prev_prefix_len=req.cache_protected_len, chunked=chunked, ) ) new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist - # there is a mamba cache in radix cache, release it if mamba_exist: - self.req_to_token_pool.mamba_pool.free(mamba_value_forked) + self.req_to_token_pool.mamba_pool.free(mamba_value_donated) # The prefix indices could be updated, reuse it match_result = self.match_prefix( @@ -689,7 +688,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ) if not mamba_exist: - assert torch.equal(new_last_node.mamba_value, mamba_value_forked) + assert torch.equal(new_last_node.mamba_value, mamba_value_donated) assert ( req.cache_protected_len <= len(new_indices) + self.page_size - 1 @@ -951,6 +950,15 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ##### Internal Helper Functions ##### + def _alloc_mamba_slot(self) -> torch.Tensor: + """Allocate one mamba pool slot, evicting if necessary.""" + slot = self.req_to_token_pool.mamba_pool.alloc(1) + if slot is None: + self.evict(EvictParams(num_tokens=0, mamba_num=1)) + slot = self.req_to_token_pool.mamba_pool.alloc(1) + assert slot is not None, "Can not alloc mamba cache" + return slot + def _match_prefix_helper( self, key: RadixKey ) -> Tuple[List[torch.Tensor], TreeNode, int]: @@ -1043,25 +1051,19 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): else: mamba_branching_seqlen = None - # Copy mamba state to req local space if cow is true + # Defer COW to forward stream: record source index, allocate destination if cow_mamba and last_node.mamba_value is not None: - # for reqs without mamba cache if req.mamba_pool_idx is None: dst_index = self.req_to_token_pool.mamba_pool.alloc(1) - # try to alloc again, protect last_node from eviction if dst_index is None: self.inc_lock_ref(last_node) self.evict(EvictParams(num_tokens=0, mamba_num=1)) dst_index = self.req_to_token_pool.mamba_pool.alloc(1) self.dec_lock_ref(last_node) assert dst_index is not None, "Can not alloc mamba cache" - src_index = last_node.mamba_value - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) req.mamba_pool_idx = dst_index[0] - else: - src_index = last_node.mamba_value - dst_index = req.mamba_pool_idx.unsqueeze(0) - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + req.mamba_cow_src_index = last_node.mamba_value + req.mamba_needs_clear = False value = value[:best_value_len] if value: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index d9bc8b190..cedbf4e90 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -364,20 +364,22 @@ class MambaPool: select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] - # clear at alloc time — expand a scalar GPU zero to the right shape, no CPU-GPU sync + return select_index + + def clear_slots(self, indices: torch.Tensor): + """Zero out mamba state at the given pool indices. Must run on forward stream.""" + need_size = len(indices) for i in range(len(self.mamba_cache.conv)): t = self.mamba_cache.conv[i] z = torch.zeros(1, dtype=t.dtype, device=t.device).expand( t.shape[0], need_size, *t.shape[2:] ) - t[:, select_index] = z + t[:, indices] = z t = self.mamba_cache.temporal z = torch.zeros(1, dtype=t.dtype, device=t.device).expand( t.shape[0], need_size, *t.shape[2:] ) - t[:, select_index] = z - - return select_index + t[:, indices] = z def free(self, free_index: torch.Tensor): if free_index.numel() == 0: @@ -389,22 +391,14 @@ class MambaPool: 1, self.size + 1, dtype=torch.int64, device=self.device ) - def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor): + def copy_from(self, src_indices: torch.Tensor, dst_indices: torch.Tensor): for i in range(len(self.mamba_cache.conv)): - self.mamba_cache.conv[i][:, dst_index] = self.mamba_cache.conv[i][ - :, src_index + self.mamba_cache.conv[i][:, dst_indices] = self.mamba_cache.conv[i][ + :, src_indices ] - self.mamba_cache.temporal[:, dst_index] = self.mamba_cache.temporal[ - :, src_index + self.mamba_cache.temporal[:, dst_indices] = self.mamba_cache.temporal[ + :, src_indices ] - return - - def fork_from(self, src_index: torch.Tensor) -> Optional[torch.Tensor]: - dst_index = self.alloc(1) - if dst_index is None: - return None - self.copy_from(src_index, dst_index) - return dst_index def get_cpu_copy(self, indices): current_platform.synchronize() @@ -548,8 +542,6 @@ class HybridReqToTokenPool(ReqToTokenPool): self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)} self.device = device - # Indexed by req_pool_idx, so size from the req pool buffer - # (self.req_to_token.shape[0]), not from the mamba state pool size. 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 @@ -558,7 +550,7 @@ class HybridReqToTokenPool(ReqToTokenPool): self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = ( torch.zeros( (req_pool_size, self.mamba_ping_pong_track_buffer_size), - dtype=torch.int32, + dtype=torch.int64, device=self.device, ) ) @@ -578,17 +570,16 @@ class HybridReqToTokenPool(ReqToTokenPool): mamba_indices: list[torch.Tensor] = [] mamba_ping_pong_track_buffers: list[torch.Tensor] = [] for req in reqs: - mid = None - if req.mamba_pool_idx is not None: # for radix cache - mid = req.mamba_pool_idx + if req.mamba_pool_idx is not None: # for radix cache / continuing chunked + pass else: mid = self.mamba_pool.alloc(1) assert ( mid is not None ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_pool.available_size()=}, {len(reqs)=}" - mid = mid[0] - req.mamba_pool_idx = mid - mamba_indices.append(mid) + req.mamba_pool_idx = mid[0] + req.mamba_needs_clear = True + mamba_indices.append(req.mamba_pool_idx) if self.enable_mamba_extra_buffer: if req.mamba_ping_pong_track_buffer is None: req.mamba_ping_pong_track_buffer = self.mamba_pool.alloc( @@ -609,9 +600,7 @@ class HybridReqToTokenPool(ReqToTokenPool): mamba_index_tensor = torch.stack(mamba_indices).to(dtype=torch.int32) self.req_index_to_mamba_index_mapping[select_index] = mamba_index_tensor if self.enable_mamba_extra_buffer: - ping_pong_tensor = torch.stack(mamba_ping_pong_track_buffers).to( - dtype=torch.int32 - ) + ping_pong_tensor = torch.stack(mamba_ping_pong_track_buffers) self.req_index_to_mamba_ping_pong_track_buffer_mapping[select_index] = ( ping_pong_tensor ) 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 7142bd9b3..7ecfec0a2 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 @@ -97,15 +97,9 @@ class MambaComponent(TreeComponent): dst_index = self.cache.req_to_token_pool.mamba_pool.alloc(1) self.cache.dec_lock_ref(last_node) assert dst_index is not None, "Can not alloc mamba cache" - self.cache.req_to_token_pool.mamba_pool.copy_from( - mamba_value, dst_index - ) req.mamba_pool_idx = dst_index[0] - else: - dst_index = req.mamba_pool_idx.unsqueeze(0) - self.cache.req_to_token_pool.mamba_pool.copy_from( - mamba_value, dst_index - ) + req.mamba_cow_src_index = mamba_value + req.mamba_needs_clear = False # HiCache: if mamba was evicted from device but has host backup, # ensure host_hit_length >= 1 so load_back is triggered. @@ -253,6 +247,15 @@ class MambaComponent(TreeComponent): self.cache.component_protected_size_[ct] -= vlen cd.lock_ref -= 1 + def _alloc_mamba_slot(self) -> torch.Tensor: + """Allocate one mamba pool slot, evicting if necessary.""" + slot = self.cache.req_to_token_pool.mamba_pool.alloc(1) + if slot is None: + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + slot = self.cache.req_to_token_pool.mamba_pool.alloc(1) + assert slot is not None, "Can not alloc mamba cache" + return slot + def prepare_for_caching_req( self, req: Req, @@ -282,27 +285,24 @@ class MambaComponent(TreeComponent): else: if cache_len is None: return 0 + # Donate the mamba index to the radix cache instead of copying. if self.enable_mamba_extra_buffer: keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_other_idx( req.mamba_next_track_idx ) - mamba_value = ( + mamba_value_donated = ( req.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone() ) - else: - mamba_value = self.cache.req_to_token_pool.get_mamba_indices( + req.mamba_ping_pong_track_buffer[keep_idx] = self._alloc_mamba_slot()[0] + self.cache.req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[ req.req_pool_idx - ).unsqueeze(-1) - mamba_value_forked = self.cache.req_to_token_pool.mamba_pool.fork_from( - mamba_value - ) - if mamba_value_forked is None: - self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) - mamba_value_forked = self.cache.req_to_token_pool.mamba_pool.fork_from( - mamba_value + ] = req.mamba_ping_pong_track_buffer + else: + mamba_value_donated = self._alloc_mamba_slot() + self.cache.req_to_token_pool.mamba_pool.copy_from( + req.mamba_pool_idx.unsqueeze(0), mamba_value_donated ) - assert mamba_value_forked is not None, "Can not alloc mamba cache" - insert_params.mamba_value = mamba_value_forked + insert_params.mamba_value = mamba_value_donated return cache_len def cleanup_after_caching_req( diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index d52414c06..08f858312 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -304,6 +304,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): mamba_track_mask: Optional[torch.Tensor] = None # shape: [b], bool # The seqlens to track mamba state if masked, prefill only. mamba_track_seqlens: Optional[torch.Tensor] = None # shape: [b], int64 + # Deferred mamba init ops: COW pairs and clear indices (performed on forward stream) + mamba_cow_src_indices: Optional[torch.Tensor] = None + mamba_cow_dst_indices: Optional[torch.Tensor] = None + mamba_clear_indices: Optional[torch.Tensor] = None # Optional seq_lens on cpu seq_lens_cpu: Optional[torch.Tensor] = None @@ -509,6 +513,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): mamba_track_indices=batch.mamba_track_indices, mamba_track_mask=batch.mamba_track_mask, mamba_track_seqlens=batch.mamba_track_seqlens, + mamba_cow_src_indices=batch.mamba_cow_src_indices, + mamba_cow_dst_indices=batch.mamba_cow_dst_indices, + mamba_clear_indices=batch.mamba_clear_indices, mm_inputs=batch.multimodal_inputs, encoder_cached=batch.encoder_cached, encoder_lens=batch.encoder_lens, @@ -544,6 +551,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): return_hidden_states_before_norm=return_hidden_states_before_norm, rids=[req.rid for req in batch.reqs], ) + device = model_runner.device if batch.extend_input_logprob_token_ids is not None: diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 68dcc12e5..63dc2068e 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -14,7 +14,10 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.schedule_batch import ( + ScheduleBatch, + set_mamba_track_indices_from_reqs, +) from sglang.srt.managers.utils import get_alloc_len_per_decode from sglang.srt.mem_cache.common import ( alloc_paged_token_slots_extend, @@ -277,22 +280,8 @@ class EagleVerifyInputV2Mixin: device=device, ) - # Set mamba_track_indices for mamba prefix-cache state tracking if get_global_server_args().enable_mamba_extra_buffer(): - mapping = ( - req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping - ) - req_pool_idx_tensor = batch.req_pool_indices.to( - device=mapping.device, dtype=torch.int64 - ) - track_col_idx = torch.tensor( - [req.mamba_next_track_idx for req in batch.reqs], - dtype=torch.int64, - pin_memory=True, - ).to(mapping.device, non_blocking=True) - batch.mamba_track_indices = mapping[ - req_pool_idx_tensor, track_col_idx - ].to(dtype=torch.int64) + set_mamba_track_indices_from_reqs(batch) batch.mamba_track_mask = None batch.mamba_track_seqlens = None