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

[AMD] Accept ROCm tensors in JIT kernel TensorMatcher + register 4 kernel tests (#29822)

Browse Source
This commit is contained in:
Michael
2026-07-02 19:29:47 -07:00
committed by GitHub
parent a6bc7fef90
commit bee0f34e68
6 changed files with 41 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/add_constant.cuh
+5 -5
View File
@@ -62,11 +62,11 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
// 1. Validate input tensors // 1. Validate input tensors
SymbolicSize N = {"num_elements"}; SymbolicSize N = {"num_elements"};
SymbolicDevice device_; SymbolicDevice device_;
TensorMatcher({N}) // 1D tensor, must be contiguous TensorMatcher({N}) // 1D tensor, must be contiguous
.with_dtype<int32_t>() // must be int32 .with_dtype<int32_t>() // must be int32
.with_device<kDLCUDA>(device_) // must be on CUDA device .with_device<kDLGPU>(device_) // must be on GPU device (CUDA or ROCm)
.verify(dst) // check tensor dst .verify(dst) // check tensor dst
.verify(src); // check tensor src .verify(src); // check tensor src
// 2. Extract required parameters, prepare for kernel launch // 2. Extract required parameters, prepare for kernel launch
const size_t num_elements = N.unwrap(); const size_t num_elements = N.unwrap();
python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh
+4 -4
View File
@@ -177,19 +177,19 @@ struct CausalConv3dCatPadKernel {
auto out_h = SymbolicSize{"out_h"}; auto out_h = SymbolicSize{"out_h"};
auto out_w = SymbolicSize{"out_w"}; auto out_w = SymbolicSize{"out_w"};
auto device = SymbolicDevice{}; auto device = SymbolicDevice{};
device.set_options<kDLCUDA>(); device.set_options<kDLGPU>();
TensorMatcher({bsz, channels, t_size, h_size, w_size}) TensorMatcher({bsz, channels, t_size, h_size, w_size})
.with_dtype<T>() .with_dtype<T>()
.template with_device<kDLCUDA>(device) .template with_device<kDLGPU>(device)
.verify(x); .verify(x);
TensorMatcher({bsz, channels, cache_t, h_size, w_size}) TensorMatcher({bsz, channels, cache_t, h_size, w_size})
.with_dtype<T>() .with_dtype<T>()
.template with_device<kDLCUDA>(device) .template with_device<kDLGPU>(device)
.verify(cache); .verify(cache);
TensorMatcher({bsz, channels, out_t, out_h, out_w}) TensorMatcher({bsz, channels, out_t, out_h, out_w})
.with_dtype<T>() .with_dtype<T>()
.template with_device<kDLCUDA>(device) .template with_device<kDLGPU>(device)
.verify(out); .verify(out);
const int64_t depth_left = pad_d_left - cache_t.unwrap(); const int64_t depth_left = pad_d_left - cache_t.unwrap();
python/sglang/jit_kernel/csrc/ngram_embedding.cuh
+26 -26
View File
@@ -209,47 +209,47 @@ struct NgramEmbeddingKernel {
// Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions
TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(ne_weights); .verify(ne_weights);
TensorMatcher({-1, -1}) // [ne_n-1, ne_k] TensorMatcher({-1, -1}) // [ne_n-1, ne_k]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_mods); .verify(ne_mods);
TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(exclusive_ne_embeder_size_sums); .verify(exclusive_ne_embeder_size_sums);
TensorMatcher({-1}) // [token_num] TensorMatcher({-1}) // [token_num]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(tokens); .verify(tokens);
TensorMatcher({-1}) // [batch_size+1] TensorMatcher({-1}) // [batch_size+1]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(exclusive_req_len_sums); .verify(exclusive_req_len_sums);
TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_token_table); .verify(ne_token_table);
TensorMatcher({-1}) // [batch_size] TensorMatcher({-1}) // [batch_size]
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(row_indices); .verify(row_indices);
TensorMatcher({-1}) // [batch_size] TensorMatcher({-1}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(column_starts); .verify(column_starts);
TensorMatcher({-1, -1}) // [token_num, (ne_n-1)*ne_k] TensorMatcher({-1, -1}) // [token_num, (ne_n-1)*ne_k]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(n_gram_ids); .verify(n_gram_ids);
const int batch_size = static_cast<int>(exclusive_req_len_sums.size(0) - 1); const int batch_size = static_cast<int>(exclusive_req_len_sums.size(0) - 1);
@@ -294,37 +294,37 @@ struct NgramEmbeddingKernel {
TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(ne_weights); .verify(ne_weights);
TensorMatcher({-1, -1}) // [ne_n-1, ne_k] TensorMatcher({-1, -1}) // [ne_n-1, ne_k]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_mods); .verify(ne_mods);
TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(exclusive_ne_embeder_size_sums); .verify(exclusive_ne_embeder_size_sums);
TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_token_table); .verify(ne_token_table);
TensorMatcher({batch_size}) // [batch_size] TensorMatcher({batch_size}) // [batch_size]
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(row_indices); .verify(row_indices);
TensorMatcher({batch_size}) // [batch_size] TensorMatcher({batch_size}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(column_starts); .verify(column_starts);
TensorMatcher({batch_size, -1}) // [batch_size, (ne_n-1)*ne_k] TensorMatcher({batch_size, -1}) // [batch_size, (ne_n-1)*ne_k]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(n_gram_ids); .verify(n_gram_ids);
const int bs = static_cast<int>(batch_size.unwrap()); const int bs = static_cast<int>(batch_size.unwrap());
@@ -371,27 +371,27 @@ struct NgramEmbeddingKernel {
// Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions
TensorMatcher({-1}) // [token_num] TensorMatcher({-1}) // [token_num]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(tokens); .verify(tokens);
TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_token_table); .verify(ne_token_table);
TensorMatcher({-1}) // [batch_size] TensorMatcher({-1}) // [batch_size]
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(row_indices); .verify(row_indices);
TensorMatcher({-1}) // [batch_size] TensorMatcher({-1}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(column_starts); .verify(column_starts);
TensorMatcher({-1}) // [batch_size] TensorMatcher({-1}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(req_lens); .verify(req_lens);
// ignore_tokens can be empty or have values // ignore_tokens can be empty or have values
@@ -400,7 +400,7 @@ struct NgramEmbeddingKernel {
if (has_ignore_tokens) { if (has_ignore_tokens) {
TensorMatcher({-1}) // [ignore_token_num] TensorMatcher({-1}) // [ignore_token_num]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ignore_tokens); .verify(ignore_tokens);
} }
@@ -447,22 +447,22 @@ struct NgramEmbeddingKernel {
TensorMatcher({batch_size}) // [batch_size] TensorMatcher({batch_size}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>(device_) .with_device<kDLGPU>(device_)
.verify(tokens); .verify(tokens);
TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(ne_token_table); .verify(ne_token_table);
TensorMatcher({batch_size}) // [batch_size] TensorMatcher({batch_size}) // [batch_size]
.with_dtype<int64_t>() .with_dtype<int64_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(row_indices); .verify(row_indices);
TensorMatcher({batch_size}) // [batch_size] TensorMatcher({batch_size}) // [batch_size]
.with_dtype<int32_t>() .with_dtype<int32_t>()
.with_device<kDLCUDA>() .with_device<kDLGPU>()
.verify(column_starts); .verify(column_starts);
const int bs = static_cast<int>(batch_size.unwrap()); const int bs = static_cast<int>(batch_size.unwrap());
test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py
+2 -1
View File
@@ -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, fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
) )
from sglang.jit_kernel.utils import get_ci_test_range 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="1-gpu-large")
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200") 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" DEVICE = "cuda"
DTYPE = torch.bfloat16 DTYPE = torch.bfloat16
test/registered/jit/test_add_constant.py
+2 -1
View File
@@ -4,10 +4,11 @@ import pytest
import torch import torch
from sglang.jit_kernel.add_constant import add_constant 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=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_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]) @pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025, 4096, 4097])
test/registered/jit/test_ngram_embedding.py
+2 -1
View File
@@ -9,9 +9,10 @@ from sglang.jit_kernel.ngram_embedding import (
update_token_table, update_token_table,
update_token_table_decode, 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_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): def _make_ngram_params(ne_n: int, ne_k: int, vocab_size: int):