diff --git a/test/registered/jit/benchmark/bench_spec_topk1.py b/test/registered/jit/benchmark/bench_spec_topk1.py index 1635a1046..fada9272c 100644 --- a/test/registered/jit/benchmark/bench_spec_topk1.py +++ b/test/registered/jit/benchmark/bench_spec_topk1.py @@ -12,11 +12,12 @@ from sglang.jit_kernel.benchmark.utils import ( run_benchmark, ) from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" ) +register_amd_ci(est_time=30, stage="jit-kernel-benchmark", runner_config="amd") BATCH_SIZE_RANGE = get_benchmark_range( diff --git a/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py b/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py index 4c5b87239..2082a1925 100644 --- a/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py +++ b/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py @@ -8,11 +8,12 @@ from sglang.kernels.ops.embeddings.vocab_parallel_embedding import ( vocab_parallel_embedding, ) from sglang.srt.layers.vocab_parallel_embedding import get_masked_input_and_mask -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" ) +register_amd_ci(est_time=10, stage="jit-kernel-benchmark", runner_config="amd") # Key order must match the perf_report x_names. DEFAULTS = dict(