Skip to content

Kolibri-1 on CUDA: prompts in the shared KV store with TF_KV_SHARE (off by default) - #1891

Closed
ashhart wants to merge 3 commits into
mainfrom
lead/s173-kolibri-kv-share
Closed

ashhart wants to merge 3 commits into
mainfrom
lead/s173-kolibri-kv-share

Conversation

@ashhart

@ashhart ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Part of #1657: Kolibri-1 on CUDA keeps its prompts in the shared paged KV store, as Gemma 4 does since #1883.

What changes

  • With TF_KV_SHARE=N (off by default, at most N kept states), Kolibri-1's kept prompt states live in the core KV pool's pages instead of private copies. A resumed prompt restores its page-aligned kept pages and its chronological ring rows, and gives the same tokens as a cold one.
  • The store's accounting matches Gemma 4's term for term: blob pages times page bytes, copy bytes, and the layout with no snapshot bytes.
  • The narrower 128-dim paged prompt kernels load only when Kolibri opts in, and in the kernel contract. Otherwise a missing kernel returns WindowShape. Spark's page loading is unchanged from Gemma 4 on CUDA: prompts in the shared KV store with TF_KV_SHARE (part of #1657) #1883.

Unset or 0 keeps today's copies, byte for byte.

Checks

  • Host: a real-loader host test for the page load, plus the store tests. The lean check shows 424 problems before and after.
  • Two independent reviews: LAND.
  • GB10, Kolibri-1, defaults (main against this change, balanced order): GREEN. Prompt tokens are identical in all 20 rows, and decode at 1 to 8 streams is within noise (8 streams: 293.1 against 293.0 tok/s).
  • GB10 with TF_KV_SHARE=8 against main, to check that replies stay equal with the store on: results will be posted here before this lands.

Use the existing TF_KV_SHARE adapter for global page references and
pool-owned sliding snapshots. Stage tree tables in active-chain order,
bind captured steps to stable slot tables, and keep the original linear
path when the switch is unset.

Add 128-dimensional prompt page addressing to the original row/replay
template, with replacement/reference/control cases and pending resource
records. Cover resumed and concurrent cache oracles, copy-draft bindings,
partial acceptance, cancellation and eviction with host regressions and
production-source faults. NVCC and checkpoint qualification remain pending.

Refs #1657
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner Author

Closing in favour of @jschmied's #1755 (with #1620), which did the same Kolibri-1 store wiring first. See the comparison and the plan on #1755.

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