[GH-ISSUE #4692] [RFC] Model-independent Engram byte tables with owned local storage and CUDA stream lookup #1195

New issue

Open

opened 2026-10-29 21:23:49 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 21:23:49 +00:00

Owner

Originally created by @magicYang1573 on GitHub (Oct 15, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4692

Motivation

The existing layer/head Engram API works for DeepSeek-shaped tables. SGLang also needs N-gram and per-layer input tables with different packing, scaling and sharding. This proposes one additional table-oriented mode on EngramStore, with no new public spec or request-handle classes.

Proposed contract

EngramStore(tables={name: {num_rows, row_bytes}}, namespace=..., local_dir=... | store_client=...).
populate(table_id, rows) copies immutable CPU byte rows into Store-owned storage.
lookup(table_ids, ids, output, *, offsets=None, stream=None) reads int32/int64 [T,K] IDs into uint8 [T,K,R]. Repeated table names and nonoverlapping token-major/head-major output are supported. ID -1 preserves output; offsets convert global IDs into table-local rows.
CPU lookup is synchronous. Local CUDA lookup enqueues complete work on the supplied stream and returns None. Consumers use ordinary streams/events. Warmup before CUDA graph capture and Store/tensor lifetime requirements are explicit.
get_table_spec, remove, close; the existing layer constructor and methods remain available.

Local implementation

A Python facade owns namespaced immutable mmap snapshots, validates layouts and performs atomic create-only publication. A private pybind/CUDA helper registers host pages and gathers byte rows directly into CUDA output. It does not depend on PyTorch headers or implement model math. This first version pins selected tables in full; it is not a host page cache.

Scope and relationship to existing work

This revisits the owned-local idea in #4083 with a model-independent interface and compatibility for the upstream layer API. It does not include the remote gather optimization in #4502 or the experimental feat/engram-pipelined-read transport. Store-backed CPU reads continue through existing ranged reads; remote CUDA/pipeline work remains separate.

A companion SGLang change keeps adapters under mem_cache/engram and changes only table construction/lookup entry points for DeepSeek V4.1, Qwen4-Exp, LongCat Ngram, and Gemma3n/Gemma4. Hashes, scales, projections, request history and collectives remain in SGLang.

Validation

Local API tests currently pass for CPU layouts, persistence/reopen, validation, CUDA stream ordering and graph replay with changing IDs. Adapter tests pass for BF16, FP8 with table scaling and FP8/E8M0 block scaling. Real-checkpoint local end-to-end comparisons and packaging validation are in progress; their results will be reported in the PR rather than inferred from these unit tests.

Review points

Is the table API appropriately scoped for upstream while retaining the existing layer API?
Is a Python facade plus private CUDA helper acceptable for the first local implementation?
Immutable snapshot semantics, full-table pinning and graph lifetime requirements should be settled before broader consumers adopt the API.

AI assistance: implementation and this RFC were prepared with Codex; human review is pending.

Originally created by @magicYang1573 on GitHub (Oct 15, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4692 ## Motivation The existing layer/head Engram API works for DeepSeek-shaped tables. SGLang also needs N-gram and per-layer input tables with different packing, scaling and sharding. This proposes one additional table-oriented mode on `EngramStore`, with no new public spec or request-handle classes. ## Proposed contract - `EngramStore(tables={name: {num_rows, row_bytes}}, namespace=..., local_dir=... | store_client=...)`. - `populate(table_id, rows)` copies immutable CPU byte rows into Store-owned storage. - `lookup(table_ids, ids, output, *, offsets=None, stream=None)` reads int32/int64 `[T,K]` IDs into uint8 `[T,K,R]`. Repeated table names and nonoverlapping token-major/head-major output are supported. ID -1 preserves output; offsets convert global IDs into table-local rows. - CPU lookup is synchronous. Local CUDA lookup enqueues complete work on the supplied stream and returns None. Consumers use ordinary streams/events. Warmup before CUDA graph capture and Store/tensor lifetime requirements are explicit. - `get_table_spec`, `remove`, `close`; the existing layer constructor and methods remain available. ## Local implementation A Python facade owns namespaced immutable mmap snapshots, validates layouts and performs atomic create-only publication. A private pybind/CUDA helper registers host pages and gathers byte rows directly into CUDA output. It does not depend on PyTorch headers or implement model math. This first version pins selected tables in full; it is not a host page cache. ## Scope and relationship to existing work This revisits the owned-local idea in #4083 with a model-independent interface and compatibility for the upstream layer API. It does not include the remote gather optimization in #4502 or the experimental feat/engram-pipelined-read transport. Store-backed CPU reads continue through existing ranged reads; remote CUDA/pipeline work remains separate. A companion SGLang change keeps adapters under mem_cache/engram and changes only table construction/lookup entry points for DeepSeek V4.1, Qwen4-Exp, LongCat Ngram, and Gemma3n/Gemma4. Hashes, scales, projections, request history and collectives remain in SGLang. ## Validation Local API tests currently pass for CPU layouts, persistence/reopen, validation, CUDA stream ordering and graph replay with changing IDs. Adapter tests pass for BF16, FP8 with table scaling and FP8/E8M0 block scaling. Real-checkpoint local end-to-end comparisons and packaging validation are in progress; their results will be reported in the PR rather than inferred from these unit tests. ## Review points - Is the table API appropriately scoped for upstream while retaining the existing layer API? - Is a Python facade plus private CUDA helper acceptable for the first local implementation? - Immutable snapshot semantics, full-table pinning and graph lifetime requirements should be settled before broader consumers adopt the API. AI assistance: implementation and this RFC were prepared with Codex; human review is pending.

bsakaguchi commented

2026-10-29 21:23:50 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 15, 2026):

Thanks for opening this issue, @magicYang1573!

Field	Value
Issue	#4692
GitHub user ID	`102229987`
Reporter	@magicYang1573

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Oct 15, 2026): Thanks for opening this issue, @magicYang1573! | Field | Value | |-------|-------| | **Issue** | #4692 | | **GitHub user ID** | `102229987` | | **Reporter** | @magicYang1573 | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.