Pre-warmed multi-adapter relay on one frozen model. Three LoRA slices (extract → reason → verify) share one base
model; each slice keeps its own pre-filled KV cache, and stages hand off token IDs on the GPU. Forked from
ring-eval/lora_static_ring.py on 2026-11-02; LoraRing remains as an alias of Baton.
| Problem (found 2026-11-01) | Fix |
|---|---|
Newer StaticLayer.update writes at its own cumulative_length and ignores cache_position. Re-running _prewarm (needed after loading trained adapter weights) left the counter at 2 × prefix, so every later write landed outside the attention mask, and long inputs wrote past the buffer: CUDA device-side assert, then a GPU driver reset. |
_prewarm zeroes the counter. _snapshot_cache_state and every forward_ids call check the real write position equals prefix_len and raise RuntimeError before any GPU write if not. |
A system prompt + examples longer than max_cache_len overflowed the cache during construction (no guard). |
_prewarm raises ValueError if the prefix does not fit. |
python -m pytest -q tests
CPU only, on a tiny random Qwen2 model (needs the cached Qwen/Qwen2.5-1.5B-Instruct tokenizer; no weights). 13 tests,
about 13 s. The core test compares cached decoding against a no-cache reference decode, including after re-pre-warming
with changed adapter weights.
Mutation-checked on 2026-11-02:
- With the counter reset removed, the re-pre-warm test fails (
cache writes at 488, expected prefix_len 244). - With the reset and the position check both removed, three tests fail, and the cached output differs from the reference from the first token.
torch 2.11.0+cu128, peft 0.19.1, transformers 5.21.0.dev0 (local fork). Not pinned: the position check is what protects against cache-API changes, and it raises instead of corrupting output.
Pipeline.last_msis only updated when all cycles run out; an early gate pass returns without setting it.- Baton still relies on transformers'
StaticCache; owning the KV buffers directly would remove the dependency. next_token.item()syncs the GPU once per generated token.- Extraction is the accuracy bottleneck (62% parse rate at the 160-token cap in the ring-eval held-out run).
MIT