[RFC] Module boundaries and incremental refactoring of the RLT inference runtime

[bjf-frz]

Summary

This RFC proposes a module-by-module refactoring of vllm-lt, grounded in detailed code walkthroughs. The goal is to establish clear responsibilities, state ownership, and interfaces while preserving recurrent looped transformer (RLT) semantics: loop-level batching, exit policies, depth-aware KV caching, and prefill/decode (PD) disaggregation.

This is the overarching architecture and work breakdown proposal. It defines module boundaries, contracts, and refactoring dependencies before individual modules receive detailed reviews and implementation changes.

vLLM V1 provides an architectural reference for the scheduler/execution feedback loop and the separation of scheduling state from device execution. This does not imply adding a vLLM dependency, reproducing its full class hierarchy, or promising API/ABI compatibility.

Motivation and scope

The implementation includes local inference, loop-level scheduling, asynchronous execution, multiple attention backends, cache management and preemption, and PD disaggregation. As these features have accumulated, several components have started accessing other components' internal queues, mutable request objects, and KV allocations.

The initial architectural review identified the following boundary concerns:

• Request objects combine logical progress with device tensors, RNG state, and asynchronous output placeholders. Multiple components mutate this state.
• The engine handles the execution loop, stage transitions, exit decisions, and output updates.
• The scheduler reads internal KV allocations, block tables, and reference counts to estimate resource costs.
• KVCacheManager combines logical block management, device storage, execution metadata, and attention dispatch.
• PD workers directly modify local scheduler queues and request stages and invoke private scheduling methods.
• Serving depends on internal scheduler/cache objects rather than a stable engine interface.

These observations identify architectural concerns, not necessarily confirmed correctness defects. Subsequent module reviews must establish concrete findings, their impact, and the evidence supporting each change.

The scope includes these capabilities as implemented or under development. Paths below identify responsibilities; some implementations may not yet be present on the default branch. Each module review must pin a reproducible commit or patch baseline before citing line numbers. This RFC does not claim that a complete line-by-line review has already been performed.

Design principles

1. Make state ownership explicit. The scheduler owns scheduling-side request state and execution-result updates. EngineCore drives the scheduling/execution loop. Worker/ModelRunner owns device execution state.
2. Separate logical resources from device operations. KV allocation, references, and release conditions are distinct from tensor storage, copies, and attention execution.
3. Use explicit contracts. Scheduling tasks, execution results, and connector metadata should replace cross-module mutation of private state.
4. Preserve RLT semantics and performance. Keep token position distinct from loop depth, preserve exit-signal timing, and retain device-resident data paths.
5. Refactor incrementally. Each step should remain runnable and verifiable. Behavioral fixes must be distinguished from structural changes.
6. Keep abstractions proportional. Introduce only the execution modes and backends actually needed. Prefer small data structures, composition, and thin interfaces.

Target architecture

flowchart TD
    A["LLM / CLI / API Server"] --> B["Engine interface and input/output processing"]
    B --> C["EngineCore"]
    C --> D["Scheduler"]
    D --> E["Logical KVCacheManager"]
    C --> F["Executor / Worker"]
    F --> G["ModelRunner"]
    G --> H["Model + Sampler"]
    H --> I["Attention backend"]
    G --> J["Device KV storage and execution metadata"]
    D --> K["Scheduler-side connector"]
    F --> L["Worker-side connector"]
    K -. "Transfer metadata / completion feedback" .-> L
    L --> M["NIXL transport"]

Loading

The main execution loop is:

Scheduler.schedule()
    -> SchedulerOutput
    -> Executor / Worker / ModelRunner
    -> ModelRunnerOutput
    -> Scheduler.update_from_output()
    -> Engine outputs

Synchronous and asynchronous paths share logical contracts. Asynchronous scheduling may advance submission ahead of GPU completion and CPU output delivery, but these progress points and their resource-release conditions must remain distinct.

The PD coordinator sits above the P/D engines and coordinates their roles. An executor dispatches work for one local engine. These are separate responsibilities.

