Skip to content

GLM paged attention and multi-sequence serving with prefix-aware scheduling - #642

Open
ClintMoody wants to merge 6 commits into
mw00:mainfrom
ClintMoody:paged-attn-main
Open

ClintMoody wants to merge 6 commits into
mw00:mainfrom
ClintMoody:paged-attn-main

Conversation

@ClintMoody

@ClintMoody ClintMoody commented Oct 28, 2026 •

Copy link
Copy Markdown

Summary

This contribution adapts GLM-5.3-Flash's latent cache to a page-pool and page-table layout, adds per-sequence state and prefix sharing, and schedules multiple requests with admission and pressure management. It includes CUDA gather routing, parity tests for pages, prefix and latent mapping, and a synthetic dynamic-sequence test.

The branch is rebased onto current main (v1.1.56, 26d53d8) and carries two commits: the paged latent cache (page table, page pool, prefix store, gather kernel) and the dynamic multi-sequence scheduler with MTP speculation. The abliteration harness (glm_ablate.cuh) is deliberately not in this PR; it belongs on its own branch.

This PR does not port GLM paged attention into DeepSeek's distinct src/ds4/ engine.

Author credit: Clint Moody contributed the paged-attention, multi-sequence serving, scheduler, and related integration work. Project Maya and Strata retain their existing author credits and copyright notices. The repository's MIT license and all upstream/vendor notices are unchanged. DeepSeek model weights and optional differently licensed drafter weights do not belong in the code PR.

What the branch contains

  • Paged latent cache: latent rows live in a GPU slab split into fixed-size pages. A logical row reaches its slab row through a page table, so a page can be freed, reused and reordered without moving the tokens that point at it. Refcounts, LRU eviction, a logical-to-physical table and a physical-to-logical mirror live in glm_pages; the canonical prefix store sits on top in glm_prefix.
  • Gather routing: dsa_select takes an optional paged table. glm_paged_gather_test drives the real kernel with an identity, a shuffled and a corrupted table and compares the device result to the host mapping. Its host mapping layer runs without a GPU.
  • Multi-sequence serving: STRATA_GLM_SEQS shares the engine between requests. A request is admitted only when its sequence state binds; a parked sequence is paged in or out; the oldest idle page is evicted under pressure.
  • Prefix reuse and MTP: repeated-prefix requests reuse cached tokens, and speculative decoding runs on the multi-sequence path.
  • Sizing and pressure: the pool sizes from STRATA_GLM_PAGE_SIZE / STRATA_GLM_PAGE_POOL_GB and shrinks by halves when the device cannot hold the request. STRATA_GLM_MULTI_WAIT bounds a request's wait for a context; 0 means wait forever.

Review fixes applied

  1. Pause/restore byte count: the table mode is now written and read as a 4-byte host value (SnapPacker::put_host, snap_unpack_host). Before, the pack wrote 4 bytes, the unpack read 8, and part_bytes counted 4, so the bounds check failed at the last field and every resume silently fell back. A static_assert now locks the width that part_bytes counts.
  2. seq_bind failure propagation: seq_bind and seq_bind_all return bool. A failed bind rejects the request instead of marking the slot busy while the model still points at another sequence.
  3. STRATA_GLM_MULTI_WAIT=0: now unbounded, as the docs say. Before, it rejected the request immediately.
  4. evict_lru: one pass plus partial_sort, same order (last use, then page index), O(N log N) instead of O(N x n).
  5. Device-side coverage: the gather test now checks that the shuffled routing is a permutation (a collision would still "move" a row), takes --require-gpu so a busy GPU cannot look like a pass, and zeroes device memory with cudaMemset (a host memset on device memory faulted). glm_pages_parity SKIPs cleanly without a usable CUDA device.
  6. Docs: docs/SETTINGS.md uses MAYA_* names, which env_alias.hpp maps to STRATA_* at startup, so the rows work under either spelling.
  7. rows_alloc gate: (n_parts_ == 1 || split_rows_setting() > 0) && fast_mode_, so the multi-GPU split-rows leg is no longer dropped when STRATA_GLM_SEQS is on.
  8. Prefill chunk clamp: STRATA_GLM_PREFIX_PAGE_ALIGN is opt-in and off by default. When unset, a chunk keeps its full size and the stateful tile is captured at the natural chunk end, so the page clamp is not a tax on every prompt.
  9. Loud refcount guards: an unref outside the pool, an unref of a freed page, and an unref of an unallocated page each print to stderr and return. The allocator-hold release is the one intentional negative count. GlmPages still relies on the caller's mutex; it holds no lock of its own.
  10. Settings rows: all 11 new variables have rows in docs/SETTINGS.md.
  11. Docs and noise: docs/SETTINGS.md uses MAYA_* names, which env_alias.hpp maps to STRATA_* at startup. No plan or task references, no run transcript, no ablation harness, no CUDA_LIBRARIES link. bench/README.md carries no generate.cpp line numbers.
  12. The 460-line pause/restore selftest stays in glm_multi.cu but is called only by glm_seq_dynamic_test.cpp; it is the only coverage of pause and restore.
  13. Commit shape: two commits, core paging first, scheduler second. They are separable into two PRs if you prefer; the scheduler needs the pool, so the order is fixed.

Test plan (performed evidence)

Build (this tree, upstream main @ 26d53d8 plus the 26 files):

cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_CUDA=ON \
      -DSTRATA_BUILD_TESTS=ON -DCMAKE_CUDA_ARCHITECTURES=86
ninja -j24

Full build, every target: EXIT 0, no errors (2m42s wall on 2x RTX 3090 Ti).

Host-runnable checks (no GPU needed):

Test Result
glm_paged_gather_test --selftest (host layer) 5 checks, 0 failures
glm_prefix_parity --selftest PASS, 0 failures
glm_lat_paged_parity --selftest 10 checks, 0 failures
glm_pages_parity --selftest SKIP without a device (it calls cudaMalloc)
tests/glm_memory_test.cpp passed
tools/check_hip_compat.py passed

Device checks (free-GPU window, engine stopped):

Test Result
glm_pages_parity --selftest 35 checks, 0 failures
glm_prefix_parity --selftest PASS, 0 failures
glm_lat_paged_parity --selftest 10 checks, 0 failures
glm_paged_gather_test --selftest --require-gpu 8 checks, 0 failures
glm_seq_dynamic_test --selftest <pack> 1116 checks, 0 failures (68 in the pause-snapshot selftest)

The fixture comes from tools/glm_synth_pack.py, which packs data/glm5-synth.gguf through tools/iq_pack.py.

CI scope, stated exactly: .github/workflows/tests.yml runs the Python suites, check_hip_compat.py, and the CPU-only C++ tests. It does not build the engine and does not run ctest, so the device tests above run on a machine with CUDA. The host layers of the gather test and the lat mapping test are the subset that can run without one.

CI's job reproduced on this exact tree, in a venv with CI's dependency list (numpy jinja2 regex pyyaml tqdm psutil jsonschema>=4.18,<5):

CI step Result
tools/check_hip_compat.py PASS, every CUDA name mapped
glm_memory_test (-Wall -Wextra -Werror) PASS
ds4_cpu_prepare_rows_test native and stub PASS
unittest discover -s serve 738 tests, OK (22 skipped), 586.9s
unittest discover -s tools 728 tests, OK (35 skipped), 34.2s

Not yet run: the CONTRIBUTING.md before/after benchmark pair (maya.py --bench on this build and on main), which needs a GLM pack and two model loads.

Earlier probe window (v1.1.33 base, not re-measured after this rebase): 1,536 of 1,775 prompt tokens were reused on the repeated-prefix probe, and the speculative path under concurrency was byte-identical to the serial MTP path.

Disclosed limitation: with paging enabled, one of three temperature-zero canaries diverged in wording from the stock answer. Sustained-load, long-context and error-path gates remain open; please do not label the paged path production-ready until they run.

Notes for maintainers

  • 26 files, +4464 / -133 against main @ 26d53d8.
  • No tracked binary fixtures; the port reuses upstream's existing fixtures.
  • No local absolute paths, hostnames, IPs or private history in the diff.
  • DeepSeek's src/ds4/ engine is deliberately untouched. No DeepSeek performance claim is made from the GLM tests.

@ClintMoody
ClintMoody force-pushed the paged-attn-main branch 2 times, most recently from 35edd9d to 6806e0f Compare October 28, 2026 22:55
@catonooka

