Fix concurrent SIGSEGV in DLRM forward() and feature extractors (#739)

Summary:
Pull Request resolved: #739

Two thread-safety bugs were causing leaf SIGSEGVs whenever requested QPS exceeded server saturation. The "collapse to 0.5 QPS" observed at over-saturation in t27/t28 was actually the leaf crashing, then the next driver phase achieving only the trickle of completions that landed before the crash.

t27 forensic analysis (2026-05-28) found crashes at BGM/Grace q=100+/inst (100% of iters), CPL q=50+/inst (17-45% of iters depending on QPS). Two distinct crash sites observed:

(A) `torch::jit::InterpreterStateImpl::runTemplate` -> `dlrmInferenceServerSide`. The existing per-clone LIFO checkout (`get_avail_thread_id` / `put_avail_thread_id`) was supposed to prevent concurrent `forward()` calls on the same Module clone (added in D105659812 / t16). It's likely that `torch::jit::Module::clone()` does not fully isolate JIT interpreter state — under high concurrency the leaf still crashed despite the LIFO. Added a per-`ThreadState` `std::mutex forward_mutex` as a safety net. Uncontended when the LIFO is working; serializes calls per clone if it's not, preventing the segfault.

(B) `dcperf::feature_extractors::generated::vc_NNNN_NNNN` -> `std::vector<IdScorePair>::_M_realloc_insert` -> `je_large_dalloc`. The generated extractor code calls `c->example->idScoreLists[i].emplace_back(...)` and `c->structData[i] += ...`. `c->example` pointed at a single shared `flat_example_` member of `FeatureExtractorSuite`, and `c->structData` pointed at a single shared `flat_struct_data_` heap buffer. `LeafNodeRank::DLRMRequestHandler` dispatches `runFeatureExtraction(*this_thread_ptr, ...)` onto the multi-threaded folly GlobalCPUThread singleton, so multiple in-flight requests sharing the same `ThreadData` can hit `runFlatExtractors` concurrently on different worker threads, racing on those shared mutable members. Fixed by making both per-call (`thread_local MockFeatureExample local_example` + `thread_local std::vector<float> local_struct`), with `local_struct` snapshotted from the read-only `flat_struct_data_` template each call.

Also converted `flat_pos_` from `size_t` to `std::atomic<size_t>` and reserve `count` slots in a single `fetch_add` so concurrent calls don't race the cursor.

Read-only state (`flat_tables_`, `flat_hash_tables_`, `flat_features_`, `flat_copies_`) stays shared — extractor code only does `.find()` on tables and never mutates the others.

Reviewed By: YifanYuan3

Differential Revision: D106731815

Author: excelle08

packages/feedsim/third_party/src/workloads/ranking/dwarfs/dlrm.cpp

@@ -37,12 +37,23 @@ struct DLRM::Impl {
   // worker thread owns a deep-cloned Module via model_copy, so
   // concurrent forward() calls touch disjoint state. See t16
   // SIGSEGV crash (2026-05-18) for the original failure.
+  //
+  // forward_mutex (t29, 2026-05-28): the LIFO checkout in
+  // get_avail_thread_id() / put_avail_thread_id() is supposed to
+  // prevent two callers from holding the same thread_id at once, but
+  // the leaf still SIGSEGVs inside torch::jit::InterpreterStateImpl::
+  // runTemplate at over-saturation (t27 q=160+/inst on BGM/Grace,
+  // q=50+ on CPL). Likely cause: torch::jit::Module::clone() does NOT
+  // fully isolate all interpreter state. Per-clone mutex is a defensive
+  // safety net — uncontended when the LIFO is doing its job, prevents
+  // the segfault if it's not.
   struct ThreadState {
     std::mt19937 rng;
     std::normal_distribution<float> dense_dist{0.0f, 1.0f};
     std::vector<float> dense_buffer;
     std::vector<int64_t> sparse_buffer;
     std::unique_ptr<torch::jit::script::Module> model_copy;
+    std::mutex forward_mutex;
   };
   std::vector<std::unique_ptr<ThreadState>> thread_states;
 
@@ -248,14 +259,17 @@ int DLRM::infer(int num_inferences, int batch_size) {
         params_.embedding_table_sizes);
 
     // Run inference on the per-thread Module clone (not the shared
-    // pimpl_->model) to avoid concurrent forward() race.
+    // pimpl_->model) to avoid concurrent forward() race. forward_mutex
+    // serializes calls per clone as a safety net on top of the LIFO
+    // checkout — see ThreadState comment.
     std::vector<torch::jit::IValue> inputs;
     inputs.push_back(dense_tensor);
     inputs.push_back(sparse_tensor);
 
     torch::NoGradGuard no_grad;
-    auto& thread_model = *pimpl_->thread_states[thread_id]->model_copy;
-    auto output = thread_model.forward(inputs).toTensor();
+    auto& ts = *pimpl_->thread_states[thread_id];
+    std::lock_guard<std::mutex> forward_lock(ts.forward_mutex);
+    auto output = ts.model_copy->forward(inputs).toTensor();
 
     // Count predictions (simulating actual work with the output)
     total_predictions += output.numel();
@@ -299,14 +313,16 @@ int DLRM::inferWithFeatures(
         {static_cast<int64_t>(batch_size), params_.num_sparse_features},
         torch::kInt64);
 
-    // Run inference on this thread's Module clone.
+    // Run inference on this thread's Module clone. forward_mutex
+    // serializes calls per clone — see ThreadState comment.
     std::vector<torch::jit::IValue> inputs;
     inputs.push_back(dense_tensor);
     inputs.push_back(sparse_tensor);
 
     torch::NoGradGuard no_grad;
-    auto& thread_model = *pimpl_->thread_states[thread_id]->model_copy;
-    auto output = thread_model.forward(inputs).toTensor();
+    auto& ts = *pimpl_->thread_states[thread_id];
+    std::lock_guard<std::mutex> forward_lock(ts.forward_mutex);
+    auto output = ts.model_copy->forward(inputs).toTensor();
 
     // Count predictions (simulating actual work with the output)
     total_predictions += output.numel();

packages/feedsim/third_party/src/workloads/ranking/feature_extractors/FeatureExtractorSuite.cpp

@@ -135,12 +135,44 @@ void FeatureExtractorSuite::runFlatExtractors(
   using namespace dcperf::feature_extractors::generated;
   if (flat_copies_.empty()) return;
 
-  flat_example_.resize(100, 50, 0);
+  // Per-call mutable state. CopyContext fields the generated extractors
+  // mutate (example via emplace_back, structData via in-place
+  // arithmetic) MUST NOT alias across concurrent invocations on the
+  // same suite — t27 forensic analysis showed SIGSEGV in
+  // vc_NNNN_NNNN -> std::vector<IdScorePair>::_M_realloc_insert ->
+  // je_large_dalloc, caused by concurrent emplace_back on the shared
+  // flat_example_.idScoreLists[i]. thread_local keeps allocator
+  // pressure low: each worker thread reuses its own buffers across
+  // calls; only initialized on first call.
+  thread_local MockFeatureExample local_example;
+  local_example.resize(100, 50, 0);
+
+  // structData is mutated by the generated code even without
+  // story_content seeding (e.g. `c->structData[off] += 1e-15f` in
+  // archetype templates), so it also needs to be per-call. We snapshot
+  // from flat_struct_data_ (the read-only template populated in
+  // initializeFlatDispatch) into a thread_local buffer once per call.
+  thread_local std::vector<float> local_struct;
+  if (local_struct.size() != static_cast<size_t>(flat_struct_size_)) {
+    local_struct.assign(flat_struct_data_.get(),
+                        flat_struct_data_.get() + flat_struct_size_);
+  } else {
+    std::copy(flat_struct_data_.get(),
+              flat_struct_data_.get() + flat_struct_size_,
+              local_struct.begin());
+  }
+  if (story_content && story_content_length > 0) {
+    int len = std::min(story_content_length, flat_struct_size_);
+    for (int i = 0; i < len; ++i) {
+      local_struct[i] = static_cast<float>(story_content[i]) / 255.0f;
+    }
+  }
+
   CopyContext ctx;
   ctx.tables = flat_tables_;
-  ctx.structData = flat_struct_data_.get();
-  ctx.structSize = flat_struct_size_;
-  ctx.example = &flat_example_;
+  ctx.structData = local_struct.data();
+  ctx.structSize = static_cast<int>(local_struct.size());
+  ctx.example = &local_example;
   ctx.features = flat_features_.data();
   ctx.numFeatures = static_cast<int>(flat_features_.size());
   ctx.queryKeys = input_sparse.data();
@@ -150,16 +182,16 @@ void FeatureExtractorSuite::runFlatExtractors(
   ctx.storyContent = story_content;
   ctx.storyContentLength = story_content_length;
 
-  // If story content is available, seed structData from it
-  if (story_content && story_content_length > 0) {
-    int len = std::min(story_content_length, flat_struct_size_);
-    for (int i = 0; i < len; ++i) {
-      flat_struct_data_[i] = static_cast<float>(story_content[i]) / 255.0f;
-    }
-  }
-
+  // Atomic cursor advance so concurrent calls don't race on
+  // flat_pos_. Reserves `count` slots in one fetch_add; each call
+  // therefore runs N consecutive copy functions from the shuffled
+  // vector, just starting at a different offset depending on
+  // interleaving. Distribution remains uniform.
+  const size_t total = flat_copies_.size();
+  size_t start = flat_pos_.fetch_add(static_cast<size_t>(count),
+                                     std::memory_order_relaxed) %
+      total;
   for (int i = 0; i < count; ++i) {
-    flat_copies_[flat_pos_](&ctx);
-    flat_pos_ = (flat_pos_ + 1) % flat_copies_.size();
+    flat_copies_[(start + static_cast<size_t>(i)) % total](&ctx);
   }
 }

packages/feedsim/third_party/src/workloads/ranking/feature_extractors/FeatureExtractorSuite.h

@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <random>
@@ -66,13 +67,27 @@ class FeatureExtractorSuite {
   std::vector<int64_t> sparse_buf_a_;
   std::vector<int64_t> sparse_buf_b_;
 
-  // Flat dispatch state
+  // Flat dispatch state. Concurrency contract: a single
+  // FeatureExtractorSuite instance can have runFlatExtractors() called
+  // from multiple threads at once (LeafNodeRank dispatches feature
+  // extraction onto the multi-threaded GlobalCPUThread pool, and
+  // multiple in-flight requests on the same ThreadData share its
+  // suite). The fields below are read-only after initializeFlatDispatch
+  // — the generated extractor code only does `.find()` on tables and
+  // never mutates flat_features_ / flat_copies_ / flat_hash_tables_,
+  // so sharing is safe.
+  //
+  // The previously-member `flat_example_` and shared writes to
+  // `flat_struct_data_` were removed: extractor code calls
+  // `c->example->idScoreLists[i].emplace_back(...)` and `c->structData[i]
+  // += ...`, both racy when shared. runFlatExtractors now uses
+  // thread_local buffers for those, with the per-call copy from
+  // flat_struct_data_ as the read-only template.
   std::vector<dcperf::feature_extractors::generated::CopyFn> flat_copies_;
-  size_t flat_pos_ = 0;
+  std::atomic<size_t> flat_pos_{0};
   std::unique_ptr<float[]> flat_struct_data_;
   int flat_struct_size_ = 0;
   std::unordered_map<int64_t, float> flat_tables_[4];
   dcperf::mock_hash::MockHashTable flat_hash_tables_[4];
-  MockFeatureExample flat_example_;
   std::vector<MockFeature> flat_features_;
 };