Skip to content

feat: native-paged decode CUDA graphs - #182

Open
drunkcoding wants to merge 26 commits into
mainfrom
plan/decode-cuda-graphs
Open

feat: native-paged decode CUDA graphs#182
drunkcoding wants to merge 26 commits into
mainfrom
plan/decode-cuda-graphs

Conversation

@drunkcoding

@drunkcoding drunkcoding commented Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Description

Opt-in decode CUDA graph capture/replay for resident ordinary-GQA Qwen3 native-paged attention only, integrated into active continuous-batching serving with explicit deny-by-default capability proof and eager execution as the authoritative fallback.

Full implementation of the Momus-approved plan in docs/superpowers/plans/2026-08-21-decode-cuda-graphs.md (12 tasks, strict TDD).

What changed

  • Authoritative KV storage (runtime/paged_kv_storage.py): one PagedKVStorage owns block allocation + per-layer native K/V tensors; both PagedKVCache and PagedAttentionBackend bind the same instance (owner-id enforced).
  • Exact Qwen3 registry (models/paged_attention_registry.py): registers only exact-type Qwen3PagedAttention; classifies DeepSeek V2/V3 as mla_layout_unsupported before any head_dim read; per-instance generated bound subclasses with unique layer_idx.
  • Allocation-free current-token KV writer (kernel/paged_kv_write.py): writes per-layer K/V before decode attention; safe under graph capture.
  • Deny-by-default capability (runtime/attention_types.py, serving/model_runner.py, runtime/model_offload.py, entrypoints/big_modeling.py): typed DecodeGraphCapability with 15 fixed reasons; composed from runtime + backend + storage-identity + device-equality + per-layer write proofs.
  • Safe lazy runner (serving/cuda_graph.py): bucketed (batch,context) keys, lazy warmup/capture, scratch reservation from the authoritative allocator, capture-failure quarantine, per-batch MOE_DISABLE_CUDA_GRAPHS kill switch, exact device-equality gates, real-row-only replay slicing.
  • Engine integration (serving/engine.py): pure decode routed through the runner only when eligible; mixed batches split eager-prefill + graph-decode + order-preserving recombine; sampler receives only real rows.
  • Lifecycle (serving/engine.py, entrypoints/openai/api_server_v2.py): invalidate_cuda_graphs on reload, idempotent shutdown, _replace_engine closes the old runner and releases scratch under a documented lock order.
  • Accounting + metrics: graph-pool vs scratch-KV bytes tracked separately; 7 fixed-name Prometheus metrics with bounded reason labels only.
  • Opt-in CLI (default disabled): --enable-decode-cuda-graphs and bucket/warmup/memory flags; the enable flag never marks a runtime safe.
  • Benchmark (benchmarks/serving/decode_cuda_graph.py, _fixture.py): paired eager/replay launch measurement; fixture mode needs no checkpoint; model mode requires --model + --offload-dir; no speedup claim.
  • Docs: docs/serving.md and docs/benchmarking.md rollout, limitations, metrics, rollback.

Scope / limitations (unchanged)

Offloaded MoE, FlashInfer plan path, DeepSeek MLA, prefill, and sampling remain eager/unsupported. No speedup is claimed or auto-enabled.

Verification

  • Plan test suites: 131 passed / 4 skipped (DeepSeek-not-installed + one CUDA-equivalence skip covered elsewhere). Full tests/python/serving: 214 passed.
  • tests/python/serving/test_cuda_graph.py: 20 passed (CUDA equivalence runs on-GPU, no skips).
  • CUDA replay parity + per-layer KV persistence: passed.
  • Rollback (MOE_DISABLE_CUDA_GRAPHS=1): passed, zero capture/replay.
  • Benchmark fixture smoke: exit 0, capability_reason=eligible, eager==replay, no speedup claim.
  • LSP diagnostics: zero errors across all 16 changed source files.
  • python -m compileall: exit 0.

Note: the single failing test_mixed_batch_flashinfer_correctness was verified to fail identically at the plan base commit (dd83676) — a pre-existing FlashInfer environment issue (missing CUDA dev headers), not a regression from this PR.

Type of Change

  • Production implementation (decode CUDA graph integration)
  • Tests (strict TDD: RED/GREEN per task)
  • Documentation update

Checklist

  • Graph-unsafe paths fail closed (deny-by-default capability)
  • Memory accounting and rollback specified and wired
  • One authoritative PagedKVStorage owner; owner/device identity enforced
  • Opt-in, disabled by default; kill switch honored per batch

drunkcoding @sisyphus-dev-ai
docs(plan): add decode CUDA graph plan
dd83676 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding

drunkcoding commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Momus plan review: [OKAY] on the first review round.

drunkcoding and others added 22 commits August 29, 2026 00:42
drunkcoding
refactor(serving): unify paged KV allocation and storage
a78e582 
Introduce PagedKVStorage as the single authoritative owner of block
allocation and per-layer native K/V tensors. Bind PagedKVCache and
PagedAttentionBackend to one instance, add an exact-type Qwen3 paged
attention registry that classifies DeepSeek V2/V3 MLA as unsupported,
and add an allocation-free current-token KV writer.

