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
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,
)
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
test/registered/jit/test_add_constant.py
+2 -1
View File
@@ -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])
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_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):