Module decomposition

Modules are functional boundaries and units of review/refactoring. They do not necessarily correspond to separate processes or elaborate inheritance hierarchies.

ID	Module	Responsibilities	Current code mapping
M1	Entrypoints and input/output processing	Python/CLI/HTTP interfaces, argument conversion, tokenization, streaming text, backpressure, and disconnect handling	entrypoints/, serving/
M2	Engine and EngineCore	Engine interface, component assembly, execution loop, in-flight task orchestration, control requests, shutdown, and error propagation	Orchestration in engine/llm_engine.py
M3	Scheduler and request state	Admission, queues, batching, priority, fairness, refill/no-refill, stage/loop progress, exit decisions, result updates, completion, and preemption decisions	core/scheduler.py, request.py, and engine-side transitions/exit logic
M4	Logical KV management	Block allocation and growth, block tables, prefix references, depth validity, transfer leases, and release conditions	Logical resource management in core/kv_cache_manager.py
M5	Executor, Worker and ModelRunner	Execution dispatch, device lifecycle, memory profiling, batch preparation, device state, KV storage, streams/events, buffers, CUDA Graphs, and asynchronous results	worker/, core/memory.py, and device storage/operations currently in KV management
M6	Model implementation and loading	Ouro architecture, weight loading, prelude/recurrent/coda, gate-score computation, and numerical semantics	models/
M7	Sampling	Greedy/top-k/top-p sampling, sampling metadata, RNG use, and logits-to-token conversion	Runner _sample*() methods and relevant contracts in sampling_params.py
M8	Attention	Backend capabilities and selection, attention metadata construction, and Torch/Triton/FlashAttention implementations	Attention kernels, backend dispatch currently in KV management, and associated metadata
M9	PD and KV transfer	P/D coordination, handoff protocol, scheduler/worker connector roles, NIXL transport, and failure cleanup	pd/

Configuration and shared types should live with their owning module or the interface they describe. Kernels are grouped by function: routing/gather/scatter belongs to the execution runtime; attention kernels belong to M8. M8 defines backend-specific attention metadata, while M5 manages its device-buffer lifecycle and execution integration.

Ownership and contracts

1. Scheduler owns logical request progression

Move stage transitions, exit policies, and result updates from the engine into the scheduler and its internal helpers, establishing a schedule() / update_from_output() feedback loop.

• The model computes gate scores.
• The runner returns scores with their request generation, token position, loop depth, and execution sequence.
• The scheduler consumes the corresponding signals, updates cumulative exit state, and decides whether to continue recurrent execution or enter coda.
• The sampler produces tokens. The scheduler applies EOS and length limits and selects the next stage.
• EngineCore manages execution submission, result collection, and control-event orchestration.

Exit policy, admission, and request transitions may use separate helpers. Scheduling-side state should retain a single owner without concentrating all logic in one large method.

2. Separate scheduler request state from worker state

Scheduler-side state	Worker-side state
Request identity, prompt, output records, and logical position	Hidden tensors, device tokens, and state slots
Stage, loop progress, and cumulative exit state	Batch rows, gather/scatter state, and device metadata
Logical records of pending execution and pending output delivery	Streams, events, buffers, and readback leases
Sampling parameters, seed, and output ordering	RNG objects and their execution/snapshot state

SchedulerOutput should progressively replace execution interfaces carrying mutable Request objects. It should describe request identity, stage, token ranges, loop depth, logical block information, and necessary state updates.

ModelRunnerOutput should distinguish gate signals, sampled tokens, completed prefill ranges, and resource/transfer completion feedback. Callers should not have to infer the meaning of an untyped result list from the current stage.

Asynchronous sampled tokens may remain on the device for the next prelude. Separating logical interfaces must not add unnecessary .item() calls, CPU round trips, or synchronization. Results arriving after cancellation must be identified by generation/execution identity so they cannot affect a new request reusing the same external request ID.

3. Logical KV management vs device storage

