feat(umbp): GPU Direct Storage (hipfile) SSD-to-GPU read path - #579
Merged
TianDi101 merged 9 commits intoSep 3, 2026
Merged
Conversation
TianDi101
force-pushed
the
feat/umbp-gds
branch
from
September 3, 2026 05:14
51f54fc to
914a901
Compare
TianDi101
force-pushed
the
refactor/umbp-backend-agnostic
branch
from
September 3, 2026 05:54
0a46875 to
4e86725
Compare
TianDi101
force-pushed
the
feat/umbp-gds
branch
from
September 3, 2026 06:10
914a901 to
d64ac4d
Compare
TransferRef gains a File kind (O_DIRECT fd + offset + opaque engine handle) beside its memory handles; MemoryRegistrar gains RegisterFile, which CompositeTransferEngine fans out to the first file-capable engine. Memory engines ignore file refs, so per-pair selection routes them unchanged. Groundwork for GdsEngine reading SSD segments straight into GPU memory.
GdsEngine implements TransferEngine over AMD hipFile: RegisterFile turns an O_DIRECT fd into a ref-counted hipFile handle, and Submit reads a segment's byte range straight into device memory with hipFileRead, no host bounce. CanHandle claims only (file source, GPU destination) -- the prefill KV-load path -- so the memory engines are untouched. Auto-enabled when hipfile is installed (UMBP_ENABLE_GDS); the build is unchanged without it.
Add TierBackend::LocateRecord returning a value's (fd, 4K-aligned offset, padded/actual size, direct_io) so a zero-copy reader can DMA it itself instead of asking the tier to copy. SSDTier reads it from the segment index; ShardedSsdTier routes to the owning drive; PeerSsdManager forwards it, skipping unknown or evicting keys. Read-only, no device IO -- GdsEngine turns it into a FileRef.
BatchResolve checks LocateRecord first: when a file engine is configured and the record is on an O_DIRECT fd, it publishes a FileRef (segment fd + the value's aligned offset) in ResolvedEntry and skips the staging arena. The fd's hipFile handle is obtained once per fd and released in Shutdown. PoolClient's local get reads a file_ref straight into the caller's device buffer. Inert until the engine is registered (next commit) -- the staging path is unchanged, existing SsdBackend tests still pass.
…o end PoolClient::Init adds GdsEngine to the composite (gated on UMBP_ENABLE_GDS), after the memory engines it never overlaps with, so an SsdBackend FileRef is routed to a hipFileRead straight into device memory. A new end-to-end test puts a key to the SSD tier, resolves it to a FileRef, and reads the segment range into GPU memory through the engine, verifying the bytes.
SsdBackend reads UMBP_ENABLE_GDS at Init; setting it to 0 keeps every SSD read on the staging arena even when the build has hipfile and a GdsEngine is registered. Default on. Correctness never depends on it (hipfile's compat mode already covers fastpath misses) -- this is for GDS-vs-staging A/B and as an operational safety valve.
Flip gds_enabled_ to false so SSD reads stay on the staging arena unless UMBP_ENABLE_GDS=1 is set; a hipfile-enabled build no longer changes read behaviour by default. The end-to-end SSD->GPU test opts in explicitly and restores the default on exit.
…able SsdBackend::Resolve published a file ref with an empty page list, but only ExecuteLocalGet ever read one. The other two local readers -- ServeLocalGets (which is what PoolClient::Get and BatchGet actually run through) and BatchGetRanges -- went straight to BuildLocalPageTransfers, hit pages.empty() and reported a miss. A peer fared worse: a file ref is never serialized, so a remote reader decoded a slot with no pages and logged it as malformed. With UMBP_ENABLE_GDS=1 that cost two medium_selection tests and four ranged-io subtests; with it off the whole feature was simply unreachable. BatchResolve now takes `allow_file_refs`, defaulted to false so that every caller that cannot consume one -- peer_service serving the wire, and the existence checks that read only `found` -- keeps getting staged pages without having to know the flag exists. Testing it before gds_enabled_ also spares those callers LocateRecord and the hipFile registration GdsHandleForFd does as a side effect. The two local readers now consume the ref the way ExecuteLocalGet already did, with BuildLocalFileRangeTransfers as the ranged counterpart: one read per requested range, GdsEngine adding src_offset to the ref's file_offset. A file ref is also only readable INTO DEVICE MEMORY -- GdsEngine claims the (file, GPU) pair and nothing claims (file, host) -- and the medium cannot know the destination at resolve time. ResolveLocalBatch therefore re-resolves any host-destined key that came back with a ref, so it lands on staged pages instead. That is a second pass rather than classifying every destination up front on purpose: with GDS off nothing publishes a ref, the pass finds nothing, and the ordinary local read pays no extra HIP call. Verified on smci355-ccs-aus-n06-21 with hipFile 0.2.0 installed and the SSD tier on ext4/nvme (O_DIRECT probe passing). Full UMBP suite, 52 tests: UMBP_ENABLE_GDS=1 went from 50/52 to 52/52, and the default-off run stays at 52/52. umbp_medium_selection (both the local round trip and the cross-node SSD peer case), umbp_ssd_ranged_io (local and remote), test_ssd_backend, test_gds_engine and test_transfer_engine each pass 3/3 with GDS on. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Nothing in the tree could measure the GDS path. Every existing bench drives Get / BatchGet / BatchGetRanges against DRAM or HBM with host destinations, and a host destination can never take a file ref -- GdsEngine claims only the (file, GPU) pair -- so they all measure the staging arm whatever the switch is set to. bench_umbp_gds_ssd_get stands up an SSD-only node, writes N objects, and reads them back into hipMalloc'd device memory through IUMBPClient::BatchGet. Mode is not a flag: run the binary twice, once with UMBP_ENABLE_GDS=1, and diff the CSV, so both arms are the same binary against the same data one environment variable apart. Two passes, because they answer different questions. `cold` is first touch, the only pass where zero-copy can win. `warm` re-reads the same keys, which now hold a read lease -- BatchResolve serves a lease before it ever considers a file ref, so the staging arm answers from the arena at memcpy speed while the GDS arm re-reads the device every time. That gap is the point: GDS does not cache, and the bench should be able to show it rather than assert it. --slots exposes the staging arena in pages. It defaults to the whole working set, because a staged page is held until its read lease expires rather than until the copy finishes; an arena sized to one batch fills after a couple of batches and every later resolve returns kBusy and sleeps, which measures the backoff instead of the data path. Pass it explicitly to study that regime on purpose -- it is real, and it is where zero-copy helps most, because a file ref occupies no arena at all. The run also verifies the bytes after the timed passes. That is not incidental: no test in the tree reads a GDS-served value into device memory through the public client API, so this is the only end-to-end check that the file ref path returns the right data. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
TianDi101
force-pushed
the
feat/umbp-gds
branch
from
September 3, 2026 06:15
d64ac4d to
ac5cd7b
Compare
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
2 participants
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.
Summary
Filetransfer endpoint (FileRef: fd / offset / padded_size / gds_handle) plusMemoryRegistrar::RegisterFile, and a hipFile-backedGdsEnginethat DMAs SSD segment ranges straight into GPU memory with no host bounce.LocateRecordonSsdTier/ShardedSsdTier/PeerSsdManagerexposes the physical O_DIRECT location of a record;SsdBackend::BatchResolvepublishes aFileReffor O_DIRECT records;PoolClientregistersGdsEngineand routes(File -> GPU)items to it.hipfileis found (UMBP_ENABLE_GDS);UMBP_ENABLE_GDS=0falls back to the existing staging path;SsdBackend::ClearLocal/Shutdownderegister cached hipFile handles to avoid stale fd reuse after segment rotation.Test plan
test_transfer_engine: File endpoint + CompositeRegisterFilefan-out + File->GPU routingtest_gds_engine: CanHandle / Plan bounds + end-to-endhipFileReadinto device memorytest_ssd_direct_io:LocateRecordpread contract (known + unknown key)test_ssd_backend: end-to-end SSD->GPU via GDS; 21 existing staging tests unchangedUMBP_ENABLE_GDS=0verified to fall back to stagingTMPDIR=/mnt/gds)