Pull requests / #3625

#3625 [Engine] release private checkpoint-copy streams when worker threads exit

open · draft · @HeimEndyd · 0 comments · View on GitHub

Server & APIMulti-GPUAMD / HIPNVIDIA / CUDAWindowsLinux

Description

## Problem and result

Checkpoint copies keep their private GPU streams in a thread-local vector of raw handles. Destroying that vector at worker exit does not destroy the streams. A layer-split prefill creates short-lived `std::async(std::launch::async)` stage workers, so checkpoint copies can create another retained stream on each worker.

This change gives the per-thread streams an owner that destroys them on their device when the worker exits, then restores the thread's previous device. Copies and reuse within a living worker are unchanged.

## Related discussion

I searched the current open issue/PR listings and relevant closed discussions on 2026-10-21, including searches for `private_stream`, `thread_local`, `StreamDestroy`, and stream leaks. I found no existing report or PR for this ownership defect. #3288 concerns snapshot staging on Windows/WDDM, and #3483 concerns expert stage-buffer fences; neither destroys these private checkpoint-copy streams. This is independent of slot-placement and checkpoint-retention work in #2384, #2526 and #2843.

## Approach

- A thread-local owner retains the existing per-device vector. On thread exit, it destroys each non-null stream on the device that created it and restores the prior device.
- A host-only regression uses the real header with mocked runtime calls. It checks joined short-lived workers, reuse within one worker, concurrent workers across devices, device restoration, and creation/copy/synchronization failures.
- The existing CUDA transfer fixture now wraps the two runtime calls made by cleanup, so it does not pass its fake stream handle to the real runtime.

## Validation

Base: `9de04cd44e927a9a3df6bbc0e94fc1be8b844ba4`; candidate: `18aa617c3fd5ea678d7fc0313deefa1718969f17`. The same patch was also replayed and host-checked on current main `011c9e600a93bbc9fd451c563501b5831c32bbca` on 2026-10-21; its stable patch-id is identical. The original base is kept here because pushing the newer upstream history to my fork would require permission to update GitHub workflows, which this fix does not need. Please keep this PR as a draft while the missing backend checks are completed.

**Passed:**

- Additional full **HIP wave32/gfx1100** engine and conversation-target build on the exact submitted commit `18aa617c` (2026-10-21), using the same pinned ROCm 7.14 / AMD clang 23 toolchain and ggml revision. Configure flags match the HIP recipe below except for `-DSTRATA_ENABLE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1100`, replacing `STRATA_HIP_GFX906`. The seven selected host CTests all passed (`100% tests passed out of 7`, 5.66 seconds), including the real-header ownership regression. Build container had no GPU devices; this does not validate execution on an RDNA GPU.

- Full engine and relevant test-target builds on the exact submitted commit `18aa617c` (2026-10-21): HIP gfx906 with ROCm 7.14 / AMD clang 23, and CUDA sm_86 with NVIDIA CUDA 13.4.92 / GCC 13. Both used Release, `STRATA_PORTABLE=ON`, `STRATA_BUILD_TESTS=OFF`, and `STRATA_BUILD_CONVERSATION_TESTS=ON`, with ggml pinned to `3cf03257f219afbe7334045ff7c6a06ac68c627d`. No GPU execution was performed by these checks.
- Integrated CTest: HIP gfx906 **7/7 passed** (cache, split failure, memory, file, save reclaim, copy lifetime, validation); CUDA sm_86 **8/8 passed**, adding the wrapped transfer test. GPU snapshot test targets also compiled, but were not executed. These are selected host tests, not a claim that the full project test suite or GPU tests passed.

Full-build command shape (source/output paths normalized):

```sh
cmake -S "$SRC" -B "$BUILD" -G Ninja \
  -DSTRATA_PORTABLE=ON -DSTRATA_BUILD_TESTS=OFF \
  -DSTRATA_BUILD_CONVERSATION_TESTS=ON -DCMAKE_BUILD_TYPE=Release \
  -DSTRATA_GGML_DIR="$PINNED_LLAMA" $BACKEND_FLAGS
# HIP:  -DSTRATA_HIP_GFX906=ON -DCMAKE_HIP_ARCHITECTURES=gfx906
#       -DCMAKE_C_COMPILER=/opt/rocm/llvm/bin/clang
#       -DCMAKE_CXX_COMPILER=/opt/rocm/llvm/bin/clang++
# CUDA: -DSTRATA_ENABLE_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=86
#       -DCMAKE_CUDA_COMPILER="$CUDA_ROOT/bin/nvcc" -DCUDAToolkit_ROOT="$CUDA_ROOT"
cmake --build "$BUILD" -j2 --target strata conversation_cache_test \
  conversation_split_failure_test conversation_memory_test conversation_file_test \
  session_save_reclaim_test conversation_snapshot_test conversation_validation_test \
  conversation_copy_lifetime_test
# CUDA additionally built conversation_transfer_test.
ctest --test-dir "$BUILD" --output-on-failure \
  -R 'conversation_(cache|split_failure|memory|file|validation|copy_lifetime|transfer)_test|session_save_reclaim_test'
```

