Category LLMs

vLLM fundamentals – paged attention and batching

Archit Vora Leave a comment

KV cache problem

  • During demoing we store key and value of past tokens in cache
    • New token need to attend to them
  • Decoding is memory bandwidth bound. It reads
    • model weights
    • KV cache
    • just to get back one token
  • How to get good throughput
    • Batch many sequence in each step
    • Helps with single read of weights
  • One model weight is solved, KV cache is the bottleneck for throughput
    • The number of sequences whose caches fit in leftover GPU memory sets your batch size.
  • Before vLLM
    • Each request got one continuous chunk size for maximum possible sequence length
    • Issues
      • memory reserved for token that were never generated
      • gaps between chunk too small to reuse
      • request with same system prompt their own copy

Paged Attention

  • PagedAttention = KV cache in pages
  • OS idea
    • doesn’t allocate large chunk of RAM
    • Gets pages on demand
    • Makes RAM look bigger
    • Maintains mapping table
  • vLLM
    • splits KV cache on blocks
    • blocks are not contiguous
    • Each sequence gets a block table for mapping
  • Standard attention kernel will not work
    • key values are no longer contiguous
  • vLLM ships PagedAttentionKernel
    • follows block table to gather block
    • It’s a memory manager and a kernel
    • Some cost of indirection compensated by increased batch size
  • More advantages
    • Sharing is easy, system prompt, beam search
    • We can preempt some sequence, move them to CPU etc

Continuous Batching

  • Shortcoming of static batching
    • Some sequence finish early, they need to wait for large one to finish
  • Continuous batching makes scheduling decision at every decode step
    • step = generating one token
  • end token can leave the batch, KV blocks are feed, new sequence can join batch
    • Mad possible by pagedAttention
  • Once concern
    • Long pre-fill is compute heavy and can stall decoding
    • Solution – chunked pre-fill
      • mixing dunked pre-fill with token budge
    • newer version stopped distinguishing them and simply decided how many tokens each request gets to process
  • Knobs to tune
    • Max sequence per batch
    • Max tokens per step ( token budget )
  • Trade off / metrics
    • Bigger batch raise throughput but each step takes longer.
      • Increase TPOT – time per output token
    • TTFT – time to first token
      • depends on suing delay and how fast refill is scheduled

Just a python process

  • vLLM is user space python process
  • At startup
    • it measure GPU weight memory needed
    • allocates one large tensor for rest
      • after that plus within this tensor
  • OS and GPU drives just see big allocation that never changes
    • They don’t know about block tables
  • Scheduling and block management runs on CPU
    • At each step
      • Engine packs blocks tables into tensors and sends the to GPU
    • Downside
      • For small model CPU overhead can become a bottleneck
        • Newer vLLM address this by running engines core in a separate process from the API server
  • PagedAttentionKernel follow that map

On RL Environments

Archit Vora Leave a comment

Where environments fit

A modern LLM goes through roughly four stages of training:

  1. Pre-training. Next-token prediction on trillions of tokens of web text, books and code. The model picks up knowledge and language.
  2. Supervised fine-tuning (SFT). Pairs of prompts and ideal answers. The model learns to follow instructions and formats.
  3. Preference tuning (RLHF). Humans pick the better of two answers, a reward model learns those preferences, and the LLM is optimized against the reward model.
  4. Reinforcement learning on verifiable rewards. The model attempts tasks whose outcome can be checked automatically, and is rewarded when it succeeds.

The fourth stage is the new one. Around late 2024 and early 2025, OpenAI’s o1 and DeepSeek-R1 showed that rewarding only checkable outcomes, like a correct math answer or a passing unit test, teaches models to reason longer and better. The natural next step was to apply the same idea to multi-step work: fixing a bug in a real codebase, operating a CRM, or preparing a report from data spread across several apps.

To train on that kind of work, you need somewhere for the model to act and a way to score what it did. That somewhere is an environment.

What vendor ship to labs 

task-0417/
  task.yaml         instruction + metadata
  Dockerfile        repo pinned at a commit, deps installed
  tools.json        what the model is allowed to call
  tests/            hidden from the model
    test_fail_to_pass.py
    test_pass_to_pass.py
  solution.patch    reference fix, proves it's solvable
  grade.sh          runs tests, returns 0 or 1
One environment is one folder. A delivery is hundreds to thousands of these.

What model sees


instruction: |
  Invoice totals are off by one cent when a line
  item ends in half a cent. 3 x 2.345 shows 7.03,
  expected 7.04. Fix it without changing the API.
repo:      invoicekit @ a91c3e2
tools:     [bash, read_file, edit_file]
max_steps: 40

# metadata, not shown to the model
category:   numeric bugs
difficulty: medium
reviewed:   2 engineers

Model’s attempt

The model only sees the visible tests in the repo. It has no idea what the hidden tests check.

> bash grep -rn "round(" invoicekit/
  totals.py:18:  return round(subtotal, 2)
> read_file invoicekit/totals.py
> edit_file invoicekit/totals.py
  - return round(subtotal, 2)
  + return Decimal(str(subtotal)).quantize(
  +     Decimal("0.01"), rounding=ROUND_HALF_UP)
> bash pytest tests/ -q
  42 passed
> done

Grading (Verifiers )

The grader labels every attempt, no human involved. Line 1 exists because models learn to edit or delete tests when they can.

grade.sh
  1. restore tests/ from sealed copy
  2. test_fail_to_pass.py    3/3 passed     bug fixed
  3. test_pass_to_pass.py    118/118 passed nothing broke
reward = 1

Calibration before shipping 

Run a current model 16 times on each task:
task-0417  5/16
task-0418 16/16
task-0419  0/16
0417  31%  keep     sometimes right = learning signal
0418 100%  drop     already solved, nothing to learn
0419   0%  inspect  too hard, or the task is broken?

What the lab receives

Why do vendor provide hold out set

  • Lab can leak the near twins with random sampling
  • Vendor will share uplift number using open weight model
  • In practice, the vendor proposes the split and ships the group labels, and the lab decides. Labs don’t fully trust vendor evals either. 
delivery-2026-09/
  tasks/          800 folders like task-0417
  report.md
    category mix       bugs, migrations, perf, ...
    baseline pass@1    per category
    removed tasks      flaky, ambiguous, hackable
  heldout/        200 tasks, eval only, never trained on
# inside the lab's training loop
for step in training:
    task    = sample(tasks)
    tries   = model.attempt(task, n=16)
    rewards = [grade(t) for t in tries]
    update(model, tries, rewards)

How does Rubric Look like

  • The example below outputs score between 0-1 and uses 1) code and 2) LLM as judge 
task: Q3 travel-expense MIS for the CFO, from Drive + Slack + email
GATES (fail any -> score 0)
  email sent to cfo@acme.test with MIS attached        
  total travel spend = 48.2L (+/- 0.5%)                
CRITERIA (points)
  +3  flags the 2 expenses approved without sign-off   [judge]
  +2  department breakdown matches the Drive sheet     [judge]
  +2  compares to Q2, explains the biggest change      [judge]
  +1  one page, readable by a CFO                      [judge]
  -3  states any number not traceable to source data   [judge]
  -2  emails anyone besides the CFO                    
score = max(0, points earned) / 8