Repository navigation
fix: initialize hybrid ATOM KV tables and per-request masks - #210
Draft
gspark-etri wants to merge 7 commits into
Draft
gspark-etri wants to merge 7 commits into
gspark-etri wants to merge 7 commits into
Conversation
The exported graph ties position_ids' batch dim to input_ids'. Prefill passed position_ids [1, S] and decode [1, 1], and decode fed the [B] sampled tokens as [1, B]. Prefill now builds [B, S] positions, decode [B, 1] positions and a [B, 1] token input. Adds tests/test_kv_cache_batch.py: greedy KV-cache generation matches HF at B=1 and B=2 (GPU only). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit d46b8ff)
save_ckpt named files by pipeline stage only, so with tp > 1 the TP
ranks overwrote each other's shards, and all DP replicas wrote and
pruned the same files. save_lora_ckpt already named TP shards but let
every DP replica write. Following save_hf_ckpt, only DP rank 0 writes
and save_ckpt/load_ckpt add a _tp{rank} suffix when tp > 1 (PP-only
names unchanged). The checkpoint directory no longer becomes
checkpoint_None_<stage> without a postfix (the "or" fallback never
triggered).
Adds tests/test_ckpt_names.py.
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
(cherry picked from commit fb09a61)
) _mps_daemon_running fell back to a host-wide pgrep, so another job's MPS daemon (other pipe directory) looked like ours; the health probe on our pipe dir then failed and the stale cleanup ran pkill -u $UID on every MPS server and control process, killing the other job's daemon. Liveness is now our control socket only, and the stale cleanup sends quit to our pipe dir and removes our socket, without pkill. Adds tests/test_mps_manager.py with a fake pipe dir. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 0e69cc9)
fix: v4.1.2 hotfix - KV cache at batch >= 2, TP/DP checkpoint names, MPS daemon scope
This was referenced Oct 23, 2026
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
With Gemma3 hybrid attention and external KV blocks, AtomBatcher supplies no local sliding-cache table, so the first request fails. Providing that table without resetting the prompt mask also makes repeated outputs depend on the previous request.
Changes
Pass lane-local cache indices alongside global block indices and reset the owning lane’s prompt mask after fresh prefill. Reject hybrid prefix/host caching and remote prefill before SDK calls until sliding-state restoration exists. Add ownership, repeated-request, staggered cancellation/refill and rejection regressions.
Verification
Baseline optimus-rm0e69cc9d54d19e336d33920298c39507d1415811; core candidate6b5d8f3. Four initial regressions fail before the change. Afterward20 tests pass;2 optional suites skip because hetero_p2p/rngd_pd are unavailable locally. New test file passes Ruff; changed directories pass sensitive-data scan and diff whitespace checks.
One allocated ATOM, rebel-compiler0.11.0/optimum-rbln0.11.0.post1, Gemma3-1B B1/context9216/three blocks:16 runs through8192 input plus128 output complete;8/8 repeat pairs exact and16/16 outputs match the experimentally validated bridge in PR209. That bridge separately matched full32-step logits with SDK-owned tables at four input lengths. Setup, immutable model revision, commands and raw outputs are in docs/validation/batching/runtime207.
Compatibility and limits
Targets optimus-rm, which owns the engine. Hardware validation requires the Gemma causal-prefill SDK correction (upstream PR827) and rebuilt artifacts. B>1 has unit coverage, not Gemma hardware coverage. Hybrid prefix/host cache and PD remain unsupported and fail explicitly. Cross-backend numerical/quality differences remain open; no general hybrid-model or serving-performance claim.
Related Issue
Related to #207 and #201; #207 remains open for sliding-state restoration. PR209 contains the exploratory evidence. Generated with Codex; maintainer review and merge required.