Skip to content

feat: connect fresh hybrid KV imports to ATOM PD scheduler #231

Closed

Description

@gspark-etri

Problem

AtomBatcher rejects hybrid Gemma prefill workers even though #225/#227 proved full-prefix RNGD/GPU KV can be imported and continued correctly in a standalone checker.

Evidence

Global layers use assigned KV blocks; sliding layers use lane-owned windows. Current AtomKVBlocks.import_span assumes every layer has global block indexing/full prefix capacity, and staging trusts source dtype instead of actual runtime allocation. Remote admission also must reset the decoder prompt mask.

Scope

Enable fresh full-prefix hybrid PD in the actual AtomBatcher using layer-specific destination/retention and validated physical staging size. Preserve rejection of unimplemented hybrid prefix/host restoration and partial-prefix imports. Add unit ownership/state tests and real Gemma HTTP PD validation, including repeated requests and cancellation/reuse where available.

Done when

The production scheduler accepts a full hybrid KV result, imports correct local/global state, resets only the admitted lane's mask, and matches native output in real PD requests. Unsupported restored/partial paths fail explicitly. Broader prefix/host cache support remains tracked under #201/#207.

Owner

Codex, maintainer review. Related #201, #207, #225, #227.

Activity

  1. gspark-etri commented on Oct 30, 2026

    @gspark-etri
    CollaboratorAuthor

    Fixed by #280, squash-merged into optimus-rm.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions