Repository navigation
fix(moe): publish rocm cpu expert inputs before notification - #764
Conversation
FlashML-org#764 FlashML-org#763 FlashML-org#708 FlashML-org#687 deferred, FlashML-org#688 kept as a feature; FlashML-org#169 kept after its upstream close, FlashML-org#494/FlashML-org#154/FlashML-org#315 replaced upstream by FlashML-org#690/FlashML-org#698; FlashML-org#278 dropped with upstream's auto->hybrid pick (FlashML-org#734). Assisted-by: Claude Opus 5.5
|
Runtime-level fix: ROCm/rocm-systems#13877. The independent HIP reproducer shows that the host can observe a stream WRITE notification before data copied earlier on the same stream becomes visible. The CLR fix adds a system-scope release before the notification and captures that barrier for every graph replay. This FreeToken PR provides the application-level workaround and startup safety checks: use mapped-memory publish/combine kernels and reject unsafe graph paths. It remains useful on unpatched runtimes. The ROCm PR is still under review and targets |
|
Thanks for tracking this down to the runtime. One thing before merge: CUDA check on one H100 80GB (driver 580.95.05, torch 2.14.1+cu130): PYTHONPATH=python pytest tests/moe/test_cpu_moe.py tests/moe/test_cpu_moe_fp8_block.py tests/engine/test_moe_cpu_layers.py
Not tested: ROCm. |
|
This update enables BF16 KernelIO and locked-fetch on ROCm under CUDA's routing conditions, without changing shared arithmetic kernels. Startup validates the model dtype against synchronous CPU execution, with ordered-copy fallback and TP-consistent Graph-to-eager fallback; eager/API errors remain fatal. The failure is consistent with Validation:
Layer latency (µs): FP8/BF16, H2048/I768/E8/K4, three layers, batch 1, four CPU workers, Hybrid accumulation, one CPU expert/token; medians across three processes.
No-event timings are diagnostic. W7900D's small event/no-event difference is within process variation. Locked-fetch decode (tokens/s): KernelIO, BF16, batch-1 Graph, 512 input/128 greedy output tokens, Triton, six CPU workers, 8 GiB expert cache,
Medians of five samples; W7900D locked-on uses four after excluding one contended sample. Its off/on ranges are 41.557–41.646 / 42.779–43.394 tokens/s. These compare candidate paths/settings, not whole-main TPS. |
ac7fb19 to
5e1f3b3
Compare
|
With identical loaded HIP/ROCr/COMGR binaries and default allocations, our no-event tests observed ready before payload visibility on R9700; we did not reproduce this on W7900D. In the traced copy path, the payload copy ends with an AGENT-scope release, which alone does not guarantee CPU visibility. Inserting a system-release event between the copy and Our controlled tests and firmware analysis suggest two possible reasons W7900D did not reproduce the issue: its UC (uncached) payload mapping, and extra GPU L2 cache writeback (GL2_WB) in the firmware's SYSTEM-acquire handling before ready is written. The corresponding R9700 path appears to lack that extra writeback. This is consistent with W7900D passing after forcing the payload to NC (non-coherent) in eager mode, but failing in Graph with AGENT acquire. These findings concern the tested stack, not a general publication guarantee. CUDA's default StreamWriteValue contract includes a pre-write fence similar to |
|
Thanks for the update. Before merging, could you refresh the PR description? It still describes the 10-27 version, and it will become the squash commit message. Parts that no longer match the head (5e1f3b3):
|
CPU expert graph replay can observe a ready notification before its D2H inputs are visible on affected ROCm stacks, silently using stale activations or routes. Qualify the IO path before model capture and establish input publication before notifying the CPU.
kernel_io_probepasses, BF16 inputs use the existing mapped-memory publish/combine kernels. The publish kernel establishes system ordering and skips CPU work for empty routes; this path needs no publication event. Shared arithmetic kernels are unchanged.hipEventReleaseToSystemwithhipEventRecordExternalbetween D2H copies and the ready notification, so every Graph replay retainsD2H → event → ready. Events belong to each(layer, batch)task and remain alive through replay/recapture. Detect native memop capture at runtime and add explicit graph nodes when needed. Eager execution uses host callbacks if the memop handshake is unavailable. No custom HIP library is required.--moe-locked-fetch/--moe-locked-max-fetcheligibility as CUDA. The gather waits through stream memops when available or a GPU flag wait otherwise; the latter can constructGpuSyncwithout requiring memops. Update AMD installation guidance for the IO selection, startup checks, fallbacks and native rebuild.Validation reported by @LZ-QWQ in the latest test/model results:
Layer latency, FP8 weights/BF16 activations, H2048/I768/E8/K4, three layers, batch 1, four CPU workers, Hybrid accumulation and one CPU expert per token; medians across three processes:
Locked-fetch decode, KernelIO/BF16, batch-1 Graph, 512 input / 128 greedy output tokens, Triton attention, six CPU workers, 8 GiB expert cache,
moe_cpu_layers=2,hybrid_max_fetch=1; enabled settings aremoe_locked_fetch=0.5,moe_locked_max_fetch=2:Medians of five samples, except W7900D locked-on uses four after excluding one contended sample; W7900D off/on ranges are 41.557–41.646 / 42.779–43.394 token/s. These compare the candidate's IO paths/settings, not whole-main throughput. Full accuracy evaluation, long-duration stability and other ASICs remain outside the reported validation. The runtime follow-up supports retaining explicit publication on both tested AMD GPUs without claiming a general HIP notification guarantee or a deployed CLR fix.
Rebuild the native extension in the matching ROCm environment before serving:
Relevant repository checks:
Related: #350, #670; runtime investigation/fix: ROCm/rocm-systems#13877.