From 436e872abe1bacd7747ff6f9d60366d8743207a3 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Wed, 8 Jul 2026 21:05:16 +0200 Subject: [PATCH] KC cache task --- .scratch/alpha-hardening/README.md | 10 +++- .../25-per-node-kv-cache-distributed.md | 48 +++++++++++++++++++ .scratch/alpha-hardening/prd.json | 25 +++++++++- 3 files changed, 80 insertions(+), 3 deletions(-) create mode 100644 .scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md diff --git a/.scratch/alpha-hardening/README.md b/.scratch/alpha-hardening/README.md index 2d2651e..4f76bfa 100644 --- a/.scratch/alpha-hardening/README.md +++ b/.scratch/alpha-hardening/README.md @@ -11,6 +11,8 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati **Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`. +**Perf follow-up (2026-07-08):** [25 — Sharded per-node KV cache for distributed generation](./issues/25-per-node-kv-cache-distributed.md) is `ready-for-agent`. The ADR-0020 mixed-topology `start_layer` bug is fixed, but the distributed generation loop still has no KV cache at all — every step re-encodes the full sequence and re-runs every layer on every node, causing quadratic tps decay observed live (22.3 → 12.6 tps over one generation on a 2-node Qwen2.5-0.5B pipeline). Must be architecture-aware: Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V. + ## Artifacts | Path | Status | @@ -18,7 +20,7 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati | [research-verifiable-inference.md](./research-verifiable-inference.md) | Complete — SOTA research, §8 layered scheme, TOPLOC adopt | | [handoff.md](./handoff.md) | Session handoff — locked decisions, env notes | | [docs/adr/0016–0019](../../docs/adr/) | Alpha scope, auth, fraud, multi-tracker design | -| [issues/](./issues/) | 22 work items (Buckets 1–3) | +| [issues/](./issues/) | 25 work items (Buckets 1–3 + perf follow-ups) | ## ADRs (this feature) @@ -78,6 +80,12 @@ Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputati | [22 MEMORY + project-status index](./issues/22-doc-memory-project-status.md) (done) | | [21 Honest-noise calibration corpus](./issues/21-honest-noise-calibration-corpus.md) (ops; prod gate for audits) | +### Phase 5 — Distributed-inference performance (post-routing-fix) + +| Issue | Depends on | +|---|---| +| [25 Sharded per-node KV cache](./issues/25-per-node-kv-cache-distributed.md) | ADR-0020 routing fix (done), [24 routing telemetry resume](./issues/24-routing-telemetry-resume.md) | + ## First 3 to implement 1. **02 + 20** — Unified auth boundary + validator service token (shared helper and roles) diff --git a/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md b/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md new file mode 100644 index 0000000..fbdef1c --- /dev/null +++ b/.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md @@ -0,0 +1,48 @@ +Status: ready-for-agent + +Scoped 2026-07-08 from a live two-machine distributed-inference debugging session (Qwen2.5-0.5B GPU+GPU pipeline, and Qwen3.6-35B-A3B mixed GPU/CPU). The ADR-0020 mixed-topology `start_layer` bug is fixed (`518c259`, `e44abc9`, `1ecc599`); this issue is the next performance blocker in the same code path. + +# 25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware) + +## What to build + +The distributed generation loop (`torch_server.py:515-612`, `_do_chat_completions` distributed path) currently has **no KV cache at all**: `model_backend.py` passes `use_cache: False` in every layer-forward call (lines 763, 768, 770-771), and each autoregressive step re-encodes the *entire* prompt-so-far from scratch (`backend.encode_prompt(current_text)`), re-running every layer on every node in the route for every generated token. + +Observed cost of this on a live 2-node Qwen2.5-0.5B GPU pipeline (layers 0-20 / 21-23): tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation — the expected quadratic-cost signature. On the Qwen3.6-35B-A3B mixed-topology case this collapses to ~0.07 tps even after the routing fix, partly for this reason. + +`X-Meshnet-Session` already exists on the wire (`torch_server.py:707`, minted fresh **per token**, not per generation) but today only labels one activation transfer for chunk reassembly/logging — it is not used to key any cached state. + +| Subtask | Owner package | Deliverable | +|---|---|---| +| Session lifecycle | `packages/node/meshnet_node/torch_server.py` | Mint session ID once per chat request (not per token); reuse across all steps of that generation; add `X-Meshnet-Seq-Len` / position header so a node can tell prefill from decode steps | +| Per-node sharded cache | `packages/node/meshnet_node/model_backend.py` | `TorchModelShard` holds a `session_id → cache_state` map scoped to *its own* layer range only (naturally sharded — no node stores another node's KV); `forward_bytes` takes `use_cache=True` and returns/reuses `past_key_values` (or `use_cache=False` for the prefill token to keep failure/eviction simple) | +| Prefill vs. decode split | `packages/node/meshnet_node/torch_server.py` | Step 0 sends the full prompt activation (current behavior); steps 1+ send only the newest token's hidden state (`[1, 1, hidden]`) with correct `position_ids`, cutting per-step payload from O(seq_len) to O(1) | +| MoE / hybrid-attention state | `packages/node/meshnet_node/model_backend.py` | Cache abstraction must hold "whatever `use_cache=True` returns for this layer range," not assume standard K/V tensors — Qwen3.6's linear-attention/hybrid layers (see `[transformers] The fast path is not available...` warning already logged at startup) cache **recurrent conv/delta state**, not K/V pairs. MoE expert routing itself is layer-local and needs no cross-token cache, but confirm no expert-choice state leaks across the stateless-vs-cached boundary when `use_cache` toggles between prefill and decode | +| Cache lifecycle | `packages/node/meshnet_node/torch_server.py` | TTL + LRU eviction per node (bounded by `max_loaded_shards`/memory budget); explicit "cache miss" response so a restarted/evicted node causes the head to fall back to a full re-prefill instead of a hard error — keep today's fully-stateless path as the recovery mode | +| Correctness parity | `tests/` | Golden-output test: distributed multi-token output with caching enabled must match the existing stateless path token-for-token (or within sampling tolerance) for a fixed prompt/seed | + +**Non-goals for first landing:** cross-node cache migration/rebalancing on route change (evict + re-prefill is acceptable initially); speculative decoding; batching multiple concurrent sessions' KV within one node beyond what eviction already requires. + +**Code refs:** + +- `packages/node/meshnet_node/torch_server.py:515-612` — distributed generation loop (`current_text = current_text + token_str`, full re-encode every step) +- `packages/node/meshnet_node/torch_server.py:690-789` — `_run_downstream_pipeline`, session minting, `X-Meshnet-Session`/`X-Meshnet-Hop-Index`/`X-Meshnet-Start-Layer` headers +- `packages/node/meshnet_node/model_backend.py:189-201, 330-351, 763-771` — `use_cache: False` call sites, `effective_start` layer-slicing logic that any cache keying must respect +- `docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md` — prerequisite routing fix this issue builds on +- `docs/adr/0021-dynamic-statistical-routing.md` — route selection this cache must stay compatible with (a route change mid-generation should trigger cache-miss fallback, not corruption) + +## Acceptance criteria + +- [ ] A session ID is stable across all steps of one chat generation (not re-minted per token) +- [ ] Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes +- [ ] Each node caches `past_key_values`/recurrent state only for its own shard's layer range; no node ever holds another node's cache +- [ ] Cache works correctly for both standard-attention shards and Qwen3.6-style hybrid linear-attention/recurrent shards (cache abstraction is not K/V-shaped-only) +- [ ] Bounded memory: TTL + LRU eviction; eviction/restart triggers a documented cache-miss response, not silent corruption or an unhandled exception +- [ ] Golden-output regression test proves cached and uncached distributed generation produce equivalent output for a fixed prompt +- [ ] Measured tps improvement recorded on the same 2-node Qwen2.5-0.5B topology used to observe the regression (target: flat tps across generation length, not decaying) +- [ ] `tests/test_two_node_pipeline.py` and `tests/test_dynamic_routing.py` still pass +- [ ] Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction + +## Notes + +MoE routing (router + expert FFN) is layer-local per token and does not itself need a cross-token cache — it was ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the ADR-0020 `start_layer` double-execution). The MoE angle that *does* matter here is architecture-awareness in the cache design: don't hardcode a K/V tensor shape assumption that breaks on Qwen3.6's hybrid attention layers. diff --git a/.scratch/alpha-hardening/prd.json b/.scratch/alpha-hardening/prd.json index af339cc..f6a20e6 100644 --- a/.scratch/alpha-hardening/prd.json +++ b/.scratch/alpha-hardening/prd.json @@ -503,9 +503,30 @@ "notes": "Source issue: .scratch/alpha-hardening/issues/24-routing-telemetry-resume.md. Resume task for interrupted 2026-07-07 Claude session; first known fix is server.py:1490 annotation crash.", "dependsOn": [], "completionNotes": "" + }, + { + "id": "AH-025", + "title": "25 — Sharded per-node KV cache for distributed generation (MoE/hybrid-attention aware)", + "description": "Status: ready-for-agent\n\nScoped 2026-07-08 from a live two-machine distributed-inference debugging session. The ADR-0020 mixed-topology start_layer bug is fixed (518c259, e44abc9, 1ecc599); this is the next performance blocker in the same path. The distributed generation loop has NO KV cache at all: model_backend.py passes use_cache: False in every layer-forward call, and each autoregressive step re-encodes the entire prompt-so-far from scratch, re-running every layer on every node in the route for every generated token. Observed on a live 2-node Qwen2.5-0.5B GPU pipeline: tps decayed from 22.3 (at 235 output tokens) to 12.6 (at 449 tokens) within a single generation, the expected quadratic-cost signature. X-Meshnet-Session already exists on the wire but is minted fresh per token and only labels one activation transfer for chunk reassembly/logging, not keyed to any cached state. Build: (1) stable per-request session lifecycle instead of per-token, (2) per-node sharded cache keyed by session scoped to that node's own layer range only, (3) prefill-vs-decode split so post-prefill steps send only the newest token's activation, (4) cache abstraction that holds whatever use_cache=True returns per layer range (not K/V-shaped-only) because Qwen3.6's hybrid linear-attention layers cache recurrent conv/delta state, not standard K/V, (5) TTL+LRU eviction with an explicit cache-miss fallback to full re-prefill so restarts/route-changes degrade gracefully instead of corrupting output. MoE expert routing itself is layer-local and was already ruled out as the cause of the earlier Qwen3.6 garbage-output bug (that was the start_layer double-execution); the MoE angle that matters here is architecture-awareness so the cache design does not hardcode a K/V shape assumption that breaks on Qwen3.6's hybrid attention layers.\n\nSource issue has full subtask table and code refs.", + "acceptanceCriteria": [ + "A session ID is stable across all steps of one chat generation (not re-minted per token)", + "Steps after the first prefill send only the new token's activation, not the full sequence, over the wire between nodes", + "Each node caches past_key_values/recurrent state only for its own shard's layer range; no node ever holds another node's cache", + "Cache works correctly for both standard-attention shards and Qwen3.6-style hybrid linear-attention/recurrent shards", + "Bounded memory: TTL + LRU eviction; eviction/restart triggers a documented cache-miss response, not silent corruption", + "Golden-output regression test proves cached and uncached distributed generation produce equivalent output for a fixed prompt", + "Measured tps improvement recorded on the same 2-node Qwen2.5-0.5B topology used to observe the regression (target: flat tps across generation length)", + "tests/test_two_node_pipeline.py and tests/test_dynamic_routing.py still pass", + "Design captured in a new ADR (or an amendment to ADR-0020/0021) covering the cache-miss/route-change interaction" + ], + "priority": 25, + "passes": false, + "notes": "Source issue: .scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md. Perf follow-up to the ADR-0020 routing fix; no prior story covered KV caching or MoE-specific caching needs.", + "dependsOn": [], + "completionNotes": "" } ], "metadata": { - "updatedAt": "2026-07-07T21:30:00.000Z" + "updatedAt": "2026-07-08T19:15:00.000Z" } -} +} \ No newline at end of file