diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 0330391a0..c351140ba 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Union import torch @@ -16,6 +17,21 @@ if TYPE_CHECKING: _is_cuda = is_cuda() _is_hip = is_hip() +# Token-buf consume tracking: init to -1, assert non-negative on gather, +# write -1 back. Catches "gather without intermediate stash" bugs. CI enables +# via the existing SGLANG_IS_IN_CI; off in production. +_DEBUG_ASSERT = os.getenv("SGLANG_IS_IN_CI", "").lower() == "true" + + +@torch.compile(dynamic=True) +def _assert_nonneg_and_invalidate( + values: torch.Tensor, buf: torch.Tensor, indices: torch.Tensor +) -> None: + """Fused: assert all `values >= 0` and scatter -1 into `buf[indices]`. + Compiled so the reduction + assert + scatter run as one kernel launch.""" + torch._assert_async((values >= 0).all()) + buf[indices] = -1 + def _resolve_future_token_ids_native(input_ids, future_token_ids_map): input_ids[:] = torch.where( @@ -59,8 +75,12 @@ class FutureMap: self.spec_algo = spec_algo self.req_pool_size = req_to_token_pool.req_to_token.shape[0] - self.output_tokens_buf = torch.empty( - (self.req_pool_size,), dtype=torch.int64, device=self.device + self.output_tokens_buf = ( + torch.full((self.req_pool_size,), -1, dtype=torch.int64, device=self.device) + if _DEBUG_ASSERT + else torch.empty( + (self.req_pool_size,), dtype=torch.int64, device=self.device + ) ) self.new_seq_lens_buf = torch.empty( (self.req_pool_size,), dtype=torch.int64, device=self.device @@ -99,6 +119,10 @@ class FutureMap: # input_ids tokens / spec extras here. if self.spec_algo.is_none(): _resolve_future_token_ids(batch.input_ids, self.output_tokens_buf) + if _DEBUG_ASSERT: + _assert_nonneg_and_invalidate( + batch.input_ids, self.output_tokens_buf, batch.req_pool_indices + ) else: self._resolve_spec_extras(batch) @@ -114,6 +138,10 @@ class FutureMap: draft_input.topk_p = self.topk_p_buf[indices] draft_input.topk_index = self.topk_index_buf[indices] draft_input.bonus_tokens = self.output_tokens_buf[indices] + if _DEBUG_ASSERT: + _assert_nonneg_and_invalidate( + draft_input.bonus_tokens, self.output_tokens_buf, indices + ) if spec_need_hidden_states(): draft_input.hidden_states = self.hidden_states_buf[indices]