Copy link
Copy Markdown

Thanks for this substantial contribution, and especially for the honest reporting.
Disclosing the temperature-zero canary divergence up front, calling out the gates you
didn't run, and avoiding DeepSeek claims from the GLM tests is exactly the right way to
submit work like this. The direction is something the project genuinely needs and the
engineering underneath is real. That said, I don't think this can merge as it stands.
Here's what I found.

What's good

The overall design is sound: a refcounted page pool with LRU eviction, a content-hashed
prefix cache, admission control with a pressure ladder, and clean gating so the
incumbent path is untouched when STRATA_GLM_PAGED is unset. The measurement
discipline (probe windows, disclosed limits) is also much appreciated.

Blockers

1. Stale base

The branch is based on v1.1.33 while main is at v1.1.53, and the PR currently shows
as conflicting. Per CONTRIBUTING.md, could you rebase onto main before we go further?

2. Correctness bugs found in review

  • D2D memcpy with a host pointer. SnapPacker::put gets called as
    packer.put(&mode, sizeof(mode)), which runs cudaMemcpy(..., cudaMemcpyDeviceToDevice) from a host stack variable, and pause_resume mirrors it
    with D2D into &mode. That's invalid CUDA usage and can silently corrupt or disable
    pause-capture at runtime.
  • seq_bind binds a null arena on allocation failure. It logs "stays parked" but
    then goes ahead with the swap using null state_/snap pointers and sets cur_seq_.
    Any caller that binds without a prior seq_ready (the selftest does) will crash.
  • page_unref underflow is silent. An unref on a refcount-0 page drives the counter
    negative and frees the page while other owners may still be gathering its rows, which
    is the cross-sequence corruption class of bug. Please make refcount violations loud
    (an assert or at least a log), and consider moving the locking and refcount contract
    into GlmPages itself instead of relying on every caller honoring a global mutex
    documented in a different file.
  • The wait gate has no terminal path. The re-push loop fix stops the tight spin,
    but a request that can never fit still retries forever, with no timeout, no
    rejection, and no signal to the operator. It's the original starvation bug at a lower
    rate.
  • The rows_alloc gate change from n_parts_ == 1 || split_rows_setting() > 0 to
    n_parts_ == 1 && fast_mode_ silently drops the multi-GPU split-rows leg.
  • Prefill chunking tax. With the pool on and prefix enabled (the default when
    paged), prefill chunks clamp to one page of 256 rows for every prompt. That's
    overhead tied to the pool rather than to actual cache hits.

3. The core claim lacks test coverage

glm_lat_paged_parity only validates the host-side lat_row_offset arithmetic; the
device gather kernels never run with a shuffled table. The model-level selftests use
PAGE_SIZE=4096, so the table is always identity. Put together, the paged gather
routing, which is the heart of this PR, has zero coverage anywhere, and the name
"bit-exact parity" oversells what the test does. Could you add a device-side test that
runs the real kernels against a shuffled non-identity table?

One more thing on this front: all four new tests are GPU-gated while CI is GPU-less,
so CI gives this feature no regression coverage at all. Worth stating that plainly in
the PR, and maybe finding a host-runnable subset.

