From 2ba94136ce21abaa069b93941fafc9beb842d290 Mon Sep 17 00:00:00 2001 From: Tarushii Goel Date: Fri, 10 Apr 2026 09:39:53 -0700 Subject: [PATCH] [sgl] improve mamba_track_indices perf in specdec (#22380) --- .../sglang/srt/speculative/eagle_info_v2.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 1be28c38f..cec0857d6 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -259,12 +259,20 @@ class EagleVerifyInputV2Mixin: # Set mamba_track_indices for mamba prefix-cache state tracking if get_global_server_args().enable_mamba_extra_buffer(): - batch.mamba_track_indices = torch.stack( - [ - req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx] - for req in batch.reqs - ] - ).to(torch.int64) + 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) batch.mamba_track_mask = None batch.mamba_track_seqlens = None