[TianDi101]

Summary

Proof of concept for the RFC in #3259: a pluggable KV pool backend, with MORI-UMBP owning the L1 (HBM) pool so DP ranks share it. The UMBP side is ROCm/mori#994.

The default backend is today's behaviour. Nothing changes unless ATOM_KV_POOL_BACKEND=umbp is set.

Opened as a draft: this is a base version to go with the RFC discussion, not a finished feature.

Changes

1. Two interfaces (atom/model_engine/kv_block_pool.py), following the scheduler / runner process split:
   • KVBlockPool, used by BlockManager: which block is next, what each block holds, what is evicted. BlockPool becomes the default implementation with its logic unchanged. Three optional hooks (prepare, fetch, flush) default to no-ops.
   • KVPoolMemory, used by ModelRunner: where the pool's bytes come from, how they divide into blocks, and when the GPU has finished writing them. The default is today's torch allocation.
2. Call sites in five existing files (+84 / −5):
   • block_manager.py: builds the pool through make_block_pool, asks a content-sharing pool for a sequence's prefix once before its first walk, and tells the pool how many fresh blocks an allocate is about to take.
   • model_runner.py: takes the paged pool's buffer from KVPoolMemory, hands it the block layout, and notes each forward's writes in _record_kv_cache_ready.
   • scheduler.py: flushes the pool at the end of Scheduler.schedule() and Scheduler.postprocess(), the scheduler that owns a BlockManager.
   • block_pool.py, envs.py: the base class, and ATOM_KV_POOL_BACKEND, ATOM_UMBP_L1_ADDRESS, ATOM_UMBP_L1_NAMESPACE, ATOM_UMBP_L1_GROUP, ATOM_UMBP_L1_FETCH.
   • docs/environment_variables.md: a "KV pool backend" section for the new variables.
3. UMBP backend (atom/model_engine/umbp_block_pool.py), imported only when selected:
   • UmbpBlockPool keeps an exact local view of this rank's pool. Slot decisions (next slot, eviction, fetch) come from UMBP, and updates are batched into one call per flush.
   • UmbpPoolMemory maps the HBM that UMBP allocated (hip_ipc.open_ipc_handle), so kernels read and write it in place. After each forward it records a stream event and reports to UMBP when the GPU has finished.

Tests

• tests/test_umbp_block_pool.py: 48 tests. The pool runs against a Python model of the service and, when mori_umbp_l1 is built, against the real one over a Unix socket. They cover step-for-step parity with BlockPool, grants, fetch, the write-completion rules, and two BlockManagers sharing a prefix.
• On this branch (rebased on current main): ruff and black are clean, and the 36 existing test files that touch BlockManager, BlockPool or the runner's pool pass (1,987 passed, 69 skipped).
• GPU, on 4 × MI355X with MiniMax-M3 and 4 DP-attention ranks: results are in the RFC. That run was on main at 724f7bf. This branch is rebased on current main and has not been re-run on GPU since.

Known limits

• One runner per rank (no TP inside a DP rank), and only models whose paged regions do not overlap.
• The fetch is synchronous: the copy (20–40 ms for a 7k-token prefix) stalls that rank's step.
• Control messages to UMBP are gRPC calls (about 28 µs each).
• An engine restart resets its pool.

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 3267 --add-label <label>