M4 owns logical resources and validity. M5 owns KV tensors and device operations. M8 consumes KV views and attention metadata.

• The scheduler uses public resource-cost and capacity interfaces rather than reading _refs or private allocations.
• M4 manages block tables, prefix references, transfer leases, and deferred release.
• The execution side performs KV writes, cross-depth finalization, snapshot copies, and restoration.
• Completion evidence establishes validity and permits reuse. Submission alone does not imply completion.
• Prepared metadata has explicit borrowing scope, layout, and lifetime. Execution code should not depend on private allocation-object structure.

Filling skipped depths after early exit is a cross-boundary protocol: determine the target range, perform the operation, and confirm completion. Existing LAST_EXITED/SHARED semantics must be preserved.

Preemption follows the same division. The scheduler selects a victim and suspends scheduling; the execution side safely saves/restores state; M4 updates logical resources. Restoration must preserve KV, hidden state, loop progress, exit state, and RNG semantics. Full-depth recomputation must not be assumed equivalent to restoring the original state.

4. Executor / Worker / ModelRunner as one initial refactor unit

• Executor: dispatch execution and control calls.
• Worker: initialize devices, load models, plan memory, initialize resources, and shut them down.
• ModelRunner: prepare and execute batches, maintain device request state, invoke sampling, and manage asynchronous submissions/results.

Start with a thin single-process execution entry point. This refactoring does not require Ray, generic RPC, or unused distributed executors. Buffers, asynchronous state, and CUDA Graphs should be reviewed together within M5 to preserve stream/event and memory-lifetime guarantees.

5. PD through scheduler-side and worker-side connectors

M9 contains three subcomponents:

1. Coordinator: P/D selection, role-specific resource credits, request routing, process lifecycle, and overall failure handling.
2. Scheduler-side connector: receive-wait state, handoff plans after allocation, completion feedback, activation conditions, and deferred release.
3. Worker-side connector / transport: memory registration, KV/hidden movement, CUDA events, NIXL handles, completion notifications, and deregistration.

Direct PD-worker manipulation of queues, Request objects, and _take() should migrate to explicit scheduling/execution protocols.

The handoff contract must specify:

• Valid KV ranges at each depth and destination block mappings.
• The final prompt token's hidden state, with coda and sampling performed by D.
• Transfer generation, chunk sequence, commit, receive completion, and acknowledgement.
• The ordering of remote-write termination and safe release during cancellation, timeout, or worker failure.
• Constraints between prefix hits, in-flight transfers, and preemption.

Commit marks a submission boundary; it is not sufficient evidence that GPU data is usable. Preserve the existing transfer direction and chunk timing rather than changing semantics merely to match an interface.

6. Stable frontend engine interface

M1 uses a common interface to submit/cancel requests, query active requests, retrieve outputs, shut down, and read resource statistics. Remove direct dependencies on engine.scheduler.requests and engine.cache_manager incrementally.

The local engine and PD facade implement the capabilities required by the frontend. PD should not need synthetic scheduler/cache objects to satisfy Serving's assumptions.

RLT-specific extensions

Feature	Primary owner	Extension
Loop-level continuous batching	M3	Scheduling tasks include stage, loop depth, and token position
Ouro / delayed / trace exits	M3, with M6 signals	Explicit signal depth, consumption count, and exit timing
Persistent recurrent hidden state	M5	Device request state retained across steps
LAST_EXITED / SHARED KV	M4 + M5 + M8	Depth mapping, validity, and physical access layout
Asynchronous loop pipeline	M2 + M3 + M5	Distinct submission, completion, and output-delivery progress
PD KV and hidden handoff	M9	Explicit state handoff and resource leases

Refactor sequence and independence

