feat: use sgl-kernel 0.0.3 in sglang (#3179)

This commit is contained in:
Yineng Zhang
2025-01-27 21:39:52 +08:00
committed by GitHub
parent 8a96f74988
commit 2f79f58873
6 changed files with 20 additions and 24 deletions

View File

@@ -19,10 +19,10 @@ from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
from sglang.srt.utils import is_flashinfer_available
from sglang.srt.utils import is_cuda_available
if is_flashinfer_available():
from flashinfer.norm import (
if is_cuda_available():
from sgl_kernel import (
fused_add_rmsnorm,
gemma_fused_add_rmsnorm,
gemma_rmsnorm,
@@ -121,8 +121,8 @@ class GemmaRMSNorm(CustomOp):
return out
if not is_flashinfer_available():
if not is_cuda_available():
logger.info(
"FlashInfer is not available on Non-NV platforms. Fallback to other kernel libraries."
"sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries."
)
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm