Migrate ngram corpus from torch cpp_extension to TVM FFI jit_kernel

[hnyls2002]

Motivation

torch.utils.cpp_extension.load() JIT cache is unreliable in CI — stale .so artifacts cause runtime errors when C++ code changes. Example failure from PR #20208 CI run:

File "python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py", line 39, in __init__
    param.min_match_window_size = min_match_window_size
AttributeError: 'ngram_corpus_cpp.Param' object has no attribute 'min_match_window_size'

The cached .so from a previous build was loaded instead of recompiling, so the new min_match_window_size field didn't exist at runtime.

Summary

• Move ngram corpus C++ source files from speculative/cpp_ngram/ to jit_kernel/csrc/ngram_corpus/
• Replace pybind11 binding with TVM FFI function wrappers using an opaque handle pattern
• Eliminates torch.utils.cpp_extension.load() JIT cache issues (hash collisions, stale caches, multi-process lock contention) by using tvm_ffi.cpp.load_inline()
• No changes to C++ implementation logic — pure build system migration
• NgramCorpus Python API unchanged, all existing callers work without modification

Test plan

• test/registered/spec/utils/test_ngram_corpus.py (17 test classes)
• test/registered/spec/test_ngram_speculative_decoding.py (E2E)

[gemini-code-assist]

Code Review

This pull request refactors the NgramCorpus implementation to use a JIT-compiled C++ FFI, replacing the previous pybind11 and torch.utils.cpp_extension setup. The changes introduce a new FFI header for instance management and a Python wrapper to interface with the JIT-loaded module. Review feedback identifies a critical data race in the global instance map access, a lack of bounds checking in memory copies that could lead to buffer overflows, and performance overhead in the CSR data conversion logic.

[hnyls2002]

/tag-and-rerun-ci

[hnyls2002]

/rerun-test test_ngram_corpus.py test_ngram_speculative_decoding.py

[hnyls2002]

/gemini review

[github-actions]

✅ 1-gpu-5090: View workflow run

cd test/ && python3 registered/spec/utils/test_ngram_corpus.py

✅ 1-gpu-h100: View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py

[gemini-code-assist]

Code Review

This pull request migrates the NgramCorpus implementation from a standard PyTorch C++ extension to a JIT-compiled kernel using a custom FFI. The changes include the introduction of a new FFI header, a Python wrapper for the JIT module, and a refactoring of the existing NgramCorpus class to utilize these new components. Feedback focuses on critical issues regarding the use of static global state in headers which could lead to linkage errors, a lack of thread safety in instance lookups, and several performance bottlenecks in the inference hot path, such as inefficient Python-to-CSR conversion and unnecessary tensor allocations and type conversions.

[gemini-code-assist]

Defining global state with static linkage in an anonymous namespace within a header file is problematic. If this header is included in multiple translation units (TUs) during the JIT compilation process (e.g., once as a source file in cpp_files and once by the generated wrapper), each TU will have its own independent copy of g_instances and g_map_mutex. This will lead to 'Invalid ngram handle' errors when trying to access an instance created in a different TU.

Consider moving these definitions to a .cpp file or using inline variables (C++17) in a named namespace to ensure a single global instance across the shared library.

[gemini-code-assist]

The get_instance function is not thread-safe. It performs a lookup on g_instances without acquiring g_map_mutex. While create and destroy are protected, concurrent read/write access to std::unordered_map is undefined behavior and can lead to crashes or memory corruption. Since this is called during the inference hot path, it must be thread-safe.

Suggested change

	inline ngram::Ngram& get_instance(int64_t handle) {
	auto it = g_instances.find(handle);
	if (it == g_instances.end()) {
	throw std::runtime_error("Invalid ngram handle: " + std::to_string(handle));
	}
	return *it->second;
	}
	inline ngram::Ngram& get_instance(int64_t handle) {
	std::lock_guard<std::mutex> lock(g_map_mutex);
	auto it = g_instances.find(handle);
	if (it == g_instances.end()) {
	throw std::runtime_error("Invalid ngram handle: " + std::to_string(handle));
	}
	return *it->second;
	}

[gemini-code-assist]

The _to_csr function is inefficient for large batches or frequent calls. Iterating through Python lists and using extend in a loop, followed by torch.tensor() on a large Python list, involves significant overhead and multiple data copies. Since this is used in the speculative decoding hot path, optimizing this conversion is recommended.

Suggested change

	def _to_csr(batch_tokens: List[List[int]]) -> Tuple[torch.Tensor, torch.Tensor]:
	flat = []
	offsets = [0]
	for seq in batch_tokens:
	flat.extend(seq)
	offsets.append(len(flat))
	tokens_flat = torch.tensor(flat, dtype=torch.int32)
	offsets_t = torch.tensor(offsets, dtype=torch.int64)
	return tokens_flat, offsets_t
	def _to_csr(batch_tokens: List[List[int]]) -> Tuple[torch.Tensor, torch.Tensor]:
	flat = [t for seq in batch_tokens for t in seq]
	offsets = [0]
	curr = 0
	for seq in batch_tokens:
	curr += len(seq)
	offsets.append(curr)
	tokens_flat = torch.tensor(flat, dtype=torch.int32)
	offsets_t = torch.tensor(offsets, dtype=torch.int64)
	return tokens_flat, offsets_t

[gemini-code-assist]

Allocating out_tokens and out_mask tensors on every call to ngram_batch_match adds unnecessary overhead to the inference path. These buffers could be pre-allocated or cached based on the batch_size and draft_token_num to improve performance and reduce memory fragmentation.

[gemini-code-assist]

The conversion to np.int64 via .astype(np.int64) creates an additional copy of the data. If the downstream components can handle int32 for tokens and uint8 for masks (which is more memory-efficient), this conversion should be removed.

[github-actions]

✅ 1-gpu-5090: View workflow run

cd test/ && python3 registered/spec/utils/test_ngram_corpus.py

✅ 1-gpu-h100: View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py

[github-actions]

✅ 1-gpu-5090: View workflow run

cd test/ && python3 registered/spec/utils/test_ngram_corpus.py

✅ 1-gpu-h100: View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py

[github-actions]

✅ 1-gpu-5090: View workflow run

cd test/ && python3 registered/spec/utils/test_ngram_corpus.py

✅ 1-gpu-h100: View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py

[github-actions]

✅ 1-gpu-5090: View workflow run

cd test/ && python3 registered/spec/utils/test_ngram_corpus.py

✅ 1-gpu-h100: View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py