Phase	Deliverable	Dependencies / boundary
R0	Module inventory, state-ownership table, and execution/resource/handoff contracts	Pin a baseline, define contracts, and migrate incrementally without changing every caller at once
R1	Scheduler–EngineCore boundary	Move state updates and exit policies; preserve shared state semantics across synchronous/asynchronous paths
R2	Logical KV–device execution–attention boundary	Design interfaces jointly, then refactor each implementation independently
R3	M5 internal structure	Worker lifecycle, runner, asynchronous state, buffers, graphs, and device state
R4	PD connector integration	Replace private-state access using stable scheduling, resource, and execution interfaces
Independent tracks	M1, M6, M7	Review, refactor, and validate independently once their contracts are stable

Refactoring batches differ from functional modules. Preemption, KV metadata, and PD handoff require joint interface design, but implementations should be independently changeable once those contracts are established. Introduce compatibility adapters, migrate callers, and then remove the old paths. Avoid maintaining two long-lived sources of state.

Required module walkthrough

Each subsequent module RFC/review must include:

1. Baseline and boundary: pinned commit/patch baseline, file inventory, callers/callees, and responsibility boundaries.
2. Classes and state: class design, field purpose, authoritative writer, state transitions, and resource lifecycle.
3. Functions and arguments: key signatures, inputs/outputs, preconditions, and side effects; tensor shape/dtype/device/layout where relevant.
4. Line-by-line coverage: review all executable logic, branches, exception paths, and cleanup paths in scope, organized by function and call chain, with file and code-range references.
5. Evidence-based findings: location, callers, trigger conditions, impact, and recommendation; distinguish confirmed defects, design concerns, open questions, and necessary complexity.
6. Target design and migration: proposed boundaries, interfaces, migration steps, compatibility strategy, and conditions for removing old paths.
7. Validation: module checks, cross-module regression coverage, relevant accuracy/performance checks, and unverified areas.

Acceptance criteria

• Preserve existing functionality and configuration semantics; document behavioral fixes separately.
• Establish clear ownership of requests, queues, device state, KV blocks, and transfer lifecycles.
• Use agreed cross-module interfaces and progressively eliminate private allocation/queue mutation.
• Cover completion, cancellation, exceptions, late results, and request-ID reuse in synchronous and asynchronous paths.
• Cover exit modes, KV layouts, chunked prefill, refill/no-refill, prefix caching, incremental growth, preemption, and applicable feature combinations.
• Compare attention/model numerics with appropriate tolerances. Do not present unverified cross-batch bitwise equivalence as an existing guarantee.
• Validate greedy and fixed-seed sampling paths, including RNG state and exit trajectories across preemption/restoration.
• Validate PD KV at every valid depth/position, hidden-state handoff, first-token semantics, cancellation during transfer, failures, and resource reclamation.
• When changing hot paths, compare throughput, TTFT, ITL/TPOT, memory usage, and relevant synchronization/copy metrics under controlled conditions. Document measurement conditions and acceptable deviations.
• Extend existing tests and benchmarks where practical. This RFC itself does not claim that tests or performance regressions have passed.

Non-goals

This effort does not expand model coverage, parallelism scale, or deployment capabilities. It does not reproduce every vLLM directory, backend, process boundary, or class hierarchy. Structural changes must not silently alter exit policies, KV layouts, or PD transfer semantics.

Assign

Phase	Status	Assignee
M3	⬜ TODO	@bjf-frz / @YzXiao101
M2	⬜ TODO	@carsontung666
M4	⬜ TODO	@BuckMulligan99
M8	⬜ TODO	@Levius-Fubuki
M5	⬜ TODO	@wuli666
M7	⬜ TODO	@ycjcx123
M6	⬜ TODO	@Tarnished-Dan / @YZJF
M9	⬜ TODO	@YouHengfei
M1	⬜ TODO	@YZJF

References

• vLLM architecture overview
• vLLM V1 EngineCore
• vLLM V1 Scheduler interface
• vLLM V1 KV cache manager
• vLLM V1 KV connector interface
• Reference cleanup RFC and discussion

References to upstream main are architectural pointers, not immutable implementation baselines. Module-level reviews must pin their own reference revisions.

[bjf-frz]

The refactor will follow this order:
M3 → M2 → M4 → M8 → M5 → M7 → M6 → M9 → M1.

