Worker: step threads on one core's two hardware threads, API result threads on another (opt-in)

docs/environment_variables.rst

@@ -496,6 +496,59 @@ Worker scheduling
        thread, which Whisper turbo's serve config sets: its 4-layer steps
        are too short to hide the planning in. Read by
        ``knobs.pre_plan_enabled``.
+   * - ``MSTAR_STEP_THREAD_CORE``
+     - config's ``knobs: step_thread_core``, else ``off``
+     - Pin each worker's main, GPU and plan threads to one core's two
+       hardware threads, and the API server's message thread and Rust
+       frontend bridge to one CPU of another core. A step hands work between
+       the three worker threads several times and every step's tokens pass
+       from the message thread to the bridge; a thread woken on an idle core
+       waits for that core to leave its idle state. The gain needs deep idle
+       states. On a 2x B300 host (Xeon 6767P, ``intel_idle``, C6 at 170 µs)
+       a handoff to another core took 130-175 µs and 4-5 µs to a running
+       one; on an H100 cloud VM without deep C-states it took 22 µs and
+       11 µs, and pinning only cost throughput there. Turn it on (``auto``)
+       where ``/sys/devices/system/cpu/cpu*/cpuidle/state*/latency`` shows a
+       C6-class state of 100 µs or more and the server is host-bound (small
+       models, high concurrency); leave it off on a VM without deep idle
+       states. Measured off → ``auto`` (all three threads on one CPU in
+       parentheses). B300 host, server ``taskset`` to 16 cores and their
+       siblings: Qwen3.5-0.8B c=1 408 → 546 tok/s, c=16 5969 → 8049;
+       Qwen3-TTS 1.7B c=1 TTFA p50 24.9 → 24.3 ms. H100 VM, 40-CPU
+       ``taskset`` (from the #919 review):
+       0.8B c=1 525 → 434 (486) tok/s, c=64 14455 → 14040 (12800); Qwen3-TTS
+       c=1 TTFA mean 21.0 → 21.0 (25.6) ms. ``auto``
+       places by the GPU's rank among the host's GPUs on its NUMA node (sysfs
+       ``/sys/bus/pci/devices``: the PCI functions of the GPU's vendor and
+       class in address order, ``numa_node`` and ``local_cpulist``), so
+       servers on different GPUs take different CPUs without knowing about
+       each other: a worker the core at that rank from the bottom of the
+       node's cores in the process mask, the API server the core at its
+       first worker's GPU's rank from the top; further workers of one
+       deployment on one GPU step by the node's GPU count. A GPU whose node
+       has no CPU in the mask picks, by its rank among the host's GPUs in
+       the same case, from the cores the node-local GPUs leave. The cores a GPU picks
+       from must hold two per GPU that picks from them (its workers count up
+       from the bottom, its API server down from the top); with fewer, or a
+       mask too small to leave a CPU unpinned, nothing is pinned and a
+       warning says why. Servers on one GPU (two MIG instances share its
+       address) would take the same CPUs, so the first ``auto`` server holds
+       ``/tmp/mstar-gpu-<PCI address>.lock`` for its life and a second one
+       pins nothing, with a warning; the lock goes with the process however
+       it exits. The lock only reaches servers that share ``/tmp``, so
+       servers in separate containers on one GPU still take the same CPUs.
+       A split config's second worker can also take a core that a server on
+       a GPU without allowed CPUs takes under the same mask. A
+       number names the first worker's core (its hardware threads in the
+       mask); the next workers take the next cores up and the API server the
+       last core of the mask. A number outside the mask, or too near its top
+       for every worker, stops the server at startup. Threads and pools a
+       pinned thread starts keep the process mask (``mstar.utils.threads``),
+       and the API server's intra-op and encode pools are sized from the
+       CPUs the pins leave (``CpuPlan`` in ``mstar.utils.cpus``). Each worker
+       and the API server log their CPUs and why. Read by
+       ``knobs.step_thread_core``; no effect where ``os.sched_setaffinity``
+       is missing (macOS, Windows).
    * - ``MSTAR_DECODE_BUCKETS``
      - ``1,2,4,8,16,32,48,64,96,128``
      - The decode batch sizes a model captures as CUDA graphs (Qwen3.5 reads

docs/serving.rst

@@ -338,8 +338,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
      - *(optional)* Engine knobs for this deployment, by name: ``device_loopback``,
        ``sampler_ingraph_scatter``, ``sampler_slots``, ``ingraph_decode_rope``,
        ``kv_chain_lazy_steps``, ``kv_plan_cache``, ``kv_pinned_indptrs``,
-       ``gc_gen0_threshold`` and ``pre_plan``
-       (``true``/``false`` or an integer). Each one's environment variable
+       ``gc_gen0_threshold``, ``pre_plan`` and ``step_thread_core``
+       (``true``/``false`` or an integer; ``step_thread_core`` takes
+       ``off``, the default, ``auto`` or a CPU number). Each one's environment variable
        (``MSTAR_DEVICE_LOOPBACK`` and so on; see
        :doc:`environment_variables`) overrides it when set, and an unknown name
        fails at startup. For example ``knobs: {kv_plan_cache: 1,

mstar/api_server/data_worker.py

@@ -9,15 +9,13 @@
 import time
 from collections import deque
 from collections.abc import Callable, Collection
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from functools import cached_property
 
 import torch
 
 from mstar.graph.base import TensorPointerInfo
 from mstar.graph.loop_indices import NestedLoopIndices
-from mstar.model.components.diffusion.image_io import PNG_WORKERS
 
 try:
     import torchaudio  # noqa: F401 — probes availability; real usage in callers
@@ -59,9 +57,8 @@
 from mstar.observability.recorders import EdgeTransfers
 from mstar.profile.format import InputInfo, RxInfo, TxInfo
 from mstar.streaming.topology import ClientStreamLimits, client_stream_limits
-from mstar.utils import profiler
+from mstar.utils import cpus, profiler
 from mstar.utils.client_errors import client_error_text
-from mstar.utils.cpus import usable_cpus
 from mstar.utils.ipc_format import (
     AbortRequest,
     ClientStreamInput,
@@ -76,6 +73,7 @@
     UnpersistTensors,
     WorkerMessageType,
 )
+from mstar.utils.threads import Thread, ThreadPool
 
 logger = logging.getLogger(__name__)
 
@@ -94,26 +92,6 @@ def pooled_output(modality: str, loop_indices: NestedLoopIndices | None) -> bool
     return modality in ARTIFACT_MODALITIES and not streamed
 
 
-def encode_workers(cpus: int | None = None, png_workers: int | None = None) -> int:
-    """Outputs encoded at once. A PNG deflates on ``png_workers`` threads of
-    its own (``image_io.PNG_WORKERS``), so the pool times that stays within
-    the CPUs the process may use."""
-    if cpus is None:
-        cpus = usable_cpus()
-    if png_workers is None:
-        png_workers = PNG_WORKERS
-    return max(1, min(4, cpus // png_workers))
-
-
-def copy_workers(cpus: int | None = None, png_workers: int | None = None) -> int:
-    """Threads a large output read copies over (``host_copy.HostCopier``): one
-    encode's share of the CPUs in the budget ``encode_workers`` sets. Reads run
-    one at a time on the data worker's thread, each ahead of its own encode."""
-    if cpus is None:
-        cpus = usable_cpus()
-    return max(1, cpus // encode_workers(cpus, png_workers))
-
-
 def _sha256(path: str) -> bytes:
     with open(path, "rb") as file:
         return hashlib.file_digest(file, "sha256").digest()
@@ -259,10 +237,10 @@ def __init__(
             communicator=self.communicator,
             tcp_transfer_device=tcp_transfer_device,
             enable_prof=enable_prof,
-            copy_workers=copy_workers(),
+            copy_workers=cpus.plan().copy_workers,
         )
 
-        self.thread = threading.Thread(
+        self.thread = Thread(
             target=_preprocess_loop,
             kwargs=dict(
                 in_queue=self.request_input_queue,
@@ -529,7 +507,7 @@ class PreprocessWorkerThread:
     # class-level so a thread built without __init__ never sends ENABLE_METRICS
     worker_metrics_requested: threading.Event | None = None
     # None (a thread built without __init__) postprocesses every output inline
-    _encode_pool: ThreadPoolExecutor | None = None
+    _encode_pool: ThreadPool | None = None
     # Appends handled per pass before output delivery gets its turn again: one
     # for each of up to 32 live calls appending at once. An append costs a file
     # decode, far less than a request's preprocess, which is one per pass.
@@ -630,8 +608,8 @@ def __init__(
         # with a bare communicator).
         self.wakeup = wakeup
         self.idle_wait_ms = 20
-        self._encode_pool = ThreadPoolExecutor(
-            max_workers=encode_workers(), thread_name_prefix="mstar-encode",
+        self._encode_pool = ThreadPool(
+            max_workers=cpus.plan().encode_workers, thread_name_prefix="mstar-encode",
         )
         # Finished encodes, appended by the pool's done callbacks and released
         # by this thread, which alone touches the per-request output state.

mstar/api_server/entrypoint.py

@@ -50,6 +50,7 @@
 from mstar.api_server.streaming_response import ReleasingStreamingResponse
 from mstar.communication.communicator import CommProtocol, make_communicator
 from mstar.communication.tensors import ensure_deployment_token
+from mstar.conductor.workers import accelerator_type, deployment_workers
 from mstar.model.multimodal import PromptPart
 from mstar.model.registry import model_init_kwargs
 from mstar.observability.catalog import STATUS_INTERVAL_S
@@ -62,14 +63,15 @@
 )
 from mstar.profile.display import pretty_print_profile
 from mstar.profile.format import OutputInfo, RequestProfile, RequestTiming
-from mstar.utils import profiler
+from mstar.utils import cpus, knobs, profiler
 from mstar.utils.client_errors import client_error_text, scrub_paths
 from mstar.utils.exitcode import describe_exitcode
 from mstar.utils.gc_freeze import freeze_after_setup
 from mstar.utils.generation_kwargs import normalize_generation_kwargs
 from mstar.utils.logging_config import quiet_noisy_loggers
 from mstar.utils.orphan import watch_parent
 from mstar.utils.socket_prefix import per_port_socket_prefix
+from mstar.utils.threads import Thread, pin_thread
 
 logger = logging.getLogger(__name__)
 
@@ -409,6 +411,11 @@ class DeadConductorError(RuntimeError):
 class APIServer:
     """Accept multimodal requests, forward to conductor, collect results."""
 
+    # The CPU the message thread and the Rust frontend's bridge share (None:
+    # not pinned) and the process mask; see ``mstar.utils.cpus``.
+    _result_cpu: int | None = None
+    _process_cpus: frozenset[int] = frozenset()
+
     def __init__(
         self,
         socket_path_prefix: str = "/tmp/mstar",
@@ -455,6 +462,7 @@ def __init__(
         ) if model is not None else frozenset()
         self.takes_stream_input = bool(self.client_streams)
         self._model_config = model_config
+        self._take_result_cpu()
 
         # Before any tensor manager; the conductor and workers inherit it.
         ensure_deployment_token()
@@ -482,7 +490,7 @@ def __init__(
         # Background thread that drains results from the conductor. Started by
         # finalize_setup() once the workers report ready — before that there's
         # no traffic to drain and the HTTP server isn't up yet.
-        self._msg_thread = threading.Thread(
+        self._msg_thread = Thread(
             target=self._process_messages, daemon=True
         )
 
@@ -528,6 +536,24 @@ def max_call_input_bytes(self) -> int:
         has accepted it, so a bad value is reported with the rest."""
         return _call_input_limit(self._model_config)
 
+    def _take_result_cpu(self) -> None:
+        """Keep the result threads' CPU from this process's plan (``main``
+        installs the deployment's) and log it."""
+        plan = cpus.plan()
+        self._result_cpu, self._process_cpus = plan.result, plan.mask
+        if plan.result is not None:
+            logger.info("API server: message thread and frontend bridge on CPU %d (%s; mask %s)",
+                        plan.result, plan.why, cpus.format_cpu_list(plan.mask))
+        elif knobs.step_thread_core() != "off":
+            logger.warning("API server: result threads not pinned (%s)", plan.why)
+
+    def pin_result_thread(self) -> None:
+        """Move the calling thread to the result threads' CPU. Every step's
+        tokens pass from the message thread to the frontend's bridge; on one
+        CPU each handoff wakes a core that is already running."""
+        if self._result_cpu is not None:
+            pin_thread({self._result_cpu}, self._process_cpus)
+
     def finalize_setup(self) -> None:
         """Block until the conductor signals that every worker has finished
         setup (weight load + warmup + CUDA-graph capture), then start draining
@@ -968,6 +994,7 @@ def _process_messages(self) -> None:
         """Drain the ZMQ pull socket and route results to pending requests.
         Also watches the conductor process. Once it exits, every pending
         request is failed and the HTTP server is told to stop."""
+        self.pin_result_thread()
         next_liveness_check = 0.0
         while self.running:
             now = time.monotonic()
@@ -2592,19 +2619,31 @@ async def shutdown_event():
 # CLI entry point
 # ------------------------------------------------------------------
 
-def default_preprocess_threads(cpus: int | None = None) -> int:
-    """Intra-op threads for the API server when neither the config nor the model sets them.
-
-    Torch's default is one per host core, which inside a container can be several
-    times the CPUs the cgroup quota lets run; a fork-join op then waits for its
-    parked threads. A quarter of the usable CPUs, at most 4: the process shares
-    them with the engine workers, and its per-request ops are a few MB.
-    """
-    if cpus is None:
-        from mstar.utils.cpus import usable_cpus
-
-        cpus = usable_cpus()
-    return max(1, min(4, cpus // 4))
+def deployment_cpu_plan(config: dict) -> cpus.CpuPlan:
+    """The deployment's CPU plan, as each worker builds it from the conductor's
+    device list: the config's ``step_thread_core`` and ``deployment_workers``.
+    Under ``auto`` this also takes the GPUs' locks, and a server that finds one
+    held pins nothing, here or in its workers. A bad knob value is left for the
+    conductor's deployment check to report; a CPU number the mask can't honour
+    raises ``CpuPlanError``."""
+    try:
+        knobs.configure(config.get("knobs"))
+    except ValueError:
+        return cpus.plan()
+    setting = knobs.step_thread_core()
+    if setting == "off":
+        return cpus.deployment_plan(setting, [])
+    devices = deployment_workers(config, accelerator_type()).devices
+    plan = cpus.devices_plan(setting, devices)
+    if setting != "auto" or not plan.pinned:
+        return plan
+    refused = cpus.claim_gpus(cpus.device_addresses(devices))
+    if refused is None:
+        return plan
+    logger.warning("step_thread_core auto: nothing pinned, %s", refused)
+    # the workers read the knob when they build their plan
+    os.environ[knobs.KNOBS["step_thread_core"]] = "off"
+    return cpus.CpuPlan(plan.mask, plan.usable, why=f"not pinned: {refused}")
 
 
 def _uvicorn_config(app, host: str, port: int, config: dict) -> uvicorn.Config:
@@ -2621,16 +2660,16 @@ def _set_preprocess_threads(model, config: dict) -> None:
     """Size torch's intra-op pool for this process's media preprocessing.
 
     ``config.yaml``'s ``preprocess_torch_threads`` wins (``null`` keeps torch's
-    default), else the model's ``PREPROCESS_TORCH_THREADS``, else
-    :func:`default_preprocess_threads`. Affects only this process: the conductor
+    default), else the model's ``PREPROCESS_TORCH_THREADS``, else the CPU
+    plan's ``preprocess_threads``. Affects only this process: the conductor
     and workers are spawned fresh.
     """
     if "preprocess_torch_threads" in config:
         threads = config["preprocess_torch_threads"]
         if threads is None:
             return
     else:
-        threads = model.PREPROCESS_TORCH_THREADS or default_preprocess_threads()
+        threads = model.PREPROCESS_TORCH_THREADS or cpus.plan().preprocess_threads
     import torch
 
     torch.set_num_threads(int(threads))
@@ -2746,6 +2785,11 @@ def main(argv: list[str] | None = None):
         cache_dir=args.cache_dir,
         **{**model_init_kwargs(model_name), **yaml_model_kwargs},
     )
+    # before any thread: the pools are sized from the CPUs it leaves unpinned
+    try:
+        cpus.install(deployment_cpu_plan(config))
+    except cpus.CpuPlanError as e:
+        raise SystemExit(f"mstar: {e}") from None
     _set_preprocess_threads(model, config)
 
     global api_server

mstar/api_server/media_io.py

@@ -31,6 +31,8 @@
 
 import numpy as np
 
+from mstar.utils.threads import Thread
+
 # MIME (top-level type or full type) -> file extension used when persisting.
 _MIME_TO_EXT: dict[str, str] = {
     "image/png": ".png",
@@ -365,7 +367,7 @@ def _cleanup() -> None:
         time.sleep(UPLOAD_CLEANUP_DELAY_S)
         remove_uploads(file_paths, upload_dir)
 
-    threading.Thread(target=_cleanup, daemon=True).start()
+    Thread(target=_cleanup, daemon=True).start()
 
 
 # ---------------------------------------------------------------------------

mstar/api_server/rust_frontend.py

@@ -72,6 +72,7 @@
 )
 from mstar.model.multimodal import PromptPart
 from mstar.utils.client_errors import client_error_text
+from mstar.utils.threads import ThreadPool
 
 logger = logging.getLogger(__name__)
 
@@ -297,6 +298,10 @@ def _handle(self, msg: dict) -> None:
     async def _serve(self) -> None:
         logger.info("Rust-frontend bridge serving")
         loop = self.loop = asyncio.get_running_loop()
+        # asyncio would start its default pool from this thread, and so on
+        # the result CPU when the bridge is pinned; handlers' to_thread work
+        # (audio planning, request decoding) belongs on the process mask
+        loop.set_default_executor(ThreadPool(thread_name_prefix="mstar-bridge"))
         # Blocking receive on an executor thread: the wait happens inside the
         # transport (GIL released), so an arriving submit is picked up
         # immediately instead of on the next tick of a sleep/poll loop — and
@@ -323,6 +328,12 @@ async def _serve(self) -> None:
 
     def run(self) -> None:
         """Blocking serve loop (the ``uvicorn.run`` replacement)."""
+        # with the message thread: tokens pass from it to this loop. The Rust
+        # process is already running, so it keeps the process mask. Test doubles
+        # stand in for the server without the pin.
+        pin = getattr(self.server, "pin_result_thread", None)
+        if pin is not None:
+            pin()
         try:
             asyncio.run(self._serve())
         except KeyboardInterrupt:

mstar/communication/arena.py

@@ -61,6 +61,7 @@
 from mstar.graph.base import GraphEdge, TensorPointerInfo
 from mstar.utils.containers import ParallelList
 from mstar.utils.host_copy import HostCopier
+from mstar.utils.threads import Thread
 
 logger = logging.getLogger(__name__)
 
@@ -945,7 +946,7 @@ def _wake_when_done(self, cuda_future) -> Future:
                 # staticmethod target closing over ONLY the queue: a bound
                 # method would pin the whole manager (and its arena — so
                 # segments never unlink) for the daemon thread's lifetime.
-                threading.Thread(
+                Thread(
                     target=self._watch_wakes, args=(self._wake_q,),
                     daemon=True,
                     name=f"arena-h2d-wake-{self.my_entity_id}").start()