diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index f65e1a721..ae3b060db 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -94,14 +94,14 @@ jobs: - "python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/**" - "python/sglang/srt/observability/**" - "python/sglang/kernels/ops/diffusion/**" - - "test/registered/jit/diffusion/**" - - "test/registered/jit/benchmark/diffusion/**" + - "test/registered/kernels/ops/diffusion/**" + - "test/registered/kernels/benchmark/diffusion/**" - "python/sglang/cli/**" jit_kernel: - ".github/workflows/pr-test.yml" - ".github/workflows/pr-test-jit-kernel.yml" - "python/pyproject.toml" - - "test/registered/jit/**" + - "test/registered/kernels/**" # sglang.kernels is the migrated kernel namespace (RFC #29630 / #30044); the # base-b-kernel suites import it directly, so kernel edits must run them. - "python/sglang/kernels/**" diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 4ba8b30b5..afdf9662f 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -199,15 +199,15 @@ jobs: - ".github/workflows/pr-test-amd-rocm720.yml" jit_kernel: - "python/sglang/kernels/**" - - "test/registered/jit/**" + - "test/registered/kernels/**" - ".github/workflows/pr-test-amd-rocm720.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - "python/sglang/kernels/ops/diffusion/**" - - "test/registered/jit/diffusion/**" - - "test/registered/jit/benchmark/diffusion/**" + - "test/registered/kernels/ops/diffusion/**" + - "test/registered/kernels/benchmark/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 7b4b15a26..566200d8e 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -187,15 +187,15 @@ jobs: - ".github/workflows/pr-test-amd.yml" jit_kernel: - "python/sglang/kernels/**" - - "test/registered/jit/**" + - "test/registered/kernels/**" - ".github/workflows/pr-test-amd.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - "python/sglang/kernels/ops/diffusion/**" - - "test/registered/jit/diffusion/**" - - "test/registered/jit/benchmark/diffusion/**" + - "test/registered/kernels/ops/diffusion/**" + - "test/registered/kernels/benchmark/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" diff --git a/python/sglang/kernels/README.md b/python/sglang/kernels/README.md index 817e34f45..7b69d06e7 100644 --- a/python/sglang/kernels/README.md +++ b/python/sglang/kernels/README.md @@ -19,21 +19,28 @@ sglang/kernels/ selector.py # heuristic select_kernel() and cached get_kernel() fused_op.py # BaseFusedOp: per-operator multi-backend contract ops/ - / # one subpackage per operator group + / # one subpackage per operator group (see list below) + jit/ # shared JIT CUDA build/runtime infra: utils/, csrc/, + # include/, __main__ (KERNEL_PATH resolves here) ``` -Groups populated in this phase: `activation`, `gemm`, `kvcache`, `layernorm`, -`moe`, `quantization`. The remaining groups (`attention`, `communication`, -`diffusion`, `grammar`, `mamba`, `memory`, `sampling`, `spatial`, -`speculative`) are reserved package placeholders whose implementations still -live in `sglang.kernels.jit` / `sgl_kernel` / `triton_ops` and will migrate in -later phases. +Operator groups (all populated): `activation`, `attention`, `communication`, +`diffusion`, `embeddings`, `gemm`, `grammar`, `kv_canary`, `kvcache`, +`layernorm`, `lplb`, `mamba`, `memory`, `model`, `moe`, `quantization`, +`sampling`, `spatial`, `speculative`. + +As of the RFC #29630 finale (#32072) the legacy `sglang.jit_kernel` package has +been **removed**: its shared build/runtime infra moved to `sglang.kernels.jit` +and each JIT-backed operator into its group as +`sglang.kernels.ops.._jit_`. Tests and benchmarks live under +`test/registered/kernels/` (`ops//` for tests, `benchmark//` for +benchmarks); shared test helpers are in `sglang.test.kernels`. ## How it works -Implementations are not moved yet. Each `ops.` function is a thin -wrapper that forwards to a chosen backend, and every backend is described by a -`KernelSpec` in the registry so alternatives can be inventoried and compared: +Each `ops.` function is a thin wrapper that forwards to a chosen +backend, and every backend is described by a `KernelSpec` in the registry so +alternatives can be inventoried and compared: - `register_kernel(KernelSpec(...))` records metadata only — an operator id (`"."`), a backend, and an import path (`"module:attr"`). No @@ -85,8 +92,8 @@ What this buys (see the - **Unified correctness testing** — a generic harness enumerates `available_backends()` and asserts each one matches `forward_native` - (`test/registered/kernels/test_fused_op_gpu_parity.py`); new backends are - picked up automatically. + (`test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py`); new + backends are picked up automatically. - **One-switch debugging** — `SGLANG_FORCE_FUSED_OP_BACKEND=torch` (or `set_fused_op_backend(KernelBackend.TORCH)`) flips *every* fused op to its reference implementation for numerical-bug bisection. diff --git a/python/sglang/kernels/ops/communication/mp.py b/python/sglang/kernels/ops/communication/mp.py index 5a961ddb5..957ac0647 100644 --- a/python/sglang/kernels/ops/communication/mp.py +++ b/python/sglang/kernels/ops/communication/mp.py @@ -1,6 +1,6 @@ """Multi-process / multi-GPU launching utilities (torchrun-based). -Shared `multigpu_launch` helper that both `sglang.kernels.jit.tests.utils` and +Shared `multigpu_launch` helper that both `sglang.test.kernels.utils` and `sglang.kernels.jit.benchmark.utils` build their domain-specific entry points on top of (`multigpu_pytest_main`, `multigpu_bench_main`). diff --git a/python/sglang/kernels/jit/tests/deepseek_v4/__init__.py b/python/sglang/test/kernels/__init__.py similarity index 100% rename from python/sglang/kernels/jit/tests/deepseek_v4/__init__.py rename to python/sglang/test/kernels/__init__.py diff --git a/python/sglang/kernels/jit/tests/kv_canary/__init__.py b/python/sglang/test/kernels/deepseek_v4/__init__.py similarity index 100% rename from python/sglang/kernels/jit/tests/kv_canary/__init__.py rename to python/sglang/test/kernels/deepseek_v4/__init__.py diff --git a/python/sglang/kernels/jit/tests/deepseek_v4/common.py b/python/sglang/test/kernels/deepseek_v4/common.py similarity index 100% rename from python/sglang/kernels/jit/tests/deepseek_v4/common.py rename to python/sglang/test/kernels/deepseek_v4/common.py diff --git a/python/sglang/test/kernels/kv_canary/__init__.py b/python/sglang/test/kernels/kv_canary/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py b/python/sglang/test/kernels/kv_canary/_canary_helpers.py similarity index 98% rename from python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py rename to python/sglang/test/kernels/kv_canary/_canary_helpers.py index 9aa1e595c..99bacfb51 100644 --- a/python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py +++ b/python/sglang/test/kernels/kv_canary/_canary_helpers.py @@ -5,21 +5,21 @@ from typing import Optional import torch -from sglang.kernels.jit.tests.kv_canary._constants import ( +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 +from sglang.kernels.ops.kv_canary.verify import VerifyPlan +from sglang.kernels.ops.kv_canary.write import WritePlan +from sglang.test.kernels.kv_canary._constants import ( _I64_SIGN_BIT, _U64_MASK, DEFAULT_NUM_SLOTS, DEFAULT_RING_CAPACITY, DEFAULT_SLOT_STRIDE_BYTES, ) -from sglang.kernels.jit.tests.kv_canary._fixtures import ( +from sglang.test.kernels.kv_canary._fixtures import ( make_real_kv_source, make_real_kv_sources, ) -from sglang.kernels.ops.kv_canary import consts -from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 -from sglang.kernels.ops.kv_canary.verify import VerifyPlan -from sglang.kernels.ops.kv_canary.write import WritePlan __all__ = [ "FakeViolationLog", diff --git a/python/sglang/kernels/jit/tests/kv_canary/_constants.py b/python/sglang/test/kernels/kv_canary/_constants.py similarity index 100% rename from python/sglang/kernels/jit/tests/kv_canary/_constants.py rename to python/sglang/test/kernels/kv_canary/_constants.py diff --git a/python/sglang/kernels/jit/tests/kv_canary/_differential.py b/python/sglang/test/kernels/kv_canary/_differential.py similarity index 99% rename from python/sglang/kernels/jit/tests/kv_canary/_differential.py rename to python/sglang/test/kernels/kv_canary/_differential.py index fd2c26148..b122a2e12 100644 --- a/python/sglang/kernels/jit/tests/kv_canary/_differential.py +++ b/python/sglang/test/kernels/kv_canary/_differential.py @@ -5,12 +5,6 @@ from typing import Any, Callable, Iterator, Optional import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - assert_canary_buf_equal, - assert_canary_state_equal, - make_log_pair, -) from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels from sglang.kernels.ops.kv_canary.plan_ref import ( @@ -30,6 +24,12 @@ from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_ke from sglang.kernels.ops.kv_canary.write_ref import ( launch_canary_write_kernel_torch_reference, ) +from sglang.test.kernels.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_log_pair, +) _DEVICE = torch.device("cuda") diff --git a/python/sglang/kernels/jit/tests/kv_canary/_fixtures.py b/python/sglang/test/kernels/kv_canary/_fixtures.py similarity index 98% rename from python/sglang/kernels/jit/tests/kv_canary/_fixtures.py rename to python/sglang/test/kernels/kv_canary/_fixtures.py index 70fbb7587..c3cd5bdb1 100644 --- a/python/sglang/kernels/jit/tests/kv_canary/_fixtures.py +++ b/python/sglang/test/kernels/kv_canary/_fixtures.py @@ -5,12 +5,12 @@ from typing import Literal, Optional import torch -from sglang.kernels.jit.tests.kv_canary._constants import DEFAULT_NUM_SLOTS from sglang.kernels.ops.kv_canary.verify import ( RealKvSource, VerifyPlan, ) from sglang.kernels.ops.kv_canary.write import WritePlan +from sglang.test.kernels.kv_canary._constants import DEFAULT_NUM_SLOTS _DEVICE = torch.device("cuda") diff --git a/python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py b/python/sglang/test/kernels/kv_canary/_fuzz_driver.py similarity index 95% rename from python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py rename to python/sglang/test/kernels/kv_canary/_fuzz_driver.py index 49919e068..566aae086 100644 --- a/python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py +++ b/python/sglang/test/kernels/kv_canary/_fuzz_driver.py @@ -3,7 +3,7 @@ from __future__ import annotations import random from typing import Any, Callable -from sglang.kernels.jit.tests.kv_canary._differential import ( +from sglang.test.kernels.kv_canary._differential import ( ShrinkResult, shrink_inputs, ) diff --git a/python/sglang/kernels/jit/tests/kv_canary/_hand_oracle.py b/python/sglang/test/kernels/kv_canary/_hand_oracle.py similarity index 100% rename from python/sglang/kernels/jit/tests/kv_canary/_hand_oracle.py rename to python/sglang/test/kernels/kv_canary/_hand_oracle.py diff --git a/python/sglang/kernels/jit/tests/kv_canary/_invariants.py b/python/sglang/test/kernels/kv_canary/_invariants.py similarity index 99% rename from python/sglang/kernels/jit/tests/kv_canary/_invariants.py rename to python/sglang/test/kernels/kv_canary/_invariants.py index 192c79152..588a242ba 100644 --- a/python/sglang/kernels/jit/tests/kv_canary/_invariants.py +++ b/python/sglang/test/kernels/kv_canary/_invariants.py @@ -11,10 +11,10 @@ from typing import Optional import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import FakeViolationLog from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan from sglang.kernels.ops.kv_canary.write import WritePlan +from sglang.test.kernels.kv_canary._canary_helpers import FakeViolationLog class PlanInvariants: diff --git a/python/sglang/kernels/jit/tests/utils.py b/python/sglang/test/kernels/utils.py similarity index 100% rename from python/sglang/kernels/jit/tests/utils.py rename to python/sglang/test/kernels/utils.py diff --git a/test/registered/jit/benchmark/bench_activation.py b/test/registered/kernels/benchmark/activation/bench_activation.py similarity index 100% rename from test/registered/jit/benchmark/bench_activation.py rename to test/registered/kernels/benchmark/activation/bench_activation.py diff --git a/test/registered/jit/benchmark/bench_add_constant.py b/test/registered/kernels/benchmark/attention/bench_add_constant.py similarity index 100% rename from test/registered/jit/benchmark/bench_add_constant.py rename to test/registered/kernels/benchmark/attention/bench_add_constant.py diff --git a/test/registered/jit/benchmark/bench_clamp_position.py b/test/registered/kernels/benchmark/attention/bench_clamp_position.py similarity index 100% rename from test/registered/jit/benchmark/bench_clamp_position.py rename to test/registered/kernels/benchmark/attention/bench_clamp_position.py diff --git a/test/registered/jit/benchmark/bench_concat_mla.py b/test/registered/kernels/benchmark/attention/bench_concat_mla.py similarity index 100% rename from test/registered/jit/benchmark/bench_concat_mla.py rename to test/registered/kernels/benchmark/attention/bench_concat_mla.py diff --git a/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py b/test/registered/kernels/benchmark/attention/bench_dsv4_fp4_indexer.py similarity index 100% rename from test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py rename to test/registered/kernels/benchmark/attention/bench_dsv4_fp4_indexer.py diff --git a/test/registered/jit/benchmark/bench_fused_qknorm_rope.py b/test/registered/kernels/benchmark/attention/bench_fused_qknorm_rope.py similarity index 100% rename from test/registered/jit/benchmark/bench_fused_qknorm_rope.py rename to test/registered/kernels/benchmark/attention/bench_fused_qknorm_rope.py diff --git a/test/registered/jit/benchmark/bench_hadamard.py b/test/registered/kernels/benchmark/attention/bench_hadamard.py similarity index 100% rename from test/registered/jit/benchmark/bench_hadamard.py rename to test/registered/kernels/benchmark/attention/bench_hadamard.py diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py b/test/registered/kernels/benchmark/attention/bench_minimax_decode_topk.py similarity index 100% rename from test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py rename to test/registered/kernels/benchmark/attention/bench_minimax_decode_topk.py diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py b/test/registered/kernels/benchmark/attention/bench_minimax_qknorm_rope.py similarity index 100% rename from test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py rename to test/registered/kernels/benchmark/attention/bench_minimax_qknorm_rope.py diff --git a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py b/test/registered/kernels/benchmark/attention/bench_mla_kv_pack_quantize_fp8.py similarity index 100% rename from test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py rename to test/registered/kernels/benchmark/attention/bench_mla_kv_pack_quantize_fp8.py diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/kernels/benchmark/attention/bench_online_c128_mtp.py similarity index 100% rename from test/registered/jit/benchmark/bench_online_c128_mtp.py rename to test/registered/kernels/benchmark/attention/bench_online_c128_mtp.py diff --git a/test/registered/jit/benchmark/bench_rope.py b/test/registered/kernels/benchmark/attention/bench_rope.py similarity index 100% rename from test/registered/jit/benchmark/bench_rope.py rename to test/registered/kernels/benchmark/attention/bench_rope.py diff --git a/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py b/test/registered/kernels/benchmark/attention/bench_sparse_mla_q8kv8_prefill_sm90.py similarity index 100% rename from test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py rename to test/registered/kernels/benchmark/attention/bench_sparse_mla_q8kv8_prefill_sm90.py diff --git a/test/registered/jit/benchmark/bench_topk.py b/test/registered/kernels/benchmark/attention/bench_topk.py similarity index 100% rename from test/registered/jit/benchmark/bench_topk.py rename to test/registered/kernels/benchmark/attention/bench_topk.py diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py similarity index 100% rename from test/registered/jit/benchmark/bench_custom_all_reduce.py rename to test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py diff --git a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py b/test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py similarity index 100% rename from test/registered/jit/benchmark/bench_symm_mem_all_gather.py rename to test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py similarity index 100% rename from test/registered/jit/benchmark/bench_tp_qknorm.py rename to test/registered/kernels/benchmark/communication/bench_tp_qknorm.py diff --git a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py b/test/registered/kernels/benchmark/diffusion/bench_causal_conv3d_cat_pad.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py rename to test/registered/kernels/benchmark/diffusion/bench_causal_conv3d_cat_pad.py diff --git a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py b/test/registered/kernels/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py rename to test/registered/kernels/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py diff --git a/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py b/test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py rename to test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py diff --git a/test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py b/test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py rename to test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py diff --git a/test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py b/test/registered/kernels/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py rename to test/registered/kernels/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py diff --git a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py b/test/registered/kernels/benchmark/diffusion/bench_norm_impls.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_norm_impls.py rename to test/registered/kernels/benchmark/diffusion/bench_norm_impls.py diff --git a/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py b/test/registered/kernels/benchmark/diffusion/bench_qknorm_rope.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py rename to test/registered/kernels/benchmark/diffusion/bench_qknorm_rope.py diff --git a/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py rename to test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py diff --git a/test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py b/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py similarity index 100% rename from test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py rename to test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py diff --git a/test/registered/jit/benchmark/diffusion/diffusion_nvfp4_shapes.json b/test/registered/kernels/benchmark/diffusion/diffusion_nvfp4_shapes.json similarity index 100% rename from test/registered/jit/benchmark/diffusion/diffusion_nvfp4_shapes.json rename to test/registered/kernels/benchmark/diffusion/diffusion_nvfp4_shapes.json diff --git a/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py b/test/registered/kernels/benchmark/embeddings/bench_vocab_parallel_embedding.py similarity index 100% rename from test/registered/jit/benchmark/bench_vocab_parallel_embedding.py rename to test/registered/kernels/benchmark/embeddings/bench_vocab_parallel_embedding.py diff --git a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py b/test/registered/kernels/benchmark/gemm/bench_dsv3_fused_a_gemm.py similarity index 100% rename from test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py rename to test/registered/kernels/benchmark/gemm/bench_dsv3_fused_a_gemm.py diff --git a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py b/test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py similarity index 100% rename from test/registered/jit/benchmark/bench_dsv3_router_gemm.py rename to test/registered/kernels/benchmark/gemm/bench_dsv3_router_gemm.py diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/kernels/benchmark/gemm/bench_fp8_blockwise_gemm.py similarity index 100% rename from test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py rename to test/registered/kernels/benchmark/gemm/bench_fp8_blockwise_gemm.py diff --git a/test/registered/jit/benchmark/kv_canary/bench_plan.py b/test/registered/kernels/benchmark/kv_canary/bench_plan.py similarity index 100% rename from test/registered/jit/benchmark/kv_canary/bench_plan.py rename to test/registered/kernels/benchmark/kv_canary/bench_plan.py diff --git a/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py b/test/registered/kernels/benchmark/kv_canary/bench_scatter_req_token_ids.py similarity index 100% rename from test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py rename to test/registered/kernels/benchmark/kv_canary/bench_scatter_req_token_ids.py diff --git a/test/registered/jit/benchmark/kv_canary/bench_verify.py b/test/registered/kernels/benchmark/kv_canary/bench_verify.py similarity index 100% rename from test/registered/jit/benchmark/kv_canary/bench_verify.py rename to test/registered/kernels/benchmark/kv_canary/bench_verify.py diff --git a/test/registered/jit/benchmark/kv_canary/bench_write.py b/test/registered/kernels/benchmark/kv_canary/bench_write.py similarity index 100% rename from test/registered/jit/benchmark/kv_canary/bench_write.py rename to test/registered/kernels/benchmark/kv_canary/bench_write.py diff --git a/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py b/test/registered/kernels/benchmark/kvcache/bench_fused_fp8_qkv_kv_cache.py similarity index 100% rename from test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py rename to test/registered/kernels/benchmark/kvcache/bench_fused_fp8_qkv_kv_cache.py diff --git a/test/registered/jit/benchmark/bench_hicache.py b/test/registered/kernels/benchmark/kvcache/bench_hicache.py similarity index 100% rename from test/registered/jit/benchmark/bench_hicache.py rename to test/registered/kernels/benchmark/kvcache/bench_hicache.py diff --git a/test/registered/jit/benchmark/bench_hisparse.py b/test/registered/kernels/benchmark/kvcache/bench_hisparse.py similarity index 100% rename from test/registered/jit/benchmark/bench_hisparse.py rename to test/registered/kernels/benchmark/kvcache/bench_hisparse.py diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py b/test/registered/kernels/benchmark/kvcache/bench_minimax_store_kv_index.py similarity index 100% rename from test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py rename to test/registered/kernels/benchmark/kvcache/bench_minimax_store_kv_index.py diff --git a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py b/test/registered/kernels/benchmark/kvcache/bench_set_mla_kv_buffer.py similarity index 100% rename from test/registered/jit/benchmark/bench_set_mla_kv_buffer.py rename to test/registered/kernels/benchmark/kvcache/bench_set_mla_kv_buffer.py diff --git a/test/registered/jit/benchmark/bench_store_cache.py b/test/registered/kernels/benchmark/kvcache/bench_store_cache.py similarity index 100% rename from test/registered/jit/benchmark/bench_store_cache.py rename to test/registered/kernels/benchmark/kvcache/bench_store_cache.py diff --git a/test/registered/jit/benchmark/bench_fused_eh_norm.py b/test/registered/kernels/benchmark/layernorm/bench_fused_eh_norm.py similarity index 100% rename from test/registered/jit/benchmark/bench_fused_eh_norm.py rename to test/registered/kernels/benchmark/layernorm/bench_fused_eh_norm.py diff --git a/test/registered/jit/benchmark/bench_norm.py b/test/registered/kernels/benchmark/layernorm/bench_norm.py similarity index 100% rename from test/registered/jit/benchmark/bench_norm.py rename to test/registered/kernels/benchmark/layernorm/bench_norm.py diff --git a/test/registered/jit/benchmark/bench_qknorm.py b/test/registered/kernels/benchmark/layernorm/bench_qknorm.py similarity index 100% rename from test/registered/jit/benchmark/bench_qknorm.py rename to test/registered/kernels/benchmark/layernorm/bench_qknorm.py diff --git a/test/registered/jit/benchmark/bench_qknorm_across_heads.py b/test/registered/kernels/benchmark/layernorm/bench_qknorm_across_heads.py similarity index 100% rename from test/registered/jit/benchmark/bench_qknorm_across_heads.py rename to test/registered/kernels/benchmark/layernorm/bench_qknorm_across_heads.py diff --git a/test/registered/jit/benchmark/bench_moe_fused_gate.py b/test/registered/kernels/benchmark/moe/bench_moe_fused_gate.py similarity index 100% rename from test/registered/jit/benchmark/bench_moe_fused_gate.py rename to test/registered/kernels/benchmark/moe/bench_moe_fused_gate.py diff --git a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py b/test/registered/kernels/benchmark/moe/bench_post_reorder_deepgemm.py similarity index 100% rename from test/registered/jit/benchmark/bench_post_reorder_deepgemm.py rename to test/registered/kernels/benchmark/moe/bench_post_reorder_deepgemm.py diff --git a/test/registered/jit/benchmark/bench_renorm.py b/test/registered/kernels/benchmark/moe/bench_renorm.py similarity index 100% rename from test/registered/jit/benchmark/bench_renorm.py rename to test/registered/kernels/benchmark/moe/bench_renorm.py diff --git a/test/registered/jit/benchmark/bench_awq_dequantize.py b/test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py similarity index 100% rename from test/registered/jit/benchmark/bench_awq_dequantize.py rename to test/registered/kernels/benchmark/quantization/bench_awq_dequantize.py diff --git a/test/registered/jit/benchmark/bench_mxfp8_moe.py b/test/registered/kernels/benchmark/quantization/bench_mxfp8_moe.py similarity index 100% rename from test/registered/jit/benchmark/bench_mxfp8_moe.py rename to test/registered/kernels/benchmark/quantization/bench_mxfp8_moe.py diff --git a/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py b/test/registered/kernels/benchmark/quantization/bench_per_tensor_quant_fp8.py similarity index 100% rename from test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py rename to test/registered/kernels/benchmark/quantization/bench_per_tensor_quant_fp8.py diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant.py b/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py similarity index 100% rename from test/registered/jit/benchmark/bench_per_token_group_quant.py rename to test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py b/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_8bit_v2.py similarity index 100% rename from test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py rename to test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_8bit_v2.py diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant_masked.py b/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_masked.py similarity index 100% rename from test/registered/jit/benchmark/bench_per_token_group_quant_masked.py rename to test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_masked.py diff --git a/test/registered/jit/benchmark/bench_ngram_compute_decode.py b/test/registered/kernels/benchmark/speculative/bench_ngram_compute_decode.py similarity index 100% rename from test/registered/jit/benchmark/bench_ngram_compute_decode.py rename to test/registered/kernels/benchmark/speculative/bench_ngram_compute_decode.py diff --git a/test/registered/jit/benchmark/bench_ngram_update_token_table.py b/test/registered/kernels/benchmark/speculative/bench_ngram_update_token_table.py similarity index 100% rename from test/registered/jit/benchmark/bench_ngram_update_token_table.py rename to test/registered/kernels/benchmark/speculative/bench_ngram_update_token_table.py diff --git a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py b/test/registered/kernels/benchmark/speculative/bench_resolve_future_token_ids.py similarity index 100% rename from test/registered/jit/benchmark/bench_resolve_future_token_ids.py rename to test/registered/kernels/benchmark/speculative/bench_resolve_future_token_ids.py diff --git a/test/registered/jit/benchmark/bench_spec_topk1.py b/test/registered/kernels/benchmark/speculative/bench_spec_topk1.py similarity index 100% rename from test/registered/jit/benchmark/bench_spec_topk1.py rename to test/registered/kernels/benchmark/speculative/bench_spec_topk1.py diff --git a/test/registered/jit/test_activation.py b/test/registered/kernels/ops/activation/test_activation.py similarity index 100% rename from test/registered/jit/test_activation.py rename to test/registered/kernels/ops/activation/test_activation.py diff --git a/test/registered/jit/test_add_constant.py b/test/registered/kernels/ops/attention/test_add_constant.py similarity index 100% rename from test/registered/jit/test_add_constant.py rename to test/registered/kernels/ops/attention/test_add_constant.py diff --git a/test/registered/jit/deepseek_v4/test_c128_v2.py b/test/registered/kernels/ops/attention/test_c128_v2.py similarity index 99% rename from test/registered/jit/deepseek_v4/test_c128_v2.py rename to test/registered/kernels/ops/attention/test_c128_v2.py index 373fd2a24..a5b0e74fa 100644 --- a/test/registered/jit/deepseek_v4/test_c128_v2.py +++ b/test/registered/kernels/ops/attention/test_c128_v2.py @@ -7,7 +7,10 @@ import pytest import torch import triton -from sglang.kernels.jit.tests.deepseek_v4.common import ( +from sglang.kernels.ops.attention.dsv4 import compress_forward +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.deepseek_v4.common import ( LegacyContext, PagedContext, make_legacy_context, @@ -15,9 +18,6 @@ from sglang.kernels.jit.tests.deepseek_v4.common import ( make_state_pool, to_seq_extend, ) -from sglang.kernels.ops.attention.dsv4 import compress_forward -from sglang.srt.utils import get_device -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=30, suite="nightly-amd-kernel-1-gpu", nightly=True) diff --git a/test/registered/jit/deepseek_v4/test_c4_v2.py b/test/registered/kernels/ops/attention/test_c4_v2.py similarity index 99% rename from test/registered/jit/deepseek_v4/test_c4_v2.py rename to test/registered/kernels/ops/attention/test_c4_v2.py index a9a7eb9d4..f05aaa903 100644 --- a/test/registered/jit/deepseek_v4/test_c4_v2.py +++ b/test/registered/kernels/ops/attention/test_c4_v2.py @@ -7,7 +7,10 @@ import pytest import torch import triton -from sglang.kernels.jit.tests.deepseek_v4.common import ( +from sglang.kernels.ops.attention.dsv4 import compress_forward +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.deepseek_v4.common import ( LegacyContext, PagedContext, make_legacy_context, @@ -15,9 +18,6 @@ from sglang.kernels.jit.tests.deepseek_v4.common import ( make_state_pool, to_seq_extend, ) -from sglang.kernels.ops.attention.dsv4 import compress_forward -from sglang.srt.utils import get_device -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=30, suite="nightly-amd-kernel-1-gpu", nightly=True) diff --git a/test/registered/jit/test_clamp_position.py b/test/registered/kernels/ops/attention/test_clamp_position.py similarity index 100% rename from test/registered/jit/test_clamp_position.py rename to test/registered/kernels/ops/attention/test_clamp_position.py diff --git a/test/registered/jit/test_concat_mla.py b/test/registered/kernels/ops/attention/test_concat_mla.py similarity index 100% rename from test/registered/jit/test_concat_mla.py rename to test/registered/kernels/ops/attention/test_concat_mla.py diff --git a/test/registered/kernels/test_cp_prefix_len_fa3_parity.py b/test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py similarity index 100% rename from test/registered/kernels/test_cp_prefix_len_fa3_parity.py rename to test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py diff --git a/test/registered/kernels/test_cute_dsl_fp8_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py similarity index 100% rename from test/registered/kernels/test_cute_dsl_fp8_paged_mqa_logits.py rename to test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py diff --git a/test/registered/jit/test_cutedsl_gdn.py b/test/registered/kernels/ops/attention/test_cutedsl_gdn.py similarity index 100% rename from test/registered/jit/test_cutedsl_gdn.py rename to test/registered/kernels/ops/attention/test_cutedsl_gdn.py diff --git a/test/registered/kernels/test_deepgemm_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py similarity index 100% rename from test/registered/kernels/test_deepgemm_paged_mqa_logits.py rename to test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py diff --git a/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py b/test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py similarity index 99% rename from test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py rename to test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py index 2658de98a..9a2f392a0 100644 --- a/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py +++ b/test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py @@ -89,16 +89,16 @@ from typing import Callable, Iterable, Literal, Optional import pytest import torch -from sglang.kernels.jit.tests.deepseek_v4.common import ( - make_legacy_context, - to_seq_extend, -) from sglang.kernels.ops.attention.dsv4 import ( CompressorDecodePlan, CompressorPrefillPlan, compress_forward, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.deepseek_v4.common import ( + make_legacy_context, + to_seq_extend, +) from sglang.utils import is_in_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/kernels/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py similarity index 100% rename from test/registered/kernels/test_dsa_indexer.py rename to test/registered/kernels/ops/attention/test_dsa_indexer.py diff --git a/test/registered/kernels/test_dsa_metadata.py b/test/registered/kernels/ops/attention/test_dsa_metadata.py similarity index 100% rename from test/registered/kernels/test_dsa_metadata.py rename to test/registered/kernels/ops/attention/test_dsa_metadata.py diff --git a/test/registered/kernels/test_dsa_transform_index.py b/test/registered/kernels/ops/attention/test_dsa_transform_index.py similarity index 100% rename from test/registered/kernels/test_dsa_transform_index.py rename to test/registered/kernels/ops/attention/test_dsa_transform_index.py diff --git a/test/registered/jit/test_dsv32_indexer_fusion.py b/test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py similarity index 100% rename from test/registered/jit/test_dsv32_indexer_fusion.py rename to test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py diff --git a/test/registered/jit/test_flash_attention_3_only_qv.py b/test/registered/kernels/ops/attention/test_flash_attention_3_only_qv.py similarity index 100% rename from test/registered/jit/test_flash_attention_3_only_qv.py rename to test/registered/kernels/ops/attention/test_flash_attention_3_only_qv.py diff --git a/test/registered/jit/test_flash_attention_4.py b/test/registered/kernels/ops/attention/test_flash_attention_4.py similarity index 100% rename from test/registered/jit/test_flash_attention_4.py rename to test/registered/kernels/ops/attention/test_flash_attention_4.py diff --git a/test/registered/kernels/test_flash_mla_backends.py b/test/registered/kernels/ops/attention/test_flash_mla_backends.py similarity index 100% rename from test/registered/kernels/test_flash_mla_backends.py rename to test/registered/kernels/ops/attention/test_flash_mla_backends.py diff --git a/test/registered/jit/deepseek_v4/test_fp4_indexer.py b/test/registered/kernels/ops/attention/test_fp4_indexer.py similarity index 100% rename from test/registered/jit/deepseek_v4/test_fp4_indexer.py rename to test/registered/kernels/ops/attention/test_fp4_indexer.py diff --git a/test/registered/jit/deepseek_v4/test_fp8_wo_a.py b/test/registered/kernels/ops/attention/test_fp8_wo_a.py similarity index 100% rename from test/registered/jit/deepseek_v4/test_fp8_wo_a.py rename to test/registered/kernels/ops/attention/test_fp8_wo_a.py diff --git a/test/registered/jit/test_fused_metadata_copy.py b/test/registered/kernels/ops/attention/test_fused_metadata_copy.py similarity index 100% rename from test/registered/jit/test_fused_metadata_copy.py rename to test/registered/kernels/ops/attention/test_fused_metadata_copy.py diff --git a/test/registered/jit/test_fused_store_index_cache.py b/test/registered/kernels/ops/attention/test_fused_store_index_cache.py similarity index 100% rename from test/registered/jit/test_fused_store_index_cache.py rename to test/registered/kernels/ops/attention/test_fused_store_index_cache.py diff --git a/test/registered/jit/test_fused_verify_triton_gdn.py b/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py similarity index 100% rename from test/registered/jit/test_fused_verify_triton_gdn.py rename to test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py diff --git a/test/registered/jit/test_hadamard_jit.py b/test/registered/kernels/ops/attention/test_hadamard_jit.py similarity index 100% rename from test/registered/jit/test_hadamard_jit.py rename to test/registered/kernels/ops/attention/test_hadamard_jit.py diff --git a/test/registered/jit/test_inkling_attn_prologue_tau.py b/test/registered/kernels/ops/attention/test_inkling_attn_prologue_tau.py similarity index 100% rename from test/registered/jit/test_inkling_attn_prologue_tau.py rename to test/registered/kernels/ops/attention/test_inkling_attn_prologue_tau.py diff --git a/test/registered/jit/test_inkling_row_scale.py b/test/registered/kernels/ops/attention/test_inkling_row_scale.py similarity index 100% rename from test/registered/jit/test_inkling_row_scale.py rename to test/registered/kernels/ops/attention/test_inkling_row_scale.py diff --git a/test/registered/jit/minimax/test_minimax_decode_topk.py b/test/registered/kernels/ops/attention/test_minimax_decode_topk.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_decode_topk.py rename to test/registered/kernels/ops/attention/test_minimax_decode_topk.py diff --git a/test/registered/jit/minimax/test_minimax_decode_topk_page_table.py b/test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_decode_topk_page_table.py rename to test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py diff --git a/test/registered/jit/test_minimax_m3_qk_norm_rope.py b/test/registered/kernels/ops/attention/test_minimax_m3_qk_norm_rope.py similarity index 100% rename from test/registered/jit/test_minimax_m3_qk_norm_rope.py rename to test/registered/kernels/ops/attention/test_minimax_m3_qk_norm_rope.py diff --git a/test/registered/jit/minimax/test_minimax_qknorm_rope.py b/test/registered/kernels/ops/attention/test_minimax_qknorm_rope.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_qknorm_rope.py rename to test/registered/kernels/ops/attention/test_minimax_qknorm_rope.py diff --git a/test/registered/kernels/test_mla_cp_fa3_parity.py b/test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py similarity index 100% rename from test/registered/kernels/test_mla_cp_fa3_parity.py rename to test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py diff --git a/test/registered/jit/test_mla_kv_pack_quantize_fp8.py b/test/registered/kernels/ops/attention/test_mla_kv_pack_quantize_fp8.py similarity index 100% rename from test/registered/jit/test_mla_kv_pack_quantize_fp8.py rename to test/registered/kernels/ops/attention/test_mla_kv_pack_quantize_fp8.py diff --git a/test/registered/jit/test_pos_enc.py b/test/registered/kernels/ops/attention/test_pos_enc.py similarity index 100% rename from test/registered/jit/test_pos_enc.py rename to test/registered/kernels/ops/attention/test_pos_enc.py diff --git a/test/registered/jit/test_rope.py b/test/registered/kernels/ops/attention/test_rope.py similarity index 100% rename from test/registered/jit/test_rope.py rename to test/registered/kernels/ops/attention/test_rope.py diff --git a/test/registered/kernels/test_sm120_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py similarity index 100% rename from test/registered/kernels/test_sm120_paged_mqa_logits.py rename to test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py diff --git a/test/registered/jit/test_sparse_mla_q8kv8_prefill_sm90.py b/test/registered/kernels/ops/attention/test_sparse_mla_q8kv8_prefill_sm90.py similarity index 100% rename from test/registered/jit/test_sparse_mla_q8kv8_prefill_sm90.py rename to test/registered/kernels/ops/attention/test_sparse_mla_q8kv8_prefill_sm90.py diff --git a/test/registered/jit/deepseek_v4/test_topk_v2.py b/test/registered/kernels/ops/attention/test_topk_v2.py similarity index 100% rename from test/registered/jit/deepseek_v4/test_topk_v2.py rename to test/registered/kernels/ops/attention/test_topk_v2.py diff --git a/test/registered/jit/test_amd_deterministic_custom_allreduce.py b/test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py similarity index 100% rename from test/registered/jit/test_amd_deterministic_custom_allreduce.py rename to test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py diff --git a/test/registered/jit/test_amd_nccl_allreduce_determinism.py b/test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py similarity index 100% rename from test/registered/jit/test_amd_nccl_allreduce_determinism.py rename to test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/kernels/ops/communication/test_custom_all_reduce.py similarity index 99% rename from test/registered/jit/test_custom_all_reduce.py rename to test/registered/kernels/ops/communication/test_custom_all_reduce.py index 38b2b0568..894b16334 100644 --- a/test/registered/jit/test_custom_all_reduce.py +++ b/test/registered/kernels/ops/communication/test_custom_all_reduce.py @@ -28,7 +28,6 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.kernels.jit.tests.utils import multigpu_pytest_main from sglang.kernels.jit.utils import cache_once, get_ci_test_range from sglang.kernels.ops.communication.all_reduce import ( AllReduceAlgo, @@ -39,6 +38,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main register_cuda_ci( est_time=300, diff --git a/test/registered/jit/test_symm_mem_all_gather.py b/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py similarity index 98% rename from test/registered/jit/test_symm_mem_all_gather.py rename to test/registered/kernels/ops/communication/test_symm_mem_all_gather.py index 816964d9b..bc73dde3d 100644 --- a/test/registered/jit/test_symm_mem_all_gather.py +++ b/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py @@ -26,13 +26,13 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.kernels.jit.tests.utils import multigpu_pytest_main from sglang.kernels.jit.utils import cache_once, get_ci_test_range from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, create_state, ) from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main register_cuda_ci(est_time=240, stage="extra-b", runner_config="8-gpu-h200") # Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps. diff --git a/test/registered/jit/test_tp_qknorm.py b/test/registered/kernels/ops/communication/test_tp_qknorm.py similarity index 99% rename from test/registered/jit/test_tp_qknorm.py rename to test/registered/kernels/ops/communication/test_tp_qknorm.py index 31f715eb0..712c24ecd 100644 --- a/test/registered/jit/test_tp_qknorm.py +++ b/test/registered/kernels/ops/communication/test_tp_qknorm.py @@ -14,7 +14,6 @@ import torch.distributed as dist import triton import sglang.srt.distributed.parallel_state as ps -from sglang.kernels.jit.tests.utils import multigpu_pytest_main from sglang.kernels.jit.utils import cache_once from sglang.kernels.ops.communication.all_reduce import ( fused_parallel_qknorm, @@ -26,6 +25,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main register_cuda_ci( est_time=300, diff --git a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py b/test/registered/kernels/ops/diffusion/test_causal_conv3d_cat_pad.py similarity index 100% rename from test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py rename to test/registered/kernels/ops/diffusion/test_causal_conv3d_cat_pad.py diff --git a/test/registered/jit/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py b/test/registered/kernels/ops/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py similarity index 100% rename from test/registered/jit/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py rename to test/registered/kernels/ops/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py diff --git a/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py b/test/registered/kernels/ops/diffusion/test_diffusion_nvfp4_scaled_mm.py similarity index 100% rename from test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py rename to test/registered/kernels/ops/diffusion/test_diffusion_nvfp4_scaled_mm.py diff --git a/test/registered/jit/diffusion/test_flydsl_fused_norm.py b/test/registered/kernels/ops/diffusion/test_flydsl_fused_norm.py similarity index 100% rename from test/registered/jit/diffusion/test_flydsl_fused_norm.py rename to test/registered/kernels/ops/diffusion/test_flydsl_fused_norm.py diff --git a/test/registered/jit/diffusion/test_fused_norm_scale_shift.py b/test/registered/kernels/ops/diffusion/test_fused_norm_scale_shift.py similarity index 100% rename from test/registered/jit/diffusion/test_fused_norm_scale_shift.py rename to test/registered/kernels/ops/diffusion/test_fused_norm_scale_shift.py diff --git a/test/registered/jit/diffusion/test_group_norm_silu.py b/test/registered/kernels/ops/diffusion/test_group_norm_silu.py similarity index 100% rename from test/registered/jit/diffusion/test_group_norm_silu.py rename to test/registered/kernels/ops/diffusion/test_group_norm_silu.py diff --git a/test/registered/jit/diffusion/test_ltx2_ada_values.py b/test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py similarity index 100% rename from test/registered/jit/diffusion/test_ltx2_ada_values.py rename to test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py diff --git a/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py b/test/registered/kernels/ops/diffusion/test_ltx2_qknorm_split_rope.py similarity index 100% rename from test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py rename to test/registered/kernels/ops/diffusion/test_ltx2_qknorm_split_rope.py diff --git a/test/registered/jit/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py similarity index 100% rename from test/registered/jit/diffusion/test_qknorm_rope.py rename to test/registered/kernels/ops/diffusion/test_qknorm_rope.py diff --git a/test/registered/jit/diffusion/test_qwen_image_modulation.py b/test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py similarity index 100% rename from test/registered/jit/diffusion/test_qwen_image_modulation.py rename to test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py diff --git a/test/registered/jit/diffusion/test_residual_gate_add.py b/test/registered/kernels/ops/diffusion/test_residual_gate_add.py similarity index 100% rename from test/registered/jit/diffusion/test_residual_gate_add.py rename to test/registered/kernels/ops/diffusion/test_residual_gate_add.py diff --git a/test/registered/jit/test_timestep_embedding.py b/test/registered/kernels/ops/diffusion/test_timestep_embedding.py similarity index 100% rename from test/registered/jit/test_timestep_embedding.py rename to test/registered/kernels/ops/diffusion/test_timestep_embedding.py diff --git a/test/registered/jit/diffusion/test_varlen_pack_pad.py b/test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py similarity index 100% rename from test/registered/jit/diffusion/test_varlen_pack_pad.py rename to test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py diff --git a/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py b/test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py similarity index 100% rename from test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py rename to test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py diff --git a/test/registered/kernels/test_vocab_parallel_embedding.py b/test/registered/kernels/ops/embeddings/test_vocab_parallel_embedding.py similarity index 100% rename from test/registered/kernels/test_vocab_parallel_embedding.py rename to test/registered/kernels/ops/embeddings/test_vocab_parallel_embedding.py diff --git a/test/registered/jit/test_cutedsl_bf16_gemm.py b/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py similarity index 100% rename from test/registered/jit/test_cutedsl_bf16_gemm.py rename to test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py diff --git a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py b/test/registered/kernels/ops/gemm/test_cutedsl_dsv3_fused_a_gemm.py similarity index 100% rename from test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py rename to test/registered/kernels/ops/gemm/test_cutedsl_dsv3_fused_a_gemm.py diff --git a/test/registered/jit/test_dsv3_fused_a_gemm.py b/test/registered/kernels/ops/gemm/test_dsv3_fused_a_gemm.py similarity index 100% rename from test/registered/jit/test_dsv3_fused_a_gemm.py rename to test/registered/kernels/ops/gemm/test_dsv3_fused_a_gemm.py diff --git a/test/registered/jit/test_dsv3_router_gemm.py b/test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py similarity index 100% rename from test/registered/jit/test_dsv3_router_gemm.py rename to test/registered/kernels/ops/gemm/test_dsv3_router_gemm.py diff --git a/test/registered/jit/test_fp8_blockwise_gemm.py b/test/registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py similarity index 100% rename from test/registered/jit/test_fp8_blockwise_gemm.py rename to test/registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py diff --git a/test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py b/test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py rename to test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py diff --git a/test/registered/jit/kv_canary/test_const_sync.py b/test/registered/kernels/ops/kv_canary/test_const_sync.py similarity index 100% rename from test/registered/jit/kv_canary/test_const_sync.py rename to test/registered/kernels/ops/kv_canary/test_const_sync.py diff --git a/test/registered/jit/kv_canary/test_kernel_config.py b/test/registered/kernels/ops/kv_canary/test_kernel_config.py similarity index 98% rename from test/registered/jit/kv_canary/test_kernel_config.py rename to test/registered/kernels/ops/kv_canary/test_kernel_config.py index 2917da076..e3d7b6d6c 100644 --- a/test/registered/jit/kv_canary/test_kernel_config.py +++ b/test/registered/kernels/ops/kv_canary/test_kernel_config.py @@ -3,7 +3,16 @@ from __future__ import annotations import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary.verify import ( + CanaryLaunchTag, + VerifyOrWriteContext, + VerifyPlan, + launch_canary_verify_kernel, +) +from sglang.kernels.ops.kv_canary.write import WritePlan +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_buf_equal, assert_canary_state_equal, @@ -15,25 +24,16 @@ from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( make_write_plan_pair, stamp_clean_chain, ) -from sglang.kernels.jit.tests.kv_canary._differential import ( +from sglang.test.kernels.kv_canary._differential import ( _assert_plans_byte_equal, _run_both_plan, _run_both_verify, _run_both_write, ) -from sglang.kernels.jit.tests.kv_canary._fixtures import ( +from sglang.test.kernels.kv_canary._fixtures import ( dummy_pseudo_tensors, empty_extras, ) -from sglang.kernels.ops.kv_canary import consts -from sglang.kernels.ops.kv_canary.verify import ( - CanaryLaunchTag, - VerifyOrWriteContext, - VerifyPlan, - launch_canary_verify_kernel, -) -from sglang.kernels.ops.kv_canary.write import WritePlan -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=60, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_pipeline_e2e.py b/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py similarity index 99% rename from test/registered/jit/kv_canary/test_pipeline_e2e.py rename to test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py index 24fa49b2a..975495fcd 100644 --- a/test/registered/jit/kv_canary/test_pipeline_e2e.py +++ b/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py @@ -5,20 +5,6 @@ from typing import Any, Optional import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - assert_canary_buf_equal, - assert_canary_state_equal, - make_canary_buf, - make_real_kv_sources, - stamp_clean_chain, - write_slot_fields, -) -from sglang.kernels.jit.tests.kv_canary._fixtures import ( - clone_real_kv_sources, - empty_extras, - make_req_to_token, -) from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels from sglang.kernels.ops.kv_canary.plan_ref import ( @@ -39,6 +25,20 @@ from sglang.kernels.ops.kv_canary.write_ref import ( launch_canary_write_kernel_torch_reference, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_canary_buf_equal, + assert_canary_state_equal, + make_canary_buf, + make_real_kv_sources, + stamp_clean_chain, + write_slot_fields, +) +from sglang.test.kernels.kv_canary._fixtures import ( + clone_real_kv_sources, + empty_extras, + make_req_to_token, +) register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_plan_fuzz.py b/test/registered/kernels/ops/kv_canary/test_plan_fuzz.py similarity index 96% rename from test/registered/jit/kv_canary/test_plan_fuzz.py rename to test/registered/kernels/ops/kv_canary/test_plan_fuzz.py index 25e3df6fc..091f406fd 100644 --- a/test/registered/jit/kv_canary/test_plan_fuzz.py +++ b/test/registered/kernels/ops/kv_canary/test_plan_fuzz.py @@ -7,20 +7,20 @@ from typing import Optional import pytest import torch -from sglang.kernels.jit.tests.kv_canary._differential import _run_both_plan -from sglang.kernels.jit.tests.kv_canary._fixtures import ( +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._differential import _run_both_plan +from sglang.test.kernels.kv_canary._fixtures import ( allocate_plan_pair, derive_plan_capacity, make_lut, make_padding_mask, make_req_to_token, ) -from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( +from sglang.test.kernels.kv_canary._fuzz_driver import ( FUZZ_SEEDS_PR, run_fuzz_combo, ) -from sglang.kernels.jit.tests.kv_canary._invariants import PlanInvariants -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._invariants import PlanInvariants register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_plan_hand.py b/test/registered/kernels/ops/kv_canary/test_plan_hand.py similarity index 99% rename from test/registered/jit/kv_canary/test_plan_hand.py rename to test/registered/kernels/ops/kv_canary/test_plan_hand.py index bbf543cec..487f26d80 100644 --- a/test/registered/jit/kv_canary/test_plan_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_plan_hand.py @@ -5,15 +5,6 @@ import random import pytest import torch -from sglang.kernels.jit.tests.kv_canary._differential import run_plan_diff -from sglang.kernels.jit.tests.kv_canary._fixtures import ( - allocate_plan_pair, - derive_plan_capacity, - empty_extras, - make_lut, - make_req_to_token, -) -from sglang.kernels.jit.tests.kv_canary._invariants import PlanInvariants from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels from sglang.kernels.ops.kv_canary.plan_ref import ( launch_canary_plan_kernels_torch_reference, @@ -21,6 +12,15 @@ from sglang.kernels.ops.kv_canary.plan_ref import ( from sglang.kernels.ops.kv_canary.verify import VerifyPlan from sglang.kernels.ops.kv_canary.write import WritePlan from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._differential import run_plan_diff +from sglang.test.kernels.kv_canary._fixtures import ( + allocate_plan_pair, + derive_plan_capacity, + empty_extras, + make_lut, + make_req_to_token, +) +from sglang.test.kernels.kv_canary._invariants import PlanInvariants register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_scatter_req_token_ids.py b/test/registered/kernels/ops/kv_canary/test_scatter_req_token_ids.py similarity index 100% rename from test/registered/jit/kv_canary/test_scatter_req_token_ids.py rename to test/registered/kernels/ops/kv_canary/test_scatter_req_token_ids.py diff --git a/test/registered/jit/kv_canary/test_utils.py b/test/registered/kernels/ops/kv_canary/test_utils.py similarity index 100% rename from test/registered/jit/kv_canary/test_utils.py rename to test/registered/kernels/ops/kv_canary/test_utils.py diff --git a/test/registered/jit/kv_canary/test_verify_fuzz.py b/test/registered/kernels/ops/kv_canary/test_verify_fuzz.py similarity index 95% rename from test/registered/jit/kv_canary/test_verify_fuzz.py rename to test/registered/kernels/ops/kv_canary/test_verify_fuzz.py index 56fa62eba..b1405a907 100644 --- a/test/registered/jit/kv_canary/test_verify_fuzz.py +++ b/test/registered/kernels/ops/kv_canary/test_verify_fuzz.py @@ -6,23 +6,6 @@ from dataclasses import dataclass import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - make_canary_buf, - make_log_pair, - make_verify_plan_pair, - stamp_clean_chain, -) -from sglang.kernels.jit.tests.kv_canary._differential import _run_both_verify -from sglang.kernels.jit.tests.kv_canary._fixtures import ( - clone_real_kv_sources, - make_real_kv_sources, -) -from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( - FUZZ_SEEDS_PR, - run_fuzz_combo, -) -from sglang.kernels.jit.tests.kv_canary._invariants import VerifyInvariants from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.verify import ( CanaryLaunchTag, @@ -30,6 +13,23 @@ from sglang.kernels.ops.kv_canary.verify import ( VerifyPlan, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( + FakeViolationLog, + make_canary_buf, + make_log_pair, + make_verify_plan_pair, + stamp_clean_chain, +) +from sglang.test.kernels.kv_canary._differential import _run_both_verify +from sglang.test.kernels.kv_canary._fixtures import ( + clone_real_kv_sources, + make_real_kv_sources, +) +from sglang.test.kernels.kv_canary._fuzz_driver import ( + FUZZ_SEEDS_PR, + run_fuzz_combo, +) +from sglang.test.kernels.kv_canary._invariants import VerifyInvariants register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_verify_hand.py b/test/registered/kernels/ops/kv_canary/test_verify_hand.py similarity index 99% rename from test/registered/jit/kv_canary/test_verify_hand.py rename to test/registered/kernels/ops/kv_canary/test_verify_hand.py index ea68821e8..266b7ae61 100644 --- a/test/registered/jit/kv_canary/test_verify_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_verify_hand.py @@ -8,33 +8,6 @@ from typing import Callable import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - assert_only_bits_set, - chain_anchor_signed, - make_canary_buf, - make_canary_buf_pair, - make_log_pair, - make_real_kv_source, - make_real_kv_sources, - make_verify_plan, - make_verify_plan_pair, - make_write_plan, - read_slot_fields, - stamp_clean_chain, - stamp_pair, - to_signed_int64, - write_slot_fields, -) -from sglang.kernels.jit.tests.kv_canary._differential import ( - _run_both_verify, - run_verify_diff, -) -from sglang.kernels.jit.tests.kv_canary._fixtures import clone_real_kv_sources -from sglang.kernels.jit.tests.kv_canary._hand_oracle import ( - _hand_fold_all, - _hand_fold_partial, -) from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 from sglang.kernels.ops.kv_canary.verify import ( @@ -52,6 +25,33 @@ from sglang.kernels.ops.kv_canary.write_ref import ( launch_canary_write_kernel_torch_reference, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( + FakeViolationLog, + assert_only_bits_set, + chain_anchor_signed, + make_canary_buf, + make_canary_buf_pair, + make_log_pair, + make_real_kv_source, + make_real_kv_sources, + make_verify_plan, + make_verify_plan_pair, + make_write_plan, + read_slot_fields, + stamp_clean_chain, + stamp_pair, + to_signed_int64, + write_slot_fields, +) +from sglang.test.kernels.kv_canary._differential import ( + _run_both_verify, + run_verify_diff, +) +from sglang.test.kernels.kv_canary._fixtures import clone_real_kv_sources +from sglang.test.kernels.kv_canary._hand_oracle import ( + _hand_fold_all, + _hand_fold_partial, +) register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_write_fuzz.py b/test/registered/kernels/ops/kv_canary/test_write_fuzz.py similarity index 96% rename from test/registered/jit/kv_canary/test_write_fuzz.py rename to test/registered/kernels/ops/kv_canary/test_write_fuzz.py index 87d8427f7..cd78d6a35 100644 --- a/test/registered/jit/kv_canary/test_write_fuzz.py +++ b/test/registered/kernels/ops/kv_canary/test_write_fuzz.py @@ -6,23 +6,6 @@ from dataclasses import dataclass import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( - FakeViolationLog, - make_canary_buf, - make_log_pair, - make_write_plan_pair, - stamp_pair, -) -from sglang.kernels.jit.tests.kv_canary._differential import _run_both_write -from sglang.kernels.jit.tests.kv_canary._fixtures import ( - clone_real_kv_sources, - make_real_kv_sources, -) -from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( - FUZZ_SEEDS_PR, - run_fuzz_combo, -) -from sglang.kernels.jit.tests.kv_canary._invariants import WriteInvariants from sglang.kernels.ops.kv_canary import consts from sglang.kernels.ops.kv_canary.verify import ( CanaryLaunchTag, @@ -30,6 +13,23 @@ from sglang.kernels.ops.kv_canary.verify import ( ) from sglang.kernels.ops.kv_canary.write import WritePlan from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( + FakeViolationLog, + make_canary_buf, + make_log_pair, + make_write_plan_pair, + stamp_pair, +) +from sglang.test.kernels.kv_canary._differential import _run_both_write +from sglang.test.kernels.kv_canary._fixtures import ( + clone_real_kv_sources, + make_real_kv_sources, +) +from sglang.test.kernels.kv_canary._fuzz_driver import ( + FUZZ_SEEDS_PR, + run_fuzz_combo, +) +from sglang.test.kernels.kv_canary._invariants import WriteInvariants register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/kv_canary/test_write_hand.py b/test/registered/kernels/ops/kv_canary/test_write_hand.py similarity index 99% rename from test/registered/jit/kv_canary/test_write_hand.py rename to test/registered/kernels/ops/kv_canary/test_write_hand.py index 9cc763daf..065633aa1 100644 --- a/test/registered/jit/kv_canary/test_write_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_write_hand.py @@ -7,7 +7,21 @@ from unittest.mock import patch import pytest import torch -from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( +from sglang.kernels.ops.kv_canary import consts +from sglang.kernels.ops.kv_canary import write as write_module +from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 +from sglang.kernels.ops.kv_canary.verify import ( + CANARY_SLOT_BYTES, + CanaryLaunchTag, + RealKvSource, + VerifyOrWriteContext, + launch_canary_verify_kernel, +) +from sglang.kernels.ops.kv_canary.write import ( + launch_canary_write_kernel, +) +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kernels.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_state_equal, assert_only_bits_set, @@ -24,32 +38,18 @@ from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( stamp_pair, to_signed_int64, ) -from sglang.kernels.jit.tests.kv_canary._differential import ( +from sglang.test.kernels.kv_canary._differential import ( _run_both_write, run_write_diff, ) -from sglang.kernels.jit.tests.kv_canary._fixtures import ( +from sglang.test.kernels.kv_canary._fixtures import ( clone_real_kv_sources, dummy_pseudo_tensors, ) -from sglang.kernels.jit.tests.kv_canary._hand_oracle import ( +from sglang.test.kernels.kv_canary._hand_oracle import ( _hand_fold_all, _hand_fold_partial, ) -from sglang.kernels.ops.kv_canary import consts -from sglang.kernels.ops.kv_canary import write as write_module -from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3 -from sglang.kernels.ops.kv_canary.verify import ( - CANARY_SLOT_BYTES, - CanaryLaunchTag, - RealKvSource, - VerifyOrWriteContext, - launch_canary_verify_kernel, -) -from sglang.kernels.ops.kv_canary.write import ( - launch_canary_write_kernel, -) -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=30, stage="jit-kernel-unit", runner_config="amd") diff --git a/test/registered/jit/test_fused_fp8_qkv_kv_cache.py b/test/registered/kernels/ops/kvcache/test_fused_fp8_qkv_kv_cache.py similarity index 100% rename from test/registered/jit/test_fused_fp8_qkv_kv_cache.py rename to test/registered/kernels/ops/kvcache/test_fused_fp8_qkv_kv_cache.py diff --git a/test/registered/jit/test_hicache.py b/test/registered/kernels/ops/kvcache/test_hicache.py similarity index 100% rename from test/registered/jit/test_hicache.py rename to test/registered/kernels/ops/kvcache/test_hicache.py diff --git a/test/registered/jit/test_hicache_page_first_write_back.py b/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py similarity index 100% rename from test/registered/jit/test_hicache_page_first_write_back.py rename to test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py diff --git a/test/registered/jit/test_hisparse.py b/test/registered/kernels/ops/kvcache/test_hisparse.py similarity index 100% rename from test/registered/jit/test_hisparse.py rename to test/registered/kernels/ops/kvcache/test_hisparse.py diff --git a/test/registered/jit/test_kvcacheio_asymmetric.py b/test/registered/kernels/ops/kvcache/test_kvcacheio_asymmetric.py similarity index 100% rename from test/registered/jit/test_kvcacheio_asymmetric.py rename to test/registered/kernels/ops/kvcache/test_kvcacheio_asymmetric.py diff --git a/test/registered/jit/minimax/test_minimax_store_kv_index.py b/test/registered/kernels/ops/kvcache/test_minimax_store_kv_index.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_store_kv_index.py rename to test/registered/kernels/ops/kvcache/test_minimax_store_kv_index.py diff --git a/test/registered/jit/test_set_mla_kv_buffer.py b/test/registered/kernels/ops/kvcache/test_set_mla_kv_buffer.py similarity index 100% rename from test/registered/jit/test_set_mla_kv_buffer.py rename to test/registered/kernels/ops/kvcache/test_set_mla_kv_buffer.py diff --git a/test/registered/jit/test_store_cache.py b/test/registered/kernels/ops/kvcache/test_store_cache.py similarity index 100% rename from test/registered/jit/test_store_cache.py rename to test/registered/kernels/ops/kvcache/test_store_cache.py diff --git a/test/registered/jit/test_fused_add_rmsnorm.py b/test/registered/kernels/ops/layernorm/test_fused_add_rmsnorm.py similarity index 100% rename from test/registered/jit/test_fused_add_rmsnorm.py rename to test/registered/kernels/ops/layernorm/test_fused_add_rmsnorm.py diff --git a/test/registered/jit/test_fused_eh_norm.py b/test/registered/kernels/ops/layernorm/test_fused_eh_norm.py similarity index 100% rename from test/registered/jit/test_fused_eh_norm.py rename to test/registered/kernels/ops/layernorm/test_fused_eh_norm.py diff --git a/test/registered/kernels/test_fused_op.py b/test/registered/kernels/ops/layernorm/test_fused_op.py similarity index 100% rename from test/registered/kernels/test_fused_op.py rename to test/registered/kernels/ops/layernorm/test_fused_op.py diff --git a/test/registered/kernels/test_fused_op_gpu_parity.py b/test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py similarity index 100% rename from test/registered/kernels/test_fused_op_gpu_parity.py rename to test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py diff --git a/test/registered/jit/test_fused_qk_gemma_rmsnorm_gate.py b/test/registered/kernels/ops/layernorm/test_fused_qk_gemma_rmsnorm_gate.py similarity index 100% rename from test/registered/jit/test_fused_qk_gemma_rmsnorm_gate.py rename to test/registered/kernels/ops/layernorm/test_fused_qk_gemma_rmsnorm_gate.py diff --git a/test/registered/kernels/test_gemma4_fused_routing.py b/test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py similarity index 100% rename from test/registered/kernels/test_gemma4_fused_routing.py rename to test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py diff --git a/test/registered/kernels/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py similarity index 100% rename from test/registered/kernels/test_kernels_namespace.py rename to test/registered/kernels/ops/layernorm/test_kernels_namespace.py diff --git a/test/registered/kernels/test_mhc_kernels.py b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py similarity index 100% rename from test/registered/kernels/test_mhc_kernels.py rename to test/registered/kernels/ops/layernorm/test_mhc_kernels.py diff --git a/python/sglang/kernels/jit/tests/test_minimax_m3_rmsnorm.py b/test/registered/kernels/ops/layernorm/test_minimax_m3_rmsnorm.py similarity index 92% rename from python/sglang/kernels/jit/tests/test_minimax_m3_rmsnorm.py rename to test/registered/kernels/ops/layernorm/test_minimax_m3_rmsnorm.py index 7eb477fec..dc46ffd90 100644 --- a/python/sglang/kernels/jit/tests/test_minimax_m3_rmsnorm.py +++ b/test/registered/kernels/ops/layernorm/test_minimax_m3_rmsnorm.py @@ -18,6 +18,9 @@ from sglang.kernels.ops.layernorm.minimax_m3_rmsnorm import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, ) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd") DEVICE = "cuda" EPS = 1e-6 @@ -76,3 +79,9 @@ def test_gemma_fused_add_rmsnorm_matches_reference(shape, dtype): torch.testing.assert_close(residual_out, ref_residual, atol=2e-2, rtol=2e-2) torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/jit/test_qknorm.py b/test/registered/kernels/ops/layernorm/test_qknorm.py similarity index 100% rename from test/registered/jit/test_qknorm.py rename to test/registered/kernels/ops/layernorm/test_qknorm.py diff --git a/test/registered/jit/test_qknorm_across_heads.py b/test/registered/kernels/ops/layernorm/test_qknorm_across_heads.py similarity index 100% rename from test/registered/jit/test_qknorm_across_heads.py rename to test/registered/kernels/ops/layernorm/test_qknorm_across_heads.py diff --git a/test/registered/jit/test_rmsnorm.py b/test/registered/kernels/ops/layernorm/test_rmsnorm.py similarity index 100% rename from test/registered/jit/test_rmsnorm.py rename to test/registered/kernels/ops/layernorm/test_rmsnorm.py diff --git a/test/registered/jit/test_rmsnorm_hf.py b/test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py similarity index 100% rename from test/registered/jit/test_rmsnorm_hf.py rename to test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py diff --git a/python/sglang/kernels/jit/tests/test_sconv_decode_metadata.py b/test/registered/kernels/ops/mamba/test_sconv_decode_metadata.py similarity index 94% rename from python/sglang/kernels/jit/tests/test_sconv_decode_metadata.py rename to test/registered/kernels/ops/mamba/test_sconv_decode_metadata.py index 6ac9f4c14..3a3cb43fc 100644 --- a/python/sglang/kernels/jit/tests/test_sconv_decode_metadata.py +++ b/test/registered/kernels/ops/mamba/test_sconv_decode_metadata.py @@ -13,6 +13,9 @@ from sglang.srt.models.inkling_common.kernels.sconv import ( fused_decode_sconv_metadata, precompute_helion_decode_metadata, ) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only") diff --git a/python/sglang/kernels/jit/tests/test_sconv_extend_metadata.py b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py similarity index 97% rename from python/sglang/kernels/jit/tests/test_sconv_extend_metadata.py rename to test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py index cecb81ff5..de7e36143 100644 --- a/python/sglang/kernels/jit/tests/test_sconv_extend_metadata.py +++ b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py @@ -18,6 +18,9 @@ from sglang.srt.models.inkling_common.kernels.sconv import ( fused_extend_sconv_metadata, precompute_helion_extend_metadata, ) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only") diff --git a/test/registered/jit/test_transfer_mamba.py b/test/registered/kernels/ops/mamba/test_transfer_mamba.py similarity index 100% rename from test/registered/jit/test_transfer_mamba.py rename to test/registered/kernels/ops/mamba/test_transfer_mamba.py diff --git a/test/registered/jit/test_inkling_rel_proj.py b/test/registered/kernels/ops/model/test_inkling_rel_proj.py similarity index 100% rename from test/registered/jit/test_inkling_rel_proj.py rename to test/registered/kernels/ops/model/test_inkling_rel_proj.py diff --git a/test/registered/kernels/test_fp4_moe.py b/test/registered/kernels/ops/moe/test_fp4_moe.py similarity index 100% rename from test/registered/kernels/test_fp4_moe.py rename to test/registered/kernels/ops/moe/test_fp4_moe.py diff --git a/test/registered/kernels/test_fused_topk_deepseek.py b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py similarity index 100% rename from test/registered/kernels/test_fused_topk_deepseek.py rename to test/registered/kernels/ops/moe/test_fused_topk_deepseek.py diff --git a/test/registered/jit/test_kpool_topk_transform.py b/test/registered/kernels/ops/moe/test_kpool_topk_transform.py similarity index 100% rename from test/registered/jit/test_kpool_topk_transform.py rename to test/registered/kernels/ops/moe/test_kpool_topk_transform.py diff --git a/test/registered/jit/test_marlin_packed_topk_unpack.py b/test/registered/kernels/ops/moe/test_marlin_packed_topk_unpack.py similarity index 100% rename from test/registered/jit/test_marlin_packed_topk_unpack.py rename to test/registered/kernels/ops/moe/test_marlin_packed_topk_unpack.py diff --git a/python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py b/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py similarity index 97% rename from python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py rename to test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py index 3038d8405..bc9f18d94 100644 --- a/python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py +++ b/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py @@ -12,7 +12,7 @@ ROCm-only. The pure quant test runs on any ROCm arch; the native MXFP8 ``dot_scaled`` linear/MoE tests are gated to CDNA4 gfx95x (the hardware microscaling matrix cores) -- gfx942 has no native ``dot_scaled`` MX path. -Run: pytest python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py -v +Run: pytest python/sglang/kernels/testing/test_minimax_m3_mxfp8.py -v """ import pytest @@ -33,6 +33,9 @@ from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import ( # noqa: E402 _mxfp8_e4m3_quantize_triton, dequant_mxfp8_to_bf16, ) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd") DEVICE = "cuda" diff --git a/test/registered/jit/minimax/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py similarity index 100% rename from test/registered/jit/minimax/test_minimax_quant_scatter.py rename to test/registered/kernels/ops/moe/test_minimax_quant_scatter.py diff --git a/test/registered/jit/test_moe_align_block_size.py b/test/registered/kernels/ops/moe/test_moe_align_block_size.py similarity index 100% rename from test/registered/jit/test_moe_align_block_size.py rename to test/registered/kernels/ops/moe/test_moe_align_block_size.py diff --git a/test/registered/jit/test_moe_fused_gate.py b/test/registered/kernels/ops/moe/test_moe_fused_gate.py similarity index 100% rename from test/registered/jit/test_moe_fused_gate.py rename to test/registered/kernels/ops/moe/test_moe_fused_gate.py diff --git a/test/registered/jit/test_moe_lora_align_block_size.py b/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py similarity index 100% rename from test/registered/jit/test_moe_lora_align_block_size.py rename to test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py diff --git a/python/sglang/kernels/jit/tests/test_moe_preprocess.py b/test/registered/kernels/ops/moe/test_moe_preprocess.py similarity index 95% rename from python/sglang/kernels/jit/tests/test_moe_preprocess.py rename to test/registered/kernels/ops/moe/test_moe_preprocess.py index 121811e82..43e12c59e 100644 --- a/python/sglang/kernels/jit/tests/test_moe_preprocess.py +++ b/test/registered/kernels/ops/moe/test_moe_preprocess.py @@ -13,6 +13,9 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import ( get_src2dst, grouped_gemm_triton, ) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only") diff --git a/python/sglang/kernels/jit/tests/test_moe_topk_sigmoid.py b/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py similarity index 99% rename from python/sglang/kernels/jit/tests/test_moe_topk_sigmoid.py rename to test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py index 769dacd92..785ef7ec3 100644 --- a/python/sglang/kernels/jit/tests/test_moe_topk_sigmoid.py +++ b/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py @@ -14,6 +14,9 @@ import pytest import torch from sglang.kernels.ops.moe.moe_topk_sigmoid import topk_sigmoid +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") try: from sgl_kernel import topk_sigmoid as topk_sigmoid_aot diff --git a/test/registered/jit/test_moe_wna16_marlin.py b/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py similarity index 100% rename from test/registered/jit/test_moe_wna16_marlin.py rename to test/registered/kernels/ops/moe/test_moe_wna16_marlin.py diff --git a/test/registered/jit/test_post_reorder_deepgemm.py b/test/registered/kernels/ops/moe/test_post_reorder_deepgemm.py similarity index 100% rename from test/registered/jit/test_post_reorder_deepgemm.py rename to test/registered/kernels/ops/moe/test_post_reorder_deepgemm.py diff --git a/test/registered/jit/test_renorm.py b/test/registered/kernels/ops/moe/test_renorm.py similarity index 100% rename from test/registered/jit/test_renorm.py rename to test/registered/kernels/ops/moe/test_renorm.py diff --git a/test/registered/jit/test_sigmoid_gate_mul.py b/test/registered/kernels/ops/moe/test_sigmoid_gate_mul.py similarity index 100% rename from test/registered/jit/test_sigmoid_gate_mul.py rename to test/registered/kernels/ops/moe/test_sigmoid_gate_mul.py diff --git a/test/registered/kernels/test_sigmoid_gate_mul.py b/test/registered/kernels/ops/moe/test_sigmoid_gate_mul_broadcast.py similarity index 100% rename from test/registered/kernels/test_sigmoid_gate_mul.py rename to test/registered/kernels/ops/moe/test_sigmoid_gate_mul_broadcast.py diff --git a/test/registered/jit/test_awq_dequantize.py b/test/registered/kernels/ops/quantization/test_awq_dequantize.py similarity index 100% rename from test/registered/jit/test_awq_dequantize.py rename to test/registered/kernels/ops/quantization/test_awq_dequantize.py diff --git a/test/registered/jit/test_awq_marlin_moe_repack.py b/test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py similarity index 100% rename from test/registered/jit/test_awq_marlin_moe_repack.py rename to test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py diff --git a/test/registered/jit/test_awq_marlin_repack.py b/test/registered/kernels/ops/quantization/test_awq_marlin_repack.py similarity index 100% rename from test/registered/jit/test_awq_marlin_repack.py rename to test/registered/kernels/ops/quantization/test_awq_marlin_repack.py diff --git a/test/registered/jit/test_gptq_marlin.py b/test/registered/kernels/ops/quantization/test_gptq_marlin.py similarity index 100% rename from test/registered/jit/test_gptq_marlin.py rename to test/registered/kernels/ops/quantization/test_gptq_marlin.py diff --git a/test/registered/jit/test_gptq_marlin_repack.py b/test/registered/kernels/ops/quantization/test_gptq_marlin_repack.py similarity index 100% rename from test/registered/jit/test_gptq_marlin_repack.py rename to test/registered/kernels/ops/quantization/test_gptq_marlin_repack.py diff --git a/test/registered/jit/test_mxfp8_moe.py b/test/registered/kernels/ops/quantization/test_mxfp8_moe.py similarity index 100% rename from test/registered/jit/test_mxfp8_moe.py rename to test/registered/kernels/ops/quantization/test_mxfp8_moe.py diff --git a/test/registered/jit/test_per_tensor_quant_fp8.py b/test/registered/kernels/ops/quantization/test_per_tensor_quant_fp8.py similarity index 100% rename from test/registered/jit/test_per_tensor_quant_fp8.py rename to test/registered/kernels/ops/quantization/test_per_tensor_quant_fp8.py diff --git a/test/registered/jit/test_per_token_group_quant.py b/test/registered/kernels/ops/quantization/test_per_token_group_quant.py similarity index 100% rename from test/registered/jit/test_per_token_group_quant.py rename to test/registered/kernels/ops/quantization/test_per_token_group_quant.py diff --git a/test/registered/jit/test_per_token_group_quant_8bit_v2.py b/test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py similarity index 100% rename from test/registered/jit/test_per_token_group_quant_8bit_v2.py rename to test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py diff --git a/test/registered/kernels/test_boundary_kv_fix_kernels.py b/test/registered/kernels/ops/speculative/test_boundary_kv_fix_kernels.py similarity index 100% rename from test/registered/kernels/test_boundary_kv_fix_kernels.py rename to test/registered/kernels/ops/speculative/test_boundary_kv_fix_kernels.py diff --git a/test/registered/kernels/test_gather_spec_extras.py b/test/registered/kernels/ops/speculative/test_gather_spec_extras.py similarity index 100% rename from test/registered/kernels/test_gather_spec_extras.py rename to test/registered/kernels/ops/speculative/test_gather_spec_extras.py diff --git a/test/registered/jit/test_ngram_embedding.py b/test/registered/kernels/ops/speculative/test_ngram_embedding.py similarity index 100% rename from test/registered/jit/test_ngram_embedding.py rename to test/registered/kernels/ops/speculative/test_ngram_embedding.py diff --git a/test/registered/jit/test_resolve_future_token_ids.py b/test/registered/kernels/ops/speculative/test_resolve_future_token_ids.py similarity index 100% rename from test/registered/jit/test_resolve_future_token_ids.py rename to test/registered/kernels/ops/speculative/test_resolve_future_token_ids.py diff --git a/test/registered/kernels/test_spec_topk1.py b/test/registered/kernels/ops/speculative/test_spec_topk1.py similarity index 100% rename from test/registered/kernels/test_spec_topk1.py rename to test/registered/kernels/ops/speculative/test_spec_topk1.py diff --git a/test/registered/unit/distributed/test_vmm_utils.py b/test/registered/unit/distributed/test_vmm_utils.py index fd3e6f5c6..23178cbe7 100644 --- a/test/registered/unit/distributed/test_vmm_utils.py +++ b/test/registered/unit/distributed/test_vmm_utils.py @@ -20,7 +20,6 @@ import torch import torch.distributed as dist from cuda.bindings import driver as drv -from sglang.kernels.jit.tests.utils import multigpu_pytest_main from sglang.kernels.jit.utils import cache_once from sglang.srt.distributed.device_communicators.vmm_utils import ( check_drv, @@ -32,6 +31,7 @@ from sglang.srt.distributed.device_communicators.vmm_utils import ( release_mappings, ) from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main register_cuda_ci(est_time=60, stage="base-b", runner_config="2-gpu-large")