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

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"]

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.

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

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

Refactor sequence and independence

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

Progress refreshed on 2026-10-22 against main at 4a350344027960a48384626aa646f1765fda6d4f. 🚧 means partial implementation or remaining integration/validation. Existing contributor claims are retained; merged increments do not complete every acceptance criterion of a module.

Cross-module status:

• [BUGFIX] Fix async speculative cleanup and serving failure isolation #115 async speculative cleanup/failure isolation is merged.
• [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118 implements compatible LAST_EXITED mixed prefill/decode recurrent batches ([RFC]: Mixed prefill/decode batching (last_exited) #81), but not multi-chunk wavefront scheduling ([RFC]: Wavefront Chunked Prefill Batching (WCPB) implementation #65). [FEAT] Implement Wavefront Chunked Prefill Batching(WCPB) #60 closed unmerged.
• Async speculative recurrent/CODA graphs and PD decode speculation exist on main. [FEAT] Add CUDA Graphs for cross-round async self-speculation #127 remains open for graph coverage/bucketing, vectorization and adaptive-draft increments; it must be reconciled with [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118/[Refactor] Move runner result application to Scheduler (M2/R1) #107.
• ci: add CPU lint and unit-test gate for pull requests #136 CPU lint/unit-test CI is merged. CPU CI does not qualify GPU correctness or performance.
• [RFC] Speculative Decoding for rlt #43 and [RFC]: Model support and feature roadmap for Ouro, Nanbeige4.2, and Huginn #69 track feature progress. Runtime/ownership references: unified execution, runtime.

No new hardware validation is claimed by this tracker refresh.

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.