From 7c505c2927728e2d5ebf7de21a46f857f8f37bba Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Fri, 19 Jun 2026 12:19:37 -0700 Subject: [PATCH] [AMD] fix(jit): port kv_canary write/verify/plan kernels to ROCm (#28357) --- .../csrc/kv_canary/canary_plan_entries.cuh | 23 +++++++----- .../csrc/kv_canary/canary_verify.cuh | 37 ++++++++++--------- .../csrc/kv_canary/canary_write.cuh | 30 +++++++-------- .../jit/kv_canary/test_const_sync.py | 3 +- .../jit/kv_canary/test_kernel_config.py | 3 +- .../jit/kv_canary/test_pipeline_e2e.py | 3 +- .../jit/kv_canary/test_plan_fuzz.py | 3 +- .../jit/kv_canary/test_plan_hand.py | 3 +- .../kv_canary/test_scatter_req_token_ids.py | 3 +- .../jit/kv_canary/test_verify_fuzz.py | 3 +- .../jit/kv_canary/test_verify_hand.py | 3 +- .../jit/kv_canary/test_write_fuzz.py | 3 +- .../jit/kv_canary/test_write_hand.py | 3 +- 13 files changed, 68 insertions(+), 52 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh b/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh index 7c6e1c7ef..781d958fa 100644 --- a/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh +++ b/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh @@ -10,7 +10,6 @@ #include #include -#include namespace { @@ -96,7 +95,11 @@ __global__ void plan_entries_persistent_kernel( static_cast(total_verify), static_cast(params.verify_capacity)); } +#ifndef USE_ROCM __trap(); +#else + __builtin_trap(); // HIP/clang device-side trap; __trap() is CUDA-only. +#endif } const int64_t tid_start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; @@ -195,31 +198,31 @@ struct PlanEntriesKernel { SymbolicSize Npool_rows = {"req_to_verify_expected_tokens_rows"}; SymbolicSize Npool_cols = {"req_to_verify_expected_tokens_cols"}; SymbolicDevice device_; - device_.set_options(); + device_.set_options(); TensorMatcher({Nbs}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(req_pool_indices) .verify(prefix_lens); TensorMatcher({Nscratch}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(verify_offsets_scratch); TensorMatcher({1}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(verify_enable); TensorMatcher({Ncap}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(out_verify_slot_indices) .verify(out_verify_expected_tokens) .verify(out_verify_expected_positions) .verify(out_verify_prev_slot_indices); TensorMatcher({Nmax_reqs, Nmax_seq_len}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(req_to_token); RuntimeCheck( full_to_swa_index_mapping.has_value() == HAS_SWA_LUT, @@ -235,17 +238,17 @@ struct PlanEntriesKernel { if constexpr (HAS_SWA_LUT) { TensorMatcher({Nlut}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(full_to_swa_index_mapping.value()); } if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) { TensorMatcher({Npool_rows, Npool_cols}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(req_to_verify_expected_tokens.value()); TensorMatcher({Nbs}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(req_to_verify_expected_tokens_valid_lens.value()); } RuntimeCheck(Nscratch.unwrap() >= Nbs.unwrap() + 1, "verify_offsets_scratch length must be >= bs_padded + 1"); diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh b/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh index 7c63016da..38840fefd 100644 --- a/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh +++ b/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh @@ -4,6 +4,7 @@ #include // For div_ceil, RuntimeCheck #include // For LaunchKernel, SGL_DEVICE +#include // For device::warp::reduce_sum #include #include @@ -124,11 +125,13 @@ __global__ void canary_verify_kernel(const VerifyKernelParams __grid_constant__ } } - uint32_t warp_active_count = local_active_count; - for (int offset = 16; offset > 0; offset >>= 1) { - warp_active_count += __shfl_down_sync(0xFFFFFFFFu, warp_active_count, offset); - } - if ((threadIdx.x & 31u) == 0u && warp_active_count != 0u) { + // Warp-size-agnostic reduction: device::warp::reduce_sum performs a width-kWarpThreads (32) xor + // all-reduce, so each 32-lane sub-group gets its own total with its own lane-0 leader. Correct on + // both CUDA (warp=32) and ROCm wave64 (two 32-lane sub-groups per wavefront); a hand-rolled + // __shfl_down over the full wavefront would cross the 32-lane boundary on wave64. + const uint32_t warp_active_count = device::warp::reduce_sum(local_active_count); + const bool is_subgroup_leader = (threadIdx.x % device::kWarpThreads) == 0u; + if (is_subgroup_leader && warp_active_count != 0u) { atomicAdd( reinterpret_cast(p.slot_run_counter), static_cast(warp_active_count)); } @@ -173,29 +176,29 @@ struct CanaryVerifyKernel { SymbolicSize N_stride = {"slot_stride_bytes"}; SymbolicSize N_verify = {"verify_capacity"}; SymbolicDevice device_; - device_.set_options(); + device_.set_options(); - TensorMatcher({N_slots, N_stride}).with_dtype().with_device(device_).verify(canary_buf); + TensorMatcher({N_slots, N_stride}).with_dtype().with_device(device_).verify(canary_buf); TensorMatcher({N_verify}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(verify_slot_indices) .verify(verify_expected_tokens) .verify(verify_expected_positions) .verify(verify_prev_slot_indices); - TensorMatcher({1}).with_dtype().with_device(device_).verify(verify_num_valid); - TensorMatcher({1}).with_dtype().with_device(device_).verify(verify_enable); + TensorMatcher({1}).with_dtype().with_device(device_).verify(verify_num_valid); + TensorMatcher({1}).with_dtype().with_device(device_).verify(verify_enable); - TensorMatcher({1}).with_dtype().with_device(device_).verify(violation_write_index); + TensorMatcher({1}).with_dtype().with_device(device_).verify(violation_write_index); SymbolicSize N_ring = {"ring_capacity"}; TensorMatcher({N_ring, static_cast(kViolationFields)}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(violation_ring); TensorMatcher({1}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(slot_run_counter) .verify(kernel_run_counter); @@ -205,25 +208,25 @@ struct CanaryVerifyKernel { SymbolicSize N_real_kv_cols_0 = {"real_kv_cols_0"}; TensorMatcher({N_real_kv_rows_0, N_real_kv_cols_0}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_0); SymbolicSize N_real_kv_rows_1 = {"real_kv_rows_1"}; SymbolicSize N_real_kv_cols_1 = {"real_kv_cols_1"}; TensorMatcher({N_real_kv_rows_1, N_real_kv_cols_1}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_1); SymbolicSize N_real_kv_rows_2 = {"real_kv_rows_2"}; SymbolicSize N_real_kv_cols_2 = {"real_kv_cols_2"}; TensorMatcher({N_real_kv_rows_2, N_real_kv_cols_2}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_2); SymbolicSize N_real_kv_rows_3 = {"real_kv_rows_3"}; SymbolicSize N_real_kv_cols_3 = {"real_kv_cols_3"}; TensorMatcher({N_real_kv_rows_3, N_real_kv_cols_3}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_3); TensorMatcher({static_cast(kMaxRealKvSources), static_cast(kRealKvSourceFieldsPerEntry)}) .with_dtype() diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh b/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh index 1ce2d0918..7ebe7dbf8 100644 --- a/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh +++ b/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh @@ -205,20 +205,20 @@ inline void canary_write_step_cuda( SymbolicSize N_write_reqs = {"write_req_capacity"}; SymbolicSize N_tokens = {"num_tokens_padded"}; SymbolicDevice device_; - device_.set_options(); + device_.set_options(); - TensorMatcher({N_slots, N_stride}).with_dtype().with_device(device_).verify(canary_buf); + TensorMatcher({N_slots, N_stride}).with_dtype().with_device(device_).verify(canary_buf); // write_offsets has shape [write_req_capacity + 1]; the length relationship is pinned by the // RuntimeCheck below, this matcher pins dtype + device. SymbolicSize N_write_offsets = {"write_offsets_len"}; - TensorMatcher({N_write_offsets}).with_dtype().with_device(device_).verify(write_offsets); - TensorMatcher({N_write_reqs}).with_dtype().with_device(device_).verify(write_seed_slot_indices); - TensorMatcher({1}).with_dtype().with_device(device_).verify(write_num_valid_reqs); + TensorMatcher({N_write_offsets}).with_dtype().with_device(device_).verify(write_offsets); + TensorMatcher({N_write_reqs}).with_dtype().with_device(device_).verify(write_seed_slot_indices); + TensorMatcher({1}).with_dtype().with_device(device_).verify(write_num_valid_reqs); TensorMatcher({N_tokens}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(input_ids) .verify(positions) .verify(out_cache_loc); @@ -232,47 +232,47 @@ inline void canary_write_step_cuda( if (enable_write_input_assert_bool) { TensorMatcher({N_tokens}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(expected_input_tokens.value()) .verify(expected_input_positions.value()); } - TensorMatcher({1}).with_dtype().with_device(device_).verify(violation_write_index); + TensorMatcher({1}).with_dtype().with_device(device_).verify(violation_write_index); SymbolicSize N_ring = {"ring_capacity"}; TensorMatcher({N_ring, static_cast(kViolationFields)}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(violation_ring); TensorMatcher({1}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(slot_run_counter) .verify(kernel_run_counter); - TensorMatcher({1}).with_dtype().with_device(device_).verify(enable_chain_position_assert); + TensorMatcher({1}).with_dtype().with_device(device_).verify(enable_chain_position_assert); SymbolicSize N_real_kv_rows_0 = {"real_kv_rows_0"}; SymbolicSize N_real_kv_cols_0 = {"real_kv_cols_0"}; TensorMatcher({N_real_kv_rows_0, N_real_kv_cols_0}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_0); SymbolicSize N_real_kv_rows_1 = {"real_kv_rows_1"}; SymbolicSize N_real_kv_cols_1 = {"real_kv_cols_1"}; TensorMatcher({N_real_kv_rows_1, N_real_kv_cols_1}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_1); SymbolicSize N_real_kv_rows_2 = {"real_kv_rows_2"}; SymbolicSize N_real_kv_cols_2 = {"real_kv_cols_2"}; TensorMatcher({N_real_kv_rows_2, N_real_kv_cols_2}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_2); SymbolicSize N_real_kv_rows_3 = {"real_kv_rows_3"}; SymbolicSize N_real_kv_cols_3 = {"real_kv_cols_3"}; TensorMatcher({N_real_kv_rows_3, N_real_kv_cols_3}) .with_dtype() - .with_device(device_) + .with_device(device_) .verify(real_kv_buf_3); TensorMatcher({static_cast(kMaxRealKvSources), static_cast(kRealKvSourceFieldsPerEntry)}) .with_dtype() diff --git a/test/registered/jit/kv_canary/test_const_sync.py b/test/registered/jit/kv_canary/test_const_sync.py index d6c800dcd..ab83a7e52 100644 --- a/test/registered/jit/kv_canary/test_const_sync.py +++ b/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 diff --git a/test/registered/jit/kv_canary/test_kernel_config.py b/test/registered/jit/kv_canary/test_kernel_config.py index 3f357e0b9..ab6d65882 100644 --- a/test/registered/jit/kv_canary/test_kernel_config.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_pipeline_e2e.py b/test/registered/jit/kv_canary/test_pipeline_e2e.py index df5d2ff72..e0bc7e094 100644 --- a/test/registered/jit/kv_canary/test_pipeline_e2e.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_plan_fuzz.py b/test/registered/jit/kv_canary/test_plan_fuzz.py index de7d6c125..8a94c9ec8 100644 --- a/test/registered/jit/kv_canary/test_plan_fuzz.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_plan_hand.py b/test/registered/jit/kv_canary/test_plan_hand.py index c0f9b1c3f..04b5612d5 100644 --- a/test/registered/jit/kv_canary/test_plan_hand.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_scatter_req_token_ids.py b/test/registered/jit/kv_canary/test_scatter_req_token_ids.py index 1bdb5bd5c..f065d3d43 100644 --- a/test/registered/jit/kv_canary/test_scatter_req_token_ids.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_verify_fuzz.py b/test/registered/jit/kv_canary/test_verify_fuzz.py index 729f23391..7594860b8 100644 --- a/test/registered/jit/kv_canary/test_verify_fuzz.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_verify_hand.py b/test/registered/jit/kv_canary/test_verify_hand.py index 6e78f34a0..f3d9f50c0 100644 --- a/test/registered/jit/kv_canary/test_verify_hand.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_write_fuzz.py b/test/registered/jit/kv_canary/test_write_fuzz.py index 3d56dc889..8fc04f0d5 100644 --- a/test/registered/jit/kv_canary/test_write_fuzz.py +++ b/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") diff --git a/test/registered/jit/kv_canary/test_write_hand.py b/test/registered/jit/kv_canary/test_write_hand.py index b3fee829c..e12dd4953 100644 --- a/test/registered/jit/kv_canary/test_write_hand.py +++ b/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")