prain: draft rows from the batch layout

prain/drafts.py

@@ -0,0 +1,34 @@
+"""Which token rows of the current forward belong to draft requests. Stdlib only.
+
+A request is a draft when its body carries "vllm_xargs": {"prain": "draft"}; vLLM hands
+that to the worker as sampling_params.extra_args. Every step, vLLM's model runner (V2, the
+default in 0.32) lays the tokens of its requests out back to back in batch order, and
+prain.experts wraps its prepare_inputs to call note_step() with that layout. The MoE
+layers of the forward that follows read spans().
+"""
+
+_DRAFTS = set()  # request ids of unfinished drafts
+_SPANS = [None]  # [(first_row, end_row), ...] of the current forward's draft tokens, or None
+
+
+def note_step(scheduler_output, req_ids, query_start_loc):
+    """scheduler_output: vLLM's SchedulerOutput of this step; req_ids: request id per batch
+    position; query_start_loc: first token row per batch position, plus the total at the end."""
+    for new in scheduler_output.scheduled_new_reqs:
+        args = (new.sampling_params.extra_args or {}) if new.sampling_params else {}
+        if args.get("prain") == "draft":
+            _DRAFTS.add(new.req_id)
+    _DRAFTS.difference_update(scheduler_output.finished_req_ids)
+    spans = [(int(query_start_loc[i]), int(query_start_loc[i + 1]))
+             for i, req_id in enumerate(req_ids) if req_id in _DRAFTS]
+    _SPANS[0] = spans or None
+
+
+def spans():
+    return _SPANS[0]
+
+
+def clear():
+    """Called after the last MoE layer: forwards without a note_step() (vLLM's warm-up and
+    profiling runs) must not see the previous step's layout."""
+    _SPANS[0] = None

tests/test_drafts.py

@@ -0,0 +1,42 @@
+import unittest
+from types import SimpleNamespace
+
+from prain import drafts
+from tests.helpers import TmpCase  # noqa: F401  (every test module imports the helpers)
+
+
+def new_request(req_id, xargs):
+    return SimpleNamespace(req_id=req_id, sampling_params=SimpleNamespace(extra_args=xargs))
+
+
+def step(new=(), finished=()):
+    return SimpleNamespace(scheduled_new_reqs=list(new), finished_req_ids=set(finished))
+
+
+class Drafts(unittest.TestCase):
+    def setUp(self):
+        drafts._DRAFTS.clear()
+        drafts.clear()
+
+    def test_draft_rows_follow_the_batch_layout(self):
+        # a: 3 prompt tokens (full); b: 2 prompt tokens (draft)
+        drafts.note_step(step([new_request("a", None), new_request("b", {"prain": "draft"})]),
+                         ["a", "b"], [0, 3, 5])
+        self.assertEqual(drafts.spans(), [(3, 5)])
+        # next step: both decode one token each, in a new order
+        drafts.note_step(step(), ["b", "a"], [0, 1, 2])
+        self.assertEqual(drafts.spans(), [(0, 1)])
+        drafts.clear()
+        self.assertIsNone(drafts.spans())
+
+    def test_finished_drafts_are_forgotten(self):
+        drafts.note_step(step([new_request("b", {"prain": "draft"})]), ["b"], [0, 4])
+        drafts.note_step(step(finished=["b"]), [], [0])
+        self.assertIsNone(drafts.spans())
+        self.assertEqual(drafts._DRAFTS, set())
+
+    def test_requests_without_the_flag_are_full(self):
+        drafts.note_step(step([new_request("a", {"other": 1}),
+                               SimpleNamespace(req_id="c", sampling_params=None)]),
+                         ["a", "c"], [0, 1, 2])
+        self.assertIsNone(drafts.spans())