[RFC]: Executor, Worker, and ModelRunner ownership and migration plan #126

Progress update — 2026-10-22

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f.

• [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 merged StageClient/StageWorker/StageEngine, a local Executor and shared ModelRuntime, execution batches and shared metadata preparation.
• [Refactor] Move runner result application to Scheduler (M2/R1) #107 merged typed ModelRunnerOutput and Scheduler-owned result/submission updates on top of [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118. The engine coordinates pending tickets, runner/connector side effects and retirement.
• [Bugfix] [M5] Restore the device token when an async request resumes in a new slot #109's resumed-device-token fix was already merged. [REFACT] [M5 1/N] Document worker and ModelRunner ownership, findings and migration plan #110 remains a draft documentation PR.
• The findings and five-step migration plan below remain pinned to 9d601ce; statements such as “no Worker or Executor” describe that historical baseline, not current main.
• Reconcile each remaining finding with the merged implementation before opening another M5 refactor. Typed results and the Worker/Executor foundation should be reused. Full runner-owned device-state isolation, explicit submission retirement and KV/storage boundary completion are not declared finished by this update.
• Preserve the agreed shared RngRegistry ownership, GPU-resident dependency paths and cancellation/reuse lifetimes. Existing M5/M2/M3/M4/M7/M9 claims and contract discussions remain in place.

Current ownership reference: unified execution. This is a progress reconciliation, not a fresh code/GPU audit of F1–F16.

Motivation.

Refs #32 (M5: Executor, Worker, and ModelRunner). This issue follows the review on #110: the code walkthrough, the findings, and the migration plan live here, and the repository keeps only stable ownership and lifecycle contracts in docs/design.md. The code walkthrough is in the appendix at the end. Everything else reflects main at 9d601ce.

M5 gives the scheduler one entry point for device execution. The scheduler describes what runs this round: which requests, which stage, which loop. The execution side manages device resources, runs the work, and reports results. The scheduler uses the results to decide what runs next, and the engine connects the two. For example, the plan asks request A to run loop 3 at position 20; the runner returns A's gate signal; the scheduler decides whether A enters loop 4 or CODA (with delayed exit, one loop later).

Done when the engine runs requests through an execution plan, submissions, and results, including cancellation, preemption, and resume, without knowing how hidden states, slots, or CUDA streams are organized.

Today this loop works, but its boundaries are implicit:

• The runner receives the scheduler's mutable Request inside each ScheduledItem and reads position, loop depth, prompt tokens, the last delivered token, output index, and sampling parameters from it while the scheduler keeps changing them.
• Per-request device state (Request.hidden_state, Request.input_token_tensor) lives on the scheduler's object and is written by six components. F1 ([Bugfix] [M5] Restore the device token when an async request resumes in a new slot #109) was a missed write in one of them: a request resumed into a new slot embedded another request's token.
• The engine reads and writes runner internals (model_runner.events, core_stream, boundary_stream, ticket.device_values), and each caller decodes Submission.values by stage.
• Prepared KV batches are built in several runner-side places from private KVCacheManager state.
• There is no Worker or Executor. Device setup, memory planning, and buffer allocation are split across the engine, core/memory.py, the cache manager, and the runner. LLMEngine.close() drains GPU work and aborts requests, but the runner has no terminal state or resource release of its own.

M5 needs agreement on these boundaries before restructuring, and on how they meet M2, M3, M4, M7 (#80), and M9 (#111).

Proposed Change.

Boundary

Ordinary ModelRunner path only. Purple: the proposed layers. Dashed: the couplings this RFC removes.

flowchart TB
  subgraph before["Before"]
    direction LR
    S1["Scheduler"] -->|"mutable Request"| E1["LLMEngine"]
    E1 -->|submit| R1["ModelRunner"]
    E1 -.->|"reads runner internals"| R1
    W1["6 writers: Scheduler, LLMEngine,<br/>preemption, PD worker,<br/>AsyncState, ModelRunner"] -.->|write| Q1[("Request.hidden_state<br/>Request.input_token_tensor")]
  end
  subgraph after["After"]
    direction LR
    S2["Scheduler"] <-->|"plan / result"| E2["LLMEngine"]
    E2 <-->|"plan / result"| X2["Executor"]
    X2 <-->|"plan / result"| R2
    subgraph worker["Worker"]
      direction LR
      R2["ModelRunner"] --> Q2[("slot table<br/>hidden, device token")]
    end
  end
  before ~~~ after
  classDef plain fill:#ffffff,stroke:#1f2328,color:#1f2328
  classDef layer fill:#efe9fb,stroke:#6f42c1,stroke-width:1.5px,color:#1f2328
  class S1,E1,R1,W1,Q1,S2,Q2 plain
  class E2,X2,R2 layer

Execution plan

RFC #32 section 2 asks SchedulerOutput to replace execution interfaces that carry mutable Request objects. The scheduler (M3) produces the plan; M5 lists what it consumes, and M2, M3, and M5 agree on the fields. On 9d601ce the runner reads these Request fields:

KV rows come from M4's public descriptions (step 3). ModelRunner.prepare already snapshots depth, position, and output index for an asynchronous submission; the plan makes that snapshot the runner's input.

Findings on current main

Sixteen findings on 9d601ce. F1 is fixed; each migration step below lists the findings it addresses.

Migration plan

No step introduces a temporary adapter that a later step deletes. Huginn requires synchronous scheduling today, so every step keeps the synchronous path working alongside CUDA async.

1. Typed results from the runners and Submission (F5, F6, F13). Following its 2026-10-12 review, [Refactor] Move runner result application to Scheduler (M2/R1) #107 makes ModelRunner and SpeculativeRunner return typed results directly; this step lands there. A submission has three phases, and M5 asks that the contract keep them separate:

   • Submit returns at once with what later device work depends on: the completion handle and the sampled token on the device. Alternatively the token stays in the runner's slot and the engine never handles it. The next PRELUDE must not wait for host collection; test_prelude_consumes_device_sample_before_cpu_delivery is an acceptance test.
   • Collect returns host values as typed fields that carry the request generation, so a stale result for a reused request ID is rejected.
   • Retire belongs to the runner: a readback buffer is reused only after its copy completes and its result has been collected into independent host data or explicitly discarded. Drop stays as a safety net.

   The engine stops reading or writing events, core_stream, boundary_stream, and device_values, and boundary overlap becomes a runner capability. _inflight splits in two: buffer lifetime moves to the runner, and the scheduling hint (prefer_recurrent: is a RECURRENT submission still running?) uses ready() on the engine's own submissions. PD stream ordering moves to the worker-side connector (M9, [RFC]: PD coordination and KV transfer contracts #111).

   M5 reviews the runner side of that PR and can validate it on GPU.

2. Runner-owned request state (F2, F4, F8, F9, F16). The runner becomes the only owner of per-request device state (hidden state and device token), for sync and async alike, through one slot table interface keyed by request generation. Sync and async may keep different physical layouts. The runner reads the execution plan and its own state, not Request. The runner does not hold generators: sampling RNG stays in the engine's RngRegistry ([REFACT] [M7 2/3] Unify the sampling module, move the sampling RNG into the engine, and split suspend from release #80), and the runner keeps calling rng.acquire where it draws randomness (sampling, and Huginn's recurrent-state initialization in PRELUDE and prefill). RNG continuity follows from that:

   • preemption calls suspend, which frees the slot and keeps the generator;
   • resume may land in a different slot; the generator is keyed by request ID, not by slot, so sampling continues from the same stream;
   • slot reuse never touches generators;
   • termination calls release, which drops the generator, so a reused request ID starts fresh.

   Scheduler.finish stops clearing device tensors; the engine calls the runner's release. Preemption asks the runner to export and import device state instead of copying Request fields (RFC [RFC] Module boundaries and incremental refactoring of the RLT inference runtime #32 section 3: the scheduler picks the victim, the execution side saves and restores). PD takes the final hidden state through a runner operation that also returns its completion event; the runner does not own or recycle transfer buffers, whose lease and teardown order are agreed in [RFC]: PD coordination and KV transfer contracts #111. Coordinate with M3 and M9.

   AsyncSpeculativeRunner follows the same request identity, release, and shutdown rules but keeps its own state (device positions, submitted and delivered rounds, up to two rounds in flight). Its payload is not merged into the slot table. Both speculative runners start their first round from the plan's input token and position, keeping today's startup order.

3. One prepared-batch builder (F3, F10, F11), after M4 exposes host row data. M5 owns one prepared-batch contract, checked against M4's public row and validity descriptions. The contract states allocation identity (today _require_live_batch), borrowing scope (one submission), the completion dependency (the consuming work's event), and reuse only after that event, including after cancellation and before a request ID is reused. Its backends keep their own data paths: host-filled (Workspace), GPU-routed (AsyncState.prepare), PrefillMetadataBank, the Huginn boundary paths, the dynamic path, and device-resident speculative positions, which are never read back to the host. Graph replay uses public operations on a prepared batch; today's standalone write() rebuilds the descriptor on every call and does not belong in the per-layer path. M8's metadata types define what the attention side consumes.

4. Thin Worker and Executor (F7, F12). The Worker initializes and shuts down devices and resources: it owns plan_cache and execution_buffer_bytes (moved out of core/) and creates the runners. The runner, buffers, graphs, and KV storage keep the tensors they use. shutdown() is idempotent, rejects new submissions, drains or discards in-flight work of both runners, and handles a partially initialized Worker. The Executor is a plain single-process class with submit, collect, suspend, release, and synchronize.

5. Measured hot-path work (F9, F10), each with a matched before/after comparison.

Invariants every step keeps

These are part of docs/design.md through #110:

• No .item(), .tolist(), or host readback of tokens, positions, or batch metadata on a device dependency path; GPU-resident paths stay on the GPU.
• With delayed exit, submit round r before reading the signal of round r − 1; never launch an extra loop.
• PREFILL and RECURRENT on the core stream; PRELUDE, CODA, and finalize on the boundary stream; routing on the copy stream; joined only by events.
• A slot, bank, readback buffer, or KV page is reused only after the event of the last work that used it.
• Padding rows never get a KV address, a signal, or a sample.
• LAST_EXITED finalize copies the exit plane into every deeper plane; SHARED finalize only validates.
• Results for a replaced request are dropped, and coda outputs are delivered in order.
• A request's release completes before its ID can be added again. Today release is synchronous, and add_request first drains a retiring speculative request with the same ID.

