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

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.