Optimize ngram decode id computation (#24757)
Co-authored-by: Codex <codex@example.com> Co-authored-by: BBuf <xiaoyu.zhang@radixark.net>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Codex
BBuf
Codex
BBuf
parent
f651b48764
commit
84e1108312
@@ -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)
|
||||
@@ -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<int64_t>(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<uint64_t>(token) * static_cast<uint64_t>(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<int>(n_gram_id);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void UpdateTokenTableKernel(
|
||||
int batch_size,
|
||||
int* tokens, // [token_num]
|
||||
@@ -213,6 +261,86 @@ struct NgramEmbeddingKernel {
|
||||
static_cast<int*>(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<int32_t>()
|
||||
.with_device<kDLCUDA>(device_)
|
||||
.verify(ne_weights);
|
||||
|
||||
TensorMatcher({-1, -1}) // [ne_n-1, ne_k]
|
||||
.with_dtype<int32_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(ne_mods);
|
||||
|
||||
TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1]
|
||||
.with_dtype<int32_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(exclusive_ne_embeder_size_sums);
|
||||
|
||||
TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len]
|
||||
.with_dtype<int32_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(ne_token_table);
|
||||
|
||||
TensorMatcher({batch_size}) // [batch_size]
|
||||
.with_dtype<int64_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(row_indices);
|
||||
|
||||
TensorMatcher({batch_size}) // [batch_size]
|
||||
.with_dtype<int32_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(column_starts);
|
||||
|
||||
TensorMatcher({batch_size, -1}) // [batch_size, (ne_n-1)*ne_k]
|
||||
.with_dtype<int32_t>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(n_gram_ids);
|
||||
|
||||
const int bs = static_cast<int>(batch_size.unwrap());
|
||||
if (bs <= 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int max_context_len = static_cast<int>(ne_token_table.size(1));
|
||||
const int num_configs = (static_cast<int>(ne_n) - 1) * static_cast<int>(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<int>(div_ceil(total_outputs, kBlockSize)));
|
||||
|
||||
LaunchKernel(grid_size, kBlockSize, stream)(
|
||||
device::ngram_embedding::ComputeNGramIdsDecodeKernel,
|
||||
bs,
|
||||
static_cast<int>(ne_n),
|
||||
static_cast<int>(ne_k),
|
||||
static_cast<const int*>(ne_weights.data_ptr()),
|
||||
static_cast<const int*>(ne_mods.data_ptr()),
|
||||
static_cast<const int*>(exclusive_ne_embeder_size_sums.data_ptr()),
|
||||
static_cast<const int*>(ne_token_table.data_ptr()),
|
||||
max_context_len,
|
||||
static_cast<const int64_t*>(row_indices.data_ptr()),
|
||||
static_cast<const int*>(column_starts.data_ptr()),
|
||||
static_cast<int*>(n_gram_ids.data_ptr()));
|
||||
}
|
||||
|
||||
static void update_token_table(
|
||||
const tvm::ffi::TensorView tokens,
|
||||
const tvm::ffi::TensorView ne_token_table,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]))
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user