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

python/sglang/jit_kernel/add_constant.py

@@ -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,

python/sglang/jit_kernel/ngram_embedding.py

@@ -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,

python/sglang/jit_kernel/utils.py

@@ -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
 
 

python/sglang/srt/layers/n_gram_embedding.py

@@ -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_embedder_size_sums[i + 1] = (
-                self.exclusive_oe_embeder_size_sums[i]
+                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_start = self.exclusive_oe_embedder_size_sums[index]
-            oe_weight_end = self.exclusive_oe_embeder_size_sums[index + 1]
+            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
                 ],