From 59cb9a9da66583b6dd7be14622d97d6fc7da710e Mon Sep 17 00:00:00 2001
From: kpham-sgl <khoa.pham@radixark.ai>
Date: Mon, 23 Mar 2026 00:37:07 -0700
Subject: [PATCH] [Spec][Ngram] 3/N: Fix synchronization issues in `Ngram.cpp`
 (#21186)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../srt/speculative/cpp_ngram/ngram.cpp       | 20 ++++++++++---------
 .../sglang/srt/speculative/cpp_ngram/ngram.h  |  9 ++++++++-
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
index 582c87659..b1d54b964 100644
--- a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
+++ b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
@@ -1,10 +1,8 @@
 #include "ngram.h"
 
-#include <chrono>
 #include <limits>
 #include <stdexcept>
 #include <string>
-#include <thread>
 
 #include "trie.h"
 
@@ -72,12 +70,10 @@ Ngram::Ngram(size_t capacity, const Param& param) : param_(param) {
 
   trie_ = std::make_unique<Trie>(capacity, param_);
 
-  quit_flag_ = false;
   insert_worker_ = std::thread(&Ngram::insertWorker, this);
 }
 
 Ngram::~Ngram() {
-  quit_flag_ = true;
   insert_queue_.close();
   if (insert_worker_.joinable()) {
     insert_worker_.join();
@@ -85,25 +81,31 @@ Ngram::~Ngram() {
 }
 
 void Ngram::synchronize() const {
-  while (!insert_queue_.empty()) {
-    std::this_thread::sleep_for(std::chrono::microseconds(10));
-  }
+  std::unique_lock<std::mutex> lock(mutex_);
+  sync_cv_.wait(lock, [this] { return pending_count_ == 0; });
 }
 
 void Ngram::asyncInsert(std::vector<std::vector<int32_t>>&& tokens) {
+  {
+    std::lock_guard<std::mutex> lock(mutex_);
+    pending_count_ += tokens.size();
+  }
   for (auto&& token : tokens) {
     insert_queue_.enqueue(std::move(token));
   }
 }
 
 void Ngram::insertWorker() {
-  while (!quit_flag_) {
+  for (;;) {
     std::vector<int32_t> data;
     if (!insert_queue_.dequeue(data)) {
-      continue;
+      break;
     }
     std::unique_lock<std::mutex> lock(mutex_);
     trie_->insert(data.data(), data.size());
+    --pending_count_;
+    lock.unlock();
+    sync_cv_.notify_all();
   }
 }
 
diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.h b/python/sglang/srt/speculative/cpp_ngram/ngram.h
index 9af65885b..377b481ae 100644
--- a/python/sglang/srt/speculative/cpp_ngram/ngram.h
+++ b/python/sglang/srt/speculative/cpp_ngram/ngram.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <condition_variable>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -18,8 +19,14 @@ class Ngram {
   std::unique_ptr<Trie> trie_;
   Param param_;
 
+  // NOTE: protects trie_ and pending_count_. Ensures batchMatch never reads
+  // trie_ while insertWorker is writing. After synchronize(), no pending
+  // inserts remain so mutex_ contention is effectively zero.
   mutable std::mutex mutex_;
-  bool quit_flag_ = false;
+  mutable std::condition_variable sync_cv_;
+  // NOTE: tracks inserts from enqueue through trie_->insert() completion,
+  // not just queue occupancy. A dequeued item may still be mid-insert.
+  size_t pending_count_ = 0;
   utils::Queue<std::vector<int32_t>> insert_queue_;
   std::thread insert_worker_;