[Feature] Provide batched fill next token mask and batched accept string/token methods.

[Seven-Streams]

This PR provides a series of methods to batched fill next token mask and batched accept string/token for GrammarMatcher. In general, there are three new methods:

• batch_fill_next_token_bitmask(matchers: List["GrammarMatcher"], bitmask: ArrayLike, index: int = 0, max_threads: int = 16, debug_print: bool = False)
• batch_accept_string(matchers: List["GrammarMatcher"], strings: List[str], debug_print: bool = False)
• batch_accept_token(matchers: List["GrammarMatcher"], tokens: List[int], debug_print: bool = False)

These methods allow users to fill multiple GrammarMatchers' token masks in a batch, and reduce the overhead of the transversion between cpp and python.

[Ubospica]

x grammar, latency

Before (sequential); After(batch)

Use one example in JSB

[Seven-Streams]

[Ubospica]

LGTM. Is y axis the average time or total time? And why is batch=1 so slow? Didn't you warm it up? And can you also show speedup with batch is large, like 300

[Seven-Streams]

LGTM. Is y axis the average time or total time? And why is batch=1 so slow? Didn't you warm it up? And can you also show speedup with batch is large, like 300

I did warmed it up. I'll test it more.

[Seven-Streams]

@Ubospica

The testing code is here:

import xgrammar as xgr
from transformers import AutoTokenizer
import time
import matplotlib.pyplot as plt
import numpy as np

tokenizer_path = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True, trust_remote_code=True)
tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer)
compiler = xgr.GrammarCompiler(tokenizer_info)
regex = "(\\s\\S){24}"
compiled_grammar = compiler.compile_regex(regex)

batch_time_list = dict()
naive_time_list = dict()

batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256]

# Warm up
for batch_size in batch_sizes:
    matchers = [
        xgr.GrammarMatcher(compiled_grammar)
        for _ in range(batch_size)
    ]
    dummy_next_token_bitmask = xgr.allocate_token_bitmask(1, tokenizer.vocab_size)
    for matcher in matchers:
        matcher.fill_next_token_bitmask(dummy_next_token_bitmask)
    
    matchers = [
        xgr.GrammarMatcher(compiled_grammar)
        for _ in range(batch_size)
    ]
    dummy_next_batch_token_bitmask = xgr.allocate_token_bitmask(batch_size, tokenizer.vocab_size)
    xgr.GrammarMatcher.batch_fill_next_token_bitmask(matchers, dummy_next_batch_token_bitmask)

# Actual benchmarking
for _ in range(10):
    for batch_size in batch_sizes:
        matchers = [
            xgr.GrammarMatcher(compiled_grammar)
            for _ in range(batch_size)
        ]
        
        next_token_bitmask = xgr.allocate_token_bitmask(batch_size, tokenizer.vocab_size)
        
        # warm up
        for i, matcher in enumerate(matchers):
            matcher.fill_next_token_bitmask(next_token_bitmask, i)

        start_time = time.time_ns()
        for i, matcher in enumerate(matchers):
            matcher.fill_next_token_bitmask(next_token_bitmask, i)
        end_time = time.time_ns()
        elapsed_time = end_time - start_time
        if batch_size not in naive_time_list:
            naive_time_list[batch_size] = []
        naive_time_list[batch_size].append(elapsed_time)
        
        matchers = [
            xgr.GrammarMatcher(compiled_grammar)
            for _ in range(batch_size)
        ]

        next_batch_token_bitmask = xgr.allocate_token_bitmask(batch_size, tokenizer.vocab_size)
        
        # warm up
        xgr.GrammarMatcher.batch_fill_next_token_bitmask(matchers, next_batch_token_bitmask)
        
        start_time = time.time_ns()
        xgr.GrammarMatcher.batch_fill_next_token_bitmask(matchers, next_batch_token_bitmask)
        end_time = time.time_ns()
        elapsed_time = end_time - start_time
        if batch_size not in batch_time_list:
            batch_time_list[batch_size] = []
        batch_time_list[batch_size].append(elapsed_time)
        
# Draw results
naive_avg_time = [np.mean(naive_time_list[bs]) for bs in batch_sizes]
single_times = []
for bs in batch_sizes:
    single_times.extend([t / bs for t in naive_time_list[bs]])
naive_avg_time.insert(0, np.mean(single_times))
batch_avg_time = [np.mean(batch_time_list[bs]) for bs in batch_sizes]
single_times = []
for bs in batch_sizes:
    single_times.extend([t / bs for t in batch_time_list[bs]])
batch_avg_time.insert(0, np.mean(single_times))
    
x = np.arange(len(batch_sizes) + 1)
width = 0.2 
fig, ax = plt.subplots()
rects1 = ax.bar(x - width/2, naive_avg_time, width, label='Naive')
rects2 = ax.bar(x + width/2, batch_avg_time, width, label='Batch')
ax.set_ylabel('Time (ns)')
ax.set_xlabel('Batch Size')
ax.set_title('Time by Batch Size and Method')
ax.set_xticks(x)

ax.set_xticklabels(["avg"] + batch_sizes)
ax.legend()
ax.bar_label(rects1, padding=3)
ax.bar_label(rects2, padding=-3)
ax.set_yscale('log')
plt.savefig("batch_vs_naive.png")
plt.show()