From 4cad864361621ca035b36756d60c93ebcb6fc2a0 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Mon, 17 Aug 2026 15:51:45 +0800 Subject: [PATCH] Fix sconv track refresh on graph capture (#35042) --- .../attention/linear/inkling_sconv_backend.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py index 2ef578591..8948a22ca 100644 --- a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py +++ b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py @@ -445,16 +445,21 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend): forward_batch.mamba_track_seqlens = self._graph_track_inert_seqlens[:rows] rows = forward_batch.batch_size query_start_loc = self.sconv_metadata.query_start_loc + # A capture batch is built directly rather than through + # ForwardBatch.init_new, so it carries no prefix lengths; its rows are + # masked off, so zeros keep the indices in bounds. A replayed or eager + # batch always has them, and must still fail rather than track against + # an invented prefix. + prefix_lens = forward_batch.extend_prefix_lens + if prefix_lens is None and on_graph_path: + prefix_lens = torch.zeros_like(forward_batch.mamba_track_seqlens) live = min( rows, forward_batch.mamba_track_seqlens.shape[0], - forward_batch.extend_prefix_lens.shape[0], + prefix_lens.shape[0], ) - lens_to_track = ( - forward_batch.mamba_track_seqlens[:live] - - forward_batch.extend_prefix_lens[:live] - ) + lens_to_track = forward_batch.mamba_track_seqlens[:live] - prefix_lens[:live] chunk_aligned = ( lens_to_track // self.mamba_cache_chunk_size ) * self.mamba_cache_chunk_size