chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[1/n jit_kernel restruct] unify cache usage and clean up naming in ngram_embedding (#20244)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-10 15:53:43 +08:00
committed by GitHub
parent 7cf0551014
commit c812504b92
4 changed files with 14 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/jit_kernel/add_constant.py
View File

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
@cache_once
def _jit_add_constant_module(constant: int) -> Module:
args = make_cpp_args(constant) # pass all the template argument
args = make_cpp_args(constant)
return load_jit(
"add_constant",
*args,

11
python/sglang/jit_kernel/ngram_embedding.py
View File

@@ -1,16 +1,15 @@
from __future__ import annotations
from functools import lru_cache
from typing import TYPE_CHECKING
from sglang.jit_kernel.utils import load_jit
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
import torch
from tvm_ffi.module import Module
@lru_cache(maxsize=None)
@cache_once
def _jit_ngram_embedding_module() -> Module:
return load_jit(
"ngram_embedding",
@@ -27,7 +26,7 @@ def compute_n_gram_ids(
ne_k: int,
ne_weights: torch.Tensor,
ne_mods: torch.Tensor,
exclusive_ne_embeder_size_sums: torch.Tensor,
exclusive_ne_embedder_size_sums: torch.Tensor,
tokens: torch.Tensor,
exclusive_req_len_sums: torch.Tensor,
ne_token_table: torch.Tensor,
@@ -43,7 +42,7 @@ def compute_n_gram_ids(
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_embeder_size_sums: exclusive sum of embedder sizes
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
@@ -57,7 +56,7 @@ def compute_n_gram_ids(
ne_k,
ne_weights,
ne_mods,
exclusive_ne_embeder_size_sums,
exclusive_ne_embedder_size_sums,
tokens,
exclusive_req_len_sums,
ne_token_table,

2
python/sglang/jit_kernel/utils.py
View File

@@ -52,7 +52,7 @@ def _resolve_kernel_path() -> pathlib.Path:
path = _environment_install() or _package_install()
if path is None:
raise RuntimeError("Cannot find sgl-kernel/jit path")
raise RuntimeError("Cannot find sglang.jit_kernel path")
return path

14
python/sglang/srt/layers/n_gram_embedding.py
View File

@@ -35,18 +35,18 @@ class NgramEmbedding(torch.nn.Module):
)
self.n_grams = (over_embedding_n - 1) * over_embedding_k
oe_hidden_dim = embedding_dim // (over_embedding_k * (over_embedding_n - 1))
self.exclusive_oe_embeder_size_sums = torch.zeros(
self.exclusive_oe_embedder_size_sums = torch.zeros(
[over_embedding_k * (over_embedding_n - 1) + 1],
dtype=torch.int32,
device="cuda",
)
for i in range(over_embedding_k * (over_embedding_n - 1)):
self.exclusive_oe_embeder_size_sums[i + 1] = (
self.exclusive_oe_embeder_size_sums[i]
self.exclusive_oe_embedder_size_sums[i + 1] = (
self.exclusive_oe_embedder_size_sums[i]
+ int(over_embedding_m + i * 2 + 1)
)
self.oe_embeder = VocabParallelEmbedding(
num_embeddings=self.exclusive_oe_embeder_size_sums[-1],
num_embeddings=self.exclusive_oe_embedder_size_sums[-1],
embedding_dim=oe_hidden_dim,
enable_tp=is_dp_attention_enabled(),
)
@@ -100,8 +100,8 @@ class NgramEmbedding(torch.nn.Module):
".weight", ""
)
)
oe_weight_start = self.exclusive_oe_embeder_size_sums[index]
oe_weight_end = self.exclusive_oe_embeder_size_sums[index + 1]
oe_weight_start = self.exclusive_oe_embedder_size_sums[index]
oe_weight_end = self.exclusive_oe_embedder_size_sums[index + 1]
assert (
oe_weight_end - oe_weight_start == loaded_weight.shape[0]
), f"{oe_weight_end - oe_weight_start=} {loaded_weight.shape[0]=}"
@@ -147,7 +147,7 @@ class NgramEmbedding(torch.nn.Module):
ne_weights=self.oe_weights,
ne_mods=self.oe_mods,
tokens=input_ids.to(torch.int32),
exclusive_ne_embeder_size_sums=self.exclusive_oe_embeder_size_sums,
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
],