From c4c405a46b5c618898ec0c9fb7dfc12d7517dbc3 Mon Sep 17 00:00:00 2001
From: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Date: Mon, 20 Jul 2026 23:51:51 -0700
Subject: [PATCH] [AMD] batch 3: register newly-added JIT kernel benchmarks for
 jit-kernel-benchmark-test-amd (#31615)

---
 test/registered/jit/benchmark/bench_spec_topk1.py              | 3 ++-
 .../registered/jit/benchmark/bench_vocab_parallel_embedding.py | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

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(