Skip to content

[Diffusion] Refactor diffusion JIT kernel test layout and narrow CI triggers #21385

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BBuf merged 4 commits into from
Mar 26, 2026
Merged

[Diffusion] Refactor diffusion JIT kernel test layout and narrow CI triggers #21385

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9c446a8
ud
BBuf Mar 25, 2026
53e8e7c
Merge branch 'main' into refactor/jit-kernel-diffusion-ci-scope
BBuf Mar 25, 2026
e7df703
Merge main into refactor/jit-kernel-diffusion-ci-scope
BBuf Mar 26, 2026
a9b0cda
Merge branch 'refactor/jit-kernel-diffusion-ci-scope' of github.com:s…
BBuf Mar 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions .claude/skills/write-sglang-test/SKILL.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,14 @@ description: Guide for writing SGLang CI/UT tests. Covers CustomTestCase, CI reg

1. **Always use `CustomTestCase`** — never raw `unittest.TestCase`. It ensures `tearDownClass` runs even when `setUpClass` fails, preventing resource leaks in CI.
2. **`tearDownClass` must be defensive** — use `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
3. **Place tests in `test/registered/<category>/`** — except JIT kernel tests and benchmarks, which live in `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/`
3. **Place tests in `test/registered/<category>/`** — except JIT kernel tests and benchmarks, which live in `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/` (nested subfolders are allowed)
4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
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 exception:
- If the task is adding or updating code under `python/sglang/jit_kernel/`, prefer the `add-jit-kernel` skill first.
- JIT kernel correctness tests use `python/sglang/jit_kernel/tests/test_*.py`.
- JIT kernel benchmarks use `python/sglang/jit_kernel/benchmark/bench_*.py`.
- JIT kernel correctness tests use `python/sglang/jit_kernel/tests/**/test_*.py`.
- JIT kernel benchmarks use `python/sglang/jit_kernel/benchmark/**/bench_*.py`.
- Those files are still executed by `test/run_suite.py`, but through dedicated kernel suites rather than `test/registered/`.

---
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/pr-test-amd-rocm720.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,8 @@ jobs:
- "python/sglang/multimodal_gen/**"
- "python/sglang/cli/**"
- "python/sglang/jit_kernel/diffusion/**"
- "python/sglang/jit_kernel/tests/diffusion/**"
- "python/sglang/jit_kernel/benchmark/diffusion/**"
- "python/pyproject_rocm.toml"
- "python/pyproject_other.toml"

Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/pr-test-amd.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,8 @@ jobs:
- "python/sglang/multimodal_gen/**"
- "python/sglang/cli/**"
- "python/sglang/jit_kernel/diffusion/**"
- "python/sglang/jit_kernel/tests/diffusion/**"
- "python/sglang/jit_kernel/benchmark/diffusion/**"
- "python/pyproject_rocm.toml"
- "python/pyproject_other.toml"

Expand Down
8 changes: 5 additions & 3 deletions .github/workflows/pr-test.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,9 @@ jobs:
- ".github/workflows/pr-test-multimodal-gen.yml"
- "python/pyproject.toml"
- "python/sglang/multimodal_gen/**/*.!(md|ipynb)"
- "python/sglang/jit_kernel/**"
- "python/sglang/jit_kernel/diffusion/**"
- "python/sglang/jit_kernel/tests/diffusion/**"
- "python/sglang/jit_kernel/benchmark/diffusion/**"
- "python/sglang/cli/**"
jit_kernel:
- ".github/workflows/pr-test.yml"
Expand Down Expand Up @@ -197,8 +199,8 @@ jobs:
echo "jit_kernel=false" >> $GITHUB_OUTPUT
fi

# Check for multimodal_gen changes
if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/multimodal_gen/|python/sglang/cli/|python/pyproject\.toml|\.github/workflows/pr-test\.yml|\.github/workflows/pr-test-multimodal-gen\.yml)"; then
# Check for multimodal_gen changes, including diffusion-specific jit_kernel coverage
if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/multimodal_gen/|python/sglang/cli/|python/sglang/jit_kernel/diffusion/|python/sglang/jit_kernel/tests/diffusion/|python/sglang/jit_kernel/benchmark/diffusion/|python/pyproject\.toml|\.github/workflows/pr-test\.yml|\.github/workflows/pr-test-multimodal-gen\.yml)"; then
echo "multimodal_gen=true" >> $GITHUB_OUTPUT
echo "Detected multimodal_gen changes"
else
Expand Down
0 ...benchmark/bench_fused_norm_scale_shift.py → ...diffusion/bench_fused_norm_scale_shift.py
View file
Open in desktop
File renamed without changes.
0 .../jit_kernel/benchmark/bench_norm_impls.py → ...l/benchmark/diffusion/bench_norm_impls.py
View file
Open in desktop
File renamed without changes.
0 .../benchmark/bench_qwen_image_modulation.py → .../diffusion/bench_qwen_image_modulation.py
View file
Open in desktop
File renamed without changes.
0 ...rnel/tests/test_fused_norm_scale_shift.py → .../diffusion/test_fused_norm_scale_shift.py
View file
Open in desktop
File renamed without changes.
0 ...ernel/tests/test_qwen_image_modulation.py → ...s/diffusion/test_qwen_image_modulation.py
View file
Open in desktop
File renamed without changes.
8 changes: 4 additions & 4 deletions test/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ Here is an illustration


## Folder organization
- `registered`: The registered test files. They are run in CI. Most tests should live in this folder. The main exception is JIT kernel coverage, which lives under `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/`.
- `registered`: The registered test files. They are run in CI. Most tests should live in this folder. The main exception is JIT kernel coverage, which lives under `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/` (including nested subfolders such as `diffusion/`).
- `manual`: Test files that CI does not run; you run them manually. Typically, these are temporary tests, deprecated tests, or tests that are not suitable for CI—such as those that take too long or require special setup. We would still like to keep some files here for anyone who wants to run them locally.
- `run_suite.py`: The launch script to run a test suite. It scans `test/registered/` and also the JIT kernel test / benchmark directories.
- Other: utility scripts and metadata folders. The `srt` folder holds our legacy CI setup and should be deprecated as soon as possible.
Expand Down Expand Up @@ -142,7 +142,7 @@ python test/run_suite.py --hw cuda --suite stage-b-test-1-gpu-small \
## CI Registry System

CI-discovered tests use a registry-based CI system for flexible backend and schedule configuration.
This includes files under `test/registered/` and, for JIT kernels, files under `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/`.
This includes files under `test/registered/` and, for JIT kernels, files under `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/`, recursively including nested subfolders.
For every CI-discovered file you add, you need to register it in a suite and provide an estimated execution time in seconds.

### Registration Functions
Expand Down Expand Up @@ -180,8 +180,8 @@ register_cuda_ci(est_time=80, suite="stage-b-test-1-gpu-small", disabled="flaky

JIT kernel files are discovered by `test/run_suite.py`, but they do not live under `test/registered/`:

- Correctness tests: `python/sglang/jit_kernel/tests/test_*.py`
- Benchmarks: `python/sglang/jit_kernel/benchmark/bench_*.py`
- Correctness tests: `python/sglang/jit_kernel/tests/**/test_*.py`
- Benchmarks: `python/sglang/jit_kernel/benchmark/**/bench_*.py`

Use dedicated kernel suites:

Expand Down
8 changes: 6 additions & 2 deletions test/run_suite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,8 +192,12 @@ def run_a_suite(args):

# JIT kernel tests and benchmarks (live alongside kernel source)
jit_kernel_dir = os.path.join(repo_root, "python", "sglang", "jit_kernel")
files += glob.glob(os.path.join(jit_kernel_dir, "tests", "test_*.py"))
files += glob.glob(os.path.join(jit_kernel_dir, "benchmark", "bench_*.py"))
files += glob.glob(
os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True
)
files += glob.glob(
os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True
)
Comment on lines +195 to +200
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To reduce code duplication, you could use a loop to handle both the 'tests' and 'benchmark' directories. This makes the code more concise and easier to maintain if more directories are added in the future.

Suggested change
files += glob.glob(
os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True
)
files += glob.glob(
os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True
)
for sub_dir, prefix in [("tests", "test_"), ("benchmark", "bench_")]:
path_pattern = os.path.join(jit_kernel_dir, sub_dir, "**", f"{prefix}*.py")
files += glob.glob(path_pattern, recursive=True)


# Strict: all discovered files must have proper registration
sanity_check = True
Expand Down
Loading