Skip to content

Proposal: streams take memory from one core pool as they grow, instead of a full window each at load #781

Closed
Closed
Proposal: streams take memory from one core pool as they grow, instead of a full window each at load#781

Description

@jschmied
jschmied
opened on Oct 13, 2026

What happens now

--parallel N reserves N whole --context windows at load, and each engine sizes them its own way:

So a 2k-token chat holds as much memory as a 128k agent session, the slots go to the prompt cache only after
every stream got its worst case, and each new family writes its own answer to "how many streams fit".

Proposal

One owner for stream memory in core, every family plugged into it:

  1. A family declares its stream layout instead of a byte total: fixed bytes per stream (recurrent states,
    Mamba/GDN, drafter state, graph buffers) and bytes per token for each paged cache (attention K/V), plus its
    page size.
  2. Core keeps the book of one pool of pages (ids, refcounts, reservations); each engine keeps the device side
    (tensors, page tables). sessions: session memory for paged-KV families (page pool, prefix trie, RAM residency budget, streamed NVMe tier) (DeepSeek-V4.1 stack 4/14) #637's sessions.Pool already splits it this way and has most of it (several caches
    per page, rows per token, dirty-range table uploads). It would implement Pages, the interface
    prompt_radix uses. Lanes admit a request when its fixed part plus the pages for prompt + max_tokens fit.
    The prompt cache's kept pages come from the same pool, so evicting cached prompts is the first answer to a
    full pool.
  3. The family's attention kernels read K/V through a per-stream page table. Families that already page (feat(core): a paged prompt cache, a radix tree over the backend's own KV pages #625's
    HIP engine, sessions: session memory for paged-KV families (page pool, prefix trie, RAM residency budget, streamed NVMe tier) (DeepSeek-V4.1 stack 4/14) #637) map onto it; the others change their cache addressing once.
  4. --parallel N becomes a cap on streams, not a reservation. Growing a stream past its admit reservation, and
    pausing the youngest stream when the pool runs dry, would be new in every engine; a later step, not the first.

What stays as it is

  • Exactness: a stream's bits do not depend on which pages it holds, so drafted == plain and concurrent == solo
    must still hold; the tests that check them now would check this too.
  • Captured graphs: a graph keeps fixed addresses for the page table and the state slots, not for the pages.

Questions before any code

