Repository navigation
Conversation
…r (hipHostRegister on ROCm)
…read batch per message pass
Every slice the fanout cuts carries the producer's uuid, so the batch's check for a tensor it is already copying handed a second slice the first one's rows. It now compares slices as the receive cache does, and tests read TP-sharded edges in a batch.
Collaborator
This branch has not been deployed
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.
Stacked on #1089 (pr-23-t5-host).
Why
The SHM arena page-locks its segments with
cudaHostRegister, which ROCm doesn't have: MI325X workers logged0 MiB pinned, so a 4 KB D2H took 26 us (one GPU, shared node). Each small device tensor also took its own arena slot and copy, about 5400 copy kernels a second on an Omni three-worker layout at c=32 and 2000 on Qwen3-TTS at c=64. Split layouts cross this path every step.What changed
mstar/platforms/host_memory.py,mstar/platforms/interface.py:Platform.host_register/host_unregister,hipHostRegisteron ROCm through the HIP runtime torch already loaded,cudaHostRegisteron CUDA (the existing call, moved). 256 MiB pinned per worker on ROCm.MSTAR_SHM_ARENA_PINstill turns it off.mstar/communication/arena.py, sender: device tensors of up to 1 MiB from oneregister_for_sendshare an arena slot of up to 8 MiB, filled by one gather and one D2H, each row 16-byte aligned.mstar/communication/arena.py,mstar/worker/worker.py, receiver: one read batch per message pass, copying each contiguous run of a peer segment once.transfer_seconds(docs/metrics.rst). No file underrust/changes: a row in a slot is a tensor with an offset, which both bookkeepers and the wire codec already carry.Verification
test_arena_transport.py,test_platforms.py,test_shm_tensor_comm.py180 passed, 5 skipped.Not in this PR
_copy_runsstill clones for a CPU-device worker inside a read batch, where API server: size the torch pool from the CPU quota, one-memcpy arena reads, pooled image encodes #903 made the plain arena read one memcpy. Only that case reaches it.arena.py(slicing large reads) after this base. Not combined yet.