Repository navigation
fix(ple): add rocm disk wait-sync support - #776
Conversation
FlashML-org#777 FlashML-org#776 FlashML-org#774 FlashML-org#781 FlashML-org#779 deferred. Assisted-by: Claude Opus 5.5
| (void)hipStreamDestroy(stream); | ||
| } | ||
| (void)hipGetLastError(); | ||
| TORCH_CHECK(rc == hipSuccess || rc == hipErrorNotSupported, |
There was a problem hiding this comment.
In auto mode a failed probe should fall back to launch-gating, as the CUDA branch of probe_wait_sync and GpuSync::probe in cpu_moe/gpu_sync_hip.cpp do. This check only accepts hipErrorNotSupported, but clr reports most capability failures as hipErrorInvalidValue (ihipStreamOperation: memory that is not signal memory, pointer lookup failure). So on such a runtime the engine fails to start, where gate would have worked. wait mode already raises when the probe returns false, so the probe itself can return false on any error.
GpuSync::probe also does two more checks that would help here:
- it reads
hipDeviceAttributeCanUseStreamWaitValuefirst; - a host thread releases the WAIT after the launch. Here the graph writes 7 and then waits for >= 7, so the WAIT is always satisfied by the graph itself. The probe never checks that a host store into signal memory releases a device WAIT, which is the path the fill relies on.
| self._worker.shutdown(wait=True) | ||
| raise | ||
|
|
||
| def close(self) -> None: |
There was a problem hiding this comment.
close() and _closed have no caller, in the engine or in the tests, so _closed is never true. Can they be dropped?
Same for the return value at L244: the HIP branch of forward_host_ctx doesn't use deferred (it awaits in its finally), so _start_fill could return None.
| histories = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs] | ||
|
|
||
| @torch.inference_mode() | ||
| def _complete() -> None: |
There was a problem hiding this comment.
This _complete is the same as the CUDA one at L315, except that histories is computed before the submit. Could both share one helper, e.g. _fill_decode(histories, bs)? Then the readback wait, the fill and the signal-on-error sequence live in one place.
| PleStore = _row_store.PleStore | ||
|
|
||
|
|
||
| def alloc_sync_flag() -> torch.Tensor: |
There was a problem hiding this comment.
Nit: the module docstring (L6-7) still says the captured lookup WAITs on a pinned flag. With this function it is pinned on CUDA and HIP signal memory on ROCm.
|
Addressed all five review comments in two commits. All 18 tests across the three test files pass on RTX4090, W7900D, and R9700, including the 16,384-kernel regression. Performance rechecks on W7900D and R9700 were broadly consistent with the previous PR revision. |
Related: #697
ROCm disk PLE currently uses
gatefallback. This PR adds HIP wait-sync and prevents deadlocks when submitting large graphs with wait-sync.It adds HIP signal memory, stream memops and capture probing (
autoprobes;waitrequires support). Use HIP signal memory for Engram's wait-sync flag.#697 uses early
cuGraphUploadto avoid the first-launch deadlock on CUDA. In our ROCm 10.1.0 large-graph checks,hipGraphUploadfollowed by stream synchronization did not prevent the deadlock. This may be related to the current HIP upload implementation:ihipGraphUploadreturns after handle validation and stream resolution, whilehipGraphUploadstill marks preparatory work as TODO. To avoid this deadlock, this PR starts a HIP-only CPU worker before replay to fill PLE rows and signal the flag independently of graph submission.CUDA's existing allocation, execution and synchronization logic is unchanged.
Validation
The baseline uses main's unchanged production code, with the new 16,384-kernel test included in the test suite.
Test files:
tests/kernels/test_row_store.pytests/models/qwen4_exp/test_ple_disk.pytests/models/deepseek_v41/test_engram.pyThe new regression and two existing Engram worker cases now run on ROCm.
Performance reference
Qwen/Qwen3.8-Flash-Next-FP8, 256 input / 128 output tokens, single GPU/request, MoE offload cache 1024, graph batch 1, ROCm 10.1.0. Compare main/gate, PR/gate and PR/wait under identical settings.
Each configuration was tested in two independent runs, with three measurements after warmup per run. The table reports the average of the two run medians.