Ds4HcMatvecPool: heap-use-after-free in dual-GPU MoE-TP decode path (fix included), plus a residual rare deadlock (unresolved)

[chrisjohnson]

Ds4HcMatvecPool: heap-use-after-free in dual-GPU MoE-TP decode path (fix included), plus a residual rare deadlock (unresolved)

Follow-up to #681 — different bug, same overall dual-GPU (R9700 Thunderbolt eGPU + Strix
Halo) DeepSeek V4 Flash setup, found while stress-testing after applying the ROCm 7.2.2
fix from that issue. This one is in dflash_server's own code, not the HIP runtime.

Environment: HIP_VISIBLE_DEVICES=1,0, DFLASH_DS4_MOE_TP_INPROC=1, --peer-access,
ROCm 7.2.2, checkout based on f0c5a5d (2026-08-26).

Part 1 — heap-use-after-free in Ds4HcMatvecPool (fixed, patch below)

Symptom: with the #681 crash fixed, a single request per server session worked
reliably, but a later request in the same session would eventually crash or hang —
not on any fixed request number, roughly 1-in-5 to 1-in-20 over a plain (non-ASan)
build. Symptoms varied across runs (std::length_error, an indefinite hang, plain
SIGSEGV), consistent with memory corruption rather than one deterministic bug.

Root cause (confirmed via a -fsanitize=address build, amdclang++/ROCm's own
clang toolchain — mixing a GCC host build with ASan against HIP translation units built
under hipcc fails outright with "incompatible ASan runtimes" even though ldd shows
both resolving to the same libasan.so.6, so the whole build has to go through one
compiler family):

hc_pre_batch(...) (server/src/deepseek4/deepseek4_graph.cpp:3918) allocates a
per-token-chunk std::vector<float> inside a closure dispatched to a persistent
worker-thread pool. It calls into Ds4HcMatvecPool::run()
(deepseek4_graph.cpp ~3521-3644), which further parallelizes the matvec's row range
across 8 dedicated worker threads and is documented as blocking until every dispatched
worker has decremented a shared completion counter. ASan caught a real,
address-symbolized heap-use-after-free: the caller freed the buffer immediately after
run() returned, while a pool worker thread it had dispatched was still reading from it
inside cpu_dot_f16_rows3_f16c — i.e. run() returned (and let its caller proceed to
free the output buffer) before every dispatched worker had actually finished.

The mechanism: each worker tracks "have I already processed the current round" via a
last-seen value assigned during its wake-up/busy-spin check, captured before it
reads the shared job struct. A worker preempted by the OS scheduler between those two
reads can resume long after several further rounds have already been dispatched and
processed by other workers, leaving its last value permanently stale relative to the
real sequence counter. On every subsequent loop iteration such a worker then treats the
current round as new (since a stale last essentially never matches the live
sequence again), reprocesses it, and issues an extra decrement against that round's
completion counter — letting the caller observe "done" and free its buffer one or more
decrements early while a real worker for that round is still writing into it.

Two earlier fix attempts (both superseded, included in the branch history for
reference) each closed part of the problem but not all of it:

• Giving each run() call its own completion counter instead of the single shared
  remaining member closed the cross-round decrement bleed, but the natural
  implementation (a stack-local std::atomic<int>) introduced a new, ASan-caught
  stack-use-after-return — a straggling worker could still write to the counter after
  the call's stack frame had been reused.
• Moving to a small ring of pool-owned (member-array) counters, round-robined per
  dispatch, fixed the stack-lifetime issue but still hit the same underlying
  use-after-free at a lower rate (26/30 clean requests before reproducing) — the ring
  reduces collision probability but doesn't address the root desync described above.
  (Widening the ring 8→4096 slots as a probe made things worse, 1/40 — informative
  negative result ruling out ring size/wraparound as the actual mechanism.)

The fix: derive each worker's "last seen round" from a round field carried inside
Job itself (read via the same struct copy as everything else, computed under
client_mu right before publishing) instead of a value captured separately before the
job read. This keeps "last seen" always exactly synchronized with the round the
worker is actually acting on. Full diff against f0c5a5d:

--- a/server/src/deepseek4/deepseek4_graph.cpp
+++ b/server/src/deepseek4/deepseek4_graph.cpp
@@ -3531,12 +3531,42 @@ struct Ds4HcMatvecPool {
         int rows;
         int cols;
         int active_workers;
+        // Index into done_slots (below) for this dispatch's completion
+        // counter. A single shared, reused counter let a worker that was
+        // still finishing a previous round decrement the NEXT round's
+        // counter after the caller had already observed completion and
+        // freed its output buffer - confirmed via AddressSanitizer
+        // heap-use-after-free. Round-robining across a small ring of
+        // pool-owned (not stack-local - a stack-local counter still let a
+        // straggling worker touch it after the owning call returned, per a
+        // second AddressSanitizer stack-use-after-return) counters gives
+        // each dispatch its own slot without that lifetime risk: a late
+        // decrement from an old round can only ever land on an old, still
+        // valid slot, never the current round's.
+        int done_slot;
+        // The seq value this Job was published under. Workers track "have I
+        // already processed the current round" via their own last-seen-seq
+        // variable - but that variable was previously assigned from a value
+        // captured BEFORE reading job (from the wake-up/busy-spin check),
+        // not from job itself. A worker preempted between those two reads
+        // could resume long after several further rounds had already been
+        // dispatched and processed by other workers; it would then record a
+        // stale, out-of-date "last seen" value, causing it to treat every
+        // subsequent round as new indefinitely and keep re-decrementing
+        // done_slots for rounds it had no real assignment in - exactly the
+        // premature-completion pattern behind the observed use-after-free.
+        // Deriving "last seen" from this field instead (read in the same
+        // struct copy as everything else) keeps it exactly synchronized
+        // with the round actually being acted on.
+        uint64_t round;
     };
+    static constexpr int kDoneSlots = 8;
     std::mutex client_mu;
     std::mutex wait_mu;
     std::condition_variable wait_cv;
     std::atomic<uint64_t> seq{0};
-    std::atomic<int> remaining{0};
+    std::atomic<int> done_slots[kDoneSlots]{};
+    int next_done_slot = 0;
     Job job{};
     std::vector<std::thread> workers;
     std::atomic<bool> stop{false};
@@ -3572,8 +3602,8 @@ struct Ds4HcMatvecPool {
                         s = seq.load(std::memory_order_acquire);
                     }
                     if (stop.load(std::memory_order_relaxed)) return;
-                    last = s;
                     const Job j = job;
+                    last = j.round;
                     if (i >= j.active_workers) continue;
                     const int chunk = (j.rows + j.active_workers - 1) / j.active_workers;
                     const int r0 = i * chunk;
@@ -3598,7 +3628,7 @@ struct Ds4HcMatvecPool {
                             j.out[r] = cpu_dot_f16_row(j.mat + (size_t) r * j.cols, j.x, j.cols);
                         }
                     }
-                    remaining.fetch_sub(1, std::memory_order_acq_rel);
+                    done_slots[j.done_slot].fetch_sub(1, std::memory_order_acq_rel);
                 }
             });
         }