Non-goals

• A multi-process or distributed Executor (Ray, RPC). RFC [RFC] Module boundaries and incremental refactoring of the RLT inference runtime #32 section 4 starts with a thin single-process entry point.
• Redesigning sampling (M7), attention backends (M8), logical KV allocation (M4), or the PD protocol (M9, [RFC]: PD coordination and KV transfer contracts #111). M5 consumes their public contracts.
• Merging the speculative runners' state into the ordinary slot table.
• Performance claims without a matched measurement.

Alternatives considered

• An engine-side adapter that converts today's Submission. Rejected in review: it adds a layer that step 1 would delete.
• Narrow accessors on the runner for the engine (events, stream identity). Keeps the engine coupled to how the runner orders work; capabilities and completion in the result remove the need.
• Runner-owned generators. Would create a second RNG owner beside RngRegistry ([REFACT] [M7 2/3] Unify the sampling module, move the sampling RNG into the engine, and split suspend from release #80).

Impact

• Performance: steps 1 to 4 intend no change. Step 2 removes the per-request hidden-state view updates in F9; the required stream dependencies remain. Any performance gain is measured in step 5.
• Memory: step 2 lowers the construction peak by not allocating the static states that async replaces (F4); steady-state memory is unchanged. Step 4 separates startup peak, resident capacity, and dynamic growth in the budget (F12).
• API: internal only. LLM, the serving entrypoints, and configuration options are unchanged.

Validation for each step

CPU and GPU suites (--run-gpu with CUDA_VISIBLE_DEVICES), plus focused regression tests for sync and async execution, delayed signals, preemption and resume, cancellation and request-ID reuse, and cleanup ordering. Every code step adds a matched before/after comparison of latency, throughput, and memory; gains for F9 and F10 stay unclaimed until measured.

Feedback Period.

One week after posting, or longer if the M2, M3, M4, M7, or M9 owners need more time.

CC List.

@bjf-frz

Module owners for contract dependencies: @carsontung666 (M2) @YzXiao101 (M3) @BuckMulligan99 (M4) @ycjcx123 (M7) @Levius-Fubuki (M8) @YouHengfei (M9)

Anything else.

Decisions requested:

1. Do you agree with the plan/result boundary, and with splitting a submission into submit, collect, and retire?
2. Do you agree that the runner owns the hidden state and device token, while the speculative runners keep their own state?
3. Do you agree with the migration order and with the interface split with M2, M3, M4, M7, and M9?
4. Every release now completes before its ID can be reused, so bare-ID release and RngRegistry.drop are safe. If release later becomes asynchronous, should release, slot ownership, and RngRegistry share one (request_id, generation) key? The RNG part needs M7's agreement ([REFACT] [M7 2/3] Unify the sampling module, move the sampling RNG into the engine, and split suspend from release #80).

Status of related fixes. Merged means on main. Proposed means in an open PR.

Before submitting a new issue...

• Make sure you already searched for relevant issues in the issue tracker, and looked for the answer in the documentation.

Appendix: code walkthrough at 9d601ce.

PREFILL    prompt chunk -> embed -> every loop depth            -> keep last hidden
PRELUDE    sampled token -> embed                               -> keep hidden
RECURRENT  kept hidden   -> one loop at each row's own depth    -> keep hidden, return gate score
CODA       kept hidden   -> LM head -> sample                   -> return token

ModelRunner
  prepare(batch)             snapshot depths/positions/output indices, build routing
  submit(batch)              pick stream, wait request events, _execute, D2H into a leased slot, record event
  _execute(batch, prepared)  four stages x four execution paths (section 4.3)
  finalize / finalize_many   early-exit KV copies into skipped depths
  suspend(request_id)        wait the request's last event, recycle its slot, keep its RNG
  release(request_id)        same as suspend, then drop its RNG
  execute(batch)             synchronous path: _execute + .tolist()

PYTHONPATH=. python -m pytest -q tests/test_async_pipeline.py tests/test_async_state.py \
  tests/test_cdb_runtime.py tests/test_engine.py tests/test_prefix_growth.py \
  tests/test_prepared_kv.py tests/test_kv_cache.py tests/test_cuda_graph.py tests/test_pd.py