popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

KC cache task

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 21:05:16 +02:00
parent 0ffd195fec
commit 436e872abe
3 changed files with 80 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
.scratch/alpha-hardening/README.md
View File

@@ -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)

48
.scratch/alpha-hardening/issues/25-per-node-kv-cache-distributed.md Normal file
View File

@@ -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.

23
.scratch/alpha-hardening/prd.json
View File

@@ -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"
}
}