minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Michael
2026-06-19 12:19:37 -07:00
committed by GitHub
parent c436a8161a
commit 7c505c2927
13 changed files with 68 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh
+13 -10
View File
@@ -10,7 +10,6 @@
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
#include <cuda_runtime.h>
namespace {
@@ -96,7 +95,11 @@ __global__ void plan_entries_persistent_kernel(
static_cast<long long>(total_verify),
static_cast<long long>(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<int64_t>(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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({Nbs}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(req_pool_indices)
.verify(prefix_lens);
TensorMatcher({Nscratch}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(verify_offsets_scratch);
TensorMatcher({1}) //
.with_dtype<int32_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(verify_enable);
TensorMatcher({Ncap}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<int32_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(full_to_swa_index_mapping.value());
}
if constexpr (HAS_VERIFY_EXPECTED_TOKEN_POOL) {
TensorMatcher({Npool_rows, Npool_cols}) //
.with_dtype<int32_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(req_to_verify_expected_tokens.value());
TensorMatcher({Nbs}) //
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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");
python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh
+20 -17
View File
@@ -4,6 +4,7 @@
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
#include <sgl_kernel/utils.cuh> // For LaunchKernel, SGL_DEVICE
#include <sgl_kernel/warp.cuh> // For device::warp::reduce_sum
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
@@ -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<device::kWarpThreads>(local_active_count);
const bool is_subgroup_leader = (threadIdx.x % device::kWarpThreads) == 0u;
if (is_subgroup_leader && warp_active_count != 0u) {
atomicAdd(
reinterpret_cast<unsigned long long*>(p.slot_run_counter), static_cast<unsigned long long>(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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLCUDA>(device_).verify(canary_buf);
TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLGPU>(device_).verify(canary_buf);
TensorMatcher({N_verify})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(verify_slot_indices)
.verify(verify_expected_tokens)
.verify(verify_expected_positions)
.verify(verify_prev_slot_indices);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(verify_num_valid);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(verify_enable);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(verify_num_valid);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(verify_enable);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(violation_write_index);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(violation_write_index);
SymbolicSize N_ring = {"ring_capacity"};
TensorMatcher({N_ring, static_cast<int64_t>(kViolationFields)})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(violation_ring);
TensorMatcher({1})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(real_kv_buf_3);
TensorMatcher({static_cast<int64_t>(kMaxRealKvSources), static_cast<int64_t>(kRealKvSourceFieldsPerEntry)})
.with_dtype<int32_t>()
python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh
+15 -15
View File
@@ -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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLCUDA>(device_).verify(canary_buf);
TensorMatcher({N_slots, N_stride}).with_dtype<uint8_t>().with_device<kDLGPU>(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<int64_t>().with_device<kDLCUDA>(device_).verify(write_offsets);
TensorMatcher({N_write_reqs}).with_dtype<int64_t>().with_device<kDLCUDA>(device_).verify(write_seed_slot_indices);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(write_num_valid_reqs);
TensorMatcher({N_write_offsets}).with_dtype<int64_t>().with_device<kDLGPU>(device_).verify(write_offsets);
TensorMatcher({N_write_reqs}).with_dtype<int64_t>().with_device<kDLGPU>(device_).verify(write_seed_slot_indices);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(write_num_valid_reqs);
TensorMatcher({N_tokens})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(expected_input_tokens.value())
.verify(expected_input_positions.value());
}
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(violation_write_index);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(device_).verify(violation_write_index);
SymbolicSize N_ring = {"ring_capacity"};
TensorMatcher({N_ring, static_cast<int64_t>(kViolationFields)})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(violation_ring);
TensorMatcher({1})
.with_dtype<int64_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(slot_run_counter)
.verify(kernel_run_counter);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device_).verify(enable_chain_position_assert);
TensorMatcher({1}).with_dtype<int32_t>().with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(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<uint8_t>()
.with_device<kDLCUDA>(device_)
.with_device<kDLGPU>(device_)
.verify(real_kv_buf_3);
TensorMatcher({static_cast<int64_t>(kMaxRealKvSources), static_cast<int64_t>(kRealKvSourceFieldsPerEntry)})
.with_dtype<int32_t>()
test/registered/jit/kv_canary/test_const_sync.py
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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")