Ultraworked with Sisyphus
drunkcoding @sisyphus-dev-ai
feat(serving): require explicit decode graph capability
c1229ce 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding
refactor(serving): prepare native paged decode for replay
503b367
drunkcoding
feat(serving): add safe lazy decode CUDA graph runner
ab81002
drunkcoding
feat(serving): replay eligible decode batches
55c2b7e
drunkcoding @sisyphus-dev-ai
fix(serving): close CUDA graphs on engine shutdown
c939ece 
Make serving-engine graph cleanup idempotent so replacement and shutdown cannot leak scratch reservations.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
fix(serving): serialize engine replacement and reload
1e543e4 
Order active steps, graph invalidation, replacement, and application shutdown under one lifecycle lock.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat(serving): account decode graph memory
c74ff20 
Report graph private-pool bytes and authoritative scratch reservations separately from existing memory budgets.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat(serving): expose decode graph runtime stats
54e86cd 
Publish bounded capability, proof coverage, storage ownership, and graph memory state in engine statistics.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat(serving): export bounded decode graph metrics
a9b0fbf 
Expose fixed-name graph counters and gauges while restricting fallback labels to the enumerated reason set.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat(serving): require explicit graph capability provider
c59098c 
Pass typed graph-safety evidence separately from the lower-level offload engine and fail closed when omitted.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat(serving): add opt-in decode graph controls
5c9c44c 
Expose disabled-by-default batch, context, warmup, and memory controls through both server initialization paths.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test(serving): prove decode CUDA graph replay parity
e7d2bc1 
Keep full bucket logits inside capture and verify real-row slicing, metadata updates, fixed devices, and per-layer KV persistence on CUDA.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
bench(serving): add resident paged decode fixture
d212d67 
Build a persistent two-layer Qwen3 fixture with one authoritative storage owner and native paged attention.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
bench(serving): compare eager and graph decode launches
fce7373 
Record paired raw CUDA timings, launch profiles, replay coverage, memory, capability proof counts, and KV checksums without a speedup threshold.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
fix(serving): preserve exact graph device identity
cf82339 
Keep explicit CUDA indices intact on CPU-only gate hosts so storage and runner mismatches fail closed.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
fix(serving): ignore non-paged mock backend storage
36ac4b2 
Adopt authoritative storage only from the concrete native paged backend, preserving legacy eager engines and mocks.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
docs(serving): add decode CUDA graph rollout guide
95783fc 
Document native-Qwen3 eligibility, bounded monitoring, paired qualification, staged rollout, and immediate rollback without a speedup claim.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test(serving): finalize CUDA graph fallback coverage
3a8b4d0 
Verify resident non-paged and offloaded runtimes never capture, including immediate kill-switch precedence.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
bench(serving): measure model-mode eager fallback
694be33 
Record raw eager-only samples and zero graph activity when real loaders report unsafe decode graph capability.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
fix(serving): route mixed decode by paged registry
c35ee6c 
Use frozen per-layer registration evidence instead of class-name heuristics when splitting mixed batches.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
fix(serving): reject negative graph memory limits
29922aa 
Fail configuration validation rather than treating a negative private-pool budget as unlimited.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding drunkcoding changed the title WIP: plan decode CUDA graph integration WIP: add native-paged decode CUDA graphs Aug 29, 2026
drunkcoding @sisyphus-dev-ai
test(serving): update watchdog engine fixture
80140bb 
Accept the explicit decode graph capability provider added to both server initialization paths.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding

drunkcoding commented Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

Implementation completion update (80140bb):

  • Branch clean and synchronized.
  • GitHub CI: build, lint, type-check, formatting/unit tests, Python 3.10 and 3.12 unit suites, and automated review all passed.
  • Decode graph work remains limited to the documented resident Qwen3/native-paged eligibility boundary.

PR remains draft; no merge/ready/force operation was performed.

drunkcoding added 2 commits August 29, 2026 14:20
drunkcoding
fix(runtime): break paged KV storage import cycle
1c98853 
Defer the BlockAllocator import into PagedKVStorage.__init__ so importing
moe_infinity.runtime.paged_kv_storage (or any module that pulls it in, e.g.
the paged KV write kernel and attention registry) before moe_infinity.serving
no longer triggers a partially-initialized circular import via
serving/__init__ -> engine -> cuda_graph -> paged_kv_storage.

Ultraworked with Sisyphus
drunkcoding
docs(plan): mark decode CUDA graph tasks complete
93f278f 
All 12 tasks implemented, tested (RED/GREEN), and committed. Mark every
plan step checkbox to reflect completed TDD work.

Ultraworked with Sisyphus
@drunkcoding
drunkcoding marked this pull request as ready for review September 2, 2026 21:53
@drunkcoding drunkcoding changed the title WIP: add native-paged decode CUDA graphs feat: native-paged decode CUDA graphs Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@drunkcoding