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

[sgl] improve mamba_track_indices perf in specdec (#22380)

Browse Source
This commit is contained in:
Tarushii Goel
2026-04-11 00:39:53 +08:00
committed by GitHub
parent f652135d52
commit 2ba94136ce
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_info_v2.py
+14 -6
View File
@@ -259,12 +259,20 @@ class EagleVerifyInputV2Mixin:
# Set mamba_track_indices for mamba prefix-cache state tracking # Set mamba_track_indices for mamba prefix-cache state tracking
if get_global_server_args().enable_mamba_extra_buffer(): if get_global_server_args().enable_mamba_extra_buffer():
batch.mamba_track_indices = torch.stack( mapping = (
[ req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping
req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx] )
for req in batch.reqs req_pool_idx_tensor = batch.req_pool_indices.to(
] device=mapping.device, dtype=torch.int64
).to(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)
batch.mamba_track_mask = None batch.mamba_track_mask = None
batch.mamba_track_seqlens = None batch.mamba_track_seqlens = None