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

[Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels (#32072)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-23 08:35:09 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 8ce68370b5
commit 99f636a86f
354 changed files with 889 additions and 875 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/add-jit-kernel/SKILL.md
+35 -35
View File
@@ -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<Op>` 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<Op>` 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
```
.claude/skills/add-sgl-kernel/SKILL.md
+1 -1
View File
@@ -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`.
.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md
+18 -18
View File
@@ -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`<br>`python/sglang/srt/distributed/communication_op.py::tensor_model_parallel_fused_allreduce_rmsnorm`<br>`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`<br>`gelu_and_mul`<br>`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`<br>`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`<br>`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`<br>`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`<br>`benchmark_combo_kernel`<br>`q_norm`<br>`k_norm`<br>`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`<br>`rms_sumsq_serial`<br>`rms_apply_serial`<br>`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`<br>`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`<br>`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*`<br>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`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`<br>`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`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/kernels/ops/attention/rope.py`<br>`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`<br>`cache_seqlens_int32`<br>`cu_seqlens_k`<br>`page_table`<br>`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`<br>`fused_metadata_copy_multi`<br>`fused_nsa_cache_seqlens`<br>`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`<br>`fused_metadata_copy_multi`<br>`fused_nsa_cache_seqlens`<br>`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`<br>`fused_qkv_a_proj_with_mqa`<br>`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`<br>`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`<br>`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`<br>`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`<br>`mrope`<br>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`<br>`fused_append_shared_experts_with_weights`<br>`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`<br>`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<br>dispatch / combine<br>grouped top-k<br>many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`<br>`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`<br>`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*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/jit_kernel/nvfp4.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`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*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`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`<br>`fused_rms_mxfp4_quant`<br>`fused_flatten_fp8_group_quant`<br>`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`<br>`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`<br>`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`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`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`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/kernels/ops/attention/fused_store_index_cache.py`<br>`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`<br>`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`<br>`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`<br>`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*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`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*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`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`<br>`softplus`<br>`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`<br>`layer_norm_gated_fwd` | `python/sglang/srt/layers/attention/fla/fused_norm_gate.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`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`<br>`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`<br>`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`<br>`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`<br>`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`<br>`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`<br>`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`<br>`tanh(gate) * rmsnorm(x)` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`residual + tanh(gate) * rmsnorm(x)`<br>`ffn_norm1(x) * scale_mlp` | `python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`get_ada_values`<br>`scale_shift_table + timestep.reshape` | `python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py`<br>`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`<br>`residual_gate_add`<br>`residual + update * gate` | `python/sglang/jit_kernel/diffusion/residual_gate_add.py`<br>`python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh`<br>`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`<br>`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`<br>`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`<br>`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`<br>`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`<br>`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`<br>`tanh(gate) * rmsnorm(x)` | `python/sglang/kernels/ops/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`residual + tanh(gate) * rmsnorm(x)`<br>`ffn_norm1(x) * scale_mlp` | `python/sglang/kernels/ops/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py`<br>`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`<br>`get_ada_values`<br>`scale_shift_table + timestep.reshape` | `python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py`<br>`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`<br>`residual_gate_add`<br>`residual + update * gate` | `python/sglang/kernels/ops/diffusion/residual_gate_add.py`<br>`python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh`<br>`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`<br>`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`<br>`fp4 kv cache` | `PR #21889`<br>`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`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`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`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`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`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`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`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`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`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`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`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`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`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`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`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`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`<br>`trtllm_fp4_block_scale_moe`<br>`FlashInfer MoE` | `PR #22918`<br>`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`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`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`<br>`GLM-5`<br>`NSA`<br>`TileLang decode` | `PR #24125`<br>`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. |
.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
+11 -9
View File
@@ -452,7 +452,8 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec(
pattern="In-place QK RMSNorm",
candidate_path=(
"python/sglang/srt/models/utils.py" "<br>python/sglang/jit_kernel/norm.py"
"python/sglang/srt/models/utils.py"
"<br>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"
"<br>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" "<br>python/sglang/srt/models/utils.py"
"python/sglang/kernels/ops/attention/rope.py"
"<br>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"
"<br>python/sglang/jit_kernel/nvfp4.py"
"<br>python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py"
"<br>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"
"<br>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"
"<br>python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py"
"<br>python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py"
"<br>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"
"<br>python/sglang/jit_kernel/diffusion/residual_gate_add.py"
"<br>python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh"
"<br>python/sglang/kernels/ops/diffusion/residual_gate_add.py"
"<br>python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh"
"<br>python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py"
),
active_keywords=(
.claude/skills/write-sglang-test/SKILL.md
+3 -3
View File
@@ -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/<category>/` (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)
.github/CODEOWNERS
+4 -4
View File
@@ -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
.github/MAINTAINER.md
+1 -1
View File
@@ -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
.github/labeler.yml
+1 -1
View File
@@ -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:
.github/workflows/_pr-test-check-changes.yml
+1 -2
View File
@@ -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.
.github/workflows/nightly-test-nvidia.yml
+1 -1
View File
@@ -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
.github/workflows/pr-test-amd-rocm720.yml
+2 -2
View File
@@ -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"
.github/workflows/pr-test-amd.yml
+2 -2
View File
@@ -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"
.github/workflows/pr-test-npu.yml
+478 -478
View File
@@ -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
.github/workflows/pr-test.yml
+1 -1
View File
@@ -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() &&
.pre-commit-config.yaml
+1 -1
View File
@@ -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
3rdparty/amd/wheel/sglang/pyproject.toml
Vendored
+1 -1
View File
@@ -186,7 +186,7 @@ sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*",
"kernels/**/*",
]
[tool.setuptools.packages.find]
benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py
+3 -1
View File
@@ -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,
)
benchmark/bench_linear_attention/bench_gdn_qkv_split.py
+3 -1
View File
@@ -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,
benchmark/hicache/bench_hicache_write_back.py
+3 -3
View File
@@ -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:
benchmark/hicache/perf.py
+1 -1
View File
@@ -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
benchmark/kernels/deepseek/benchmark_cute_dsl_fp8_paged_mqa_logits.py
+2 -2
View File
@@ -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,
docs_new/docs/developer_guide/development_jit_kernel_guide.mdx
+12 -12
View File
@@ -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 <sgl_kernel/tensor.h> // 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.
python/pyproject.toml
+1 -1
View File
@@ -192,7 +192,7 @@ killall_sglang = "sglang.cli.killall:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*",
"kernels/**/*",
"multimodal_gen/apps/realtime_webui/**/*"
]
python/pyproject_cpu.toml
+1 -1
View File
@@ -123,7 +123,7 @@ sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
"kernels/**/*"
]
[tool.setuptools.packages.find]
python/pyproject_npu.toml
+1 -1
View File
@@ -122,7 +122,7 @@ sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
"kernels/**/*"
]
[tool.setuptools.packages.find]
python/pyproject_other.toml
+1 -1
View File
@@ -202,7 +202,7 @@ binding = "PyO3"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
"kernels/**/*"
]
[tool.setuptools.packages.find]
python/pyproject_xpu.toml
+1 -1
View File
@@ -130,7 +130,7 @@ sglang = "sglang.cli.main:main"
[tool.setuptools.package-data]
"sglang" = [
"srt/**/*",
"jit_kernel/**/*"
"kernels/**/*"
]
[tool.setuptools.packages.find]
python/sglang/jit_kernel/activation.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/dsv3_fused_a_gemm.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/dsv3_router_gemm.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/norm.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/per_tensor_quant_fp8.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/per_token_group_quant.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/per_token_group_quant_8bit_v2.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/set_mla_kv_buffer.py
-5
View File
@@ -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("__")})
python/sglang/jit_kernel/tests/kv_canary/__init__.py
View File
python/sglang/kernels/README.md
+3 -3
View File
@@ -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.
python/sglang/kernels/__init__.py
+2 -2
View File
@@ -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.
"""
python/sglang/kernels/fused_op.py
+1 -1
View File
@@ -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_<backend>`` methods.
"""
python/sglang/jit_kernel/.clang-format → python/sglang/kernels/jit/.clang-format
View File
python/sglang/kernels/jit/__init__.py
+1 -1
View File
@@ -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.
"""
python/sglang/jit_kernel/__main__.py → python/sglang/kernels/jit/__main__.py
View File
python/sglang/jit_kernel/benchmark/kv_canary/utils.py → python/sglang/kernels/jit/benchmark/kv_canary/utils.py
View File
python/sglang/jit_kernel/benchmark/marker.py → python/sglang/kernels/jit/benchmark/marker.py
View File
python/sglang/jit_kernel/benchmark/utils.py → python/sglang/kernels/jit/benchmark/utils.py
View File
python/sglang/jit_kernel/csrc/add_constant.cuh → python/sglang/kernels/jit/csrc/add_constant.cuh
View File
python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh → python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh
View File
python/sglang/jit_kernel/csrc/attention/fused_fp8_qkv_kv_cache.cuh → python/sglang/kernels/jit/csrc/attention/fused_fp8_qkv_kv_cache.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh → python/sglang/kernels/jit/csrc/deepseek_v32/indexer_k.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c128.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c128_online.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c128_online_v2.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c128_v2.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c4.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c4_v2.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/common.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/online_c128_mtp.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/paged_mqa_metadata.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/rope.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/rope.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/store.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/store.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh
View File
python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh → python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh → python/sglang/kernels/jit/csrc/diffusion/causal_conv3d_cat_pad.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh → python/sglang/kernels/jit/csrc/diffusion/ltx2_qknorm_split_rope.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/norm_scale_shift.cuh → python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh → python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh → python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh
View File
python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh → python/sglang/kernels/jit/csrc/diffusion/timestep_embedding.cuh
View File
python/sglang/jit_kernel/csrc/distributed/communicator.cuh → python/sglang/kernels/jit/csrc/distributed/communicator.cuh
View File
python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh → python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh
View File
python/sglang/jit_kernel/csrc/distributed/ipc.cuh → python/sglang/kernels/jit/csrc/distributed/ipc.cuh
View File
python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh → python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh
View File
python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh → python/sglang/kernels/jit/csrc/dsa/fused_store_index_cache.cuh
View File
python/sglang/jit_kernel/csrc/dsa/kpool_topk_transform.cuh → python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/activation.cuh → python/sglang/kernels/jit/csrc/elementwise/activation.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh → python/sglang/kernels/jit/csrc/elementwise/clamp_position.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh → python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh → python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh → python/sglang/kernels/jit/csrc/elementwise/fused_eh_norm.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh → python/sglang/kernels/jit/csrc/elementwise/fused_metadata_copy.cuh
+1 -1
View File
@@ -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,
python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh → python/sglang/kernels/jit/csrc/elementwise/fused_qknorm_rope.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh → python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh → python/sglang/kernels/jit/csrc/elementwise/pos_enc.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh → python/sglang/kernels/jit/csrc/elementwise/qknorm.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh → python/sglang/kernels/jit/csrc/elementwise/qknorm_across_heads.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh → python/sglang/kernels/jit/csrc/elementwise/resolve_future_token_ids.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh → python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/rmsnorm_hf.cuh → python/sglang/kernels/jit/csrc/elementwise/rmsnorm_hf.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/rope.cuh → python/sglang/kernels/jit/csrc/elementwise/rope.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/set_mla_kv_buffer.cuh → python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh
View File
python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py → python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py
View File
python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h → python/sglang/kernels/jit/csrc/fast-hadamard-transform/fast_hadamard_transform.h
View File

Some files were not shown because too many files have changed in this diff Show More