python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (77-82)

Accessing staging_sam_ here is thread-unsafe. If clearExternalCorpus is called on the main thread while this background operation is running, staging_sam_ could be reset to nullptr between the check on line 78 and the dereference on line 81, leading to a crash.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (84-92)

Multiple unsynchronized accesses to staging_sam_ occur here before the lock is acquired on line 94. Specifically, finalize(), empty(), and reset() are called on staging_sam_ while the main thread might be concurrently resetting it in clearExternalCorpus.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (103-107)

Resetting staging_sam_ while holding mutex_ does not protect against crashes in appendExternalCorpusTokens or finishExternalCorpusLoad, as those methods access staging_sam_ without acquiring the same lock. To fix this, all accesses to staging_sam_ should be protected by the mutex, or the staging logic should be refactored to avoid shared state during the loading process.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (70-75)

The staging_sam_ member is accessed here without holding mutex_. While the comment suggests it is only accessed from a background thread, clearExternalCorpus (which runs on the main thread) also modifies staging_sam_. This creates a race condition if a load is started while a clear operation is in progress.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (215-219)

Iterating over std::unordered_map results in a non-deterministic order of SAM merging. If multiple SAMs provide different candidates for the same prefix, the final merged result may vary between runs depending on the iteration order. For consistent speculative decoding, consider using a deterministic container like std::map or sorting the keys before iteration.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp (204-210)

The condition per_sam_budget == 0 is overly restrictive. If total_sam_budget is less than num_sams, all SAMs are completely bypassed even if some budget is available. It would be more efficient to distribute the available budget to a subset of SAMs rather than falling back to trie-only matching.