On HIP the CTest selector omitted `transfer`; its link wrapper is CUDA-specific. The selected runs reported `100% tests passed out of 7` (HIP) and `100% tests passed out of 8` (CUDA).

- The new regression compiled and ran without GPU devices using the HIP SDK headers and both the wave32 and gfx906 compatibility include paths. The same test against the unmodified base fails on its first joined worker (`FAIL: exited worker retains a stream`); this change passes worker-exit, reuse, concurrent-device and failure-cleanup checks.
- `conversation_snapshot.cpp` and `conversation_state.cpp` compiled as host objects against both HIP compatibility include paths, using ROCm Clang, C++20, `-O2 -UNDEBUG -pthread -D__HIP_PLATFORM_AMD__=1` (plus `-DSTRATA_HIP_GFX906=1` for that path). This is compilation of the affected host units, not a full engine build.
- `git diff --check`.
- Separate, earlier deployment evidence: the same stream-ownership change, backported into the older `0.1.44.1` helper on Linux 6.8 with four Radeon VII/gfx906 cards, passed a production 35,235-token prompt and a repeat using 35,228 cached tokens. A passive probe during real prefill recorded 8 calls to the owner destructor and 8 successful `hipStreamDestroy` calls from it. That deployment carries other local patches; it is not GPU execution evidence for this exact main-based PR and not a performance comparison.

- Additional passive production follow-up of that older backport after about 10.5 hours: 1,039,110 generated tokens, no restarts/OOM, and sampled shared memory flat at 1,871,532,032 bytes for about 9 hours 40 minutes after warmup (5-minute samples). The 44 GiB memory limit was still reached during the run; this is evidence about stream-memory growth, not a claim that all memory pressure or slow-prefill causes are fixed.

Host-check commands, with checkout/output paths normalized (executed in a pinned ROCm toolchain container without GPU devices):

```sh
CXX=/opt/rocm/llvm/bin/clang++
flags="-std=c++20 -O2 -UNDEBUG -pthread -D__HIP_PLATFORM_AMD__=1 -Iinclude -isystem include/strata/platform/hip_compat -isystem /opt/rocm/include -isystem /opt/rocm/core-7.14/include -DSTRATA_HIP_GFX906=1"
$CXX $flags src/core/conversation_copy_lifetime_test.cpp -o "$OUT/lifetime-test"
"$OUT/lifetime-test"
$CXX $flags -c src/core/conversation_snapshot.cpp -o "$OUT/conversation_snapshot.o"
$CXX $flags -c src/core/conversation_state.cpp -o "$OUT/conversation_state.o"
```

The second pass used `include/strata/hip_compat` instead of `include/strata/platform/hip_compat` and omitted `-DSTRATA_HIP_GFX906=1`. The baseline test used the same test source next to the unmodified base's `conversation_copy.hpp` and was required to exit 1; the patched test exited 0 on both passes.

**Failed:**

- The intentionally unpatched regression fails as described above.
- During initial test development, the gfx906 shim lacked the test's selected CUDA error/type aliases. The fixture now uses a typed error value and the type of `cudaMemcpyDefault`; both header-path checks pass.
- A separate early standalone GPU copy fixture during the older-version rollout failed before model startup. It did not reproduce the production environment (memlock and mapped runtime-library identity were not established); its cause remains unresolved. It is not counted as validation. The later model/cache and passive destructor checks above are independent evidence.

**Not run:**

- CUDA GPU execution and Metal builds/hardware validation. This Linux environment has no Metal build environment; targeted help with that build is welcome. SYCL has its own converted conversation sources and does not include this helper; its build was not run. No backend hardware-correctness claim is made from the host tests.
- Model-backed execution and long-duration memory/latency tests on this exact main-based branch. Production evidence above is from the older backport only.

## Limits

This fixes private-stream ownership only. It does not change cache eviction, slot placement, checkpoint retention, copy ordering, or configured memory limits. The separate thread-local event vector in `order_after_queued()` is unchanged. Cleanup is best effort when the runtime/device is already unavailable; failed device selection skips destruction of that device's handles. No claim is made that all observed shared memory was leaked stream storage, or that all causes of slow prefill are fixed.

Related on this site

Quick links to related pages.