[bjf-frz]

M3 first increment: reusable refactoring reference

Implementation: Scheduler responsibilities, admission flow, and state contracts — PR #34, based on main commit 37575cff0493c2e82b5008807321064077f89d06.

The Scheduler walkthrough documents the motivation, priority and capacity examples, bounded bypass, batch construction, callback side effects, stage/admission/batch flowcharts, and function input/output contracts. The KV calculation note preserves the worked SHARED/LAST_EXITED examples. Both documents use technical English.

Refactoring type	Example in this change	Reusable review question
Responsibility separation	Extract refill/no-refill stage selection from admission and batch construction	Does one component mix policy, resource allocation, and execution?
State ownership	Policy owns fairness counters; NoRefillPolicy owns phase	Who initializes, advances, and resets each state? Is runtime state confused with configuration?
Type safety	FinishReason, NoRefillPhase, ResumeResult	Do strings or booleans hide a finite state or multiple outcomes?
Named data contracts	AdmissionPlan describes allocation inputs and block budgets	Are related values named, unit-aware, and explicit about lifetime?
Function interfaces and arguments	Explicit active_count, token_budget, frontier; keyword-only reserve_outputs; optional and enum returns	Are dependencies visible, boolean arguments readable, and zero distinguishable from failure?
Control-flow simplification and helper extraction	Admission and batch construction follow named business steps	Can the main path be understood before reading implementation details?
Decision/mechanism separation	Separate preemption eligibility and victim ranking from snapshot/copy operations	Can selection rules change without rewriting resource operations?
Unified cleanup	abort delegates to finish	Do termination paths duplicate cleanup or risk drifting apart?
Naming and units	selected_request_ids; watermark_ratio, watermark_blocks, admission_headroom_blocks	Does the name reveal purpose, scope, and measurement unit?
Rationale and walkthrough documentation	Priority, budget, bypass, batch-packing, callback examples and flowcharts	Are the reasons and failure paths documented, not just the happy path?
Regression validation	Fairness failures, cancellation cleanup, output serialization	Which invariants must survive the restructuring?
Explicit behavioral fixes	Empty prefill/recurrent batches do not alter fairness accounting	Is each behavior change identified separately and tested?

Applying this to later increments

For each module, record the original problem, final responsibility/state owner, input/output and side-effect contract, representative successful and blocked cases, migration impact, and validation evidence. Use the categories above where applicable; do not require every module to introduce policy classes or data containers. Keep structural changes and intentional behavioral fixes explicit and separate.

Function-interface review is part of the refactor, including argument meaning and units, keyword-only switches, enum/optional return contracts, and the distinction between caller configuration and internal runtime state. This increment renames the Python watermark keyword to watermark_ratio while preserving CLI --kv-watermark and public finish-reason strings.

The deliberate fairness fix counts only nonempty scheduled batches: empty prefill does not increment the counter, and empty recurrent does not reset it. CPU validation: 150 passed, 40 skipped; changed-file Ruff and formatting checks passed. GPU execution, quality, and performance remain unverified. This increment does not complete M3 or remove Scheduler/KV private access, Engine result progression, or PD private scheduling calls.

[Levius-Fubuki]

Hi @bjf-frz, I'd like to take ownership of M8 — Attention, starting with the required module walkthrough and interface proposal, followed by incremental implementation once the M4/M5 boundaries are agreed.

Relevant work: merged vllm#55629 (Triton attention, numerical validation and fallback coverage), merged vllm#55285 (BLIP-2 Q-Former SDPA), and ongoing vllm-omni#7895 (CUDA SDPA execution-path/capability contract).

Proposed scope:

1. Pin a baseline and review the attention dispatch/metadata code in core/kv_cache_manager.py, the Torch/Triton/FlashAttention paths, and their callers/tests. Document state writers, tensor/layout contracts, supported configurations, fallback/error behavior, and evidence-backed findings.
2. Agree on a small backend selection/capability and metadata interface with M4/M5: M4 retains logical allocation/depth validity; M5 retains device storage and buffer/stream lifetimes; M8 owns attention-specific metadata semantics and execution. Preserve LAST_EXITED/SHARED and existing numerical behavior.
3. Migrate in focused PRs, separating structural changes from bug fixes, with matched-input numerical regressions and controlled performance checks for affected hot paths.

