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 4d03f9cdb..a4d0677d3 100644 --- a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py +++ b/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( diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 0aa5437e6..d87e54353 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 ( - self.model_runner.server_args.enable_mamba_extra_buffer() - and not self.model_runner.server_args.disable_radix_cache - and self.model_runner.spec_algorithm.is_none() + return self.model_runner.server_args.enable_mamba_extra_buffer() and ( + not self.model_runner.server_args.disable_radix_cache ) def _cache_loc_dtype(self): diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 97c7702c7..23db4547b 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/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 ( - self.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) - if ( - self.mamba_track_mask is not None - and forward_batch.mamba_track_mask is not None - ): - self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask) - if ( - self.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) + if self.mamba_track_indices is not None: + if 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_() + if self.mamba_track_mask is not None: + if forward_batch.mamba_track_mask is not None: + self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask) + else: + self.mamba_track_mask[:bs].zero_() + self.mamba_track_mask[bs:].zero_() + if self.mamba_track_seqlens is not None: + if 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_() if forward_batch.mrope_positions is not None: self.mrope_positions[:, :raw_num_tokens].copy_(