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

[VLM] Adopt jit qk_norm kernel in VLM (#16171)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-01-01 10:10:36 +08:00
committed by GitHub
parent 12b89e51d8
commit 3a42c5e341
4 changed files with 26 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/jit_kernel/benchmark/bench_qknorm.py
View File

@@ -5,6 +5,10 @@ from typing import Tuple
import torch
import triton
import triton.testing
from sgl_kernel import rmsnorm
from sglang.jit_kernel.norm import fused_inplace_qknorm
from sglang.srt.utils import get_current_device_stream_fast
IS_CI = (
os.getenv("CI", "false").lower() == "true"
@@ -20,13 +24,12 @@ def sglang_aot_qknorm(
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
from sgl_kernel import rmsnorm
head_dim = q.shape[-1]
q = q.view(-1, head_dim)
k = k.view(-1, head_dim)
current_stream = torch.cuda.current_stream()
current_stream = get_current_device_stream_fast()
alt_stream.wait_stream(current_stream)
rmsnorm(q, q_weight, out=q)
with torch.cuda.stream(alt_stream):
@@ -40,7 +43,6 @@ def sglang_jit_qknorm(
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
from sglang.jit_kernel.norm import fused_inplace_qknorm
fused_inplace_qknorm(q, k, q_weight, k_weight)
@@ -51,7 +53,6 @@ def flashinfer_qknorm(
q_weight: torch.Tensor,
k_weight: torch.Tensor,
) -> None:
from flashinfer.norm import rmsnorm
rmsnorm(q, q_weight, out=q)
rmsnorm(k, k_weight, out=k)