Move DP-attention adapter methods to SchedulerDPAttnAdapter (#25612)

2026-05-18 18:32:16 +08:00
parent 8f37a8a3f3
commit 4d6eec7b32
8 changed files with 272 additions and 289 deletions
@@ -87,7 +87,10 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
    s.running_batch.batch_is_full = False
    s.running_batch.reqs = []
    s.get_new_batch_prefill = MagicMock(return_value=None)
-    s.maybe_prepare_mlp_sync_batch = MagicMock(side_effect=lambda batch, **_: batch)
+    s.dp_attn_adapter = MagicMock()
+    s.dp_attn_adapter.maybe_prepare_mlp_sync_batch = MagicMock(
+        side_effect=lambda batch, **_: batch
+    )
    s._maybe_prepare_ngram_embedding = MagicMock(side_effect=lambda batch: batch)
    s.update_running_batch = MagicMock(side_effect=lambda batch: batch)
    s.tree_cache = tree_cache