diff --git a/src/cco/cco_init.cpp b/src/cco/cco_init.cpp index b8f1f57ab..622cc685b 100644 --- a/src/cco/cco_init.cpp +++ b/src/cco/cco_init.cpp @@ -99,7 +99,29 @@ static hipMemAllocationType CcoWindowAllocType() { const char* e = getenv("CCO_UNCACHED_WINDOW"); return e && atoi(e) == 0; }(); - return cached ? hipMemAllocationTypePinned : hipMemAllocationTypeUncached; + if (cached) return hipMemAllocationTypePinned; + // hipMemAllocationTypeUncached was added on 2025-09-11 (HIP_VERSION > 70051831). + // Guard the symbol so CCO still builds on older ROCm (e.g. 7.0.0); mirror the + // shmem path in symmetric_memory.cpp: use the raw value 0x40000000 on the + // 70051831 build where the symbol may be missing (excluding the buggy + // 7c9236b16 build), and fall back to Pinned on anything older. +#if HIP_VERSION > 70051831 + return hipMemAllocationTypeUncached; +#elif HIP_VERSION == 70051831 + if (strcmp(HIP_VERSION_GITHASH, "7c9236b16") != 0) { + return static_cast(0x40000000); // hipMemAllocationTypeUncached + } + MORI_SHMEM_WARN( + "CCO uncached window requested but ROCm build {} has a known " + "hipMemAllocationTypeUncached issue; falling back to Pinned memory", + HIP_VERSION_GITHASH); + return hipMemAllocationTypePinned; +#else + MORI_SHMEM_WARN( + "CCO uncached window requested but ROCm version does not support " + "hipMemAllocationTypeUncached; falling back to Pinned memory"); + return hipMemAllocationTypePinned; +#endif } // Local slot base = the VA where this rank's slice of the flat VA starts. diff --git a/tests/python/ops/dispatch_combine_v2/bench_dispatch_combine.py b/tests/python/ops/dispatch_combine_v2/bench_dispatch_combine.py index edd28c88c..8653930e8 100644 --- a/tests/python/ops/dispatch_combine_v2/bench_dispatch_combine.py +++ b/tests/python/ops/dispatch_combine_v2/bench_dispatch_combine.py @@ -279,6 +279,11 @@ def mock_fmoe(): def time_eager(fn, ct): for _ in range(WARMUP): fn(ct) + sync() + comm.barrier() # lock-step warmup: sync(device)+barrier so no + # rank starts call N+1 before all finish N (avoids the cco cross- + # device barrier flag-overwrite deadlock at small token counts). + # Untimed, so it does not affect the measured replay bandwidth. sync() comm.barrier() s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event( @@ -295,6 +300,11 @@ def time_eager(fn, ct): def time_graph(fn, ct): for _ in range(WARMUP): fn(ct) + sync() + comm.barrier() # lock-step warmup: sync(device)+barrier so no + # rank starts call N+1 before all finish N (avoids the cco cross- + # device barrier flag-overwrite deadlock at small token counts). + # Untimed, so it does not affect the measured replay bandwidth. sync() gr = torch.cuda.CUDAGraph() with torch.cuda.graph(gr):