Skip to content

feat: support fresh hybrid KV imports in ATOM PD scheduler - #232

Draft
gspark-etri wants to merge 3 commits into
fix/229-serving-eosfrom
feat/231-atom-hybrid-pd
Draft

gspark-etri wants to merge 3 commits into
fix/229-serving-eosfrom
feat/231-atom-hybrid-pd

Conversation

@gspark-etri

@gspark-etri gspark-etri commented Oct 23, 2026 •

Copy link
Copy Markdown
Collaborator

Background

Gemma hybrid caches were rejected whenever a remote prefill worker was configured. Fresh remote prefixes need different placement for sliding and global attention, and source dtype metadata can disagree with physical KV storage.

Changes

  • Enable full-prefix hybrid PD imports, place sliding tails by lane and global KV by assigned block, and validate geometry and physical storage.
  • Reset the imported lane's decode mask; fail partial-prefix results per request while retaining lane reuse.
  • Add CPU regressions and a real ATOM HTTP validator using recorded RNGD KV exports.

Verification

  • Base f84823a; hardware-validated change c76e03a; recovery follow-up 1e54786. PYTHONPATH=opt_prime python -m pytest -q opt_prime/tests/test_atom_hybrid.py opt_prime/tests/test_atom_batching.py opt_prime/tests/test_pd_recovery.py opt_prime/tests/test_serve_batched.py: 37 passed, one existing warning. Initial regression failed before the fix; partial-result and malformed KV geometry recovery verified red/green. Hardware/runtime errors still poison the scheduler; only validation errors before device writes are request-scoped. The exception follow-up also passed the long-input hardware rerun described below.
  • One allocated ATOM, Gemma3-1B B1/context9216/k3, corrected experiment: isolate Gemma native global-prefill causality discrepancy #225 artifact, pinned checkpoint 5b11413a10db4e486ef16a20101fd028f8f2499c. Four questions at short/medium lengths, two passes, cancellation/disconnect and reuse: 18 normal/reuse responses exactly match native tokens through EOS; 20 remote / 0 local prefills including cancellations. Final free lane [0].
  • Latest hardware follow-up on 1e54786 (evidence commit69bca73): four6423–6428token inputs, two passes, cancel/disconnect and reuse:10normal/reuse responses match native ATOM,12remote/0local prefills. This remains recorded RNGD KV replay. The shared arithmetic answer25 is wrong, so this proves behavior parity, not quality.
  • Reproduction commands, selected JSON and experiment record under opt_prime/examples/pd/http_matrix/.
  • Scoped Ruff and diff check pass. Sensitive checker: changed files clean; four unchanged baseline findings in unrelated documentation.

Compatibility and limits

Stacked on #230. Fresh full-prefix PD only; hybrid prefix/host cache restoration remains rejected. Hardware producer is recorded RNGD KV replay, not live prefill. No overlap, hardware batch >1 or speedup claim. Vendor SIGPIPE backtraces on socket disconnect persist despite successful reuse; separate #214/#215 addresses that issue. Corrected Gemma artifact depends on upstream optimum-rbln #827.

Related Issue

Closes #231. Parent validation #201/#227 remains active.

gspark-etri added a commit that referenced this pull request Oct 30, 2026
Lane-local sliding cache tables and global block tables for hybrid (Gemma3)
external-cache serving; fresh full-prefix hybrid PD; batch-1 single-block PD
with prefix/host caching off. Hybrid prefix and host caching stay rejected.

From feat/236-atom-single-pd (#210, #232, #237): atom/batching.py,
tests/test_atom_hybrid.py, tests/test_atom_batching.py. ATOM.md section
rewritten to the supported scope.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant