[Mamba] Fix spec-v2 + extra_buffer crash (guard None mamba_next_track_idx) (#27998)

This commit is contained in:
Serge Panev
2026-07-15 18:23:14 -07:00
committed by GitHub
parent 7647a9d260
commit 1c4892d7bb
+7 -1
View File
@@ -1654,9 +1654,15 @@ def set_mamba_track_indices_from_reqs(batch):
all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[ all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[
batch.req_pool_indices batch.req_pool_indices
] # (bs, ping_pong_size), int64, on device ] # (bs, ping_pong_size), int64, on device
# Guard: mamba_next_track_idx may be None for requests that haven't
# gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
# Default to 0 (first ping-pong slot) to avoid TypeError.
idx = ( idx = (
torch.tensor( torch.tensor(
[req.mamba_next_track_idx for req in batch.reqs], [
req.mamba_next_track_idx if req.mamba_next_track_idx is not None else 0
for req in batch.reqs
],
dtype=torch.int64, dtype=torch.int64,
pin_memory=True, pin_memory=True,
) )