diff --git a/python/sglang/jit_kernel/csrc/add_constant.cuh b/python/sglang/jit_kernel/csrc/add_constant.cuh index f42be433c..754ec27cb 100644 --- a/python/sglang/jit_kernel/csrc/add_constant.cuh +++ b/python/sglang/jit_kernel/csrc/add_constant.cuh @@ -62,11 +62,11 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { // 1. Validate input tensors SymbolicSize N = {"num_elements"}; SymbolicDevice device_; - TensorMatcher({N}) // 1D tensor, must be contiguous - .with_dtype() // must be int32 - .with_device(device_) // must be on CUDA device - .verify(dst) // check tensor dst - .verify(src); // check tensor src + TensorMatcher({N}) // 1D tensor, must be contiguous + .with_dtype() // must be int32 + .with_device(device_) // must be on GPU device (CUDA or ROCm) + .verify(dst) // check tensor dst + .verify(src); // check tensor src // 2. Extract required parameters, prepare for kernel launch const size_t num_elements = N.unwrap(); diff --git a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh index a0adf7cd3..c17093902 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh +++ b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh @@ -177,19 +177,19 @@ struct CausalConv3dCatPadKernel { auto out_h = SymbolicSize{"out_h"}; auto out_w = SymbolicSize{"out_w"}; auto device = SymbolicDevice{}; - device.set_options(); + device.set_options(); TensorMatcher({bsz, channels, t_size, h_size, w_size}) .with_dtype() - .template with_device(device) + .template with_device(device) .verify(x); TensorMatcher({bsz, channels, cache_t, h_size, w_size}) .with_dtype() - .template with_device(device) + .template with_device(device) .verify(cache); TensorMatcher({bsz, channels, out_t, out_h, out_w}) .with_dtype() - .template with_device(device) + .template with_device(device) .verify(out); const int64_t depth_left = pad_d_left - cache_t.unwrap(); diff --git a/python/sglang/jit_kernel/csrc/ngram_embedding.cuh b/python/sglang/jit_kernel/csrc/ngram_embedding.cuh index 4eac7a946..e44a4a36e 100644 --- a/python/sglang/jit_kernel/csrc/ngram_embedding.cuh +++ b/python/sglang/jit_kernel/csrc/ngram_embedding.cuh @@ -209,47 +209,47 @@ struct NgramEmbeddingKernel { // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] .with_dtype() - .with_device(device_) + .with_device(device_) .verify(ne_weights); TensorMatcher({-1, -1}) // [ne_n-1, ne_k] .with_dtype() - .with_device() + .with_device() .verify(ne_mods); TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] .with_dtype() - .with_device() + .with_device() .verify(exclusive_ne_embeder_size_sums); TensorMatcher({-1}) // [token_num] .with_dtype() - .with_device() + .with_device() .verify(tokens); TensorMatcher({-1}) // [batch_size+1] .with_dtype() - .with_device() + .with_device() .verify(exclusive_req_len_sums); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() - .with_device() + .with_device() .verify(ne_token_table); TensorMatcher({-1}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(row_indices); TensorMatcher({-1}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(column_starts); TensorMatcher({-1, -1}) // [token_num, (ne_n-1)*ne_k] .with_dtype() - .with_device() + .with_device() .verify(n_gram_ids); const int batch_size = static_cast(exclusive_req_len_sums.size(0) - 1); @@ -294,37 +294,37 @@ struct NgramEmbeddingKernel { TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] .with_dtype() - .with_device(device_) + .with_device(device_) .verify(ne_weights); TensorMatcher({-1, -1}) // [ne_n-1, ne_k] .with_dtype() - .with_device() + .with_device() .verify(ne_mods); TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] .with_dtype() - .with_device() + .with_device() .verify(exclusive_ne_embeder_size_sums); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() - .with_device() + .with_device() .verify(ne_token_table); TensorMatcher({batch_size}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(row_indices); TensorMatcher({batch_size}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(column_starts); TensorMatcher({batch_size, -1}) // [batch_size, (ne_n-1)*ne_k] .with_dtype() - .with_device() + .with_device() .verify(n_gram_ids); const int bs = static_cast(batch_size.unwrap()); @@ -371,27 +371,27 @@ struct NgramEmbeddingKernel { // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions TensorMatcher({-1}) // [token_num] .with_dtype() - .with_device(device_) + .with_device(device_) .verify(tokens); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() - .with_device() + .with_device() .verify(ne_token_table); TensorMatcher({-1}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(row_indices); TensorMatcher({-1}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(column_starts); TensorMatcher({-1}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(req_lens); // ignore_tokens can be empty or have values @@ -400,7 +400,7 @@ struct NgramEmbeddingKernel { if (has_ignore_tokens) { TensorMatcher({-1}) // [ignore_token_num] .with_dtype() - .with_device() + .with_device() .verify(ignore_tokens); } @@ -447,22 +447,22 @@ struct NgramEmbeddingKernel { TensorMatcher({batch_size}) // [batch_size] .with_dtype() - .with_device(device_) + .with_device(device_) .verify(tokens); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() - .with_device() + .with_device() .verify(ne_token_table); TensorMatcher({batch_size}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(row_indices); TensorMatcher({batch_size}) // [batch_size] .with_dtype() - .with_device() + .with_device() .verify(column_starts); const int bs = static_cast(batch_size.unwrap()); diff --git a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py index 2d93e652a..25c2a3b0b 100644 --- a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py +++ b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py @@ -10,10 +10,11 @@ from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import ( fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton, ) from sglang.jit_kernel.utils import get_ci_test_range -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=45, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd") DEVICE = "cuda" DTYPE = torch.bfloat16 diff --git a/test/registered/jit/test_add_constant.py b/test/registered/jit/test_add_constant.py index 48f2476bf..bb9f0be50 100644 --- a/test/registered/jit/test_add_constant.py +++ b/test/registered/jit/test_add_constant.py @@ -4,10 +4,11 @@ import pytest import torch from sglang.jit_kernel.add_constant import add_constant -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=45, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True) +register_amd_ci(est_time=8, stage="jit-kernel-unit", runner_config="amd") @pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025, 4096, 4097]) diff --git a/test/registered/jit/test_ngram_embedding.py b/test/registered/jit/test_ngram_embedding.py index ee8f4def4..dd054fe29 100644 --- a/test/registered/jit/test_ngram_embedding.py +++ b/test/registered/jit/test_ngram_embedding.py @@ -9,9 +9,10 @@ from sglang.jit_kernel.ngram_embedding import ( update_token_table, update_token_table_decode, ) -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-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=8, stage="jit-kernel-unit", runner_config="amd") def _make_ngram_params(ne_n: int, ne_k: int, vocab_size: int):