We can take the CUDA side (Nemotron first, then the shared attention modules in #746) once the shape is agreed.

Edited: #625's Pages is an interface, not a pool; neither paged pool grows a stream past its admit reservation.

Activity

  1. ashhart commented on Oct 14, 2026

    @ashhart
    ashhart
    on Oct 14, 2026
    Owner

    Hello @jschmied,

    Thank you for writing the questions before any code.

    On main, CUDA prices every stream at the full window. Nemotron's seqBytes is the device bytes of one sequence at max_len (zig/src/families/nemotron/cuda_engine.zig:309), and admit multiplies that by the stream count (zig/src/native/cuda.zig:340). Qwen3.8-27B sizes each slot from the full context (zig/src/native/qwen27_served.zig:257). GLM does the same unless its pool is on, in which case the per-stream state drops the token cap (zig/src/native/glm_host.zig:45). prompt_cache.Pages is the interface a paged store would call (zig/src/core/prompt_cache.zig:62). There is no DeepSeek family on main to bring a page pool with it.

    Whether core owns the pool, and whether Pages is the base, is the maintainer's call. A CUDA PR that starts with Nemotron, after that shape is agreed, is the piece that helps.

    Thank you again.

  2. ashhart commented on Oct 15, 2026

    @ashhart
    ashhart
    on Oct 15, 2026
    Owner

    Thank you, @jschmied, for separating fixed stream state from token pages and putting cached prompts under the same memory accounting.

    We'll agree the interface first, then pilot it with Nemotron on CUDA. Please send the layout and reservation contract, including ownership through the existing Pages interface and how cached prompts are evicted. We'll review that contract before the pilot implementation. The first slice reserves prompt plus max_tokens at admission and checks page lifetimes, allocation failures, drafted versus plain replies, resumed versus fresh prompts, and concurrent versus solo requests. Growth beyond the reservation and preemption remain a separate decision.

    Posted by Ash's AI Assistant Synthena

  3. jschmied commented on Oct 15, 2026

    @jschmied
    jschmied
    on Oct 15, 2026
    ContributorAuthor

    @ashhart, here is the contract, checked against main at c7137bb. Nemotron numbers: 23 Mamba and 6 attention layers, 2 KV heads x 128, one MTP layer.

    1. Layout

    A family declares fixed_bytes (per stream, any length), token_bytes (all paged caches and layers), page_tokens and snap_bytes (one kept recurrent snapshot's slot, section 4). Nemotron:

    field contents bytes
    fixed ssm 46.0 MiB + conv_base 0.8 MiB (mambaPlanes, cuda_state.zig:70), raw/xc/dt verify-window replay rows 17.4 MiB, io rows, the head's chain words (cuda_state.zig:75-82, cuda_mtp.zig:137), page table ~64.3 MiB
    token K+V of 6 attention layers (6 KiB) + the MTP head's K+V (1 KiB); kvRowBytes = 512 B (cuda_state.zig:65) 7 KiB
    page 512 rows = chunk_keys, the attention kernels' fixed-position key chunk (cuda_state.zig:15) 3.5 MiB
    snap ssm + conv_base (what cuda_reuse.zig:45 already copies), 49,082,368 B, kept in a slot of whole 2 MiB driver pages (allocBytes, cuda_reuse.zig:40) 48 MiB a slot, outside the pages

    Core computes pages(n) = ceil(n / page_tokens). It also computes the pool, which is the room left after the weights (today's room, native/cuda.zig:338) minus --parallel x fixed_bytes, minus the snapshot slots (Limits.snaps x snap_bytes). At a 128k window, one stream costs 64 MiB plus 899 MiB of KV today, all of it up front (seqBytes, cuda_engine.zig:306). Under this contract a 2,048-token prompt with the default 4,096 max_tokens reserves 13 pages (45.5 MiB): 2,048 + 4,096 + 16 = 6,160 rows. A 2,000-token prompt fits in 12.

    In the pilot, the fixed part of every lane is allocated at load (N x 64 MiB, N <= 8, cuda_engine.zig:68). Graphs keep fixed addresses. The engine stores the KV planes as [page][layer][K|V][row], and Forward.cache (cuda_forward.zig:103) becomes a page-table lookup per 512-key chunk.

    2. Reservation

    • Need: pages(prompt + max_tokens + max_rows). max_rows is 16, the verify window (cuda_state.zig:11). Nemotron checks the same bound today (cuda_lanes.zig:241); GLM reserves it (glm/backend.zig:68). max_tokens is the server's value after it clamps to the window (server/chat.zig:100-122).
    • Check: Info.admit.fits (engine_api.zig:166), called by lane_host_admit.pick under the host lock. A request fits when available() + reclaimable >= need. Pages adopted from a cache hit still count toward the need, so the check is conservative.
    • Taking the pages: the backend's prefill promises need - adopted pages after lookupRewind/begin (lane_host.zig:312), as GLM does (glm/backend.zig:97). If the promise fails it returns error.LaneMemory, and the request goes back to the head of the queue (lane_host.zig:362, 447). If nothing else is running it fails with alone_words (lane_host.zig:475).
    • Refusing: only at load. The pool must hold pages(--context + max_rows) once, or startup refuses as admit does today (cuda.zig:340). After that any single request fits on its own, so a request either runs or queues.
    • Release: on finish, cancel and error, releaseFn (cuda_lanes.zig:433) gives back the unused promise and releases every page the stream holds. Pages the store retained stay with the store.

    3. Ownership through Pages

    • Core owns page ids, refcounts, the free list and the promised count. The pool implements prompt_cache.Pages (prompt_cache.zig:62), and available() means "free and not promised", as documented there. promise/unpromise/take sit beside the vtable for lanes and engine.
    • The engine owns the device arena (page_bytes x count) and each stream's page table: one device buffer per lane at a fixed address, rewritten before replay.
    • Sharing:
      • begin retains the matched path into adopt (prompt_radix_tree.zig:134). The stream writes those ids into its table.
      • Marks and resumes sit on page edges only (prompt_radix.zig:58, 185). A shared page is therefore always full, and a stream starts writing in a page of its own. No copy-on-write is needed.
      • The engine never writes a page whose holder count is above 1, and asserts this in debug builds.
    • On keep, Kept.shared (prompt_radix.zig:183) names the leading pages the tree already held. The stream retains the tree's id, rewrites its table and releases its duplicate. Safe only because resumed == fresh bit for bit (section 5).
    • MTP lookahead: Nemotron's head row p reads token p+1, so today's store runs with lookahead = 1 (cuda.zig:360), and prompt_radix asserts lookahead == 0 (prompt_radix.zig:34). In this contract the head's row for (hidden p, token p+1) is stored at slot p+1. A page then depends only on tokens up to its own. The snapshot carries hidden row at-1 (5 KiB), and the resuming stream computes that one head row itself.
    • Pinning means holding the page: a stream retains every page it maps. idle() (prompt_radix_tree.zig:230) skips any node whose pages have more than one holder.

    4. Evicting cached prompts

    • On demand only, with no background reaper:
      • A reservation that does not fit calls reclaim(need) (prompt_radix.zig:267). That evicts the least valuable idle leaf: never-resumed before resumed, then least recently used.
      • A keep that passes Limits evicts by victim (prompt_radix.zig:259): superseded turns first.
    • Never evicted: a node that a stream holds, and the entry being restored (the skip in reserve).
    • Snapshots are stored in slots, not in pages. At load the engine allocates Limits.snaps slots of snap_bytes (48 MiB for Nemotron), all in one device buffer, and keeps a free list of them. That buffer is the snapshots' backing storage; withholding page ids would leave them none.
    • A slot is one contiguous region because an armed mark's chunk writes the SSM and conv states straight into it (arm, cuda_reuse.zig:141); 14 scattered pages would break that write. The slots replace today's spare Mamba buffers (spares_max 4, cuda_marks.zig:8) and the cuMemAlloc a save makes when no spare waits (cuda_reuse.zig:133), so no snapshot allocates after load.
    • Charge: --prompt-cache-gib is split at load into Limits.snaps x 48 MiB for the slots and the rest as Limits.pages. Both come out of the pool before any stream's pages, as above. save takes a free slot and drop returns it. A keep with no free slot evicts by victim or is refused (prompt_radix.zig:200), and never takes from a stream's promise.
    • Limits.snaps keeps its meaning, a count of snapshots, now also the number of slots.

    5. Invariants and checks for the first slice

    • Page lifetimes: a core test over a fake pool, in the style of prompt_radix_test.zig, runs random admit, keep, resume, finish, cancel and error sequences. After each one, with the store cleared, every refcount is 0 and free == total; releasing a page at 0 panics. On CUDA, debug builds fill released pages with NaN, and the idle server logs promised 0, free + tree == total and free slots + held slots == Limits.snaps.
    • Allocation failures: fail_after injected in the fake pool (like memory_lanes, lanes/fake.zig:60): a failed promise queues and runs later (lane_host_admit_test.zig:130, 167); a failed snapshot is refused with the reply unchanged (prompt_radix_test.zig:214); a failed table write rolls back; nothing leaks.
    • Drafted == plain: tools/zig/reuse_served.py compare (token SHA, draft: false arm) on Nemotron CUDA, with greedy and seeded sampling.
    • Resumed == fresh: reuse_served.py record with the cache on against replay with --prompt-cache-gib 0, plus resumes across a page edge and one row on either side of it.
    • Concurrent == solo: tools/zig/structured_served.py and reuse_served.py's concurrent group, run at --parallel 4 against 1, with prompts that share pages.
    • Admission: reservations past the pool: the second request waits, then runs; plus reuse_served.py evict.

    6. Out of scope

    • Growing a stream past its reservation. A reply that reaches max_tokens ends with length, as it does today.
    • Preemption or pausing of running streams.
    • The fixed part in the pool; other families, Metal, the cluster mirror.

    7. Choices for you

    • Store for the pilot: Nemotron uses the byte-copy store today, which keeps prompts at any row. Nothing on main implements Pages yet. Moving Nemotron to prompt_radix means resume points every 512 tokens and the head-slot shift above. The cheaper alternative is to keep the copy store and charge its bytes to the pool.
    • Allocation calls: Pages has no take or reserve call, so the draft puts promise/unpromise/take beside it. They could go into the interface instead.
    • The engine's own sequence: its KV sits in full-window buffers that the single-stream fast path uses. It can move to pages, or keep a full window outside the pool.

  4. ashhart commented on Oct 16, 2026

    @ashhart
    ashhart
    on Oct 16, 2026
    Owner

    Thank you, @jschmied, for the full reservation contract, including adopted pages, idle-only eviction and the MTP lookahead transition.

    We checked the accounting and lifetime rules. One detail to make explicit is the backing storage for snapshots: withholding ids from an already allocated device arena does not pay for an additional copied snapshot allocation. Also, a 2,048-token prompt plus 4,096 reply tokens and 16 verify rows needs thirteen 512-row pages, or 45.5 MiB; a 2,000-token prompt fits twelve.

    Please update the contract with the snapshot's actual storage and charge. We'll come back on the copy-store versus radix choice, reservation API placement and the engine's own sequence before the pilot implementation.

    Posted by Ash's AI Assistant Synthena

  5. jschmied commented on Oct 16, 2026

    @jschmied
    jschmied
    on Oct 16, 2026
    ContributorAuthor

    @ashhart, thank you. I've updated the contract above in place:

    • Snapshot storage and charge (section 4): snapshots get their own storage instead of withholding page ids. At load the engine allocates Limits.snaps slots of 48 MiB (Nemotron's 49,082,368-byte Mamba state in whole 2 MiB driver pages) in one device buffer. They come out of the pool beside the streams' fixed parts, and --prompt-cache-gib is split into the slots and Limits.pages. A slot stays contiguous because an armed mark's chunk writes the SSM and conv states straight into it (cuda_reuse.zig:141). The slots replace today's spares and the cuMemAlloc a save makes when none waits (cuda_reuse.zig:133), so no snapshot allocates after load.
    • Arithmetic (section 1): corrected. A 2,048-token prompt with 4,096 reply tokens and 16 verify rows needs 6,160 rows, which is 13 pages or 45.5 MiB; a 2,000-token prompt fits in 12.

  6. ashhart commented on Oct 16, 2026

    @ashhart
    ashhart
    on Oct 16, 2026
    Owner

    Thank you, @jschmied, for giving snapshots their own backed 48 MiB slots and correcting the reservation arithmetic.

    For the first pilot, keep the copy store with the current MTP lookahead. That preserves the existing head positions and resume behavior while we check paging. Charge copied KV storage to the same budget alongside the snapshot slots.

    Put reservation leases beside Pages, under one core owner. This keeps the cache's page-ownership interface small while admission and allocation use the same counters.

    Keep one full-window sequence for the engine, charged separately before sizing the pool. That preserves the existing graph addresses and single-stream fast path; count its memory once.

    Please start with a model-free core layout and reservation owner, with fake-pool tests for page lifetimes, competing reservations and cleanup after allocation or table-upload failures, then pilot it on Nemotron CUDA with drafted/plain, resumed/fresh and concurrent/solo checks.

    Posted by Ash's AI Assistant Synthena

  7. jschmied commented on Oct 16, 2026

    @jschmied
    jschmied
    on Oct 16, 2026
    ContributorAuthor

    @ashhart, the first slice is up as draft #935, on main 46a08c1 (1.0.12). It follows your 07:28Z choices: the copy store with the current MTP lookahead, leases beside Pages under one core owner, and the engine's own full-window sequence charged once. The model-free core and its fake-pool tests are in the first commit, the Nemotron CUDA pilot in the second. Checks and what was not run are in the PR description.

  8. ashhart commented on Oct 17, 2026

    @ashhart
    ashhart
    on Oct 17, 2026
    Owner

    Thank you, @jschmied, #935 turns the three chosen pilot interfaces into a core owner and a Nemotron CUDA consumer with failure controls.

    We checked the supplied pilot description and current 197db0c guard: prompt_cache.zig:464 now applies the entry-count limit to the reuses shortcut too. The earlier generic count-limit finding closes; your current Nemotron adapter did not use that shortcut.

    We will review the page/slot ownership and integrated native gates together, retaining the copy store, MTP lookahead and separately charged full-window sequence. Please keep Super VL, multi-rank and structured-output cases explicitly open until their own device runs; the existing Lightning and host receipts remain useful.

    Posted by Ash's AI Assistant Synthena

  9. ashhart
    added 2 commits that reference this issue on Oct 18, 2026

    Core: one owner for KV pages, with reservation leases beside Pages (#781

    7e10ed1

    Nemotron CUDA: streams beside the engine's own take KV pages from the…

    a9dfd38

  10. ashhart commented on Oct 18, 2026

    @ashhart
    ashhart
    on Oct 18, 2026
    Owner

    Thank you, @jschmied, for the core reservation contract and corrected 13-page/45.5 MiB example. 1.1.0 carries #935's agreed Nemotron CUDA pilot: one core page owner, leases beside Pages, the existing copy store/MTP lookahead and one separately charged full-window engine sequence. Snapshot and copied-KV storage stay budgeted; host-state placement is explicit when device cache space does not fit. After its final ownership, admission, draft/cache/concurrency and multi-rank gates pass, this closes the agreed pilot. Other families, preemption and growth beyond reservations remain follow-ups.

    Host ownership and GB10 serving passed; RTX PRO 6000 cache counts matched 12/12 on-device, 4/12 with oversized host fallback.

    Posted by Ash's AI Assistant Synthena

  11. ashhart
    closed this as completedon Oct 18, 2026
  12. BlivionIaG
    added 2 commits that reference this issue on Oct 20, 2026

    Core: one owner for KV pages, with reservation leases beside Pages (a…

    bcf2256

    Nemotron CUDA: streams beside the engine's own take KV pages from the…

    fa4506d
  13. Bizuayeu
    added 3 commits that reference this issue on Oct 21, 2026

    glm cuda: streams past the engine's own take their KV pages from the …

    804c4da

    glm cuda: streams past the engine's own take their KV pages from the …

    f28a57f

    glm cuda: streams past the engine's own take their KV pages from the …

    26d5e5e
  14. ashhart
    added a commit that references this issue on Oct 25, 2026

    glm cuda: streams past the engine's own take their KV pages from the …

    44b1e84
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

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions