Repository navigation
Feat/umbp hbm ssd backends - #542
Merged
TianDi101 merged 5 commits intoAug 11, 2026
Merged
Conversation
Exercises the Phase 1-6 backend/transfer abstraction by actually adding a
medium and an engine through it, and writes down the resulting recipe. Two
findings came out of doing it rather than describing it.
FINDING 1: the interface was right, its one implementation was not.
MediumBackend was already medium-agnostic, and PageBackend's class comment
already claimed to serve "DRAM or HBM" — truthfully for the slot lifecycle,
bitmap allocator, event outbox, reaper, read leases and copy pins, none of
which consult the tier. But Init/Shutdown/BuildBufferRefs hardcoded
HostMemAllocator, MemoryLocationType::CPU and device=-1, so the claim did not
survive contact with a second medium.
PageMemorySource Allocate/Release/LocationType/Device/Name
Five methods holding exactly what differs per medium. HostPageMemorySource is
the old inline body, moved; the OwnershipConfig constructor now funnels into
the PageMemorySource one, so DRAM and HBM share ONE Init path rather than two
kept in step. LocationType/Device are the load-bearing pair: PageBackend
mirrors them into every published TransferRef, and that is what selects the
engine. Getting them wrong fails at neither build nor init.
So HbmPageMemorySource IS the HBM backend — no HbmBackend type, ~100 lines,
zero duplication of the 800 that already worked.
FINDING 2: HBM's local path was unservable by any engine in tree.
LocalCopyEngine requires BOTH endpoints loc == CPU
MoriIoEngine `if (src_remote == dst_remote) return false;`
A both-local pair with a GPU endpoint was claimed by nobody, so an HBM
backend's local Put/Get could not complete at all. The REMOTE side already
worked unmodified, exactly as §2 predicted ("why the remote path already works
against an HBM peer, and why the local fast path does not") — this closes the
local half.
HbmCopyEngine H2D + D2H + D2D, hipMemcpy, ~200 lines
No new TransferRef field needed: host_ptr is documented as the "process-local
view", and for hipMalloc'd memory the device pointer IS that view. The three
engines now partition (src, dst) with no overlap, so composite order stays
documentation rather than a tie-break — asserted, not assumed. Synchronous by
choice: the parallelism that matters is across keys and already exists in
PoolClient's executors, and a shared stream would need its own synchronization
for no win at KV-block sizes. Restores the caller's current device, since
Submit runs on threads that are not ours to leave re-pointed.
SSD: staged, and deliberately NOT a FileRef.
transfer_engine.h reserves a file endpoint for SsdBackend. Not taken. SSD bytes
are not addressable, and closing that gap in the transfer layer means a new
TransferRef kind, a PosixFileEngine, AND chaining in CompositeTransferEngine
for the remote reader — which that class explicitly does not implement. Instead
SsdBackend publishes ordinary registered host DRAM and spills behind it:
Allocate lends a staging page, Commit spills it to PeerSsdManager and returns
the page, Resolve fills a page under a read lease and the reaper reclaims it.
Cost is one host copy per side; benefit is SSD reaching the data plane with
ZERO transfer-layer change, and a remote peer needing no code at all. A FileRef
backend stays the right answer for GDS; this does not block it.
PeerSsdManager needed one accessor (SizeOf) — PrepareRead picks its staging
buffer before the read, so the reader must know the size while holding only the
key. Its header anticipated this adapter; Phase 0's dormancy ends here.
Two honest limits, both pinned by tests so a fix changes them deliberately:
- one key, one page. PrepareRead takes a single (ptr, cap), not a scatter
list. Fine while master's page_size IS the KV block size.
- staging exhaustion degrades a Get to found=false, which makes the client
retry another peer for a key this node does hold. This is the "not ready,
retry here" state medium_backend.h records as proposed and rejected.
Skills, written from what the work actually required:
.claude/skills/umbp-add-backend two shapes: 5-method
PageMemorySource, or full
MediumBackend + staging
.claude/skills/umbp-add-transfer-engine pair dispatch, disjointness,
Plan/Submit/Wait, PeerDirectory
Both carry the contracts with no compile-time protection (one event bundle per
seq, full-sync clearing the outbox in the same critical section, kFailedNoSpace
vs kFailed, Evict's positional results) and the build's real constraint: umbp
needs protoc/grpc_cpp_plugin from the mori image, not the bare host.
Verified: ctest 38/38 (36 pre-existing + 2 new), including test_page_backend
and test_transfer_engine unchanged. test_hbm_backend runs real H2D/D2H/D2D and
a Put/Get round trip on an MI355X; test_ssd_backend runs 21 cases against a
real posix SSD tier. clang-format clean.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ibutedClient
Migrated from the worktree-umbp-bench-tier branch (commit b634ae4a plus the
uncommitted work that sat on top of it).
- UMBPHbmConfig (enabled/device/capacity_bytes) on UMBPDistributedConfig,
lowered by ToPoolClientConfig into PoolClientConfig::hbm, so PoolClient::Init
registers an HBM PageBackend. Unlike dram/ssd there is no ownership struct to
thread through: hipMalloc has no hugepage/NUMA/prefault dimension.
- IUMBPClient::RegisterMemory grows loc/device describing the CALLER's
allocation, so a GPU-resident src/dst routes through HbmCopyEngine instead of
being assumed host memory. StandaloneProcessClient rejects non-CPU rather
than silently mis-registering.
- umbp_bench.py: --tier {dram,hbm,ssd}, --dst-loc {host,gpu}, --hbm-device.
DeviceBuffer allocates the read destination via raw HIP (ctypes on
libamdhip64) so --tier hbm --dst-loc gpu exercises the D2D hipMemcpy path
without pulling in a torch dependency.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…peers
Migrated from the uncommitted work-in-progress on feat/umbp-hbm-ssd-backends
in the main checkout, so HBM/SSD can actually be benchmarked: PoolClient::Init
always registers DRAM, so --tier hbm/ssd produces a two-backend node, which is
exactly the case this fixes.
buffer_index is backend-local, but EnsurePeerServiceConnection folded every
backend's buffers into one flat dram_memory_descs list. One medium's buffers
were published, the rest unreachable -- yet HasRemoteBuffers still went true,
so the next resolve asked the peer to omit descriptors and the missing media's
pages were read against the published medium's memory. Now each desc names its
backend, and Build{Put,Get}Transfers snapshot per backend the batch touches
rather than once per peer.
pool_client.{h,cpp} carried changes from both this work and the bench-tier
migration; both are preserved (SlotPlan::backend_id plus the loc/device
RegisterMemory signature).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ier lying SsdBackend was already registered by PoolClient::Init from PoolClientConfig::ssd -- the only thing missing was DistributedClient lowering UMBPConfig::ssd into it, so the tier stayed dark however it was configured. The bench's "SSD is unwired" refusal was describing the old PeerSsdManager world and is now wrong, so it is gone. Opt-in is UMBPDistributedConfig::enable_ssd_tier, deliberately NOT UMBPConfig::ssd.enabled -- that defaults to true, so keying off it would make every existing distributed deployment start advertising SSD capacity. Also fixes the benchmark measuring the wrong medium: DRAM is always registered, and routing has no tier order (kMostAvailable picks whichever medium has the most free bytes), so --tier hbm against the 8 GiB DRAM default and a 4 GiB HBM pool routed every put to DRAM and labelled it HBM. Non-DRAM tiers now shrink DRAM to a 256 MiB floor so most-available can only pick the requested medium. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e re-cache Two things that made a tier benchmark measure something other than the tier. SSD reads stage through a registered host arena of ssd_staging_buffer_slots pages; the default 16 caps read concurrency at 16 keys, and a slot stays pinned for UMBP_SSD_READ_LEASE_MS when the reader's best-effort ReleaseSsdLease is lost. A 1024-key sweep therefore returned hits=17 misses=1007 -- the overflow degrades to MISS, not to a slow hit, so it reads as a correctness failure rather than a capacity limit. ssd_backend.cpp says it directly: "Sizing the arena for read concurrency is the mitigation." Slots and arena bytes are now env-tunable and default to 512 / 2 GiB, which takes the same sweep to 1024 hits / 0 misses / 0 mismatches. cache_remote_fetches defaults on, and the reader gets its own pool of the tier under test, so after pass 1 most reads were served from the reader's own memory instead of RDMA out of the writer's tier. UMBP_CACHE_REMOTE_FETCHES=0 pins every pass to a genuine remote read; the default keeps production behavior. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
TianDi101
added a commit
that referenced
this pull request
Sep 3, 2026
* feat(umbp): HBM + SSD backends, HbmCopyEngine, and the two how-to skills
Exercises the Phase 1-6 backend/transfer abstraction by actually adding a
medium and an engine through it, and writes down the resulting recipe. Two
findings came out of doing it rather than describing it.
FINDING 1: the interface was right, its one implementation was not.
MediumBackend was already medium-agnostic, and PageBackend's class comment
already claimed to serve "DRAM or HBM" — truthfully for the slot lifecycle,
bitmap allocator, event outbox, reaper, read leases and copy pins, none of
which consult the tier. But Init/Shutdown/BuildBufferRefs hardcoded
HostMemAllocator, MemoryLocationType::CPU and device=-1, so the claim did not
survive contact with a second medium.
PageMemorySource Allocate/Release/LocationType/Device/Name
Five methods holding exactly what differs per medium. HostPageMemorySource is
the old inline body, moved; the OwnershipConfig constructor now funnels into
the PageMemorySource one, so DRAM and HBM share ONE Init path rather than two
kept in step. LocationType/Device are the load-bearing pair: PageBackend
mirrors them into every published TransferRef, and that is what selects the
engine. Getting them wrong fails at neither build nor init.
So HbmPageMemorySource IS the HBM backend — no HbmBackend type, ~100 lines,
zero duplication of the 800 that already worked.
FINDING 2: HBM's local path was unservable by any engine in tree.
LocalCopyEngine requires BOTH endpoints loc == CPU
MoriIoEngine `if (src_remote == dst_remote) return false;`
A both-local pair with a GPU endpoint was claimed by nobody, so an HBM
backend's local Put/Get could not complete at all. The REMOTE side already
worked unmodified, exactly as §2 predicted ("why the remote path already works
against an HBM peer, and why the local fast path does not") — this closes the
local half.
HbmCopyEngine H2D + D2H + D2D, hipMemcpy, ~200 lines
No new TransferRef field needed: host_ptr is documented as the "process-local
view", and for hipMalloc'd memory the device pointer IS that view. The three
engines now partition (src, dst) with no overlap, so composite order stays
documentation rather than a tie-break — asserted, not assumed. Synchronous by
choice: the parallelism that matters is across keys and already exists in
PoolClient's executors, and a shared stream would need its own synchronization
for no win at KV-block sizes. Restores the caller's current device, since
Submit runs on threads that are not ours to leave re-pointed.
SSD: staged, and deliberately NOT a FileRef.
transfer_engine.h reserves a file endpoint for SsdBackend. Not taken. SSD bytes
are not addressable, and closing that gap in the transfer layer means a new
TransferRef kind, a PosixFileEngine, AND chaining in CompositeTransferEngine
for the remote reader — which that class explicitly does not implement. Instead
SsdBackend publishes ordinary registered host DRAM and spills behind it:
Allocate lends a staging page, Commit spills it to PeerSsdManager and returns
the page, Resolve fills a page under a read lease and the reaper reclaims it.
Cost is one host copy per side; benefit is SSD reaching the data plane with
ZERO transfer-layer change, and a remote peer needing no code at all. A FileRef
backend stays the right answer for GDS; this does not block it.
PeerSsdManager needed one accessor (SizeOf) — PrepareRead picks its staging
buffer before the read, so the reader must know the size while holding only the
key. Its header anticipated this adapter; Phase 0's dormancy ends here.
Two honest limits, both pinned by tests so a fix changes them deliberately:
- one key, one page. PrepareRead takes a single (ptr, cap), not a scatter
list. Fine while master's page_size IS the KV block size.
- staging exhaustion degrades a Get to found=false, which makes the client
retry another peer for a key this node does hold. This is the "not ready,
retry here" state medium_backend.h records as proposed and rejected.
Skills, written from what the work actually required:
.claude/skills/umbp-add-backend two shapes: 5-method
PageMemorySource, or full
MediumBackend + staging
.claude/skills/umbp-add-transfer-engine pair dispatch, disjointness,
Plan/Submit/Wait, PeerDirectory
Both carry the contracts with no compile-time protection (one event bundle per
seq, full-sync clearing the outbox in the same critical section, kFailedNoSpace
vs kFailed, Evict's positional results) and the build's real constraint: umbp
needs protoc/grpc_cpp_plugin from the mori image, not the bare host.
Verified: ctest 38/38 (36 pre-existing + 2 new), including test_page_backend
and test_transfer_engine unchanged. test_hbm_backend runs real H2D/D2H/D2D and
a Put/Get round trip on an MI355X; test_ssd_backend runs 21 cases against a
real posix SSD tier. clang-format clean.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* feat(umbp-bench): --tier/--dst-loc flags and HBM wiring through DistributedClient
Migrated from the worktree-umbp-bench-tier branch (commit b634ae4a plus the
uncommitted work that sat on top of it).
- UMBPHbmConfig (enabled/device/capacity_bytes) on UMBPDistributedConfig,
lowered by ToPoolClientConfig into PoolClientConfig::hbm, so PoolClient::Init
registers an HBM PageBackend. Unlike dram/ssd there is no ownership struct to
thread through: hipMalloc has no hugepage/NUMA/prefault dimension.
- IUMBPClient::RegisterMemory grows loc/device describing the CALLER's
allocation, so a GPU-resident src/dst routes through HbmCopyEngine instead of
being assumed host memory. StandaloneProcessClient rejects non-CPU rather
than silently mis-registering.
- umbp_bench.py: --tier {dram,hbm,ssd}, --dst-loc {host,gpu}, --hbm-device.
DeviceBuffer allocates the read destination via raw HIP (ctypes on
libamdhip64) so --tier hbm --dst-loc gpu exercises the D2D hipMemcpy path
without pulling in a torch dependency.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* feat(umbp): per-backend buffer indices (backend_id) for multi-medium peers
Migrated from the uncommitted work-in-progress on feat/umbp-hbm-ssd-backends
in the main checkout, so HBM/SSD can actually be benchmarked: PoolClient::Init
always registers DRAM, so --tier hbm/ssd produces a two-backend node, which is
exactly the case this fixes.
buffer_index is backend-local, but EnsurePeerServiceConnection folded every
backend's buffers into one flat dram_memory_descs list. One medium's buffers
were published, the rest unreachable -- yet HasRemoteBuffers still went true,
so the next resolve asked the peer to omit descriptors and the missing media's
pages were read against the published medium's memory. Now each desc names its
backend, and Build{Put,Get}Transfers snapshot per backend the batch touches
rather than once per peer.
pool_client.{h,cpp} carried changes from both this work and the bench-tier
migration; both are preserved (SlotPlan::backend_id plus the loc/device
RegisterMemory signature).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* feat(umbp): wire the SSD tier through DistributedClient, and stop --tier lying
SsdBackend was already registered by PoolClient::Init from
PoolClientConfig::ssd -- the only thing missing was DistributedClient lowering
UMBPConfig::ssd into it, so the tier stayed dark however it was configured.
The bench's "SSD is unwired" refusal was describing the old PeerSsdManager
world and is now wrong, so it is gone.
Opt-in is UMBPDistributedConfig::enable_ssd_tier, deliberately NOT
UMBPConfig::ssd.enabled -- that defaults to true, so keying off it would make
every existing distributed deployment start advertising SSD capacity.
Also fixes the benchmark measuring the wrong medium: DRAM is always
registered, and routing has no tier order (kMostAvailable picks whichever
medium has the most free bytes), so --tier hbm against the 8 GiB DRAM default
and a 4 GiB HBM pool routed every put to DRAM and labelled it HBM. Non-DRAM
tiers now shrink DRAM to a 256 MiB floor so most-available can only pick the
requested medium.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* fix(umbp-bench): size the SSD staging arena, and let tier runs disable re-cache
Two things that made a tier benchmark measure something other than the tier.
SSD reads stage through a registered host arena of ssd_staging_buffer_slots
pages; the default 16 caps read concurrency at 16 keys, and a slot stays
pinned for UMBP_SSD_READ_LEASE_MS when the reader's best-effort
ReleaseSsdLease is lost. A 1024-key sweep therefore returned hits=17
misses=1007 -- the overflow degrades to MISS, not to a slow hit, so it reads
as a correctness failure rather than a capacity limit. ssd_backend.cpp says it
directly: "Sizing the arena for read concurrency is the mitigation." Slots and
arena bytes are now env-tunable and default to 512 / 2 GiB, which takes the
same sweep to 1024 hits / 0 misses / 0 mismatches.
cache_remote_fetches defaults on, and the reader gets its own pool of the tier
under test, so after pass 1 most reads were served from the reader's own
memory instead of RDMA out of the writer's tier. UMBP_CACHE_REMOTE_FETCHES=0
pins every pass to a genuine remote read; the default keeps production
behavior.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist