minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Mamba] Fix extra_buffer overlap schedule races (#24954)

Browse Source
This commit is contained in:
Hanming Lu
2026-05-19 12:13:20 +08:00
committed by GitHub
parent f0763859ed
commit 862d39e06c
8 changed files with 176 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+25
View File
@@ -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,
python/sglang/srt/managers/schedule_batch.py
+66 -16
View File
@@ -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]
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+2 -6
View File
@@ -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:
python/sglang/srt/mem_cache/mamba_radix_cache.py
+30 -28
View File
@@ -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:
python/sglang/srt/mem_cache/memory_pool.py
+19 -30
View File
@@ -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
)
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+21 -21
View File
@@ -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(
python/sglang/srt/model_executor/forward_batch_info.py
+8
View File
@@ -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:
python/sglang/srt/speculative/eagle_info_v2.py
+5 -16
View File
@@ -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