Summary

This PR keeps the Qwen select01 Triton kernel version that showed a stable end-to-end win in Qwen-Image denoise.

The final change set:

• switches modulation loads to pointer-select so each row only loads the chosen scale/shift/gate branch
• pins both select01 launches to num_warps=4, num_stages=4
• drops later scalar-base / 8w1s / residual-only experiments from the active code path because they did not produce stable model-level gains

Made with Codex(AKO4ALL framework and SGLang Diffusion SKILL).

Implementation

The optimized kernels are:

• fuse_layernorm_scale_shift_gate_select01_kernel
• fuse_residual_layernorm_scale_shift_gate_select01_kernel

Main code changes:

• build branch-specific pointer tensors for scale0/1, shift0/1, and gate0/1
• select pointers with tl.where(idx, ...)
• issue only one load for each modulation tensor on the chosen branch
• keep the validated 4w4s launch config for both kernels

Validation

Correctness:

• python -m py_compile python/sglang/jit_kernel/diffusion/triton/scale_shift.py
• pytest -q python/sglang/jit_kernel/tests/test_qwen_image_modulation.py -q

Performance:

Nsight Compute

A representative ncu check on the layernorm select01 kernel at (2,2048,3072) shows:

• gpu__time_duration.avg: 35.744 us -> 28.704 us
• launch__registers_per_thread: 96 -> 72

The optimized kernel reduces single-launch latency from the baseline by 19.7% (28.70 us here), while also increasing L2 and DRAM throughput. This indicates that the win is not just a launch-parameter artifact: the kernel is doing less wasted work and using the memory hierarchy more effectively.

Executed and issued instruction counts both drop by about 18.5%, which is consistent with the kernel rewrite: the pointer-select path avoids loading and computing both modulation branches before selecting one.

This kernel is register-limited, and lowering register pressure improves occupancy materially. In this tuned version, registers per thread drop from 96 to 72, which raises both theoretical and achieved occupancy and improves latency hiding.

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.