feat(pd): first-token router forwarding + P-side partial tail-block save

[xiaguan]

feat(pd): first-token router forwarding + P-side partial tail-block save

P-side half of the cross-engine GLM5.2 P/D deployment (vLLM TP8 prefill on one 8×H200 node, openinfer EP8 strict no-prefill decode on another; companion PR in the openinfer repo). Two pieces:

Connector: pegaflow.pd_tail_save

vLLM only hashes full blocks, so a prompt's partial tail block never enters the tier — a strict no-prefill decode peer would have to recompute it. When enabled, the step whose schedule finishes the prompt also saves the partial tail page under a key derived with vLLM's own hash function over (last_full_block_hash, tail_prompt_token_ids, None) — independently derivable by the decode peer. The page may already carry the racing first generated token's row; harmless, the key covers only prompt positions and the peer recomputes everything past them.

Trigger correctness (reviewed, tested):

• written = num_computed_tokens + num_scheduled_tokens straight from the scheduler's per-step output. The first version kept a connector-side accumulator; that lies across preempt/resume cycles and can save a torn tail page under a valid content-addressed key — and the tier's late-duplicate dedup would then reject the correct bytes forever. num_computed_tokens also covers the local prefix-cache hit case (every multi-turn turn ≥ 2 on P).
• Requests with lora_request, cache_salt, or multimodal features never tail-save: the tail key carries no extra_keys, so it would alias differently-salted prompts onto one key.
• Unit tests cover the trigger boundary, preempt/resume, the prefix-hit admission step, sub-block prompts (NONE_HASH parent), and the salt/lora refusal.

Router: --pd-first-token (variant A)

P is called with max_tokens=1 + return_token_ids; its single token goes back to the client and is appended to the token-id prompt D receives on /v1/completions — the decode node's first forward is a one-token decode, zero prompt positions computed. D failures before any bytes reach the client retry the whole P→D flow (the P leg is idempotent: same prompt, same content-addressed KV).

API-contract hardening from review:

• Explicit max_tokens on the D leg — engines default an absent value to 16, silently truncating every open-ended chat request (--pd-max-model-len bounds it).
• Chat-only field shapes (logprobs bool, tools, response_format, …) are stripped before the completions call instead of 400-ing deterministically through every retry; deterministic 4xx no longer consumes flow retries at all.
• n > 1 and batched prompts get an upfront 400 instead of a silent choices[0] splice.
• A decode stream that dies mid-flight emits an error SSE frame (pd_decode_interrupted) instead of a clean-looking EOF.
• Internal token-id fields are stripped from client-facing responses; in-flight counters no longer leak on parse failures.

Verified

Full-chain verification lives with the openinfer-side PR and docs/models/glm52/pd-m2-execution.md there: token-identical output vs direct vLLM, zero-prefill admission throughout, gsm8k/NIAH accuracy parity, failure injection (kill P / kill metaserver ⇒ clean errors + automatic recovery), and an equal-card multi-turn acceptance benchmark vs 2× mixed vLLM.

Follow-up (tracked)

Metaserver metadata is in-memory; on restart, pegaflow-server does not re-publish its block catalog, so previously saved prefixes stay undiscoverable until the producer re-saves them. The clean fix is a catalog re-publish on metaserver reconnect (registry-failover semantics).

🤖 Generated with Claude Code

[xiaguan]

Review side-finding (pre-existing on master, NOT introduced by this PR): the read-mode full-block save path has the same disease this PR's tail-save trigger was built to avoid — a connector-side accumulator that lies across preempt/resume and can save a torn block under a valid content-addressed key.

_scheduled_tokens[req_id] += num_tokens (read mode, scheduler.py) keeps accumulating across a preemption: resume only resets _allocated_blocks, never the accumulator. The write-only mode uses max(num_computed_tokens + num_tokens) and self-heals when num_computed_tokens resets to 0; read mode does not.

Repro timeline (vbs=16, prompt=50 → full-block hashes exist upfront for idx 0,1,2):

step 1  chunked prefill schedules tokens 0..31          acc += 32 → 32
        blocks idx 0,1 saved (content correct ✓)        next_stored = 2

step 2  KV pressure → PREEMPT
        physical blocks freed, num_computed_tokens = 0  acc NOT reset, still 32  ← the lie

step 3  RESUME with fresh physical blocks
        token budget this step is shared with ~100 running decodes,
        so the prefill chunk is truncated to 40 tokens (0..39, NOT block-aligned)
        allocated = ceil(40/16) = 3 blocks
        acc += 40 → 72
        _consume_full_block_saves:
          local_saveable   = min(len(allocated)=3, 72//16=4) = 3
          saveable_blk_idx = min(len(block_hashes)=3, 3)     = 3
          next_stored=2 → emits save for block idx 2 (tokens 32..47)
        but the GPU only computed up to token 39 this step —
        rows 40..47 of that freshly-allocated block are uninitialized memory
        → garbage bytes enter the tier under the legitimate content hash

step 4  tokens 40..49 scheduled, block idx 2 is finally fully written
        but next_stored already advanced to 3 → never re-saved;
        the tier's late-duplicate dedup rejects the correction forever

Impact: any later request that warm-hits that prefix (local read, cross-node fetch, or a decode peer) loads garbage KV for those rows — corrupted output, not a perf issue — and every hit refreshes LRU, keeping the poisoned entry alive until eviction/restart.

Trigger requires all of: read/read_write mode + preemption mid-prefill + a post-resume chunk truncated mid-block by the shared token budget. Narrow, but all three co-occur exactly under high load.

The fix is the same medicine this PR already applied to the tail: written = num_computed_tokens + num_scheduled_tokens is now available at the _consume_save_intent call site — threading it into _consume_full_block_saves as a cap (saveable_block_idx = min(..., base_block_idx + written // vbs)) closes it. Out of scope here per review discussion; filing as a follow-up.

[xiaguan]

Ran the repo's own vLLM correctness E2E gate on this branch (default connector path, pegaflow.pd_tail_save off) to confirm the _consume_save_intent refactor introduces no regression on the standard save/load path. Server binary and Python extension both rebuilt from branch head (version handshake enforced), single-GPU dev box:

cd python && uv run --extra test pytest -m e2e tests/test_vllm_e2e_correctness.py \
  --model Qwen3-4B --max-model-len 4096

test_execution_plan_outputs_match_baseline PASSED
test_cache_metrics                         PASSED
test_no_data_path_rpc_failures             PASSED
test_no_kv_load_failures                   PASSED
4 passed in 368.86s

CI lint failures (ruff C401 / rustfmt / clippy allow-without-reason) are fixed as of 3302677 — all checks green. Feature-on behavior remains covered by the cross-engine verification tracked with the companion PR.

[GentleCold]

LGTM