diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index c690bd8368..892875c60a 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -236,6 +236,13 @@ Read by the ``mstar-server`` binary and its bridge * - ``MSTAR_MAX_BODY_MB`` - ``128`` - Request body limit (multipart uploads included). + * - ``MSTAR_TOKENIZE_PARALLEL`` + - ``1`` + - When a long prompt is encoded as parallel chunks, set + ``TOKENIZERS_PARALLELISM=true`` for the process first: the Rust + tokenizer only uses its thread pool while that variable reads true, + and serving environments often export it false. ``0`` leaves the + variable as it is (the chunked encode then runs serially). * - ``MSTAR_TOKENIZER`` - unset - Path to a HuggingFace ``tokenizer.json`` enabling frontend @@ -315,6 +322,32 @@ Worker scheduling - ``1024`` - Cap on back-to-back speculative steps before the leader yields to other ready work. + * - ``MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS`` + - ``32,64,128,256,512,1024,2048,4096,8192`` + - Total tokens of a packed Qwen3.5 text-prefill step that get a captured + graph (one per bucket and batch size). A prompt past the top bucket + prefills eagerly, which is launch-bound: an eager 8k prefill on 0.8B + spends 78 ms of wall time on 36 ms of kernels. The default now reaches + 8192; narrow it to save graph memory and capture time, or extend it. + * - ``MSTAR_DEVICE_LOOPBACK`` + - ``1`` + - Whether a decode node may keep its loop-back token on the device: + the sampler writes each request's last token to a slot-addressed + master and the next step reads its input ids from it, so the worker + routes the signal with no tensor (no per-request store, hold or + cleanup per step). Only nodes that opt in (Qwen3.5 decode) and only + when the sampler has graph buffers; ``0`` keeps every node on the + per-request tensor path. + * - ``MSTAR_SAMPLER_INGRAPH_SCATTER`` + - ``0`` + - The captured sampler gathers its RNG offsets from the slot masters + and scatters the advanced offsets and the sampled tokens back inside + the graph, so a replayed decode step issues none of those launches on + the gpu thread afterwards (each one is a GIL release the main thread's + postprocess wins). Padding rows then address a trash master row and + the masters are allocated once for ``MSTAR_SAMPLER_SLOTS`` concurrent + requests (default 4096, a hard cap with this on). Only with the device + loop-back on. An experiment knob; measure before turning it on. * - ``MSTAR_RUST_SEND_TIMING`` - unset - ``1`` makes the Rust runtime print, every 500 calls, how long the diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index f919336ea6..4802a9140d 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -25,6 +25,7 @@ ResultTensors, ) from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator +from mstar.communication.event import EventWakeup from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import chain @@ -138,6 +139,8 @@ def __init__( self.output_queue = queue.Queue() self.profile_queue = queue.Queue() self.stop_event = threading.Event() + # Wakes the thread's idle wait: every producer below signals it. + self.wakeup = EventWakeup() self.per_request_reading_tensors = {} self.output_loop_idxs: dict[str, NameToLoopIndices] = {} @@ -174,6 +177,7 @@ def __init__( reads_done_queue=self.reads_done_queue, discard_tensor_queue=self.discard_tensor_queue, stop_event=self.stop_event, + wakeup=self.wakeup, communicator=self.communicator, tensor_manager=self.tensor_manager, model=model, @@ -184,10 +188,18 @@ def __init__( ) self.thread.start() + def _signal(self): + """End the thread's idle wait: new work is queued. getattr: tests + build this facade without its constructor.""" + wakeup = getattr(self, "wakeup", None) + if wakeup is not None: + wakeup.signal() + def new_request(self, input: PreprocessInput): self.output_loop_idxs[input.request_id] = {} self.per_request_reading_tensors[input.request_id] = 0 self.request_input_queue.put(input) + self._signal() def abort_request(self, request_id: str): # Forward the abort to the conductor and begin draining our own reads @@ -196,6 +208,7 @@ def abort_request(self, request_id: str): # REMOVE_REQUEST, so a worker still reading them can't be unlinked # out from under it. Drop only the main-thread bookkeeping here. self.abort_request_queue.put(request_id) + self._signal() self.output_loop_idxs.pop(request_id, None) self.per_request_reading_tensors.pop(request_id, None) @@ -207,6 +220,7 @@ def finished_reading(self, request_id: str, drained: bool = True): when delivery was abandoned (TTL, client gone) and the ACK must wait for any in-flight read.""" self.reads_done_queue.put((request_id, drained)) + self._signal() self.output_loop_idxs.pop(request_id, None) self.per_request_reading_tensors.pop(request_id, None) @@ -229,13 +243,17 @@ def new_result_tensors(self, input: ResultTensors): input.request_id, self.per_request_reading_tensors[input.request_id] ) self.result_tensor_input_queue.put(input) + self._signal() def new_result_token( self, request_id: str, value: int, loop_indices: NestedLoopIndices | None, signal: str, modality: str, + wake: bool = True, ): """A value that rode inline: same accounting as a result tensor, no - transport read and nothing to ack.""" + transport read and nothing to ack. ``wake=False`` queues without + ending the thread's idle wait: a frame's caller hands over every + request's value first and calls ``wake`` once.""" if request_id not in self.output_loop_idxs: logger.debug("Late token for cleaned-up request %s dropped", request_id) return @@ -247,6 +265,12 @@ def new_result_token( self.result_tensor_input_queue.put( InlineToken(request_id, value, loop_indices, modality) ) + if wake: + self._signal() + + def wake(self): + """End the thread's idle wait after a batch of ``wake=False`` puts.""" + self._signal() def discard_result_tensors(self, input: ResultTensors): """Ack and drop result tensors for an already-removed request. @@ -255,6 +279,7 @@ def discard_result_tensors(self, input: ResultTensors): producing worker gets its TENSOR_RECEIVED ack and frees the buffers. """ self.discard_tensor_queue.put(input) + self._signal() def has_pending_tensors(self, request_id: str): return self.per_request_reading_tensors.get(request_id, 0) > 0 @@ -324,11 +349,13 @@ def get_rx_info(self, request_id: str) -> list[RxInfo]: def cleanup_request(self, request_id: str): self.cleanup_request_queue.put(request_id) + self._signal() self.output_loop_idxs.pop(request_id, None) self.per_request_reading_tensors.pop(request_id, None) def shutdown(self): self.stop_event.set() + self._signal() if self.thread.is_alive(): self.thread.join() @@ -384,6 +411,7 @@ def __init__( enable_prof: bool=False, model_config: dict | None = None, enable_nvtx: bool=False, + wakeup: EventWakeup | None = None, ): # keying a stream needs the deployment's page size as well as the # model's declaration, so a worker built without a config keys no stream @@ -428,6 +456,12 @@ def __init__( # Owned by PreprocessWorker (main thread); used only from this thread. self.communicator = communicator self.tensor_manager = tensor_manager + # Set by the producers (queue puts, shutdown) and by finished tensor + # reads; run() blocks on the communicator's poll with it instead of + # sleeping and sweeping the queues. None keeps the sleep (tests built + # with a bare communicator). + self.wakeup = wakeup + self.idle_wait_ms = 20 def _cleanup_request_state(self, request_id: str, *, force: bool = False) -> None: """Release transport and postprocessing state owned by this thread. @@ -677,10 +711,14 @@ def _read_result_tensor( self, result: ResultTensors ): result.graph_edge.name = f"{result.modality}_output" - self.tensor_manager.start_read_tensors( + futures = self.tensor_manager.start_read_tensors( request_id=result.request_id, graph_edges=[result.graph_edge], ) + wakeup = getattr(self, "wakeup", None) + if wakeup is not None and futures: + # a finished read ends the idle wait, so the chunk goes out at once + wakeup.register_futures(futures) if result.request_id not in self.tensor_uuid_to_metadata_per_request: self.tensor_uuid_to_metadata_per_request[result.request_id] = {} state = self.request_output_state.setdefault( @@ -930,11 +968,28 @@ def _hard_cleanup(self, request_id: str) -> None: self.tensor_manager.force_cleanup_request(request_id) self._drop_request_state(request_id) + def _idle_wait(self): + """How a pass that found nothing waits: on the communicator's poll, + ended by a message, a producer's signal or a finished read, with + ``idle_wait_ms`` as the backstop; or a 1 ms sleep when the thread has + no wakeup or its communicator cannot poll one.""" + # getattr: teardown tests build a thread without its constructor + wakeup = getattr(self, "wakeup", None) + comm = getattr(self, "communicator", None) + register = getattr(comm, "register_event_for_poll", None) + wait = getattr(comm, "wait_for_work", None) + if wakeup is None or not callable(register) or not callable(wait): + return lambda: time.sleep(0.001) + register(wakeup) + idle_ms = getattr(self, "idle_wait_ms", 20) + return lambda: wait(idle_ms) + def run(self): # A thread told to stop before it starts (teardown tests build one with # only its queues) has nothing to warm up. if not self.stop_event.is_set(): warm_up_model(getattr(self, "model", None)) + idle_wait = self._idle_wait() while not self.stop_event.is_set(): did_work = False try: @@ -1018,7 +1073,7 @@ def run(self): logger.exception("PreprocessWorkerThread error") if not did_work: - time.sleep(0.001) + idle_wait() # Stopping, and nothing will send the RemoveRequest for what is still # tracked (the conductor is gone or going), so drop it here rather than diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 6c9e28e361..8ab6f3c866 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -34,6 +34,7 @@ from mstar.profile.format import OutputInfo, RequestProfile, RequestTiming from mstar.utils import profiler from mstar.utils.exitcode import describe_exitcode +from mstar.utils.gc_freeze import freeze_after_setup from mstar.utils.logging_config import quiet_noisy_loggers from mstar.utils.orphan import watch_parent @@ -366,6 +367,9 @@ def finalize_setup(self) -> None: and message.message_type == "setup_done" ): logger.info("All workers ready") + # the set-up heap is complete: keep gen-2 collections off + # it (measured 200 ms per pass on this process under load) + freeze_after_setup("API server") self._msg_thread.start() return logger.warning( @@ -578,18 +582,24 @@ def _process_messages(self) -> None: # One frame for the step: every request's value. # A request that is gone needs nothing acked back. body = message.body + queued = False with self.request_lock: for i, rid in enumerate(body.request_ids): if rid in self.pending_requests: + # one wake-up for the frame, after every + # request's value is queued self.preprocess_worker.new_result_token( rid, body.values[i], body.loop_indices[i], - body.signal, body.modality, + body.signal, body.modality, wake=False, ) + queued = True elif rid not in self.recently_completed: logger.warning( "Token for unknown request %s dropped", rid, ) + if queued: + self.preprocess_worker.wake() continue rid = message.body.request_id @@ -733,14 +743,31 @@ def _process_messages(self) -> None: # ---------------------------------------------------------- async def iter_result_chunks(self, request_id: str): - """Yield raw :class:`ResultChunk` objects as they arrive. + """Yield raw :class:`ResultChunk` objects as they arrive, one at a time. Shared source for both output surfaces: ``/generate`` formats each chunk as NDJSON (via :meth:`async_stream_results`), while the - OpenAI-compatible endpoints translate the same chunks into SSE. The - per-request timeout, incremental drain, and final flush behave exactly - as before — only the yielded type changed (``ResultChunk`` instead of a - pre-serialized line). + OpenAI-compatible endpoints translate the same chunks into SSE. This + is a flattening wrapper over :meth:`iter_result_chunk_batches`, which + hands out the chunks the way they arrive (several per wake under load) + and owns the timeout, the drain and the abort on an early exit. + """ + batches = self.iter_result_chunk_batches(request_id) + try: + async for batch in batches: + for chunk in batch: + yield chunk + finally: + await batches.aclose() + + async def iter_result_chunk_batches(self, request_id: str): + """Yield the :class:`ResultChunk` objects as they arrive, as one list + per drain: everything the message thread appended since the last + wake. Under load several tokens of a request are often waiting when + the handler runs, and sending them as one body frame costs one write + instead of one per token; when the loop keeps up the lists are of + length one and nothing is delayed. The final flush and an in-band + error come as the last list. """ start = time.time() finished = False @@ -773,10 +800,11 @@ async def iter_result_chunks(self, request_id: str): else: done = True - for chunk in new_chunks: + if new_chunks: if self.enable_nvtx: - profiler.mark("apiserver.chunk_available") - yield chunk + for _ in new_chunks: + profiler.mark("apiserver.chunk_available") + yield new_chunks if done: logger.info("Async stream results received finish for %s", request_id) @@ -792,19 +820,20 @@ async def iter_result_chunks(self, request_id: str): # lock; the popped request is no longer shared with other threads. if finished_req is not None: self._finalize_profile(finished_req) - for chunk in remaining: - yield chunk + tail = list(remaining) # A request can fail after the stream is already open # (preprocess error, result-delivery timeout); the HTTP # status is committed by then, so the error must travel # in-band as the final chunk. if finished_req is not None and finished_req.error is not None: - yield ResultChunk( + tail.append(ResultChunk( request_id=request_id, modality="error", data=str(finished_req.error).encode("utf-8"), metadata={"status": finished_req.error_status}, - ) + )) + if tail: + yield tail finished = True break diff --git a/mstar/api_server/openai/_util.py b/mstar/api_server/openai/_util.py index 7ae0f7f240..3481d1b1eb 100644 --- a/mstar/api_server/openai/_util.py +++ b/mstar/api_server/openai/_util.py @@ -18,7 +18,18 @@ def rid(prefix: str) -> str: def sse(obj: dict) -> str: - return f"data: {json.dumps(obj)}\n\n" + return f"data: {_dumps(obj)}\n\n" + + +try: + import orjson + + def _dumps(obj: dict) -> str: + # ~5x json.dumps; one event per token at tens of thousands a second + return orjson.dumps(obj).decode() +except ImportError: # pragma: no cover + def _dumps(obj: dict) -> str: + return json.dumps(obj) def error_type(status: int) -> str: diff --git a/mstar/api_server/openai/serving_chat.py b/mstar/api_server/openai/serving_chat.py index 8e007de3d4..cbb6a24d2d 100644 --- a/mstar/api_server/openai/serving_chat.py +++ b/mstar/api_server/openai/serving_chat.py @@ -97,25 +97,37 @@ def error(message: str, status: int) -> str: yield chunk({"role": "assistant"}) failed = False try: - async for c in api.iter_result_chunks(request_id): - if c.modality == "text": - yield chunk({"content": c.data.decode("utf-8", "replace")}) - elif c.modality == "audio": - # Streaming audio deltas are base64 16-bit PCM at the model rate. - yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) - elif c.modality == "image": - yield chunk({"content": media_io.png_to_data_url(c.data)}) - elif c.modality == "error": - # The request failed after the stream opened (an engine error - # mid generation, a preprocess error); the HTTP status is - # committed, so the failure travels in-band the way the - # non-streaming path's error body does, not as a normal - # ``stop`` a client would take for a complete answer. The error - # is the iterator's last chunk, so the loop ends on its own; - # returning here instead would trip the iterator's abort on a - # request that is already gone. - failed = True - yield error(c.data.decode("utf-8", "replace"), int(c.metadata.get("status", 500))) + # One yield per batch of text chunks: the tokens a request had waiting + # when its handler ran go out as one body frame (one write), which is + # what keeps the delivery loop ahead at high token rates. A batch of + # one is the common case when the loop keeps up. + async for batch in api.iter_result_chunk_batches(request_id): + text_events: list[str] = [] + for c in batch: + if c.modality == "text": + text_events.append(chunk({"content": c.data.decode("utf-8", "replace")})) + continue + if text_events: + yield "".join(text_events) + text_events = [] + if c.modality == "audio": + # Streaming audio deltas are base64 16-bit PCM at the model rate. + yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) + elif c.modality == "image": + yield chunk({"content": media_io.png_to_data_url(c.data)}) + elif c.modality == "error": + # The request failed after the stream opened (an engine error + # mid generation, a preprocess error); the HTTP status is + # committed, so the failure travels in-band the way the + # non-streaming path's error body does, not as a normal + # ``stop`` a client would take for a complete answer. The error + # is the iterator's last chunk, so the loop ends on its own; + # returning here instead would trip the iterator's abort on a + # request that is already gone. + failed = True + yield error(c.data.decode("utf-8", "replace"), int(c.metadata.get("status", 500))) + if text_events: + yield "".join(text_events) except HTTPException as exc: # The delivery timeout raises out of the iterator (which aborts the # request on its way out); report it the same way. diff --git a/mstar/communication/event.py b/mstar/communication/event.py index 352c07954c..82a8b00d51 100644 --- a/mstar/communication/event.py +++ b/mstar/communication/event.py @@ -9,6 +9,10 @@ def __init__(self): def _wake(self, _fut): # runs on whatever thread finished the future os.eventfd_write(self.event, 1) # one syscall, thread-safe, async-signal-safe + def signal(self): + """Wake the poller from any thread (a producer with new work).""" + os.eventfd_write(self.event, 1) + def register_future(self, future: Future): if future.done(): return diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 3569707ef3..b8aab3e747 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -33,6 +33,7 @@ from mstar.profile.format import RxInfo, TxInfo from mstar.profile.worker import GraphTimings from mstar.utils.exitcode import describe_exitcode +from mstar.utils.gc_freeze import freeze_after_setup from mstar.utils.ipc_format import ( ConductorMessageType, DrainRequest, @@ -1537,6 +1538,8 @@ def run(self): from mstar.utils.profiler import range_pop, range_push self._wait_for_workers_ready() + # the set-up heap is complete: keep gen-2 collections off it + freeze_after_setup("Conductor") self.communicator.send( "api_server", diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index c682b01cad..fedddc475f 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -295,6 +295,12 @@ class ExecutingBatch: # the resource that ran out, so an eviction can be scoped to it failed_resource: str | None = None + # Every output signal of this step stays on the device (the node's + # ``device_loopback_signals``): collect builds no per-request output + # dicts and the per-request ``postprocess`` is skipped; the stop check, + # the inline emit and the speculation splice read the rows instead. + rows_only: bool = False + # This step's outputs, published as soon as the forward has been # submitted — the tensors exist then, even though their values land later. outputs: BatchedModelOutput = field(default_factory=BatchedModelOutput) @@ -497,6 +503,9 @@ def load_model( for node in relevant_nodes: node_to_resources.setdefault(node, []).append(spec.resource_key) + # rid -> (node, walk) of its last step here when that step was rows-only, + # i.e. the sampler's slot master holds the token its next step reads + self._loopback_steps: dict[int, tuple[str, str]] = {} self._runner = StepRunner( self._resources, # every node, including one that owns nothing (Code2Wav) — an @@ -754,7 +763,22 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: # read leniently, as the probe itself was only reached per prepared row keyed = getattr(self, "_keyed_walks", None) probe_prefix = bool(keyed) and walk in keyed.get(batch.node_name, ()) - for rid in batch.request_ids: + loopback = self._loopback_rows(batch) + if not probe_prefix: + uniform = submodule.uniform_row_inputs( + walk, has_loopback=loopback is not None and all(loopback), + ) + if uniform is not None: + # one object for all the step's rows: the per-row call would + # return an equal one each time (128 calls at a full decode + # batch, on the gpu thread before the launch) + node_inputs = [uniform] * len(batch.request_ids) + batch.register_prepare_batch(node_inputs) + batch.running_batched = submodule.can_batch( + batch=batch, model_inputs=node_inputs + ) + return + for i, rid in enumerate(batch.request_ids): try: req_inputs = prepare( graph_walk=walk, @@ -766,6 +790,7 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: # look-ahead frames a token encoder withholds) is_final_stream_chunk=rid in final_rids, input_metadata=metadata.get(rid, EMPTY_INPUT_METADATA), + has_loopback=loopback is not None and loopback[i], ) if req_inputs is not None and probe_prefix: req_inputs = self._skip_cached_prefix(batch, rid, req_inputs) @@ -787,6 +812,26 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: batch=batch, model_inputs=node_inputs ) + def _loopback_rows(self, batch: ExecutingBatch) -> list[bool] | None: + """Per row, whether its input token is on the device (see + ``_loopback_steps``); None when no row's is.""" + steps = self._loopback_steps + if not steps: + return None + key = (batch.node_name, batch.step_context.graph_walk) + get = steps.get + return [get(rid) == key for rid in batch.request_ids] + + def _record_loopback(self, batch: ExecutingBatch) -> None: + steps = self._loopback_steps + if batch.rows_only: + steps.update(dict.fromkeys( + batch.request_ids, (batch.node_name, batch.step_context.graph_walk), + )) + elif steps: + for rid in batch.request_ids: + steps.pop(rid, None) + def inline_client_signals( self, node_name: str, graph_walk: str, ) -> dict[str, str]: @@ -796,6 +841,16 @@ def inline_client_signals( return {} return mgmt.submodule.inline_client_signals(graph_walk) + def device_loopback_signals( + self, node_name: str, graph_walk: str, + ) -> frozenset[str]: + """The node's loop-back signals it keeps on the device this walk, so + the worker routes them without a tensor (see the submodule hook).""" + mgmt = self._submodules.get(node_name) + if mgmt is None: + return frozenset() + return mgmt.submodule.device_loopback_signals(graph_walk) + def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], host_rows: HostRows | None = None, @@ -885,10 +940,15 @@ def exec( # admit/plan/commit assume the whole batch reaches the forward # in order; an unbatchable walk with >1 request can't, so each # request runs its own full cycle and merges. + if not batch.running_batched and lease is None: + # the unbatched forward has per-request outputs: no rows to stay on + batch.rows_only = False if not batch.running_batched and len(batch.request_ids) > 1 and lease is None: batch.outputs = self._exec_per_request(batch) else: batch.outputs = self._exec_single(batch) + if batch.admit_error is None: + self._record_loopback(batch) batch.outputs_ready.set() return batch.outputs except Exception: @@ -963,6 +1023,7 @@ def _exec_single( submodule_mgmt, lease, raw, inputs, req_info, request_ids=batch.request_ids, step_request_ids=batch.step_context.padded_request_ids, + rows_only=batch.rows_only, ) finally: if PHASE_PERIOD: @@ -1382,8 +1443,17 @@ def exec_and_postprocess( outputs = self.exec(batch) # A failed admit means no forward ran and every rid's outputs are empty, # so the tail has nothing to consume; the worker re-drives the step. + # A rows-only step has no per-request outputs; only a batch tail runs. if batch.admit_error is None: - self.postprocess_batch(batch, outputs) + if not batch.rows_only: + self.postprocess_batch(batch, outputs) + else: + submodule = self._submodules[batch.node_name].submodule + # no submodule overrides it as of now + if type(submodule).batch_postprocess is not NodeSubmodule.batch_postprocess: + submodule.batch_postprocess( + outputs, batch.per_request_info_wrapped, batch.inputs, + ) return outputs def check_stop_for_batch( @@ -1527,15 +1597,39 @@ def _collect_outputs( req_info: Mapping[str, CurrentForwardPassInfo], request_ids: list[str], step_request_ids: tuple[str, ...], + rows_only: bool = False, ) -> BatchedModelOutput: """Per-rid outputs for the real requests: drop the padding rows and map a captured graph's keys back to real ids. + ``rows_only`` (see ``ExecutingBatch.rows_only``) keeps the row views + and the stop buffers but builds no per-rid dicts: nothing downstream + reads them on such a step, and at 128 rows they were a few hundred + allocations on the gpu thread. + A captured forward emits its per-rid entries under the slot's padding ids (those were the batch at capture time), so entry ``i`` belongs to ``request_ids[i]`` on either path. """ submodule = submodule_mgmt.submodule + if ( + rows_only + and raw_outputs.check_stop_buffers is not None + and not raw_outputs.per_rid_outputs + and not raw_outputs.packed_outputs + and type(submodule).filter_batched_output + is NodeSubmodule.filter_batched_output + ): + # No clone either: the worker takes the host copy on this thread, + # on this stream, right behind the step (``_stage_host_rows``), so + # a later replay cannot overwrite the rows first, and nothing else + # reads them on the device. Padded rows past the real ones ride + # along and are never read. + return BatchedModelOutput( + check_stop_buffers=raw_outputs.check_stop_buffers, + row_request_ids=tuple(request_ids), + rows_only=True, + ) out_ids = ( step_request_ids if lease is None else submodule_mgmt.cuda_graph_runner.slot_for(lease).dummy_rids @@ -1544,6 +1638,7 @@ def _collect_outputs( row_clones, row_views = self._merge_per_rid( outputs, raw_outputs, request_ids, out_ids, submodule, req_info, + rows_only=rows_only, ) self._merge_unpacked( outputs, raw_outputs, request_ids, submodule, @@ -1551,6 +1646,7 @@ def _collect_outputs( ) return BatchedModelOutput( per_rid_outputs=outputs, + rows_only=rows_only and not outputs, # the stop buffer that is a row output shares its clone: one copy # of the sampled tokens per step, not two check_stop_buffers=raw_outputs.clone_check_stop_buffers( @@ -1570,6 +1666,7 @@ def _merge_per_rid( out_ids: list[str], submodule: NodeSubmodule, req_info: Mapping[str, CurrentForwardPassInfo], + rows_only: bool = False, ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[torch.Tensor, ...]] | None]: """Fold the forward's per-rid entries into ``outputs``. @@ -1596,7 +1693,7 @@ def _merge_per_rid( for name, tensor in row_clones.items() if isinstance(tensor, torch.Tensor) and tensor.dim() } - if row_views: + if row_views and not rows_only: for i, rid in enumerate(request_ids): merged = outputs.setdefault(rid, {}) for name, views in row_views.items(): @@ -1844,6 +1941,7 @@ def _pre_plan_for_batch( batch.step_context.set_padded_rids( cg_runner.step_ids(lease, batch.request_ids) ) + t0 = perf_counter() if PHASE_PERIOD else 0.0 step = submodule_mgmt.submodule.declare_step( graph_walk=batch.step_context.graph_walk, request_ids=batch.step_context.padded_request_ids, @@ -1851,6 +1949,8 @@ def _pre_plan_for_batch( slot_lease=lease, per_request_info=batch.per_request_info, ) + if PHASE_PERIOD: + phase_record("engine.pre_declare", perf_counter() - t0) if step is None: return False if batch.inputs is not None or ( @@ -1859,14 +1959,17 @@ def _pre_plan_for_batch( batch.step_context.graph_walk, ) ): - # exec drives this step as declared here (its admit ran as - # pre_admit), instead of declaring and admitting again + # exec drives this step as declared here, and its admit skips + # the resources pre_admit covered batch.step = step batch.step_context.is_preplan = True step.set_ctx(batch.step_context) try: + t0 = perf_counter() if PHASE_PERIOD else 0.0 admit_outcome = self._runner.pre_admit(step) + if PHASE_PERIOD: + phase_record("engine.pre_admit", perf_counter() - t0) if not admit_outcome.ok: # exec re-drives the step and reports the failure from there self.reset_pre_plan_for_batch(batch) @@ -2150,6 +2253,7 @@ def add_request( self._runner.ingest_request(request_id, overrides) def remove_request(self, request_id: str) -> None: + self._loopback_steps.pop(request_id, None) self._runner.remove_request(request_id) for submodule_mgmt in self._submodules.values(): submodule_mgmt.submodule.cleanup_request(request_id) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 0147596208..8938eea4cf 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -364,6 +364,9 @@ def __init__( ) self._streams: dict[str, LabelToStream] = {} self._overrides: dict[str, KVReqConfig] = {} + # (node, walk) -> (label, stop buffer name) per request, for the + # decode-step chain extension; dropped with the request + self._chain_lookup: dict[str, dict[tuple[str, str], tuple]] = {} # its entity id is the worker id, and one worker is one copy of a node self._replica = ( transfer_engine_info.my_entity_id @@ -717,13 +720,26 @@ def extend_prefix_chains_batch( buffers = host_rows.buffers rows_of: dict[str, list] = {} page_size = self.config.page_size + key = (node_name, graph_walk) + lookup = self._chain_lookup with self._lock: for rid in request_ids: - label = self._keyed_label(rid, node_name, graph_walk) - if label is None: - continue - tensor = (self._overrides[rid].prefix_decode or {}).get(label) - if not tensor: + # the label and buffer name do not change over a request's + # life (its overrides are set at ingest), so each is looked up + # once instead of on every decode step + per_rid = lookup.get(rid) + if per_rid is None: + per_rid = lookup[rid] = {} + entry = per_rid.get(key) + if entry is None: + label = self._keyed_label(rid, node_name, graph_walk) + tensor = ( + (self._overrides[rid].prefix_decode or {}).get(label) + if label is not None else None + ) + entry = per_rid[key] = (label, tensor) + label, tensor = entry + if label is None or not tensor: continue rows = rows_of.get(tensor) if rows is None: @@ -2007,6 +2023,7 @@ def remove_request(self, rid: str): self._forget_reservations(rid) self._streams.pop(rid, None) self._overrides.pop(rid, None) + self._chain_lookup.pop(rid, None) self._transfer.remove_request(rid) if _DEBUG_ASSERTS: self.assert_pages_conserved() diff --git a/mstar/engine/resources/kv/plan.py b/mstar/engine/resources/kv/plan.py index c0be0d33a9..3c2549f89d 100644 --- a/mstar/engine/resources/kv/plan.py +++ b/mstar/engine/resources/kv/plan.py @@ -128,7 +128,7 @@ def build_paged_indptrs( paged_kv_indptr=_int32_tensor(kv_indptr), paged_kv_indices=_int32_tensor(all_pages), paged_kv_last_page_len=_int32_tensor(last_page_lens), - kv_lens=_int32_tensor(kv_lens), + kv_lens=_host_lens(kv_lens), ) @@ -136,6 +136,26 @@ def _int32_tensor(values: list[int]) -> torch.Tensor: return torch.from_numpy(np.array(values, dtype=np.int32)) +class HostLens(torch.Tensor): + """A host int32 row whose iteration yields numpy scalars. + + FlashInfer's decode ``plan`` takes the longest row with the Python builtin + ``max(seq_lens)`` and ``.item()`` on the result; iterating a tensor hands + out one 0-d tensor per row, about 2 us each (0.35 ms at 128 rows, on the + plan thread every decode step). Iterating the numpy view costs 0.1 us a + row and the numpy scalar has ``.item()`` too. Everything else FlashInfer + does with the row (``.cpu()``, ``.min()``, ``len``, a ``copy_`` into its + device buffer) is plain tensor behaviour. + """ + + def __iter__(self): + return iter(self.numpy()) + + +def _host_lens(values: list[int]) -> torch.Tensor: + return _int32_tensor(values).as_subclass(HostLens) + + @dataclass class KVPlanOutput: """ diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index dbc1670c46..56787ad5bf 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -106,6 +106,8 @@ def __init__( # the step whose pre-plan is staged across the pre-planning resources, # as `_step_key` describes it; None when nothing is staged self._staged: tuple | None = None + # (step key, ready) of the last pre_admit whose stage is still up + self._pre_admitted: tuple | None = None def resolve_cached_prefix( self, rid: str, node_name: str, graph_walk: str, @@ -261,11 +263,23 @@ def ingest_request( self, rid: str, overrides: Mapping[str, ResourceReqConfig] | None = None, ) -> None: - """open state on all resources; `overrides` are per-resource""" - for key in self._order: - self._resources[key].ingest_request( - rid, None if overrides is None else overrides.get(key) - ) + """open state on all resources; `overrides` are per-resource. All or + nothing: a resource that raises leaves the request on none of them.""" + for i, key in enumerate(self._order): + try: + self._resources[key].ingest_request( + rid, None if overrides is None else overrides.get(key) + ) + except Exception: + # the raising resource may have taken part of it too + for done in self._order[:i + 1]: + try: + self._resources[done].remove_request(rid) + except Exception: + logger.exception( + "rolling back request %s on resource %s failed", rid, done, + ) + raise def remove_request(self, rid: str) -> None: for key in self._order: @@ -333,9 +347,10 @@ def _step_key(self, step: SubmoduleStep): ), ) - def _drop_stale_preplan(self, step: SubmoduleStep) -> None: + def _drop_stale_preplan(self, step: SubmoduleStep) -> tuple | None: + """Returns the step's key when the stage is for this very step.""" if self._staged is None or step.ctx.is_preplan: - return + return None key = self._step_key(step) if self._staged != key: # Not the normal path: the worker meant to run the staged step @@ -346,6 +361,8 @@ def _drop_stale_preplan(self, step: SubmoduleStep) -> None: self._staged[0], len(self._staged[1]), key[0], len(key[1]), ) self.clear_preplan() + return None + return key @property def staged(self) -> bool: @@ -355,6 +372,7 @@ def staged(self) -> bool: def clear_preplan(self) -> None: """Drop the staged pre-plan on every resource, and the record of it.""" self._staged = None + self._pre_admitted = None for key in self._order: self._resources[key].clear_preplan() @@ -362,10 +380,21 @@ def admit(self, step: SubmoduleStep) -> FullAdmitOutcome: """Reserve capacity for step. A refused admit gives back everything the step reserved. """ - self._drop_stale_preplan(step) + staged_key = self._drop_stale_preplan(step) + # A step whose ``pre_admit`` ran a step ahead, and whose stage is still + # up for it, skips the resources that sweep covered: each would only + # find its own reservation and no-op (0.1 ms of Python per step at + # 128 rows, on the gpu thread). A refusal below still unwinds them. + covered: list[str] = [] ready = True + pre = self._pre_admitted + if pre is not None and staged_key is not None and pre[0] == staged_key: + covered = self._preplan_keys_for(step) + ready = pre[1] admitted: list[str] = [] for key in self._keys_for(step): + if key in covered: + continue if self._nvtx: range_push(f"res.admit.{key}") t0 = perf_counter() if PHASE_PERIOD else 0.0 @@ -381,7 +410,7 @@ def admit(self, step: SubmoduleStep) -> FullAdmitOutcome: "Admit for resource %s failed with error: %s", key, outcome.reason.message ) - for done in reversed([*admitted, key]): + for done in reversed([*covered, *admitted, key]): self._resources[done].rollback_admit( step.get(done), step.ctx, ) @@ -400,6 +429,7 @@ def plan(self, step: SubmoduleStep) -> dict[str, Any]: # the resources promote their staged share below (or plan afresh # if nothing was staged): either way nothing stays staged self._staged = None + self._pre_admitted = None results = step.ctx.plan_results results.clear() for key in self._keys_for(step): @@ -442,9 +472,11 @@ def pre_admit(self, step: SubmoduleStep) -> FullAdmitOutcome: self._resources[done].rollback_admit( step.get(done), step.ctx, ) + self._pre_admitted = None return FullAdmitOutcome(outcome, key) admitted.append(key) ready = ready and outcome.ready + self._pre_admitted = (self._step_key(step), ready) return FULL_ADMIT_OK if ready else FULL_ADMIT_NOT_READY def pre_plan(self, step: SubmoduleStep) -> dict[str, Any]: diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index 522123adbf..1b86a60446 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -1,6 +1,7 @@ from dataclasses import asdict +from time import perf_counter import torch @@ -13,6 +14,8 @@ ) from mstar.engine.resources.sampler.utils import CudaGraphableSampler, Sampler, SamplerBuffers from mstar.engine.resources.step import SlotLease, StepContext +from mstar.utils.knobs import device_loopback_enabled, sampler_ingraph_scatter +from mstar.utils.profiler import PHASE_PERIOD, phase_record class SamplerResource(Resource): @@ -70,6 +73,11 @@ def __init__( self._preplan_cg_sampler: CudaGraphableSampler | None = None self._preplan_key = None self._preplanned = False + # (cg slot, padded bs) of the step planned last, for the last-token + # read in ``loopback_tokens``; None on an eager step + self._step_slot: tuple[int, int] | None = None + # whether the last sampled token is persisted by slot at all + self._keep_last_token = device_loopback_enabled() # rid -> the prompt tokens a cache hit kept out of this step's inputs self._cached_prefix: dict[str, torch.Tensor] = {} # this step's `SamplerStep.kept_rids`, for the eager draw @@ -120,6 +128,7 @@ def build_cuda_graph_buffers( vocab_size=self._vocab_size, cg_slots=self._cg_slots, enable_min_p=self._enable_min_p, + ingraph_scatter=self._keep_last_token and sampler_ingraph_scatter(), ) def ingest_request(self, rid: str, overrides: SamplingReqConfig | None=None): @@ -263,12 +272,14 @@ def plan(self, step: SamplerStep, ctx: StepContext): self._cg_sampler = self._preplan_cg_sampler self._preplan_cg_sampler = None self._preplanned = False + self._step_slot = (ctx.slot_lease.slot, ctx.slot_lease.bucket.bs) return # invalidated on the inline path here (not in commit, which now runs # before output collection); a preplan must leave the in-flight one be if not ctx.is_preplan: self._cg_sampler = None + self._step_slot = None lease = ctx.slot_lease if self._cg_buffers is None or lease is None: return @@ -290,6 +301,7 @@ def plan(self, step: SamplerStep, ctx: StepContext): # fresh inline (capture / no preplan): gather the per-step state too self._gather_dynamic(ctx, lease) self._cg_sampler = sampler + self._step_slot = (cg_slot, padded_bs) def _gather_dynamic(self, ctx: StepContext, lease: SlotLease): """Gather the per-step RNG offset + seen-token mask into this step's @@ -316,7 +328,24 @@ def commit(self, step: SamplerStep, ctx: StepContext): return kept = step.kept_rids keep = None if kept is None else [rid in kept for rid in ctx.request_ids] - self._cg_buffers.scatter_offset(ctx.slot_lease.slot, keep=keep) + if self._cg_buffers.ingraph_scatter: + # the captured sample scattered both rows itself + if self._penalty_needed_this_step: + self._cg_sampler.sync_seen_token_masks( + [self._sampler.get_token_mask(rid) for rid in ctx.request_ids], + keep=keep, + ) + return + t0 = perf_counter() if PHASE_PERIOD else 0.0 + keep_gpu = self._cg_buffers.keep_mask(keep) + self._cg_buffers.scatter_offset(ctx.slot_lease.slot, keep=keep_gpu) + if PHASE_PERIOD: + phase_record("res.commit.sampler.offset", perf_counter() - t0) + t0 = perf_counter() + if self._keep_last_token: + self._cg_buffers.scatter_last_token(ctx.slot_lease.slot, keep=keep_gpu) + if PHASE_PERIOD: + phase_record("res.commit.sampler.last_token", perf_counter() - t0) # Skipped when nothing read the mask this step: there is then nothing to # copy back, and the rows go stale only for requests at penalty 1.0, # which never read them. See `_penalty_live` for why that stays sound. @@ -348,4 +377,33 @@ def sample( request_ids, logits, apply_penalty=self._apply_penalty_this_step ) - return self._sampler.sample(request_ids, logits, kept_rids=self._kept_rids) + tokens = self._sampler.sample(request_ids, logits, kept_rids=self._kept_rids) + # An eager step has no per-step row to scatter from: write the masters + # by request, so the next step can still read its token on the device. + bufs = self._cg_buffers + if ( + self._keep_last_token and bufs is not None + and all(bufs.has_slot(rid) for rid in request_ids) + ): + bufs.write_last_tokens(request_ids, tokens) + return tokens + + @property + def has_slot_masters(self) -> bool: + """Whether the last sampled token of every request is kept on the + device by slot (the graph-buffer path): what the device-side + loop-back needs. False on an eager-only node or with the knob off.""" + return self._cg_buffers is not None and self._keep_last_token + + def loopback_tokens(self, request_ids: list[str]) -> torch.Tensor: + """The last sampled token of each row of this step, as its input ids, + read on the device. Under a slot lease the planned slot's index row + orders them (padding rows read slot 0); an eager step looks its + requests up by slot. Call after ``plan``.""" + bufs = self._cg_buffers + if bufs is None: + raise RuntimeError("device loop-back needs the sampler's slot masters") + if self._cg_sampler is not None and self._step_slot is not None: + cg_slot, padded_bs = self._step_slot + return bufs.gather_last_tokens(cg_slot, padded_bs) + return bufs.last_tokens_for(request_ids) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index ea17baaff7..c5206acae3 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -28,6 +28,7 @@ import triton.language as tl from mstar.utils.h2d import PinnedStager +from mstar.utils.knobs import sampler_slots logger = logging.getLogger(__name__) @@ -1032,17 +1033,45 @@ class CudaGraphableSampler(BaseSampler): seen_tokens_buf: torch.Tensor | None = None # [bs, V] bool # ``None`` for submodules whose ``SamplerSpec`` leaves ``enable_min_p`` off. min_p_buf: torch.Tensor | None = None + # ``[bs]`` per-step row of the slot's ``last_token`` buffer: the captured + # sample writes the (TP-agreed) token here, and ``commit`` scatters it to + # the slot master so the next step can read the token on the device. + last_token_buf: torch.Tensor | None = None tp_group: "CommGroup | None" = None # noqa: F821 + # In-graph gather/scatter (``MSTAR_SAMPLER_INGRAPH_SCATTER``): the slot + # masters and this slot's padded index row, all static, so the captured + # sample reads its offsets from the masters and writes the advanced + # offsets and the sampled tokens back without a launch after the replay. + offset_master: torch.Tensor | None = None + last_token_master: torch.Tensor | None = None + slot_idx_view: torch.Tensor | None = None + ingraph_scatter: bool = False # Set during graph capture, and used by the cuda graph runner to determine # whether requests' seen token buffers should be synced post-replay applied_penalty_in_graph: bool = False + def gather_in_graph(self) -> None: + """Offsets of this step's rows off the master, inside the graph.""" + torch.index_select( + self.offset_master, 0, self.slot_idx_view, out=self.offset_buf, + ) + + def scatter_in_graph(self, codes: torch.Tensor) -> None: + """Advanced offsets and sampled tokens back to the masters, inside the + graph: every row, padding rows into the trash row.""" + self.offset_master.index_copy_(0, self.slot_idx_view, self.offset_buf) + self.last_token_master.index_copy_( + 0, self.slot_idx_view, codes.to(self.last_token_master.dtype), + ) + @torch.compiler.disable def sample( self, request_ids: list[str], logits: torch.Tensor, apply_penalty: bool = False, ): + if self.ingraph_scatter: + self.gather_in_graph() codes = sample_cuda_graphable_gpu( logits, self.temperature_buf, self.top_k_buf, self.top_p_buf, @@ -1054,6 +1083,10 @@ def sample( ) self.offset_buf += 1 codes = self._broadcast_tokens(codes) + if self.last_token_buf is not None: + self.last_token_buf.copy_(codes) + if self.ingraph_scatter: + self.scatter_in_graph(codes) if apply_penalty and self.seen_tokens_buf is not None: self.applied_penalty_in_graph = True # Record the (broadcast, TP-agreed) token in the seen-token buffer so @@ -1108,6 +1141,7 @@ class Buffer: master: torch.Tensor default: float dtype: torch.dtype + _slot_views: dict = field(default_factory=dict) @classmethod def allocate( @@ -1136,8 +1170,18 @@ def grow_master(self, new_capacity: int) -> None: def slot_view(self, cg_slot: int, bs: int) -> torch.Tensor: """This slot's per-step row (clamped: a single-buffered buffer — the - RNG offset, gathered inline and used serialized — ignores ``cg_slot``).""" - return self.buf[cg_slot if self.buf.shape[0] > 1 else 0, :bs] + RNG offset, gathered inline and used serialized — ignores ``cg_slot``). + + Memoized per (slot, rows): ``buf`` is never reallocated, and the view + is taken on the gpu thread every step, where each tensor call is a + GIL release the main thread's postprocess can win.""" + key = (cg_slot, bs) + view = self._slot_views.get(key) + if view is None: + view = self._slot_views[key] = self.buf[ + cg_slot if self.buf.shape[0] > 1 else 0, :bs + ] + return view def gather(self, idx_view: torch.Tensor, padded_bs: int, cg_slot: int) -> None: torch.index_select(self.master, 0, idx_view, out=self.slot_view(cg_slot, padded_bs)) @@ -1154,11 +1198,10 @@ def scatter( ``keep`` ([real_bs] bool, on device) leaves the False rows' masters as they were. """ - idx = idx_view[:real_bs] rows = self.slot_view(cg_slot, real_bs) if keep is not None: - rows = torch.where(keep.view(-1, *([1] * (rows.dim() - 1))), rows, self.master[idx]) - self.master.index_copy_(0, idx, rows) + rows = torch.where(keep.view(-1, *([1] * (rows.dim() - 1))), rows, self.master[idx_view]) + self.master.index_copy_(0, idx_view, rows) @dataclass @@ -1284,6 +1327,14 @@ class SamplerBuffers: # ``_scalar_buffers`` (never written from a SamplingConfig) and reset to 0 # on register instead. offset: Buffer + # Per-request last sampled token, by slot: the captured sample writes the + # per-step row (double-buffered, one row per cg slot), ``scatter_last_token`` + # persists it to the master after the replay, and a decode step reads its + # input ids straight off the master (``gather_last_tokens``) instead of a + # per-request tensor routed through the graph. An eager step writes the + # master by request (``write_last_tokens``). Never gathered into the + # per-step row: only the sample writes it. + last_token: Buffer # TP communicator for the submodule that owns these buffers. Passed # through ``slice_for_bs`` into every per-step ``CudaGraphableSampler`` @@ -1304,6 +1355,12 @@ class SamplerBuffers: # Master cache capacity (grown by doubling when more requests are # concurrently registered than the per-step buffer holds). _master_capacity: int = field(default=0, repr=False) + # Rows of the device masters; lags ``_master_capacity`` until the gpu thread grows them. + _device_capacity: int = field(default=0, repr=False) + # the in-graph scatter (knob): padding rows address ``_trash_slot``, one + # master row past the capacity that no request is ever given + ingraph_scatter: bool = False + _trash_slot: int = 0 # Number of double-buffer slots (per-step buffers carry this leading dim). cg_slots: int = 1 # Per-step slot-index staging, per cg slot. ``_slot_idx_cpu`` is pinned so @@ -1318,6 +1375,8 @@ class SamplerBuffers: # Real (unpadded) batch size of the last gather per cg slot — the # scatter-back writes only these rows (padding rows all map to slot 0). _last_real_bs: list[int] = field(default_factory=list, repr=False) + # memoized per (slot, real rows): see ``_scatter_idx`` + _scatter_idx_views: dict = field(default_factory=dict, repr=False) # cg slot -> the (rids, padded_bs) its pinned index row currently holds, so # gather_dynamic can re-issue the H2D without redoing the O(bs) CPU fill _staged_rids: dict[int, tuple[tuple[str, ...], int]] = field( @@ -1334,6 +1393,8 @@ class SamplerBuffers: # thread, default stream) drains it. _pending_offset_reset: set[int] = field(default_factory=set, repr=False) _offset_reset_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + # Host master-row writes vs their regrow; a plan-thread write between copy and swap was lost. + _host_master_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) # Last-known config per rid — change-detect for ``update_request_config`` # so steady-state per-step calls do zero GPU work (for the scalar rows). _cached_config: dict[str, SamplingConfig] = field(default_factory=dict, repr=False) @@ -1367,9 +1428,15 @@ def allocate( vocab_size: int | None = None, cg_slots: int = 1, enable_min_p: bool = False, + ingraph_scatter: bool = False, + slots: int | None = None, ) -> "SamplerBuffers": """Allocate sampling buffers for ``max_batch_size``. + With ``ingraph_scatter`` the masters hold ``slots`` rows (the knob's + 4096 by default) plus one trash row for padding, allocated once: the + captured graphs keep their addresses, so they never grow. + ``vocab_size`` (when not None) enables the seen-token mask buffer for the repetition penalty. The master rows default to a ``SamplingConfig()`` row (temp=1, top_k=0, top_p=1, rep_penalty=1) — what an unregistered slot @@ -1378,16 +1445,21 @@ def allocate( """ pinned = torch.cuda.is_available() and device.type == "cuda" cap = max_batch_size + if ingraph_scatter: + cap = max(max_batch_size, slots if slots is not None else sampler_slots()) + # one row past the capacity is the trash row for padding when the + # in-graph scatter is on (harmless, never read for a request) + rows = cap + 1 if ingraph_scatter else cap def mk(dtype: torch.dtype, default: float) -> HostBuffer: return HostBuffer.allocate( - max_batch_size, cap, device, dtype, default, cg_slots + max_batch_size, rows, device, dtype, default, cg_slots ) # seen token mask is not double-buffered, as it depends on the GPU # value of the previous step. Same goes for offset seen_tokens = ( - MaskBuffer.allocate(max_batch_size, cap, vocab_size, device, cg_slots=1) + MaskBuffer.allocate(max_batch_size, rows, vocab_size, device, cg_slots=1) if vocab_size is not None else None ) return cls( @@ -1397,11 +1469,15 @@ def mk(dtype: torch.dtype, default: float) -> HostBuffer: top_p=mk(torch.float32, 1.0), seed=mk(torch.long, 0), rep_penalty=mk(torch.float32, 1.0), - offset=Buffer.allocate(max_batch_size, cap, device, torch.long, 0, 1), + offset=Buffer.allocate(max_batch_size, rows, device, torch.long, 0, 1), + last_token=Buffer.allocate(max_batch_size, rows, device, torch.long, 0, cg_slots), tp_group=tp_group, seen_tokens=seen_tokens, min_p=mk(torch.float32, 0.0) if enable_min_p else None, _master_capacity=cap, + _device_capacity=cap, + ingraph_scatter=ingraph_scatter, + _trash_slot=cap if ingraph_scatter else 0, cg_slots=cg_slots, _slot_idx_cpu=torch.zeros(cg_slots, max_batch_size, dtype=torch.long, pin_memory=pinned), _slot_idx_gpu=torch.zeros(cg_slots, max_batch_size, dtype=torch.long, device=device), @@ -1423,7 +1499,12 @@ def slice_for_bs(self, bs: int, cg_slot: int = 0) -> dict[str, Any]: "rep_penalty_buf": self.rep_penalty.slot_view(cg_slot, bs), "seen_tokens_buf": self.seen_tokens.slot_view(cg_slot, bs) if self.seen_tokens is not None else None, "min_p_buf": self.min_p.slot_view(cg_slot, bs) if self.min_p is not None else None, + "last_token_buf": self.last_token.slot_view(cg_slot, bs), "tp_group": self.tp_group, + "offset_master": self.offset.master, + "last_token_master": self.last_token.master, + "slot_idx_view": self._slot_idx_gpu[cg_slot, :bs], + "ingraph_scatter": self.ingraph_scatter, } # ------------------------------------------------------------------ @@ -1447,13 +1528,16 @@ def _write_master_row(self, slot: int, cfg: SamplingConfig) -> None: # argmax (first index on ties), which from_probs returns for any # seed. See ``sample_cuda_graphable_gpu``. t, k, p = 0.0, 0, 1.0 - self.temperature.write_master_row(slot, t) - self.top_k.write_master_row(slot, k) - self.top_p.write_master_row(slot, p) - self.seed.write_master_row(slot, cfg.seed) - self.rep_penalty.write_master_row(slot, float(cfg.repetition_penalty)) - if self.min_p is not None: - self.min_p.write_master_row(slot, float(cfg.min_p) if cfg.temperature > 0 else 0.0) + with self._host_master_lock: + self.temperature.write_master_row(slot, t) + self.top_k.write_master_row(slot, k) + self.top_p.write_master_row(slot, p) + self.seed.write_master_row(slot, cfg.seed) + self.rep_penalty.write_master_row(slot, float(cfg.repetition_penalty)) + if self.min_p is not None: + self.min_p.write_master_row( + slot, float(cfg.min_p) if cfg.temperature > 0 else 0.0 + ) self._config_version += 1 def _grow_master(self, new_capacity: int) -> None: @@ -1462,19 +1546,44 @@ def _grow_master(self, new_capacity: int) -> None: Triggered when concurrently-registered requests exceed the current master capacity. Per-step buffers (sized to the cuda-graph max_bs) are NOT resized — the gather only reads ``padded_bs`` rows from master. + + Host rows only, under the lock ``_write_master_row`` takes, so a plan-thread + ``_init_slot`` cannot land in the old array; device masters grow on the gpu thread. """ - for buf in self._scalar_buffers(): - buf.grow_master(new_capacity) - self.offset.grow_master(new_capacity) - if self.seen_tokens is not None: - self.seen_tokens.grow_master(new_capacity) - self._free_slots.extend(range(self._master_capacity, new_capacity)) + if self.ingraph_scatter: + raise RuntimeError( + f"more than {self._master_capacity} concurrent requests on the " + "sampler with the in-graph scatter on; raise MSTAR_SAMPLER_SLOTS " + "or set MSTAR_SAMPLER_INGRAPH_SCATTER=0" + ) + with self._host_master_lock: + for buf in self._scalar_buffers(): + buf.grow_master(new_capacity) + old_capacity = self._master_capacity self._master_capacity = new_capacity + self._free_slots.extend(range(old_capacity, new_capacity)) + + def _sync_device_capacity(self) -> None: + """Grow the device masters on the gpu thread, behind its own scatters. + + Swapping them on the main thread let an in-flight scatter land in the old tensor. + """ + cap = self._master_capacity + if self._device_capacity >= cap: + return + self.offset.grow_master(cap) + self.last_token.grow_master(cap) + if self.seen_tokens is not None: + self.seen_tokens.grow_master(cap) + self._device_capacity = cap def register_request( self, rid: str, sampling_config: SamplingConfig | None = None, ) -> None: """Allocate a slot for ``rid`` and seed its master row.""" + # a capture padding handle (negative) takes no slot; it reads slot 0 + if isinstance(rid, int) and rid < 0: + return if rid in self._rid_to_slot: # Re-registration: just refresh the config in place. if sampling_config is not None: @@ -1534,6 +1643,7 @@ def stage_seen_token_masks( """ if self.seen_tokens is None: return + self._sync_device_capacity() # Batched for the same reason as sync_seen_token_masks: bs launches # become one. _init_slot still has to run per slot (it is CPU-side # bookkeeping plus a master-row write) before anything is staged. @@ -1598,7 +1708,7 @@ def _stage_slot_idx( self._init_slot(slot, rid) rows.append(slot) if len(rows) < padded_bs: - rows.extend([0] * (padded_bs - len(rows))) + rows.extend([self._trash_slot] * (padded_bs - len(rows))) self._slot_idx_np[cg_slot, :padded_bs] = rows self._staged_rids[cg_slot] = (tuple(request_ids), padded_bs) @@ -1654,6 +1764,7 @@ def gather_dynamic( The seen-token mask is large ([bs, V] bool); only gathered when the caller's graph actually applies the penalty in-graph (the Talker).""" + self._sync_device_capacity() if self._staged_rids.get(cg_slot) != (tuple(request_ids), padded_bs): self._stage_slot_idx(request_ids, padded_bs, cg_slot) with self._offset_reset_lock: @@ -1661,7 +1772,9 @@ def gather_dynamic( for slot in resets: self.offset.master[slot:slot + 1].zero_() idx_view = self._upload_slot_idx(padded_bs, cg_slot) - self.offset.gather(idx_view, padded_bs, cg_slot) + if not self.ingraph_scatter: + # with it on, the captured sample gathers off the master itself + self.offset.gather(idx_view, padded_bs, cg_slot) if self.seen_tokens is not None and gather_seen_tokens: self.seen_tokens.gather(idx_view, padded_bs, cg_slot) @@ -1680,21 +1793,81 @@ def sampler_for(self, padded_bs: int, cg_slot: int) -> "CudaGraphableSampler": sampler.applied_penalty_in_graph = False return sampler + def _scatter_idx(self, cg_slot: int) -> tuple[torch.Tensor, int]: + """The slot's device index row narrowed to the real rows, memoized per + (slot, rows) like ``Buffer.slot_view`` (the index buffer is static).""" + real_bs = self._last_real_bs[cg_slot] + key = (cg_slot, real_bs) + idx = self._scatter_idx_views.get(key) + if idx is None: + idx = self._scatter_idx_views[key] = self._slot_idx_gpu[cg_slot, :real_bs] + return idx, real_bs + + def keep_mask(self, keep: list[bool] | None) -> torch.Tensor | None: + """``keep`` (request order) on the device for ``scatter_*``; None when + every row is kept.""" + if keep is None or all(keep): + return None + # fresh pinned tensor: the host allocator holds it until the copy lands + device = self._slot_idx_gpu.device + return torch.tensor( + keep, dtype=torch.bool, pin_memory=device.type == "cuda", + ).to(device, non_blocking=True) + def scatter_offset( - self, cg_slot: int = 0, keep: list[bool] | None = None, + self, cg_slot: int = 0, keep: torch.Tensor | None = None, ) -> None: """Persist the (in-graph advanced) per-step offsets back to their slot masters. Call once AFTER the graph replay for the gather on ``cg_slot``; - GPU-only, real rows only (padding rows all map to slot 0). ``keep``, - in request order, skips the rows whose draw was dropped.""" - keep_gpu = None - if keep is not None and not all(keep): - # fresh pinned tensor: the host allocator holds it until the copy lands - device = self._slot_idx_gpu.device - keep_gpu = torch.tensor( - keep, dtype=torch.bool, pin_memory=device.type == "cuda", - ).to(device, non_blocking=True) - self.offset.scatter( - self._slot_idx_gpu[cg_slot], self._last_real_bs[cg_slot], cg_slot, - keep=keep_gpu, + GPU-only, real rows only (padding rows all map to slot 0). ``keep`` + (from ``keep_mask``) skips the rows whose draw was dropped.""" + idx, real_bs = self._scatter_idx(cg_slot) + self.offset.scatter(idx, real_bs, cg_slot, keep=keep) + + # ------------------------------------------------------------------ + # Last sampled token: written by the sample, read by the next step + # ------------------------------------------------------------------ + def scatter_last_token( + self, cg_slot: int = 0, keep: torch.Tensor | None = None, + ) -> None: + """Persist the tokens the captured sample wrote into ``cg_slot``'s + per-step row to their slot masters. Same contract as ``scatter_offset``: + after the replay, real rows only.""" + idx, real_bs = self._scatter_idx(cg_slot) + self.last_token.scatter(idx, real_bs, cg_slot, keep=keep) + + def gather_last_tokens(self, cg_slot: int, padded_bs: int) -> torch.Tensor: + """``[padded_bs]`` last tokens of the batch staged on ``cg_slot``, in + its row order (padding rows read slot 0, like every other gather). + A fresh tensor on the current stream; the runner stages it into the + slot's static ``input_ids``. Call after ``gather_dynamic`` uploaded the + slot's index row for this step.""" + return self.last_token.master.index_select( + 0, self._slot_idx_gpu[cg_slot, :padded_bs] ) + + def _slots_of(self, request_ids: list[str]) -> torch.Tensor: + slots = [ + 0 if isinstance(rid, int) and rid < 0 else self._rid_to_slot[rid] + for rid in request_ids + ] + return torch.tensor(slots, dtype=torch.long, device=self.last_token.master.device) + + def write_last_tokens(self, request_ids: list[str], tokens: torch.Tensor) -> None: + """Eager-step write of ``tokens`` (``[len(request_ids)]``) to the + requests' master rows. Every request must be registered.""" + if not request_ids: + return + self._sync_device_capacity() + self.last_token.master.index_copy_( + 0, self._slots_of(request_ids), tokens.to(self.last_token.master.dtype) + ) + + def last_tokens_for(self, request_ids: list[str]) -> torch.Tensor: + """``[len(request_ids)]`` last tokens of registered requests, for an + eager step that has no staged slot row.""" + self._sync_device_capacity() + return self.last_token.master.index_select(0, self._slots_of(request_ids)) + + def has_slot(self, rid: str) -> bool: + return rid in self._rid_to_slot diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index 15844215f5..ddd36b9160 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -60,6 +60,7 @@ Qwen3_5VisionConfig, ) from mstar.model.submodule_base import NodeSubmodule +from mstar.utils.fast_tokenize import encode_ids logger = logging.getLogger(__name__) @@ -421,8 +422,8 @@ def process_prompt( add_generation_prompt=True, enable_thinking=kwargs.get("enable_thinking", True), ) - input_ids = self.tokenizer(text, return_tensors="pt").input_ids[0] - + # a long prompt is encoded as parallel chunks, identical ids + input_ids = torch.tensor(encode_ids(self.tokenizer, text), dtype=torch.long) spans = find_media_spans(input_ids, self._placeholder_specs()) segments = split_around_spans(input_ids, spans) check_plan(prefill_plan(parts), spans, len(segments)) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 883d5deec9..2acf405560 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -9,6 +9,7 @@ import logging import os from collections.abc import Mapping +from time import perf_counter from typing import Any import torch @@ -69,13 +70,25 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + device_loopback_enabled, ) +from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) class LLMSubmodule(ARNodeSubmodule): - PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024, 2048] + # Total tokens of a packed text-prefill step, one captured graph per + # (bucket, batch size). A prompt past the top bucket prefills eagerly, and + # an eager 8k prefill is launch-bound (~1,100 kernel launches, the GPU + # idle half the time: 78 ms of wall for 36 ms of kernels on 0.8B), so the + # ladder runs to 8192. MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS narrows or + # extends it (static buffers are sized by the largest bucket). + PREFILL_TOKEN_BUCKETS = [ + int(x) for x in os.environ.get( + "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", "32,64,128,256,512,1024,2048,4096,8192", + ).split(",") + ] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] # Capture rows and a replay's padding rows address the pool's sink and # hold no slot, so these buckets do not size `gdn_state.max_slots`; that is @@ -273,14 +286,38 @@ def prepare_inputs( graph_walk: str, fwd_info: CurrentForwardPassInfo, inputs: NameToTensorList, + has_loopback: bool = False, **kwargs: Any, ) -> ARNodeInputs: if graph_walk == "prefill_vision": return self._vision_inputs(fwd_info, inputs) + if graph_walk == "decode" and has_loopback and self._device_loopback(): + # one token per row, read off the sampler's slot master in + # `preprocess`; the routed signal is empty (or ignored) here + return ARNodeInputs(input_seq_len=1) input_ids = inputs["text_inputs"][0] return ARNodeInputs(input_seq_len=input_ids.shape[0], input_ids=input_ids) + def uniform_row_inputs( + self, graph_walk: str, has_loopback: bool = False, + ) -> ARNodeInputs | None: + # the loop-back decode row above, built once for the step's rows + if graph_walk == "decode" and has_loopback and self._device_loopback(): + return ARNodeInputs(input_seq_len=1) + return None + + def _device_loopback(self) -> bool: + return ( + device_loopback_enabled() + and self.node_resources[SAMPLER].has_slot_masters + ) + + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: + if graph_walk == "decode" and self._device_loopback(): + return frozenset({"text_inputs"}) + return frozenset() + def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: """``[3, total_tokens]`` for the step, in packed request order. @@ -319,7 +356,29 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: out: dict[str, torch.Tensor | Any] = {} - if inputs[0].input_ids is not None: + first = inputs[0] + if first.input_embeds is None and ( + first.input_ids is None or any(inp.input_ids is None for inp in inputs) + ): + # device loop-back: every row of a decode step is its request's + # last sampled token, read by slot (padding rows read slot 0) + sampler: SamplerResource = engine_inputs.resources[SAMPLER] + t0 = perf_counter() if PHASE_PERIOD else 0.0 + ids = sampler.loopback_tokens(engine_inputs.request_ids) + # rows whose last step ran elsewhere brought their token + routed = [ + i for i, inp in enumerate(inputs) + if inp.input_ids is not None and inp.input_seq_len + ] + if routed: + ids.index_copy_( + 0, torch.tensor(routed, device=ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + out["input_ids"] = ids + if PHASE_PERIOD: + phase_record("engine.preprocess.loopback", perf_counter() - t0) + elif inputs[0].input_ids is not None: out["input_ids"] = torch.cat([inp.input_ids for inp in inputs], dim=0) else: out["input_embeds"] = torch.cat( @@ -531,23 +590,31 @@ def check_stop_batched( if not torch.is_tensor(tokens) or tokens.dim() == 0: return None values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + rows = host_rows.request_ids + if len(rows) >= len(request_ids) and all( + a == b for a, b in zip(request_ids, rows, strict=False) + ): + # the common case: the batch in forward order, no row map needed + pairs = zip(request_ids, values, strict=False) + else: + row_of = {rid: i for i, rid in enumerate(rows)} + pairs = ( + (rid, values[i]) for rid in request_ids + if (i := row_of.get(rid)) is not None and i < len(values) + ) # no row, no output: the per-request path skips it too stop_ids = self.config.stop_token_ids stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too + for rid, value in pairs: info = request_infos[rid] - hit_eos = ( - not info.resource_configs[SAMPLER].ignore_eos - and values[i] in stop_ids - ) - out_of_budget = ( + # the set test first: a stop id is rare, the config read behind it + # goes through the wrapper + if ( + value in stop_ids + and not info.resource_configs[SAMPLER].ignore_eos + ) or ( info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens - ) - if hit_eos or out_of_budget: + ): stops[rid] = {"decode_loop"} return stops diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 7a32432592..238b6037f5 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -13,6 +13,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep +from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) class HostRows(NamedTuple): @@ -84,6 +85,13 @@ class BatchedModelOutput: # of a batch clone: name -> the views, row i for ``row_request_ids[i]``. # The store describes such rows without reading each view. row_views: dict[str, tuple[torch.Tensor, ...]] | None = None + # The stop buffers' host copy, staged by the gpu thread on the execution + # stream behind the step so the step's completion event covers it; the + # main thread reads it after waiting on that event. + host_rows: HostRows | None = None + # No per-rid entries were built: every output signal of the step stays on + # the device (see ``ExecutingBatch.rows_only``). Readers use the rows. + rows_only: bool = False @classmethod def coerce(cls, output: BatchedModelOutput | dict[str, Any]) -> BatchedModelOutput: @@ -688,6 +696,20 @@ def filter_batched_output( a fixed key set for graph compat, so the filtering happens here.""" return outputs + def uniform_row_inputs( + self, graph_walk: str, has_loopback: bool = False, + ) -> "NodeInputs | None": + """The inputs every row of this walk gets, when ``prepare_inputs`` + would return an equal object for each regardless of the request, its + routed signals and its metadata (a decode walk that reads its token + off the device). The engine then prepares the batch with one call and + shares the object across the step's rows, so it must not be mutated + per row. ``has_loopback``: every row's last step on this node kept its + loop-back signals on the device (``device_loopback_signals``). None + (the default) prepares row by row.""" + del graph_walk, has_loopback + return None + def step_is_input_free(self, graph_walk: str) -> bool: """Whether ``declare_step`` for this walk reads nothing from the inputs but their spans, and every real row has the span of the template row. @@ -697,6 +719,22 @@ def step_is_input_free(self, graph_walk: str) -> bool: del graph_walk return False + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: + """Loop-back input signals of this walk whose per-row value the node + keeps on the device itself (a sampler's slot master), so the worker + routes them with no tensor: the next step reads the value in + ``preprocess`` and every fallback batch arrives with the signal + empty. ``prepare_inputs`` is told per row (``has_loopback``) whether + its last step here kept the value on the device; a row without it + (e.g. its previous step ran on another rank) brings the tensor. + On a walk whose output signals all stay on the device the + engine builds no per-request output dicts and does not call + ``postprocess``: the stop check, the inline emit and the next step's + speculation read the step's rows instead. Empty by default; + ``MSTAR_DEVICE_LOOPBACK=0`` keeps every node on the tensor path.""" + del graph_walk + return frozenset() + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: """Output signals whose client-facing value is one scalar per row that the stop check already copies to the host, as signal name -> the name @@ -845,6 +883,15 @@ def postprocess( """ return + def batch_postprocess( + self, outputs: BatchedModelOutput, + request_infos: dict[int, CurrentForwardPassInfo], + inputs: list[NodeInputs], + ) -> None: + """``postprocess`` for a rows-only step, which has no per-request + outputs: works on the rows of ``outputs`` in place. Optional; only + called when overridden.""" + def check_stop( self, request_id: str, request_info: CurrentForwardPassInfo, diff --git a/mstar/utils/fast_tokenize.py b/mstar/utils/fast_tokenize.py new file mode 100644 index 0000000000..88e6fddca6 --- /dev/null +++ b/mstar/utils/fast_tokenize.py @@ -0,0 +1,93 @@ +"""Tokenize a long prompt as parallel chunks, byte-identical to one encode. + +An 8k-token prompt costs the fast tokenizer ~15 ms single-threaded, a quarter +of the request's time to first token. The Rust backend encodes a batch in +parallel, so a long text is cut into chunks and encoded as one batch; the +cuts are chosen so every pre-token stays whole, which keeps the BPE merges, +and so the ids, identical to encoding the whole text. + +The cut rule assumes a GPT-style pre-tokenizer (the Qwen2 family, GPT-2, +Llama 3): a space starts a pre-token unless it sits inside a run of +whitespace, so a cut is placed only right before a space whose previous +character is not whitespace. Runs of whitespace, newline runs and the +punctuation-plus-newline pre-tokens are never cut, and special-token strings +contain no space. The identity is checked by test/modular/test_fast_tokenize.py +against the real tokenizer over the benchmark prompts. +""" +from __future__ import annotations + +import os + +# Below this many characters a single encode is as fast; above it chunking +# pays. ~4 KB of text is roughly 1k tokens. +DEFAULT_CHUNK_CHARS = int(os.environ.get("MSTAR_TOKENIZE_CHUNK_CHARS", "4096")) +DEFAULT_MAX_CHUNKS = int(os.environ.get("MSTAR_TOKENIZE_MAX_CHUNKS", "8")) + + +def cut_points( + text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS, +) -> list[int]: + """Indices where ``text`` may be cut: each is a space that starts a + whitespace run, at or after the next chunk target. Empty when the text is + short or no safe cut exists.""" + n = len(text) + if chunk_chars <= 0 or n <= chunk_chars: + return [] + want = min(max_chunks, max(2, n // chunk_chars)) + size = n // want + cuts: list[int] = [] + at = 0 + for _ in range(want - 1): + target = at + size + if target >= n: + break + i = text.find(" ", target) + while i > 0 and text[i - 1].isspace(): + i = text.find(" ", i + 1) + if i <= 0 or i >= n - 1: + break + cuts.append(i) + at = i + return cuts + + +def split_at(text: str, cuts: list[int]) -> list[str]: + parts = [] + at = 0 + for c in cuts: + parts.append(text[at:c]) + at = c + parts.append(text[at:]) + return parts + + +def _parallelism_on() -> None: + """The Rust backend encodes a batch on its thread pool only while + ``TOKENIZERS_PARALLELISM`` reads true; serving environments often export + it false (to silence the fork warning), which makes the chunked encode + no faster than one encode (18 vs 19 ms for 8.8k tokens on H200 hosts). + ``MSTAR_TOKENIZE_PARALLEL=0`` leaves the variable alone.""" + if os.environ.get("MSTAR_TOKENIZE_PARALLEL", "1") != "1": + return + if os.environ.get("TOKENIZERS_PARALLELISM", "").lower() not in ("true", "1"): + os.environ["TOKENIZERS_PARALLELISM"] = "true" + + +def encode_ids( + tokenizer, text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, + max_chunks: int = DEFAULT_MAX_CHUNKS, +) -> list[int]: + """Token ids of ``text`` with the tokenizer's defaults for special + tokens: a long text goes through the Rust backend as a batch of chunks, + a short one (or a tokenizer without a Rust backend) through ``tokenizer`` + itself.""" + backend = getattr(tokenizer, "backend_tokenizer", None) + cuts = cut_points(text, chunk_chars, max_chunks) if backend is not None else [] + if not cuts: + return list(tokenizer(text).input_ids) + _parallelism_on() + encodings = backend.encode_batch(split_at(text, cuts), add_special_tokens=False) + ids: list[int] = [] + for enc in encodings: + ids.extend(enc.ids) + return ids diff --git a/mstar/utils/gc_freeze.py b/mstar/utils/gc_freeze.py new file mode 100644 index 0000000000..53be3c2c02 --- /dev/null +++ b/mstar/utils/gc_freeze.py @@ -0,0 +1,25 @@ +"""Move a process's set-up heap out of the garbage collector's reach. + +A serving process builds a large, long-lived heap while it starts (configs, +tokenizers, weights' Python wrappers, capture buffers). Every gen-2 +collection afterwards walks all of it: measured at 200 ms per pass in the +API server process, every ~20 s under load, each pass stalling request +admission and token delivery for that long. Collecting once and freezing +moves those objects to the permanent generation, so later collections only +walk what the process created since. Refcounting still frees frozen objects; +only a reference cycle alive at the freeze would be retained for good. +""" +import gc +import logging + +logger = logging.getLogger(__name__) + + +def freeze_after_setup(who: str) -> int: + """``gc.collect()`` then ``gc.freeze()`` for a process that finished its + set-up. Returns the number of objects frozen.""" + gc.collect() + gc.freeze() + frozen = gc.get_freeze_count() + logger.info("%s: gc.freeze() after setup, %d objects in the permanent generation", who, frozen) + return frozen diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py new file mode 100644 index 0000000000..e6e0df4388 --- /dev/null +++ b/mstar/utils/knobs.py @@ -0,0 +1,28 @@ +"""Environment knobs read by more than one package.""" +import os + + +def device_loopback_enabled() -> bool: + """``MSTAR_DEVICE_LOOPBACK`` (default 1): whether a node may keep its + loop-back token on the device (the sampler's slot master) instead of + routing a tensor per request; also whether the sampler persists the + last token at all.""" + return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" + + +def sampler_ingraph_scatter() -> bool: + """``MSTAR_SAMPLER_INGRAPH_SCATTER`` (default 0): the captured sampler + gathers its RNG offsets from the slot masters and scatters the advanced + offsets and the sampled tokens back inside the graph, so a replayed step + issues none of those launches on the gpu thread (each is a GIL release + the main thread's postprocess can win). Padding rows then address a trash + master row, and the masters are allocated once for + ``MSTAR_SAMPLER_SLOTS`` (default 4096) concurrent requests, since a + captured graph holds their addresses.""" + return os.environ.get("MSTAR_SAMPLER_INGRAPH_SCATTER", "0") == "1" + + +def sampler_slots() -> int: + """``MSTAR_SAMPLER_SLOTS`` (default 4096): master rows allocated up front + when the in-graph scatter is on (see ``sampler_ingraph_scatter``).""" + return int(os.environ.get("MSTAR_SAMPLER_SLOTS", "4096")) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 6478e15e1a..1af9ab1e45 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -22,6 +22,7 @@ from mstar.communication.tensors import ( LocalTransferEngine, NameToTensorList, + StoredOutputs, create_tensor_communication_manager, ) from mstar.conductor.request_info import CurrentForwardPassInfo @@ -148,6 +149,14 @@ def _parse_tp_async_sched(raw: str) -> tuple[bool, frozenset[str] | None]: return True, frozenset(n.strip() for n in raw.split(",") if n.strip()) +@dataclass +class _SpecFairness: + """The decode loop's fairness settings for ``_build_speculation``: whether + to peek for other ready work, and the consecutive-step cap.""" + peek: bool + max_consecutive: int + + @dataclass class PendingBatch: batch: ScheduledBatch @@ -206,6 +215,10 @@ class EvictionPolicy(Enum): # resource to consult. +# pinned host-row buffers in flight at once (see Worker._stage_host_rows) +_HOST_ROWS_RING = 4 + + class _LazyPerRid: """``Worker._rows_to_per_rid`` on demand: the per-rid host dicts are built the first time a consumer reads them. The batched stop check, the row store @@ -259,6 +272,11 @@ class Worker: computation via engines. """ + # Rids that own a stream buffer (a streaming model's chunked inputs). + # Empty for a text model, which lets a speculation build skip its three + # per-row stream polls; None (a bare instance) means unknown, so poll. + _stream_buffer_rids: set[int] | None = None + def __init__( self, worker_id: str, @@ -283,6 +301,10 @@ def __init__( self.worker_id = worker_id self.device = device self.enable_nvtx = nvtx_enabled(enable_nvtx, device) + # (node, walk) -> the loop-back signals the node keeps on the device + self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} + self._host_rows_slot = 0 + self._stream_buffer_rids = set() # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -664,12 +686,19 @@ def _add_new_request(self, body: NewRequest) -> None: self._last_active[(request_id, node_name)] = now self.request_state.add_request(request_id, body.request_info) - self.engine_manager.add_request( - request_id, body.request_info.resource_configs, - ) self.tensor_manager.register_request( request_id, self._graph_runtime.get_sharding_config(request_id), ) + try: + self.engine_manager.add_request( + request_id, body.request_info.resource_configs, + ) + except Exception as err: + # the engine rolled the ingest back; the conductor's REMOVE tears down the rest + logger.exception("Worker %s: ingest of request %s failed", self.worker_id, body.request_id) + self._fail_requests({request_id: f"Error in worker: {type(err).__name__}: {err}"}) + self._process_unprocessed(body.request_id) + return # Create StreamBuffers for consumer connections on this worker. A request # with several partitions here arrives once per partition: keep the @@ -687,6 +716,8 @@ def _add_new_request(self, body: NewRequest) -> None: req_info.stream_buffers[conn.edge_name] = sbuf consumer = self._consumer_node_cache.get(conn.edge_name, "") req_info.stream_buffers_by_consumer.setdefault(consumer, {})[conn.edge_name] = sbuf + if self._stream_buffer_rids is not None: + self._stream_buffer_rids.add(request_id) if conn.to_partition == body.request_info.partition_name: lead = body.request_info.stream_lead_items.get(conn.edge_name) if lead: @@ -707,10 +738,13 @@ def _add_new_request(self, body: NewRequest) -> None: self._graph_runtime.ingest_inputs_batch( self._one_rids_edges(request_id, signal_only), can_buffer=True, ) - # process messages that may have came in out-of-order - if body.request_id in self._unprocessed_messages: - self._process_message_list(self._unprocessed_messages[body.request_id]) - del self._unprocessed_messages[body.request_id] + self._process_unprocessed(body.request_id) + + def _process_unprocessed(self, wire_rid: str) -> None: + """Process messages for ``wire_rid`` that came in before its NEW_REQUEST.""" + if wire_rid in self._unprocessed_messages: + self._process_message_list(self._unprocessed_messages[wire_rid]) + del self._unprocessed_messages[wire_rid] def _remove_request(self, body: RemoveRequest) -> None: @@ -782,6 +816,8 @@ def _remove_request(self, body: RemoveRequest) -> None: self._pending_drains.discard(body.request_id) self._reads_done_sent.discard(body.request_id) self.engine_manager.remove_request(request_id) + if self._stream_buffer_rids: + self._stream_buffer_rids.discard(request_id) self.request_state.remove_request(request_id) self.tensor_manager.force_cleanup_request(request_id) self.profile_info.pop_request(request_id) @@ -1160,6 +1196,12 @@ def _mark_stream_ingested( if sbuf is not None: sbuf.ingested_chunk = streaming_edge.chunk + def _streams_possible(self) -> bool: + """False when no live request owns a stream buffer: the per-row stream + polls of a speculation build then have nothing to find.""" + rids = getattr(self, "_stream_buffer_rids", None) + return rids is None or bool(rids) + def _stream_chunks_for( self, request_id: int, node_name: str, inputs: NameToTensorList ) -> dict[str, StreamChunkInfo] | None: @@ -1327,11 +1369,12 @@ def _build_input_metadata_and_fwd_info( ) -> tuple[dict[str, InputMetadata], dict[str, CurrentForwardPassInfo]]: meta = {} fwd_info = {} + streams = self._streams_possible() for rid, inputs in per_request_inputs.items(): req_info = self.request_state.get_fwd_info(rid, batch_partition) fwd_info[rid] = req_info stream_chunks = {} - if chunks := self._stream_chunks_for(rid, node_name, inputs): + if streams and (chunks := self._stream_chunks_for(rid, node_name, inputs)): stream_chunks = chunks meta[rid] = InputMetadata( stream_chunks=stream_chunks, @@ -1805,6 +1848,12 @@ def _execute_on_gpu_thread( from mstar.utils.profiler import range_pop, range_push engine = self.engine_manager.get_engine(batch.node_name) + # Every output of this step stays on the device: the engine keeps the + # rows and skips the per-request dicts and tail (cached per node and + # walk, a dict read here). + node_batch.rows_only = self._all_on_device( + engine, batch.node_name, batch.graph_walk, batch.output_signals, + ) if logger.isEnabledFor(logging.DEBUG): # test/waypoint/serve_rollout.py parses this line (wire ids, logged # before prepare_inputs runs) to recover the DiT schedule. @@ -1859,6 +1908,7 @@ def _execute_on_gpu_thread( with self._span("worker.gpu_thread.exec"): outputs = engine.exec_and_postprocess(node_batch) if execution_stream is not None: + self._stage_host_rows(outputs) event = torch.Event() event.record(execution_stream) node_batch.completion_event = event @@ -2141,7 +2191,7 @@ def _assemble_speculation( per_request_input_metadata = { rid: InputMetadata(stream_chunks=chunks) for rid in request_ids if (chunks := self._stream_chunks_for(rid, spec_node, per_request_inputs[rid])) - } + } if self._streams_possible() else {} spec_batch = ScheduledBatch( node_name=spec_node, graph_walk=pending.graph_walk, @@ -2316,10 +2366,13 @@ def _try_speculate_next( # Polling the StreamBuffers stays on this side: they hold real tensors. polled: list[tuple[int, StreamingEdge]] = [] per_rid_counts: list[int] = [] - for r in candidates: - edges = self._poll_stream_buffers_for_speculation(r, spec_node_name) - polled.extend((r, e) for e in edges) - per_rid_counts.append(len(edges)) + if self._streams_possible(): + for r in candidates: + edges = self._poll_stream_buffers_for_speculation(r, spec_node_name) + polled.extend((r, e) for e in edges) + per_rid_counts.append(len(edges)) + else: + per_rid_counts = [0] * len(candidates) prep = self._graph_runtime.prep_spec_rids(SpeculationPrepInput( spec_node_name=spec_node_name, @@ -2416,15 +2469,30 @@ def _thread_outputs_to_speculative( ): """Splice batch N's outputs into the spec batch's inputs. - Runs on the plan thread as soon as N's outputs are published, so it - must not read a tensor VALUE: it only moves tensor lists and tests - which keys are present. N's kernels may still be running. + Runs on the main thread right after N's future resolves, which is + when the GPU thread has enqueued N, so it must not read a tensor + VALUE: it only moves tensor lists and tests which keys are present. + N's kernels may still be running. """ threaded_continuing: set[int] = set() dropped: set[int] = set() + # A rows-only step kept every output on the device: a rid that ran is + # present, and its loop-back signal arrives empty, the way every + # fallback batch of such a walk does (``device_loopback_signals``). + rows_only = getattr(outputs_N, "rows_only", False) + ran = frozenset(outputs_N.row_request_ids or ()) if rows_only else None for rid in list(speculation.node_batch.request_ids): if rid not in speculation.continuing_rids: continue # fresh rid — inputs already gathered. + if ran is not None: + if rid in ran: + inputs = speculation.node_batch.per_request_input_tensors[rid] + for input_name, _ in speculation.consumed_edges: + inputs[input_name] = [] + threaded_continuing.add(rid) + else: + dropped.add(rid) + continue rid_outputs = outputs_N.get(rid, {}) ok = True for input_name, _ in speculation.consumed_edges: @@ -2873,7 +2941,9 @@ def _pp_stage(name: str) -> None: range_pop(synchronize=False) return - _pp_stage("check_stop") + # the stop check itself is the direct window above; this stage is its + # tail (prefix chains, stream-terminated loops, failed rids) + _pp_stage("check_stop_tail") if self.enable_nvtx: range_pop(synchronize=False) range_push("worker.postprocess.stop_loops", synchronize=False) @@ -2938,8 +3008,22 @@ def _pp_stage(name: str) -> None: # reuses; safe because the sends below are their last reader. # Row views of one batch clone (a decode step) are described from one # row; anything else goes through the general per-tensor path. + # The client copy first: whether it rides inline decides whether the + # tensors are needed at all. + inline_signal, inline_values = self._inline_emit_values( + engine, batch_N, host_rows, rids, signals, + ) stored = None - if outputs.row_views and outputs.row_request_ids is not None: + if inline_signal is not None and self._all_on_device( + engine, batch_N.node_name, batch_N.graph_walk, signals, + ): + # Device loop-back: the node keeps every routed signal's value on + # the device (the sampler's slot master) and the client copy rides + # inline, so nothing is stored, held or freed. The route still + # runs per request for the loop counters and the fallbacks; it + # carries zero tensors per (request, signal). + stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) + elif outputs.row_views and outputs.row_request_ids is not None: stored = self.tensor_manager.store_row_outputs_batch( rids, outputs, signals, outputs.row_views, outputs.row_request_ids, node_name=batch_N.node_name, @@ -2963,7 +3047,8 @@ def _pp_stage(name: str) -> None: # one per tensor -- with a Rust bookkeeper each is a boundary crossing, # and a 128-request batch has hundreds of them. The count is uniform, # so it crosses as a scalar rather than a list built per batch. - self.tensor_manager.increment_ref_batch_uniform(flat_uuids, 1) + if flat_uuids: + self.tensor_manager.increment_ref_batch_uniform(flat_uuids, 1) if self._phase_period: self._phase_record( "worker.postprocess.store_tensors", @@ -2972,9 +3057,6 @@ def _pp_stage(name: str) -> None: # The graph runtime's own share of postprocess: the routing call. _t_route = _time.perf_counter() if self._phase_period else 0.0 - inline_signal, inline_values = self._inline_emit_values( - engine, batch_N, host_rows, rids, signals, - ) route_output = self._graph_runtime.complete_and_route_batch( RouteInput( partition=batch_N.partition, @@ -3009,8 +3091,7 @@ def _pp_stage(name: str) -> None: if self.enable_nvtx: range_pop(synchronize=False) range_push("worker.postprocess.register_outputs", synchronize=False) - with self._span("worker.postprocess.register_outputs"): - self._register_outputs(route_output) + self._register_outputs(route_output) _pp_stage("register_outputs") # send outputs @@ -3135,6 +3216,21 @@ def _pp_stage(name: str) -> None: if self.enable_nvtx: range_pop(synchronize=False) + def _all_on_device(self, engine, node_name: str, graph_walk: str, signals) -> bool: + """Whether every output signal of this step is one the node keeps on + the device (cached per node and walk: it depends on the node's + resources, not on the step).""" + # without inline emit, a rows-only step has no path for the client's tokens + if not signals or not self._graph_runtime.supports_inline_emit: + return False + key = (node_name, graph_walk) + on_device = self._device_loopback.get(key) + if on_device is None: + on_device = self._device_loopback[key] = frozenset( + engine.device_loopback_signals(node_name, graph_walk) + ) + return all(s in on_device for s in signals) + def _inline_emit_values( self, engine, batch_N: PendingBatch, host_rows: "HostRows | None", rids: list[int], signals, @@ -3166,6 +3262,34 @@ def _inline_emit_values( return None, [] return None, [] + def _stage_host_rows(self, outputs: "BatchedModelOutput") -> None: + """Start the device-to-host copy of the step's row-addressed stop + buffers here on the gpu thread, on the execution stream right behind + the clone that made them, so the completion event recorded next + covers it. The main thread's stop check then finds the rows landed + once it has waited on that event: no side stream, no second sync, and + one fewer GIL release on its critical path. + + A small ring of pinned buffers: step N's rows are read by the main + thread while this thread already stages N+1's (the pipeline is one + step deep, the ring leaves slack). + """ + buffers = outputs.check_stop_buffers + if not buffers or outputs.row_request_ids is None: + return + slot = self._host_rows_slot = (self._host_rows_slot + 1) % _HOST_ROWS_RING + host: dict[str, torch.Tensor] = {} + for index, (name, tensor) in enumerate(buffers.items()): + if not (torch.is_tensor(tensor) and tensor.is_cuda): + host[name] = tensor + continue + buf = self._get_pinned_d2h_buffer( + f"host_rows{slot}", tensor.shape, tensor.dtype, index, + ) + buf.copy_(tensor, non_blocking=True) + host[name] = buf + outputs.host_rows = HostRows(tuple(outputs.row_request_ids), host) + def _get_pinned_d2h_buffer( self, purpose: str, @@ -3265,6 +3389,11 @@ def _prematerialize_for_check_stop( list(outputs.row_request_ids) if outputs.row_request_ids is not None else request_ids ) + landed = outputs.host_rows + if landed is not None and completion_event is not None: + # staged by the gpu thread behind the step; the caller has waited + # on the completion event, so the rows are here + return _LazyPerRid(landed.buffers, list(landed.request_ids)), landed if not torch.cuda.is_available() or completion_event is None: if outputs.check_stop_buffers is not None and row_rids is not None: # host tensors already (a CPU device): only the re-keying @@ -3498,6 +3627,99 @@ def _fail_requests(self, errors: dict[int, str]) -> None: # Note: we do not cleanup the request right now; we wait for the conductor # to officially send a removal message + def _build_speculation( + self, pending: PendingBatch | None, consecutive: int, + fair: "_SpecFairness", + ) -> tuple[Speculation | None, tuple[str, str] | None]: + """Speculate the step after ``pending`` (N+1 from N), or a + yield-away batch when fairness says so. Returns the speculation + and the (node, walk) to exclude on the non-speculative path.""" + speculation = None + yield_away_from_target = None + batch = None + # Build nothing on a refusal; the caller's admit fence is what + # makes it knowable here. Ordinary speculation would be dropped anyway, + # but a yield-away batch survives ``_maybe_clear_spec`` and is + # already broadcast — so the eviction below would land after its + # head went out, and the follower would admit it a delta behind. + # Skipping keeps that eviction in the gap between steps. + # + # The marker below stays outside this: a follower settles N on + # exactly one of {head, marker}, and it moves no pages. + admit_refused = ( + pending is not None + and pending.node_batch.admit_error is not None + ) + + if pending is not None and self._can_speculate(pending.batch): + # Fairness check (peek-based, replaces the old iter- + # counter cap): only break the spec chain when there's + # another (node, walk) actually ready to schedule on + # this worker. On single-walk workers (Orpheus LLM, + # Orpheus SNAC) this returns False and we always speculate. + other_ready = ( + fair.peek + and consecutive >= 1 + and self.scheduler.has_ready_excluding( + self.request_state, + (pending.node_name, pending.graph_walk), + ) + ) + must_yield_away = consecutive >= fair.max_consecutive or other_ready + if not must_yield_away and not admit_refused: + if self.enable_nvtx: + range_push("worker.speculate", synchronize=False) + _t0 = _time.perf_counter() if self._phase_period else 0.0 + speculation = self._try_speculate_next(pending) + if self._phase_period: + self._phase_record("speculate", _time.perf_counter() - _t0) + if self.enable_nvtx: + range_pop(synchronize=False) + if speculation is not None: + # Broadcast the head now, during forward N, so + # followers build it during theirs (-1: not parallel). + speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( + speculation.node_batch, + speculative=True, spec_from_seq=pending.tp_seq, + ) + if self._tp_lead_needs_marker(pending, speculation): + self._broadcast_tp_nospec(pending) + # ``yield_away_from_target`` stays None when the admit + # refused, so the non-speculative path below schedules + # without the fairness exclusion — it should be free to pick + # up the batch ``_handle_admit_failure`` just pushed back. + if speculation is None and not admit_refused: + yield_away_from_target = ( + pending.node_name, + pending.graph_walk, + ) if must_yield_away else None + with self._span("worker.schedule_yield_away"): + batch = self.scheduler.get_next_batch( + self.request_state, + exclude_target=yield_away_from_target, + ) + if batch is not None: + node_batch = self._build_executing_batch(batch) + batch_partition = self.request_state.get_partition_for_node(batch.node_name) + logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") + speculation = Speculation( + scheduled_batch=batch, + node_batch=node_batch, + consumed_edges=set(), + continuing_rids=set(), # n/a + partition=batch_partition, + is_new_iter=False, + is_same_node=False, + is_yield_away=True + ) + + # A leader stamps the seq it sends; a follower's + # batch keeps the one it came off the FIFO with. + ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) + speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq + + return speculation, yield_away_from_target + def run(self) -> None: switch_interval = os.environ.get("MSTAR_PY_SWITCH_INTERVAL_SEC", "") if switch_interval: @@ -3620,9 +3842,22 @@ def run(self) -> None: spec_peek_for_fairness = ( os.environ.get("MSTAR_SPEC_PEEK_FOR_FAIRNESS", "1") == "1" ) + fair = _SpecFairness( + peek=spec_peek_for_fairness, max_consecutive=max_consecutive_spec, + ) consecutive_spec_steps = 0 yield_away_from_target: tuple[str, str] | None = None + def _arm_speculation(spec: Speculation) -> None: + # Hand N+1 to the plan thread NOW. It waits on N's + # commit event, then reserves a slot and pre-plans — + # while the main thread sits in await_gpu with the GIL + # released and N's kernels are still running. + if plan_executor is not None: + spec.plan_future = plan_executor.submit( + self._preplan_spec, pending, spec, + ) + def _set_pending(p: PendingBatch): nonlocal pending pending = p @@ -3727,103 +3962,9 @@ def _phase_flush() -> None: # Only when (a) there's a pending step and (b) it's AR-engine. # For non-AR or non-loop-body steps, falls through to the # non-speculative path below (drain, then schedule). - speculation = None - yield_away_from_target = None - - # Build nothing on a refusal; the fence above is what makes it - # knowable here. Ordinary speculation would be dropped anyway, - # but a yield-away batch survives ``_maybe_clear_spec`` and is - # already broadcast — so the eviction below would land after its - # head went out, and the follower would admit it a delta behind. - # Skipping keeps that eviction in the gap between steps. - # - # The marker below stays outside this: a follower settles N on - # exactly one of {head, marker}, and it moves no pages. - admit_refused = ( - pending is not None - and pending.node_batch.admit_error is not None + speculation, yield_away_from_target = self._build_speculation( + pending, consecutive_spec_steps, fair, ) - - if pending is not None and self._can_speculate(pending.batch): - # Fairness check (peek-based, replaces the old iter- - # counter cap): only break the spec chain when there's - # another (node, walk) actually ready to schedule on - # this worker. On single-walk workers (Orpheus LLM, - # Orpheus SNAC) this returns False and we always speculate. - must_yield_for_fairness = ( - spec_peek_for_fairness - and consecutive_spec_steps >= 1 - and self.scheduler.has_ready_excluding( - self.request_state, - (pending.node_name, pending.graph_walk), - ) - ) - must_yield_away = ( - consecutive_spec_steps >= max_consecutive_spec - or must_yield_for_fairness - ) - if not must_yield_away and not admit_refused: - if self.enable_nvtx: - range_push("worker.speculate", synchronize=False) - _t0 = _time.perf_counter() if phase_period else 0.0 - speculation = self._try_speculate_next(pending) - if phase_period: - _phase_record("speculate", _time.perf_counter() - _t0) - if self.enable_nvtx: - range_pop(synchronize=False) - if speculation is not None: - # Broadcast the head now, during forward N, so - # followers build it during theirs (-1: not parallel). - speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( - speculation.node_batch, - speculative=True, spec_from_seq=pending.tp_seq, - ) - if self._tp_lead_needs_marker(pending, speculation): - self._broadcast_tp_nospec(pending) - # ``yield_away_from_target`` stays None when the admit - # refused, so the non-speculative path below schedules - # without the fairness exclusion — it should be free to pick - # up the batch ``_handle_admit_failure`` just pushed back. - if speculation is None and not admit_refused: - yield_away_from_target = ( - pending.node_name, - pending.graph_walk, - ) if must_yield_away else None - with self._span("worker.schedule_yield_away"): - batch = self.scheduler.get_next_batch( - self.request_state, - exclude_target=yield_away_from_target, - ) - if batch is not None: - node_batch = self._build_executing_batch(batch) - batch_partition = self.request_state.get_partition_for_node(batch.node_name) - logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") - speculation = Speculation( - scheduled_batch=batch, - node_batch=node_batch, - consumed_edges=set(), - continuing_rids=set(), # n/a - partition=batch_partition, - is_new_iter=False, - is_same_node=False, - is_yield_away=True - ) - - # A leader stamps the seq it sends; a follower's - # batch keeps the one it came off the FIFO with. - ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) - speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq - - def _arm_speculation(spec: Speculation) -> None: - # Hand N+1 to the plan thread NOW. It waits on N's - # commit event, then reserves a slot and pre-plans — - # while the main thread sits in await_gpu with the GIL - # released and N's kernels are still running. - if plan_executor is not None: - spec.plan_future = plan_executor.submit( - self._preplan_spec, pending, spec, - ) - if speculation is not None: _arm_speculation(speculation) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 60b471ee67..e8816f0d7d 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -62,6 +62,15 @@ test/modular/test_publish_snapshot_batch.py test/modular/test_sampler_for_cache.py test/modular/test_send_prep_entries.py test/modular/test_kv_prefix_batch.py +test/modular/test_data_worker_wakeup.py +test/modular/test_sampler_last_token.py +test/modular/test_qwen3_5_device_loopback.py +test/modular/test_worker_device_loopback.py +test/modular/test_gc_freeze.py +test/modular/test_qwen3_5_bucket_knobs.py +test/modular/test_fast_tokenize.py test/modular/test_wire_codec.py test/modular/test_wire_security.py test/modular/test_model_kwargs_tensor_transport.py +test/modular/test_engine_rows_only_tail.py +test/modular/test_ingest_rollback.py diff --git a/test/modular/test_admit_failure_handling.py b/test/modular/test_admit_failure_handling.py index 97fcdedd34..8eb987489e 100644 --- a/test/modular/test_admit_failure_handling.py +++ b/test/modular/test_admit_failure_handling.py @@ -74,7 +74,7 @@ def postprocess_batch(self, batch, outputs): def _batch(): - return SimpleNamespace(request_ids=["r0", "r1"], admit_error=None) + return SimpleNamespace(request_ids=["r0", "r1"], admit_error=None, rows_only=False) @pytest.mark.parametrize( diff --git a/test/modular/test_api_result_delivery.py b/test/modular/test_api_result_delivery.py index d90109f3fa..a0f3786651 100644 --- a/test/modular/test_api_result_delivery.py +++ b/test/modular/test_api_result_delivery.py @@ -303,3 +303,50 @@ def call_soon_threadsafe(self, fn, *args): assert len(loop.calls) == 1 assert all(req.wake.is_set() for req in reqs[:3]) assert not reqs[3].wake.is_set() + + +def test_chunks_waiting_together_come_as_one_batch(): + """Everything the message thread appended since the last wake is drained + as one list, so a streaming handler can send it as one body frame; a + chunk that lands later comes on its own, and the finish ends the stream + without an empty batch.""" + pw = _StubPreprocessWorker() + server = _api_server_stub(pw) + req = _pending_request() + for i in range(3): + req.chunks.append(ResultChunk(request_id="r5", modality="text", data=f"t{i}".encode())) + server.pending_requests["r5"] = req + + async def _collect(): + batches = server.iter_result_chunk_batches("r5") + got = [await batches.__anext__()] + with server.request_lock: + req.chunks.append(ResultChunk(request_id="r5", modality="text", data=b"t3")) + req.notify() + got.append(await batches.__anext__()) + with server.request_lock: + req.event.set() + req.notify() + async for batch in batches: + got.append(batch) + return got + + got = asyncio.run(_collect()) + assert [[c.data for c in b] for b in got] == [[b"t0", b"t1", b"t2"], [b"t3"]] + assert "r5" not in server.pending_requests + + +def test_error_tail_is_its_own_last_batch(): + pw = _StubPreprocessWorker() + server = _api_server_stub(pw) + req = _pending_request() + req.chunks.append(ResultChunk(request_id="r6", modality="text", data=b"partial")) + req.error = "result delivery timed out; response is incomplete" + req.error_status = 500 + req.event.set() + server.pending_requests["r6"] = req + + async def _collect(): + return [[c.modality for c in b] async for b in server.iter_result_chunk_batches("r6")] + + assert asyncio.run(_collect()) == [["text"], ["error"]] diff --git a/test/modular/test_batched_model_output.py b/test/modular/test_batched_model_output.py index 37542448f9..ba35142600 100644 --- a/test/modular/test_batched_model_output.py +++ b/test/modular/test_batched_model_output.py @@ -199,3 +199,62 @@ def test_merging_two_row_addressed_outputs_drops_the_buffers(): )) assert merged.check_stop_buffers is None and merged.row_request_ids is None assert merged.get_check_stop_input()["b"]["new_token"][0].item() == 2 + + +class _StagingWorker(_StubWorker): + """The gpu-thread side of the host rows: stage behind the step, hand + the landed rows to the stop check without a sync of its own.""" + + _host_rows_slot = 0 + _d2h_stream = None + device = torch.device("cpu") + _stage_host_rows = Worker._stage_host_rows + _prematerialize_for_check_stop = Worker._prematerialize_for_check_stop + + +def test_landed_host_rows_skip_the_side_stream_copy(): + """Rows staged by the gpu thread are what the stop check reads: no + side stream, no second copy, row i still request i.""" + rids = ("a", "b", "c") + tokens = torch.tensor([5, 6, 7]) + out = BatchedModelOutput( + check_stop_buffers={"new_token": tokens}, row_request_ids=rids, + ) + worker = _StagingWorker() + worker._stage_host_rows(out) + assert out.host_rows is not None and out.host_rows.request_ids == rids + # a host tensor is passed through as is + assert out.host_rows.buffers["new_token"] is tokens + + per_rid, host_rows = worker._prematerialize_for_check_stop( + out, completion_event=object(), request_ids=list(rids), + ) + assert host_rows is out.host_rows + assert per_rid["b"]["new_token"][0].item() == 6 + assert worker._d2h_stream is None, "no side stream was needed" + + +def test_nothing_to_stage_without_row_buffers(): + out = BatchedModelOutput(per_rid_outputs={"a": {"new_token": [torch.tensor([1])]}}) + worker = _StagingWorker() + worker._stage_host_rows(out) + assert out.host_rows is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_staged_rows_land_behind_the_stream_and_rotate_buffers(): + rids = tuple(f"r{i}" for i in range(6)) + worker = _StagingWorker() + seen = set() + for step in range(3): + tokens = torch.arange(100 * step, 100 * step + 6, device="cuda") + out = BatchedModelOutput( + check_stop_buffers={"new_token": tokens}, row_request_ids=rids, + ) + worker._stage_host_rows(out) + buf = out.host_rows.buffers["new_token"] + assert not buf.is_cuda and buf.is_pinned() + torch.cuda.synchronize() + assert buf.tolist() == list(range(100 * step, 100 * step + 6)) + seen.add(buf.data_ptr()) + assert len(seen) == 3, "consecutive steps use different pinned buffers" diff --git a/test/modular/test_collect_outputs_fast_path.py b/test/modular/test_collect_outputs_fast_path.py index db3a049aa6..afb9bc8328 100644 --- a/test/modular/test_collect_outputs_fast_path.py +++ b/test/modular/test_collect_outputs_fast_path.py @@ -118,3 +118,39 @@ def test_stop_buffer_shares_the_row_clone(): ) stop2 = other.clone_check_stop_buffers(reuse=other.clone_row_outputs(2)) assert torch.equal(stop2["new_token"], torch.ones(2, dtype=torch.int64)) + + +def test_rows_only_keeps_the_views_and_builds_no_per_rid_dicts(): + """A step whose outputs all stay on the device: the row views and the + stop buffer are what its readers use, so no per-rid dict is built.""" + rids = [1, 2, 3, 4] + raw = _raw(bs=4) + eng = Engine.__new__(Engine) + outputs = {} + row_clones, row_views = eng._merge_per_rid( + outputs, raw, rids, rids, _Sub(), {}, rows_only=True, + ) + assert outputs == {} + assert set(row_views) == {"new_token", "aux"} and len(row_views["new_token"]) == 4 + assert row_clones["new_token"].tolist() == [100, 101, 102, 103] + # the ordinary path still fills them + outputs = {} + eng._merge_per_rid(outputs, _raw(bs=4), rids, rids, _Sub(), {}) + assert set(outputs) == set(rids) + + +def test_rows_only_collect_takes_no_clone(): + """On a rows-only step the gpu thread copies the rows to the host right + behind the step, so the graph's own output buffer is handed over as is.""" + from types import SimpleNamespace + + rids = [1, 2, 3, 4] + raw = _raw(bs=4) + eng = Engine.__new__(Engine) + out = eng._collect_outputs( + SimpleNamespace(submodule=_Sub(), cuda_graph_runner=None), None, raw, + [], {}, request_ids=rids, step_request_ids=tuple(rids), rows_only=True, + ) + assert out.rows_only and out.per_rid_outputs == {} and out.row_views is None + assert out.check_stop_buffers["new_token"] is raw.check_stop_buffers["new_token"] + assert out.row_request_ids == (1, 2, 3, 4) diff --git a/test/modular/test_data_worker_wakeup.py b/test/modular/test_data_worker_wakeup.py new file mode 100644 index 0000000000..4b4c9e68ba --- /dev/null +++ b/test/modular/test_data_worker_wakeup.py @@ -0,0 +1,170 @@ +"""The data worker thread's idle wait: it blocks on the communicator's poll +and is woken by a producer's signal, a finished tensor read or shutdown, +instead of sleeping 1 ms and sweeping its queues.""" +from __future__ import annotations + +import queue +import select +import threading +import time +from concurrent.futures import Future + +from mstar.api_server.data_worker import InlineToken, PreprocessWorkerThread +from mstar.communication.event import EventWakeup +from mstar.graph.loop_indices import NestedLoopIndices + + +class _Model: + def postprocess(self, output, modality, request_kwargs=None): + return f"<{int(output[0])}>".encode() + + +class _PollingComm: + """A communicator whose idle wait really blocks on the wakeup's fd.""" + + def __init__(self): + self.event = None + self.waits = 0 + + def register_event_for_poll(self, event): + self.event = event + + def wait_for_work(self, timeout_ms): + self.waits += 1 + ready, _, _ = select.select([self.event.fd], [], [], timeout_ms / 1000.0) + if ready: + self.event.drain() + + def get_all_new_messages(self): + return [] + + def send(self, entity, msg): + pass + + +class _BareComm: + def get_all_new_messages(self): + return [] + + def send(self, entity, msg): + pass + + +class _TM: + def __init__(self): + self.futures: list[Future] = [] + + def start_read_tensors(self, request_id, graph_edges, graph_walk=None): + return list(self.futures) + + def get_ready_tensors(self): + return {} + + def has_inflight_reads(self, request_id): + return False + + def cleanup_request(self, request_id): + pass + + def force_cleanup_request(self, request_id): + pass + + def ack_unread_tensors(self, request_id, graph_edges): + pass + + +def _thread(comm, wakeup, idle_wait_ms=2000): + stop = threading.Event() + worker = PreprocessWorkerThread( + in_queue=queue.Queue(), result_tensor_queue=queue.Queue(), out_queue=queue.Queue(), + profile_queue=queue.Queue(), cleanup_request_queue=queue.Queue(), + abort_request_queue=queue.Queue(), reads_done_queue=queue.Queue(), + discard_tensor_queue=queue.Queue(), stop_event=stop, communicator=comm, + tensor_manager=_TM(), model=_Model(), wakeup=wakeup, + ) + worker.idle_wait_ms = idle_wait_ms + t = threading.Thread(target=worker.run, daemon=True) + t.start() + return worker, stop, t + + +def _loop(i): + return NestedLoopIndices( + loop_name_order=["decode_loop"], loop_indices={"decode_loop": i}, wg_fwd_pass_idx=i + 1, + ) + + +def _wait_until(pred, timeout_s): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(0.001) + return pred() + + +def test_a_signalled_producer_ends_the_idle_wait_at_once(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=2000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) # the thread is parked + t0 = time.monotonic() + worker.result_tensor_queue.put(InlineToken("req", 7, _loop(0), "text")) + wakeup.signal() + chunk = worker.out_queue.get(timeout=1.0) + assert time.monotonic() - t0 < 0.5 # far below the 2 s backstop + assert chunk.request_id == "req" and chunk.data == b"<7>" + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + + +def test_shutdown_wakes_a_parked_thread(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=5000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) + t0 = time.monotonic() + stop.set() + wakeup.signal() + t.join(timeout=2.0) + assert not t.is_alive() and time.monotonic() - t0 < 1.0 + + +def test_an_idle_thread_does_not_spin(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=50) + time.sleep(0.3) + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + # ~6 waits of 50 ms in 0.3 s, not ~300 sweeps of 1 ms + assert comm.waits <= 12 + + +def test_a_finished_read_wakes_the_thread(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=2000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) + fut: Future = Future() + worker.tensor_manager.futures = [fut] + worker.wakeup.register_futures(worker.tensor_manager.start_read_tensors("r", [])) + waits = comm.waits + time.sleep(0.05) + assert comm.waits == waits # still parked + fut.set_result(None) # the reader thread finishing wakes the poll + assert _wait_until(lambda: comm.waits > waits, 1.0) + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + + +def test_without_a_wakeup_the_thread_keeps_the_sleep(): + worker, stop, t = _thread(_BareComm(), None) + worker.result_tensor_queue.put(InlineToken("req", 3, _loop(0), "text")) + chunk = worker.out_queue.get(timeout=1.0) + assert chunk.data == b"<3>" + stop.set() + t.join(timeout=1.0) + assert not t.is_alive() diff --git a/test/modular/test_engine_rows_only_tail.py b/test/modular/test_engine_rows_only_tail.py new file mode 100644 index 0000000000..2663183f6a --- /dev/null +++ b/test/modular/test_engine_rows_only_tail.py @@ -0,0 +1,84 @@ +"""A rows-only step runs only a submodule's batch tail, and only when it has +one; a step on the unbatched forward (per request or a single row) is not +rows-only.""" +import threading +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.engine import Engine # noqa: E402 +from mstar.model.submodule_base import BatchedModelOutput, NodeSubmodule # noqa: E402 + + +class _Plain(NodeSubmodule): + def prepare_inputs(self, *a, **kw): + return None + + def forward(self, *a, **kw): + return None + + +class _WithTail(_Plain): + def batch_postprocess(self, outputs, request_infos, inputs): + self.seen = (outputs, request_infos, inputs) + + +def _engine(sub): + eng = Engine.__new__(Engine) + eng._submodules = {"llm": SimpleNamespace(submodule=sub)} + eng._loopback_steps = {} + eng.per_request_tails = [] + eng.postprocess_batch = lambda batch, outputs: eng.per_request_tails.append(batch) + return eng + + +def _batch(rows_only): + return SimpleNamespace( + node_name="llm", admit_error=None, rows_only=rows_only, + per_request_info_wrapped={1: object()}, inputs=[object()], + ) + + +@pytest.mark.parametrize("sub_cls", [_Plain, _WithTail]) +def test_rows_only_runs_the_batch_tail_only_when_overridden(sub_cls): + sub = sub_cls.__new__(sub_cls) + eng = _engine(sub) + out = BatchedModelOutput() + eng.exec = lambda batch: out + batch = _batch(rows_only=True) + assert eng.exec_and_postprocess(batch) is out + assert eng.per_request_tails == [] + if sub_cls is _WithTail: + assert sub.seen == (out, batch.per_request_info_wrapped, batch.inputs) + eng.exec_and_postprocess(_batch(rows_only=False)) + assert len(eng.per_request_tails) == 1 + + +@pytest.mark.parametrize("rids", [[1, 2], [1]]) +def test_an_unbatched_step_clears_rows_only(rids): + eng = _engine(_Plain.__new__(_Plain)) + eng._enable_nvtx = False + eng._enable_profile = False + eng._device = SimpleNamespace(type="cpu") + eng._autocast_dtype_for = lambda sub: None + eng.reserve_replay_slot = lambda batch: None + seen = {} + + def _per_request(batch): + seen["rows_only"] = batch.rows_only + return BatchedModelOutput() + + eng._exec_per_request = _per_request + eng._exec_single = _per_request + batch = SimpleNamespace( + node_name="llm", request_ids=rids, rows_only=True, running_batched=False, + step_context=SimpleNamespace(slot_lease=None, graph_walk="decode"), + exec_timings=SimpleNamespace(start=None), admit_error=None, + outputs_ready=threading.Event(), release_waiters=lambda: None, + preplanned_rids=None, preplan_event=None, + ) + eng.exec(batch) + assert seen["rows_only"] is False and batch.rows_only is False + assert eng._loopback_steps == {} diff --git a/test/modular/test_engine_uniform_inputs.py b/test/modular/test_engine_uniform_inputs.py new file mode 100644 index 0000000000..9cc184caf1 --- /dev/null +++ b/test/modular/test_engine_uniform_inputs.py @@ -0,0 +1,102 @@ +"""Engine.prepare_inputs: a walk whose rows all get the same inputs is +prepared with one call, the object shared by the rows; any other walk still +prepares row by row.""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.engine import Engine # noqa: E402 + + +class _Batch: + def __init__(self, rids, walk="decode"): + self.node_name = "llm" + self.request_ids = list(rids) + self.step_context = SimpleNamespace(graph_walk=walk) + self.per_request_info_wrapped = {r: SimpleNamespace() for r in rids} + self.per_request_input_tensors = {} + self.per_request_input_metadata = {} + self.final_stream_rids = set() + self.skipped_rids = set() + self.failed_requests = {} + self.inputs = None + self.running_batched = None + + def register_prepare_batch(self, inputs): + self.inputs = inputs + + def drop_rids(self, rids): + self.request_ids = [r for r in self.request_ids if r not in rids] + + def register_failure(self, rid, error): + self.failed_requests[rid] = str(error) + + +class _Sub: + def __init__(self, uniform=None): + self.uniform = uniform + self.calls = [] + self.row_loopback = [] + self.uniform_loopback = None + + def uniform_row_inputs(self, graph_walk, has_loopback=False): + self.uniform_loopback = has_loopback + return self.uniform + + def prepare_inputs(self, **kw): + self.calls.append(kw["graph_walk"]) + self.row_loopback.append(kw["has_loopback"]) + return object() + + def can_batch(self, batch, model_inputs): + return len(model_inputs) > 0 + + +def _engine(sub): + eng = Engine.__new__(Engine) + eng._submodules = {"llm": SimpleNamespace(submodule=sub, resources={})} + eng._enable_nvtx = False + eng._loopback_steps = {} + return eng + + +def test_a_uniform_walk_is_prepared_with_one_object(): + row = object() + sub = _Sub(uniform=row) + batch = _Batch([1, 2, 3]) + _engine(sub).prepare_inputs(batch) + assert batch.inputs == [row, row, row] and all(x is row for x in batch.inputs) + assert sub.calls == [] and batch.running_batched is True + + +def test_other_walks_still_prepare_row_by_row(): + sub = _Sub(uniform=None) + batch = _Batch([1, 2]) + _engine(sub).prepare_inputs(batch) + assert sub.calls == ["decode", "decode"] and len(batch.inputs) == 2 + + +def test_loopback_is_per_row_and_tracks_each_request_s_last_step(): + sub = _Sub(uniform=None) + eng = _engine(sub) + ran = _Batch([1, 2]) + ran.rows_only = True + eng._record_loopback(ran) + eng.prepare_inputs(_Batch([1, 2, 3])) + # 3 never ran rows-only here (e.g. its prefill ran on another rank) + assert sub.uniform_loopback is False and sub.row_loopback == [True, True, False] + eng.prepare_inputs(_Batch([1, 2])) + assert sub.uniform_loopback is True + # a step that is not rows-only clears the record; so does another walk's + other = _Batch([2]) + other.rows_only = False + eng._record_loopback(other) + sub.row_loopback.clear() + eng.prepare_inputs(_Batch([1, 2])) + assert sub.row_loopback == [True, False] + sub.row_loopback.clear() + eng.prepare_inputs(_Batch([1], walk="prefill")) + assert sub.row_loopback == [False] diff --git a/test/modular/test_fast_tokenize.py b/test/modular/test_fast_tokenize.py new file mode 100644 index 0000000000..1e4da0f51d --- /dev/null +++ b/test/modular/test_fast_tokenize.py @@ -0,0 +1,89 @@ +"""Chunked tokenization: cuts only at the start of a whitespace run, chunks +concatenate to the text, and the ids are identical to one encode (checked +against the real Qwen tokenizer when it is in the local cache).""" +import json +import os + +import pytest + +from mstar.utils.fast_tokenize import cut_points, encode_ids, split_at + + +def test_cuts_start_whitespace_runs_and_chunks_rebuild_the_text(): + text = ("word " * 300 + "a b \n\n c. \n d" + " tail " * 300) * 4 + cuts = cut_points(text, chunk_chars=512, max_chunks=8) + assert 1 <= len(cuts) <= 7 + for c in cuts: + assert text[c] == " " and not text[c - 1].isspace() + assert "".join(split_at(text, cuts)) == text + assert cut_points("short text", chunk_chars=4096) == [] + assert cut_points("x" * 10000, chunk_chars=4096) == [] # no space: no safe cut + + +def test_a_short_text_or_a_slow_tokenizer_takes_the_plain_path(): + class _Slow: + def __call__(self, text): + class _R: + input_ids = [len(text)] + return _R() + + assert encode_ids(_Slow(), "x" * 10000) == [10000] + + +def _tokenizer(): + cache = os.environ.get("MSTAR_TEST_HF_HOME") + if not cache: + pytest.skip("set MSTAR_TEST_HF_HOME to a cache holding Qwen/Qwen3.5-0.8B for the identity check") + os.environ.setdefault("HF_HOME", cache) + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B") + + +def test_chunked_ids_are_identical_to_one_encode_on_the_benchmark_prompts(): + tok = _tokenizer() + prompts = os.environ.get("MSTAR_TEST_PROMPTS") + texts = [] + if prompts and os.path.exists(prompts): + d = json.load(open(prompts)) + for key in ("long8k", "long", "short"): + texts += [x if isinstance(x, str) else x.get("prompt", "") for x in d.get(key, [])] + texts += [ + "a b \n\n c. \n d" * 2000, ("<|im_start|>user\nhello world<|im_end|>\n" * 400), + "x,y;z!! \n\nq" * 1500, "tab\tsep arated\r\nlines " * 1200, "数字 123 456 文本 " * 1000, + ] + for t in texts: + whole = tok(t).input_ids + assert encode_ids(tok, t, chunk_chars=2048, max_chunks=8) == whole + assert encode_ids(tok, t, chunk_chars=512, max_chunks=16) == whole + + +def test_chunked_encode_turns_the_crate_parallelism_on(monkeypatch): + """The Rust backend only runs a batch in parallel while + TOKENIZERS_PARALLELISM reads true; a serving environment that exported + it false gets it flipped before the batch, unless the knob says no.""" + from mstar.utils.fast_tokenize import encode_ids + + class _Enc: + def __init__(self, ids): + self.ids = ids + + class _Backend: + def encode_batch(self, parts, add_special_tokens=False): + return [_Enc([len(p)]) for p in parts] + + class _Tok: + backend_tokenizer = _Backend() + + def __call__(self, text): + raise AssertionError("the chunked path was expected") + + text = "word " * 4000 + monkeypatch.setenv("TOKENIZERS_PARALLELISM", "false") + monkeypatch.delenv("MSTAR_TOKENIZE_PARALLEL", raising=False) + ids = encode_ids(_Tok(), text, chunk_chars=4096, max_chunks=8) + assert sum(ids) == len(text) and os.environ["TOKENIZERS_PARALLELISM"] == "true" + + monkeypatch.setenv("TOKENIZERS_PARALLELISM", "false") + monkeypatch.setenv("MSTAR_TOKENIZE_PARALLEL", "0") + encode_ids(_Tok(), text, chunk_chars=4096, max_chunks=8) + assert os.environ["TOKENIZERS_PARALLELISM"] == "false" diff --git a/test/modular/test_gc_freeze.py b/test/modular/test_gc_freeze.py new file mode 100644 index 0000000000..ff4dce574d --- /dev/null +++ b/test/modular/test_gc_freeze.py @@ -0,0 +1,41 @@ +"""The set-up heap freeze and the data worker's per-frame wake-up.""" +import gc + +from mstar.api_server.data_worker import PreprocessWorker +from mstar.utils.gc_freeze import freeze_after_setup + + +def test_freeze_after_setup_moves_the_heap(): + gc.unfreeze() + keep = [[i] for i in range(100)] # live objects the collector tracks + frozen = freeze_after_setup("test") + try: + assert frozen > 0 and gc.get_freeze_count() == frozen + assert keep[0] == [0] # frozen objects are still alive and usable + finally: + gc.unfreeze() + assert gc.get_freeze_count() == 0 + + +def test_a_frame_of_tokens_wakes_the_thread_once(): + facade = PreprocessWorker.__new__(PreprocessWorker) + facade.output_loop_idxs = {"r1": {}, "r2": {}} + facade.per_request_reading_tensors = {"r1": 0, "r2": 0} + + class _Q: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + facade.result_tensor_input_queue = _Q() + signals = [] + facade._signal = lambda: signals.append(1) + facade.new_result_token("r1", 5, None, "text_inputs", "text", wake=False) + facade.new_result_token("r2", 6, None, "text_inputs", "text", wake=False) + assert len(facade.result_tensor_input_queue.items) == 2 and signals == [] + facade.wake() + assert signals == [1] + facade.new_result_token("r1", 7, None, "text_inputs", "text") + assert signals == [1, 1] diff --git a/test/modular/test_ingest_rollback.py b/test/modular/test_ingest_rollback.py new file mode 100644 index 0000000000..9fc68919b9 --- /dev/null +++ b/test/modular/test_ingest_rollback.py @@ -0,0 +1,75 @@ +"""A request whose ingest raises on one resource is left on none of them, and +the worker fails only that request instead of the batch in flight.""" +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.resources.runner import StepRunner # noqa: E402 +from mstar.worker.worker import Worker # noqa: E402 + + +class _Res: + def __init__(self, name, log, fail=False): + self.name, self.log, self.fail = name, log, fail + + def ingest_request(self, rid, overrides=None): + self.log.append(("ingest", self.name, rid)) + if self.fail: + raise RuntimeError("over capacity") + + def remove_request(self, rid): + self.log.append(("remove", self.name, rid)) + + +def _runner(resources): + runner = StepRunner.__new__(StepRunner) + runner._resources = {r.name: r for r in resources} + runner._order = [r.name for r in resources] + return runner + + +def test_a_raising_ingest_rolls_back_every_resource_it_reached(): + log = [] + runner = _runner([_Res("kv", log), _Res("sampler", log, fail=True), _Res("pos", log)]) + with pytest.raises(RuntimeError, match="over capacity"): + runner.ingest_request(7) + assert log == [ + ("ingest", "kv", 7), ("ingest", "sampler", 7), + ("remove", "kv", 7), ("remove", "sampler", 7), + ] + + +def test_the_worker_fails_only_the_request_whose_ingest_raised(): + w = Worker.__new__(Worker) + w.worker_id = "w0" + w._draining_rids = set() + w._unprocessed_messages = {"wire": ["parked"]} + w._last_active = {} + processed, failed, registered = [], [], [] + w._graph_runtime = SimpleNamespace( + add_request=lambda **kw: 5, get_sharding_config=lambda rid: None, + ) + w.engine_manager = SimpleNamespace( + evictable_nodes=lambda: [], + add_request=lambda rid, cfg: (_ for _ in ()).throw(RuntimeError("over capacity")), + ) + w.request_state = SimpleNamespace(add_request=lambda rid, info: None) + w.tensor_manager = SimpleNamespace( + register_request=lambda rid, sharding: registered.append(rid), + start_read_tensors=lambda *a, **kw: pytest.fail("no reads for a failed request"), + ) + w._fail_requests = failed.append + w._process_message_list = processed.extend + body = SimpleNamespace( + request_id="wire", partition_worker_graph_ids=None, worker_graph_to_workers={}, + request_info=SimpleNamespace( + partition_name="p", graph_walk="g", resource_configs=None, rid_handle=None, + ), + ) + Worker._add_new_request(w, body) + assert registered == [5] and len(failed) == 1 and list(failed[0]) == [5] + assert "over capacity" in failed[0][5] + # messages parked before the NEW still go through (a parked REMOVE among them) + assert processed == ["parked"] and "wire" not in w._unprocessed_messages diff --git a/test/modular/test_kv_plan_host_lens.py b/test/modular/test_kv_plan_host_lens.py new file mode 100644 index 0000000000..05fbc8a1bc --- /dev/null +++ b/test/modular/test_kv_plan_host_lens.py @@ -0,0 +1,28 @@ +"""The KV plan's host row of sequence lengths iterates as numpy scalars, so +FlashInfer's ``max(seq_lens).item()`` in its decode plan stays cheap.""" + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from mstar.engine.resources.kv.plan import HostLens, _host_lens # noqa: E402 + + +def test_host_lens_iterates_as_numpy_scalars(): + lens = _host_lens([3, 17, 9]) + assert isinstance(lens, HostLens) and lens.dtype == torch.int32 + items = list(lens) + assert [int(x) for x in items] == [3, 17, 9] + assert all(isinstance(x, np.integer) for x in items) + # what FlashInfer does with it + assert max(lens).item() == 17 and int(lens.min()) == 3 and len(lens) == 3 + assert lens.cpu() is lens + + +def test_host_lens_copies_into_a_plain_buffer(): + lens = _host_lens([5, 6, 7, 8]) + buf = torch.zeros(8, dtype=torch.int32) + buf[: len(lens)].copy_(lens) + assert buf.tolist() == [5, 6, 7, 8, 0, 0, 0, 0] + assert torch.equal(lens, torch.tensor([5, 6, 7, 8], dtype=torch.int32)) diff --git a/test/modular/test_multi_stream_final_flag.py b/test/modular/test_multi_stream_final_flag.py index f8bd466bf3..ea712c308c 100644 --- a/test/modular/test_multi_stream_final_flag.py +++ b/test/modular/test_multi_stream_final_flag.py @@ -160,3 +160,38 @@ def test_a_dropped_speculative_rid_neither_flushes_nor_reports_done(): assert speculation.dropped == {"r"} assert (batch.final_stream_rids, batch.stream_partition_done_rids) == (set(), set()) assert _step(worker, "vocoder", {"r": {"pitch"}}) == ({"r"}, {"r"}) + + +def test_a_rows_only_step_threads_its_loop_back_signal_as_present(): + """With every output of the step on the device, a continuing rid that + ran is spliced with an empty loop-back signal (what the walk's fallback + batches carry too); one that did not run is dropped.""" + worker = _worker() + batch = Worker._make_executing_batch( + worker, node_name="vocoder", graph_walk="chunk", request_ids=["r", "s"], + per_request_input_tensors={"r": {}, "s": {}}, + per_request_info={"r": None, "s": None}, + ) + speculation = SimpleNamespace( + node_batch=batch, continuing_rids={"r", "s"}, consumed_edges=[("loop", None)], + scheduled_batch=SimpleNamespace(request_to_worker_graph={"r": 0, "s": 0}), + consumed_streaming_edges={}, spec_id=0, + ) + outputs_N = SimpleNamespace(rows_only=True, row_request_ids=("r",)) + + Worker._thread_outputs_to_speculative(worker, speculation, outputs_N) + + assert speculation.continuing_rids == {"r"} + assert batch.per_request_input_tensors["r"] == {"loop": []} + assert batch.request_ids == ["r"] and "s" not in batch.per_request_input_tensors + + +def test_stream_polls_are_skipped_only_when_no_request_owns_a_buffer(): + worker = _worker() + assert Worker._streams_possible(worker), "unknown (a bare instance): poll" + worker._stream_buffer_rids = None + assert Worker._streams_possible(worker) + worker._stream_buffer_rids = set() + assert not Worker._streams_possible(worker) + worker._stream_buffer_rids.add(7) + assert Worker._streams_possible(worker) diff --git a/test/modular/test_openai_router.py b/test/modular/test_openai_router.py index 267ddb4a79..2fd333f38f 100644 --- a/test/modular/test_openai_router.py +++ b/test/modular/test_openai_router.py @@ -5,6 +5,7 @@ paths are all exercised without the engine. """ +import asyncio import base64 import json import sys @@ -57,13 +58,16 @@ async def collect_results(self, request_id, raw_request=None): self.last_raw_request = raw_request return self._chunks.get(request_id, []) - async def iter_result_chunks(self, request_id): + async def iter_result_chunk_batches(self, request_id): # Same contract as the real one: a consumer that stops early, or an - # exception, aborts the request; a fully drained stream does not. + # exception, aborts the request; a fully drained stream does not. The + # chunks come the way the server hands them out, in batches (here all + # at once), so the streaming handlers see the coalesced form. finished = False try: - for c in self._chunks.get(request_id, []): - yield c + chunks = self._chunks.get(request_id, []) + if chunks: + yield list(chunks) if self.raise_after is not None: raise self.raise_after finished = True @@ -71,6 +75,15 @@ async def iter_result_chunks(self, request_id): if not finished: self.aborted.append(request_id) + async def iter_result_chunks(self, request_id): + batches = self.iter_result_chunk_batches(request_id) + try: + async for batch in batches: + for c in batch: + yield c + finally: + await batches.aclose() + @pytest.fixture def client_and_stub(monkeypatch): @@ -300,6 +313,40 @@ def test_videos_stream_ndjson(client_and_stub): assert _b64.b64decode(body["data"][0]["b64_json"]) == b"mp4-full" +def test_chat_stream_sends_a_batch_of_tokens_as_one_frame(): + """Tokens that were waiting together leave as one body frame (one write + on the socket), still one SSE event per token for the client; a chunk + of another modality gets a frame of its own.""" + from mstar.api_server.openai.serving_chat import _stream + + class _Batches: + async def iter_result_chunk_batches(self, request_id): + yield [_Chunk("text", b"a"), _Chunk("text", b"b"), _Chunk("text", b"c")] + yield [_Chunk("text", b"d"), _Chunk("audio", b"\x00\x01"), _Chunk("text", b"e")] + + async def _collect(): + return [frame async for frame in _stream(_Batches(), "bagel", "rid", 24000)] + + frames = asyncio.run(_collect()) + # role, abc, d, audio, e, stop, DONE + assert [f.count("data: ") for f in frames] == [1, 3, 1, 1, 1, 1, 1] + events = [json.loads(l[6:]) for f in frames[:-1] for l in f.splitlines() if l.startswith("data: ")] + assert "".join(e["choices"][0]["delta"].get("content", "") for e in events) == "abcde" + assert events[5]["choices"][0]["delta"]["audio"]["data"] == base64.b64encode(b"\x00\x01").decode() + assert events[-1]["choices"][0]["finish_reason"] == "stop" + assert frames[-1] == "data: [DONE]\n\n" + + +def test_sse_event_is_one_json_line(): + from mstar.api_server.openai._util import sse + + obj = {"choices": [{"delta": {"content": "h\u00e9llo \"q\""}}], "n": 1} + line = sse(obj) + assert line.startswith("data: ") and line.endswith("\n\n") + assert "\n" not in line[:-2] + assert json.loads(line[6:-2]) == obj + + def test_chat_stream_reports_a_failed_request_in_band(client_and_stub): """A request that fails after the stream opened ends with an error event, not a ``finish_reason: stop`` that reads as a complete answer.""" diff --git a/test/modular/test_preplan_error_paths.py b/test/modular/test_preplan_error_paths.py index da26c00d07..b8d6a8817c 100644 --- a/test/modular/test_preplan_error_paths.py +++ b/test/modular/test_preplan_error_paths.py @@ -131,7 +131,7 @@ def test_gpu_thread_drops_the_stage_when_prepare_inputs_raises() -> None: request_ids=("a",), release_waiters=lambda: released.append(True), per_request_info={"a": SimpleNamespace(request_id="a")}, ) - batch = SimpleNamespace(node_name="dit", graph_walk="decode") + batch = SimpleNamespace(node_name="dit", graph_walk="decode", output_signals=()) with pytest.raises(RuntimeError, match="bad inputs"): Worker._execute_on_gpu_thread(w, batch, node_batch, _done_future()) assert resets == [node_batch] diff --git a/test/modular/test_preplan_guard.py b/test/modular/test_preplan_guard.py index d60b98e2a4..39532cc5bc 100644 --- a/test/modular/test_preplan_guard.py +++ b/test/modular/test_preplan_guard.py @@ -208,3 +208,41 @@ def test_runner_clear_preplan_reaches_every_resource() -> None: runner.clear_preplan() assert kv.events[-1] == "clear" and attn.events[-1] == "clear" assert runner.plan(_step("a")) == {"kv": "fresh", "attn": "fresh"} + + +class _Counting(_Blind): + """Records admits as well, the way a real resource re-sweeps its rows.""" + + def admit(self, step, ctx): + self.events.append("pre_admit" if ctx.is_preplan else "admit") + return super().admit(step, ctx) + + +def test_admit_of_a_pre_admitted_step_skips_the_resources_that_sweep_covered() -> None: + kv, attn = _Counting(), _Counting(deps=("kv",)) + runner = StepRunner({"kv": kv, "attn": attn}) + staged = _step("a", preplan=True) + assert runner.pre_admit(staged).ok + runner.pre_plan(staged) + + # the same step reaches the GPU thread: nothing to admit again + live = _step("a") + assert runner.admit(live).ok + assert kv.events == ["pre_admit", "pre_plan"] + assert attn.events == ["pre_admit", "pre_plan"] + assert runner.plan(live) == {"kv": "staged", "attn": "staged"} + + # the next step of the same shape was not pre-admitted: a full sweep + assert runner.admit(_step("a")).ok + assert kv.events[-1] == "admit" and attn.events[-1] == "admit" + + +def test_admit_of_a_foreign_step_still_sweeps_everything() -> None: + kv, attn = _Counting(), _Counting(deps=("kv",)) + runner = StepRunner({"kv": kv, "attn": attn}) + staged = _step("a", preplan=True) + assert runner.pre_admit(staged).ok + runner.pre_plan(staged) + assert runner.admit(_step("b", span=17, slot=None)).ok + assert kv.events == ["pre_admit", "pre_plan", "clear", "admit"] + assert attn.events == ["pre_admit", "pre_plan", "clear", "admit"] diff --git a/test/modular/test_qwen3_5_bucket_knobs.py b/test/modular/test_qwen3_5_bucket_knobs.py new file mode 100644 index 0000000000..877f36dc73 --- /dev/null +++ b/test/modular/test_qwen3_5_bucket_knobs.py @@ -0,0 +1,28 @@ +"""The Qwen3.5 graph-bucket ladders come from the environment.""" +import importlib + + +def _reload(monkeypatch, name, value): + if value is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, value) + import mstar.model.qwen3_5.submodules as m + return importlib.reload(m) + + +def test_prefill_ladder_default_reaches_8k_and_the_knob_narrows_it(monkeypatch): + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", None) + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS[-1] == 8192 + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS[0] == 32 + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", "64,2048") + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS == [64, 2048] + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", None) + assert len(m.LLMSubmodule.PREFILL_TOKEN_BUCKETS) == 9 + + +def test_decode_ladder_knob(monkeypatch): + m = _reload(monkeypatch, "MSTAR_QWEN35_DECODE_BUCKETS", "1,2,4,8,16,32") + assert m.LLMSubmodule.DECODE_CAPTURE_BATCH_SIZES == [1, 2, 4, 8, 16, 32] + m = _reload(monkeypatch, "MSTAR_QWEN35_DECODE_BUCKETS", None) + assert m.LLMSubmodule.DECODE_CAPTURE_BATCH_SIZES[-1] == 128 diff --git a/test/modular/test_qwen3_5_check_stop_batched.py b/test/modular/test_qwen3_5_check_stop_batched.py index 7822615118..87e89994fe 100644 --- a/test/modular/test_qwen3_5_check_stop_batched.py +++ b/test/modular/test_qwen3_5_check_stop_batched.py @@ -88,3 +88,30 @@ def test_engine_falls_back_and_attributes_failures(): ) assert stops == {"a": {"decode_loop"}} assert list(failures) == ["x"] + + +def test_rows_out_of_batch_order_still_map_by_request(): + """A batch listed in another order than the forward ran it reads each + request's own row, and a request without a row is skipped.""" + from types import SimpleNamespace + + from mstar.model.qwen3_5.config import SAMPLER + from mstar.model.qwen3_5.submodules import LLMSubmodule + from mstar.model.submodule_base import HostRows + + sub = LLMSubmodule.__new__(LLMSubmodule) + sub.config = SimpleNamespace(stop_token_ids={7}) + rows = HostRows(("a", "b", "c"), {"new_token": torch.tensor([1, 7, 1])}) + + def info(max_tokens, iters=0): + return SimpleNamespace( + resource_configs={SAMPLER: SimpleNamespace(ignore_eos=False)}, + dynamic_loop_iter_counts={"decode_loop": iters}, max_tokens=max_tokens, + ) + + infos = {"a": info(100), "b": info(100), "c": info(5, iters=4), "d": info(100)} + stops = LLMSubmodule.check_stop_batched(sub, ["c", "d", "b", "a"], infos, rows) + assert stops == {"c": {"decode_loop"}, "b": {"decode_loop"}} + # forward order: the same answer without the row map + stops = LLMSubmodule.check_stop_batched(sub, ["a", "b", "c"], infos, rows) + assert stops == {"b": {"decode_loop"}, "c": {"decode_loop"}} diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py new file mode 100644 index 0000000000..89b5dcc3f5 --- /dev/null +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -0,0 +1,126 @@ +"""Qwen3.5 decode with the device-side loop-back: the step's input ids come +from the sampler's slot master, the routed signal carries no tensor, and the +knob or a sampler without masters keeps the tensor path.""" +from types import SimpleNamespace + +import pytest +import torch + +from mstar.model.qwen3_5.config import SAMPLER +from mstar.model.qwen3_5.submodules import LLMSubmodule + + +class _Sampler: + def __init__(self, masters=True, tokens=None): + self.has_slot_masters = masters + self.tokens = tokens + self.asked = None + + def loopback_tokens(self, request_ids): + self.asked = list(request_ids) + return self.tokens + + +def _sub(sampler): + sub = LLMSubmodule.__new__(LLMSubmodule) + sub.node_resources = {SAMPLER: sampler} + return sub + + +def _prepare(sub, walk, inputs, has_loopback=True): + return LLMSubmodule.prepare_inputs( + sub, graph_walk=walk, fwd_info=None, inputs=inputs, has_loopback=has_loopback, + ) + + +def test_decode_takes_no_input_ids_when_the_sampler_keeps_them(): + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {}) + assert out.input_ids is None and out.input_embeds is None and out.input_seq_len == 1 + # a routed tensor is ignored on this path: the master is the source + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}) + assert out.input_ids is None + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset({"text_inputs"}) + assert LLMSubmodule.device_loopback_signals(sub, "prefill_text") == frozenset() + + +def test_a_row_without_loopback_reads_its_routed_token(): + # e.g. its prefill ran on another rank: the master does not hold its token + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}, has_loopback=False) + assert out.input_ids.tolist() == [7] and out.input_seq_len == 1 + + +def test_without_masters_or_with_the_knob_off_the_tensor_path_stays(monkeypatch): + sub = _sub(_Sampler(masters=False)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}) + assert out.input_ids.tolist() == [7] and out.input_seq_len == 1 + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset() + with pytest.raises(KeyError): + _prepare(sub, "decode", {}) # the tensor path needs the signal + monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([9])]}) + assert out.input_ids.tolist() == [9] + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset() + + +def _stub_model(sub, seen): + def _position_ids_3d(inputs): + seen["n"] = len(inputs) + return torch.zeros(3, len(inputs)) + + sub.model = SimpleNamespace(model=SimpleNamespace( + build_cos_sin=lambda pos, dtype: (pos, pos), + embed_tokens=SimpleNamespace(weight=torch.zeros(1)), + )) + sub._position_ids_3d = _position_ids_3d + + +def test_preprocess_mixes_routed_rows_into_the_sampler_tokens(): + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 5, 0])) + sub = _sub(sampler) + rows = [ + _prepare(sub, "decode", {"text_inputs": [torch.tensor([9])]}, has_loopback=False), + _prepare(sub, "decode", {}), + _prepare(sub, "decode", {"text_inputs": [torch.tensor([8])]}, has_loopback=False), + _prepare(sub, "decode", {}), + ] + engine_inputs = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2", "r3", "pad1"], + ) + _stub_model(sub, {}) + out = LLMSubmodule.preprocess(sub, "decode", engine_inputs, rows) + assert out["input_ids"].tolist() == [9, 4, 8, 0] + + +def test_preprocess_reads_the_step_tokens_off_the_sampler(): + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 0, 0])) + sub = _sub(sampler) + rows = [_prepare(sub, "decode", {}) for _ in range(4)] + engine_inputs = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2", "pad1", "pad2"], + ) + seen = {} + _stub_model(sub, seen) + out = LLMSubmodule.preprocess(sub, "decode", engine_inputs, rows) + assert out["input_ids"].tolist() == [3, 4, 0, 0] + assert sampler.asked == ["r1", "r2", "pad1", "pad2"] + assert "input_embeds" not in out and seen["n"] == 4 + + +def test_the_loop_back_decode_row_is_uniform(monkeypatch): + from mstar.model.submodule_base import ARNodeInputs + + sub = _sub(_Sampler(masters=True)) + row = LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=True) + assert isinstance(row, ARNodeInputs) and row.input_seq_len == 1 + assert row.input_ids is None and row.input_embeds is None + # a row without loopback in the batch: row by row + assert LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=False) is None + assert LLMSubmodule.uniform_row_inputs(sub, "prefill_text", has_loopback=True) is None + assert LLMSubmodule.uniform_row_inputs( + _sub(_Sampler(masters=False)), "decode", has_loopback=True, + ) is None + monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") + assert LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=True) is None diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py new file mode 100644 index 0000000000..69845b47b4 --- /dev/null +++ b/test/modular/test_sampler_last_token.py @@ -0,0 +1,181 @@ +"""The sampler's last-token master: the device-side loop-back's store. The +captured sample writes a per-step row, commit scatters it to the slot master, +the next step gathers its input ids by slot; an eager step writes and reads +the master by request.""" +import threading + +import numpy as np +import pytest +import torch + +from mstar.engine.resources.sampler.utils import HostBuffer, SamplerBuffers, SamplingConfig + + +def _buffers(max_bs=4, cg_slots=2): + return SamplerBuffers.allocate( + max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, + ) + + +def test_eager_write_and_read_go_by_request(): + bufs = _buffers() + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b"], torch.tensor([11, 12])) + assert bufs.last_tokens_for(["b", "a"]).tolist() == [12, 11] + bufs.write_last_tokens(["c"], torch.tensor([13], dtype=torch.int32)) + assert bufs.last_tokens_for(["c", "a"]).tolist() == [13, 11] + assert bufs.has_slot("a") and not bufs.has_slot("zz") + + +def test_the_captured_row_scatters_to_the_slot_master_and_gathers_back(): + bufs = _buffers(max_bs=4, cg_slots=2) + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b", "c"], torch.tensor([1, 2, 3])) + # a step of [b, a] padded to 4 on cg slot 1: static then dynamic gather + bufs.gather_static(["b", "a"], 4, 1) + bufs.gather_dynamic(["b", "a"], 4, 1) + # what the captured sample does: write the step's tokens into the slot row + bufs.slice_for_bs(4, 1)["last_token_buf"].copy_(torch.tensor([22, 21, 99, 99])) + bufs.scatter_last_token(1) + # real rows reached their masters, padding rows did not touch slot 0's + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [21, 22, 3] + # the next step's input ids, in the slot row's order; padding rows read + # slot 0 (whatever request holds it, or the default when it is free) + slot0_tok = bufs.last_token.master[0].item() + assert bufs.gather_last_tokens(1, 4).tolist() == [22, 21, slot0_tok, slot0_tok] + # the other cg slot is untouched by this step + bufs.gather_static(["c"], 4, 0) + bufs.gather_dynamic(["c"], 4, 0) + assert bufs.gather_last_tokens(0, 4).tolist()[0] == 3 + + +def test_growth_and_slot_reuse_keep_the_tokens(): + bufs = _buffers(max_bs=2, cg_slots=1) + bufs.register_request("a") + bufs.register_request("b") + bufs.write_last_tokens(["a", "b"], torch.tensor([5, 6])) + bufs.register_request("c") # beyond the capacity of 2: the masters grow + bufs.write_last_tokens(["c"], torch.tensor([7])) + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [5, 6, 7] + bufs.unregister_request("a") + bufs.register_request("d") # reuses a's slot; its first step writes it + bufs.write_last_tokens(["d"], torch.tensor([8])) + assert bufs.last_tokens_for(["d", "b", "c"]).tolist() == [8, 6, 7] + + +def test_registration_growth_leaves_the_device_masters_to_the_gpu_thread(): + """A scatter that read the master before a registration grew it must still land.""" + bufs = _buffers(max_bs=2, cg_slots=1) + bufs.register_request("a") + bufs.register_request("b") + bufs.gather_static(["a", "b"], 2, 0) + bufs.gather_dynamic(["a", "b"], 2, 0) + bufs.slice_for_bs(2, 0)["last_token_buf"].copy_(torch.tensor([5, 6])) + # commit's scatter, split where the GIL can switch: it reads the master... + master = bufs.last_token.master + idx, real_bs = bufs._scatter_idx(0) + bufs.register_request("c") # ...the main thread grows past the capacity of 2... + assert bufs.last_token.master is master + master.index_copy_(0, idx, bufs.last_token.slot_view(0, real_bs)) # ...then writes + # the next step grows the device rows on the gpu thread, keeping the tokens + bufs.gather_static(["a", "b"], 2, 0) + bufs.gather_dynamic(["a", "b"], 2, 0) + assert bufs.last_token.master.shape[0] >= 3 and bufs.offset.master.shape[0] >= 3 + assert bufs.gather_last_tokens(0, 2).tolist() == [5, 6] + + +def test_plan_thread_row_init_survives_a_concurrent_host_regrow(monkeypatch): + """A plan-thread ``_init_slot`` between the regrow's copy and swap must still land.""" + bufs = _buffers(max_bs=2, cg_slots=1) + copied, written = threading.Event(), threading.Event() + paused = [] + + def grow_master(self, new_capacity): + new = np.full(new_capacity, self.default, dtype=self.master.dtype) + new[: self.master.shape[0]] = self.master + if not paused: + paused.append(self) + copied.set() + written.wait(timeout=0.3) # times out once the write waits on the lock + self.master = new + + monkeypatch.setattr(HostBuffer, "grow_master", grow_master) + bufs.register_request("a", SamplingConfig(temperature=0.0)) + bufs.register_request("b") + + def plan(): + copied.wait() + bufs.gather_static(["a"], 1, 0) # inits a's row on the plan thread + written.set() + + t = threading.Thread(target=plan) + t.start() + bufs.register_request("c") # grows the host masters on the main thread + t.join() + assert bufs.temperature.master[bufs._rid_to_slot["a"]] == 0.0 + + +def _ingraph(max_bs=4, cg_slots=2, slots=8): + return SamplerBuffers.allocate( + max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, + ingraph_scatter=True, slots=slots, + ) + + +def test_ingraph_scatter_rows_land_by_slot_and_padding_hits_the_trash_row(): + """What the captured sample does with the knob on, step by step on the + CPU views: gather the offsets off the master, advance, scatter offsets + and tokens back; padding rows read and write the trash row only.""" + bufs = _ingraph() + assert bufs.ingraph_scatter and bufs._trash_slot == 8 + assert bufs.last_token.master.shape[0] == 9 and bufs.offset.master.shape[0] == 9 + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b", "c"], torch.tensor([1, 2, 3])) + bufs.gather_static(["b", "a"], 4, 1) + bufs.gather_dynamic(["b", "a"], 4, 1) # uploads the index row, no offset gather + # (a fresh slot's offset reset ran there; an advanced offset survives) + bufs.offset.master[bufs._rid_to_slot["a"]] = 10 + assert bufs._slot_idx_gpu[1, :4].tolist()[2:] == [8, 8] + s = bufs.sampler_for(4, 1) + assert s.ingraph_scatter and s.slot_idx_view.tolist()[:2] == [bufs._rid_to_slot[r] for r in ("b", "a")] + s.gather_in_graph() + assert s.offset_buf.tolist()[1] == 10 + s.offset_buf += 1 + s.scatter_in_graph(torch.tensor([22, 21, 99, 98])) + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [21, 22, 3] + assert bufs.offset.master[bufs._rid_to_slot["a"]].item() == 11 + assert bufs.last_token.master[8].item() in (99, 98) + # the next step's input ids: real rows by slot, padding rows the trash row + got = bufs.gather_last_tokens(1, 4).tolist() + assert got[:2] == [22, 21] and got[2] == got[3] == bufs.last_token.master[8].item() + + +def test_ingraph_scatter_masters_never_grow(): + bufs = _ingraph(max_bs=2, slots=2) + bufs.register_request("a") + bufs.register_request("b") + with pytest.raises(RuntimeError, match="MSTAR_SAMPLER_SLOTS"): + bufs.register_request("c") + # the default path still grows + plain = SamplerBuffers.allocate(max_batch_size=2, device=torch.device("cpu")) + for rid in ("a", "b", "c"): + plain.register_request(rid) + plain.write_last_tokens(["c"], torch.tensor([7])) # a gpu-thread access grows the device rows + assert plain.last_token.master.shape[0] >= 3 and not plain.ingraph_scatter + + +def test_capture_padding_rows_take_no_slot(): + bufs = _buffers() + free = len(bufs._free_slots) + for rid in (-2, -3, -4): + bufs.register_request(rid) + assert len(bufs._free_slots) == free and not bufs.has_slot(-2) + bufs.register_request("a") + bufs.write_last_tokens(["a"], torch.tensor([5])) + # padding rows read slot 0 + assert bufs.last_tokens_for([-2, "a"]).tolist() == [ + bufs.last_token.master[0].item(), 5, + ] diff --git a/test/modular/test_worker_device_loopback.py b/test/modular/test_worker_device_loopback.py new file mode 100644 index 0000000000..b362954dca --- /dev/null +++ b/test/modular/test_worker_device_loopback.py @@ -0,0 +1,51 @@ +"""The worker's side of the device loop-back: a step whose every routed signal +the node keeps on the device, with the client copy inline, stores and holds +nothing and routes zero tensors per (request, signal).""" +from types import SimpleNamespace + +from mstar.communication.tensors import StoredOutputs +from mstar.worker.worker import Worker + + +class _Engine: + def __init__(self, on_device): + self.on_device = on_device + self.calls = 0 + + def device_loopback_signals(self, node, walk): + self.calls += 1 + return self.on_device.get((node, walk), frozenset()) + + +def _worker(inline_emit=True): + w = Worker.__new__(Worker) + w._device_loopback = {} + w._graph_runtime = SimpleNamespace(supports_inline_emit=inline_emit) + return w + + +def test_all_on_device_is_per_step_signals_and_cached_per_node_and_walk(): + engine = _Engine({("LLM", "decode"): frozenset({"text_inputs"})}) + w = _worker() + assert w._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + assert w._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + assert engine.calls == 1, "the node's answer is cached per (node, walk)" + # a step that also routes a signal the node does not keep stays on tensors + assert not w._all_on_device(engine, "LLM", "decode", ["text_inputs", "kv"]) + assert not w._all_on_device(engine, "LLM", "prefill_text", ["new_token"]) + assert not w._all_on_device(engine, "LLM", "decode", []) + assert engine.calls == 2 + + +def test_no_rows_only_step_on_a_runtime_without_inline_emit(): + """Rows-only steps send the client's tokens inline; without it they were dropped.""" + engine = _Engine({("LLM", "decode"): frozenset({"text_inputs"})}) + assert not _worker(inline_emit=False)._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + + +def test_the_zero_tensor_store_shape_matches_the_route_layout(): + rids, signals = [3, 5, 8], ["text_inputs"] + stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) + assert stored.flat_uuids == [] and stored.flat_rids == [] and stored.signal_idxs == [] + assert stored.num_tensors == [0, 0, 0] + assert SimpleNamespace(**stored._asdict()).num_tensors == [0, 0, 0] if hasattr(stored, "_asdict") else True diff --git a/test/rust/test_rust_graph_runtime.py b/test/rust/test_rust_graph_runtime.py index 1faec35751..569b65f59d 100644 --- a/test/rust/test_rust_graph_runtime.py +++ b/test/rust/test_rust_graph_runtime.py @@ -1025,3 +1025,38 @@ def test_inline_values_must_cover_every_rid_or_are_ignored(): )) assert 201 in out.register_uuids, "no values, so the tensor path" rt.send_outputs(_send_input(rid, out.completion_id)) + + +def test_a_loop_back_routed_with_no_tensor_keeps_the_loop_going(): + """The device-side loop-back: the worker routes the decode signal with + zero tensors per request (the token stays on the device) and the client + copy rides inline. The node stays ready, the loop counter advances, the + next pop reports the signal with no uuids, and cleanup frees nothing.""" + rt = _client_loop_runtime() + rid = _admit(rt) + rt._communicator = _Recorder() + _reach_dec(rt, rid, 300) + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[], num_tensors=[0], inline_signal="token", inline_values=[7], + )) + assert out.register_uuids == [], "nothing to stage for the client" + rt.send_outputs(_send_input(rid, out.completion_id)) + popped = rt.pop_rids("dec", WALK, [rid]) + assert popped is not None, "the loop-back with no tensor still fed the node" + inputs = popped.input_edges.to_input_tensors(lambda u: (_ for _ in ()).throw(KeyError(u))) + assert inputs.by_rid[rid] == {"token": []} + before = dict(rt.get_dynamic_loop_iters([rid], "default").values[0]) + # another zero-tensor step: still ready, nothing freed by the cleanup, the + # loop counter moved + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[], num_tensors=[0], inline_signal="token", inline_values=[8], + )) + assert out.freed_inputs.uuids == [] + rt.send_outputs(_send_input(rid, out.completion_id)) + assert rt.pop_rids("dec", WALK, [rid]) is not None + after = dict(rt.get_dynamic_loop_iters([rid], "default").values[0]) + assert after["loop"] == before["loop"] + 1