Skip to content

Pinned SHM arena on ROCm and one staging copy per step - #1095

Draft
merceod wants to merge 5 commits into
users/atindra/pr-23-t5-hostfrom
users/atindra/pr-29-s5-pinned-arena
Draft

merceod wants to merge 5 commits into
users/atindra/pr-23-t5-hostfrom
users/atindra/pr-29-s5-pinned-arena

Conversation

@merceod

@merceod merceod commented Nov 1, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #1089 (pr-23-t5-host).

Why

The SHM arena page-locks its segments with cudaHostRegister, which ROCm doesn't have: MI325X workers logged 0 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, hipHostRegister on ROCm through the HIP runtime torch already loaded, cudaHostRegister on CUDA (the existing call, moved). 256 MiB pinned per worker on ROCm. MSTAR_SHM_ARENA_PIN still turns it off.
  • mstar/communication/arena.py, sender: device tensors of up to 1 MiB from one register_for_send share 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.
  • A read batch keeps the slices of a TP-sharded tensor apart, keyed as [Transport][Bug fix] Key received slices by offset; refuse sharded streams across TP sizes #623 does. Without that, a second slice read in one batch came back with the first slice's rows. Read-batch variants of [Transport][Bug fix] Key received slices by offset; refuse sharded streams across TP sizes #623's TP tests cover it.
  • The slots and the read batch run on every platform with a device, so CUDA takes them too. Pinning is the only per-platform part.
  • A batch's copy time is shared across its edge reads in transfer_seconds (docs/metrics.rst). No file under rust/ changes: a row in a slot is a tensor with an offset, which both bookkeepers and the wire codec already carry.

Verification

  • CPU: ruff clean. r9b's targeted set 764 passed, 35 skipped. This branch's tip with that set plus the arena, platform, SHM, scheduler and metrics tests 1073 passed, 40 skipped, 0 failed, none failing only here. test_arena_transport.py, test_platforms.py, test_shm_tensor_comm.py 180 passed, 5 skipped.
  • GPU, one MI325X: test_arena_transport 71, test_platforms 56, test_shm_tensor_comm 57 passed, 1 skipped.
  • Whole week-2 stack, not this PR alone, window G, 8x MI325X, exclusive, Qwen3-TTS 1.7B split bf16, two timed runs a cell, tip f7122c3 against r9b f906c7d:
arm c=1 audio s/s c=16 c=64 TTFA p50 c=1 / c=64 ms WER
tip 13.5 115.1 177.7 16.4 / 105.6 1.31% PASS
r9b 13.4 112.5 172.5 18.2 / 152.9 1.965% WARN
vllm-omni 10.0 85.0 164.0 34.4 / 236.7 1.528% PASS

Not in this PR

merceod and others added 5 commits October 31, 2026 23:37
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.
@vasilevklart

Copy link
Copy Markdown
Collaborator

"In #1095 batched staging does host.copy_(torch.cat(pieces)) so a batch of host sources would go back through the intra-op pool" from #1129

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants