-
-
Real 4090 log replay, task 454: compute, cache reuse and H1 canonicalization across separate generated sessions.
-
Fixed four-turn cost sweep with untrained adapters: 1k to 8k shared context, five repeats per size. This measures cost, not quality.
-
All six methods on 100 MBPP tasks: computed tokens versus multi-turn final task success. Equivalent quality is not established.
-
RoleKV homepage: shared KV-cache visualization with a recorded, fixed four-turn cost replay.
Inspiration
Multi-role coding agents (a coder, a reviewer, a planner) keep re-processing the same context. A popular cheap setup is one small open model with a LoRA adapter per role. The roles share almost all their weights, yet every role switch throws away the KV cache and re-prefills the task, the repository code, the test output and every earlier turn. As sessions grow, most of the GPU time before each answer goes into re-reading text the model has already read. We wanted to know whether that waste is avoidable without giving up the benefit of separate role adapters.
What it does
RoleKV trains each role as an Activated LoRA (aLoRA): an adapter that only switches on after the role's header. Everything before the header is processed with plain base weights, so the Coder and the Reviewer share one KV cache. A role switch only computes the new tokens. When a role hands off, we recompute its short output with base weights ("canonicalization", H1), which keeps the shared cache exactly what each role saw in training.
We compare this against per-role standard LoRAs that re-prefill at every switch (A2), a naive cache-reuse baseline (A3), a single shared LoRA (A1), the base model (A0), and an aLoRA variant that skips canonicalization (H2), all on MBPP coding tasks where code is checked by real unit tests.
The dashboard makes the cost visible:
- Session replay plays a Coder ↔ Reviewer session turn by turn, showing how many tokens each method recomputes, reuses or canonicalizes.
- Cost vs quality puts every method on one chart: prefill tokens per task against final task success.
- Context sweep shows how the savings change as the shared context grows from ≈1k to ≈8k tokens.
The dashboard computes its headline numbers in the browser from the loaded event logs.
How we built it
- Model: Qwen2.5-Coder-1.5B-Instruct (0.5B for CPU development and tests), trained with Hugging Face PEFT
(
peft==0.21.2, aLoRA viaalora_invocation_tokens) on PyTorch + transformers. - Runtime: a shared-session engine that keeps one KV cache, switches adapters per role, crops and recomputes the cache at H1 handoffs, and logs prefill, reuse and canonicalization tokens per turn.
- Data: MBPP. Teacher dialogues (Claude) for training, with every tool output coming from actually running the tests, never invented. 280 verified trajectories: 338 Coder rows and 280 Reviewer rows (222 APPROVE, 58 REVISE, with the REVISE rows oversampled ×3 in training).
- Correctness tests: cache equivalence against a full aLoRA forward pass, canonicalization equal to a fresh base prefill, adapter-on during decode, and a hand-computed cost-accounting check.
- Dashboard: a static HTML/CSS/JS page with hand-drawn SVG charts and no build step or dependencies, designed first in Paper. It validates every log line against our frozen event-log contract, audits token accounting, and keeps mock data from ever mixing with real results.
Challenges we ran into
- Exact cache semantics. aLoRA's activation point is computed from the input ids of each forward call. In a hand-written decode loop, the adapter silently turns off unless you pass the offsets explicitly. We wrote a test specifically for "adapter on during decode".
- Tokenization boundaries. The role header is the adapter's trigger, so it must tokenize identically in training and at inference. We tokenize every prompt segment separately and never let BPE merge across a boundary.
- Handoff policy. The previous role's output was computed with its adapter. Deciding whether to recompute it (H1) or keep it (H2), and accounting for that cost honestly, shaped the whole runtime.
- Honest cost accounting. It's easy to make savings look bigger by forgetting the canonicalization tokens. Our logs count them, and the dashboard audits every task total against its turns.
Comparable timing. We trained and evaluated across laptops with different GPUs; timing is compared only on one RTX 4090, and token counts (deterministic) are the primary metric.
Library overhead that looked like a result. On our first 4090 run, aLoRA decoded 15% slower per token than plain LoRA, so short sessions took longer with RoleKV. The overhead came from PEFT's aLoRA path, not from aLoRA itself:
- a boolean-mask row selection that waits for the GPU in every layer;
- per-call hook registration that walks the whole model.
We replaced them with an equivalent slice and persistent hooks, both checked against PEFT's own output. That brought decoding to within 1–4% of plain LoRA. Moving every adapter to bf16 (next point) then made plain LoRA cheaper, so the Python overhead left in PEFT's aLoRA path shows as 8% slower decoding. That overhead counts against RoleKV, not for it.
- A default that flattered us. PEFT upcasts adapter weights to fp32. Plain LoRA applies its adapter to every token, so the upcast slowed it far more than aLoRA, which is active from the last role header onward: +39% vs +7% on a 4k-token prefill. We run every method's adapters in bf16. That made our latency ratios smaller, and fair.
- Training data that didn't match eval. The generated SFT data used a different system prompt than the runtime. We re-exported it without new API calls, and a test now guards the data contract.
Accomplishments that we're proud of
- In a fixed 4-turn session, A5-H1 computes 3.6× fewer tokens than A2 at ~1k tokens of shared context and
3.9× fewer at ~8k. On the RTX 4090, its first token after a role switch arrives in ~50 ms at every size,
while A2's grows from 94 to 583 ms: 1.9× to 11.7× faster. Source:
bench/results/cost_sweep_4090.md. This sweep forces the same conversation with untrained adapters; it measures cost, not coding quality. - In the eval, A5-H1 computes 3.2× fewer prefill tokens per task than A2: 2,197 vs 7,013, 68.7% fewer, on the
same 100 MBPP tasks with ≈1.8k tokens of shared context. Its role switches are 3.1× faster (median TTFT 45 vs
139 ms). Source:
eval/results/eval-1.5b-4090-202610180231.jsonl. - Final task success is 61/100 for A5-H1 vs 59/100 for per-role LoRAs (A2). Paired by task, only A2 passes 8 and only A5-H1 passes 10. The exact paired test gives p = 0.81; this does not establish equivalent quality.
- Every runtime correctness test passes on the Mac CPU in fp32: cache equivalence, handoff canonicalization, adapter on during decode, and cost accounting checked against an independent embedding-layer counter. On the RTX 4090 the same 60 tests pass in fp32 on CUDA.
- A dashboard that makes a KV-cache idea understandable in seconds, and refuses to show unverified numbers as results.
These are multi-turn results with feedback from all MBPP asserts after each eligible Coder turn, not standard held-out pass-at-one. Different generated histories and stopping behavior contribute to session costs: A5-H1 averages 2.94 turns versus A2's 3.27. A false APPROVE can end a session before another repair attempt; we have not measured the counterfactual quality effect.
What we learned
- Cache sharing is only as good as its correctness: a single silently-disabled adapter would have produced great-looking savings and wrong outputs. Equivalence tests had to come before any experiment.
- The savings come from the shared context, so they depend on how long it is relative to what each role writes. In a fixed 4-turn session, the token saving is 3.6× at ~1k tokens of shared context and 3.9× at ~8k. It approaches the number of turns, because per-role LoRAs re-read the context once per turn. Switch latency diverges much faster. A2's grows linearly with the context (94 → 583 ms from 1k to 8k tokens), while A5-H1's stays at ~50 ms. Whole sessions are dominated by decoding, so they gain less: from 4% slower at 1k to 13% faster at 8k.
- H1 vs H2: 61 vs 60 tasks passed. The H1 handoff costs 138 extra tokens per task; these counts do not establish equivalent quality. aLoRA uses r=32 and standard LoRA r=16.
- One shared LoRA (A1) passes 61/100 and mismatched KV reuse (A3) passes 62/100. Both compute fewer tokens and have lower median switch TTFT than H1. We did not demonstrate a task-success, token-cost or switch-latency advantage over these baselines. Exact H1 cache semantics and separate role adapters remain the design contribution.
- H1's Reviewer approves 41 of 45 second reviews, including 35 whose final code fails; A2 has zero such false approvals. The training data lacks second reviews, but that does not prove the cause. A separate diagnostic matched an independent PEFT reference on 12 fp32 inputs; it does not establish bf16 equivalence or a training cause. The team froze the evaluated v1 and deferred repair. Multi-round generation and pilot validation tooling is implemented separately; no repaired adapters were trained or evaluated.
Hypotheses we did not test are under What's next.
What's next
Future work, not done in this hackathon:
- Evaluate whether multi-round training improves second-review behavior, using a separately versioned dataset and new adapter evaluations. The existing results remain the frozen v1 baseline.
- More roles (a planner) and longer agent loops, where switches are more frequent.
- Larger models (7B) and repository-level tasks (SWE-bench-style issues).
- Head-to-head comparison with inference-time cache-sharing methods for standard LoRAs (e.g. LRAgent, DroidSpeak-style partial recompute).
- Serving with an aLoRA-aware engine and measuring throughput under concurrent sessions.
- Training with simulated H2 contexts to see whether the cheaper handoff can match H1.
Built With
- claude
- javascript
- mbpp
- paper
- peft
- python
- pytorch
- qwen2.5-coder
- svg
- transformers
Log in or sign up for Devpost to join the conversation.