[Refactor] Use is_in_ci() utility in JIT kernel benchmarks (#17118)
This commit is contained in:
@@ -1,10 +1,10 @@
|
|||||||
import os
|
|
||||||
from typing import Optional, Tuple
|
from typing import Optional, Tuple
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import triton
|
import triton
|
||||||
import triton.testing
|
import triton.testing
|
||||||
|
|
||||||
|
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||||
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
|
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -22,10 +22,7 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
_is_hip = False
|
_is_hip = False
|
||||||
|
|
||||||
IS_CI = (
|
IS_CI = is_in_ci()
|
||||||
os.getenv("CI", "false").lower() == "true"
|
|
||||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
|
||||||
)
|
|
||||||
|
|
||||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
import itertools
|
import itertools
|
||||||
import os
|
|
||||||
from typing import Tuple
|
from typing import Tuple
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
@@ -7,13 +6,11 @@ import triton
|
|||||||
import triton.testing
|
import triton.testing
|
||||||
from sgl_kernel import rmsnorm
|
from sgl_kernel import rmsnorm
|
||||||
|
|
||||||
|
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||||
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
from sglang.jit_kernel.norm import fused_inplace_qknorm
|
||||||
from sglang.srt.utils import get_current_device_stream_fast
|
from sglang.srt.utils import get_current_device_stream_fast
|
||||||
|
|
||||||
IS_CI = (
|
IS_CI = is_in_ci()
|
||||||
os.getenv("CI", "false").lower() == "true"
|
|
||||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
|
||||||
)
|
|
||||||
|
|
||||||
alt_stream = torch.cuda.Stream()
|
alt_stream = torch.cuda.Stream()
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
import itertools
|
import itertools
|
||||||
import os
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import triton
|
import triton
|
||||||
@@ -7,12 +6,10 @@ import triton.testing
|
|||||||
from flashinfer import rmsnorm as fi_rmsnorm
|
from flashinfer import rmsnorm as fi_rmsnorm
|
||||||
from sgl_kernel import rmsnorm
|
from sgl_kernel import rmsnorm
|
||||||
|
|
||||||
|
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
||||||
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
|
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
|
||||||
|
|
||||||
IS_CI = (
|
IS_CI = is_in_ci()
|
||||||
os.getenv("CI", "false").lower() == "true"
|
|
||||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def sglang_aot_rmsnorm(
|
def sglang_aot_rmsnorm(
|
||||||
|
|||||||
Reference in New Issue
Block a user