Pluggable routers - #189
Open
stephen-dwq wants to merge 8 commits into
Open
Conversation
* bug fixes * fix dynamic loop stop bug
…190) MStarClient.stream() yielded zero events when the server sent a Content-Type without a charset (e.g. application/x-ndjson). requests only honours decode_unicode=True when resp.encoding is set, which is derived from that charset, so iter_lines() produced bytes and _stream() skipped every line via its isinstance(line, str) guard. Non-streaming was unaffected because resp.json() ignores the charset. Default resp.encoding to utf-8 when unset, and add a regression test covering a charset-less NDJSON stream. Closes #163
…195) * fix * Properly merge unpacked output with rid outputs
…2) (#168) * shm arena: tensor transport over persistent mmapped segments (Step 2) Replaces the per-tensor file open/write/read/unlink in SharedMemoryCommunicationManager with a Rust segmented /dev/shm arena (rust/src/shm.rs: persistent mmaps + first-fit allocator, buffer-protocol views for zero-copy staging). - Producer reserves (segment, offset) per tensor and D2H-copies straight into the mapped segment on the dedicated copy stream; one stream sync covers the batch before the control message ships. - The location rides the existing descriptors (TensorPointerInfo grows optional shm_segment/shm_offset) — no wire-shape change. - Consumer opens the named segment once, views the bytes zero-copy, and H2D-copies on its dedicated stream; the stream is synchronized before ACKing, since the producer reclaims the slot on ACK (the file transport's f.read() made that implicit). - Pinning: each mapped segment is cudaHostRegister-ed once per process, both sides (MSTAR_SHM_ARENA_PIN, default on) — copies through the side streams then run at page-locked bandwidth and stay asynchronous, which is what preserves the copy/compute overlap the side streams exist for. Measured on the standalone bench: 256 MiB x20 D2H = 466 ms pageable mmap vs 187 ms registered (== torch pinned); registration is a one-time ~21 ms per 256 MiB segment. - Capacity: the arena grows by segments up to MSTAR_SHM_ARENA_MAX_SEGMENTS (mappings never move — registrations and open consumer views stay valid; oversized tensors get a dedicated segment). At the cap, sends backpressure until consumers ACK, failing loudly after MSTAR_SHM_ARENA_FULL_TIMEOUT_S. - Reclaim is uuid-keyed on the sender (cleanup maps uuid -> free), same lifecycle as the file transport's unlink. - Selection: MSTAR_SHM_ARENA = 0 (default, files) / 1 / AUTO on the SHM protocol, mirroring MSTAR_RUST_ZMQ; documented in docs/environment_variables.rst. Tests: cargo tests for the arena/allocator; pytest covers the producer->consumer roundtrip through two managers, descriptor stamping, uuid reclaim, growth + backpressure-then-fail, and factory selection. * arena: defer the consumer ACK on a CUDA event instead of a host sync The edge's H2D copies must complete before the ACK lets the producer reclaim the slot. A blocking h2d-stream synchronize enforced that on the host thread — serializing the transport tick on the slowest copy, which compounds under concurrency. get_ready_tensors already polls a future per pending edge, so the constraint moves onto the device timeline: one CUDA event covers the batch (all copies queue on the h2d stream in program order) and the edge reports ready when the event fires. No host block; the wait_stream ordering for downstream kernels is unchanged. The event path needs CUDA, so it rides the same c16-32 gate run as the pinning claim; on CPU (and in CI) edges stay future=None as before. * arena: fail loudly on transport mismatch in both directions An arena consumer already rejected file-producer descriptors explicitly; the reverse (file consumer, arena producer) surfaced as a bare FileNotFoundError on a path that never existed. Both directions now raise the same explicit message: MSTAR_SHM_ARENA must match across the deployment. Test covers both. * ci: run the full test/rust suite (the arena tests were not exercised) cpu torch + triton + numpy cover the tensor-transport import chain; the job previously pinned the Step-1 communicator test only. * docs: drop migration-step jargon from arena module docs * arena: fragmentation gauge, spill-to-file at the cap, pinned-bytes budget Review follow-ups on the arena's behavior under sustained heterogeneous load: - Fragmentation observability: the allocator (already a sorted, coalescing free-list) now exposes largest_free_block; SegmentedShmArena and the manager surface (total, free, largest, pinned) via stats(). Growth logs the snapshot, and the exact fragmentation signature — a reserve failing while total free covers it — logs a warning naming the collapsed largest block. - Graceful spill: at the segment cap, a send now backpressures briefly (MSTAR_SHM_ARENA_SPILL_AFTER_S, 0.05 s) and then stages the tensor through the per-uuid file protocol instead of failing — slower, never fails, the file transport's saturation behavior. Descriptors keep shm_segment=None for spilled tensors; the consumer reads those through a file fallback (which also makes a file-producer + arena-consumer deployment interoperate); reclaim unlinks spilled files. MSTAR_SHM_ARENA_SPILL=0 restores strict backpressure + timeout. - Pinned budget: MSTAR_SHM_ARENA_PIN_MAX_MB (4096) caps TOTAL cudaHostRegister-ed bytes, distinct from the segment cap — pinned pages come out of the OS's pageable pool system-wide. Segments past the budget stay unpinned (copies work, no async overlap), as do oversized dedicated segments, whose one-shot transfer doesn't amortize the registration. - Wake on read completion: start_read_tensors now returns a real Future (completed by a watcher thread when the h2d-stream event fires) so the worker's eventfd wakes the moment copies land, instead of an otherwise-idle worker discovering them on its next 10 ms poll tick. - GIL released across the arena FFI's slow paths (reserve's segment-growth mmap, multi-allocation uuid frees). Tests updated to the new contract (spill roundtrip + file reclaim + stats gauge; strict mode still fails loudly; the mixed-transport test now proves file->arena interop instead of a refusal). Docs: env vars for the new knobs; installation's Rust section reframed as a running component list; arena module docstring condensed. * test: mixed arena+spill edge; deterministic fragmentation signature Two saturation corners the spill tests didn't pin down: (1) a single edge mixing arena-staged and spilled tensors — the consumer dispatches per descriptor, one start_read_tensors call reads both; (2) the fragmentation warning itself, constructed deterministically: fill to the cap, free alternate slots (total free ample, largest block small), then a reserve that fits the total but no block — asserts the warning fires and the tensor spills while a small tensor still lands in a hole. * docs: frame the rust/ crate as general components, not transport-only The lib-target comment explained the rlib in terms of specific future consumers, and the crate docs described the crate as exactly its first two modules. Both now say what is structurally true: one crate of independent, individually opt-in components with a Python surface and a plain-library form — new capabilities land as new modules. * arena: round-3 review fixes — borrow, GIL, event, race, pin policy, stats - SegmentedShmArena gets interior mutability (segments behind a Mutex), so reserve takes &self: the binding releases the GIL across the growth path, and a &mut PyO3 borrow held there made any concurrent &self call (stats, free, free_uuid) raise "Already mutably borrowed". Growth now serializes on the segments lock instead of the PyCell borrow. - cudaHostRegister goes through ctypes, which releases the GIL for the duration — registering a 256 MiB segment on the send path no longer stalls the process's other Python threads once per growth (the torch cudart binding holds the GIL). Torch-binding fallback kept for environments without a loadable libcudart. - _CudaEventFuture uses Event(blocking=True): the default busy-waits, so the watcher thread turning events into wake futures burned a full core per wait. - _wake_when_done guards its queue/thread creation with a lock (two concurrent first-callers could create two watchers and lose wakes). - Oversized-segment pinning REVERSED per review: dedicated segments are reused for later large tensors, so an unpinned one degrades every subsequent transfer through it — they now pin like any segment, within the budget. - Manager stats() renamed stats_summary() (the raw arena tuple keeps stats()), and under --log-stats the producer path logs the snapshot at most once per MSTAR_SHM_ARENA_STATS_INTERVAL_S (default 60 s), so a long soak leaves an occupancy/fragmentation time series in the logs. * arena: per-entity ceiling checks + rebase to the new register_for_send Ceilings are per-entity and multiply across a node (every entity — workers + the api-server data worker — creates its own arena, and a consumer pins peer segments beyond its own): construction now fails fast with the sizing formula when ONE entity's MAX_SEGMENTS x SEGMENT_MB already exceeds /dev/shm (statvfs), warns when it exceeds current free space, and warns when the per-process pin budget is an outsized share of physical RAM. Both x(num entities) formulas documented in the module docstring and env-vars table. Rebased on main: register_for_send now receives TensorPointerInfos directly; the arena override stamps the passed info as well as the store-tracked ones, and the tests call with infos. * arena: instance-unique names, orphan sweep, immediate spill, dead API out Deep-review fixes (verified by reviewer across 7 models x 9 configs with bitwise gates): - Segment base names are now instance-unique (mstar_arena_{entity}_{pid}_{token}): a fixed per-entity name in the global /dev/shm namespace let a second server's create() truncate the first's LIVE segments (silent cross-server corruption, observed on a shared cluster) and made cross-user startups die on Permission denied. Wire-compatible: consumers open whatever name the descriptor carries. Test: two same-entity-id managers coexist and the first's staged data survives the second's creation. - Startup sweep of orphaned segments: SIGKILL never runs Drop, leaving up to a full arena per kill until reboot. The pid embedded in the new names gives the sweep its liveness check (/proc/<pid>); files it cannot judge or cannot remove are left with a debug note. Test: a dead-owner file is reclaimed, a live-owner file is untouched. - Spill grace defaults to 0: on a worker, the TENSOR_RECEIVED ACKs that free slots are processed by the same thread that would sit in the grace wait, so waiting was pure dead time per at-cap tensor (and strict mode a guaranteed stall-then-fail). Docs now state plainly that strict backpressure is only meaningful where another thread drains ACKs (the threaded api-server). - The uuid-ledger arena API (reserve_for/free_uuid) had no Python users — removed along with its bindings; the module docs no longer describe a reclaim flow this PR does not implement. The _infos_by_uuid side-table is gone too (register_for_send stamps the infos it is passed). * arena: evict dead peer segments (instance-unique names made the cache a leak) Self-review follow-up: the consumer's peer-segment cache never evicted, and the instance-unique naming fix turned that from stale-reuse into an unbounded leak — every producer restart mints NEW segment names, so a long-lived consumer mapped and pinned each generation while the old unlinked-but-mapped segments' memory (and registered pinned bytes) stayed resident forever. Cached entries now record their pinned size and are evicted once the backing /dev/shm file is gone (producer finished or restarted): cudaHostUnregister via ctypes (GIL released), mapping dropped, pinned accounting decremented. Eviction is gated on `pending` being empty — an mmap must outlive any in-flight h2d copy reading from it, and the pending futures are exactly those copies — and time-gated off the hot path. Test: a consumer's cache entry disappears after its producer cleans up and exits. * arena: adversarial-audit fixes — races, leaks, watcher lifecycle, ACK gate From a two-lens self-audit (concurrency; lifecycle/protocol) ahead of the soak: - Pin accounting was a torn read-modify-write: _pin releases the GIL (ctypes), so concurrent pins corrupted _pinned_bytes, and _peer_view's before/after delta could attribute another thread's bytes to the wrong segment (mis-eviction later). One _pin_lock now makes budget check + register + accounting atomic, and closes the check-then-insert race that let two threads map+pin the same peer segment (loser leaked). - Eviction is now safe under threaded reads: gated on being the sole active reader (a second thread may have queued an async H2D whose future has not reached `pending` yet), and unregisters BEFORE unmapping — keeping the entry on unregister failure so accounting stays truthful. - register_for_send can no longer orphan slots: an exception between reserve and the _arena_locs record frees the slot on unwind, and a concurrent-duplicate registration returns its slot instead of leaking it. The D2H copy falls back to blocking when no copy stream exists, so a descriptor can never ship ahead of its bytes. - The wake watcher survives a cancelled/already-resolved future (its death silently downgraded every later wake to the poll tick), closes over only its queue (a bound method pinned the whole manager — and the arena's segments — for the daemon's lifetime), and close() stops it via sentinel. - Double-ACK gate: a re-surfaced (retried) edge no longer re-ACKs — the double-decrement could free a slot while another consumer in the fanout still held the descriptor (silent wrong-bytes on the arena; loud failure on the file transport). - Sibling arena bindings (free/stats/num_segments) release the GIL: a growth holds the segments mutex across an mmap, and blocking on it with the GIL held froze every Python thread for the duration. - Growth warns when this entity's real segment bytes exceed 80% of /dev/shm (oversized dedicated segments grow past the static ceiling), and the periodic stats line carries live_slots/spill_files — the leak canary for deferred reclaim (aborts without ACKs) during soaks. * arena: TTL backstop for abort-orphaned slots (default off) A request aborted after staging but before every consumer ACKs defers reclaim forever (cleanup_request waits for ACKs that will never come). MSTAR_SHM_ARENA_SLOT_TTL_S (default 0 = off) adds a bounded backstop: slots older than the TTL are force-freed, with a loud warning carrying the running total. Safety argument: a slot older than the REQUEST timeout cannot have a legitimate reader — the request is dead by the system's own contract — so a bound safely above it (recommend >= 2x the request timeout) cannot race a real consumer. Reclaims run under capacity pressure (retrying the reserve before spilling) and with the periodic stats sweep. Left off by default pending review discussion; precedent across the ecosystem: vllm-omni's RDMA sender uses a 300 s TTL sweep as its abort backstop (with an acknowledged TTL-vs-in-flight caveat), sglang's omni pipeline a 600 s central watermark. The stronger long-term fix is event-driven reclaim wired to the abort control message (the shape sglang's LLM disaggregation uses), which touches shared cleanup_request semantics and deserves its own change. * docs: repair the garbled reclaim sentence in the shm module header The removal of the stale conductor-driven-reclaim description cut mid-sentence; the header now states the actual contract (embedder frees on consumer ACKs), matching the SegmentedShmArena docstring. * arena: fix the soak's two findings — ACK over-suppression, exit unlink Soak finding 1 (one leaked slot per completed request on the LLM worker): my double-ACK gate keyed suppression by uuid alone, so a tensor legitimately consumed by TWO nodes with staggered readiness had its second reference's ACK suppressed — the producer's fanout refcount (which correctly counted both) never reached zero. The gate now keys by (uuid, edge name, destination node): a re-DELIVERED edge repeats its triple and stays suppressed (the corruption the gate exists for), while a distinct consumer of the same uuid ACKs normally; partial edges ACK only their fresh infos, and the per-request triple set is dropped with the request. Test drives both cases through the real ACK plumbing. Soak finding 2 (worker segments survive graceful shutdown): workers exit with the manager still referenced — no cleanup path runs, the interpreter never collects it, and the Rust Drop that unlinks never fires (the api entity has an explicit cleanup, which is why only its segments disappeared). A weakref.finalize callback — capturing only the mutable segment-path list, never the manager — now unlinks at interpreter shutdown regardless. Test proves it from a subprocess that exits holding a global reference to the manager. * Fix the persisted-tensor reference leak and worker SIGTERM cleanup Root cause of the soak's one-slot-per-request leak, and it is neither the abort deferral nor the ACK gate — it is pre-existing shared bookkeeping that the arena's live_slots canary made visible for the first time: `set_output_ref_counts` counts `routing.persist` in the fanout, so a persisted tensor carries one reference PER PERSIST EDGE on top of its persist flag. Nothing ever released those references. Clearing the flag (unpersist, or request cleanup) left the count above zero, so `can_gc` stayed false and the tensor was never collectable — `cleanup_request` logged "Deferring cleanup ... awaiting TENSOR_RECEIVED ACK" and never revisited. BAGEL's end-of-prefill token is both emitted and persisted, which is exactly one leaked slot per request. The file transport leaks the same way (a per-uuid /dev/shm file per persisted tensor), which is likely the manual cleanup maintainers have been doing for a while. Fix: the manager tracks the persist references it holds (one per `set_persist(True)` — the same edges the fanout counted) and releases them when the flag clears, in step with the flag that was set alongside them. `cleanup_request` now clears via `set_persist` for the same reason. Regression tests cover both the request-end path and the unpersist-then-consume re-route path. Second finding (worker segments surviving shutdown): SIGTERM — what `Conductor.shutdown`'s `p.terminate()` sends — defaults to immediate death, so no unwinding, no atexit, no finalizers, and the segments were never unlinked. The main process gets this for free from SIGINT -> KeyboardInterrupt, which is why only the workers leaked. The worker target now turns SIGTERM into SystemExit, and shutdown escalates to SIGKILL after the join window so a worker stuck in a C call cannot hang teardown (its segments are then reclaimed by the next start's sweep). Also: document the ACK gate's known limitation (a graph routing the same tensor to the same node under the same edge name twice in one request is indistinguishable from a re-delivery), and `ruff --fix`. test/modular failure set is byte-identical to upstream/main before and after (45 pre-existing GPU-dependent failures on this box); arena suite 15/15. * Defer the persist-leak fix to #183; drop the double-ACK guard @NSagan271's #183 is the canonical fix for the persisted-tensor leak, in the conductor/worker domain: it stops counting routing.persist in set_output_ref_counts at the source, leaving the persist flag + the conductor's unpersist-with-K accounting to balance the references. My previous persist-hold mechanism here (tracking references per set_persist and releasing them on flag clear) was a second, redundant fix for the same bug — and once #183 removes the count, it would over-dereference and free persisted tensors early. Reverted, so #168 makes no changes to the shared persist / ref-count semantics and rebases cleanly onto #183. Also dropping the double-ACK guard from the self-audit: as NSagan271 noted, a persisted tensor legitimately re-sent to the same node across graph walks looks identical to a re-delivery, so the guard could suppress a real ACK and undercount against the conductor's K — exactly what #183's accounting must not lose. Removing it restores the pre-audit ACK behavior the conductor logic is designed around. Kept: the worker SIGTERM -> SystemExit graceful-exit handler and the SIGKILL-after-join escalation in Conductor.shutdown (matches the code NSagan271 verified), which is what lets the arena's exit finalizer run and unlink worker segments. The three tests tied to the reverted shared code are removed with it; arena suite green.
Collaborator
|
@stephen-dwq I think your PR somehow includes all of the Rust SHM arena commits as diffs... not sure if there will be major conflicts as a result when trying to merge. Also are the issues with running qwen3 resolved? |
NSagan271
reviewed
Jul 31, 2026
NSagan271
left a comment
NSagan271
left a comment
Collaborator
There was a problem hiding this comment.
Router code overall looks good. Once you fix the git diff issue I can sanity check on a few qwen3-omni configurations.
| hidden_states: torch.Tensor, | ||
| router_states: torch.Tensor | None = None, | ||
| *, | ||
| return_router_states: bool = False, |
Collaborator
There was a problem hiding this comment.
Just to make sure I understand, the state machine operates on a per-layer level (i.e., the previous layer's state feeds into the current layer, and so on?
| return result | ||
|
|
||
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: | ||
| def forward( |
Collaborator
There was a problem hiding this comment.
Nit: maybe worth a comment that router states returned are per-shard; no reduction occurs in this function (which might be self-explanatory given that router states are arbitrary)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds the passing of custom
nn.Modules in place of the standardTopKRouterto the variousMoeBlocks. Adds handling custom routers that are stateful e.g.router(hidden_states, router_states) -> (routing_weights, selected_experts, router_states_next)inside theMoeBlocks.Closes #182
How was it tested?
This PR has passed
test_qwen3_omni_fused_moe.py. I am running into problems when I try to run Qwen3, which is the only model that currently uses theMoeBlocks, so no smoketest has been run.Checklist
ruff check .passes