mirror of
https://github.com/kvcache-ai/sglang.git
synced 2026-07-15 11:53:52 +00:00
29 lines
690 B
Python
29 lines
690 B
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm_ffi.module import Module
|
|
|
|
|
|
@cache_once
|
|
def _jit_add_constant_module(constant: int) -> Module:
|
|
args = make_cpp_args(constant)
|
|
return load_jit(
|
|
"add_constant",
|
|
*args,
|
|
cuda_files=["add_constant.cuh"],
|
|
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
|
|
)
|
|
|
|
|
|
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
|
|
dst = torch.empty_like(src)
|
|
module = _jit_add_constant_module(constant)
|
|
module.add_constant(dst, src)
|
|
return dst
|