[srt] Fix sconv state memory corruption on specdec (#34043)

python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py

@@ -140,6 +140,17 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
         self._graph_track_conv_indices = torch.zeros(
             (max_bs, self.conv_state_len), dtype=torch.int64, device=dev
         )
+        # Inert track fields for graph capture: a capture warmup batch that
+        # carries no tracking metadata must still LAUNCH the track scatter
+        # (all rows masked off), or the python-level `if` specializes the
+        # scatter out of the captured graph.
+        self._graph_track_inert_mask = torch.zeros(max_bs, dtype=torch.bool, device=dev)
+        self._graph_track_inert_indices = torch.zeros(
+            max_bs, dtype=torch.int64, device=dev
+        )
+        self._graph_track_inert_seqlens = torch.zeros(
+            max_bs, dtype=torch.int64, device=dev
+        )
         # Same address-stability requirement; the base only sizes this from
         # init_cuda_graph_state, which the prefill graph never calls.
         self._alloc_cache_indices_buf(max_bs)
@@ -419,7 +430,15 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
         every row it may read must index inside *this* replay's token buffer.
         """
         if forward_batch.mamba_track_mask is None:
-            return
+            if not on_graph_path:
+                self.sconv_metadata.track_conv_indices = None
+                return
+            # Graph capture must still launch the track scatter (see the inert
+            # buffers in __init__).
+            rows = forward_batch.batch_size
+            forward_batch.mamba_track_mask = self._graph_track_inert_mask[:rows]
+            forward_batch.mamba_track_indices = self._graph_track_inert_indices[:rows]
+            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
         live = min(

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -537,10 +537,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self.raw_bs = 0
 
     def _is_mamba_track_enabled(self) -> bool:
-        return (
+        return self.model_runner.server_args.enable_mamba_extra_buffer() and (
-            self.model_runner.server_args.enable_mamba_extra_buffer()
+            not self.model_runner.server_args.disable_radix_cache
-            and not self.model_runner.server_args.disable_radix_cache
-            and self.model_runner.spec_algorithm.is_none()
         )
 
     def _cache_loc_dtype(self):

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -421,21 +421,20 @@ class PrefillInputBuffers(ForwardInputBuffers):
         self.positions[:raw_num_tokens].copy_(forward_batch.positions)
         self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)
 
-        if (
+        if self.mamba_track_indices is not None:
-            self.mamba_track_indices is not None
+            if forward_batch.mamba_track_indices is not None:
-            and forward_batch.mamba_track_indices is not None
+                self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
-        ):
+            self.mamba_track_indices[bs:].zero_()
-            self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
+        if self.mamba_track_mask is not None:
-        if (
+            if forward_batch.mamba_track_mask is not None:
-            self.mamba_track_mask is not None
+                self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
-            and forward_batch.mamba_track_mask is not None
+            else:
-        ):
+                self.mamba_track_mask[:bs].zero_()
-            self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
+            self.mamba_track_mask[bs:].zero_()
-        if (
+        if self.mamba_track_seqlens is not None:
-            self.mamba_track_seqlens is not None
+            if forward_batch.mamba_track_seqlens is not None:
-            and forward_batch.mamba_track_seqlens is not None
+                self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
-        ):
+            self.mamba_track_seqlens[bs:].zero_()
-            self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
 
         if forward_batch.mrope_positions is not None:
             self.mrope_positions[:, :raw_num_tokens].copy_(