neuron-tai/.scratch/distributed-gguf-runtime/issues/02-stable-session-and-distributed-kv-pytorch.md

02 — Stable session and distributed KV in PyTorch path

Status: ready-for-agent

Goal

Fix the existing distributed PyTorch path so it does not recompute the full growing prompt for every output token.

Scope

• Introduce stable session_id for one request/session.
• Add per-node session cache keyed by session_id.
• Split /forward semantics into prefill and decode-step.
• Use model cache objects / past_key_values where supported.
• Keep hot KV local to each shard node.
• Add cleanup/TTL for abandoned sessions.

Current Problem

The current distributed path:

• calls encode_prompt(current_text) for every generated token
• sends full-sequence activations through the route
• calls layers with use_cache=False
• creates a fresh UUID inside _run_downstream_pipeline()

Acceptance

• Decode seam payload is one token / one hidden state after prefill.
• Per-shard cache grows locally with generated tokens.
• Regression test proves layer calls use cache after prefill.
• Fallback error is explicit for models whose manual cache API is unsupported.