From 99f636a86fb80d5e399eda93902e772fba1957b9 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 23 Jul 2026 08:35:09 +0800 Subject: [PATCH] [Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels (#32072) Co-authored-by: Claude Opus 4.8 --- .claude/skills/add-jit-kernel/SKILL.md | 70 +- .claude/skills/add-sgl-kernel/SKILL.md | 2 +- .../references/fuse-overlap-catalog.md | 36 +- .../scripts/triage_kernel_helpers.py | 20 +- .claude/skills/write-sglang-test/SKILL.md | 6 +- .github/CODEOWNERS | 8 +- .github/MAINTAINER.md | 2 +- .github/labeler.yml | 2 +- .github/workflows/_pr-test-check-changes.yml | 3 +- .github/workflows/nightly-test-nvidia.yml | 2 +- .github/workflows/pr-test-amd-rocm720.yml | 4 +- .github/workflows/pr-test-amd.yml | 4 +- .github/workflows/pr-test-npu.yml | 956 +++++++++--------- .github/workflows/pr-test.yml | 2 +- .pre-commit-config.yaml | 2 +- 3rdparty/amd/wheel/sglang/pyproject.toml | 2 +- .../bench_cutedsl_kda_decode.py | 4 +- .../bench_gdn_qkv_split.py | 4 +- benchmark/hicache/bench_hicache_write_back.py | 6 +- benchmark/hicache/perf.py | 2 +- ...benchmark_cute_dsl_fp8_paged_mqa_logits.py | 4 +- .../development_jit_kernel_guide.mdx | 24 +- python/pyproject.toml | 2 +- python/pyproject_cpu.toml | 2 +- python/pyproject_npu.toml | 2 +- python/pyproject_other.toml | 2 +- python/pyproject_xpu.toml | 2 +- python/sglang/jit_kernel/activation.py | 5 - python/sglang/jit_kernel/dsv3_fused_a_gemm.py | 5 - python/sglang/jit_kernel/dsv3_router_gemm.py | 5 - python/sglang/jit_kernel/norm.py | 5 - .../sglang/jit_kernel/per_tensor_quant_fp8.py | 5 - .../jit_kernel/per_token_group_quant.py | 5 - .../per_token_group_quant_8bit_v2.py | 5 - python/sglang/jit_kernel/set_mla_kv_buffer.py | 5 - .../jit_kernel/tests/kv_canary/__init__.py | 0 python/sglang/kernels/README.md | 6 +- python/sglang/kernels/__init__.py | 4 +- python/sglang/kernels/fused_op.py | 2 +- .../{jit_kernel => kernels/jit}/.clang-format | 0 python/sglang/kernels/jit/__init__.py | 2 +- .../{jit_kernel => kernels/jit}/__main__.py | 0 .../jit}/benchmark/kv_canary/utils.py | 0 .../jit}/benchmark/marker.py | 0 .../jit}/benchmark/utils.py | 0 .../jit}/csrc/add_constant.cuh | 0 .../jit}/csrc/attention/fixup_zero_kv.cuh | 0 .../csrc/attention/fused_fp8_qkv_kv_cache.cuh | 0 .../jit}/csrc/deepseek_v32/indexer_k.cuh | 0 .../jit}/csrc/deepseek_v4/c128.cuh | 0 .../jit}/csrc/deepseek_v4/c128_online.cuh | 0 .../jit}/csrc/deepseek_v4/c128_online_v2.cuh | 0 .../jit}/csrc/deepseek_v4/c128_v2.cuh | 0 .../jit}/csrc/deepseek_v4/c4.cuh | 0 .../jit}/csrc/deepseek_v4/c4_v2.cuh | 0 .../jit}/csrc/deepseek_v4/c_plan.cuh | 0 .../jit}/csrc/deepseek_v4/common.cuh | 0 .../fp8_wo_a_group_major_quant.cuh | 0 .../jit}/csrc/deepseek_v4/fused_norm_rope.cuh | 0 .../csrc/deepseek_v4/fused_norm_rope_v2.cuh | 0 .../jit}/csrc/deepseek_v4/hash_topk.cuh | 0 .../jit}/csrc/deepseek_v4/main_norm_rope.cuh | 0 .../deepseek_v4/mega_moe_pre_dispatch.cuh | 0 .../jit}/csrc/deepseek_v4/online_c128_mtp.cuh | 0 .../csrc/deepseek_v4/paged_mqa_metadata.cuh | 0 .../jit}/csrc/deepseek_v4/rope.cuh | 0 .../silu_and_mul_masked_post_quant.cuh | 0 .../jit}/csrc/deepseek_v4/store.cuh | 0 .../jit}/csrc/deepseek_v4/topk_v1.cuh | 0 .../jit}/csrc/deepseek_v4/topk_v2.cuh | 0 .../csrc/diffusion/causal_conv3d_cat_pad.cuh | 0 .../csrc/diffusion/ltx2_qknorm_split_rope.cuh | 0 .../jit}/csrc/diffusion/norm_scale_shift.cuh | 0 .../jit}/csrc/diffusion/qknorm_rope.cuh | 0 .../jit}/csrc/diffusion/residual_gate_add.cuh | 0 .../csrc/diffusion/timestep_embedding.cuh | 0 .../jit}/csrc/distributed/communicator.cuh | 0 .../csrc/distributed/custom_all_reduce.cuh | 0 .../jit}/csrc/distributed/ipc.cuh | 0 .../jit}/csrc/distributed/tp_qknorm.cuh | 0 .../jit}/csrc/dsa/fused_store_index_cache.cuh | 0 .../jit}/csrc/dsa/kpool_topk_transform.cuh | 0 .../jit}/csrc/elementwise/activation.cuh | 0 .../jit}/csrc/elementwise/clamp_position.cuh | 0 .../jit}/csrc/elementwise/concat_mla.cuh | 0 .../csrc/elementwise/fused_add_rmsnorm.cuh | 0 .../jit}/csrc/elementwise/fused_eh_norm.cuh | 0 .../csrc/elementwise/fused_metadata_copy.cuh | 2 +- .../csrc/elementwise/fused_qknorm_rope.cuh | 0 .../jit}/csrc/elementwise/kvcache.cuh | 0 .../jit}/csrc/elementwise/pos_enc.cuh | 0 .../jit}/csrc/elementwise/qknorm.cuh | 0 .../csrc/elementwise/qknorm_across_heads.cuh | 0 .../elementwise/resolve_future_token_ids.cuh | 0 .../jit}/csrc/elementwise/rmsnorm.cuh | 0 .../jit}/csrc/elementwise/rmsnorm_hf.cuh | 0 .../jit}/csrc/elementwise/rope.cuh | 0 .../csrc/elementwise/set_mla_kv_buffer.cuh | 0 .../csrc/fast-hadamard-transform/code_gen.py | 0 .../fast_hadamard_transform.h | 0 .../fast_hadamard_transform_common.h | 0 .../fast_hadamard_transform_special.h | 0 .../fast-hadamard-transform/hadamard_jit.cuh | 0 .../fast-hadamard-transform/static_switch.h | 0 .../jit}/csrc/gemm/awq_dequantize.cuh | 0 .../jit}/csrc/gemm/dsv3_fused_a_gemm.cuh | 0 .../jit}/csrc/gemm/dsv3_router_gemm.cuh | 0 .../fp8_blockwise_scaled_mm_entry.cuh | 0 .../fp8_blockwise_scaled_mm_sm120.cuh | 0 .../csrc/gemm/marlin/awq_marlin_repack.cuh | 0 .../jit}/csrc/gemm/marlin/dequant.h | 0 .../jit}/csrc/gemm/marlin/gptq_marlin.cuh | 0 .../csrc/gemm/marlin/gptq_marlin_repack.cuh | 0 .../jit}/csrc/gemm/marlin/kernel.h | 0 .../jit}/csrc/gemm/marlin/marlin.cuh | 0 .../jit}/csrc/gemm/marlin/marlin_dtypes.cuh | 0 .../jit}/csrc/gemm/marlin/marlin_template.h | 0 .../jit}/csrc/gemm/marlin_moe/kernel.h | 0 .../csrc/gemm/marlin_moe/marlin_template.h | 0 .../csrc/gemm/marlin_moe/moe_wna16_marlin.cuh | 0 .../jit}/csrc/gemm/per_tensor_quant_fp8.cuh | 0 .../jit}/csrc/gemm/per_token_group_quant.cuh | 0 .../gemm/per_token_group_quant_8bit_v2.cuh | 0 .../jit}/csrc/hisparse.cuh | 0 .../jit}/csrc/inkling/causal_conv1d.cuh | 0 .../jit}/csrc/inkling/draft_extend_sconv.cuh | 0 .../jit}/csrc/inkling/fused_decode_update.cuh | 0 .../csrc/inkling/gather_scatter_sconv.cuh | 0 .../jit}/csrc/inkling/inkling_all_reduce.cuh | 0 .../jit}/csrc/inkling/inkling_ar_barrier.cuh | 0 .../csrc/inkling/inkling_ar_fused_decode.cuh | 0 .../inkling/inkling_ar_scattered_sconv.cuh | 0 .../inkling/inkling_attn_prologue_fused.cuh | 0 .../jit}/csrc/inkling/inkling_rel_proj.cuh | 0 .../jit}/csrc/inkling/inkling_row_scale.cuh | 0 .../jit}/csrc/inkling/update_sconv_cache.cuh | 0 .../jit}/csrc/kv_canary/canary_common.cuh | 0 .../csrc/kv_canary/canary_plan_entries.cuh | 0 .../jit}/csrc/kv_canary/canary_verify.cuh | 2 +- .../jit}/csrc/kv_canary/canary_write.cuh | 2 +- .../jit}/csrc/kv_canary/consts.cuh | 0 .../jit}/csrc/kvcacheio/hicache.cuh | 2 +- .../jit}/csrc/kvcacheio/relayout.cuh | 0 .../jit}/csrc/kvcacheio/staged_write_back.cuh | 0 .../jit}/csrc/kvcacheio/transfer_mamba.cuh | 0 .../jit}/csrc/lora/moe_lora_align_kernel.cu | 0 .../jit}/csrc/lplb/dispatch_probability.cuh | 0 .../jit}/csrc/lplb/ipm.cuh | 0 .../jit}/csrc/lplb/lp_post.cuh | 0 .../jit}/csrc/lplb/lp_prep.cuh | 0 .../csrc/minimax/fused_gemma_qknorm_rope.cuh | 0 .../csrc/minimax/fused_store_kv_index.cuh | 0 .../jit}/csrc/minimax/minimax_decode_topk.cuh | 0 .../csrc/minimax/per_token_quant_ue8m0.cuh | 0 ...es_sm100_mxfp8_blockscaled_group_quant.cuh | 0 ...sm100_mxfp8_blockscaled_moe_group_gemm.cuh | 0 ...fp8_blockscaled_moe_group_gemm_functor.cuh | 0 ...xfp8_blockscaled_moe_group_gemm_traits.cuh | 0 .../csrc/moe/inkling_gate_topk_renorm.cuh | 0 .../jit}/csrc/moe/moe_align_kernel.cu | 0 .../jit}/csrc/moe/moe_finalize_fuse_shared.cu | 0 .../jit}/csrc/moe/moe_fused_gate.cuh | 0 .../jit}/csrc/moe/moe_permute_prepare.cu | 0 .../jit}/csrc/moe/moe_topk_sigmoid.cuh | 0 .../jit}/csrc/moe/tvm_ffi_utils.h | 0 .../jit}/csrc/ngram_corpus/ngram.cpp | 0 .../jit}/csrc/ngram_corpus/ngram.h | 0 .../csrc/ngram_corpus/ngram_corpus_ffi.cpp | 0 .../jit}/csrc/ngram_corpus/param.h | 0 .../jit}/csrc/ngram_corpus/queue.h | 0 .../jit}/csrc/ngram_corpus/result.cpp | 0 .../jit}/csrc/ngram_corpus/result.h | 0 .../csrc/ngram_corpus/suffix_automaton.cpp | 0 .../jit}/csrc/ngram_corpus/suffix_automaton.h | 0 .../jit}/csrc/ngram_corpus/trie.cpp | 0 .../jit}/csrc/ngram_corpus/trie.h | 0 .../jit}/csrc/ngram_embedding.cuh | 0 .../sparse_mla_q8kv8_prefill_sm90/config.h | 0 .../sparse_mla_q8kv8_prefill_sm90/defines.h | 0 .../dense_fp8_transpose_v.h | 0 .../dense_fp8_utils.h | 0 .../sparse_mla_q8kv8_prefill_sm90/entry.cuh | 0 .../sparse_mla_q8kv8_prefill_sm90/helpers.h | 0 .../sparse_mla_q8kv8_prefill_sm90/kernel.cuh | 0 .../sparse_mla_q8kv8_prefill_sm90/params.h | 0 .../kimi_k2_moe_fused_gate.cuh | 0 .../moe_lora_merged_align_kernel.cu | 0 .../trtllm_lora_temp/topk_softmax_pack.cuh | 2 +- .../jit}/include/sgl_kernel/atomic.cuh | 0 .../jit}/include/sgl_kernel/cta.cuh | 0 .../sgl_kernel/deepseek_v4/compress.cuh | 0 .../sgl_kernel/deepseek_v4/compress_v2.cuh | 0 .../sgl_kernel/deepseek_v4/fp8_utils.cuh | 0 .../sgl_kernel/deepseek_v4/kvcacheio.cuh | 0 .../sgl_kernel/deepseek_v4/topk_impl.cuh | 0 .../sgl_kernel/distributed/communicator.cuh | 0 .../jit}/include/sgl_kernel/ffi.h | 0 .../jit}/include/sgl_kernel/impl/norm.cuh | 0 .../jit}/include/sgl_kernel/math.cuh | 0 .../jit}/include/sgl_kernel/runtime.cuh | 0 .../jit}/include/sgl_kernel/scalar_type.hpp | 0 .../jit}/include/sgl_kernel/source_location.h | 0 .../jit}/include/sgl_kernel/tensor.h | 0 .../jit}/include/sgl_kernel/tile.cuh | 0 .../jit}/include/sgl_kernel/type.cuh | 0 .../jit}/include/sgl_kernel/utils.cuh | 0 .../jit}/include/sgl_kernel/utils.h | 0 .../jit}/include/sgl_kernel/vec.cuh | 0 .../jit}/include/sgl_kernel/warp.cuh | 0 .../jit/tests/deepseek_v4}/__init__.py | 0 .../jit}/tests/deepseek_v4/common.py | 0 .../jit/tests/kv_canary}/__init__.py | 0 .../jit}/tests/kv_canary/_canary_helpers.py | 4 +- .../jit}/tests/kv_canary/_constants.py | 0 .../jit}/tests/kv_canary/_differential.py | 2 +- .../jit}/tests/kv_canary/_fixtures.py | 2 +- .../jit}/tests/kv_canary/_fuzz_driver.py | 2 +- .../jit}/tests/kv_canary/_hand_oracle.py | 0 .../jit}/tests/kv_canary/_invariants.py | 2 +- .../jit}/tests/test_minimax_m3_mxfp8.py | 2 +- .../jit}/tests/test_minimax_m3_rmsnorm.py | 0 .../jit}/tests/test_moe_preprocess.py | 0 .../jit}/tests/test_moe_topk_sigmoid.py | 0 .../jit}/tests/test_sconv_decode_metadata.py | 0 .../jit}/tests/test_sconv_extend_metadata.py | 0 .../jit}/tests/utils.py | 0 python/sglang/kernels/jit/utils/compile.py | 6 +- .../sglang/kernels/ops/activation/__init__.py | 8 +- .../decode/flash_with_topk_idx.py | 2 +- .../kernels/ops/communication/__init__.py | 2 +- python/sglang/kernels/ops/communication/mp.py | 4 +- .../sglang/kernels/ops/diffusion/__init__.py | 6 +- python/sglang/kernels/ops/gemm/__init__.py | 4 +- .../sglang/kernels/ops/gemm/fused_a_gemm.py | 6 +- .../sglang/kernels/ops/layernorm/__init__.py | 8 +- python/sglang/kernels/ops/moe/__init__.py | 2 +- .../kernels/ops/quantization/__init__.py | 2 +- python/sglang/kernels/spec.py | 6 +- .../existing-fast-paths.md | 28 +- .../runtime/layers/activation.py | 2 +- .../runtime/layers/layernorm.py | 5 +- python/sglang/srt/layers/activation.py | 2 +- python/sglang/srt/layers/attention/vision.py | 4 +- python/sglang/srt/layers/layernorm.py | 6 +- python/sglang/srt/layers/moe/cutlass_moe.py | 2 +- .../sglang/srt/layers/moe/cutlass_w4a8_moe.py | 10 +- .../moe/fused_moe_triton/fused_marlin_moe.py | 2 +- .../fused_moe_triton/triton_kernels_moe.py | 2 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 4 +- .../moe/moe_runner/triton_utils/fused_moe.py | 2 +- python/sglang/srt/layers/moe/topk.py | 2 +- python/sglang/srt/layers/quantization/gguf.py | 2 +- .../srt/lora/trtllm_lora_temp/environ.py | 2 +- python/sglang/srt/models/deepseek_v2.py | 2 +- python/sglang/srt/models/utils.py | 5 +- scripts/code_sync/utils.py | 2 +- sgl-kernel/benchmark/bench_fp8_gemm.py | 4 +- .../benchmark/bench_fp8_gemm_swap_ab.py | 4 +- .../benchmark/bench_per_tensor_quant_fp8.py | 4 +- sgl-kernel/csrc/common_extension.cc | 2 +- sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu | 2 +- sgl-kernel/python/sgl_kernel/moe.py | 2 +- .../sgl_kernel/testing/rotary_embedding.py | 6 +- sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py | 4 +- test/README.md | 2 +- .../attention/unittests/KNOWN_FAILURES.md | 2 +- .../attention/unittests/dsa/README.md | 2 +- .../attention/unittests/dual_chunk/README.md | 2 +- .../jit/benchmark/bench_activation.py | 18 +- .../jit/benchmark/bench_add_constant.py | 2 +- .../jit/benchmark/bench_awq_dequantize.py | 2 +- .../jit/benchmark/bench_clamp_position.py | 2 +- .../jit/benchmark/bench_concat_mla.py | 2 +- .../jit/benchmark/bench_custom_all_reduce.py | 4 +- .../jit/benchmark/bench_dsv3_fused_a_gemm.py | 4 +- .../jit/benchmark/bench_dsv3_router_gemm.py | 8 +- .../jit/benchmark/bench_dsv4_fp4_indexer.py | 2 +- .../jit/benchmark/bench_fp8_blockwise_gemm.py | 2 +- .../jit/benchmark/bench_fused_eh_norm.py | 2 +- .../benchmark/bench_fused_fp8_qkv_kv_cache.py | 2 +- .../jit/benchmark/bench_fused_qknorm_rope.py | 2 +- .../jit/benchmark/bench_hadamard.py | 2 +- .../registered/jit/benchmark/bench_hicache.py | 2 +- .../jit/benchmark/bench_hisparse.py | 2 +- .../bench_mla_kv_pack_quantize_fp8.py | 2 +- .../jit/benchmark/bench_moe_fused_gate.py | 4 +- .../jit/benchmark/bench_mxfp8_moe.py | 2 +- .../benchmark/bench_ngram_compute_decode.py | 2 +- .../bench_ngram_update_token_table.py | 2 +- test/registered/jit/benchmark/bench_norm.py | 8 +- .../jit/benchmark/bench_online_c128_mtp.py | 2 +- .../benchmark/bench_per_tensor_quant_fp8.py | 6 +- .../benchmark/bench_per_token_group_quant.py | 10 +- .../bench_per_token_group_quant_8bit_v2.py | 6 +- .../bench_per_token_group_quant_masked.py | 10 +- .../benchmark/bench_post_reorder_deepgemm.py | 2 +- test/registered/jit/benchmark/bench_qknorm.py | 6 +- .../benchmark/bench_qknorm_across_heads.py | 4 +- test/registered/jit/benchmark/bench_renorm.py | 2 +- .../bench_resolve_future_token_ids.py | 2 +- test/registered/jit/benchmark/bench_rope.py | 2 +- .../jit/benchmark/bench_set_mla_kv_buffer.py | 6 +- .../bench_sparse_mla_q8kv8_prefill_sm90.py | 2 +- .../jit/benchmark/bench_spec_topk1.py | 2 +- .../jit/benchmark/bench_store_cache.py | 4 +- .../benchmark/bench_symm_mem_all_gather.py | 4 +- test/registered/jit/benchmark/bench_topk.py | 2 +- .../jit/benchmark/bench_tp_qknorm.py | 4 +- .../bench_vocab_parallel_embedding.py | 2 +- .../diffusion/bench_causal_conv3d_cat_pad.py | 2 +- .../bench_diffusion_nvfp4_scaled_mm.py | 4 +- .../diffusion/bench_fused_norm_scale_shift.py | 2 +- .../benchmark/diffusion/bench_norm_impls.py | 8 +- .../benchmark/diffusion/bench_qknorm_rope.py | 4 +- .../diffusion/bench_qwen_image_modulation.py | 2 +- .../jit/benchmark/kv_canary/bench_plan.py | 4 +- .../kv_canary/bench_scatter_req_token_ids.py | 2 +- .../jit/benchmark/kv_canary/bench_verify.py | 4 +- .../jit/benchmark/kv_canary/bench_write.py | 4 +- .../minimax/bench_minimax_decode_topk.py | 2 +- .../minimax/bench_minimax_qknorm_rope.py | 2 +- .../minimax/bench_minimax_store_kv_index.py | 2 +- .../jit/deepseek_v4/test_c128_v2.py | 2 +- test/registered/jit/deepseek_v4/test_c4_v2.py | 2 +- .../jit/diffusion/test_qknorm_rope.py | 2 +- .../jit/kv_canary/test_const_sync.py | 8 +- .../jit/kv_canary/test_kernel_config.py | 6 +- .../jit/kv_canary/test_pipeline_e2e.py | 4 +- .../jit/kv_canary/test_plan_fuzz.py | 8 +- .../jit/kv_canary/test_plan_hand.py | 6 +- test/registered/jit/kv_canary/test_utils.py | 2 +- .../jit/kv_canary/test_verify_fuzz.py | 10 +- .../jit/kv_canary/test_verify_hand.py | 8 +- .../jit/kv_canary/test_write_fuzz.py | 10 +- .../jit/kv_canary/test_write_hand.py | 8 +- test/registered/jit/test_activation.py | 4 +- test/registered/jit/test_custom_all_reduce.py | 2 +- test/registered/jit/test_cutedsl_gdn.py | 2 +- ...epseek_v4_compress_state_runtime_shapes.py | 2 +- test/registered/jit/test_dsv3_fused_a_gemm.py | 2 +- test/registered/jit/test_dsv3_router_gemm.py | 2 +- test/registered/jit/test_fused_add_rmsnorm.py | 2 +- test/registered/jit/test_hadamard_jit.py | 8 +- .../jit/test_per_tensor_quant_fp8.py | 4 +- .../jit/test_per_token_group_quant.py | 4 +- .../jit/test_per_token_group_quant_8bit_v2.py | 4 +- test/registered/jit/test_qknorm.py | 2 +- .../jit/test_qknorm_across_heads.py | 2 +- test/registered/jit/test_rmsnorm.py | 6 +- test/registered/jit/test_set_mla_kv_buffer.py | 4 +- .../jit/test_symm_mem_all_gather.py | 2 +- test/registered/jit/test_tp_qknorm.py | 2 +- .../kernels/test_kernels_namespace.py | 5 +- .../unit/distributed/test_vmm_utils.py | 2 +- 354 files changed, 889 insertions(+), 875 deletions(-) delete mode 100644 python/sglang/jit_kernel/activation.py delete mode 100644 python/sglang/jit_kernel/dsv3_fused_a_gemm.py delete mode 100644 python/sglang/jit_kernel/dsv3_router_gemm.py delete mode 100644 python/sglang/jit_kernel/norm.py delete mode 100644 python/sglang/jit_kernel/per_tensor_quant_fp8.py delete mode 100644 python/sglang/jit_kernel/per_token_group_quant.py delete mode 100644 python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py delete mode 100644 python/sglang/jit_kernel/set_mla_kv_buffer.py delete mode 100644 python/sglang/jit_kernel/tests/kv_canary/__init__.py rename python/sglang/{jit_kernel => kernels/jit}/.clang-format (100%) rename python/sglang/{jit_kernel => kernels/jit}/__main__.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/benchmark/kv_canary/utils.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/benchmark/marker.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/benchmark/utils.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/add_constant.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/attention/fixup_zero_kv.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/attention/fused_fp8_qkv_kv_cache.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v32/indexer_k.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c128.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c128_online.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c128_online_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c128_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c4.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c4_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/c_plan.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/common.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/fused_norm_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/fused_norm_rope_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/hash_topk.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/main_norm_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/online_c128_mtp.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/paged_mqa_metadata.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/store.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/topk_v1.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/deepseek_v4/topk_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/causal_conv3d_cat_pad.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/ltx2_qknorm_split_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/norm_scale_shift.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/qknorm_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/residual_gate_add.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/diffusion/timestep_embedding.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/distributed/communicator.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/distributed/custom_all_reduce.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/distributed/ipc.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/distributed/tp_qknorm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/dsa/fused_store_index_cache.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/dsa/kpool_topk_transform.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/activation.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/clamp_position.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/concat_mla.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/fused_add_rmsnorm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/fused_eh_norm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/fused_metadata_copy.cuh (99%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/fused_qknorm_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/kvcache.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/pos_enc.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/qknorm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/qknorm_across_heads.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/resolve_future_token_ids.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/rmsnorm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/rmsnorm_hf.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/elementwise/set_mla_kv_buffer.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/code_gen.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/fast_hadamard_transform.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/hadamard_jit.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/fast-hadamard-transform/static_switch.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/awq_dequantize.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/dsv3_fused_a_gemm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/dsv3_router_gemm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/awq_marlin_repack.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/dequant.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/gptq_marlin.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/gptq_marlin_repack.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/kernel.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/marlin.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/marlin_dtypes.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin/marlin_template.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin_moe/kernel.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin_moe/marlin_template.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/per_tensor_quant_fp8.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/per_token_group_quant.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/gemm/per_token_group_quant_8bit_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/hisparse.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/causal_conv1d.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/draft_extend_sconv.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/fused_decode_update.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/gather_scatter_sconv.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_all_reduce.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_ar_barrier.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_ar_fused_decode.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_ar_scattered_sconv.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_attn_prologue_fused.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_rel_proj.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/inkling_row_scale.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/inkling/update_sconv_cache.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kv_canary/canary_common.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kv_canary/canary_plan_entries.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kv_canary/canary_verify.cuh (99%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kv_canary/canary_write.cuh (99%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kv_canary/consts.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kvcacheio/hicache.cuh (99%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kvcacheio/relayout.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kvcacheio/staged_write_back.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/kvcacheio/transfer_mamba.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/lora/moe_lora_align_kernel.cu (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/lplb/dispatch_probability.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/lplb/ipm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/lplb/lp_post.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/lplb/lp_prep.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/minimax/fused_gemma_qknorm_rope.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/minimax/fused_store_kv_index.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/minimax/minimax_decode_topk.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/minimax/per_token_quant_ue8m0.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/inkling_gate_topk_renorm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/moe_align_kernel.cu (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/moe_finalize_fuse_shared.cu (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/moe_fused_gate.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/moe_permute_prepare.cu (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/moe_topk_sigmoid.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/moe/tvm_ffi_utils.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/ngram.cpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/ngram.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/ngram_corpus_ffi.cpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/param.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/queue.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/result.cpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/result.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/suffix_automaton.cpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/suffix_automaton.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/trie.cpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_corpus/trie.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/ngram_embedding.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/config.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/defines.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_transpose_v.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_utils.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/entry.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/helpers.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/kernel.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/sparse_mla_q8kv8_prefill_sm90/params.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu (100%) rename python/sglang/{jit_kernel => kernels/jit}/csrc/trtllm_lora_temp/topk_softmax_pack.cuh (99%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/atomic.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/cta.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/deepseek_v4/compress.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/deepseek_v4/compress_v2.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/deepseek_v4/fp8_utils.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/deepseek_v4/kvcacheio.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/deepseek_v4/topk_impl.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/distributed/communicator.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/ffi.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/impl/norm.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/math.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/runtime.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/scalar_type.hpp (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/source_location.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/tensor.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/tile.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/type.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/utils.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/utils.h (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/vec.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit}/include/sgl_kernel/warp.cuh (100%) rename python/sglang/{jit_kernel => kernels/jit/tests/deepseek_v4}/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/deepseek_v4/common.py (100%) rename python/sglang/{jit_kernel/tests/deepseek_v4 => kernels/jit/tests/kv_canary}/__init__.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_canary_helpers.py (98%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_constants.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_differential.py (99%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_fixtures.py (98%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_fuzz_driver.py (95%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_hand_oracle.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/kv_canary/_invariants.py (99%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_minimax_m3_mxfp8.py (99%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_minimax_m3_rmsnorm.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_moe_preprocess.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_moe_topk_sigmoid.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_sconv_decode_metadata.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/test_sconv_extend_metadata.py (100%) rename python/sglang/{jit_kernel => kernels/jit}/tests/utils.py (100%) diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index c36f9eeff..6b9758041 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -23,7 +23,7 @@ Add a new operation that scales each element of a tensor by a scalar factor: --- -## Common Abstractions in `python/sglang/jit_kernel/include/sgl_kernel/` +## Common Abstractions in `python/sglang/kernels/jit/include/sgl_kernel/` **Always prefer these abstractions over raw CUDA primitives.** They provide safety, readability, and consistency with the rest of the codebase. @@ -191,16 +191,16 @@ LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); ## Step 0 (optional): Generate a `.clangd` config for better IDE support ```bash -python -m sglang.jit_kernel -h # for verbose help info about clangd configuration -python -m sglang.jit_kernel -python -m sglang.jit_kernel --dep cutlass flashinfer # with cutlass/flashinfer dependency +python -m sglang.kernels.jit -h # for verbose help info about clangd configuration +python -m sglang.kernels.jit +python -m sglang.kernels.jit --dep cutlass flashinfer # with cutlass/flashinfer dependency ``` --- -## Step 1: Implement the CUDA kernel in `jit_kernel/csrc/` +## Step 1: Implement the CUDA kernel in `kernels/jit/csrc/` -Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`. +Create `python/sglang/kernels/jit/csrc/elementwise/scale.cuh`. The implementation fully uses the project abstractions described above: @@ -330,9 +330,9 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { --- -## Step 2: Add the Python wrapper in `jit_kernel/` +## Step 2: Add the Python wrapper in `kernels/jit/` -Create `python/sglang/jit_kernel/scale.py`: +Create `python/sglang/kernels/jit/scale.py`: ```python from __future__ import annotations @@ -341,7 +341,7 @@ from typing import TYPE_CHECKING import torch -from sglang.jit_kernel.utils import ( +from sglang.kernels.jit.utils import ( cache_once, is_arch_support_pdl, load_jit, @@ -438,10 +438,10 @@ if torch.cuda.get_device_capability()[0] < 9: ## Step 4: Write tests (required) -JIT kernel correctness tests and benchmarks live under `test/registered/jit/` and `test/registered/jit/benchmark/` (NOT inside the `sglang` package -- a `register_*_ci(...)` call anywhere under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook). Only their test-only helpers (e.g. `benchmark/marker.py`) stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path. **CI does not run `pytest` in those directories directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` and `bench_*.py` under `test/registered/`, collects `register_*_ci(...)` calls by **statically parsing each file's AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check. +JIT kernel correctness tests and benchmarks live under `test/registered/jit/` and `test/registered/jit/benchmark/` (NOT inside the `sglang` package -- a `register_*_ci(...)` call anywhere under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook). Only their test-only helpers (e.g. `benchmark/marker.py`) stay alongside the kernel source under `python/sglang/kernels/jit/` and are imported by absolute path. **CI does not run `pytest` in those directories directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` and `bench_*.py` under `test/registered/`, collects `register_*_ci(...)` calls by **statically parsing each file's AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check. - **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `base-b-kernel-unit-test-1-gpu-large` on H100 and `base-b-kernel-unit-test-4-gpu-b200` on B200/SM100 paths (see `.github/workflows/pr-test-jit-kernel.yml`). Multi-GPU JIT tests use `base-b-kernel-unit-test-8-gpu-h200`. -- **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/jit_kernel/utils.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`). +- **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/kernels/jit/utils/compile.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`). Registration pattern (module level, **literal** `est_time`, `stage`, and `runner_config` values — required for AST parsing): @@ -477,7 +477,7 @@ Create `test/registered/jit/test_scale.py`: ```python import pytest import torch -from sglang.jit_kernel.scale import scale +from sglang.kernels.jit.scale import scale from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -527,7 +527,7 @@ if __name__ == "__main__": Benchmarks are `bench_*.py` files under `test/registered/jit/benchmark/`. They are picked up by the same `run_suite.py` machinery as unit tests. Register them for **`base-b-kernel-benchmark-test-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large`). -Benchmarks use the project's own `marker` framework (in `python/sglang/jit_kernel/benchmark/marker.py`) — **do not** use `triton.testing.perf_report` / `triton.testing.do_bench` directly. The marker framework provides (public names: `benchmark`, `parametrize`, `do_bench`, `skip`, `BenchResult`, `BenchSkip`): +Benchmarks use the project's own `marker` framework (in `python/sglang/kernels/jit/benchmark/marker.py`) — **do not** use `triton.testing.perf_report` / `triton.testing.do_bench` directly. The marker framework provides (public names: `benchmark`, `parametrize`, `do_bench`, `skip`, `BenchResult`, `BenchSkip`): - **`@marker.benchmark(line_arg, line_vals, *, unit="us")`** — the **innermost** decorator (bottom of the stack, directly above `def benchmark`). Declares the column axis: each value in `line_vals` becomes a result column, and `line_arg` is the parameter name passed into the benchmark function. `unit` is one of `"us" | "ms" | "s"`. - **`@marker.parametrize(names, vals, ci_vals=None)`** — stackable decorator that adds a row axis (pytest-style). Each `@parametrize` adds one (or more, correlated) parameter the benchmark is swept over (Cartesian product across all `parametrize` decorators). `names` may be a single name (`"size"`) or a comma-separated correlated tuple axis (`"h,d"`, with `vals` then a list of tuples like `[(1, 64), (2, 128)]`). Pass the optional third `ci_vals` for a smaller sweep that is auto-selected under `is_in_ci()` — this is the built-in CI-shrinking mechanism, so you usually don't need `get_benchmark_range` for swept axes. @@ -547,9 +547,9 @@ Create `test/registered/jit/benchmark/bench_scale.py`: ```python import torch -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random -from sglang.jit_kernel.scale import scale as jit_scale +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.jit.scale import scale as jit_scale from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large") @@ -614,7 +614,7 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1 ## Troubleshooting - **`No CI registry found in ...` from `run_suite.py`**: add a module-level `register_cuda_ci(...)` with literal `est_time`, `stage`, and `runner_config` (and optional `nightly=True`); starred args and non-literal values break AST collection -- **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/jit_kernel/csrc/`; reduce template argument combinations +- **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/kernels/jit/csrc/`; reduce template argument combinations - **CUDA crash / illegal memory access**: `CUDA_LAUNCH_BLOCKING=1`; `compute-sanitizer --tool memcheck python ...` - **Unstable benchmark results**: `marker.do_bench` uses CUDA-graph-based timing by default; set `use_cuda_graph=False` only if the kernel can't be captured. Make sure `graph_clone_args` covers every *read* tensor — reusing a single buffer keeps it L2-hot and skews results - **Missing GB/s column**: the column is on by default; check that `SGLANG_KERNEL_DISABLE_LOG_BANDWIDTH` is not `1` and `disable_log_bandwidth` is not `True`. For in-place kernels (return `None`) the `memory_output="out"` default counts nothing — pass the written tensors via `memory_output=(...)` @@ -626,30 +626,30 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1 - `docs_new/docs/developer_guide/development_jit_kernel_guide.mdx` - `test/run_suite.py` — suite names, discovery of `test/registered/`, execution entrypoint for CI - `python/sglang/test/ci/ci_register.py` — `register_cuda_ci` and AST registration rules -- `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args`, `should_run_full_tests`, `get_ci_test_range` -- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device` -- `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` -- `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector` -- `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory` -- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `DTypeTrait`, `packed_t`, `device::cast`, `device::unpack`, `ReductionTrait` -- `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::` -- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce` and `reduce_sum/max/min` wrappers -- `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max` -- `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max` -- `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers -- `python/sglang/jit_kernel/csrc/add_constant.cuh` — minimal runnable reference -- `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh` — real example using `TensorMatcher` + `LaunchKernel` + `tile::Memory` -- `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern -- `python/sglang/jit_kernel/benchmark/marker.py` — `benchmark`, `parametrize`, `do_bench`, `BenchResult` -- `python/sglang/jit_kernel/benchmark/utils.py` — `create_random` / `create_empty` / `get_benchmark_range` helpers and `DEFAULT_DTYPE` / `DEFAULT_DEVICE` +- `python/sglang/kernels/jit/utils/compile.py` — `cache_once`, `load_jit`, `make_cpp_args`, `should_run_full_tests`, `get_ci_test_range` +- `python/sglang/kernels/jit/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device` +- `python/sglang/kernels/jit/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` +- `python/sglang/kernels/jit/include/sgl_kernel/vec.cuh` — `AlignedVector` +- `python/sglang/kernels/jit/include/sgl_kernel/tile.cuh` — `tile::Memory` +- `python/sglang/kernels/jit/include/sgl_kernel/type.cuh` — `DTypeTrait`, `packed_t`, `device::cast`, `device::unpack`, `ReductionTrait` +- `python/sglang/kernels/jit/include/sgl_kernel/math.cuh` — `device::math::` +- `python/sglang/kernels/jit/include/sgl_kernel/warp.cuh` — `warp::reduce` and `reduce_sum/max/min` wrappers +- `python/sglang/kernels/jit/include/sgl_kernel/cta.cuh` — `cta::reduce_max` +- `python/sglang/kernels/jit/include/sgl_kernel/atomic.cuh` — `atomic::max` +- `python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers +- `python/sglang/kernels/jit/csrc/add_constant.cuh` — minimal runnable reference +- `python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh` — real example using `TensorMatcher` + `LaunchKernel` + `tile::Memory` +- `python/sglang/kernels/jit/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern +- `python/sglang/kernels/jit/benchmark/marker.py` — `benchmark`, `parametrize`, `do_bench`, `BenchResult` +- `python/sglang/kernels/jit/benchmark/utils.py` — `create_random` / `create_empty` / `get_benchmark_range` helpers and `DEFAULT_DTYPE` / `DEFAULT_DEVICE` - `test/registered/jit/benchmark/bench_qknorm.py` — real example: multi-axis `parametrize` (with `ci_vals`) + in-place `memory_output` - `test/registered/jit/benchmark/bench_store_cache.py` — real example: scoped `memory_args` / `memory_output` + selective `graph_clone_args` ## Summary of Files Created ``` -python/sglang/jit_kernel/csrc/elementwise/scale.cuh # NEW: CUDA kernel -python/sglang/jit_kernel/scale.py # NEW: Python wrapper +python/sglang/kernels/jit/csrc/elementwise/scale.cuh # NEW: CUDA kernel +python/sglang/kernels/jit/scale.py # NEW: Python wrapper test/registered/jit/test_scale.py # NEW: Tests test/registered/jit/benchmark/bench_scale.py # NEW: Benchmark ``` diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index 559b8751f..7bd17d589 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -18,7 +18,7 @@ Add a new operation that scales each element of a tensor by a scalar factor: ## Two rules of thumb (must follow) -1. **Prefer `python/sglang/jit_kernel` first** when the kernel does **not** depend on CUTLASS or another large C++ project. This is the default path for lightweight kernels that benefit from rapid iteration. +1. **Prefer `python/sglang/kernels/jit` first** when the kernel does **not** depend on CUTLASS or another large C++ project. This is the default path for lightweight kernels that benefit from rapid iteration. 2. **Prefer `sgl-kernel`** when the kernel **does** depend on CUTLASS or another large C++ project, or when it should be part of the AOT wheel / torch op registration flow. 3. **Exception**: if the dependency is `flashinfer`, or CUTLASS that is already provided through `flashinfer`, the kernel can still be implemented as `jit_kernel`. diff --git a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md index 97cc097b7..c15542dca 100644 --- a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md @@ -50,14 +50,14 @@ in-flight row as shipped. | AITER allreduce fusion | ROCm all-reduce plus RMSNorm still split | `python/sglang/srt/layers/layernorm.py::forward_with_allreduce_fusion`
`python/sglang/srt/distributed/communication_op.py::tensor_model_parallel_fused_allreduce_rmsnorm`
`python/sglang/srt/layers/communicator.py::apply_aiter_all_reduce_fusion` | ROCm-side fused TP all-reduce + RMSNorm with fallback to plain all-reduce plus norm | On AMD, rule out existing AITER fusion before proposing a new communication fusion. | | Fused activation-and-mul (`SwiGLU` / `GeGLU`) | `silu_and_mul`
`gelu_and_mul`
`npu_swiglu` | `python/sglang/srt/layers/activation.py` | Single op covers activation plus elementwise multiply across CUDA / CPU / NPU / XPU backends | Treat separate activation + mul on packed MLP outputs as missing existing fusion. | | Fused dual residual RMSNorm | residual add plus two RMSNorm-like kernels around Grok blocks | `python/sglang/srt/layers/elementwise.py::fused_dual_residual_rmsnorm`
`python/sglang/srt/models/grok.py` | One Triton kernel computes intermediate residual update and next RMSNorm output together | On Grok-like residual layouts, treat split residual + norm as missing existing fusion. | -| In-place QK RMSNorm | split `q_norm` / `k_norm` kernels | `python/sglang/srt/models/utils.py::apply_qk_norm`
`python/sglang/jit_kernel/norm.py::fused_inplace_qknorm` | In-place JIT QK norm plus optional `alt_stream` overlap for K | Check shape, dtype, deterministic mode, and in-place legality before proposing a new QK fuse. | +| In-place QK RMSNorm | split `q_norm` / `k_norm` kernels | `python/sglang/srt/models/utils.py::apply_qk_norm`
`python/sglang/kernels/ops/layernorm/_jit_norm.py::fused_inplace_qknorm` | In-place JIT QK norm plus optional `alt_stream` overlap for K | Check shape, dtype, deterministic mode, and in-place legality before proposing a new QK fuse. | | TorchInductor horizontal Q/K norm combo-kernels | `combo_kernels`
`benchmark_combo_kernel`
`q_norm`
`k_norm`
`split_with_sizes` | `torch._inductor.config.combo_kernels` | TorchInductor can horizontally fuse sibling Q-norm and K-norm kernels in compiled traces, often deleting `split_with_sizes` / `clone` ladders | Treat separate Q/K norm ladders in compile-heavy traces as an existing compiler-fusion family first. | | MiniMax TP fused QK RMSNorm | `MiniMaxM2RMSNormTP`
`rms_sumsq_serial`
`rms_apply_serial`
`forward_qk` | `python/sglang/srt/models/minimax_m2.py` | Triton kernels compute Q / K sumsq together, TP all-reduces shared stats, then apply both RMSNorms together | On MiniMax traces, separate Q norm and K norm are usually a missed model-specific Triton fusion. | -| Fused QK RMSNorm + RoPE | `qknorm*` + `rope*` + `rotary*` as separate steps | `python/sglang/jit_kernel/fused_qknorm_rope.py`
`python/sglang/srt/models/qwen3_moe.py` | One JIT kernel applies QK RMSNorm and RoPE in-place on packed QKV | For compatible LLMs, classify split QK norm + RoPE as a missing existing fusion. | +| Fused QK RMSNorm + RoPE | `qknorm*` + `rope*` + `rotary*` as separate steps | `python/sglang/kernels/ops/attention/fused_qknorm_rope.py`
`python/sglang/srt/models/qwen3_moe.py` | One JIT kernel applies QK RMSNorm and RoPE in-place on packed QKV | For compatible LLMs, classify split QK norm + RoPE as a missing existing fusion. | | Fused QK RoPE reshape + KV cache write | `fused_qk_rope_reshape_and_cache*`
RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. | -| Fused RoPE + KV cache store | `fused_set_kv_buffer`
RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`
`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. | +| Fused RoPE + KV cache store | `fused_set_kv_buffer`
RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/kernels/ops/attention/rope.py`
`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. | | Fused decode metadata setup | `normal_decode_set_metadata`
`cache_seqlens_int32`
`cu_seqlens_k`
`page_table`
`swa_page_table` | `python/sglang/srt/layers/attention/flashattention_backend.py::normal_decode_set_metadata` | Triton decode path fuses seq-len cast/add, prefix-sum, req-to-token gather, page-table divide, and optional SWA metadata build into 1-2 kernels | If decode exposes multiple tiny metadata kernels before attention, first compare against this existing fused metadata-prep path. | -| NSA fused metadata copy for graph replay | `fused_metadata_copy`
`fused_metadata_copy_multi`
`fused_nsa_cache_seqlens`
`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around NSA replay as a missed existing replay fusion. | +| NSA fused metadata copy for graph replay | `fused_metadata_copy`
`fused_metadata_copy_multi`
`fused_nsa_cache_seqlens`
`fused_flashmla_metadata` | `python/sglang/kernels/ops/attention/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around NSA replay as a missed existing replay fusion. | | DeepSeek MLA fused projection + norm + RoPE | `qkv_proj_with_rope_fused_weight`
`fused_qkv_a_proj_with_mqa`
`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`
`python/sglang/srt/models/deepseek_v2.py` | CPU / ROCm paths fuse DeepSeek MLA projection packing with q / k norm, RoPE, and cache-oriented MLA prep | For DeepSeek MLA, split proj / norm / rope prep is usually an existing backend-specific fuse that did not fire. | | Fused QK RoPE concat + MLA cache write | `fused_qk_rope_cat_and_cache_mla`
`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | ROCm MLA path can fuse Q / K RoPE packing, concat, and MLA cache write in one backend-specific op | On DeepSeek / MLA traces, separate RoPE-cat-cache steps are not automatically novel. | | Qwen3 decode fused QK norm + 3D mRoPE + KV cache write | `fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`
`mrope`
decode cache write | `python/sglang/srt/models/qwen3.py` | ROCm / AITER decode path fuses QK norm, 3D mRoPE, and paged KV cache write | On Qwen3-style decode, separate norm + mRoPE + cache-store kernels are not a novel opportunity. | @@ -69,13 +69,13 @@ in-flight row as shipped. | Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`
`fused_append_shared_experts_with_weights`
`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`
`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. | | Fused MoE dispatch / permute / combine | token permutation
dispatch / combine
grouped top-k
many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | `FusedMoE` plus DeepEP / FlashInfer / FuseEP / standard dispatch backends and `permute_fusion=True` | First ask whether the model is missing an existing `FusedMoE`-style path or backend-specific dispatcher path. | | Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. | -| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | +| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | | DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`
`fused_rms_mxfp4_quant`
`fused_flatten_fp8_group_quant`
`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. | | NSA fused top-k transform / page-table build | `fast_topk_transform_fused`
`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/nsa_backend.py` | NSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If NSA top-k metadata work is split, check `SGLANG_NSA_FUSE_TOPK` and backend support first. | -| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`
`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. | +| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/kernels/ops/attention/fused_store_index_cache.py`
`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. | | Fused sampling temperature + softmax | `fused_temperature_softmax*` | `python/sglang/srt/layers/fused_sampling.py`
`python/sglang/srt/layers/sampler.py` | Triton single-pass / multi-pass kernels fuse temperature scaling and softmax during decode | Separate temp-divide + softmax at decode batch sizes is often a missed existing fusion. | | Fused logit softcap | `fused_softcap`
`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`
`python/sglang/srt/layers/logits_processor.py` | Triton kernels fuse cast-to-float and softcap / tanh math for logits or generic elementwise softcapping | Treat exposed cast + softcap ladders as an existing Triton fuse family. | -| Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`
`qkvz_proj`
`ba_proj`
`qkvabz_proj`
`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`
`python/sglang/srt/models/qwen3_next.py`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/kimi_linear.py`
`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. | +| Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`
`qkvz_proj`
`ba_proj`
`qkvabz_proj`
`fused_qkvbfg_a_proj` | `python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py`
`python/sglang/srt/models/qwen3_next.py`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/kimi_linear.py`
`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. | | Fused GDN gating prep | `fused_gdn_gating`
`softplus`
`beta_output` | `python/sglang/srt/layers/attention/fla/fused_gdn_gating.py` | Triton kernel computes GDN gate preparation such as `-exp(A_log) * softplus(...)` and `sigmoid(b)` together | On GDN traces, treat split gate-prep elementwise kernels as missing existing fusion first. | | Fused RMSNorm-gated linear-attention output | `FusedRMSNormGated`
`layer_norm_gated_fwd` | `python/sglang/srt/layers/attention/fla/fused_norm_gate.py`
`python/sglang/srt/models/qwen3_next.py`
`python/sglang/srt/models/kimi_linear.py` | One Triton op covers residual-aware (RMS)Norm plus sigmoid / swish gating | If norm and output gate appear as separate kernels in GDN / Kimi-like blocks, first suspect a missing existing fusion. | | Fused gated RMSNorm / LayerNorm | `rms_norm_gated`
`layer_norm_gated` | `python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py` | Mamba-derived kernels can fuse normalization with the gating branch `z * sigmoid(z)` | Treat split norm and gate post-processing on Mamba-style blocks as an existing fusion family. | @@ -111,15 +111,15 @@ in-flight row as shipped. | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | --- | --- | --- | --- | --- | -| Fused residual + norm + scale + shift | residual add, norm, scale, shift, gate around DiT blocks | `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_scale_residual_norm_scale_shift(...)` | Treat split residual + norm + modulation as a missing existing diffusion fusion first. | -| Fused norm + scale + shift | norm followed by scale / shift elementwise kernels | `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_norm_scale_shift(...)` | Existing modulation fusion already covers this family. | -| Triton scale / shift and gate-select kernels | tiny scale / shift or gate-select kernels dominate modulation blocks | `python/sglang/jit_kernel/diffusion/triton/scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/elementwise.py` | `fuse_scale_shift_kernel(...)` and `fuse_layernorm_scale_shift_gate_select01_kernel(...)` | Check whether the runtime is missing these existing Triton fusions. | -| Fused add-RMSNorm and one-pass RMSNorm | residual add plus RMSNorm still split on short hidden sizes | `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
`python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` | `fused_add_rmsnorm(...)` and `triton_one_pass_rms_norm(...)` | For short hidden-size diffusion blocks, this is already an established fusion family. | -| Fused diffusion QK norm + RoPE | split QK norm and RoPE in diffusion attention blocks | `python/sglang/jit_kernel/diffusion/qknorm_rope.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py::apply_qk_norm_rope` | `fused_inplace_qknorm_rope(...)`, with fallback to QK norm plus `apply_flashinfer_rope_qk_inplace(...)` | Distinguish between missing fused qknorm + rope and the existing FlashInfer RoPE fallback. | -| Z-Image fused `norm(x) * tanh(scale) + shift` | `fused_norm_tanh_mul_add`
`tanh(gate) * rmsnorm(x)` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | CuTeDSL kernel plus runtime helper for Z-Image residual-form modulation | Treat split Z-Image residual-form modulation as a missing existing diffusion fusion, not a novel idea. | -| Z-Image fused residual modulation + next norm-scale | `fused_norm_tanh_mul_add_norm_scale`
`residual + tanh(gate) * rmsnorm(x)`
`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/models/dits/zimage.py` | One CuTeDSL kernel fuses the first residual-form modulation and the next normalization / scale stage | If you see this chain split in Z-Image traces, report it as a missing existing mainline fusion family. | -| LTX2 fused Ada values | `ltx2_ada_values9`
`get_ada_values`
`scale_shift_table + timestep.reshape` | `python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py`
`python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py` | PR `#29390` fuses LTX-2.3 Ada value materialization for video/audio streams and reuses the 9 Ada tensors across self-attention, MLP, and prompt-cross-attention blocks | Treat repeated Ada add/reshape/slice ladders in LTX2 traces as a missing shipped SGLang fusion first. | -| LTX2 residual-gate add | `diffusion_residual_gate_add`
`residual_gate_add`
`residual + update * gate` | `python/sglang/jit_kernel/diffusion/residual_gate_add.py`
`python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh`
`python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py` | PR `#29361` fuses LTX2 `residual + update * gate` sites for attention, cross-attention, and feed-forward updates into one CUDA custom op when dtype, shape, device, and contiguity guards pass | Treat split add/mul gate ladders in LTX2 traces as a missing shipped SGLang fusion first. | +| Fused residual + norm + scale + shift | residual add, norm, scale, shift, gate around DiT blocks | `python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_scale_residual_norm_scale_shift(...)` | Treat split residual + norm + modulation as a missing existing diffusion fusion first. | +| Fused norm + scale + shift | norm followed by scale / shift elementwise kernels | `python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | `fused_norm_scale_shift(...)` | Existing modulation fusion already covers this family. | +| Triton scale / shift and gate-select kernels | tiny scale / shift or gate-select kernels dominate modulation blocks | `python/sglang/kernels/ops/diffusion/triton/scale_shift.py`
`python/sglang/multimodal_gen/runtime/layers/elementwise.py` | `fuse_scale_shift_kernel(...)` and `fuse_layernorm_scale_shift_gate_select01_kernel(...)` | Check whether the runtime is missing these existing Triton fusions. | +| Fused add-RMSNorm and one-pass RMSNorm | residual add plus RMSNorm still split on short hidden sizes | `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
`python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py` | `fused_add_rmsnorm(...)` and `triton_one_pass_rms_norm(...)` | For short hidden-size diffusion blocks, this is already an established fusion family. | +| Fused diffusion QK norm + RoPE | split QK norm and RoPE in diffusion attention blocks | `python/sglang/kernels/ops/diffusion/qknorm_rope.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py::apply_qk_norm_rope` | `fused_inplace_qknorm_rope(...)`, with fallback to QK norm plus `apply_flashinfer_rope_qk_inplace(...)` | Distinguish between missing fused qknorm + rope and the existing FlashInfer RoPE fallback. | +| Z-Image fused `norm(x) * tanh(scale) + shift` | `fused_norm_tanh_mul_add`
`tanh(gate) * rmsnorm(x)` | `python/sglang/kernels/ops/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/layers/layernorm.py` | CuTeDSL kernel plus runtime helper for Z-Image residual-form modulation | Treat split Z-Image residual-form modulation as a missing existing diffusion fusion, not a novel idea. | +| Z-Image fused residual modulation + next norm-scale | `fused_norm_tanh_mul_add_norm_scale`
`residual + tanh(gate) * rmsnorm(x)`
`ffn_norm1(x) * scale_mlp` | `python/sglang/kernels/ops/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`
`python/sglang/multimodal_gen/runtime/models/dits/zimage.py` | One CuTeDSL kernel fuses the first residual-form modulation and the next normalization / scale stage | If you see this chain split in Z-Image traces, report it as a missing existing mainline fusion family. | +| LTX2 fused Ada values | `ltx2_ada_values9`
`get_ada_values`
`scale_shift_table + timestep.reshape` | `python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py`
`python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py` | PR `#29390` fuses LTX-2.3 Ada value materialization for video/audio streams and reuses the 9 Ada tensors across self-attention, MLP, and prompt-cross-attention blocks | Treat repeated Ada add/reshape/slice ladders in LTX2 traces as a missing shipped SGLang fusion first. | +| LTX2 residual-gate add | `diffusion_residual_gate_add`
`residual_gate_add`
`residual + update * gate` | `python/sglang/kernels/ops/diffusion/residual_gate_add.py`
`python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh`
`python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py` | PR `#29361` fuses LTX2 `residual + update * gate` sites for attention, cross-attention, and feed-forward updates into one CUDA custom op when dtype, shape, device, and contiguity guards pass | Treat split add/mul gate ladders in LTX2 traces as a missing shipped SGLang fusion first. | | Nunchaku fused GELU MLP | `_fused_gelu_mlp`
`fused_gelu_mlp` | `python/sglang/multimodal_gen/runtime/models/dits/flux.py` | Nunchaku path fuses `fc1 GEMM + GELU + shift + re-quant + fc2.lora_down` before the second GEMM | Treat split GELU-MLP on Nunchaku checkpoints as an existing fused family, not a new discovery. | ## 5. Diffusion kernel-overlap and async-communication families @@ -143,10 +143,10 @@ Stable entries should be folded into the mainline family rows above. | PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`
`fp4 kv cache` | `PR #21889`
`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 NSA KV pages directly while quantizing BF16 input | If NSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. | | PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`
`WRITE_PT`
`dequant_fp4_paged_decode` | `PR #21889`
`python/sglang/srt/layers/attention/nsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. | | PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`
`trtllm_fp8_block_scale_moe` | `PR #21491`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. | -| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`
`per_token_quant_fp8` | `PR #22005`
`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`
`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. | +| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`
`per_token_quant_fp8` | `PR #22005`
`python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`
`python/sglang/kernels/jit/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. | | PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`
`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`
`rotary_dim` | `PR #20667`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. | | PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`
`fp8_scaled_mm`
`nvjet`
`cudaMemsetAsync` | `PR #22392`
`sgl-kernel/python/sgl_kernel/gemm.py`
`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. | -| PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`
`nvfp4 expert quant`
`cutlass moe` | `PR #18612`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/jit_kernel/nvfp4.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. | +| PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`
`nvfp4 expert quant`
`cutlass moe` | `PR #18612`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. | | PR `#22918` FlashInfer per-token NVFP4 MoE | `per_token_nvfp4`
`trtllm_fp4_block_scale_moe`
`FlashInfer MoE` | `PR #22918`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | Adds FlashInfer-backed per-token NVFP4 MoE execution so expert quant/dequant work can move into the fused MoE backend | Treat standalone per-token NVFP4 MoE support kernels as a candidate missing backend-selection path, not an automatically novel kernel idea. | | PR `#22851` NSA top-k backend and FlashInfer / PyTorch top-k split | `nsa topk`
`flashinfer_topk`
`pytorch_topk`
`fast_topk_transform` | `PR #22851`
`python/sglang/srt/layers/attention/nsa_backend.py` | Makes NSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When NSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. | | PR `#24125` GLM5 NSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`
`GLM-5`
`NSA`
`TileLang decode` | `PR #24125`
`python/sglang/srt/layers/attention/nsa_backend.py` | Skips redundant cat/copy work in the GLM5 NSA TileLang decode path | Treat cat/copy bursts in GLM5 NSA decode as a concrete in-flight cleanup opportunity. | diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py index 7fddd29b5..11f3a06b7 100644 --- a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py @@ -452,7 +452,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="In-place QK RMSNorm", candidate_path=( - "python/sglang/srt/models/utils.py" "
python/sglang/jit_kernel/norm.py" + "python/sglang/srt/models/utils.py" + "
python/sglang/kernels/ops/layernorm/_jit_norm.py" ), active_keywords=("fused_inplace_qknorm", "minimaxm2rmsnormtp"), split_groups=(("apply_qk_norm", "q_norm", "k_norm", "qknorm"),), @@ -466,7 +467,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="Fused QK RMSNorm + RoPE", candidate_path=( - "python/sglang/jit_kernel/fused_qknorm_rope.py" + "python/sglang/kernels/ops/attention/fused_qknorm_rope.py" "
python/sglang/srt/models/qwen3_moe.py" ), active_keywords=("fused_qknorm_rope", "fused_qk_norm_rope"), @@ -499,7 +500,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="Fused RoPE + KV cache store", candidate_path=( - "python/sglang/jit_kernel/rope.py" "
python/sglang/srt/models/utils.py" + "python/sglang/kernels/ops/attention/rope.py" + "
python/sglang/srt/models/utils.py" ), active_keywords=("fused_set_kv_buffer",), split_groups=( @@ -531,7 +533,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ), FusionPatternSpec( pattern="NSA fused metadata copy for graph replay", - candidate_path="python/sglang/jit_kernel/fused_metadata_copy.py", + candidate_path="python/sglang/kernels/ops/attention/fused_metadata_copy.py", active_keywords=( "fused_metadata_copy", "fused_metadata_copy_multi", @@ -698,7 +700,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( pattern="Fused MoE activation + quant / re-quant", candidate_path=( "python/sglang/srt/layers/moe/ep_moe/kernels.py" - "
python/sglang/jit_kernel/nvfp4.py" + "
python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py" "
python/sglang/srt/layers/moe/cutlass_w4a8_moe.py" ), active_keywords=( @@ -759,7 +761,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="NSA fused quantize + indexed K-cache store", candidate_path=( - "python/sglang/jit_kernel/fused_store_index_cache.py" + "python/sglang/kernels/ops/attention/fused_store_index_cache.py" "
python/sglang/srt/layers/attention/nsa/nsa_indexer.py" ), active_keywords=("fused_store_index_k_cache",), @@ -859,7 +861,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( pattern="SGLang LTX2 fused Ada values", candidate_path=( "PR #29390" - "
python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py" + "
python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py" "
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py" ), active_keywords=( @@ -886,8 +888,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( pattern="SGLang LTX2 residual-gate add CUDA fast path", candidate_path=( "PR #29361" - "
python/sglang/jit_kernel/diffusion/residual_gate_add.py" - "
python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh" + "
python/sglang/kernels/ops/diffusion/residual_gate_add.py" + "
python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh" "
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py" ), active_keywords=( diff --git a/.claude/skills/write-sglang-test/SKILL.md b/.claude/skills/write-sglang-test/SKILL.md index d6bbd86e9..2c22f99c4 100644 --- a/.claude/skills/write-sglang-test/SKILL.md +++ b/.claude/skills/write-sglang-test/SKILL.md @@ -16,7 +16,7 @@ This skill covers **how to write and register tests**. For CI pipeline internals 5. **Prefer mock over real server** — when testing logic that doesn't need a server / engine launch (middleware, request routing, config validation, argument parsing), use `unittest.mock.patch` / `MagicMock` and place tests in `test/registered/unit/`. Only launch a real server when the test genuinely needs inference results or server lifecycle behavior. JIT kernel notes: -- If the task is adding or updating code under `python/sglang/jit_kernel/`, prefer the `add-jit-kernel` skill first. +- If the task is adding or updating code under `python/sglang/kernels/jit/`, prefer the `add-jit-kernel` skill first. - JIT kernel correctness tests use `test/registered/jit/**/test_*.py`. - JIT kernel benchmarks use `test/registered/jit/benchmark/**/bench_*.py`. - Those files are executed by `test/run_suite.py` through dedicated kernel suites (`base-b-kernel-*`); a `register_*_ci(...)` call placed under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook. @@ -391,7 +391,7 @@ test/ ├── manual/ # Non-CI: debugging, one-off, manual verification └── run_suite.py # CI runner (scans registered/ plus jit_kernel test/benchmark files) -python/sglang/jit_kernel/ +python/sglang/kernels/jit/ ├── tests/ # JIT kernel correctness tests (CI-discovered by test/run_suite.py) └── benchmark/ # JIT kernel benchmarks (CI-discovered by test/run_suite.py) ``` @@ -443,7 +443,7 @@ Before submitting a test: - [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`) - [ ] Has `register_*_ci(...)` call at module level - [ ] Placed in `test/registered//` (JIT kernel test/benchmark → `test/registered/jit/` or `test/registered/jit/benchmark/`) -- [ ] JIT kernel work: test files live in `test/registered/jit/`; only test-only helpers stay under `python/sglang/jit_kernel/` +- [ ] JIT kernel work: test files live in `test/registered/jit/`; only test-only helpers stay under `python/sglang/kernels/jit/` - [ ] Backend-independent tests: `register_cuda_ci` only + smallest model - [ ] Logic that doesn't need a server / engine launch → unit test in `registered/unit/` (see Unit Tests section) - [ ] `setUpClass` launches server, `tearDownClass` kills it (if server-based) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 5fe76d3af..d930516af 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -4,8 +4,8 @@ /docs @wisclmy0611 @zijiexia @sogalin /docs_new @wisclmy0611 @zijiexia @Richardczl98 @JustinTong0323 @sogalin /python/pyproject.toml @merrymercy @Fridge003 @ispobock -/python/sglang/jit_kernel @DarkSharpness @BBuf @celve @HydraQYH @yuan-luo -/python/sglang/jit_kernel/diffusion @yingluosanqian @BBuf @mickqian +/python/sglang/kernels @DarkSharpness @BBuf @celve @HydraQYH @yuan-luo +/python/sglang/kernels/ops/diffusion @yingluosanqian @BBuf @mickqian /python/sglang/kernels/ops/attention/fla @yizhang2077 @hebiao064 @yuan-luo /python/sglang/multimodal_gen @mickqian @ping1jing2 @HaiShaw @yichiche @AgainstEntropy @BBuf /python/sglang/multimodal_gen/runtime/cache @DefTruth @@ -101,5 +101,5 @@ /python/sglang/srt/speculative/adaptive_*.py @Qiaolin-Yu @alphabetc1 /python/sglang/srt/speculative/cpp_ngram @hnyls2002 @Qiaolin-Yu @kpham-sgl /python/sglang/srt/speculative/frozen_kv_mtp_*.py @hnyls2002 @Qiaolin-Yu @kpham-sgl @pyc96 -/python/sglang/jit_kernel/ngram_*.py @hnyls2002 @Qiaolin-Yu @kpham-sgl -/python/sglang/jit_kernel/csrc/ngram_corpus @hnyls2002 @Qiaolin-Yu @kpham-sgl +/python/sglang/kernels/ops/speculative/ngram_*.py @hnyls2002 @Qiaolin-Yu @kpham-sgl +/python/sglang/kernels/jit/csrc/ngram_corpus @hnyls2002 @Qiaolin-Yu @kpham-sgl diff --git a/.github/MAINTAINER.md b/.github/MAINTAINER.md index ee4678d4e..4aac1716c 100644 --- a/.github/MAINTAINER.md +++ b/.github/MAINTAINER.md @@ -80,7 +80,7 @@ related files [@BBuf](https://github.com/BBuf) (BBuf) related files -- python/sglang/jit_kernel +- python/sglang/kernels - sgl-kernel ### Speculative decoding diff --git a/.github/labeler.yml b/.github/labeler.yml index 21a07d908..618083f0c 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -14,7 +14,7 @@ sgl-kernel: # JIT kernel specific jit-kernel: - changed-files: - - any-glob-to-any-file: 'python/sglang/jit_kernel/**/*' + - any-glob-to-any-file: 'python/sglang/kernels/**/*' # Documentation documentation: diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 763cb6c35..f65e1a721 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -93,7 +93,7 @@ jobs: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/**" - "python/sglang/srt/observability/**" - - "python/sglang/jit_kernel/**" + - "python/sglang/kernels/ops/diffusion/**" - "test/registered/jit/diffusion/**" - "test/registered/jit/benchmark/diffusion/**" - "python/sglang/cli/**" @@ -101,7 +101,6 @@ jobs: - ".github/workflows/pr-test.yml" - ".github/workflows/pr-test-jit-kernel.yml" - "python/pyproject.toml" - - "python/sglang/jit_kernel/**" - "test/registered/jit/**" # sglang.kernels is the migrated kernel namespace (RFC #29630 / #30044); the # base-b-kernel suites import it directly, so kernel edits must run them. diff --git a/.github/workflows/nightly-test-nvidia.yml b/.github/workflows/nightly-test-nvidia.yml index dc9ac1f2b..e5ec6b4fb 100644 --- a/.github/workflows/nightly-test-nvidia.yml +++ b/.github/workflows/nightly-test-nvidia.yml @@ -92,7 +92,7 @@ jobs: runs-on: 1-gpu-h100 timeout-minutes: 60 env: - # Full jit_kernel test grids (see sglang.jit_kernel.utils.should_run_full_tests) + # Full jit_kernel test grids (see sglang.kernels.jit.utils.should_run_full_tests) SGLANG_JIT_KERNEL_RUN_FULL_TESTS: "1" # Match pr-test-jit-kernel workflow for consistent JIT warmup behavior SGLANG_JIT_DEEPGEMM_FAST_WARMUP: true diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 529fad125..4ba8b30b5 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -198,14 +198,14 @@ jobs: - "sgl-kernel/**/!(*.md|THIRDPARTYNOTICES.txt|LICENSE)" - ".github/workflows/pr-test-amd-rocm720.yml" jit_kernel: - - "python/sglang/jit_kernel/**" + - "python/sglang/kernels/**" - "test/registered/jit/**" - ".github/workflows/pr-test-amd-rocm720.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - - "python/sglang/jit_kernel/diffusion/**" + - "python/sglang/kernels/ops/diffusion/**" - "test/registered/jit/diffusion/**" - "test/registered/jit/benchmark/diffusion/**" - "python/pyproject_rocm.toml" diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 307aab6b0..7b4b15a26 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -186,14 +186,14 @@ jobs: - "sgl-kernel/**/!(*.md|THIRDPARTYNOTICES.txt|LICENSE)" - ".github/workflows/pr-test-amd.yml" jit_kernel: - - "python/sglang/jit_kernel/**" + - "python/sglang/kernels/**" - "test/registered/jit/**" - ".github/workflows/pr-test-amd.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - - "python/sglang/jit_kernel/diffusion/**" + - "python/sglang/kernels/ops/diffusion/**" - "test/registered/jit/diffusion/**" - "test/registered/jit/benchmark/diffusion/**" - "python/pyproject_rocm.toml" diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index a1a2cf228..d0186b204 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -1,478 +1,478 @@ -name: PR Test (NPU) - -on: - push: - branches: [ main ] - pull_request: - workflow_dispatch: - workflow_call: - inputs: - ref: - description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' - required: false - type: string - default: '' - run_all_tests: - description: "Run all tests (for releasing or testing purpose)" - required: false - type: boolean - default: false - -concurrency: - group: pr-test-npu-${{ inputs.ref || github.ref }} - cancel-in-progress: ${{ github.event_name != 'workflow_call' }} - -jobs: - # ==================== Check Changes ==================== # - check-changes: - runs-on: ubuntu-latest - outputs: - changes_exist: ${{ steps.filter.outputs.main_package == 'true' || steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true'}} - main_package: ${{ steps.filter.outputs.main_package == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} - multimodal_gen: ${{ steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - ref: ${{ inputs.ref || github.ref }} - - - name: Determine run mode - id: run-mode - run: | - # Run all tests for workflow_call (when ref input is provided) - # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref - if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then - echo "run_all_tests=true" >> $GITHUB_OUTPUT - echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" - else - echo "run_all_tests=false" >> $GITHUB_OUTPUT - echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" - fi - - - name: Detect file changes - id: filter - uses: dorny/paths-filter@v3 - if: steps.run-mode.outputs.run_all_tests != 'true' - with: - filters: | - main_package: - - "python/sglang/!(multimodal_gen)/**/!(*.md)" - - "python/pyproject_npu.toml" - - "scripts/ci/npu/npu_ci_install_dependency.sh" - - "test/registered/ascend/**" - - ".github/workflows/pr-test-npu.yml" - multimodal_gen: - - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - - "python/sglang/jit_kernel/diffusion/triton/npu_fallback.py" - - "python/sglang/srt/**" - - "python/pyproject_npu.toml" - - "scripts/ci/npu/npu_ci_install_dependency.sh" - - ".github/workflows/pr-test-npu.yml" - - # ==================== PR Gate ==================== # - pr-gate: - needs: check-changes - if: needs.check-changes.outputs.changes_exist == 'true' - uses: ./.github/workflows/pr-gate.yml - secrets: inherit - - set-image-config: - runs-on: ubuntu-latest - outputs: - CANN_image_a3: ${{ steps.set-vars.outputs.CANN_image_a3 }} - CANN_image_910b: ${{ steps.set-vars.outputs.CANN_image_910b }} - steps: - # When triggered by PR, no inputs parameters are used. The latest community code is tested by default. - - name: Set image config - id: set-vars - run: | - echo "CANN_image_a3=swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-a3-ubuntu22.04-py3.11" >> $GITHUB_OUTPUT - echo "CANN_image_910b=swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-910b-ubuntu22.04-py3.11" >> $GITHUB_OUTPUT - - stage-b-test-1-npu-a2: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.main_package == 'true' - runs-on: linux-aarch64-a2-1 - strategy: - fail-fast: false - matrix: - part: [ 0, 1 ] - container: - image: ${{ needs.set-image-config.outputs.CANN_image_910b }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - ref: ${{ inputs.ref || github.ref }} - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh 910b - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - run: | - cd test - python3 run_suite.py --hw npu --suite stage-b-test-1-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 - - stage-b-test-2-npu-a2: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.main_package == 'true' - runs-on: linux-aarch64-a2-2 - strategy: - fail-fast: true - matrix: - part: [0, 1] - container: - image: ${{ needs.set-image-config.outputs.CANN_image_910b }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - ref: ${{ inputs.ref || github.ref }} - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh 910b - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - run: | - cd test - python3 run_suite.py --hw npu --suite stage-b-test-2-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 - - stage-b-test-4-npu-a3: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.main_package == 'true' - runs-on: linux-aarch64-a3-4 - container: - image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - ref: ${{ inputs.ref || github.ref }} - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh a3 - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - run: | - cd test - python3 run_suite.py --hw npu --suite stage-b-test-4-npu-a3 --timeout-per-file 3600 - - - stage-b-test-16-npu-a3: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.main_package == 'true' - runs-on: linux-aarch64-a3-16 - container: - image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - ref: ${{ inputs.ref || github.ref }} - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh a3 - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - run: | - cd test - python3 run_suite.py --hw npu --suite stage-b-test-16-npu-a3 --timeout-per-file 3600 - - multimodal-gen-test-1-npu-a3: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.multimodal_gen == 'true' - runs-on: linux-aarch64-a3-2 - container: - image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh a3 - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures - run: | - cd python - python3 sglang/multimodal_gen/test/run_suite.py --suite 1-npu - - - name: Upload diffusion failure artifacts - if: always() - uses: actions/upload-artifact@v4 - with: - name: diffusion-failures-npu-1-${{ github.run_attempt }} - path: diffusion-failures/ - if-no-files-found: ignore - retention-days: 7 - - multimodal-gen-test-2-npu-a3: - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.multimodal_gen == 'true' - runs-on: linux-aarch64-a3-16 - container: - image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} - steps: - - name: Checkout code - uses: actions/checkout@v4 - - - name: Mark repository safe - run: | - git config --system --add safe.directory ${GITHUB_WORKSPACE} - - - name: Install dependencies - env: - TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" - PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" - GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" - RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" - run: | - # speed up by using infra cache services - CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" - sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list - pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host "${CACHING_URL}" - - bash scripts/ci/npu/npu_ci_install_dependency.sh a3 - # copy required file from our daily cache - cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp - # copy gsm8k dataset - cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp - - - name: Run test - timeout-minutes: 60 - env: - SGLANG_USE_MODELSCOPE: true - SGLANG_IS_IN_CI: true - HF_ENDPOINT: https://hf-mirror.com - TORCH_EXTENSIONS_DIR: /tmp/torch_extensions - PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" - STREAMS_PER_DEVICE: 32 - SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures - run: | - cd python - python3 sglang/multimodal_gen/test/run_suite.py --suite 2-npu - - - name: Upload diffusion failure artifacts - if: always() - uses: actions/upload-artifact@v4 - with: - name: diffusion-failures-npu-2-${{ github.run_attempt }} - path: diffusion-failures/ - if-no-files-found: ignore - retention-days: 7 - - pr-single-node-tests: - name: single-node-poc - needs: [check-changes, pr-gate, set-image-config] - if: needs.check-changes.outputs.main_package == 'true' - strategy: - fail-fast: false - max-parallel: 6 - matrix: - test_config: - # qwen3_6_27b performance tests - - name: qwen3_6_27b_w8a8_1p_in64k_out1k_50ms - runner: linux-aarch64-a3-2 - test_case: test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py - test_type: 'perf' - uses: ./.github/workflows/nightly-test-npu-e2e-single-node.yml - with: - runner: ${{ matrix.test_config.runner }} - test_type: ${{ matrix.test_config.test_type }} - test_config_name: ${{ matrix.test_config.name }} - test_case: ${{ matrix.test_config.test_case }} - image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} - install_sglang_from_source: false - install_sglang_deps: true - device_type_for_deps: 'a3' - transformers_version: '' - - pr-test-npu-finish: - needs: - [ - check-changes, - - stage-b-test-1-npu-a2, - stage-b-test-2-npu-a2, - stage-b-test-4-npu-a3, - stage-b-test-16-npu-a3, - - multimodal-gen-test-1-npu-a3, - multimodal-gen-test-2-npu-a3, - - pr-single-node-tests, - ] - if: always() - runs-on: ubuntu-latest - steps: - - name: Check all dependent job statuses - run: | - # Convert the 'needs' context to a JSON string - json_needs='${{ toJson(needs) }}' - - # Get a list of all job names from the JSON keys - job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') - - for job in $job_names; do - # For each job, extract its result - result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') - - # Print the job name and its result - echo "$job: $result" - - # Check for failure or cancellation and exit if found - if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then - echo "The above jobs failed." - exit 1 - fi - done - # If the loop completes, all jobs were successful - echo "All jobs completed successfully" - exit 0 +name: PR Test (NPU) + +on: + push: + branches: [ main ] + pull_request: + workflow_dispatch: + workflow_call: + inputs: + ref: + description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' + required: false + type: string + default: '' + run_all_tests: + description: "Run all tests (for releasing or testing purpose)" + required: false + type: boolean + default: false + +concurrency: + group: pr-test-npu-${{ inputs.ref || github.ref }} + cancel-in-progress: ${{ github.event_name != 'workflow_call' }} + +jobs: + # ==================== Check Changes ==================== # + check-changes: + runs-on: ubuntu-latest + outputs: + changes_exist: ${{ steps.filter.outputs.main_package == 'true' || steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true'}} + main_package: ${{ steps.filter.outputs.main_package == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} + multimodal_gen: ${{ steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Determine run mode + id: run-mode + run: | + # Run all tests for workflow_call (when ref input is provided) + # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref + if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then + echo "run_all_tests=true" >> $GITHUB_OUTPUT + echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" + else + echo "run_all_tests=false" >> $GITHUB_OUTPUT + echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" + fi + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + if: steps.run-mode.outputs.run_all_tests != 'true' + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**/!(*.md)" + - "python/pyproject_npu.toml" + - "scripts/ci/npu/npu_ci_install_dependency.sh" + - "test/registered/ascend/**" + - ".github/workflows/pr-test-npu.yml" + multimodal_gen: + - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" + - "python/sglang/kernels/ops/diffusion/triton/npu_fallback.py" + - "python/sglang/srt/**" + - "python/pyproject_npu.toml" + - "scripts/ci/npu/npu_ci_install_dependency.sh" + - ".github/workflows/pr-test-npu.yml" + + # ==================== PR Gate ==================== # + pr-gate: + needs: check-changes + if: needs.check-changes.outputs.changes_exist == 'true' + uses: ./.github/workflows/pr-gate.yml + secrets: inherit + + set-image-config: + runs-on: ubuntu-latest + outputs: + CANN_image_a3: ${{ steps.set-vars.outputs.CANN_image_a3 }} + CANN_image_910b: ${{ steps.set-vars.outputs.CANN_image_910b }} + steps: + # When triggered by PR, no inputs parameters are used. The latest community code is tested by default. + - name: Set image config + id: set-vars + run: | + echo "CANN_image_a3=swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-a3-ubuntu22.04-py3.11" >> $GITHUB_OUTPUT + echo "CANN_image_910b=swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0-910b-ubuntu22.04-py3.11" >> $GITHUB_OUTPUT + + stage-b-test-1-npu-a2: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-aarch64-a2-1 + strategy: + fail-fast: false + matrix: + part: [ 0, 1 ] + container: + image: ${{ needs.set-image-config.outputs.CANN_image_910b }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh 910b + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + run: | + cd test + python3 run_suite.py --hw npu --suite stage-b-test-1-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + stage-b-test-2-npu-a2: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-aarch64-a2-2 + strategy: + fail-fast: true + matrix: + part: [0, 1] + container: + image: ${{ needs.set-image-config.outputs.CANN_image_910b }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh 910b + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + run: | + cd test + python3 run_suite.py --hw npu --suite stage-b-test-2-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + stage-b-test-4-npu-a3: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-aarch64-a3-4 + container: + image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh a3 + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + run: | + cd test + python3 run_suite.py --hw npu --suite stage-b-test-4-npu-a3 --timeout-per-file 3600 + + + stage-b-test-16-npu-a3: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-aarch64-a3-16 + container: + image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh a3 + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + run: | + cd test + python3 run_suite.py --hw npu --suite stage-b-test-16-npu-a3 --timeout-per-file 3600 + + multimodal-gen-test-1-npu-a3: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.multimodal_gen == 'true' + runs-on: linux-aarch64-a3-2 + container: + image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh a3 + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py --suite 1-npu + + - name: Upload diffusion failure artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-failures-npu-1-${{ github.run_attempt }} + path: diffusion-failures/ + if-no-files-found: ignore + retention-days: 7 + + multimodal-gen-test-2-npu-a3: + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.multimodal_gen == 'true' + runs-on: linux-aarch64-a3-16 + container: + image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Mark repository safe + run: | + git config --system --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install dependencies + env: + TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" + PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + UV_INDEX_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" + GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" + RUSTUP_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local:8082" + run: | + # speed up by using infra cache services + CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" + sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list + pip config set global.index-url http://${CACHING_URL}/pypi/simple + pip config set global.trusted-host "${CACHING_URL}" + + bash scripts/ci/npu/npu_ci_install_dependency.sh a3 + # copy required file from our daily cache + cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp + # copy gsm8k dataset + cp ~/.cache/modelscope/hub/datasets/tmp/test.jsonl /tmp + + - name: Run test + timeout-minutes: 60 + env: + SGLANG_USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + HF_ENDPOINT: https://hf-mirror.com + TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 + SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py --suite 2-npu + + - name: Upload diffusion failure artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-failures-npu-2-${{ github.run_attempt }} + path: diffusion-failures/ + if-no-files-found: ignore + retention-days: 7 + + pr-single-node-tests: + name: single-node-poc + needs: [check-changes, pr-gate, set-image-config] + if: needs.check-changes.outputs.main_package == 'true' + strategy: + fail-fast: false + max-parallel: 6 + matrix: + test_config: + # qwen3_6_27b performance tests + - name: qwen3_6_27b_w8a8_1p_in64k_out1k_50ms + runner: linux-aarch64-a3-2 + test_case: test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py + test_type: 'perf' + uses: ./.github/workflows/nightly-test-npu-e2e-single-node.yml + with: + runner: ${{ matrix.test_config.runner }} + test_type: ${{ matrix.test_config.test_type }} + test_config_name: ${{ matrix.test_config.name }} + test_case: ${{ matrix.test_config.test_case }} + image: ${{ needs.set-image-config.outputs.CANN_image_a3 }} + install_sglang_from_source: false + install_sglang_deps: true + device_type_for_deps: 'a3' + transformers_version: '' + + pr-test-npu-finish: + needs: + [ + check-changes, + + stage-b-test-1-npu-a2, + stage-b-test-2-npu-a2, + stage-b-test-4-npu-a3, + stage-b-test-16-npu-a3, + + multimodal-gen-test-1-npu-a3, + multimodal-gen-test-2-npu-a3, + + pr-single-node-tests, + ] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + # Convert the 'needs' context to a JSON string + json_needs='${{ toJson(needs) }}' + + # Get a list of all job names from the JSON keys + job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') + + for job in $job_names; do + # For each job, extract its result + result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') + + # Print the job name and its result + echo "$job: $result" + + # Check for failure or cancellation and exit if found + if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then + echo "The above jobs failed." + exit 1 + fi + done + # If the loop completes, all jobs were successful + echo "All jobs completed successfully" + exit 0 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index e7ff56f62..7b99b548e 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -228,7 +228,7 @@ jobs: call-jit-kernel-tests: needs: [check-changes, call-gate, sgl-kernel-build-wheels] # Run on scheduled/parallel-dispatch runs (same pattern as the base-* stages) so the - # jit_kernel suite is exercised on main 3x daily, not only on PRs that touch jit_kernel/**. + # jit_kernel suite is exercised on main 3x daily, not only on PRs that touch kernels/**. # check-changes already forces jit_kernel='true' on scheduled runs (run_all_tests). if: | always() && diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a4064717d..483f950c4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_stages: [pre-commit, pre-push, manual] -exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/kernels/ops/diffusion/render|python/sglang/jit_kernel/flash_attention/cute) +exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/kernels/ops/diffusion/render|python/sglang/kernels/ops/attention/flash_attn/cute) repos: - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index 9dc3dd71d..bbc9f464b 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -186,7 +186,7 @@ sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*", + "kernels/**/*", ] [tool.setuptools.packages.find] diff --git a/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py b/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py index af138d5d0..f501d5454 100644 --- a/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py +++ b/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py @@ -17,7 +17,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch import triton -from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update +from sglang.kernels.ops.attention.cutedsl_kda import ( + cutedsl_fused_sigmoid_gating_kda_update, +) from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/benchmark/bench_linear_attention/bench_gdn_qkv_split.py b/benchmark/bench_linear_attention/bench_gdn_qkv_split.py index 30aaf9e26..5d9278670 100644 --- a/benchmark/bench_linear_attention/bench_gdn_qkv_split.py +++ b/benchmark/bench_linear_attention/bench_gdn_qkv_split.py @@ -4,7 +4,9 @@ import argparse import torch -from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkv_split_gdn_prefill +from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + fused_qkv_split_gdn_prefill, +) DTYPES = { "bf16": torch.bfloat16, diff --git a/benchmark/hicache/bench_hicache_write_back.py b/benchmark/hicache/bench_hicache_write_back.py index 6dd6e7d76..54e0eace2 100644 --- a/benchmark/hicache/bench_hicache_write_back.py +++ b/benchmark/hicache/bench_hicache_write_back.py @@ -12,7 +12,7 @@ from sgl_kernel.kvcacheio import ( transfer_kv_all_layer_mla_lf_pf, ) -from sglang.jit_kernel.hicache import ( +from sglang.kernels.ops.kvcache.hicache import ( can_use_hicache_jit_kernel, transfer_hicache_all_layer_mla_staged_lf_pf, transfer_hicache_all_layer_staged_lf_pf, @@ -327,8 +327,8 @@ Environment: Comparison target: -- MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_staged_lf_pf` -- MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.jit_kernel.hicache.transfer_hicache_all_layer_mla_staged_lf_pf` +- MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_staged_lf_pf` +- MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_mla_staged_lf_pf` Metric: diff --git a/benchmark/hicache/perf.py b/benchmark/hicache/perf.py index 2349af4b1..4ffc1f807 100644 --- a/benchmark/hicache/perf.py +++ b/benchmark/hicache/perf.py @@ -15,7 +15,7 @@ def jit_hicache_impl( item_bytes: int, block_quota: int, ) -> None: - from sglang.jit_kernel.hicache import transfer_hicache_one_layer + from sglang.kernels.ops.kvcache.hicache import transfer_hicache_one_layer _ = item_bytes diff --git a/benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py b/benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py index b157a583c..4e57efaac 100644 --- a/benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py +++ b/benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py @@ -8,8 +8,8 @@ import sys import numpy as np import torch -import sglang.jit_kernel.dsa.cutedsl_paged_mqa_logits # noqa: F401 -from sglang.jit_kernel.dsa import pick_dsl_expand +import sglang.kernels.ops.attention.dsa.cutedsl_paged_mqa_logits # noqa: F401 +from sglang.kernels.ops.attention.dsa import pick_dsl_expand from sglang.srt.layers.attention.dsa.utils import ( fp8_mqa_logits_ceil_to_ue8m0, fp8_mqa_logits_make_fused_kv, diff --git a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx index 28511be09..cdbfb8288 100644 --- a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx +++ b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx @@ -10,18 +10,18 @@ We strongly recommend using `clangd` as the language server for JIT kernel devel For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/). If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration. -All JIT-related files are located in `python/sglang/jit_kernel`. +All JIT-related files are located in `python/sglang/kernels/jit`. Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime. Consequently, a static `compile_commands.json` cannot be generated. -To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory. +To enable code completion with `clangd`, run `python -m sglang.kernels.jit` to generate a `.clangd` configuration file in your current directory. After generating the file, restart the clangd language server. It should now recognize all JIT kernel files. ## Code Structure ### C++ Implementation -C++ source code is located in `python/sglang/jit_kernel/csrc`. -Reusable functions should be placed in `python/sglang/jit_kernel/include`. +C++ source code is located in `python/sglang/kernels/jit/csrc`. +Reusable functions should be placed in `python/sglang/kernels/jit/include`. We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings. Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects. @@ -29,12 +29,12 @@ Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from ### Python Interface -Python interfaces are defined in `python/sglang/jit_kernel`. -The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module. +Python interfaces are defined in `python/sglang/kernels/jit`. +The `load_jit` utility function in `python/sglang/kernels/jit/utils/compile.py` loads and returns the compiled module. To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`. The function can then be called in Python as `module.func`. -For caching compiled modules, prefer `sglang.jit_kernel.utils.cache_once` over `functools.lru_cache`. +For caching compiled modules, prefer `sglang.kernels.jit.utils.cache_once` over `functools.lru_cache`. `functools.lru_cache` is not compatible with `torch.compile`. ### C++ Utilities @@ -161,7 +161,7 @@ def add_constant(src: torch.Tensor, c: int): ### STEP 1: Write the C++ kernel -Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter. +Write your CUDA kernel in [kernels/jit/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter. ```cpp Example #include // For TensorMatcher, SymbolicSize, SymbolicDevice @@ -224,7 +224,7 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { ### STEP 2: Create Python Interfaces Next, expose the kernel through a Python wrapper. -Create a new file at [jit_kernel/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces. +Create a new file at [kernels/ops/attention/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/ops/attention/add_constant.py) and expose the needed interfaces. ```python Example from __future__ import annotations @@ -232,7 +232,7 @@ from typing import TYPE_CHECKING import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @@ -268,7 +268,7 @@ Keep the Python wrapper thin, but still validate the basic invariants such as de Finally, import and use the kernel like a regular Python function: ```python Example -from sglang.jit_kernel.add_constant import add_constant +from sglang.kernels.jit.add_constant import add_constant ``` For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/test/registered/jit/test_add_constant.py). @@ -276,7 +276,7 @@ For a complete, runnable example, refer to [test_add_constant.py](https://github ## C++ Include Library Reference The JIT kernel framework provides a set of reusable C++ headers in -`python/sglang/jit_kernel/include/sgl_kernel/`. Each header is designed +`python/sglang/kernels/jit/include/sgl_kernel/`. Each header is designed to be lightweight and self-contained. Below is a summary of each header and its key APIs. diff --git a/python/pyproject.toml b/python/pyproject.toml index 4131966a2..9a6b3d930 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -192,7 +192,7 @@ killall_sglang = "sglang.cli.killall:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*", + "kernels/**/*", "multimodal_gen/apps/realtime_webui/**/*" ] diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 36d4568c0..31ae12a42 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -123,7 +123,7 @@ sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*" + "kernels/**/*" ] [tool.setuptools.packages.find] diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 09f43088e..727c89de9 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -122,7 +122,7 @@ sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*" + "kernels/**/*" ] [tool.setuptools.packages.find] diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 8ad576f62..b6727f358 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -202,7 +202,7 @@ binding = "PyO3" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*" + "kernels/**/*" ] [tool.setuptools.packages.find] diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index 28aeb9f6f..bcd3d1a72 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -130,7 +130,7 @@ sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", - "jit_kernel/**/*" + "kernels/**/*" ] [tool.setuptools.packages.find] diff --git a/python/sglang/jit_kernel/activation.py b/python/sglang/jit_kernel/activation.py deleted file mode 100644 index 520e78415..000000000 --- a/python/sglang/jit_kernel/activation.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.activation._jit_activation.""" - -from sglang.kernels.ops.activation import _jit_activation as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py deleted file mode 100644 index 56f770684..000000000 --- a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm.""" - -from sglang.kernels.ops.gemm import _jit_dsv3_fused_a_gemm as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/dsv3_router_gemm.py b/python/sglang/jit_kernel/dsv3_router_gemm.py deleted file mode 100644 index fac174e38..000000000 --- a/python/sglang/jit_kernel/dsv3_router_gemm.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.gemm._jit_dsv3_router_gemm.""" - -from sglang.kernels.ops.gemm import _jit_dsv3_router_gemm as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py deleted file mode 100644 index b1a334507..000000000 --- a/python/sglang/jit_kernel/norm.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.layernorm._jit_norm.""" - -from sglang.kernels.ops.layernorm import _jit_norm as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/per_tensor_quant_fp8.py b/python/sglang/jit_kernel/per_tensor_quant_fp8.py deleted file mode 100644 index b2b987ec0..000000000 --- a/python/sglang/jit_kernel/per_tensor_quant_fp8.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8.""" - -from sglang.kernels.ops.quantization import _jit_per_tensor_quant_fp8 as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/per_token_group_quant.py b/python/sglang/jit_kernel/per_token_group_quant.py deleted file mode 100644 index 314c8c056..000000000 --- a/python/sglang/jit_kernel/per_token_group_quant.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_token_group_quant.""" - -from sglang.kernels.ops.quantization import _jit_per_token_group_quant as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py b/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py deleted file mode 100644 index 0b24d721f..000000000 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2.""" - -from sglang.kernels.ops.quantization import _jit_per_token_group_quant_8bit_v2 as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/set_mla_kv_buffer.py b/python/sglang/jit_kernel/set_mla_kv_buffer.py deleted file mode 100644 index cdda777b2..000000000 --- a/python/sglang/jit_kernel/set_mla_kv_buffer.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility shim (RFC #29630 Phase 4) -> sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer.""" - -from sglang.kernels.ops.kvcache import _jit_set_mla_kv_buffer as _impl - -globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")}) diff --git a/python/sglang/jit_kernel/tests/kv_canary/__init__.py b/python/sglang/jit_kernel/tests/kv_canary/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/python/sglang/kernels/README.md b/python/sglang/kernels/README.md index 30602f032..817e34f45 100644 --- a/python/sglang/kernels/README.md +++ b/python/sglang/kernels/README.md @@ -26,7 +26,7 @@ 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.jit_kernel` / `sgl_kernel` / `triton_ops` and will migrate in +live in `sglang.kernels.jit` / `sgl_kernel` / `triton_ops` and will migrate in later phases. ## How it works @@ -104,7 +104,7 @@ What this buys (see the > SGLang runtime code and tests should import callable kernels from > `sglang.kernels.ops.*`. -Implementation work can still happen in `sglang.jit_kernel` or `sgl_kernel`. +Implementation work can still happen in `sglang.kernels.jit` or `sgl_kernel`. When a PR adds a new callable kernel, add a `sglang.kernels.ops.*` entry point -for it, and avoid growing `sglang.jit_kernel` as a long-term public operator +for it, and avoid growing `sglang.kernels.jit` as a long-term public operator namespace. diff --git a/python/sglang/kernels/__init__.py b/python/sglang/kernels/__init__.py index aaaffc4e2..40ef0629a 100644 --- a/python/sglang/kernels/__init__.py +++ b/python/sglang/kernels/__init__.py @@ -7,7 +7,7 @@ SGLang runtime code and tests should import callable kernels from from sglang.kernels.ops.activation import silu_and_mul from sglang.kernels.ops.kvcache import reshape_and_cache_flash -Implementations still live in ``sglang.jit_kernel`` (JIT CUDA), the +Implementations still live in ``sglang.kernels.jit`` (JIT CUDA), the ``sgl_kernel`` wheel (AOT CUDA/C++), Triton op modules, etc. The ``ops.*`` functions are thin wrappers that forward to a chosen backend; the :data:`~sglang.kernels.registry.registry` provides an inventory of every @@ -18,7 +18,7 @@ with a required pure-``torch`` ``forward_native`` reference and a ``SGLANG_FORCE_FUSED_OP_BACKEND`` global switch. Importing this package (and any ``ops.*`` group) does not import a kernel -backend (``sgl_kernel`` / ``sglang.jit_kernel``) or trigger JIT compilation: +backend (``sgl_kernel`` / ``sglang.kernels.jit``) or trigger JIT compilation: registration is metadata-only and backends are imported lazily on first call. This keeps the namespace usable for inventory tooling on a CPU-only box. """ diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py index f896bec08..d39ed4d19 100644 --- a/python/sglang/kernels/fused_op.py +++ b/python/sglang/kernels/fused_op.py @@ -28,7 +28,7 @@ implementations with a single switch. Like the rest of ``sglang.kernels``, importing this module (and instantiating subclasses) never imports a kernel backend (``sgl_kernel`` / -``sglang.jit_kernel``) or triggers JIT compilation; backends are imported +``sglang.kernels.jit``) or triggers JIT compilation; backends are imported lazily inside the ``forward_`` methods. """ diff --git a/python/sglang/jit_kernel/.clang-format b/python/sglang/kernels/jit/.clang-format similarity index 100% rename from python/sglang/jit_kernel/.clang-format rename to python/sglang/kernels/jit/.clang-format diff --git a/python/sglang/kernels/jit/__init__.py b/python/sglang/kernels/jit/__init__.py index 3a456a260..277002c78 100644 --- a/python/sglang/kernels/jit/__init__.py +++ b/python/sglang/kernels/jit/__init__.py @@ -1,6 +1,6 @@ """Internal JIT home under ``sglang.kernels`` (RFC #29630). -Mirrors the legacy ``sglang.jit_kernel`` tree; shared build/runtime +Mirrors the legacy ``sglang.kernels.jit`` tree; shared build/runtime infrastructure lives in :mod:`sglang.kernels.jit.utils`. csrc / include / operators migrate here in later phases. """ diff --git a/python/sglang/jit_kernel/__main__.py b/python/sglang/kernels/jit/__main__.py similarity index 100% rename from python/sglang/jit_kernel/__main__.py rename to python/sglang/kernels/jit/__main__.py diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/utils.py b/python/sglang/kernels/jit/benchmark/kv_canary/utils.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/kv_canary/utils.py rename to python/sglang/kernels/jit/benchmark/kv_canary/utils.py diff --git a/python/sglang/jit_kernel/benchmark/marker.py b/python/sglang/kernels/jit/benchmark/marker.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/marker.py rename to python/sglang/kernels/jit/benchmark/marker.py diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/kernels/jit/benchmark/utils.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/utils.py rename to python/sglang/kernels/jit/benchmark/utils.py diff --git a/python/sglang/jit_kernel/csrc/add_constant.cuh b/python/sglang/kernels/jit/csrc/add_constant.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/add_constant.cuh rename to python/sglang/kernels/jit/csrc/add_constant.cuh diff --git a/python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh b/python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh rename to python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh diff --git a/python/sglang/jit_kernel/csrc/attention/fused_fp8_qkv_kv_cache.cuh b/python/sglang/kernels/jit/csrc/attention/fused_fp8_qkv_kv_cache.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/attention/fused_fp8_qkv_kv_cache.cuh rename to python/sglang/kernels/jit/csrc/attention/fused_fp8_qkv_kv_cache.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh b/python/sglang/kernels/jit/csrc/deepseek_v32/indexer_k.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v32/indexer_k.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c128.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c128.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c128_online.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c128_online.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c128_online.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c128_online_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c128_online_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c128_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c128_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c4.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c4.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c4_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c4_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/common.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/common.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/online_c128_mtp.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/online_c128_mtp.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/paged_mqa_metadata.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/paged_mqa_metadata.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/rope.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/rope.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/rope.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/store.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/store.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/store.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/store.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh rename to python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh b/python/sglang/kernels/jit/csrc/diffusion/causal_conv3d_cat_pad.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh rename to python/sglang/kernels/jit/csrc/diffusion/causal_conv3d_cat_pad.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/ltx2_qknorm_split_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh rename to python/sglang/kernels/jit/csrc/diffusion/ltx2_qknorm_split_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/norm_scale_shift.cuh b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/norm_scale_shift.cuh rename to python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh rename to python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh b/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh rename to python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh diff --git a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh b/python/sglang/kernels/jit/csrc/diffusion/timestep_embedding.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh rename to python/sglang/kernels/jit/csrc/diffusion/timestep_embedding.cuh diff --git a/python/sglang/jit_kernel/csrc/distributed/communicator.cuh b/python/sglang/kernels/jit/csrc/distributed/communicator.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/distributed/communicator.cuh rename to python/sglang/kernels/jit/csrc/distributed/communicator.cuh diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh rename to python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh diff --git a/python/sglang/jit_kernel/csrc/distributed/ipc.cuh b/python/sglang/kernels/jit/csrc/distributed/ipc.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/distributed/ipc.cuh rename to python/sglang/kernels/jit/csrc/distributed/ipc.cuh diff --git a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh b/python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh rename to python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh diff --git a/python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh b/python/sglang/kernels/jit/csrc/dsa/fused_store_index_cache.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh rename to python/sglang/kernels/jit/csrc/dsa/fused_store_index_cache.cuh diff --git a/python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh b/python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh rename to python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh b/python/sglang/kernels/jit/csrc/elementwise/activation.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/activation.cuh rename to python/sglang/kernels/jit/csrc/elementwise/activation.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh b/python/sglang/kernels/jit/csrc/elementwise/clamp_position.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh rename to python/sglang/kernels/jit/csrc/elementwise/clamp_position.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh b/python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh rename to python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh b/python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh rename to python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh b/python/sglang/kernels/jit/csrc/elementwise/fused_eh_norm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh rename to python/sglang/kernels/jit/csrc/elementwise/fused_eh_norm.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh b/python/sglang/kernels/jit/csrc/elementwise/fused_metadata_copy.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh rename to python/sglang/kernels/jit/csrc/elementwise/fused_metadata_copy.cuh index 1758c38da..af2d1905f 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh +++ b/python/sglang/kernels/jit/csrc/elementwise/fused_metadata_copy.cuh @@ -1,6 +1,6 @@ /* * Fused metadata copy kernel for DSA backend CUDA graph replay. - * JIT-compiled version for python/sglang/jit_kernel. + * JIT-compiled version for python/sglang/kernels/jit. * * OVERVIEW: * This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k, diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/elementwise/fused_qknorm_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh rename to python/sglang/kernels/jit/csrc/elementwise/fused_qknorm_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh b/python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh rename to python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh b/python/sglang/kernels/jit/csrc/elementwise/pos_enc.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh rename to python/sglang/kernels/jit/csrc/elementwise/pos_enc.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh b/python/sglang/kernels/jit/csrc/elementwise/qknorm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh rename to python/sglang/kernels/jit/csrc/elementwise/qknorm.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh b/python/sglang/kernels/jit/csrc/elementwise/qknorm_across_heads.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh rename to python/sglang/kernels/jit/csrc/elementwise/qknorm_across_heads.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh b/python/sglang/kernels/jit/csrc/elementwise/resolve_future_token_ids.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh rename to python/sglang/kernels/jit/csrc/elementwise/resolve_future_token_ids.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh b/python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh rename to python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/rmsnorm_hf.cuh b/python/sglang/kernels/jit/csrc/elementwise/rmsnorm_hf.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/rmsnorm_hf.cuh rename to python/sglang/kernels/jit/csrc/elementwise/rmsnorm_hf.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh b/python/sglang/kernels/jit/csrc/elementwise/rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/rope.cuh rename to python/sglang/kernels/jit/csrc/elementwise/rope.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/set_mla_kv_buffer.cuh b/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/elementwise/set_mla_kv_buffer.cuh rename to python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform.h similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform.h diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/hadamard_jit.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/hadamard_jit.cuh diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h b/python/sglang/kernels/jit/csrc/fast-hadamard-transform/static_switch.h similarity index 100% rename from python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h rename to python/sglang/kernels/jit/csrc/fast-hadamard-transform/static_switch.h diff --git a/python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh b/python/sglang/kernels/jit/csrc/gemm/awq_dequantize.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh rename to python/sglang/kernels/jit/csrc/gemm/awq_dequantize.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/dsv3_fused_a_gemm.cuh rename to python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/dsv3_router_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/dsv3_router_gemm.cuh rename to python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh b/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh rename to python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh b/python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh rename to python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin/awq_marlin_repack.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin/awq_marlin_repack.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/dequant.h b/python/sglang/kernels/jit/csrc/gemm/marlin/dequant.h similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/dequant.h rename to python/sglang/kernels/jit/csrc/gemm/marlin/dequant.h diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin_repack.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin_repack.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/kernel.h b/python/sglang/kernels/jit/csrc/gemm/marlin/kernel.h similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/kernel.h rename to python/sglang/kernels/jit/csrc/gemm/marlin/kernel.h diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin/marlin.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin/marlin.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/marlin_dtypes.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin/marlin_dtypes.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/marlin_dtypes.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin/marlin_dtypes.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/marlin_template.h b/python/sglang/kernels/jit/csrc/gemm/marlin/marlin_template.h similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin/marlin_template.h rename to python/sglang/kernels/jit/csrc/gemm/marlin/marlin_template.h diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/kernel.h b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin_moe/kernel.h rename to python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h rename to python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh rename to python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh b/python/sglang/kernels/jit/csrc/gemm/per_tensor_quant_fp8.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh rename to python/sglang/kernels/jit/csrc/gemm/per_tensor_quant_fp8.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/per_token_group_quant.cuh rename to python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh rename to python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh diff --git a/python/sglang/jit_kernel/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/hisparse.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/hisparse.cuh rename to python/sglang/kernels/jit/csrc/hisparse.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/causal_conv1d.cuh b/python/sglang/kernels/jit/csrc/inkling/causal_conv1d.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/causal_conv1d.cuh rename to python/sglang/kernels/jit/csrc/inkling/causal_conv1d.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/draft_extend_sconv.cuh b/python/sglang/kernels/jit/csrc/inkling/draft_extend_sconv.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/draft_extend_sconv.cuh rename to python/sglang/kernels/jit/csrc/inkling/draft_extend_sconv.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/fused_decode_update.cuh b/python/sglang/kernels/jit/csrc/inkling/fused_decode_update.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/fused_decode_update.cuh rename to python/sglang/kernels/jit/csrc/inkling/fused_decode_update.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/gather_scatter_sconv.cuh b/python/sglang/kernels/jit/csrc/inkling/gather_scatter_sconv.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/gather_scatter_sconv.cuh rename to python/sglang/kernels/jit/csrc/inkling/gather_scatter_sconv.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_all_reduce.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_all_reduce.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_all_reduce.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_all_reduce.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_ar_barrier.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_ar_barrier.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_ar_barrier.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_ar_barrier.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_ar_fused_decode.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_ar_fused_decode.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_ar_fused_decode.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_ar_fused_decode.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_ar_scattered_sconv.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_ar_scattered_sconv.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_ar_scattered_sconv.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_ar_scattered_sconv.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_attn_prologue_fused.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_attn_prologue_fused.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_attn_prologue_fused.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_attn_prologue_fused.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_rel_proj.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_rel_proj.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_rel_proj.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_rel_proj.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/inkling_row_scale.cuh b/python/sglang/kernels/jit/csrc/inkling/inkling_row_scale.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/inkling_row_scale.cuh rename to python/sglang/kernels/jit/csrc/inkling/inkling_row_scale.cuh diff --git a/python/sglang/jit_kernel/csrc/inkling/update_sconv_cache.cuh b/python/sglang/kernels/jit/csrc/inkling/update_sconv_cache.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/inkling/update_sconv_cache.cuh rename to python/sglang/kernels/jit/csrc/inkling/update_sconv_cache.cuh diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_common.cuh b/python/sglang/kernels/jit/csrc/kv_canary/canary_common.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kv_canary/canary_common.cuh rename to python/sglang/kernels/jit/csrc/kv_canary/canary_common.cuh diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh b/python/sglang/kernels/jit/csrc/kv_canary/canary_plan_entries.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kv_canary/canary_plan_entries.cuh rename to python/sglang/kernels/jit/csrc/kv_canary/canary_plan_entries.cuh diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh b/python/sglang/kernels/jit/csrc/kv_canary/canary_verify.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh rename to python/sglang/kernels/jit/csrc/kv_canary/canary_verify.cuh index 38840fefd..4f6d2f30b 100644 --- a/python/sglang/jit_kernel/csrc/kv_canary/canary_verify.cuh +++ b/python/sglang/kernels/jit/csrc/kv_canary/canary_verify.cuh @@ -139,7 +139,7 @@ __global__ void canary_verify_kernel(const VerifyKernelParams __grid_constant__ } // namespace -// API source of truth: docstring of canary_verify_step in python/sglang/jit_kernel/kv_canary/verify.py. +// API source of truth: docstring of canary_verify_step in python/sglang/kernels/ops/kv_canary/verify.py. // // ABI notes: // - real_kv_buf_0 .. real_kv_buf_3 are 4 fixed uint8 tensor slots. Unused slots are dummy 1-byte tensors; diff --git a/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh b/python/sglang/kernels/jit/csrc/kv_canary/canary_write.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh rename to python/sglang/kernels/jit/csrc/kv_canary/canary_write.cuh index 7ebe7dbf8..d93b16ef0 100644 --- a/python/sglang/jit_kernel/csrc/kv_canary/canary_write.cuh +++ b/python/sglang/kernels/jit/csrc/kv_canary/canary_write.cuh @@ -166,7 +166,7 @@ __global__ void canary_write_kernel(const WriteKernelParams __grid_constant__ p) } // namespace -// API source of truth: docstring of canary_write_step in python/sglang/jit_kernel/kv_canary/write.py. +// API source of truth: docstring of canary_write_step in python/sglang/kernels/ops/kv_canary/write.py. // // ABI notes (same as verify): // - real_kv_buf_0 .. real_kv_buf_3 are 4 fixed uint8 tensor slots. diff --git a/python/sglang/jit_kernel/csrc/kv_canary/consts.cuh b/python/sglang/kernels/jit/csrc/kv_canary/consts.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kv_canary/consts.cuh rename to python/sglang/kernels/jit/csrc/kv_canary/consts.cuh diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh rename to python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh index ee16fba0a..740b5f458 100644 --- a/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh @@ -41,7 +41,7 @@ using PackageType = decltype(get_mem_package()); // has no equivalent PTX, but non-temporal (streaming) loads/stores express the // same intent for one-shot HiCache write-back traffic that should not pollute // the cache. Guard the PTX behind USE_ROCM so the JIT module also compiles with -// hipcc; see python/sglang/jit_kernel/utils.py for the ROCm build flags. +// hipcc; see python/sglang/kernels/jit/utils/compile.py for the ROCm build flags. #ifdef USE_ROCM // Native Clang vector types so a single __builtin_nontemporal_{load,store} maps // to one vectorized global_{load,store}_dwordx{2,4}. Issuing N independent diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/relayout.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kvcacheio/relayout.cuh rename to python/sglang/kernels/jit/csrc/kvcacheio/relayout.cuh diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/staged_write_back.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh rename to python/sglang/kernels/jit/csrc/kvcacheio/staged_write_back.cuh diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/kvcacheio/transfer_mamba.cuh rename to python/sglang/kernels/jit/csrc/kvcacheio/transfer_mamba.cuh diff --git a/python/sglang/jit_kernel/csrc/lora/moe_lora_align_kernel.cu b/python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu similarity index 100% rename from python/sglang/jit_kernel/csrc/lora/moe_lora_align_kernel.cu rename to python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu diff --git a/python/sglang/jit_kernel/csrc/lplb/dispatch_probability.cuh b/python/sglang/kernels/jit/csrc/lplb/dispatch_probability.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/lplb/dispatch_probability.cuh rename to python/sglang/kernels/jit/csrc/lplb/dispatch_probability.cuh diff --git a/python/sglang/jit_kernel/csrc/lplb/ipm.cuh b/python/sglang/kernels/jit/csrc/lplb/ipm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/lplb/ipm.cuh rename to python/sglang/kernels/jit/csrc/lplb/ipm.cuh diff --git a/python/sglang/jit_kernel/csrc/lplb/lp_post.cuh b/python/sglang/kernels/jit/csrc/lplb/lp_post.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/lplb/lp_post.cuh rename to python/sglang/kernels/jit/csrc/lplb/lp_post.cuh diff --git a/python/sglang/jit_kernel/csrc/lplb/lp_prep.cuh b/python/sglang/kernels/jit/csrc/lplb/lp_prep.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/lplb/lp_prep.cuh rename to python/sglang/kernels/jit/csrc/lplb/lp_prep.cuh diff --git a/python/sglang/jit_kernel/csrc/minimax/fused_gemma_qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/minimax/fused_gemma_qknorm_rope.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/minimax/fused_gemma_qknorm_rope.cuh rename to python/sglang/kernels/jit/csrc/minimax/fused_gemma_qknorm_rope.cuh diff --git a/python/sglang/jit_kernel/csrc/minimax/fused_store_kv_index.cuh b/python/sglang/kernels/jit/csrc/minimax/fused_store_kv_index.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/minimax/fused_store_kv_index.cuh rename to python/sglang/kernels/jit/csrc/minimax/fused_store_kv_index.cuh diff --git a/python/sglang/jit_kernel/csrc/minimax/minimax_decode_topk.cuh b/python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/minimax/minimax_decode_topk.cuh rename to python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh diff --git a/python/sglang/jit_kernel/csrc/minimax/per_token_quant_ue8m0.cuh b/python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/minimax/per_token_quant_ue8m0.cuh rename to python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh b/python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh rename to python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh b/python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh rename to python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh b/python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh rename to python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh b/python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh rename to python/sglang/kernels/jit/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/inkling_gate_topk_renorm.cuh b/python/sglang/kernels/jit/csrc/moe/inkling_gate_topk_renorm.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/inkling_gate_topk_renorm.cuh rename to python/sglang/kernels/jit/csrc/moe/inkling_gate_topk_renorm.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu b/python/sglang/kernels/jit/csrc/moe/moe_align_kernel.cu similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu rename to python/sglang/kernels/jit/csrc/moe/moe_align_kernel.cu diff --git a/python/sglang/jit_kernel/csrc/moe/moe_finalize_fuse_shared.cu b/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/moe_finalize_fuse_shared.cu rename to python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu diff --git a/python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh b/python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh rename to python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/moe_permute_prepare.cu b/python/sglang/kernels/jit/csrc/moe/moe_permute_prepare.cu similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/moe_permute_prepare.cu rename to python/sglang/kernels/jit/csrc/moe/moe_permute_prepare.cu diff --git a/python/sglang/jit_kernel/csrc/moe/moe_topk_sigmoid.cuh b/python/sglang/kernels/jit/csrc/moe/moe_topk_sigmoid.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/moe_topk_sigmoid.cuh rename to python/sglang/kernels/jit/csrc/moe/moe_topk_sigmoid.cuh diff --git a/python/sglang/jit_kernel/csrc/moe/tvm_ffi_utils.h b/python/sglang/kernels/jit/csrc/moe/tvm_ffi_utils.h similarity index 100% rename from python/sglang/jit_kernel/csrc/moe/tvm_ffi_utils.h rename to python/sglang/kernels/jit/csrc/moe/tvm_ffi_utils.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp b/python/sglang/kernels/jit/csrc/ngram_corpus/ngram.cpp similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp rename to python/sglang/kernels/jit/csrc/ngram_corpus/ngram.cpp diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h b/python/sglang/kernels/jit/csrc/ngram_corpus/ngram.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/ngram.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp b/python/sglang/kernels/jit/csrc/ngram_corpus/ngram_corpus_ffi.cpp similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp rename to python/sglang/kernels/jit/csrc/ngram_corpus/ngram_corpus_ffi.cpp diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/param.h b/python/sglang/kernels/jit/csrc/ngram_corpus/param.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/param.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/param.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/queue.h b/python/sglang/kernels/jit/csrc/ngram_corpus/queue.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/queue.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/queue.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp b/python/sglang/kernels/jit/csrc/ngram_corpus/result.cpp similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp rename to python/sglang/kernels/jit/csrc/ngram_corpus/result.cpp diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/result.h b/python/sglang/kernels/jit/csrc/ngram_corpus/result.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/result.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/result.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.cpp b/python/sglang/kernels/jit/csrc/ngram_corpus/suffix_automaton.cpp similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.cpp rename to python/sglang/kernels/jit/csrc/ngram_corpus/suffix_automaton.cpp diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.h b/python/sglang/kernels/jit/csrc/ngram_corpus/suffix_automaton.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/suffix_automaton.h diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/trie.cpp b/python/sglang/kernels/jit/csrc/ngram_corpus/trie.cpp similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/trie.cpp rename to python/sglang/kernels/jit/csrc/ngram_corpus/trie.cpp diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/trie.h b/python/sglang/kernels/jit/csrc/ngram_corpus/trie.h similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_corpus/trie.h rename to python/sglang/kernels/jit/csrc/ngram_corpus/trie.h diff --git a/python/sglang/jit_kernel/csrc/ngram_embedding.cuh b/python/sglang/kernels/jit/csrc/ngram_embedding.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/ngram_embedding.cuh rename to python/sglang/kernels/jit/csrc/ngram_embedding.cuh diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/config.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/config.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/config.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/config.h diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/defines.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/defines.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/defines.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/defines.h diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_transpose_v.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_transpose_v.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_transpose_v.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_transpose_v.h diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_utils.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_utils.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_utils.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/dense_fp8_utils.h diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/entry.cuh b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/entry.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/entry.cuh rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/entry.cuh diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/helpers.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/helpers.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/helpers.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/helpers.h diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/kernel.cuh b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/kernel.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/kernel.cuh rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/kernel.cuh diff --git a/python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/params.h b/python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/params.h similarity index 100% rename from python/sglang/jit_kernel/csrc/sparse_mla_q8kv8_prefill_sm90/params.h rename to python/sglang/kernels/jit/csrc/sparse_mla_q8kv8_prefill_sm90/params.h diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh b/python/sglang/kernels/jit/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh rename to python/sglang/kernels/jit/csrc/trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu b/python/sglang/kernels/jit/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu similarity index 100% rename from python/sglang/jit_kernel/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu rename to python/sglang/kernels/jit/csrc/trtllm_lora_temp/moe_lora_merged_align_kernel.cu diff --git a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh b/python/sglang/kernels/jit/csrc/trtllm_lora_temp/topk_softmax_pack.cuh similarity index 99% rename from python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh rename to python/sglang/kernels/jit/csrc/trtllm_lora_temp/topk_softmax_pack.cuh index 7a85f612a..f470b4dd9 100644 --- a/python/sglang/jit_kernel/csrc/trtllm_lora_temp/topk_softmax_pack.cuh +++ b/python/sglang/kernels/jit/csrc/trtllm_lora_temp/topk_softmax_pack.cuh @@ -58,7 +58,7 @@ __device__ float convert_to_float(T x) { } } -// Reference pack (bit-identical to jit_kernel/flashinfer_trtllm_moe/topk_pack.py): +// Reference pack (bit-identical to kernels/ops/moe/trtllm_lora_temp/topk_pack.py): // low 16 bits = bf16(weight) bits (round-to-nearest-even, same as torch/triton // `.to(bfloat16)`), high 16 bits = int16 expert id. __device__ __forceinline__ int32_t pack_routed(int32_t id, float w) { diff --git a/python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh b/python/sglang/kernels/jit/include/sgl_kernel/atomic.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/atomic.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/cta.cuh b/python/sglang/kernels/jit/include/sgl_kernel/cta.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/cta.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/cta.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/compress.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/compress.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/compress_v2.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/compress_v2.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/kvcacheio.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/kvcacheio.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh b/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/ffi.h b/python/sglang/kernels/jit/include/sgl_kernel/ffi.h similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/ffi.h rename to python/sglang/kernels/jit/include/sgl_kernel/ffi.h diff --git a/python/sglang/jit_kernel/include/sgl_kernel/impl/norm.cuh b/python/sglang/kernels/jit/include/sgl_kernel/impl/norm.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/impl/norm.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/impl/norm.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh b/python/sglang/kernels/jit/include/sgl_kernel/math.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/math.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/math.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh b/python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/runtime.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/scalar_type.hpp b/python/sglang/kernels/jit/include/sgl_kernel/scalar_type.hpp similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/scalar_type.hpp rename to python/sglang/kernels/jit/include/sgl_kernel/scalar_type.hpp diff --git a/python/sglang/jit_kernel/include/sgl_kernel/source_location.h b/python/sglang/kernels/jit/include/sgl_kernel/source_location.h similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/source_location.h rename to python/sglang/kernels/jit/include/sgl_kernel/source_location.h diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/kernels/jit/include/sgl_kernel/tensor.h similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/tensor.h rename to python/sglang/kernels/jit/include/sgl_kernel/tensor.h diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh b/python/sglang/kernels/jit/include/sgl_kernel/tile.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/tile.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/tile.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh b/python/sglang/kernels/jit/include/sgl_kernel/type.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/type.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/type.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/utils.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/utils.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/kernels/jit/include/sgl_kernel/utils.h similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/utils.h rename to python/sglang/kernels/jit/include/sgl_kernel/utils.h diff --git a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh b/python/sglang/kernels/jit/include/sgl_kernel/vec.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/vec.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/vec.cuh diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh similarity index 100% rename from python/sglang/jit_kernel/include/sgl_kernel/warp.cuh rename to python/sglang/kernels/jit/include/sgl_kernel/warp.cuh diff --git a/python/sglang/jit_kernel/__init__.py b/python/sglang/kernels/jit/tests/deepseek_v4/__init__.py similarity index 100% rename from python/sglang/jit_kernel/__init__.py rename to python/sglang/kernels/jit/tests/deepseek_v4/__init__.py diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/common.py b/python/sglang/kernels/jit/tests/deepseek_v4/common.py similarity index 100% rename from python/sglang/jit_kernel/tests/deepseek_v4/common.py rename to python/sglang/kernels/jit/tests/deepseek_v4/common.py diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/__init__.py b/python/sglang/kernels/jit/tests/kv_canary/__init__.py similarity index 100% rename from python/sglang/jit_kernel/tests/deepseek_v4/__init__.py rename to python/sglang/kernels/jit/tests/kv_canary/__init__.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py b/python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py similarity index 98% rename from python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py rename to python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py index c7e67ace1..9aa1e595c 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py +++ b/python/sglang/kernels/jit/tests/kv_canary/_canary_helpers.py @@ -5,14 +5,14 @@ from typing import Optional import torch -from sglang.jit_kernel.tests.kv_canary._constants import ( +from sglang.kernels.jit.tests.kv_canary._constants import ( _I64_SIGN_BIT, _U64_MASK, DEFAULT_NUM_SLOTS, DEFAULT_RING_CAPACITY, DEFAULT_SLOT_STRIDE_BYTES, ) -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +from sglang.kernels.jit.tests.kv_canary._fixtures import ( make_real_kv_source, make_real_kv_sources, ) diff --git a/python/sglang/jit_kernel/tests/kv_canary/_constants.py b/python/sglang/kernels/jit/tests/kv_canary/_constants.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/_constants.py rename to python/sglang/kernels/jit/tests/kv_canary/_constants.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/_differential.py b/python/sglang/kernels/jit/tests/kv_canary/_differential.py similarity index 99% rename from python/sglang/jit_kernel/tests/kv_canary/_differential.py rename to python/sglang/kernels/jit/tests/kv_canary/_differential.py index 8a0ac1a7e..fd2c26148 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_differential.py +++ b/python/sglang/kernels/jit/tests/kv_canary/_differential.py @@ -5,7 +5,7 @@ from typing import Any, Callable, Iterator, Optional import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_buf_equal, assert_canary_state_equal, diff --git a/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py b/python/sglang/kernels/jit/tests/kv_canary/_fixtures.py similarity index 98% rename from python/sglang/jit_kernel/tests/kv_canary/_fixtures.py rename to python/sglang/kernels/jit/tests/kv_canary/_fixtures.py index 2a04cb97a..70fbb7587 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_fixtures.py +++ b/python/sglang/kernels/jit/tests/kv_canary/_fixtures.py @@ -5,7 +5,7 @@ from typing import Literal, Optional import torch -from sglang.jit_kernel.tests.kv_canary._constants import DEFAULT_NUM_SLOTS +from sglang.kernels.jit.tests.kv_canary._constants import DEFAULT_NUM_SLOTS from sglang.kernels.ops.kv_canary.verify import ( RealKvSource, VerifyPlan, diff --git a/python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py b/python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py similarity index 95% rename from python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py rename to python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py index 6af7c305b..49919e068 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_fuzz_driver.py +++ b/python/sglang/kernels/jit/tests/kv_canary/_fuzz_driver.py @@ -3,7 +3,7 @@ from __future__ import annotations import random from typing import Any, Callable -from sglang.jit_kernel.tests.kv_canary._differential import ( +from sglang.kernels.jit.tests.kv_canary._differential import ( ShrinkResult, shrink_inputs, ) diff --git a/python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py b/python/sglang/kernels/jit/tests/kv_canary/_hand_oracle.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/_hand_oracle.py rename to python/sglang/kernels/jit/tests/kv_canary/_hand_oracle.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/_invariants.py b/python/sglang/kernels/jit/tests/kv_canary/_invariants.py similarity index 99% rename from python/sglang/jit_kernel/tests/kv_canary/_invariants.py rename to python/sglang/kernels/jit/tests/kv_canary/_invariants.py index 2bc7621c6..192c79152 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_invariants.py +++ b/python/sglang/kernels/jit/tests/kv_canary/_invariants.py @@ -11,7 +11,7 @@ from typing import Optional import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import FakeViolationLog +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 diff --git a/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py b/python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py similarity index 99% rename from python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py rename to python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py index 3fff5425e..3038d8405 100644 --- a/python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py +++ b/python/sglang/kernels/jit/tests/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/jit_kernel/tests/test_minimax_m3_mxfp8.py -v +Run: pytest python/sglang/kernels/jit/tests/test_minimax_m3_mxfp8.py -v """ import pytest diff --git a/python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py b/python/sglang/kernels/jit/tests/test_minimax_m3_rmsnorm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_minimax_m3_rmsnorm.py rename to python/sglang/kernels/jit/tests/test_minimax_m3_rmsnorm.py diff --git a/python/sglang/jit_kernel/tests/test_moe_preprocess.py b/python/sglang/kernels/jit/tests/test_moe_preprocess.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_moe_preprocess.py rename to python/sglang/kernels/jit/tests/test_moe_preprocess.py diff --git a/python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py b/python/sglang/kernels/jit/tests/test_moe_topk_sigmoid.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_moe_topk_sigmoid.py rename to python/sglang/kernels/jit/tests/test_moe_topk_sigmoid.py diff --git a/python/sglang/jit_kernel/tests/test_sconv_decode_metadata.py b/python/sglang/kernels/jit/tests/test_sconv_decode_metadata.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_sconv_decode_metadata.py rename to python/sglang/kernels/jit/tests/test_sconv_decode_metadata.py diff --git a/python/sglang/jit_kernel/tests/test_sconv_extend_metadata.py b/python/sglang/kernels/jit/tests/test_sconv_extend_metadata.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_sconv_extend_metadata.py rename to python/sglang/kernels/jit/tests/test_sconv_extend_metadata.py diff --git a/python/sglang/jit_kernel/tests/utils.py b/python/sglang/kernels/jit/tests/utils.py similarity index 100% rename from python/sglang/jit_kernel/tests/utils.py rename to python/sglang/kernels/jit/tests/utils.py diff --git a/python/sglang/kernels/jit/utils/compile.py b/python/sglang/kernels/jit/utils/compile.py index f054eaa56..e44821d25 100644 --- a/python/sglang/kernels/jit/utils/compile.py +++ b/python/sglang/kernels/jit/utils/compile.py @@ -73,7 +73,9 @@ def _local_jit_source_hash(source_files: List[str]) -> str: @cache_once def _resolve_kernel_path() -> pathlib.Path: # Resolve via the package spec so the lookup is location-independent. - spec = importlib.util.find_spec("sglang.jit_kernel") + # The C++/CUDA sources (``csrc/``) and headers (``include/``) live next to + # this package under ``sglang.kernels.jit`` (RFC #29630). + spec = importlib.util.find_spec("sglang.kernels.jit") assert spec is not None and spec.origin is not None cur_dir = pathlib.Path(spec.origin).parent.resolve() @@ -90,7 +92,7 @@ def _resolve_kernel_path() -> pathlib.Path: path = _environment_install() or _package_install() if path is None: - raise RuntimeError("Cannot find sglang.jit_kernel path") + raise RuntimeError("Cannot find sglang.kernels.jit path") return path diff --git a/python/sglang/kernels/ops/activation/__init__.py b/python/sglang/kernels/ops/activation/__init__.py index 4450dee8f..fe5248c3c 100644 --- a/python/sglang/kernels/ops/activation/__init__.py +++ b/python/sglang/kernels/ops/activation/__init__.py @@ -119,7 +119,7 @@ class SiluAndMulOp(_GatedActivationOp): } descriptions = { KernelBackend.AOT: "silu_and_mul (sgl_kernel wheel).", - KernelBackend.JIT: "silu_and_mul (sglang.jit_kernel).", + KernelBackend.JIT: "silu_and_mul (sglang.kernels.jit).", KernelBackend.AITER: "silu_and_mul (aiter, ROCm).", KernelBackend.TORCH: "silu_and_mul (pure-torch reference).", } @@ -153,7 +153,7 @@ class GeluAndMulOp(_GatedActivationOp): kernel_attr = "gelu_and_mul" descriptions = { KernelBackend.AOT: "gelu_and_mul (sgl_kernel wheel).", - KernelBackend.JIT: "gelu_and_mul (sglang.jit_kernel).", + KernelBackend.JIT: "gelu_and_mul (sglang.kernels.jit).", KernelBackend.TORCH: "gelu_and_mul (pure-torch reference).", } @@ -170,7 +170,7 @@ class GeluTanhAndMulOp(_GatedActivationOp): kernel_attr = "gelu_tanh_and_mul" descriptions = { KernelBackend.AOT: "gelu_tanh_and_mul (sgl_kernel wheel).", - KernelBackend.JIT: "gelu_tanh_and_mul (sglang.jit_kernel).", + KernelBackend.JIT: "gelu_tanh_and_mul (sglang.kernels.jit).", KernelBackend.TORCH: "gelu_tanh_and_mul (pure-torch reference).", } @@ -195,7 +195,7 @@ class ReLU2Op(BaseFusedOp): description="relu(x) ** 2; returns tensor", ) descriptions = { - KernelBackend.JIT: "relu(x)**2 (sglang.jit_kernel).", + KernelBackend.JIT: "relu(x)**2 (sglang.kernels.jit).", KernelBackend.TORCH: "relu(x)**2 (pure-torch reference).", } diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py index 06088e136..90f6de2fb 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py @@ -959,7 +959,7 @@ def flash_decode_with_topk_idx( # Single-stage JIT radix-select: one kernel, no intermediate buffers. # Equivalent output to the 2-stage path (set of block ids, front-packed, # -1 padded); ~2-16x faster for long context. See - # sglang/jit_kernel/minimax_decode_topk.py. + # sglang/kernels/ops/attention/minimax_decode_topk.py. from sglang.kernels.ops.attention.minimax_decode_topk import minimax_decode_topk minimax_decode_topk(score, seq_lens, block_size, topk, out=topk_idx) diff --git a/python/sglang/kernels/ops/communication/__init__.py b/python/sglang/kernels/ops/communication/__init__.py index 1d20cb8ff..b932b763d 100644 --- a/python/sglang/kernels/ops/communication/__init__.py +++ b/python/sglang/kernels/ops/communication/__init__.py @@ -5,7 +5,7 @@ wrappers are exposed here: the collective ops (custom all-reduce and friends) are stateful — they manage workspaces / IPC handles and are driven through a ``CustomAllreduce``-style object and ``torch.ops.sgl_kernel.*`` bindings rather than standalone callable kernels, so a thin ``sglang.kernels.ops`` forwarder -would be misleading. Import them from ``sgl_kernel`` / ``sglang.jit_kernel`` +would be misleading. Import them from ``sgl_kernel`` / ``sglang.kernels.jit`` directly until a proper stateful-op interface is designed. """ diff --git a/python/sglang/kernels/ops/communication/mp.py b/python/sglang/kernels/ops/communication/mp.py index 61da78e6c..5a961ddb5 100644 --- a/python/sglang/kernels/ops/communication/mp.py +++ b/python/sglang/kernels/ops/communication/mp.py @@ -1,7 +1,7 @@ """Multi-process / multi-GPU launching utilities (torchrun-based). -Shared `multigpu_launch` helper that both `sglang.jit_kernel.tests.utils` and -`sglang.jit_kernel.benchmark.utils` build their domain-specific entry points on +Shared `multigpu_launch` helper that both `sglang.kernels.jit.tests.utils` and +`sglang.kernels.jit.benchmark.utils` build their domain-specific entry points on top of (`multigpu_pytest_main`, `multigpu_bench_main`). When a script that calls one of those wrappers is run with plain `python`, the diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 85345459f..4baeb0ce5 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -29,7 +29,7 @@ register_kernel( target="sglang.kernels.ops.diffusion.group_norm_silu:apply_group_norm_silu", capabilities=_CUDA, format_signature=FormatSignature(description="fused GroupNorm + SiLU"), - description="Fused group-norm + SiLU (sglang.jit_kernel).", + description="Fused group-norm + SiLU (sglang.kernels.jit).", ) ) register_kernel( @@ -39,7 +39,7 @@ register_kernel( target="sglang.kernels.ops.diffusion.residual_gate_add:residual_gate_add_cuda", capabilities=_CUDA, format_signature=FormatSignature(description="residual + gate * update"), - description="Fused residual gate-add (sglang.jit_kernel).", + description="Fused residual gate-add (sglang.kernels.jit).", ) ) register_kernel( @@ -51,7 +51,7 @@ register_kernel( format_signature=FormatSignature( in_place=True, description="fused in-place QK-norm + RoPE" ), - description="Fused QK-norm + RoPE (sglang.jit_kernel).", + description="Fused QK-norm + RoPE (sglang.kernels.jit).", ) ) diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index a07ae0671..5fafd7652 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/python/sglang/kernels/ops/gemm/__init__.py @@ -65,7 +65,7 @@ register_kernel( supported_dtypes=("bfloat16",), description="DeepSeek-V3 fused QKV-A GEMM (drop-in with AOT signature)", ), - description="DeepSeek-V3 fused-A GEMM (sglang.jit_kernel).", + description="DeepSeek-V3 fused-A GEMM (sglang.kernels.jit).", ) ) register_kernel( @@ -78,7 +78,7 @@ register_kernel( supported_dtypes=("bfloat16",), description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]", ), - description="DeepSeek-V3 router GEMM (sglang.jit_kernel, JIT-only).", + description="DeepSeek-V3 router GEMM (sglang.kernels.jit, JIT-only).", ) ) diff --git a/python/sglang/kernels/ops/gemm/fused_a_gemm.py b/python/sglang/kernels/ops/gemm/fused_a_gemm.py index 72f012b32..4671e48d9 100644 --- a/python/sglang/kernels/ops/gemm/fused_a_gemm.py +++ b/python/sglang/kernels/ops/gemm/fused_a_gemm.py @@ -2,7 +2,7 @@ Dispatches to one of two interchangeable implementations via ``backend``: -- ``"jit"``: runtime-compiled CUDA C++ (``sglang.jit_kernel.dsv3_fused_a_gemm``). +- ``"jit"``: runtime-compiled CUDA C++ (``sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm``). - ``"cutedsl"``: CuTe DSL (``sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm``). - ``"auto"``: CuTe DSL on SM120+, otherwise the JIT kernel. @@ -69,7 +69,9 @@ def dsv3_fused_a_gemm( backend = _AUTO_BACKEND if backend == FusedAGemmBackend.JIT: - from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm as impl + from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import ( + dsv3_fused_a_gemm as impl, + ) else: from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import ( dsv3_fused_a_gemm as impl, diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 44f3630c0..8a70c0bd1 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -66,7 +66,7 @@ class RMSNormOp(BaseFusedOp): ) descriptions = { KernelBackend.AOT: "RMS normalization (sgl_kernel wheel).", - KernelBackend.JIT: "RMS normalization (sglang.jit_kernel).", + KernelBackend.JIT: "RMS normalization (sglang.kernels.jit).", KernelBackend.AITER: "RMS normalization (aiter rmsnorm2d_fwd, ROCm).", KernelBackend.TORCH_NPU: "RMS normalization (torch_npu, Ascend).", KernelBackend.TORCH: "RMS normalization (pure-torch reference).", @@ -180,7 +180,7 @@ class FusedAddRMSNormOp(BaseFusedOp): "Fused residual-add + RMS normalization (sgl_kernel wheel)." ), KernelBackend.JIT: ( - "Fused residual-add + RMS normalization (sglang.jit_kernel)." + "Fused residual-add + RMS normalization (sglang.kernels.jit)." ), KernelBackend.AITER: ("Fused residual-add + RMS normalization (aiter, ROCm)."), KernelBackend.TORCH_NPU: ( @@ -289,7 +289,7 @@ class GemmaRMSNormOp(BaseFusedOp): descriptions = { KernelBackend.AOT: "Gemma-style RMS normalization (sgl_kernel wheel).", KernelBackend.JIT: ( - "Gemma-style RMS normalization (rocm-triton, sglang.jit_kernel)." + "Gemma-style RMS normalization (rocm-triton, sglang.kernels.jit)." ), KernelBackend.TORCH_NPU: ("Gemma-style RMS normalization (torch_npu, Ascend)."), KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).", @@ -382,7 +382,7 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): KernelBackend.AOT: ("Gemma-style fused residual-add + RMS normalization."), KernelBackend.JIT: ( "Gemma-style fused residual-add + RMS normalization " - "(rocm-triton, sglang.jit_kernel)." + "(rocm-triton, sglang.kernels.jit)." ), KernelBackend.TORCH: ( "Gemma-style fused residual-add + RMS normalization " diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index cb09aef67..492e171fd 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -40,7 +40,7 @@ register_kernel( in_place=True, description="MoE align-block-size (JIT variant, AOT signature)", ), - description="MoE align-block-size (sglang.jit_kernel).", + description="MoE align-block-size (sglang.kernels.jit).", ) ) register_kernel( diff --git a/python/sglang/kernels/ops/quantization/__init__.py b/python/sglang/kernels/ops/quantization/__init__.py index a744b97ba..2e377121c 100644 --- a/python/sglang/kernels/ops/quantization/__init__.py +++ b/python/sglang/kernels/ops/quantization/__init__.py @@ -67,7 +67,7 @@ register_kernel( "layouts, optional fused silu_and_mul and masked EP-MoE schedule" ), ), - description="Unified per-token-group quantization (sglang.jit_kernel).", + description="Unified per-token-group quantization (sglang.kernels.jit).", ) ) diff --git a/python/sglang/kernels/spec.py b/python/sglang/kernels/spec.py index 1db35d2ec..dba590ac1 100644 --- a/python/sglang/kernels/spec.py +++ b/python/sglang/kernels/spec.py @@ -2,7 +2,7 @@ This module defines small, dependency-free descriptors used to *inventory* kernel implementations and drive a simple, heuristic dispatch. It intentionally -does not import ``torch``, ``sgl_kernel`` or ``sglang.jit_kernel`` at module +does not import ``torch``, ``sgl_kernel`` or ``sglang.kernels.jit`` at module import time so that ``import sglang.kernels`` stays cheap and works on a CPU-only box (see RFC #29630, Phase 2). @@ -29,7 +29,7 @@ import msgspec class KernelBackend(str, Enum): """Provenance of a kernel implementation (how it is built), not its device. - ``JIT`` (``sglang.jit_kernel``, compiles under nvcc *and* hipcc) and ``AOT`` + ``JIT`` (``sglang.kernels.jit``, compiles under nvcc *and* hipcc) and ``AOT`` (the ``sgl_kernel`` wheel, built for CUDA *and* ROCm) are both cross-device; which devices a given op supports is expressed by its :class:`CapabilityRequirement` list. Platform-specific libraries (e.g. @@ -40,7 +40,7 @@ class KernelBackend(str, Enum): TORCH = "torch" # pure-torch reference (forward_native) TORCH_COMPILE = "torch_compile" # torch.compile(forward_native) TRITON = "triton" - JIT = "jit" # sglang.jit_kernel (nvcc / hipcc) + JIT = "jit" # sglang.kernels.jit (nvcc / hipcc) AOT = "aot" # sgl_kernel wheel (CUDA / ROCm builds) CUTE_DSL = "cute_dsl" FLASHINFER = "flashinfer" diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index f73f5aa6c..1acb738d4 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -10,17 +10,17 @@ framework-specific optimization workflow. - `python/sglang/multimodal_gen/runtime/layers/elementwise.py` - `python/sglang/multimodal_gen/runtime/layers/fused_scale_shift_gate.py` - `python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py` -- `python/sglang/jit_kernel/diffusion/triton/scale_shift.py` -- `python/sglang/jit_kernel/diffusion/group_norm_silu.py` -- `python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py` -- `python/sglang/jit_kernel/diffusion/triton/norm.py` -- `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` -- `python/sglang/jit_kernel/diffusion/triton/rotary.py` -- `python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py` -- `python/sglang/jit_kernel/diffusion/residual_gate_add.py` -- `python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh` -- `python/sglang/jit_kernel/diffusion/triton/varlen_pack_pad.py` -- `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py` +- `python/sglang/kernels/ops/diffusion/triton/scale_shift.py` +- `python/sglang/kernels/ops/diffusion/group_norm_silu.py` +- `python/sglang/kernels/ops/diffusion/triton/group_norm_silu.py` +- `python/sglang/kernels/ops/diffusion/triton/norm.py` +- `python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py` +- `python/sglang/kernels/ops/diffusion/triton/rotary.py` +- `python/sglang/kernels/ops/diffusion/triton/ltx2_rotary.py` +- `python/sglang/kernels/ops/diffusion/residual_gate_add.py` +- `python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh` +- `python/sglang/kernels/ops/diffusion/triton/varlen_pack_pad.py` +- `python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py` - `test/registered/jit/diffusion/test_qwen_image_modulation.py` - `test/registered/jit/diffusion/test_group_norm_silu.py` - `test/registered/jit/diffusion/test_residual_gate_add.py` @@ -29,7 +29,7 @@ framework-specific optimization workflow. - `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py` - `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py` - `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py` -- `python/sglang/jit_kernel/norm.py` +- `python/sglang/kernels/ops/layernorm/_jit_norm.py` - `python/sglang/multimodal_gen/runtime/platforms/cuda.py` - `python/sglang/multimodal_gen/runtime/layers/attention/selector.py` - `docs_new/docs/sglang-diffusion/attention_backends.mdx` (repo root) @@ -137,7 +137,7 @@ framework-specific optimization workflow. **QK Norm Optimization** - Entry point: `apply_qk_norm` in `layernorm.py`. -- Fast path: JIT fused inplace QK norm from `python/sglang/jit_kernel/norm.py` via `fused_inplace_qknorm`. +- Fast path: JIT fused inplace QK norm from `python/sglang/kernels/ops/layernorm/_jit_norm.py` via `fused_inplace_qknorm`. - Preconditions for fused path: - CUDA only. - `allow_inplace=True` and `q_eps == k_eps`. @@ -149,7 +149,7 @@ framework-specific optimization workflow. **QK Norm + RoPE Optimization** - Entry point: `apply_qk_norm_rope` in `layernorm.py`. -- Fast path: JIT fused inplace QK norm + RoPE from `python/sglang/jit_kernel/diffusion/qknorm_rope.py` via `fused_inplace_qknorm_rope`. +- Fast path: JIT fused inplace QK norm + RoPE from `python/sglang/kernels/ops/diffusion/qknorm_rope.py` via `fused_inplace_qknorm_rope`. - Toggle: `SGLANG_ENABLE_FUSED_QKNORM_ROPE=1` keeps the fused path enabled by default. - Preconditions for fused path: - CUDA only. diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py index 868120d77..1e116dc10 100644 --- a/python/sglang/multimodal_gen/runtime/layers/activation.py +++ b/python/sglang/multimodal_gen/runtime/layers/activation.py @@ -19,7 +19,7 @@ _is_npu = current_platform.is_npu() _is_xpu = current_platform.is_xpu() if _is_cuda: - from sglang.jit_kernel.activation import silu_and_mul + from sglang.kernels.ops.activation._jit_activation import silu_and_mul elif _is_hip or _is_xpu: from sgl_kernel import silu_and_mul diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 2b98d67af..57f05d68b 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -11,13 +11,16 @@ import torch import torch.nn as nn import torch.nn.functional as F -from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm from sglang.kernels.ops.diffusion.qknorm_rope import ( can_use_fused_inplace_qknorm_rope, fused_inplace_qknorm_rope, ) from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm from sglang.kernels.ops.diffusion.triton.scale_shift import fuse_scale_shift_kernel +from sglang.kernels.ops.layernorm._jit_norm import ( + can_use_fused_inplace_qknorm, + fused_inplace_qknorm, +) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index a7db9d27a..2c9a2a123 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -57,7 +57,7 @@ _is_xpu = is_xpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: - from sglang.jit_kernel.activation import ( + from sglang.kernels.ops.activation._jit_activation import ( gelu_and_mul, gelu_tanh_and_mul, relu2, diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index b996632b6..3722af7eb 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -12,7 +12,9 @@ import torch.nn as nn import torch.nn.functional as F from einops import rearrange -from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm as can_use_jit_qk_norm +from sglang.kernels.ops.layernorm._jit_norm import ( + can_use_fused_inplace_qknorm as can_use_jit_qk_norm, +) from sglang.srt.environ import envs from sglang.srt.models.utils import apply_qk_norm from sglang.srt.runtime_context import get_parallel diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index d57c5cf8f..b23a7cce6 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -147,8 +147,10 @@ if _is_cuda: _jit_rmsnorm_hf = None - from sglang.jit_kernel.norm import fused_add_rmsnorm as _jit_fused_add_rmsnorm - from sglang.jit_kernel.norm import ( + from sglang.kernels.ops.layernorm._jit_norm import ( + fused_add_rmsnorm as _jit_fused_add_rmsnorm, + ) + from sglang.kernels.ops.layernorm._jit_norm import ( is_supported_jit_fused_add_rmsnorm_hidden_size, ) diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index ad16e3c7e..dc80ed9f5 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -18,7 +18,7 @@ if _is_cuda: shuffle_rows, ) - from sglang.jit_kernel.activation import silu_and_mul + from sglang.kernels.ops.activation._jit_activation import silu_and_mul def cutlass_fused_experts_fp8( diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index ee6332367..3e965580a 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -18,14 +18,10 @@ if _is_cuda_alike: ) if _is_cuda: - from sglang.jit_kernel.activation import silu_and_mul + from sglang.kernels.ops.activation._jit_activation import silu_and_mul else: from sgl_kernel import silu_and_mul -from sglang.jit_kernel.per_tensor_quant_fp8 import ( - per_tensor_absmax_fp8, - per_tensor_quant_fp8, -) from sglang.kernels.ops.moe.ep_moe_kernels import ( cutlass_w4_run_moe_ep_preproess, deepep_ll_get_cutlass_w4a8_moe_mm_data, @@ -39,6 +35,10 @@ from sglang.kernels.ops.moe.ep_moe_kernels import ( silu_mul_dynamic_tensorwise_quant_for_cutlass_moe, silu_mul_static_tensorwise_quant_for_cutlass_moe, ) +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_absmax_fp8, + per_tensor_quant_fp8, +) def cutlass_w4a8_moe( diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index f7a07c012..c7779b61f 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -11,7 +11,7 @@ _is_cuda = is_cuda() if _is_cuda: from sgl_kernel import moe_sum_reduce - from sglang.jit_kernel.activation import silu_and_mul + from sglang.kernels.ops.activation._jit_activation import silu_and_mul from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py index f36d77b20..f0e8ee6a1 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py @@ -30,7 +30,7 @@ if is_sm120_supported(): update_opt_flags_constraints({"is_persistent": False}) if is_cuda(): - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul + from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul else: from sgl_kernel import gelu_and_mul, silu_and_mul diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index a0e84f732..16a7697ab 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -56,7 +56,9 @@ _is_musa = is_musa() # Imported only for the SGLANG_OPT_FIX_MEGA_MOE_MEMORY=False fallback path. if not (_is_npu or _is_hip) and _is_cuda: - from sglang.jit_kernel.activation import silu_and_mul as _legacy_silu_and_mul + from sglang.kernels.ops.activation._jit_activation import ( + silu_and_mul as _legacy_silu_and_mul, + ) elif _is_musa: _silu_and_mul_musa = torch.nn.SwishGLU() else: diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index 9b456f122..90f940495 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -60,7 +60,7 @@ _is_musa = is_musa() if _is_cuda: from sgl_kernel import moe_sum_reduce - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul + from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index b87276287..accbc5d17 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -2230,5 +2230,5 @@ def select_experts( # NOTE: the AOT sgl_kernel::moe_fused_gate and sgl_kernel::kimi_k2_moe_fused_gate # ops (and their torch.compile fake impls) were retired here — both CUDA gate -# paths now route through the unified Triton router (jit_kernel/moe_fused_gate.py), +# paths now route through the unified Triton router (kernels/ops/moe/moe_fused_gate.py), # whose Python impl is traceable directly, so no register_fake shim is needed. diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index e5fbc24ee..e4eda93fa 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -51,7 +51,7 @@ if _is_cuda: ggml_mul_mat_vec_a8, ) - from sglang.jit_kernel.activation import gelu_and_mul, silu_and_mul + from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul elif _is_musa: from sgl_kernel import gelu_and_mul, moe_align_block_size, moe_sum, silu_and_mul from sgl_kernel.quantization import ( diff --git a/python/sglang/srt/lora/trtllm_lora_temp/environ.py b/python/sglang/srt/lora/trtllm_lora_temp/environ.py index 0a0495606..1ab0611dc 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/environ.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/environ.py @@ -108,7 +108,7 @@ class _LoraEnvs: # diagnostics / tuning: SGLANG_OPT_LORA_SHRINK_TUNE = _GatedBool("SGLANG_OPT_LORA_SHRINK_TUNE", False) - # kimi NVFP4 permute+quant fuse — read in jit_kernel/trtllm_lora_temp/core.py (Python) to pass + # kimi NVFP4 permute+quant fuse — read in kernels/ops/moe/trtllm_lora_temp/core.py (Python) to pass # a bool to the kernel, AND C++-side via getenv in the launcher. Default off (kimi-only). SGLANG_OPT_FUSED_PERMUTE_QUANT = _GatedBool("SGLANG_OPT_FUSED_PERMUTE_QUANT", False) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 46512b69c..312be7fc0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -213,7 +213,7 @@ if _use_aiter: pass if _is_cuda: - from sglang.jit_kernel.dsv3_router_gemm import ( + from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import ( dsv3_router_gemm as _jit_dsv3_router_gemm, ) elif _is_npu: diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index dd767b2f5..97e0f2325 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -24,8 +24,11 @@ import torch import triton import triton.language as tl -from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm from sglang.kernels.ops.attention.rope import FusedSetKVBufferArg +from sglang.kernels.ops.layernorm._jit_norm import ( + can_use_fused_inplace_qknorm, + fused_inplace_qknorm, +) from sglang.srt.environ import envs from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled diff --git a/scripts/code_sync/utils.py b/scripts/code_sync/utils.py index ed9be8323..31b3cc067 100644 --- a/scripts/code_sync/utils.py +++ b/scripts/code_sync/utils.py @@ -18,7 +18,7 @@ FOLDER_NAMES = [ "docs", "examples", "python/sglang/lang", - "python/sglang/jit_kernel", + "python/sglang/kernels/jit", "python/sglang/srt", "python/sglang/test", "python/sglang/utils.py", diff --git a/sgl-kernel/benchmark/bench_fp8_gemm.py b/sgl-kernel/benchmark/bench_fp8_gemm.py index e37653d51..35e10cb2e 100644 --- a/sgl-kernel/benchmark/bench_fp8_gemm.py +++ b/sgl-kernel/benchmark/bench_fp8_gemm.py @@ -8,7 +8,9 @@ import torch import triton from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.utils import is_in_ci # Optional vLLM import diff --git a/sgl-kernel/benchmark/bench_fp8_gemm_swap_ab.py b/sgl-kernel/benchmark/bench_fp8_gemm_swap_ab.py index 99803a46a..2d2127dd2 100644 --- a/sgl-kernel/benchmark/bench_fp8_gemm_swap_ab.py +++ b/sgl-kernel/benchmark/bench_fp8_gemm_swap_ab.py @@ -18,7 +18,9 @@ import torch import triton from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.utils import is_in_ci IS_CI = is_in_ci() diff --git a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py index 107123465..fee2e4158 100644 --- a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py @@ -8,7 +8,9 @@ import torch import triton import triton.testing -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.utils import is_in_ci # Optional imports diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 4748c3cd1..37c1a8fa8 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -174,7 +174,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // moe_fused_gate / kimi_k2_moe_fused_gate (AOT) retired: the CUDA gate/topk path // now routes through the unified Triton router - // (python/sglang/jit_kernel/moe_fused_gate.py). + // (python/sglang/kernels/ops/moe/moe_fused_gate.py). m.def( "fp8_blockwise_scaled_grouped_mm(Tensor output, Tensor a_ptrs, Tensor b_ptrs, Tensor out_ptrs, Tensor " diff --git a/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu b/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu index 935278d4b..138daf094 100644 --- a/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu +++ b/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu @@ -14,7 +14,7 @@ limitations under the License. ==============================================================================*/ // DeepSeek-V4 fused norm + RoPE kernels, ported from JIT kernel -// python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh +// python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh // to sgl-kernel AOT compilation with CUDA + HIP (ROCm) support. #ifndef USE_ROCM diff --git a/sgl-kernel/python/sgl_kernel/moe.py b/sgl-kernel/python/sgl_kernel/moe.py index 30bf0fad6..c01a0ac37 100755 --- a/sgl-kernel/python/sgl_kernel/moe.py +++ b/sgl-kernel/python/sgl_kernel/moe.py @@ -106,7 +106,7 @@ def moe_sum( # moe_fused_gate / kimi_k2_moe_fused_gate (AOT gate kernels) retired — the gate/topk # path is consolidated onto the unified Triton router in -# python/sglang/jit_kernel/moe_fused_gate.py (sglang issue #26771). +# python/sglang/kernels/ops/moe/moe_fused_gate.py (sglang issue #26771). def fp8_blockwise_scaled_grouped_mm( diff --git a/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py index 6d319f843..1a4d90a9d 100644 --- a/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py +++ b/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py @@ -3,8 +3,10 @@ from typing import Optional, Tuple, Union import torch -from sglang.jit_kernel.rope import FusedSetKVBufferArg as _JitFusedSetKVBufferArg -from sglang.jit_kernel.rope import ( +from sglang.kernels.ops.attention.rope import ( + FusedSetKVBufferArg as _JitFusedSetKVBufferArg, +) +from sglang.kernels.ops.attention.rope import ( apply_rope_with_cos_sin_cache_inplace as _jit_apply_rope_with_cos_sin_cache_inplace, ) diff --git a/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py b/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py index 5a6b109e7..969775fdd 100644 --- a/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py +++ b/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py @@ -5,7 +5,9 @@ import torch from sgl_kernel import cutlass_w4a8_moe_mm from utils import is_hopper -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) def pack_int4_values_to_int8(int4_values_interleaved: torch.Tensor) -> torch.Tensor: diff --git a/test/README.md b/test/README.md index cfa439061..b11a7671d 100644 --- a/test/README.md +++ b/test/README.md @@ -72,7 +72,7 @@ Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and ru Keep `est_time`, `stage`, `runner_config` as **literal values** — `run_suite.py` collects them by AST parsing. -JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path): +JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/kernels/jit/` and are imported by absolute path): - Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large` - Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large` diff --git a/test/registered/attention/unittests/KNOWN_FAILURES.md b/test/registered/attention/unittests/KNOWN_FAILURES.md index bec1b0340..eefbf2d5c 100644 --- a/test/registered/attention/unittests/KNOWN_FAILURES.md +++ b/test/registered/attention/unittests/KNOWN_FAILURES.md @@ -45,7 +45,7 @@ ImportError: cannot import name 'flash_attn_varlen_func' from 'flash_attn' ``` **Root cause**: `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` -via `sglang.jit_kernel.flash_attention`. On SM 8.x / 9.x that resolves to +via `sglang.kernels.ops.attention.flash_attention`. On SM 8.x / 9.x that resolves to sgl-kernel's FA3 build (works on H200). On other SMs, the JIT kernel falls back to the upstream `flash_attn` (FA2) wheel — but the `lmsysorg/sglang:nightly-dev-cu13` container's `flash_attn` package on diff --git a/test/registered/attention/unittests/dsa/README.md b/test/registered/attention/unittests/dsa/README.md index 4ba5a6c4d..cf6a82d90 100644 --- a/test/registered/attention/unittests/dsa/README.md +++ b/test/registered/attention/unittests/dsa/README.md @@ -29,7 +29,7 @@ hardware/SDK. The variant tests live in `test_dsa.py` as |---|---|---|---|---| | `flashmla_sparse` | ✓ | ✓ | ✓ | SM>=9.0 + `sgl_kernel.flash_mla` | | `flashmla_kv` | ✓ | ✓ | ✓ | SM>=9.0 + `sgl_kernel.flash_mla` | -| `fa3` | ✓ | ✓ | ✓ | SM>=9.0 + `sglang.jit_kernel.flash_attention` | +| `fa3` | ✓ | ✓ | ✓ | SM>=9.0 + `sglang.kernels.ops.attention.flash_attention` | | `tilelang` | ✓ (topk=2048 dedicated fixture) | ✓ (topk=2048 dedicated fixture) | skipped: not yet wired into CG runner | `tilelang_sparse_fwd` asserts `topk == 2048`; the topk=2048 fixture instance (`build_dsa_sparse_attention_fixture(..., index_topk=2048)`) is used by `test_sparse_tilelang_prefill_case` / `test_sparse_tilelang_decode_case`. The default-topk impl-variant matrix still skips tilelang with the same reason. **SM10.x container gate**: on Blackwell the tilelang JIT generates `wait_wgmma` which the container's MMA template doesn't ship (`KNOWN_FAILURES.md §2`); `dsa_impl_capability("tilelang")` skips on `major >= 10`. Set `SGLANG_TEST_DSA_TILELANG_FORCE=1` to override after re-imaging. | | `trtllm` | skipped: SM<10 | skipped: SM<10 | skipped: SM<10 | TRT-LLM Gen FMHA/MLA requires Blackwell (SM>=10.0). | | `aiter` | skipped: not HIP | skipped: not HIP | skipped: not HIP | AMD-only kernel library. | diff --git a/test/registered/attention/unittests/dual_chunk/README.md b/test/registered/attention/unittests/dual_chunk/README.md index d9c027cd7..18e6231ac 100644 --- a/test/registered/attention/unittests/dual_chunk/README.md +++ b/test/registered/attention/unittests/dual_chunk/README.md @@ -46,7 +46,7 @@ Columns are runner modes; rows are kernel-path modes of the single ## Container Gate (SM10.x) `DualChunkFlashAttentionBackend` calls `flash_attn_varlen_func` via -`sglang.jit_kernel.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's +`sglang.kernels.ops.attention.flash_attention`. On SM8.x / SM9.x that resolves to sgl-kernel's FA3 build; on SM != {8, 9} (notably SM10.x / GB300) the JIT kernel falls back to the upstream `flash_attn` (FA2) wheel, which the `lmsysorg/sglang:nightly-dev-cu13` container ships without an SM10.x-compiled diff --git a/test/registered/jit/benchmark/bench_activation.py b/test/registered/jit/benchmark/bench_activation.py index 773b8581f..c80d8bab0 100644 --- a/test/registered/jit/benchmark/bench_activation.py +++ b/test/registered/jit/benchmark/bench_activation.py @@ -4,12 +4,18 @@ from sgl_kernel import gelu_and_mul as gelu_and_mul_aot from sgl_kernel import gelu_tanh_and_mul as gelu_tanh_and_mul_aot from sgl_kernel import silu_and_mul as silu_and_mul_aot -from sglang.jit_kernel.activation import gelu_and_mul as gelu_and_mul_jit -from sglang.jit_kernel.activation import gelu_tanh_and_mul as gelu_tanh_and_mul_jit -from sglang.jit_kernel.activation import relu2 as relu2_jit -from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.ops.activation._jit_activation import ( + gelu_and_mul as gelu_and_mul_jit, +) +from sglang.kernels.ops.activation._jit_activation import ( + gelu_tanh_and_mul as gelu_tanh_and_mul_jit, +) +from sglang.kernels.ops.activation._jit_activation import relu2 as relu2_jit +from sglang.kernels.ops.activation._jit_activation import ( + silu_and_mul as silu_and_mul_jit, +) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_add_constant.py b/test/registered/jit/benchmark/bench_add_constant.py index 2a86dad20..e0d20985d 100644 --- a/test/registered/jit/benchmark/bench_add_constant.py +++ b/test/registered/jit/benchmark/bench_add_constant.py @@ -2,7 +2,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark_no_cudagraph, diff --git a/test/registered/jit/benchmark/bench_awq_dequantize.py b/test/registered/jit/benchmark/bench_awq_dequantize.py index bd3842acc..6e547a4fb 100644 --- a/test/registered/jit/benchmark/bench_awq_dequantize.py +++ b/test/registered/jit/benchmark/bench_awq_dequantize.py @@ -4,7 +4,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import run_benchmark +from sglang.kernels.jit.benchmark.utils import run_benchmark from sglang.kernels.ops.quantization.awq_dequantize import ( awq_dequantize as jit_awq_dequantize, ) diff --git a/test/registered/jit/benchmark/bench_clamp_position.py b/test/registered/jit/benchmark/bench_clamp_position.py index 01d698857..393479683 100644 --- a/test/registered/jit/benchmark/bench_clamp_position.py +++ b/test/registered/jit/benchmark/bench_clamp_position.py @@ -4,7 +4,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/bench_concat_mla.py b/test/registered/jit/benchmark/bench_concat_mla.py index 314742d34..fb6183f65 100644 --- a/test/registered/jit/benchmark/bench_concat_mla.py +++ b/test/registered/jit/benchmark/bench_concat_mla.py @@ -6,7 +6,7 @@ import triton.testing from sgl_kernel import concat_mla_absorb_q as aot_absorb_q from sgl_kernel import concat_mla_k as aot_k -from sglang.jit_kernel.benchmark.utils import run_benchmark +from sglang.kernels.jit.benchmark.utils import run_benchmark from sglang.kernels.ops.attention.concat_mla import concat_mla_absorb_q as jit_absorb_q from sglang.kernels.ops.attention.concat_mla import concat_mla_k as jit_k from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py index 667aa4375..aa918ef69 100644 --- a/test/registered/jit/benchmark/bench_custom_all_reduce.py +++ b/test/registered/jit/benchmark/bench_custom_all_reduce.py @@ -10,8 +10,8 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl from sglang.kernels.ops.communication.mp import register_comm_cleanup from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py index 429737b54..5de153d71 100644 --- a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py @@ -8,9 +8,9 @@ import torch import torch.nn.functional as F import triton.testing -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm +from sglang.kernels.jit.benchmark import marker from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import dsv3_fused_a_gemm from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import ( dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm, ) diff --git a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py index d5f04a9de..e85af0b8b 100644 --- a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py @@ -1,16 +1,16 @@ """Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch). Run on a Hopper (SM90+) GPU: - python -m sglang.jit_kernel.benchmark.bench_dsv3_router_gemm + python -m sglang.kernels.jit.benchmark.bench_dsv3_router_gemm """ import torch import torch.nn.functional as F -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random -from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime +from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import dsv3_router_gemm from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py b/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py index bb7569e72..1e11568a1 100644 --- a/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py +++ b/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py @@ -6,7 +6,7 @@ import torch import triton from sglang.benchmark.bench_utils import run_bench -from sglang.jit_kernel.benchmark.utils import get_benchmark_range +from sglang.kernels.jit.benchmark.utils import get_benchmark_range from sglang.srt.utils import is_sm100_supported from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py index 514c155c2..1df2da96d 100644 --- a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py +++ b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py @@ -5,7 +5,7 @@ import sys import torch import triton -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark from sglang.kernels.ops.gemm.fp8_blockwise_gemm import fp8_blockwise_scaled_mm from sglang.srt.utils import is_sm120_supported from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_fused_eh_norm.py b/test/registered/jit/benchmark/bench_fused_eh_norm.py index baff79839..8848a5c78 100644 --- a/test/registered/jit/benchmark/bench_fused_eh_norm.py +++ b/test/registered/jit/benchmark/bench_fused_eh_norm.py @@ -2,7 +2,7 @@ from __future__ import annotations import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py b/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py index 24d4e5350..8011e16ff 100644 --- a/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py +++ b/test/registered/jit/benchmark/bench_fused_fp8_qkv_kv_cache.py @@ -1,6 +1,6 @@ import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.kvcache.fused_fp8_qkv_kv_cache import fused_fp8_qkv_kv_cache from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_fused_qknorm_rope.py b/test/registered/jit/benchmark/bench_fused_qknorm_rope.py index 17408d7f7..51942890e 100644 --- a/test/registered/jit/benchmark/bench_fused_qknorm_rope.py +++ b/test/registered/jit/benchmark/bench_fused_qknorm_rope.py @@ -14,7 +14,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark from sglang.kernels.ops.attention.fused_qknorm_rope import ( fused_qk_norm_rope as fused_qk_norm_rope_jit, ) diff --git a/test/registered/jit/benchmark/bench_hadamard.py b/test/registered/jit/benchmark/bench_hadamard.py index 76beebbd8..5a2da9f23 100644 --- a/test/registered/jit/benchmark/bench_hadamard.py +++ b/test/registered/jit/benchmark/bench_hadamard.py @@ -7,7 +7,7 @@ import torch.nn.functional as F import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, diff --git a/test/registered/jit/benchmark/bench_hicache.py b/test/registered/jit/benchmark/bench_hicache.py index 9bd45e535..054f3e2ea 100644 --- a/test/registered/jit/benchmark/bench_hicache.py +++ b/test/registered/jit/benchmark/bench_hicache.py @@ -25,7 +25,7 @@ import triton import triton.testing from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer -from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range +from sglang.kernels.jit.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range from sglang.kernels.ops.kvcache.hicache import ( can_use_hicache_jit_kernel, transfer_hicache_all_layer, diff --git a/test/registered/jit/benchmark/bench_hisparse.py b/test/registered/jit/benchmark/bench_hisparse.py index cb589ce57..051dcba21 100644 --- a/test/registered/jit/benchmark/bench_hisparse.py +++ b/test/registered/jit/benchmark/bench_hisparse.py @@ -5,7 +5,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE +from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py index 64a2c3520..cc0da748d 100644 --- a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py +++ b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py @@ -8,7 +8,7 @@ import triton import triton.language as tl import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, DEFAULT_QUANTILES, diff --git a/test/registered/jit/benchmark/bench_moe_fused_gate.py b/test/registered/jit/benchmark/bench_moe_fused_gate.py index 5c4365cf6..e220cd490 100644 --- a/test/registered/jit/benchmark/bench_moe_fused_gate.py +++ b/test/registered/jit/benchmark/bench_moe_fused_gate.py @@ -1,7 +1,7 @@ import torch -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_mxfp8_moe.py b/test/registered/jit/benchmark/bench_mxfp8_moe.py index d00fa1142..6b75424a2 100644 --- a/test/registered/jit/benchmark/bench_mxfp8_moe.py +++ b/test/registered/jit/benchmark/bench_mxfp8_moe.py @@ -6,7 +6,7 @@ from typing import Any import torch import triton -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark from sglang.kernels.ops.quantization.mxfp8 import ( es_sm100_mxfp8_blockscaled_grouped_quant, es_sm100_mxfp8_blockscaled_moe_grouped_gemm, diff --git a/test/registered/jit/benchmark/bench_ngram_compute_decode.py b/test/registered/jit/benchmark/bench_ngram_compute_decode.py index 934c51063..7f60bb4c7 100644 --- a/test/registered/jit/benchmark/bench_ngram_compute_decode.py +++ b/test/registered/jit/benchmark/bench_ngram_compute_decode.py @@ -2,7 +2,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark_no_cudagraph, diff --git a/test/registered/jit/benchmark/bench_ngram_update_token_table.py b/test/registered/jit/benchmark/bench_ngram_update_token_table.py index 3198713b6..bb179fbac 100644 --- a/test/registered/jit/benchmark/bench_ngram_update_token_table.py +++ b/test/registered/jit/benchmark/bench_ngram_update_token_table.py @@ -2,7 +2,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark_no_cudagraph, diff --git a/test/registered/jit/benchmark/bench_norm.py b/test/registered/jit/benchmark/bench_norm.py index 790e0c62b..f4a52596d 100644 --- a/test/registered/jit/benchmark/bench_norm.py +++ b/test/registered/jit/benchmark/bench_norm.py @@ -6,9 +6,11 @@ import triton.testing from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm from flashinfer.norm import rmsnorm as fi_rmsnorm -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm -from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.ops.layernorm._jit_norm import ( + fused_add_rmsnorm as jit_fused_add_rmsnorm, +) +from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/jit/benchmark/bench_online_c128_mtp.py index 71137aa88..80376b8dd 100644 --- a/test/registered/jit/benchmark/bench_online_c128_mtp.py +++ b/test/registered/jit/benchmark/bench_online_c128_mtp.py @@ -9,7 +9,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py b/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py index e9e049f3e..f9c9cef0b 100644 --- a/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py +++ b/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py @@ -4,8 +4,10 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci( diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant.py b/test/registered/jit/benchmark/bench_per_token_group_quant.py index 873ffe992..8ba53f2cc 100644 --- a/test/registered/jit/benchmark/bench_per_token_group_quant.py +++ b/test/registered/jit/benchmark/bench_per_token_group_quant.py @@ -1,10 +1,12 @@ -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_empty, create_random +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_empty, create_random # per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the # kernel is kept only as the perf baseline for this benchmark. -from sglang.jit_kernel.per_token_group_quant import per_token_group_quant -from sglang.jit_kernel.per_token_group_quant_8bit_v2 import ( +from sglang.kernels.ops.quantization._jit_per_token_group_quant import ( + per_token_group_quant, +) +from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import ( per_token_group_quant_8bit_v2, ) from sglang.kernels.ops.quantization.fp8_kernel import ( diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py b/test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py index af2efd29f..7b6b940f8 100644 --- a/test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py +++ b/test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py @@ -1,9 +1,9 @@ import torch from sgl_kernel import sgl_per_token_group_quant_8bit -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random -from sglang.jit_kernel.per_token_group_quant_8bit_v2 import ( +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import ( per_token_group_quant_8bit_v2, ) from sglang.kernels.ops.quantization.fp8_kernel import ( diff --git a/test/registered/jit/benchmark/bench_per_token_group_quant_masked.py b/test/registered/jit/benchmark/bench_per_token_group_quant_masked.py index 461130a50..a34c578e8 100644 --- a/test/registered/jit/benchmark/bench_per_token_group_quant_masked.py +++ b/test/registered/jit/benchmark/bench_per_token_group_quant_masked.py @@ -2,13 +2,15 @@ import math import torch -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_empty, create_random +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_empty, create_random # per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the # kernel is kept only as the perf baseline for this benchmark. -from sglang.jit_kernel.per_token_group_quant import per_token_group_quant -from sglang.jit_kernel.per_token_group_quant_8bit_v2 import ( +from sglang.kernels.ops.quantization._jit_per_token_group_quant import ( + per_token_group_quant, +) +from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import ( per_token_group_quant_8bit_v2, ) from sglang.kernels.ops.quantization.fp8_kernel import ( diff --git a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py b/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py index 3964f7e34..1c8a4d5a4 100644 --- a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py +++ b/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py @@ -1,6 +1,6 @@ import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.moe.ep_moe_kernels import ( post_reorder_deepgemm, post_reorder_triton_kernel, diff --git a/test/registered/jit/benchmark/bench_qknorm.py b/test/registered/jit/benchmark/bench_qknorm.py index 381ae2d7a..fdc738b1b 100644 --- a/test/registered/jit/benchmark/bench_qknorm.py +++ b/test/registered/jit/benchmark/bench_qknorm.py @@ -1,8 +1,8 @@ import torch -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import create_random -from sglang.jit_kernel.norm import fused_inplace_qknorm +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm from sglang.srt.utils import get_current_device_stream_fast from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_qknorm_across_heads.py b/test/registered/jit/benchmark/bench_qknorm_across_heads.py index 8429111ef..be6726f20 100644 --- a/test/registered/jit/benchmark/bench_qknorm_across_heads.py +++ b/test/registered/jit/benchmark/bench_qknorm_across_heads.py @@ -6,8 +6,8 @@ import triton import triton.testing from sgl_kernel import rmsnorm -from sglang.jit_kernel.benchmark.utils import run_benchmark -from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads +from sglang.kernels.jit.benchmark.utils import run_benchmark +from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm_across_heads from sglang.srt.utils import get_current_device_stream_fast from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci diff --git a/test/registered/jit/benchmark/bench_renorm.py b/test/registered/jit/benchmark/bench_renorm.py index a84427085..9195c032e 100644 --- a/test/registered/jit/benchmark/bench_renorm.py +++ b/test/registered/jit/benchmark/bench_renorm.py @@ -5,7 +5,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph +from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci diff --git a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py b/test/registered/jit/benchmark/bench_resolve_future_token_ids.py index 06c51c2d4..08ba11530 100644 --- a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py +++ b/test/registered/jit/benchmark/bench_resolve_future_token_ids.py @@ -4,7 +4,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/bench_rope.py b/test/registered/jit/benchmark/bench_rope.py index 0cbe7a69b..8091e7e65 100644 --- a/test/registered/jit/benchmark/bench_rope.py +++ b/test/registered/jit/benchmark/bench_rope.py @@ -4,7 +4,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, diff --git a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py index 42da45253..a6d369c0e 100644 --- a/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py +++ b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py @@ -14,14 +14,16 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, DEFAULT_QUANTILES, get_benchmark_range, ) -from sglang.jit_kernel.set_mla_kv_buffer import set_mla_kv_buffer as jit_set from sglang.kernels.jit.utils import is_arch_support_pdl +from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import ( + set_mla_kv_buffer as jit_set, +) from sglang.srt.mem_cache.utils import set_mla_kv_buffer_kernel as sglang_triton_kernel from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton as sglang_wrapper from sglang.test.ci.ci_register import register_cuda_ci diff --git a/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py b/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py index bfa56fbe1..e18964e0e 100644 --- a/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py +++ b/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py @@ -6,7 +6,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph +from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import ( sparse_mla_q8kv8_prefill_fwd, ) diff --git a/test/registered/jit/benchmark/bench_spec_topk1.py b/test/registered/jit/benchmark/bench_spec_topk1.py index fada9272c..d44cdaadd 100644 --- a/test/registered/jit/benchmark/bench_spec_topk1.py +++ b/test/registered/jit/benchmark/bench_spec_topk1.py @@ -6,7 +6,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/bench_store_cache.py b/test/registered/jit/benchmark/bench_store_cache.py index d70ba5f26..94564441b 100644 --- a/test/registered/jit/benchmark/bench_store_cache.py +++ b/test/registered/jit/benchmark/bench_store_cache.py @@ -1,7 +1,7 @@ import torch -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, create_empty, create_random, diff --git a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py b/test/registered/jit/benchmark/bench_symm_mem_all_gather.py index a06d72c53..51d8addbe 100644 --- a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py +++ b/test/registered/jit/benchmark/bench_symm_mem_all_gather.py @@ -26,8 +26,8 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main from sglang.kernels.jit.utils import cache_once from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, diff --git a/test/registered/jit/benchmark/bench_topk.py b/test/registered/jit/benchmark/bench_topk.py index f88d4fa1b..575fa1e6d 100644 --- a/test/registered/jit/benchmark/bench_topk.py +++ b/test/registered/jit/benchmark/bench_topk.py @@ -1,6 +1,6 @@ import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.attention.dsv4.topk import ( plan_topk_v2, topk_transform_512, diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py index 02673ed86..0895d65a3 100644 --- a/test/registered/jit/benchmark/bench_tp_qknorm.py +++ b/test/registered/jit/benchmark/bench_tp_qknorm.py @@ -23,8 +23,8 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.benchmark import marker -from sglang.jit_kernel.benchmark.utils import multigpu_bench_main +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import multigpu_bench_main from sglang.kernels.jit.utils import cache_once, get_ci_test_range from sglang.kernels.ops.communication.all_reduce import ( fused_parallel_qknorm, diff --git a/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py b/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py index 2082a1925..6b332f875 100644 --- a/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py +++ b/test/registered/jit/benchmark/bench_vocab_parallel_embedding.py @@ -3,7 +3,7 @@ import torch.nn.functional as F import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark from sglang.kernels.ops.embeddings.vocab_parallel_embedding import ( vocab_parallel_embedding, ) diff --git a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py b/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py index a259e1d98..44a864d59 100644 --- a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py +++ b/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py @@ -2,7 +2,7 @@ from dataclasses import dataclass import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import ( fused_causal_conv3d_cat_pad_cuda, ) diff --git a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py index 495ff93ed..3e4daa3ce 100644 --- a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py @@ -10,7 +10,7 @@ from typing import Any, Callable import flashinfer import torch -from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE +from sglang.kernels.jit.benchmark.utils import DEFAULT_DTYPE from sglang.kernels.jit.utils import KERNEL_PATH from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci @@ -26,7 +26,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent REPO_ROOT = ( Path(os.environ["SGLANG_NVFP4_REPO_ROOT"]) if os.environ.get("SGLANG_NVFP4_REPO_ROOT") - # Anchor on the installed jit_kernel package (python/sglang/jit_kernel) so + # Anchor on the installed jit_kernel package (python/sglang/kernels/jit) so # this stays correct regardless of where the benchmark file lives. else KERNEL_PATH.parents[2] ) diff --git a/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py b/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py index f96da7f72..56628bb37 100644 --- a/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py +++ b/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py @@ -8,7 +8,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph +from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNormScaleShift, diff --git a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py index b87320a96..3ff6af016 100644 --- a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py +++ b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py @@ -13,12 +13,14 @@ from typing import Callable import torch import torch.nn.functional as F -from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE -from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm -from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm +from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE from sglang.kernels.jit.utils import KERNEL_PATH from sglang.kernels.ops.diffusion.triton.norm import norm_infer, rms_norm_fn from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm +from sglang.kernels.ops.layernorm._jit_norm import ( + fused_add_rmsnorm as jit_fused_add_rmsnorm, +) +from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.utils import is_in_ci diff --git a/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py b/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py index 93e48d832..6ed86c581 100644 --- a/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py +++ b/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py @@ -5,7 +5,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, @@ -131,7 +131,7 @@ def clone_inputs( def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None: from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace - from sglang.jit_kernel.norm import fused_inplace_qknorm + from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm q = inputs["q"] k = inputs["k"] diff --git a/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py b/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py index 9204b5bbb..5b956a286 100644 --- a/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py +++ b/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py @@ -3,7 +3,7 @@ from typing import Tuple import torch import triton.testing -from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph +from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph from sglang.kernels.ops.diffusion.triton.norm import norm_infer from sglang.kernels.ops.diffusion.triton.scale_shift import ( fuse_layernorm_scale_shift_gate_select01_kernel, diff --git a/test/registered/jit/benchmark/kv_canary/bench_plan.py b/test/registered/jit/benchmark/kv_canary/bench_plan.py index 9f8d5cba4..9bf64fb67 100644 --- a/test/registered/jit/benchmark/kv_canary/bench_plan.py +++ b/test/registered/jit/benchmark/kv_canary/bench_plan.py @@ -7,7 +7,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.kv_canary.utils import ( +from sglang.kernels.jit.benchmark.kv_canary.utils import ( POOL_AXIS, SWA_WINDOW, BenchCase, @@ -15,7 +15,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import ( build_full_matrix_cases, naive_cumsum_fn, ) -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py b/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py index 0d9778b59..5eea09a34 100644 --- a/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py +++ b/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py @@ -6,7 +6,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark_no_cudagraph, diff --git a/test/registered/jit/benchmark/kv_canary/bench_verify.py b/test/registered/jit/benchmark/kv_canary/bench_verify.py index a34ff21f1..350754183 100644 --- a/test/registered/jit/benchmark/kv_canary/bench_verify.py +++ b/test/registered/jit/benchmark/kv_canary/bench_verify.py @@ -6,7 +6,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.kv_canary.utils import ( +from sglang.kernels.jit.benchmark.kv_canary.utils import ( RING_CAPACITY, SWA_WINDOW, BenchCase, @@ -16,7 +16,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import ( make_real_kv_sources, naive_slot_copy_fn, ) -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/kv_canary/bench_write.py b/test/registered/jit/benchmark/kv_canary/bench_write.py index 0366a2fda..79963e0bd 100644 --- a/test/registered/jit/benchmark/kv_canary/bench_write.py +++ b/test/registered/jit/benchmark/kv_canary/bench_write.py @@ -6,7 +6,7 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.kv_canary.utils import ( +from sglang.kernels.jit.benchmark.kv_canary.utils import ( RING_CAPACITY, SWA_WINDOW, BenchCase, @@ -16,7 +16,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import ( make_real_kv_sources, naive_slot_copy_fn, ) -from sglang.jit_kernel.benchmark.utils import ( +from sglang.kernels.jit.benchmark.utils import ( DEFAULT_DEVICE, get_benchmark_range, run_benchmark, diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py b/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py index e368e4dcc..215279cdd 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py +++ b/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py @@ -9,7 +9,7 @@ intermediate buffers; the baseline is two launches with split-K partials. import torch import triton -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.attention.minimax_decode_topk import minimax_decode_topk from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx import ( _topk_index_merge_kernel, diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py b/test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py index 61b49f6d0..fdb98a365 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py +++ b/test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py @@ -5,7 +5,7 @@ launch) vs the unfused path (GemmaRMSNorm(q) + GemmaRMSNorm(k) + rotary_emb, import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.attention.minimax_qknorm_rope import ( minimax_qknorm_rope, minimax_qknorm_rope_grouped, diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py b/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py index 29632bcc0..7bdd4c1e0 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py +++ b/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py @@ -3,7 +3,7 @@ per-buffer index_put_ stores (main K, main V, index K, optional index V).""" import torch -from sglang.jit_kernel.benchmark import marker +from sglang.kernels.jit.benchmark import marker from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/jit/deepseek_v4/test_c128_v2.py b/test/registered/jit/deepseek_v4/test_c128_v2.py index e09c68d89..373fd2a24 100644 --- a/test/registered/jit/deepseek_v4/test_c128_v2.py +++ b/test/registered/jit/deepseek_v4/test_c128_v2.py @@ -7,7 +7,7 @@ import pytest import torch import triton -from sglang.jit_kernel.tests.deepseek_v4.common import ( +from sglang.kernels.jit.tests.deepseek_v4.common import ( LegacyContext, PagedContext, make_legacy_context, diff --git a/test/registered/jit/deepseek_v4/test_c4_v2.py b/test/registered/jit/deepseek_v4/test_c4_v2.py index f193330ff..a9a7eb9d4 100644 --- a/test/registered/jit/deepseek_v4/test_c4_v2.py +++ b/test/registered/jit/deepseek_v4/test_c4_v2.py @@ -7,7 +7,7 @@ import pytest import torch import triton -from sglang.jit_kernel.tests.deepseek_v4.common import ( +from sglang.kernels.jit.tests.deepseek_v4.common import ( LegacyContext, PagedContext, make_legacy_context, diff --git a/test/registered/jit/diffusion/test_qknorm_rope.py b/test/registered/jit/diffusion/test_qknorm_rope.py index 3048569c0..bd617adb2 100644 --- a/test/registered/jit/diffusion/test_qknorm_rope.py +++ b/test/registered/jit/diffusion/test_qknorm_rope.py @@ -48,7 +48,7 @@ def split_qknorm_rope( ) -> None: from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace - from sglang.jit_kernel.norm import fused_inplace_qknorm + from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm fused_inplace_qknorm(q, k, q_weight, k_weight) apply_rope_with_cos_sin_cache_inplace( diff --git a/test/registered/jit/kv_canary/test_const_sync.py b/test/registered/jit/kv_canary/test_const_sync.py index 120f60db6..de39c8f74 100644 --- a/test/registered/jit/kv_canary/test_const_sync.py +++ b/test/registered/jit/kv_canary/test_const_sync.py @@ -3,7 +3,7 @@ from __future__ import annotations import re from pathlib import Path -import sglang.jit_kernel +import sglang.kernels.jit from sglang.kernels.ops.kv_canary import consts from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci @@ -11,10 +11,10 @@ register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-la register_amd_ci(est_time=5, stage="jit-kernel-unit", runner_config="amd") -# Resolve the kernel source against the installed jit_kernel package rather -# than this file's location, so the test stays correct wherever it lives. +# Resolve the kernel source against the installed sglang.kernels.jit package +# rather than this file's location, so the test stays correct wherever it lives. _CONSTS_CUH: Path = ( - Path(sglang.jit_kernel.__file__).resolve().parent + Path(sglang.kernels.jit.__file__).resolve().parent / "csrc" / "kv_canary" / "consts.cuh" diff --git a/test/registered/jit/kv_canary/test_kernel_config.py b/test/registered/jit/kv_canary/test_kernel_config.py index cbe407faa..2917da076 100644 --- a/test/registered/jit/kv_canary/test_kernel_config.py +++ b/test/registered/jit/kv_canary/test_kernel_config.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_buf_equal, assert_canary_state_equal, @@ -15,13 +15,13 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( make_write_plan_pair, stamp_clean_chain, ) -from sglang.jit_kernel.tests.kv_canary._differential import ( +from sglang.kernels.jit.tests.kv_canary._differential import ( _assert_plans_byte_equal, _run_both_plan, _run_both_verify, _run_both_write, ) -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +from sglang.kernels.jit.tests.kv_canary._fixtures import ( dummy_pseudo_tensors, empty_extras, ) diff --git a/test/registered/jit/kv_canary/test_pipeline_e2e.py b/test/registered/jit/kv_canary/test_pipeline_e2e.py index 86e86c595..24fa49b2a 100644 --- a/test/registered/jit/kv_canary/test_pipeline_e2e.py +++ b/test/registered/jit/kv_canary/test_pipeline_e2e.py @@ -5,7 +5,7 @@ from typing import Any, Optional import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_buf_equal, assert_canary_state_equal, @@ -14,7 +14,7 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( stamp_clean_chain, write_slot_fields, ) -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +from sglang.kernels.jit.tests.kv_canary._fixtures import ( clone_real_kv_sources, empty_extras, make_req_to_token, diff --git a/test/registered/jit/kv_canary/test_plan_fuzz.py b/test/registered/jit/kv_canary/test_plan_fuzz.py index 67385ee66..25e3df6fc 100644 --- a/test/registered/jit/kv_canary/test_plan_fuzz.py +++ b/test/registered/jit/kv_canary/test_plan_fuzz.py @@ -7,19 +7,19 @@ from typing import Optional import pytest import torch -from sglang.jit_kernel.tests.kv_canary._differential import _run_both_plan -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +from sglang.kernels.jit.tests.kv_canary._differential import _run_both_plan +from sglang.kernels.jit.tests.kv_canary._fixtures import ( allocate_plan_pair, derive_plan_capacity, make_lut, make_padding_mask, make_req_to_token, ) -from sglang.jit_kernel.tests.kv_canary._fuzz_driver import ( +from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( FUZZ_SEEDS_PR, run_fuzz_combo, ) -from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants +from sglang.kernels.jit.tests.kv_canary._invariants import PlanInvariants 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") diff --git a/test/registered/jit/kv_canary/test_plan_hand.py b/test/registered/jit/kv_canary/test_plan_hand.py index d6ab5cfed..bbf543cec 100644 --- a/test/registered/jit/kv_canary/test_plan_hand.py +++ b/test/registered/jit/kv_canary/test_plan_hand.py @@ -5,15 +5,15 @@ import random import pytest import torch -from sglang.jit_kernel.tests.kv_canary._differential import run_plan_diff -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +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.jit_kernel.tests.kv_canary._invariants import PlanInvariants +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, diff --git a/test/registered/jit/kv_canary/test_utils.py b/test/registered/jit/kv_canary/test_utils.py index 9f74f64d0..82736140a 100644 --- a/test/registered/jit/kv_canary/test_utils.py +++ b/test/registered/jit/kv_canary/test_utils.py @@ -1,6 +1,6 @@ from __future__ import annotations -from sglang.jit_kernel.benchmark.kv_canary.utils import ( +from sglang.kernels.jit.benchmark.kv_canary.utils import ( MAX_EXTEND_TOKENS_PER_FORWARD, build_fast_matrix_cases, build_full_matrix_cases, diff --git a/test/registered/jit/kv_canary/test_verify_fuzz.py b/test/registered/jit/kv_canary/test_verify_fuzz.py index c26068f0b..56fa62eba 100644 --- a/test/registered/jit/kv_canary/test_verify_fuzz.py +++ b/test/registered/jit/kv_canary/test_verify_fuzz.py @@ -6,23 +6,23 @@ from dataclasses import dataclass import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +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.jit_kernel.tests.kv_canary._differential import _run_both_verify -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +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.jit_kernel.tests.kv_canary._fuzz_driver import ( +from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( FUZZ_SEEDS_PR, run_fuzz_combo, ) -from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants +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, diff --git a/test/registered/jit/kv_canary/test_verify_hand.py b/test/registered/jit/kv_canary/test_verify_hand.py index f28b6accb..ea68821e8 100644 --- a/test/registered/jit/kv_canary/test_verify_hand.py +++ b/test/registered/jit/kv_canary/test_verify_hand.py @@ -8,7 +8,7 @@ from typing import Callable import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( FakeViolationLog, assert_only_bits_set, chain_anchor_signed, @@ -26,12 +26,12 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( to_signed_int64, write_slot_fields, ) -from sglang.jit_kernel.tests.kv_canary._differential import ( +from sglang.kernels.jit.tests.kv_canary._differential import ( _run_both_verify, run_verify_diff, ) -from sglang.jit_kernel.tests.kv_canary._fixtures import clone_real_kv_sources -from sglang.jit_kernel.tests.kv_canary._hand_oracle import ( +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, ) diff --git a/test/registered/jit/kv_canary/test_write_fuzz.py b/test/registered/jit/kv_canary/test_write_fuzz.py index 8951eff42..87d8427f7 100644 --- a/test/registered/jit/kv_canary/test_write_fuzz.py +++ b/test/registered/jit/kv_canary/test_write_fuzz.py @@ -6,23 +6,23 @@ from dataclasses import dataclass import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +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.jit_kernel.tests.kv_canary._differential import _run_both_write -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +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.jit_kernel.tests.kv_canary._fuzz_driver import ( +from sglang.kernels.jit.tests.kv_canary._fuzz_driver import ( FUZZ_SEEDS_PR, run_fuzz_combo, ) -from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants +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, diff --git a/test/registered/jit/kv_canary/test_write_hand.py b/test/registered/jit/kv_canary/test_write_hand.py index 009e4da31..9cc763daf 100644 --- a/test/registered/jit/kv_canary/test_write_hand.py +++ b/test/registered/jit/kv_canary/test_write_hand.py @@ -7,7 +7,7 @@ from unittest.mock import patch import pytest import torch -from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( +from sglang.kernels.jit.tests.kv_canary._canary_helpers import ( FakeViolationLog, assert_canary_state_equal, assert_only_bits_set, @@ -24,15 +24,15 @@ from sglang.jit_kernel.tests.kv_canary._canary_helpers import ( stamp_pair, to_signed_int64, ) -from sglang.jit_kernel.tests.kv_canary._differential import ( +from sglang.kernels.jit.tests.kv_canary._differential import ( _run_both_write, run_write_diff, ) -from sglang.jit_kernel.tests.kv_canary._fixtures import ( +from sglang.kernels.jit.tests.kv_canary._fixtures import ( clone_real_kv_sources, dummy_pseudo_tensors, ) -from sglang.jit_kernel.tests.kv_canary._hand_oracle import ( +from sglang.kernels.jit.tests.kv_canary._hand_oracle import ( _hand_fold_all, _hand_fold_partial, ) diff --git a/test/registered/jit/test_activation.py b/test/registered/jit/test_activation.py index dd9975f7b..e8244c5b1 100644 --- a/test/registered/jit/test_activation.py +++ b/test/registered/jit/test_activation.py @@ -4,12 +4,12 @@ import pytest import torch import torch.nn.functional as F -from sglang.jit_kernel.activation import ( +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.activation._jit_activation import ( SUPPORTED_ACTIVATIONS, relu2, run_activation, ) -from sglang.kernels.jit.utils import get_ci_test_range from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py index 38ad20e95..38b2b0568 100644 --- a/test/registered/jit/test_custom_all_reduce.py +++ b/test/registered/jit/test_custom_all_reduce.py @@ -28,7 +28,7 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.tests.utils import multigpu_pytest_main +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, diff --git a/test/registered/jit/test_cutedsl_gdn.py b/test/registered/jit/test_cutedsl_gdn.py index 584d645b5..b7fce1054 100644 --- a/test/registered/jit/test_cutedsl_gdn.py +++ b/test/registered/jit/test_cutedsl_gdn.py @@ -13,7 +13,7 @@ try: import cutlass # noqa: F401 from cutlass.cute.runtime import from_dlpack - from sglang.jit_kernel import cutedsl_gdn + from sglang.kernels.ops.attention import cutedsl_gdn CUTEDSL_AVAILABLE = True except ImportError: diff --git a/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py b/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py index a08be320b..2658de98a 100644 --- a/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py +++ b/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py @@ -89,7 +89,7 @@ from typing import Callable, Iterable, Literal, Optional import pytest import torch -from sglang.jit_kernel.tests.deepseek_v4.common import ( +from sglang.kernels.jit.tests.deepseek_v4.common import ( make_legacy_context, to_seq_extend, ) diff --git a/test/registered/jit/test_dsv3_fused_a_gemm.py b/test/registered/jit/test_dsv3_fused_a_gemm.py index cc6489ebf..128e8734a 100644 --- a/test/registered/jit/test_dsv3_fused_a_gemm.py +++ b/test/registered/jit/test_dsv3_fused_a_gemm.py @@ -6,12 +6,12 @@ import pytest import torch import torch.nn.functional as F -from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm from sglang.kernels.jit.utils import ( get_ci_test_range, get_jit_cuda_arch, is_hip_runtime, ) +from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import dsv3_fused_a_gemm from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_dsv3_router_gemm.py b/test/registered/jit/test_dsv3_router_gemm.py index 928318ac3..00204776f 100644 --- a/test/registered/jit/test_dsv3_router_gemm.py +++ b/test/registered/jit/test_dsv3_router_gemm.py @@ -6,12 +6,12 @@ import sys import pytest import torch -from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm from sglang.kernels.jit.utils import ( get_ci_test_range, get_jit_cuda_arch, is_hip_runtime, ) +from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import dsv3_router_gemm from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_fused_add_rmsnorm.py b/test/registered/jit/test_fused_add_rmsnorm.py index afdabb241..1bb41dd92 100644 --- a/test/registered/jit/test_fused_add_rmsnorm.py +++ b/test/registered/jit/test_fused_add_rmsnorm.py @@ -20,7 +20,7 @@ def sglang_jit_fused_add_rmsnorm( *, cast_x_before_out_mul: bool = False, ) -> None: - from sglang.jit_kernel.norm import fused_add_rmsnorm + from sglang.kernels.ops.layernorm._jit_norm import fused_add_rmsnorm fused_add_rmsnorm( input, residual, weight, eps, cast_x_before_out_mul=cast_x_before_out_mul diff --git a/test/registered/jit/test_hadamard_jit.py b/test/registered/jit/test_hadamard_jit.py index fad99752d..dec9981fc 100644 --- a/test/registered/jit/test_hadamard_jit.py +++ b/test/registered/jit/test_hadamard_jit.py @@ -19,7 +19,7 @@ from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=128, stage="base-b-kernel-unit", runner_config="1-gpu-large") # Exact M×N Hadamard matrices (±1 entries) copied from -# python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py. +# python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py. # These are non-power-of-2 Hadamard matrices constructed via Paley/Williamson methods. # "+" = +1, "-" = -1. Used by the _12n/_20n/_28n/_40n kernel variants. @@ -145,7 +145,7 @@ def _parse_hadamard_str(s): # Parsed M×M special Hadamard matrices, keyed by M (the "multiple"). -# Copied from python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py +# Copied from python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py # (had_12_paley, had_20_will, had_28_will, had_40_tpal) _SPECIAL_MATRICES = { 12: _parse_hadamard_str(_HAD_12_STR), @@ -187,7 +187,7 @@ def hadamard_transform_mn_ref(x, multiple, scale=1.0): x = x.reshape(-1, dim) # The kernel requires dim % (4*M) == 0 (for vectorized memory access). - # See python/sglang/jit_kernel/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc. + # See python/sglang/kernels/ops/attention/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc. pad_multiple = 4 * multiple if dim % pad_multiple != 0: pad_size = pad_multiple - dim % pad_multiple @@ -319,7 +319,7 @@ def test_hadamard_transform_scale_one(dtype): # Test dimensions for M×N variants: dim = M * N where N = 2^k. # M = 12/20/28/40 are the non-power-of-2 Hadamard sizes registered in -# python/sglang/jit_kernel/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel). +# python/sglang/kernels/ops/attention/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel). # range(2,9) gives N = 4,8,...,256 so dims cover a practical range. _12N_DIMS = [12 * (2**k) for k in range(2, 9)] # 48, 96, ... , 3072 _20N_DIMS = [20 * (2**k) for k in range(2, 9)] # 80, 160, ... , 5120 diff --git a/test/registered/jit/test_per_tensor_quant_fp8.py b/test/registered/jit/test_per_tensor_quant_fp8.py index 589b0ae69..0af83ef3d 100644 --- a/test/registered/jit/test_per_tensor_quant_fp8.py +++ b/test/registered/jit/test_per_tensor_quant_fp8.py @@ -5,8 +5,10 @@ from typing import Optional, Tuple import pytest import torch -from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_per_token_group_quant.py b/test/registered/jit/test_per_token_group_quant.py index 485ab432f..d5ae71e33 100644 --- a/test/registered/jit/test_per_token_group_quant.py +++ b/test/registered/jit/test_per_token_group_quant.py @@ -21,8 +21,10 @@ import itertools import pytest import torch -from sglang.jit_kernel.per_token_group_quant import per_token_group_quant from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.quantization._jit_per_token_group_quant import ( + per_token_group_quant, +) from sglang.kernels.ops.quantization.fp8_kernel import ( create_per_token_group_quant_fp8_output_scale, fp8_dtype, diff --git a/test/registered/jit/test_per_token_group_quant_8bit_v2.py b/test/registered/jit/test_per_token_group_quant_8bit_v2.py index 3fcbca018..4d7403c13 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit_v2.py +++ b/test/registered/jit/test_per_token_group_quant_8bit_v2.py @@ -3,10 +3,10 @@ import itertools import pytest import torch -from sglang.jit_kernel.per_token_group_quant_8bit_v2 import ( +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import ( per_token_group_quant_8bit_v2, ) -from sglang.kernels.jit.utils import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_qknorm.py b/test/registered/jit/test_qknorm.py index 136cd189e..0669782ee 100644 --- a/test/registered/jit/test_qknorm.py +++ b/test/registered/jit/test_qknorm.py @@ -34,7 +34,7 @@ def sglang_jit_qknorm( q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: - from sglang.jit_kernel.norm import fused_inplace_qknorm + from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm fused_inplace_qknorm(q, k, q_weight, k_weight) diff --git a/test/registered/jit/test_qknorm_across_heads.py b/test/registered/jit/test_qknorm_across_heads.py index 8ab741d0d..e2efb28ab 100644 --- a/test/registered/jit/test_qknorm_across_heads.py +++ b/test/registered/jit/test_qknorm_across_heads.py @@ -19,7 +19,7 @@ def sglang_jit_qknorm_across_heads( q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: - from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads + from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm_across_heads fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight) diff --git a/test/registered/jit/test_rmsnorm.py b/test/registered/jit/test_rmsnorm.py index 206942240..e4d726408 100644 --- a/test/registered/jit/test_rmsnorm.py +++ b/test/registered/jit/test_rmsnorm.py @@ -26,7 +26,7 @@ def sglang_jit_rmsnorm( output: torch.Tensor | None = None, eps: float = EPS, ) -> None: - from sglang.jit_kernel.norm import rmsnorm + from sglang.kernels.ops.layernorm._jit_norm import rmsnorm rmsnorm(input, weight, out=output, eps=eps) @@ -127,7 +127,7 @@ def test_rmsnorm( @pytest.mark.parametrize("hidden_size", [64, 128, 256, 512, 8192, 8704, 16384]) def test_rmsnorm_hidden_size_support(hidden_size: int) -> None: - from sglang.jit_kernel.norm import _is_supported_rmsnorm_hidden_size + from sglang.kernels.ops.layernorm._jit_norm import _is_supported_rmsnorm_hidden_size assert _is_supported_rmsnorm_hidden_size(hidden_size) @@ -148,7 +148,7 @@ def test_rmsnorm_hidden_size_support(hidden_size: int) -> None: ], ) def test_rmsnorm_kernel_dispatch(hidden_size: int, expected: str) -> None: - from sglang.jit_kernel.norm import _rmsnorm_kernel_class + from sglang.kernels.ops.layernorm._jit_norm import _rmsnorm_kernel_class assert _rmsnorm_kernel_class(hidden_size) == expected diff --git a/test/registered/jit/test_set_mla_kv_buffer.py b/test/registered/jit/test_set_mla_kv_buffer.py index 2015a04eb..6ff6adec4 100644 --- a/test/registered/jit/test_set_mla_kv_buffer.py +++ b/test/registered/jit/test_set_mla_kv_buffer.py @@ -3,11 +3,11 @@ import sys import pytest import torch -from sglang.jit_kernel.set_mla_kv_buffer import ( +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import ( can_use_set_mla_kv_buffer, set_mla_kv_buffer, ) -from sglang.kernels.jit.utils import get_ci_test_range from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") diff --git a/test/registered/jit/test_symm_mem_all_gather.py b/test/registered/jit/test_symm_mem_all_gather.py index ea6b7e650..816964d9b 100644 --- a/test/registered/jit/test_symm_mem_all_gather.py +++ b/test/registered/jit/test_symm_mem_all_gather.py @@ -26,7 +26,7 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.tests.utils import multigpu_pytest_main +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, diff --git a/test/registered/jit/test_tp_qknorm.py b/test/registered/jit/test_tp_qknorm.py index aa1434ea6..31f715eb0 100644 --- a/test/registered/jit/test_tp_qknorm.py +++ b/test/registered/jit/test_tp_qknorm.py @@ -14,7 +14,7 @@ import torch.distributed as dist import triton import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.tests.utils import multigpu_pytest_main +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, diff --git a/test/registered/kernels/test_kernels_namespace.py b/test/registered/kernels/test_kernels_namespace.py index 79f19da23..4e13e7230 100644 --- a/test/registered/kernels/test_kernels_namespace.py +++ b/test/registered/kernels/test_kernels_namespace.py @@ -183,11 +183,12 @@ def test_platform_detect_does_not_raise(): def test_import_stays_metadata_only(): - # Importing the namespace must not pull in sgl_kernel / sglang.jit_kernel. + # Importing the namespace must not pull in the AOT backend (sgl_kernel) or + # the JIT compilation infra (sglang.kernels.jit), which import torch / nvcc. code = ( "import sys, sglang.kernels.ops; " "print('DIRTY' if 'sgl_kernel' in sys.modules or any(" - "m.startswith('sglang.jit_kernel') for m in sys.modules) else 'CLEAN')" + "m.startswith('sglang.kernels.jit') for m in sys.modules) else 'CLEAN')" ) r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) assert r.returncode == 0, r.stderr diff --git a/test/registered/unit/distributed/test_vmm_utils.py b/test/registered/unit/distributed/test_vmm_utils.py index f1cd47551..fd3e6f5c6 100644 --- a/test/registered/unit/distributed/test_vmm_utils.py +++ b/test/registered/unit/distributed/test_vmm_utils.py @@ -20,7 +20,7 @@ import torch import torch.distributed as dist from cuda.bindings import driver as drv -from sglang.jit_kernel.tests.utils import multigpu_pytest_main +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,