minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[MLX] Fix FutureMap relay unit test to use RelayPayload (#29311)

Browse Source
This commit is contained in:
Lijuan Tang
2026-06-25 16:53:58 -07:00
committed by GitHub
parent 38e857f423
commit 623300a589
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+4 -1
View File
@@ -384,6 +384,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
def test_mlx_scheduler_init_overlap_keeps_future_map_relay(self): def test_mlx_scheduler_init_overlap_keeps_future_map_relay(self):
from sglang.srt.managers import scheduler as scheduler_module from sglang.srt.managers import scheduler as scheduler_module
from sglang.srt.managers.overlap_utils import RelayPayload
from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -417,7 +418,9 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertIsNotNone(scheduler.future_map) self.assertIsNotNone(scheduler.future_map)
indices = torch.tensor([1], dtype=torch.int64) indices = torch.tensor([1], dtype=torch.int64)
scheduler.future_map.stash(indices, torch.tensor([7], dtype=torch.int64)) scheduler.future_map.stash(
indices, RelayPayload(bonus_tokens=torch.tensor([7], dtype=torch.int64))
)
self.assertEqual(int(scheduler.future_map.output_tokens_buf[1].item()), 7) self.assertEqual(int(scheduler.future_map.output_tokens_buf[1].item()), 7)
def test_decode_finalize_does_not_snapshot_auxiliary_state(self): def test_decode_finalize_does_not_snapshot_auxiliary_state(self):