@@ -3619,8 +3649,14 @@ struct Ds4HcMatvecPool {
         std::lock_guard<std::mutex> lk(client_mu);
         row_fn = nullptr;
         const int active_workers = std::min(nth, rows);
-        job = {mat, x, out, rows, cols, active_workers};
-        remaining.store(active_workers, std::memory_order_release);
+        const int slot = next_done_slot;
+        next_done_slot = (next_done_slot + 1) % kDoneSlots;
+        std::atomic<int> & round_done = done_slots[slot];
+        round_done.store(active_workers, std::memory_order_release);
+        // Only this call (holding client_mu) can be advancing seq right
+        // now, so this predicts exactly what fetch_add below will produce.
+        const uint64_t new_round = seq.load(std::memory_order_relaxed) + 1;
+        job = {mat, x, out, rows, cols, active_workers, slot, new_round};
         {
             // Publish the new generation while holding wait_mu so a worker
             // cannot miss the transition between its predicate check and
@@ -3630,7 +3666,7 @@ struct Ds4HcMatvecPool {
         }
         wait_cv.notify_all();
         int spins = 0;
-        while (remaining.load(std::memory_order_acquire) != 0) {
+        while (round_done.load(std::memory_order_acquire) != 0) {
             if (++spins < 65536) { cpu_relax(); }
             else { std::this_thread::yield(); spins = 0; }
         }
@@ -3644,15 +3680,19 @@ struct Ds4HcMatvecPool {
         std::lock_guard<std::mutex> lk(client_mu);
         row_fn = std::move(fn);
         const int active_workers = std::min(nth, rows);
-        job = {nullptr, nullptr, nullptr, rows, 0, active_workers};
-        remaining.store(active_workers, std::memory_order_release);
+        const int slot = next_done_slot;
+        next_done_slot = (next_done_slot + 1) % kDoneSlots;
+        std::atomic<int> & round_done = done_slots[slot];
+        round_done.store(active_workers, std::memory_order_release);
+        const uint64_t new_round = seq.load(std::memory_order_relaxed) + 1;
+        job = {nullptr, nullptr, nullptr, rows, 0, active_workers, slot, new_round};
         {
             std::lock_guard<std::mutex> wake_lk(wait_mu);
             seq.fetch_add(1, std::memory_order_release);
         }
         wait_cv.notify_all();
         int spins = 0;
-        while (remaining.load(std::memory_order_acquire) != 0) {
+        while (round_done.load(std::memory_order_acquire) != 0) {
             if (++spins < 65536) { cpu_relax(); }
             else { std::this_thread::yield(); spins = 0; }
         }

Validation: rebuilt under the same ASan+amdclang++ configuration, ran real
/v1/chat/completions decode requests against the actual DS4 target model (not a
synthetic unit test) — 48/50 requests clean with zero ASan errors on the first batch
(the 2 "failures" were confirmed client-side 30s timeout artifacts, not crashes — server
kept processing normally). Continued testing since with the plain (non-ASan) build: 235+
consecutive real requests, 0 crashes, as of this writing. Every prior build (pre-fix)
crashed within 1-30 requests, every single test run across this investigation, so this
is a large, clearly measurable improvement — but see Part 2 below for a residual issue
found during extended testing.

Part 2 — residual rare deadlock (NOT fixed, root cause unknown)

During extended stress testing of the Part 1 fix, hit a genuine deadlock once (around
request 65-70 of a cumulative ~90 at the time; has not recurred since, across 235+
total real requests in the current session). Confirmed via a live gdb -p <pid>
attach: every one of Ds4HcMatvecPool's worker threads was idle in wait_cv.wait()
with nothing to do, while the calling thread was parked in run()'s spin-wait loop,
waiting on a completion counter that never reached zero. In other words: some worker
that should have decremented the round's counter did not, but by the time of the
attach every worker had already gone back to sleep — so this is not the Part 1
use-after-free (no memory-safety violation, and ASan reported nothing at the time of
the hang; a hang isn't something ASan detects).

I don't have a confirmed root cause for this one. It's genuinely unclear whether it's a
new symptom introduced by the Part 1 fix, or a separate, much rarer, pre-existing bug in
the same class that the far-more-frequent crash simply never had a chance to manifest
before (the crash always happened first). A standalone ~90-line repro harness that
exercises Ds4HcMatvecPool's exact threading pattern in isolation (extracted class body,
synthetic 24×16384 matvec calls at realistic dimensions, no model loading) has run
270M+ iterations (fast build) and 80M+ iterations (ASan build) without reproducing it,
which is itself informative — whatever triggers this needs either real GPU/HIP driver
timing or some interaction with the rest of dflash_server that the pure-CPU synthetic
repro doesn't replicate.

Flagging this as open/unresolved rather than holding the Part 1 fix back on it — the
fix is a clear, large improvement on its own, but this residual deadlock means the pool
still isn't provably 100% correct under all interleavings. Happy to share the repro
harness or dig further if it's useful to narrow down; wanted to get the confirmed fix +
known residual risk in front of you rather than sit on it.

[chrisjohnson]

For anyone who wants to pull this directly: the fix (commits 233c8b8..176c0ae on top of f0c5a5d, including the two superseded intermediate attempts for reference) is pushed as a branch here: https://github.com/chrisjohnson/lucebox/tree/m137-ds4-matvec-pool-uaf-fix