tiny

Author: yizhang2077

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -515,7 +515,7 @@ def cache_unfinished_req(self, req: Req, chunked=False) -> None:
         )
 
         if not mamba_exist:
-            assert new_last_node.mamba_value == mamba_value_forked
+            assert torch.equal(new_last_node.mamba_value, mamba_value_forked)
         assert len(req.prefix_indices) <= len(
             new_indices
         ), f"{req.prefix_indices=}, {new_indices=}"

python/sglang/srt/mem_cache/memory_pool.py

@@ -216,7 +216,7 @@ def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor):
         self.mamba_cache[1][:, dst_index] = self.mamba_cache[1][:, src_index]
         return
 
-    def fork_from(self, src_index: torch.Tensor) -> Optional[int]:
+    def fork_from(self, src_index: torch.Tensor) -> Optional[torch.Tensor]:
         dst_index = self.alloc(1)
         if dst_index == None:
             return None