Merge pull request #694 from Aharrypotter/perf/gdn-lane8-product

perf(cpu): raise GDN recurrence lane cap 4 -> 8

Author: Aharrypotter

mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp

@@ -23,10 +23,11 @@ namespace mllm::cpu::gdn {
 namespace {
 
 constexpr int kMaxStackNormalizedHeadDim = 256;
-// GDN state updates are bandwidth-heavy on heterogeneous mobile CPUs. A small
-// fixed lane cap avoids making every efficiency core part of the per-layer
-// completion barrier.
-constexpr int kMaxParallelGDNLanes = 4;
+// GDN state updates are bandwidth-heavy on heterogeneous mobile CPUs. The lane
+// cap bounds how many tasks share the per-layer completion barrier. 8 lanes
+// matches the 8-core phones this kernel targets (4B has 32 recurrence tasks,
+// so all 8 cores participate); bitwise-safe because tasks are disjoint.
+constexpr int kMaxParallelGDNLanes = 8;
 // Scalar state elements updated across all [batch, value_head] tasks.
 constexpr std::size_t kMinParallelGDNWork = 65536;
 

mllm/engine/HpcThreadPool.cpp

@@ -50,7 +50,7 @@ void HpcThreadPool::wakeup() {
 
 int HpcThreadPool::acquireTaskSlot() {
   std::lock_guard<std::mutex> _l(queue_mutex_);
-  for (int i = 0; i < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++i) {
+  for (int i = 0; i < kHpcThreadPoolTaskLimit; ++i) {
     if (task_available_[i]) {
       task_available_[i] = false;
       return i;
@@ -60,7 +60,7 @@ int HpcThreadPool::acquireTaskSlot() {
 }
 
 void HpcThreadPool::releaseTaskSlot(int task_slot_idx) {
-  if (task_slot_idx < 0 || task_slot_idx >= MLLM_HPC_THREAD_POOL_TASK_LIMITS) { return; }
+  if (task_slot_idx < 0 || task_slot_idx >= kHpcThreadPoolTaskLimit) { return; }
   std::lock_guard<std::mutex> _l(queue_mutex_);
   task_available_[task_slot_idx] = true;
 }
@@ -95,9 +95,14 @@ void HpcThreadPool::splitTask(HpcThreadPoolTask&& task, int task_slot_idx) {
   // e.g.: threads is 4, tiles_name is 12.
   //  0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
   if (tiles_num > thread_cnt_) {
+    // Capture task and true_idx BY VALUE: splitTask's parameters (a moved-in
+    // rvalue task and a stack vector) die when this function returns, but the
+    // worker threads execute this lambda asynchronously — capturing by
+    // reference leaves a dangling reference (use-after-free). This was latent
+    // with 4 lanes and became reachable at 8 lanes.
     tasks_[task_slot_idx].first = {
         .func =
-            [tiles_num, &task, &true_idx, this](int thread_idx) {
+            [tiles_num, task, true_idx, this](int thread_idx) {
               for (int v = thread_idx; v < tiles_num; v += thread_cnt_) { task.func(true_idx[v]); }
             },
         .start = 0,
@@ -107,7 +112,7 @@ void HpcThreadPool::splitTask(HpcThreadPoolTask&& task, int task_slot_idx) {
     tiles_num = thread_cnt_;
   } else {
     tasks_[task_slot_idx].first = {
-        .func = [tiles_num, &task, &true_idx, this](int thread_idx) { task.func(true_idx[thread_idx]); },
+        .func = [tiles_num, task, true_idx, this](int thread_idx) { task.func(true_idx[thread_idx]); },
         .start = 0,
         .end = tiles_num,
         .step = 1,
@@ -140,11 +145,11 @@ HpcThreadPool::HpcThreadPool(int thread_cnt) {
   thread_cnt_ = thread_cnt;
   available_task_slots_ = 0;
   available_task_slots_old_ = 0;
-  task_available_.resize(MLLM_HPC_THREAD_POOL_TASK_LIMITS);
-  tasks_.resize(MLLM_HPC_THREAD_POOL_TASK_LIMITS);
+  task_available_.resize(kHpcThreadPoolTaskLimit);
+  tasks_.resize(kHpcThreadPoolTaskLimit);
 
   // Each task should hold some thread ok flag that mark this thread's work is done.
-  for (int t = 0; t < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++t) {
+  for (int t = 0; t < kHpcThreadPoolTaskLimit; ++t) {
     task_available_[t] = true;
     for (int i = 0; i < thread_cnt_; ++i) { tasks_[t].second.emplace_back(new std::atomic_bool{false}); }
   }
@@ -154,7 +159,7 @@ HpcThreadPool::HpcThreadPool(int thread_cnt) {
     workers_.emplace_back([this, thread_idx]() {
       while (!stop_) {
         while (available_task_slots_ > 0) {
-          for (int i = 0; i < MLLM_HPC_THREAD_POOL_TASK_LIMITS; ++i) {
+          for (int i = 0; i < kHpcThreadPoolTaskLimit; ++i) {
             if (*tasks_[i].second[thread_idx]) {
               tasks_[i].first.func(thread_idx);
               { *tasks_[i].second[thread_idx] = false; }

mllm/engine/HpcThreadPool.hpp

@@ -13,10 +13,13 @@
 #include <functional>
 #include <condition_variable>
 
-#define MLLM_HPC_THREAD_POOL_TASK_LIMITS 2
-
 namespace mllm {
 
+// One task slot per expected concurrent op. 8 covers the GDN recurrence
+// 8-lane cap plus sibling parallel ops; too small a limit makes splitTask fall
+// back to serial on the main thread under multi-layer concurrency.
+inline constexpr int kHpcThreadPoolTaskLimit = 8;
+
 struct HpcThreadPoolTask {
   std::function<void(int)> func;
   int start = 0;

tests/cpu/Qwen35GDNTest.cpp

@@ -18,8 +18,11 @@ using mllm::cpu::gdn::gatedDeltaRuleF32;
 class ScopedCpuOpThreads {
  public:
   explicit ScopedCpuOpThreads(int32_t thread_count) : original_thread_count_(mllm::Context::instance().getCpuOpThreads()) {
+    // initializeContext() registers the CPU backend; SymbolTable::reg aborts on
+    // a duplicate key, so call it exactly once (the tests have no fixture init).
+    static const bool kContextInitialized = [] { mllm::initializeContext(); return true; }();
+    (void)kContextInitialized;
     mllm::Context::instance().setCpuOpThreads(thread_count);
-    mllm::initializeContext();
   }
 
   ~ScopedCpuOpThreads() { mllm::Context::instance().setCpuOpThreads(original_thread_count_); }
@@ -218,7 +221,10 @@ TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) {
   constexpr int kValueHeads = 32;
   constexpr int kKeyDim = 128;
   constexpr int kValueDim = 128;
-  constexpr int kThreadCount = 4;
+  // Exercises the parallel lane partition up to the 8-lane cap
+  // (kMaxParallelGDNLanes); tasks are disjoint so output must be bitwise
+  // identical regardless of how many lanes the scheduler picks.
+  constexpr int kThreadCount = 8;
 
   std::vector<float> q(kBatch * kSequence * kKeyHeads * kKeyDim);
   std::vector<float> k(q.size());
@@ -266,4 +272,78 @@ TEST(Qwen35GDNTest, ParallelBatchValueHeadsMatchSerialBitwise) {
   }
 }
 
+// 4B real geometry (B=1, S=69, 16 key heads, 32 value heads, 128 dims) at the
+// 8-lane cap — exercises the full task fan-out (32 tasks) that the small
+// geometry above does not. Guards against the device crash observed on
+// OnePlus with the 8-lane product build.
+TEST(Qwen35GDNTest, FourBGeometry8LaneDoesNotCrash) {
+  constexpr int kBatch = 1;
+  constexpr int kSequence = 69;
+  constexpr int kKeyHeads = 16;
+  constexpr int kValueHeads = 32;
+  constexpr int kKeyDim = 128;
+  constexpr int kValueDim = 128;
+  constexpr int kThreadCount = 8;
+
+  std::vector<float> q(kBatch * kSequence * kKeyHeads * kKeyDim);
+  std::vector<float> k(q.size());
+  std::vector<float> v(kBatch * kSequence * kValueHeads * kValueDim);
+  std::vector<float> a(kBatch * kSequence * kValueHeads);
+  std::vector<float> b(a.size());
+  std::vector<float> a_log(kValueHeads);
+  std::vector<float> dt_bias(kValueHeads);
+
+  for (std::size_t i = 0; i < q.size(); ++i) {
+    q[i] = 0.01F * static_cast<float>(static_cast<int>(i % 7) - 3);
+    k[i] = 0.01F * static_cast<float>(static_cast<int>(i % 5) - 2);
+  }
+  for (std::size_t i = 0; i < v.size(); ++i) { v[i] = 0.01F * static_cast<float>(static_cast<int>(i % 11) - 5); }
+  for (std::size_t i = 0; i < a.size(); ++i) {
+    a[i] = 0.001F * static_cast<float>(static_cast<int>(i % 3));
+    b[i] = 0.001F * static_cast<float>(static_cast<int>(i % 9));
+  }
+  for (int i = 0; i < kValueHeads; ++i) { a_log[i] = -1.0F; dt_bias[i] = 0.0F; }
+
+  std::vector<float> state(kBatch * kValueHeads * kValueDim * kKeyDim, 0.0F);
+  std::vector<float> output(v.size());
+  std::vector<float> ref_output(v.size());
+  std::vector<float> ref_state = state;
+
+  // Serial reference, then 8-lane parallel — must be bitwise identical.
+  gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), ref_state.data(),
+                    ref_output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim,
+                    /*thread_count=*/1);
+  const ScopedCpuOpThreads scoped_threads(kThreadCount);
+  gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), state.data(),
+                    output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount);
+
+  for (std::size_t i = 0; i < output.size(); ++i) {
+    ASSERT_EQ(ref_output[i], output[i]) << "output index " << i;
+  }
+  for (std::size_t i = 0; i < state.size(); ++i) {
+    ASSERT_EQ(ref_state[i], state[i]) << "state index " << i;
+  }
+
+  // Repeat the full 4B GDN pass 24 times (one per layer) to mimic the real
+  // model's layer loop, which interleaves the recurrence with other parallel
+  // ops on the shared thread pool. Context init is now once-only (see
+  // ScopedCpuOpThreads), so this exercises multi-call thread-pool reuse.
+  // Run the recurrence 24 times on a FRESH copy of the initial state each
+  // time (mirroring one GDN layer per model layer from the same prefill input),
+  // and compare each run's output to the serial reference for that same input.
+  // This exercises repeated thread-pool push/acquire/release cycles — the
+  // multi-call reuse pattern that crashed on device.
+  for (int layer = 0; layer < 24; ++layer) {
+    std::vector<float> layer_state(state.size(), 0.0F);
+    std::vector<float> layer_ref_state(state.size(), 0.0F);
+    gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), layer_ref_state.data(),
+                      ref_output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, /*thread_count=*/1);
+    gatedDeltaRuleF32(q.data(), k.data(), v.data(), a.data(), b.data(), a_log.data(), dt_bias.data(), layer_state.data(),
+                      output.data(), kBatch, kSequence, kKeyHeads, kValueHeads, kKeyDim, kValueDim, kThreadCount);
+    for (std::size_t i = 0; i < output.size(); ++i) {
+      ASSERT_EQ(ref_output[i], output[i]) << "layer " << layer << " output index " << i;
+    }
+  }
+}
+
 }  // namespace