[aoshen02]

Purpose

Split from #62044 (by @faresobeid): only the store eviction fix.

BlockObjectStore kept referenced and unreferenced objects in one LRU, and eviction scanned it from the oldest key, skipping referenced ones. Keys referenced by running requests pile up at the front, so every put() that had to evict rescanned all of them on the writer thread. #62044 measured about 9 ms per put with 300k keys, and reports that removing this rescan cut the writer thread from 281 s to 2.3 s over 20 RL steps on GLM-4.5-Air, for +15.6% throughput.

The store now follows BlockPool's free block queue: the LRU holds only unreferenced objects.

• _slots maps every stored key to its slot.
• A key leaves the LRU when its reference count goes from 0 to 1, and goes back to its tail when the count drops to 0. Within one put(), released keys end up after the batch's objects, in release order, as before.
• Eviction takes keys from the front of the LRU. As before, it skips keys stored by the same put() unless that call also releases them; they are at the tail by then.

With balanced retains and releases, eviction order, victims and errors are the same as before. Releasing a key that is not referenced raises KeyError, as before; keys released earlier in that call go to the tail of the LRU instead of keeping their old position, so later victims can differ. The writer thread stops after that error.

Test Plan

pytest tests/distributed/aux_output_connector/test_store.py

test_store_eviction_does_not_rescan_referenced_keys fills a store with 900 referenced and 100 unreferenced objects, then makes 100 put() calls. The referenced keys are a str subclass that counts its own hashes, so any dict or set lookup of a referenced key is counted, whatever container the store uses. The test also checks that the 100 unreferenced objects are the ones evicted and that every referenced object stays readable. It does not touch store internals.

Test Result

• test_store.py: 65 passed.
• The new test counts 0 hashes of referenced keys with this change and 180,000 on main (limit 900). It also fails on main with _references or _lru renamed.
• Differential fuzz against main (not part of the PR), 5,000 random sequences of put (duplicate keys, varying payloads, retains, releases) and get on stores of 1 to 10 slots:
  • balanced references: stored keys, allocated keys, reference counts and eviction order are identical after every operation;
  • with unreferenced releases: the store invariant holds after every operation; 404 sequences pick different victims after the KeyError.
  • Evicting from the wrong end of the LRU differs on the first sequence.

Benchmark

GLM-4.5-Air, one GB200 node (TP4), gpu_memory_utilization=0.9 (2,434,064 KV tokens, 167,341 store objects), max_model_len=20480, routed experts enabled. Each step is an RL-style rollout: 64 prompts x 8 samples (512 sequences), 1k-16k output tokens with ignore_eos, seeded per step so every run generates exactly the same 13,448,609 tokens over 3 steps. Both arms use the same 2026-10-16 nightly image and differ only in store.py; runs are ordered main, PR, PR, main. The store is set to 1.1x the KV-sized default so that main does not hit the full-KV crash fixed by #62044. Writer-thread time is the sum of BlockObjectStore.put wall time.

Run	Step 0	Step 1	Step 2	Throughput	Writer put time
main	510.5 s	487.8 s	425.5 s	9,446 tok/s	526.0 s
PR	417.2 s	415.0 s	414.4 s	10,788 tok/s	4.4 s
PR	417.5 s	415.1 s	413.9 s	10,789 tok/s	7.3 s
main	420.0 s	545.5 s	549.4 s	8,878 tok/s	617.7 s

Throughput is +17.7% (9,162 to 10,788 tok/s on average). On main, the step time follows the writer-thread time: about 420 s at 140 s of put, about 545 s at 240 s. With gpu_memory_utilization=0.5 (753,552 KV tokens, about 52k objects) the writer was busy for 17% of each step on main and throughput was unchanged (5,868 vs 5,897 tok/s), so the gain needs a store large enough for the rescan to delay the engine.

Reproduce

Run the script below once with main and once with this PR, on one node with 4 GPUs (max_bytes is 1.1x the default store size for this model and KV size):

import random
import sys
import time

from vllm import LLM, SamplingParams

WORDS = ["alpha", "beta", "gamma", "delta", "prime", "integer", "triangle", "sum",
         "product", "remainder", "circle", "radius", "matrix", "vector", "angle"]


def main():
    model, steps = sys.argv[1], int(sys.argv[2])
    llm = LLM(
        model=model,
        tensor_parallel_size=4,
        gpu_memory_utilization=0.9,
        max_model_len=20480,
        enable_return_routed_experts=True,
        aux_output_config={"max_bytes": 985_303_808},
        seed=0,
    )
    for step in range(steps):
        rng = random.Random(1000 + step)
        prompts, params = [], []
        for g in range(64):
            body = " ".join(rng.choice(WORDS) for _ in range(rng.randint(150, 900)))
            text = f"Problem {step}-{g}: {body}. Solve it step by step."
            for _ in range(8):
                prompts.append(text)
                params.append(SamplingParams(temperature=1.0, ignore_eos=True,
                                             max_tokens=rng.randint(1024, 16384)))
        t0 = time.perf_counter()
        outs = llm.generate(prompts, params, use_tqdm=False)
        dt = time.perf_counter() - t0
        gen = sum(len(c.token_ids) for o in outs for c in o.outputs)
        print(f"step {step}: {dt:.1f} s, {gen} tokens, {gen / dt:.0f} tok/s", flush=True)


if __name__ == "__main__":
    main()

python repro.py zai-org/GLM-4.5-Air 3

A single step with this script gave 459.3 s on main and 417.7 s with this PR for the same 4,543,356 tokens.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95851 for commit 262ceef39b2c.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 262ceef39b2c before starting CI. Pull the branch before pushing more changes.