4. CONTRIBUTING.md compliance

  • About 11 new STRATA_GLM_* variables have no rows in docs/SETTINGS.md (two are
    documented in bench/README.md instead).
  • The new parity targets link ${CUDA_LIBRARIES}, which this project never defines, so
    a HIP configure fails to link them. The sibling targets use
    ${_strata_gpu_runtime_target}. The CUDA API calls themselves all look covered by
    existing hip_compat defines, so it's just the CMake line.
  • One topic per PR: the 145-line residual-ablation harness in glm_ablate.cuh plus its
    three env vars is unrelated here. Could you split it out or drop it?
  • The README "fork provenance" note can't merge upstream, since the repo would end up
    describing itself as a fork of itself. The honest caveats in it are great material
    for the PR description, though. The bench/README.md framing ("maya-paged-attention
    mirror") and its hard-coded generate.cpp line numbers need to go too.

5. Process noise in the code

The diff carries 136 internal "Task N.M" plan references, dated run transcripts used as
comment justification ("RUN 2: 6,143,865 wait lines..."), remediation archaeology, a
roughly 460-line env-mutating selftest compiled into the production glm_multi.cu
translation unit, test-only public members, and printf diagnostics on hot paths. None
of this survives contact with the next maintainer. Comments should explain the code,
not narrate the history of writing it.

Suggested path forward

  1. Rebase onto current main, and split into core-paging and scheduler-serving PRs,
    with ablation separate or dropped.
  2. Fix the correctness items in section 2 and add the loud refcount checks.
  3. Add the shuffled-table device-side parity test from section 3.
  4. Add the SETTINGS.md rows, switch the link line to _strata_gpu_runtime_target, and
    drop the README note.
  5. Do a comment cleanup pass per section 5.

None of this diminishes the amount of real work you put in. It's just that the last
30% is what makes this safe to run in production. Happy to re-review once the rebase
lands.

@ClintMoody

Copy link
Copy Markdown
Author

Rebased, reworked, and pushed. The branch is rewritten onto main @ 26d53d8 (v1.1.56), head ff7fc0d, and GitHub now reports it as mergeable: 26 files, +4464 / -133. The PR body has the full map; the short version follows.

Blockers

  1. Stale base: gone. Old v1.1.33 history dropped; two commits on current main (core paging, then scheduler).
  2. Correctness
    • D2D host pointer: the table mode is now a 4-byte host value (put_host / snap_unpack_host), with a static_assert on the width that part_bytes counts. While fixing it I found the old staging code wrote 4 bytes and read 8, so every pause_resume was already failing its bounds check and falling back.
    • seq_bind: returns bool now; seq_bind_all ANDs the parts, and a failed bind rejects the request instead of marking the slot busy.
    • page_unref: out-of-range, unallocated and double-free unrefs print to stderr and return. Locking still belongs to the caller's mutex; moving it into GlmPages is a separate change I did not make.
    • Wait gate: STRATA_GLM_MULTI_WAIT bounds the wait and rejects with a diagnostic; 0 means wait forever, as the docs say.
    • rows_alloc gate: (n_parts_ == 1 || split_rows_setting() > 0) && fast_mode_, so split rows survive.
    • Prefill clamp: STRATA_GLM_PREFIX_PAGE_ALIGN is opt-in and off by default, so the chunk tax is only paid when the reuse is requested.
  3. Test coverage: glm_paged_gather_test drives the real dsa_select on an identity, a shuffled and a corrupted table, and checks the shuffled routing is a permutation (a collision would still "move" a row). Its host mapping layer runs without a GPU, and --require-gpu stops a busy GPU from looking like a pass. On device: 8 checks, 0 failures.
  4. CONTRIBUTING: all 11 new variables have docs/SETTINGS.md rows (under MAYA_*; env_alias.hpp maps them to STRATA_*); the targets link ${_strata_gpu_runtime_target}, no CUDA_LIBRARIES; glm_ablate.cuh is out of this PR; the README fork note is gone; bench/README.md has no line numbers.
  5. Process noise: no plan references, no run transcripts, no ablation harness, no test-only public members. The pause/restore selftest stays but is reached only from glm_seq_dynamic_test.cpp; it is the only coverage of pause and restore, so I would rather keep it than lose it.

Evidence on this tree (engine stopped for the device window, then restarted):

Test Result
Full build, all targets, cmake + ninja -j24 0 errors
glm_pages_parity --selftest 35 checks, 0 failures
glm_prefix_parity --selftest 0 failures
glm_lat_paged_parity --selftest 10 checks, 0 failures
glm_paged_gather_test --selftest --require-gpu 8 checks, 0 failures
glm_seq_dynamic_test --selftest 1116 checks, 0 failures

Stated plainly, as you asked: CI does not build the engine and does not run ctest. The host layers above (gather, lat) and the Python suite are what run GPU-less; the rest ran on a 2x RTX 3090 Ti box.

Still open: the CONTRIBUTING.md before/after maya.py --bench pair (needs a GLM pack and two loads), and the disclosed canary divergence with paging on. On the two-PR suggestion: the two commits are already separable and I can open them as two PRs if you prefer; the scheduler needs the pool, so the order is fixed.

@ClintMoody

Copy link
Copy Markdown
Author

One addition to the evidence, since you asked for the CI claim stated exactly: I reproduced CI's Python job on this tree in a venv with CI's dependency list.

CI step Result
tools/check_hip_compat.py PASS
glm_memory_test (-Wall -Wextra -Werror) PASS
ds4_cpu_prepare_rows_test native and stub PASS
unittest discover -s serve 738 tests, OK (22 skipped)
unittest discover -s tools 728 tests, OK (35 skipped)

The workflow still does not build the engine and does not run ctest, so the device checks (pages 35, gather 8, lat 10, prefix, 1116 seq checks) remain a local run on a CUDA box.

@ClintMoody

Copy link
Copy Markdown
Author

Two-GPU (2x RTX 3090 Ti) run of this head, with the shared prompt sent as a real system message, which is what a client sends. Traffic: 7,000-token system prompt, 600 tokens of user text per stream, up to 1,200-token answers, thinking at max, temperature 0.

config streams wall worst first token tok/s in all
8 seqs, 262K parked, no interleave 8 109.9 s 54.9 s 19.0
8 seqs, 64K parked, 2,048-token pieces, MAYA_GLM_MTP_CONC=0 8 42.2 s 39.2 s 31.2

The same head and the same settings, with the shared text in the user message instead of a system message: 8 streams at 102.6 s wall, 93.9 s worst first token, 21.3 tok/s.

Why the second row is faster: the slot cut saved the 5,587-token system prompt (glm slots: took back a 5587-token conversation from disk (0.0 s)), so each stream read 546 tokens in one 576-token chunk (glm prefill: 546 tokens in 2162.0 ms (3.96 ms/token, 253 tok/s)).

On the 256-row clamp you noted: on this head, a read from token 0 goes in one 6,144-token chunk at 798 to 827 tok/s, against 286 tok/s on the v1.1.48 port. STRATA_GLM_PREFIX_PAGE_ALIGN was off, as the setting describes.

Still open: first token grows with stream count (2.5 s at 1, 4.8 s at 2, 17.9 s at 4, 39.2 s at 8). One serve thread reads prompts one at a time, so eight reads of 2.15 s are already 17 s of queue, and a 546-token read runs at 253 tok/s because of the stateful tile at the chunk end. Cutting that per-chunk overhead, and reading while other streams decode, would be the next step.

@ClintMoody

Copy link
Copy Markdown
Author

Rebased onto main (7bf15ac). New head: dcaa5df. The two commits were replayed with no conflicts and no changes: git diff ff7fc0d dcaa5df -- src include CMakeLists.txt is empty, so the engine code and the measurements above are unchanged. The move only brings in the server side fixes that landed after v1.1.56.

@mw00

mw00 commented Nov 9, 2026

Copy link
Copy Markdown
Owner

Thanks, parity tests and a rebase on v1.1.56 are appreciated. Missing: (1) the GPUs, CPU, RAM, model+quant you tested on; (2) single-request prefill and decode before/after (must be unchanged for normal use), and multi-request throughput with its numbers; (3) a greedy-output match between this and main for the same prompts; (4) which behaviour is on by default - anything that changes answers should be opt-in.

@ClintMoody

Copy link
Copy Markdown
Author

Rebased onto v1.1.57 (0065548). New head f397014 = e697ec2 (paged latent cache) + e67fd33-equivalent (dynamic multi-seq scheduler + MTP speculation + paged wiring). Content diff vs the old head dcaa5df is empty: the rebase changes parents only, the PR's 2 commits apply onto v1.1.57 with the same resolved content.

@ClintMoody

Copy link
Copy Markdown
Author

Correction to my note above: the rebase is not content-identical this time. v1.1.57 already merged much of the scheduler work (glm-dsa B-series: GlmStateLayout, state_carve, mtp_carry_forget), so the rebased second commit (f397014) carries fewer own lines than dcaa5df did (3707 vs 3762 patch lines): HEAD's layout rewrite wins where the PR overlapped it, the PR keeps its includes/helpers and its own scheduler/speculation/paged wiring, and the overlap is resolved in upstream's favor. Verified by patch-text diff of dcaa5df vs f397014: 171 lines vanish, 116 appear, all accounted for by upstream's absorbed features plus context shifts.

@ClintMoody

Copy link
Copy Markdown
Author

One more commit landed on the PR head during the upgrade window's test runs: a669456. The PR's glm_seq_dynamic_test exposed a real divergence in the rebased stack: upstream's B4.1 rewrite of Glm5Model::seq_layout recomputes the sequence layout with Fp16 + the fast ik ring + kv-resident latents unconditionally, while the load path (glm_model.cu) carves SlowF32 with the plain ring and full-context latents when fast_mode_ is off (the STRATA_GLM_SLOW legs), so the guard in seq_alloc_arenas rejected every pack it did not carve itself. The fix mirrors the load's mode pick in seq_layout; fast mode is byte-identical before and after, so production behavior does not change. Verified: 1116 checks, 0 failures after the fix (the port tree's pre-fix binary had passed window 2 only because it was still the Nov 1 build).

…ore and the gather kernel

The latent rows of the GLM fast path move into a paged pool. A logical
row reaches its slab row through a page table, so a page can be freed,
reused and reordered without moving the tokens that point at it.

- glm_pages: the pool (bump-allocated pages, refcounts, LRU eviction,
  logical->physical table and the physical->logical mirror)
- glm_prefix: the canonical prefix store on top of the pool
- glm_batch: dsa_select reads the page table (paged_lat_table)
- parity: pages (35 checks), prefix (13 boundaries), lat mapping (10),
  and a device-side gather over an identity, a shuffled and a corrupted
  table (8 checks); the gather test's host layer runs without a GPU

The pool sizes from STRATA_GLM_PAGE_SIZE / STRATA_GLM_PAGE_POOL_GB and
shrinks by halves when the device cannot hold the request.
…d path's wiring

STRATA_GLM_SEQS lets several requests share the engine: a request is
admitted only when its sequence state binds, a parked sequence is paged
in or out, and the oldest idle page is evicted under pressure.

- generate: admission, eviction and a terminal path for a request that
  waits too long (STRATA_GLM_MULTI_WAIT=0 means wait forever, as the
  docs say); a failed bind rejects the request instead of marking the
  slot busy while the model points at another sequence
- glm_multi: sequence bind/release, the pause snapshot capture/restore
  (the table mode is written and read as a 4-byte host value, the width
  part_bytes counts) and the split-rows gate for n_parts > 1
- glm_prefill: the chunk/page clamp is opt-in (STRATA_GLM_PREFIX_PAGE_ALIGN)
- glm_rows, glm_fast, glm_window, glm_model: the paged table reaches
  dsa_select and the rows arena
- MTP: speculative decoding on the multi-sequence path
- tests: glm_seq_dynamic_test (1116 checks on a real pack fixture) and
  tools/glm_synth_pack.py, the fixture generator the test uses
…recompute diverges in the slow legs (STRATA_GLM_SLOW=1: the load carves SlowF32 with the plain ring and full-ctx latents, the B4.1 recompute assumed Fp16 + the fast ring + kv-resident rows)
clintm and others added 3 commits November 12, 2026 12:08
v1.1.58 renamed the engine's env family STRATA_ -> MAYA_ and added a
one-way alias (STRATA_X promotes to MAYA_X once at program start), so the
paged knobs must be read as MAYA_ names: MAYA_GLM_PAGED=1 otherwise fell
back to OFF and MAYA_GLM_SEQ_FLOOR to the full cap.

Renames MAYA_GLM_PAGED, MAYA_GLM_PAGE_SIZE, MAYA_GLM_PAGE_POOL_GB,
MAYA_GLM_SEQ_FLOOR, MAYA_GLM_SEQ_VRAM_FIT and the seq test's env names.
The lifecycle test now passes (1116 checks, 0 failures).
The runtime setenv sites do not get the one-way STRATA_ -> MAYA_ alias (it
runs once at program start).
v1.1.58 renamed the engine env family STRATA_ -> MAYA_ and made STRATA_X a
one-way alias for MAYA_X. v1.1.61 code reads MAYA_X, so the branch's remaining
STRATA_GLM_* reads ignored the config's MAYA_ names: paged attention, the page
pool, the state floor and the pause knobs all silently stayed at their
defaults. The reads now match the alias direction.
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.

4 participants