I'll follow the stated refactor dependencies and coordinate with the existing inactive-row/buffer work (#27/#29) and speculative verification (#44). I am also volunteering for #43's evaluation feature, which can supply shared validation cases without expanding M8 into scheduler or speculative algorithm ownership.

Please record me as the M8 owner if this boundary fits, or let me know if someone is already working on it so we can agree on the split.

[BuckMulligan99]

Hi @bjf-frz, I'd like to take M4 — Logical KV management. Could you assign it to me?

Plan: (1) walkthrough with pinned baseline and state owners; (2) public interfaces to replace _refs / _get_allocation / _prepare_batch access from scheduler, preemption and runners; (3) split device storage out, coordinated with M5/M8; (4) rollback contract for speculative decoding and PD leases.

If there are any specific requirements or conventions I should follow, just let me know.

[bjf-frz]

@BuckMulligan99 @Levius-Fubuki @ycjcx123

You can refer to the M3 section, which lists several points we have already identified. These should be treated as a starting point rather than an exhaustive list. If you find any behavior that does not match the expected implementation, or identify opportunities for further performance improvements, feel free to make the necessary changes.

[YzXiao101]

Hi @bjf-frz, I'd like to take on a focused M3 follow-up: make Scheduler own the queue transition when a request is preempted.

Proposed scope:

1. Move removal from execution queues and re-entry into WAITING from PreemptionManager to Scheduler.
2. Keep victim selection, CPU snapshots, KV/device release and restoration in PreemptionManager; keep resume through the existing Scheduler.enqueue interface.

I've opened PR #50 with this implementation. Local validation: 245 passed, 114 skipped; Ruff and diff checks pass.

[carsontung666]

Hi @bjf-frz , I'd like to take M2. Could you assign it to me?

Preliminary plan:

1. Compare the sync and async step paths and record where submission, GPU completion, and CPU output delivery diverge.
2. Make both paths share one loop: schedule work, run it, then apply the results in one place.

I may revise this after reading the code more carefully.

[bjf-frz]

Hi @bjf-frz, I'd like to take on a focused M3 follow-up: make Scheduler own the queue transition when a request is preempted.

Proposed scope:

1. Move removal from execution queues and re-entry into WAITING from PreemptionManager to Scheduler.
2. Keep victim selection, CPU snapshots, KV/device release and restoration in PreemptionManager; keep resume through the existing Scheduler.enqueue interface.

I've opened PR #50 with this implementation. Local validation: 245 passed, 114 skipped; Ruff and diff checks pass.

Sounds good, RFC list has been updated.

You can draw the timing diagram of the schedule and update it in the document. Showing it in full may look quite complex. I suggest you understand it step by step, starting from simple cases and moving to complex ones. In the process of understanding and sorting it out, you may discover some pleasant surprises or other ideas to improve performance.

[bjf-frz]

Hi @bjf-frz , I'd like to take M2. Could you assign it to me?

Preliminary plan:

1. Compare the sync and async step paths and record where submission, GPU completion, and CPU output delivery diverge.
2. Make both paths share one loop: schedule work, run it, then apply the results in one place.

I may revise this after reading the code more carefully.

The first point looks good. For the second, rather than forcing sync and async into one loop, make them share the same scheduling/execution/result contracts and result-application framework, while preserving their different submission, completion, and delivery timing.

Look forward to your job, RFC list has been updated.

[Levius-Fubuki]

Thanks for the guidance. I have opened draft PR #46 for the M8 walkthrough and first backend/metadata extraction. I will use the M3 review points as a starting checklist, coordinate the remaining M4/M5 boundaries, and document and validate any behavioral fixes or performance changes separately.