100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sglang.jit_kernel.utils import cache_once, load_jit
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
from tvm_ffi.module import Module
|
|
|
|
|
|
@cache_once
|
|
def _jit_ngram_embedding_module() -> Module:
|
|
return load_jit(
|
|
"ngram_embedding",
|
|
cuda_files=["ngram_embedding.cuh"],
|
|
cuda_wrappers=[
|
|
("compute_n_gram_ids", "&NgramEmbeddingKernel::compute_n_gram_ids"),
|
|
("update_token_table", "&NgramEmbeddingKernel::update_token_table"),
|
|
],
|
|
)
|
|
|
|
|
|
def compute_n_gram_ids(
|
|
ne_n: int,
|
|
ne_k: int,
|
|
ne_weights: torch.Tensor,
|
|
ne_mods: torch.Tensor,
|
|
exclusive_ne_embedder_size_sums: torch.Tensor,
|
|
tokens: torch.Tensor,
|
|
exclusive_req_len_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 embedding.
|
|
|
|
Args:
|
|
ne_n: n value for n-gram
|
|
ne_k: k value for n-gram configurations
|
|
ne_weights: weights tensor with shape [ne_n-1, ne_k, ne_n]
|
|
ne_mods: mods tensor with shape [ne_n-1, ne_k]
|
|
exclusive_ne_embedder_size_sums: exclusive sum of embedder sizes
|
|
tokens: input token ids
|
|
exclusive_req_len_sums: exclusive sum of request lengths
|
|
ne_token_table: token table for all requests
|
|
row_indices: row indices for each request
|
|
column_starts: column start positions for each request
|
|
n_gram_ids: output tensor for n-gram ids
|
|
"""
|
|
module = _jit_ngram_embedding_module()
|
|
module.compute_n_gram_ids(
|
|
ne_n,
|
|
ne_k,
|
|
ne_weights,
|
|
ne_mods,
|
|
exclusive_ne_embedder_size_sums,
|
|
tokens,
|
|
exclusive_req_len_sums,
|
|
ne_token_table,
|
|
row_indices,
|
|
column_starts,
|
|
n_gram_ids,
|
|
)
|
|
|
|
|
|
def update_token_table(
|
|
tokens: torch.Tensor,
|
|
ne_token_table: torch.Tensor,
|
|
row_indices: torch.Tensor,
|
|
column_starts: torch.Tensor,
|
|
req_lens: torch.Tensor,
|
|
ignore_tokens: torch.Tensor | None = None,
|
|
) -> None:
|
|
"""
|
|
Update the token table with new tokens.
|
|
|
|
Args:
|
|
tokens: input token ids
|
|
ne_token_table: token table for all requests
|
|
row_indices: row indices for each request
|
|
column_starts: column start positions for each request
|
|
req_lens: request lengths
|
|
ignore_tokens: tokens to be ignored (marked as negative in table)
|
|
"""
|
|
module = _jit_ngram_embedding_module()
|
|
if ignore_tokens is None:
|
|
# Create an empty tensor for ignore_tokens
|
|
ignore_tokens = tokens.new_empty(0, dtype=tokens.dtype)
|
|
module.update_token_table(
|
|
tokens,
|
|
ne_token_table,
|
|
row_indices,
|
|
column_starts,
|
|
req_lens,
|
|
ignore_tokens,
|
|
)
|