Add Inkling model support (#31681)

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Qiaolin Yu <qiaolin.yu@radixark.ai>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Aurick Qiao <aurick@thinkingmachines.ai>
Co-authored-by: Joseph <jk@thinkingmachines.ai>

python/sglang/srt/managers/detokenizer_manager.py

@@ -124,6 +124,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
     def init_tokenizer(self, server_args: ServerArgs):
         if server_args.skip_tokenizer_init:
             self.tokenizer = None
+            self.vocab_size = None
         else:
             self.tokenizer = get_tokenizer(
                 server_args.tokenizer_path,
@@ -132,6 +133,10 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
                 revision=server_args.revision,
                 tokenizer_backend=server_args.tokenizer_backend,
             )
+            try:
+                self.vocab_size = len(self.tokenizer)
+            except TypeError:
+                self.vocab_size = getattr(self.tokenizer, "vocab_size", None)
 
     def init_running_status(self, server_args: ServerArgs):
         self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES)
@@ -204,6 +209,20 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
         # If it is embedding model, no detokenization is needed.
         return recv_obj
 
+    @staticmethod
+    def _clamp_decode_ids(ids: List[int], vocab_size: Optional[int]) -> List[int]:
+        """Map out-of-range token ids to 0 so the tokenizer can decode them.
+
+        Multimodal placeholder ids (e.g. Inkling's negative -101/-102, or radix-cache
+        pad-value hashes) are not real vocab tokens; tiktoken-style backends raise
+        OverflowError on negative / out-of-range ids. These only appear in the
+        surrogate-context prefix (before read_offset) and carry no text, and the clamp
+        is applied identically to surr_ids and read_ids, so the incremental
+        (read-minus-surr) output text is unchanged.
+        """
+        hi = vocab_size if vocab_size else None
+        return [t if (0 <= t and (hi is None or t < hi)) else 0 for t in ids]
+
     def _grouped_batch_decode(
         self,
         ids_list: List[List[int]],
@@ -270,6 +289,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
 
     def _decode_batch_token_id_output(self, recv_obj: BatchTokenIDOutput):
         bs = len(recv_obj.rids)
+        vocab_size = self.vocab_size
 
         # Initialize decode status
         read_ids, surr_ids = [], []
@@ -278,14 +298,18 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
             if rid not in self.decode_status:
                 s = DecodeStatus(
                     decoded_text=recv_obj.decoded_texts[i],
-                    decode_ids=list(recv_obj.decode_ids[i]),
+                    decode_ids=self._clamp_decode_ids(
+                        recv_obj.decode_ids[i], vocab_size
+                    ),
                     surr_offset=0,
                     read_offset=recv_obj.read_offsets[i],
                 )
                 self.decode_status[rid] = s
             else:
                 s = self.decode_status[rid]
-                s.decode_ids.extend(recv_obj.decode_ids[i])
+                s.decode_ids.extend(
+                    self._clamp_decode_ids(recv_obj.decode_ids[i], vocab_size)
+                )
 
             read_ids.append(
                 self.trim_matched_stop(

python/sglang/srt/managers/io_struct.py

@@ -2041,6 +2041,7 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True):
     added_tokens_config: Optional[Dict[str, int]] = None
     lora_id: Optional[str] = None
     load_format: Optional[str] = None
+    expected_checksums: Optional[Dict[str, str]] = None
 
     def to_ref(self) -> LoRARef:
         return LoRARef(

python/sglang/srt/managers/schedule_policy.py

@@ -636,10 +636,33 @@ class PrefillAdder:
             alloc = min(extend_input_len, self.rem_chunk_tokens)
         else:
             alloc = extend_input_len
-        budget = max(alloc, self.tree_cache.sliding_window_size) + self.page_size
+        window = self.tree_cache.sliding_window_size
+        return max(alloc - window, 0) + self._swa_reserved_tokens(swa_host_hit_length)
+
+    def _swa_reserved_tokens(self, swa_host_hit_length: int = 0) -> int:
+        """SWA tokens a request needs regardless of extend length: the sliding
+        window (decode headroom) + allocator page slack + the load-back window
+        charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap."""
+        reserved = self.tree_cache.sliding_window_size + self.page_size
         if swa_host_hit_length > 0:
-            budget += self.ceil_paged_tokens(swa_host_hit_length)
-        return budget
+            reserved += self.ceil_paged_tokens(swa_host_hit_length)
+        return reserved
+
+    def _swa_chunk_cap(self, swa_host_hit_length: int = 0) -> int:
+        """Largest page-aligned extend chunk the SWA pool can admit right now,
+        keeping a sliding window of headroom below rem_swa_tokens; 0 if not
+        even one page fits. Only valid when is_hybrid_swa is True.
+
+        Escape hatch for a request whose budget can never pass the
+        _swa_budget_for_req gate (extend near/above the pool size, or a large
+        load-back charge): without shrinking its chunk it would be rejected
+        forever (head-of-line livelock). Shrinking is sound because past a
+        chunk boundary only the sliding window stays locked — the rest turns
+        evictable — so each pass's transient footprint fits the pool."""
+        cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(swa_host_hit_length)
+        if cap <= 0:
+            return 0
+        return cap // self.page_size * self.page_size
 
     def _mamba_gap_budget_for_req(self, req: Req) -> int:
         """Shared-gap reservation (full-token-equivalents) for a request's new
@@ -1023,12 +1046,19 @@ class PrefillAdder:
         if total_tokens >= self.rem_total_tokens:
             return AddReqResult.NO_TOKEN
 
+        chunk_tokens_limit = self.rem_chunk_tokens
         if self.is_hybrid_swa:
+            # host-hit prefix is loaded back, not re-prefilled, so the SWA peak is
+            # driven only by the freshly-prefilled tail (the loaded window is
+            # charged separately via swa_host_hit_length).
             swa_needed = self._swa_budget_for_req(
-                cand_extend_input_len, swa_host_hit_length=req.swa_host_hit_length
+                real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
             )
             if swa_needed >= self.rem_swa_tokens:
-                return AddReqResult.NO_TOKEN
+                swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
+                if self.rem_chunk_tokens is None or swa_cap <= 0:
+                    return AddReqResult.NO_TOKEN
+                chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
 
         if (
             self.rem_chunk_tokens is None
@@ -1046,11 +1076,15 @@ class PrefillAdder:
                 return AddReqResult.NO_TOKEN
 
             if self.is_hybrid_swa:
+                # self.rem_swa_tokens may decrease after the lock acquisition
                 swa_needed = self._swa_budget_for_req(
-                    cand_extend_input_len, swa_host_hit_length=req.swa_host_hit_length
+                    real_input_tokens, swa_host_hit_length=req.swa_host_hit_length
                 )
                 if swa_needed >= self.rem_swa_tokens:
-                    return AddReqResult.NO_TOKEN
+                    swa_cap = self._swa_chunk_cap(req.swa_host_hit_length)
+                    if self.rem_chunk_tokens is None or swa_cap <= 0:
+                        return AddReqResult.NO_TOKEN
+                    chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap)
 
             if req.needs_host_load_back():
                 new_indices, req.last_node = self.tree_cache.init_load_back(
@@ -1088,7 +1122,7 @@ class PrefillAdder:
 
                 self._add_dllm_req(req, prefix_len)
                 self._req_inc_lock_ref(req)
-            elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
+            elif chunk_tokens_limit is None or input_tokens <= chunk_tokens_limit:
                 # Non-chunked prefill — the whole sequence is committed this iter.
                 req.set_extend_range(
                     len(req.prefix_indices), len(req.full_untruncated_fill_ids)
@@ -1108,7 +1142,7 @@ class PrefillAdder:
                 )
             else:
                 # Make sure at least one page is available
-                trunc_len = self.rem_chunk_tokens // self.page_size * self.page_size
+                trunc_len = chunk_tokens_limit // self.page_size * self.page_size
 
                 if trunc_len <= 0:
                     return AddReqResult.OTHER

python/sglang/srt/managers/scheduler.py

@@ -1481,6 +1481,18 @@ class Scheduler(
         self.schedule_stream = self.device_module.Stream(priority=0)
         if self.device == "cpu":
             self.schedule_stream.synchronize = lambda: None  # No-op for CPU
+        elif is_cuda() or _is_hip:
+            # CUDA/HIP streams come from a fixed round-robin pool. Redraw if this
+            # stream aliases forward_stream, which would eliminate scheduler
+            # overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle;
+            # other accelerators (e.g. NPU/XPU) skip the alias check.
+            _redraws = 0
+            while (
+                self.schedule_stream.cuda_stream == self.forward_stream.cuda_stream
+                and _redraws < 64
+            ):
+                self.schedule_stream = self.device_module.Stream(priority=0)
+                _redraws += 1
         # The global WAR barrier fences the scheduler's next shared-buffer write
         # on the previous forward's read of the unified memory pool.
         self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -160,11 +160,14 @@ class SchedulerInvariantChecker:
             self.req_to_token_pool.mamba_pool.size,
         )
         if leak:
-            # Page-level leak diagnosis for mamba
-            free_full_pages = set(
-                self.token_to_kv_pool_allocator.free_pages.tolist()
-                + self.token_to_kv_pool_allocator.release_pages.tolist()
-            )
+            # Page-level leak diagnosis for mamba. Allocator flavors without
+            # page free-lists (free_pages is None) skip the page census — the
+            # dump must never crash the watchdog thread that calls it.
+            free_pages = self.token_to_kv_pool_allocator.free_pages
+            release_pages = self.token_to_kv_pool_allocator.release_pages
+            if free_pages is None or release_pages is None:
+                return leak, msg
+            free_full_pages = set(free_pages.tolist() + release_pages.tolist())
             cached_full_pages = set(self.tree_cache.all_values_flatten().tolist())
             expected_full_pages = set(
                 range(1, self.token_to_kv_pool_allocator.size + 1)

python/sglang/srt/managers/tp_worker.py

@@ -221,6 +221,37 @@ class BaseTpWorker(ABC):
             tensors = dict(bucket.reconstruct_tensors())
         else:
             tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors)
+        if recv_req.expected_checksums is not None:
+            import hashlib
+
+            exp = recv_req.expected_checksums
+            mismatch, missing = [], []
+            for name, want in exp.items():
+                if name not in tensors:
+                    missing.append(name)
+                    continue
+                got = hashlib.sha256(
+                    tensors[name]
+                    .detach()
+                    .cpu()
+                    .contiguous()
+                    .flatten()
+                    .view(torch.uint8)
+                    .numpy()
+                    .tobytes()
+                ).hexdigest()
+                if got != want:
+                    mismatch.append(name)
+            extra = [n for n in tensors if n not in exp]
+            if mismatch or missing or extra:
+                raise RuntimeError(
+                    f"[LORA-CHECK] rank{self.tp_rank} adapter sync MISMATCH of {len(exp)} expected: "
+                    f"{len(mismatch)} value-diff {mismatch[:5]}, {len(missing)} missing {missing[:5]}, "
+                    f"{len(extra)} extra {extra[:5]}"
+                )
+            logger.info(
+                f"[LORA-CHECK] rank{self.tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)"
+            )
         result = self.model_runner.load_lora_adapter_from_tensors(
             recv_req.to_ref(),
             tensors,