biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[vLLM IR] rework gemma_rms_norm (#39014)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
This commit is contained in:
Jiangyun Zhu
2026-04-07 16:37:00 +08:00
committed by GitHub
parent da4c0e4db9
commit 8060bb0333
8 changed files with 106 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/kernels/aiter_ops.py
View File

@@ -36,13 +36,11 @@ AITER_SUPPORTED = is_aiter_found()
rms_no_var_16bit_only = (
lambda x, weight, epsilon, variance_size=None: variance_size is None
and x.dtype
in (
torch.float16,
torch.bfloat16,
)
and x.dtype in (torch.float16, torch.bfloat16)
and (weight is None or weight.dtype == x.dtype)
)
"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override."""
"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
and requires weight dtype to match x dtype."""
@ir.ops.rms_norm.register_impl(

7
vllm/kernels/vllm_c.py
View File

@@ -11,8 +11,11 @@ current_platform.import_kernels()
CUDA_ALIKE = current_platform.is_cuda_alike()
"""Most kernels in this file are supported on all CUDA-alike platforms."""
rms_no_var_size = lambda x, weight, epsilon, variance_size=None: variance_size is None
"""vLLM kernel does not support variance_size parameter."""
rms_no_var_size = (
lambda x, weight, epsilon, variance_size=None: variance_size is None
and (weight is None or weight.dtype == x.dtype)
)
"""vLLM kernel requires no variance_size override and matching input/weight dtype."""
@ir.ops.rms_norm.register_impl(

4
vllm/kernels/xpu_ops.py
View File

@@ -18,7 +18,9 @@ def is_xpu_kernels_found() -> bool:
XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
"""Kernels in this file are supported if vLLM XPU kernels are installed."""
rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and (
weight is None or weight.dtype == x.dtype
)
@ir.ops.rms_norm.register_impl(