Skip to content

vllm-rbln

Reference: v0.11.3a21, an alpha. Its dependency is vLLM 0.26.0+cpu; the plugin's version is separate. Only the native RBLNScheduler and RBLNKVCacheManager full-attention paths are considered.

Tagged behavior

RBLNScheduler.schedule constrains a batch to decodes or a lone prefill. Accepting a waiting prefill can discard an already selected decode batch. DecodeBatchBudget separates a PP-dependent hard cap from a soft cap based on current demand. Selection, decode budget

The waiting path calls allocate_slots with full_sequence_must_fit=True: admission tests whether the current full sequence fits even when the allocation is only for the next chunk. Admission capacity and actual allocation are different quantities. Scheduler

Sub-block caching finds a prefix inside a larger physical block and creates copy operations into another block. apply_sub_block_match transfers source references to the copy operations; release_copy_ops releases them. Lookup granularity, allocation granularity and copy lifetime must remain separate. Cache manager

Compiled inputs and padding

The native runner compiles decode shapes up to the PP-stage ceiling max_num_seqs // pipeline_parallel_size, then rounds active decode requests up to the smallest configured bucket. Buckets may be linear, exponential or manual; a missing covering bucket is an error, not an arbitrary new shape. Runner bucket configuration, bucket lookup

For text inputs the runner stages a lone prefill with query dimension max_num_tokens, even when its actual chunk is shorter. Decode uses padded request rows and the step's uniform query length. InputStager fills dummy rows/positions and copies only the actual rectangle. Padding changes device work and buffers; it must not become extra generated/computed request tokens. Input layout, staging

determine_batch_execution_and_padding requires a uniform query length (num_tokens % num_reqs == 0). Its specialized DP path can force decoding peers to the top bucket and prefill token dimension when another rank prefills. This applies even with ordinary one-token full-attention decode; local phase isolation does not remove cross-rank shape coordination. Shape routes

For a local non-speculative model, a cost expression can charge fixed prefill width or bucketed decoders while the logical work stays unchanged. The current example uses illustrative costs and does not emulate these compiled shapes.

Executable serQ approximation

vllm-rbln deployment in serQ

examples/vendors/rbln.sq
// Executable approximation: one prefill or a decode-only batch.
// PP caps, remote-KV admission and sub-block copying are not modeled.
// The native RBLNScheduler keeps vLLM's prefix cache: a waiting request
// looks up its hit and the step's blocks are cached once it is final, and a
// failed allocation preempts the last running request under FCFS
// (docs/use-cases/rbln.md links the lines at tag v0.11.3a21).
use "../../lib/vllm.sq";
fn main() {
  device npu { kv cap 8192; }
  engine llm on npu {
    reqs cap 4;
    tokens cap 256;
    // What the RBLN scheduler does with one request, as vLLM does: prefill
    // what the cache does not hold and decode. A preempted request keeps
    // the tokens it generated (`kept`) and prefills them again.
    request (prompt, output) {
      set kept = prompt;
      at admission {
        known = max(kept, computed + 1);
        left = budget_left(llm);
        reused = min(cachedin(kv), reusable(known, blocksize(kv)));
      }
      set kept = known;
      set c = min(cached, reusable(known, blocksize(kv)));
      prefill (known - c) growing kv;
      decode (output - (known - prompt)) growing kv;
      cache (prompt + output);
    }
    schedule {
      exclusive prefill each at most (128);
      // the reused prefix plus a chunk of at most 128, if the whole
      // sequence fits
      admit waiting while (running.preempted == 0) each at most (128)
        taking reqs (1), kv (min(known, reused + min(128, left))) reserve (known);
    }
    execute (0.002 + 0.00001 * batch.tokens);
  }
  pool kv on npu { block 16; evict lru; preempt lifo; }
  pool reqs on llm { queue fifo; }
  workload {
    arrive batch(6);
    init { set prompt = 256; set out = 16; }
  }
  server {
    llm (prompt, out);
  }
}

exclusive prefill selects either one prefill or a decode-only batch. A fitting waiting prefill can replace tentative resident decodes and use the full token budget. Cancelled decode work does not advance computed KV; its already acquired allocation stays held. A selected prefill stops further waiting admission. Ordinary capacity, budget-exhaustion and preemption gates still apply. See the exclusive prefill reference.

reserve (known) tests the current sequence's capacity (the prompt, or what a resumed request had computed), while taking allocates only the initial chunk and growing kv extends it as computation advances. The live token budget is bound at admission, rather than captured before queuing.

The native scheduler keeps vLLM's prefix cache and preemption. A waiting request looks up its prefix hit (get_computed_blocks), the step's blocks are cached once scheduling is final (cache_blocks), and a failed allocation preempts the last running request under FCFS (running.pop()). The program says the same for a session's own prefix, with the hit bound at admission, cache (prompt + out) and preempt lifo on kv. RBLNScheduler hands a preemption to upstream vLLM (_preempt_request), which keeps the generated tokens and drops their KV, so a resumed request prefills from what it had computed (known, as lib/vllm.sq cites). In this workload neither fires: the six requests are separate sessions, and serQ's cache is keyed by session where vLLM's is shared by content hash, so none hits another's prefix; four slots of at most 272 tokens never fill 8192 tokens (512 blocks) of KV.

The example covers local phase isolation. PP caps, remote-KV decode-ready admission guards and sub-block copy semantics remain outside it. The hand-derived interpreter regressions do not establish native scheduler equivalence.

Remaining gaps

Requirement Needed refinement
PP hard/soft decode caps Per-request selection and batch constraints, separate from resident slot count
Sub-block hit and copy Independent lookup/allocation units and source/destination objects with copy leases

Reducing block to the sub-block size would also reduce physical allocation. It can match hit counts while predicting the wrong memory pressure.

Validation

this reduced example links and completes its six requests. All 44 iteration assignments were checked to contain either a lone prefill or decodes only. The interpreter policy also has six hand-derived regression tests. No native RBLN differential test has been run.