From 84e1108312b52f8e00032845af2d85a3073d8aae Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 2 Jun 2026 17:37:34 +0800 Subject: [PATCH] Optimize ngram decode id computation (#24757) Co-authored-by: Codex Co-authored-by: BBuf --- .../benchmark/bench_ngram_compute_decode.py | 121 +++++++++++++++++ .../jit_kernel/csrc/ngram_embedding.cuh | 128 ++++++++++++++++++ python/sglang/jit_kernel/ngram_embedding.py | 33 +++++ .../jit_kernel/tests/test_ngram_embedding.py | 97 +++++++++++++ python/sglang/srt/layers/n_gram_embedding.py | 60 +++++--- 5 files changed, 417 insertions(+), 22 deletions(-) create mode 100644 python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py create mode 100644 python/sglang/jit_kernel/tests/test_ngram_embedding.py diff --git a/python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py b/python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py new file mode 100644 index 000000000..63ecfaf6a --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py @@ -0,0 +1,121 @@ +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + get_benchmark_range, + run_benchmark_no_cudagraph, +) +from sglang.jit_kernel.ngram_embedding import ( + compute_n_gram_ids, + compute_n_gram_ids_decode, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large") + +NE_N = 8 +NE_K = 2 +VOCAB_SIZE = 32000 +MAX_CONTEXT_LEN = 1024 +BATCH_SIZE_LIST = get_benchmark_range( + full_range=[1, 2, 8, 32, 128, 512, 1024, 2048, 4096], + ci_range=[32, 1024], +) + + +def _make_ngram_params(): + ne_weights = torch.zeros([NE_N - 1, NE_K, NE_N], dtype=torch.int32) + ne_mods = torch.zeros([NE_N - 1, NE_K], dtype=torch.int32) + exclusive_sums = torch.zeros([(NE_N - 1) * NE_K + 1], dtype=torch.int32) + + for n in range(2, NE_N + 1): + for k in range(NE_K): + config_id = (n - 2) * NE_K + k + mod = 65537 + 2 * config_id + ne_mods[n - 2][k] = mod + exclusive_sums[config_id + 1] = exclusive_sums[config_id] + mod + for delta in range(NE_N): + ne_weights[n - 2][k][delta] = pow(VOCAB_SIZE, delta, mod) + + return ( + ne_weights.to(DEFAULT_DEVICE), + ne_mods.to(DEFAULT_DEVICE), + exclusive_sums.to(DEFAULT_DEVICE), + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=BATCH_SIZE_LIST, + line_arg="provider", + line_vals=["general", "decode"], + line_names=["general compute_n_gram_ids", "decode fast path"], + styles=[("blue", "-"), ("orange", "-")], + ylabel="us", + plot_name="ngram-compute-decode", + args={}, + ) +) +def benchmark(batch_size: int, provider: str): + num_configs = (NE_N - 1) * NE_K + max_running_reqs = batch_size + 8 + ne_weights, ne_mods, exclusive_sums = _make_ngram_params() + ne_token_table = torch.randint( + 0, + VOCAB_SIZE, + (max_running_reqs, MAX_CONTEXT_LEN), + dtype=torch.int32, + device=DEFAULT_DEVICE, + ) + row_indices = torch.arange(batch_size, dtype=torch.int64, device=DEFAULT_DEVICE) + column_starts = torch.randint( + 0, MAX_CONTEXT_LEN, (batch_size,), dtype=torch.int32, device=DEFAULT_DEVICE + ) + n_gram_ids = torch.empty( + (batch_size, num_configs), dtype=torch.int32, device=DEFAULT_DEVICE + ) + + if provider == "general": + tokens = torch.empty(batch_size, dtype=torch.int32, device=DEFAULT_DEVICE) + exclusive_req_len_sums = torch.arange( + batch_size + 1, dtype=torch.int32, device=DEFAULT_DEVICE + ) + + def fn(): + compute_n_gram_ids( + NE_N, + NE_K, + ne_weights, + ne_mods, + exclusive_sums, + tokens, + exclusive_req_len_sums, + ne_token_table, + row_indices, + column_starts, + n_gram_ids, + ) + + else: + + def fn(): + compute_n_gram_ids_decode( + NE_N, + NE_K, + ne_weights, + ne_mods, + exclusive_sums, + ne_token_table, + row_indices, + column_starts, + n_gram_ids, + ) + + return run_benchmark_no_cudagraph(fn) + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/ngram_embedding.cuh b/python/sglang/jit_kernel/csrc/ngram_embedding.cuh index 733365443..0683e0154 100644 --- a/python/sglang/jit_kernel/csrc/ngram_embedding.cuh +++ b/python/sglang/jit_kernel/csrc/ngram_embedding.cuh @@ -13,6 +13,9 @@ namespace device::ngram_embedding { +constexpr int kDecodeBlockSize = 256; +constexpr int kMaxComputeNGramIdsDecodeBlocks = 65535; + __global__ void ComputeNGramIdsKernel( int batch_size, int ne_n, @@ -84,6 +87,51 @@ __global__ void ComputeNGramIdsKernel( } } +__global__ void ComputeNGramIdsDecodeKernel( + int batch_size, + int ne_n, + int ne_k, + const int* __restrict__ ne_weights, // [ne_n-1,ne_k,ne_n] + const int* __restrict__ ne_mods, // [ne_n-1,ne_k] + const int* __restrict__ exclusive_ne_embeder_size_sums, // [(ne_n-1)*ne_k] + const int* __restrict__ ne_token_table, // [max_running_reqs, max_context_len] + int max_context_len, // max_context_len + const int64_t* __restrict__ row_indices, // [batch_size] + const int* __restrict__ column_starts, // [batch_size] + int* __restrict__ n_gram_ids // [batch_size, (ne_n-1)*ne_k] +) { + const int num_configs = (ne_n - 1) * ne_k; + const int total_outputs = batch_size * num_configs; + + for (int output_idx = blockIdx.x * blockDim.x + threadIdx.x; output_idx < total_outputs; + output_idx += blockDim.x * gridDim.x) { + const int req_id = output_idx / num_configs; + const int config_idx = output_idx - req_id * num_configs; + const int k_idx = config_idx % ne_k; + const int n_idx = config_idx / ne_k; + const int weight_offset = n_idx * ne_k * ne_n + k_idx * ne_n; + const int ne_mod = ne_mods[n_idx * ne_k + k_idx]; + + uint64_t n_gram_id = 0; + const int64_t req_token_table_offset = row_indices[req_id] * static_cast(max_context_len); + const int64_t current_token_table_offset = req_token_table_offset + column_starts[req_id]; + for (int j = 0; j < n_idx + 2; j++) { + if (current_token_table_offset - j < req_token_table_offset) { + break; + } + const int token = ne_token_table[current_token_table_offset - j]; + if (token < 0) { + break; + } + const uint64_t term = static_cast(token) * static_cast(ne_weights[weight_offset + j]); + n_gram_id += term % ne_mod; + } + n_gram_id %= ne_mod; + n_gram_id += exclusive_ne_embeder_size_sums[n_idx * ne_k + k_idx]; + n_gram_ids[output_idx] = static_cast(n_gram_id); + } +} + __global__ void UpdateTokenTableKernel( int batch_size, int* tokens, // [token_num] @@ -213,6 +261,86 @@ struct NgramEmbeddingKernel { static_cast(n_gram_ids.data_ptr())); } + static void compute_n_gram_ids_decode( + const int64_t ne_n, + const int64_t ne_k, + const tvm::ffi::TensorView ne_weights, + const tvm::ffi::TensorView ne_mods, + const tvm::ffi::TensorView exclusive_ne_embeder_size_sums, + const tvm::ffi::TensorView ne_token_table, + const tvm::ffi::TensorView row_indices, + const tvm::ffi::TensorView column_starts, + const tvm::ffi::TensorView n_gram_ids) { + using namespace host; + + auto device_ = SymbolicDevice{}; + auto batch_size = SymbolicSize{"batch_size"}; + + TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] + .with_dtype() + .with_device(device_) + .verify(ne_weights); + + TensorMatcher({-1, -1}) // [ne_n-1, ne_k] + .with_dtype() + .with_device() + .verify(ne_mods); + + TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] + .with_dtype() + .with_device() + .verify(exclusive_ne_embeder_size_sums); + + TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] + .with_dtype() + .with_device() + .verify(ne_token_table); + + TensorMatcher({batch_size}) // [batch_size] + .with_dtype() + .with_device() + .verify(row_indices); + + TensorMatcher({batch_size}) // [batch_size] + .with_dtype() + .with_device() + .verify(column_starts); + + TensorMatcher({batch_size, -1}) // [batch_size, (ne_n-1)*ne_k] + .with_dtype() + .with_device() + .verify(n_gram_ids); + + const int bs = static_cast(batch_size.unwrap()); + if (bs <= 0) { + return; + } + + const int max_context_len = static_cast(ne_token_table.size(1)); + const int num_configs = (static_cast(ne_n) - 1) * static_cast(ne_k); + const int total_outputs = bs * num_configs; + const auto stream = LaunchKernel::resolve_device(device_.unwrap()); + + constexpr int kBlockSize = device::ngram_embedding::kDecodeBlockSize; + const int grid_size = std::min( + device::ngram_embedding::kMaxComputeNGramIdsDecodeBlocks, + static_cast(div_ceil(total_outputs, kBlockSize))); + + LaunchKernel(grid_size, kBlockSize, stream)( + device::ngram_embedding::ComputeNGramIdsDecodeKernel, + bs, + static_cast(ne_n), + static_cast(ne_k), + static_cast(ne_weights.data_ptr()), + static_cast(ne_mods.data_ptr()), + static_cast(exclusive_ne_embeder_size_sums.data_ptr()), + static_cast(ne_token_table.data_ptr()), + max_context_len, + static_cast(row_indices.data_ptr()), + static_cast(column_starts.data_ptr()), + static_cast(n_gram_ids.data_ptr())); + } + static void update_token_table( const tvm::ffi::TensorView tokens, const tvm::ffi::TensorView ne_token_table, diff --git a/python/sglang/jit_kernel/ngram_embedding.py b/python/sglang/jit_kernel/ngram_embedding.py index f07937787..1d9613f98 100644 --- a/python/sglang/jit_kernel/ngram_embedding.py +++ b/python/sglang/jit_kernel/ngram_embedding.py @@ -17,6 +17,10 @@ def _jit_ngram_embedding_module() -> Module: cuda_files=["ngram_embedding.cuh"], cuda_wrappers=[ ("compute_n_gram_ids", "&NgramEmbeddingKernel::compute_n_gram_ids"), + ( + "compute_n_gram_ids_decode", + "&NgramEmbeddingKernel::compute_n_gram_ids_decode", + ), ("update_token_table", "&NgramEmbeddingKernel::update_token_table"), ], ) @@ -68,6 +72,35 @@ def compute_n_gram_ids( ) +@debug_kernel_api +def compute_n_gram_ids_decode( + ne_n: int, + ne_k: int, + ne_weights: torch.Tensor, + ne_mods: torch.Tensor, + exclusive_ne_embedder_size_sums: torch.Tensor, + ne_token_table: torch.Tensor, + row_indices: torch.Tensor, + column_starts: torch.Tensor, + n_gram_ids: torch.Tensor, +) -> None: + """ + Compute n-gram IDs for decode, where each request contributes one token. + """ + module = _jit_ngram_embedding_module() + module.compute_n_gram_ids_decode( + ne_n, + ne_k, + ne_weights, + ne_mods, + exclusive_ne_embedder_size_sums, + ne_token_table, + row_indices, + column_starts, + n_gram_ids, + ) + + @debug_kernel_api def update_token_table( tokens: torch.Tensor, diff --git a/python/sglang/jit_kernel/tests/test_ngram_embedding.py b/python/sglang/jit_kernel/tests/test_ngram_embedding.py new file mode 100644 index 000000000..fe0b0162f --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_ngram_embedding.py @@ -0,0 +1,97 @@ +import sys + +import pytest +import torch + +from sglang.jit_kernel.ngram_embedding import ( + compute_n_gram_ids, + compute_n_gram_ids_decode, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + + +def _make_ngram_params(ne_n: int, ne_k: int, vocab_size: int): + ne_weights = torch.zeros([ne_n - 1, ne_k, ne_n], dtype=torch.int32) + ne_mods = torch.zeros([ne_n - 1, ne_k], dtype=torch.int32) + exclusive_sums = torch.zeros([(ne_n - 1) * ne_k + 1], dtype=torch.int32) + + for n in range(2, ne_n + 1): + for k in range(ne_k): + config_id = (n - 2) * ne_k + k + mod = 65537 + 2 * config_id + ne_mods[n - 2][k] = mod + exclusive_sums[config_id + 1] = exclusive_sums[config_id] + mod + for delta in range(ne_n): + ne_weights[n - 2][k][delta] = pow(vocab_size, delta, mod) + + return ( + ne_weights.cuda(), + ne_mods.cuda(), + exclusive_sums.cuda(), + ) + + +@pytest.mark.parametrize("batch_size", [1, 2, 17, 128, 1024]) +def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None: + ne_n = 8 + ne_k = 2 + vocab_size = 32000 + max_context_len = 1024 + max_running_reqs = batch_size + 8 + num_configs = (ne_n - 1) * ne_k + + ne_weights, ne_mods, exclusive_sums = _make_ngram_params(ne_n, ne_k, vocab_size) + ne_token_table = torch.randint( + 0, + vocab_size, + (max_running_reqs, max_context_len), + dtype=torch.int32, + device="cuda", + ) + row_indices = torch.randperm(max_running_reqs, device="cuda")[:batch_size].to( + torch.int64 + ) + column_starts = torch.randint( + 0, max_context_len, (batch_size,), dtype=torch.int32, device="cuda" + ) + tokens = torch.empty(batch_size, dtype=torch.int32, device="cuda") + exclusive_req_len_sums = torch.arange( + batch_size + 1, dtype=torch.int32, device="cuda" + ) + n_gram_ids_general = torch.empty( + (batch_size, num_configs), dtype=torch.int32, device="cuda" + ) + n_gram_ids_decode = torch.empty_like(n_gram_ids_general) + + compute_n_gram_ids( + ne_n=ne_n, + ne_k=ne_k, + ne_weights=ne_weights, + ne_mods=ne_mods, + exclusive_ne_embedder_size_sums=exclusive_sums, + tokens=tokens, + exclusive_req_len_sums=exclusive_req_len_sums, + ne_token_table=ne_token_table, + row_indices=row_indices, + column_starts=column_starts, + n_gram_ids=n_gram_ids_general, + ) + compute_n_gram_ids_decode( + ne_n=ne_n, + ne_k=ne_k, + ne_weights=ne_weights, + ne_mods=ne_mods, + exclusive_ne_embedder_size_sums=exclusive_sums, + ne_token_table=ne_token_table, + row_indices=row_indices, + column_starts=column_starts, + n_gram_ids=n_gram_ids_decode, + ) + + torch.testing.assert_close(n_gram_ids_decode, n_gram_ids_general, atol=0, rtol=0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/python/sglang/srt/layers/n_gram_embedding.py b/python/sglang/srt/layers/n_gram_embedding.py index e6ac56326..0ddbd529d 100644 --- a/python/sglang/srt/layers/n_gram_embedding.py +++ b/python/sglang/srt/layers/n_gram_embedding.py @@ -2,7 +2,10 @@ import torch from torch import nn from torch.nn import Parameter -from sglang.jit_kernel.ngram_embedding import compute_n_gram_ids +from sglang.jit_kernel.ngram_embedding import ( + compute_n_gram_ids, + compute_n_gram_ids_decode, +) from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -135,27 +138,40 @@ class NgramEmbedding(torch.nn.Module): or forward_batch.forward_mode.is_decode() ): ngram_embedding_info = forward_batch.ngram_embedding_info - torch.cumsum( - ngram_embedding_info.req_lens, - dim=0, - dtype=torch.int32, - out=self.exclusive_req_len_sums[1 : 1 + forward_batch.batch_size], - ) - compute_n_gram_ids( - ne_n=self.over_embedding_n, - ne_k=self.over_embedding_k, - ne_weights=self.oe_weights, - ne_mods=self.oe_mods, - tokens=input_ids.to(torch.int32), - exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums, - exclusive_req_len_sums=self.exclusive_req_len_sums[ - : forward_batch.batch_size + 1 - ], - ne_token_table=ngram_embedding_info.token_table, - row_indices=forward_batch.req_pool_indices, - column_starts=ngram_embedding_info.column_starts, - n_gram_ids=self.oe_n_gram_ids[: len(input_ids)], - ) + if forward_batch.forward_mode.is_decode(): + compute_n_gram_ids_decode( + ne_n=self.over_embedding_n, + ne_k=self.over_embedding_k, + ne_weights=self.oe_weights, + ne_mods=self.oe_mods, + exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums, + ne_token_table=ngram_embedding_info.token_table, + row_indices=forward_batch.req_pool_indices, + column_starts=ngram_embedding_info.column_starts, + n_gram_ids=self.oe_n_gram_ids[: len(input_ids)], + ) + else: + torch.cumsum( + ngram_embedding_info.req_lens, + dim=0, + dtype=torch.int32, + out=self.exclusive_req_len_sums[1 : 1 + forward_batch.batch_size], + ) + compute_n_gram_ids( + ne_n=self.over_embedding_n, + ne_k=self.over_embedding_k, + ne_weights=self.oe_weights, + ne_mods=self.oe_mods, + tokens=input_ids.to(torch.int32), + exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums, + exclusive_req_len_sums=self.exclusive_req_len_sums[ + : forward_batch.batch_size + 1 + ], + ne_token_table=ngram_embedding_info.token_table, + row_indices=forward_batch.req_pool_indices, + column_starts=ngram_embedding_info.column_starts, + n_gram_ids=self.oe_n_gram_ids[: len(input_ids)], + ) # [13, seq_len, hidden_dim] all_hidden_states = torch.empty(