[AMD] fix(jit): port kv_canary write/verify/plan kernels to ROCm (#28357)

test/registered/jit/kv_canary/test_const_sync.py

@@ -5,9 +5,10 @@ from pathlib import Path
 
 import sglang.jit_kernel
 from sglang.jit_kernel.kv_canary import consts
-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=5, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=5, suite="jit-kernel-unit-test-amd")
 
 
 # Resolve the kernel source against the installed jit_kernel package rather

test/registered/jit/kv_canary/test_kernel_config.py

@@ -33,9 +33,10 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     dummy_pseudo_tensors,
     empty_extras,
 )
-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=60, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=60, suite="jit-kernel-unit-test-amd")
 
 _DEVICE = torch.device("cuda")
 

test/registered/jit/kv_canary/test_pipeline_e2e.py

@@ -38,9 +38,10 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     empty_extras,
     make_req_to_token,
 )
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_plan_fuzz.py

@@ -20,9 +20,10 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
     run_fuzz_combo,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_plan_hand.py

@@ -20,9 +20,10 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
     make_req_to_token,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_scatter_req_token_ids.py

@@ -10,10 +10,11 @@ from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
     launch_scatter_req_token_ids_kernel,
     scatter_req_token_ids_torch_reference,
 )
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=10, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_verify_fuzz.py

@@ -29,9 +29,10 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
     run_fuzz_combo,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_verify_hand.py

@@ -51,9 +51,10 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_write_fuzz.py

@@ -29,9 +29,10 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
     run_fuzz_combo,
 )
 from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")

test/registered/jit/kv_canary/test_write_hand.py

@@ -49,9 +49,10 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
     _hand_fold_all,
     _hand_fold_partial,
 )
-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, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
 
 
 _DEVICE = torch.device("cuda")